From 783f8a2f6cb9a932c3df9ac74d070f5e7e2142a5 Mon Sep 17 00:00:00 2001 From: cc Date: Wed, 30 Sep 2026 22:58:35 -0700 Subject: [PATCH 01/11] feat: assemble complete CDR runtime and minimum storage compatibility --- .pre-commit-config.yaml | 6 + AGENTS.md | 2 +- Cargo.lock | 147 +-- README.md | 24 +- data/.lfs/astar_corner_min_cost.png.tar.gz | 4 +- data/.lfs/astar_min_cost.png.tar.gz | 4 +- data/.lfs/gradient_simple.png.tar.gz | 4 +- data/.lfs/make_path_mask_full.png.tar.gz | 4 +- .../.lfs/make_path_mask_two_meters.png.tar.gz | 4 +- data/.lfs/resample_path_simple.png.tar.gz | 4 +- data/.lfs/resample_path_smooth.png.tar.gz | 4 +- data/.lfs/unitree_go2_detection_cdr.tar.gz | 3 + .../visualize_occupancy_rainbow.png.tar.gz | 4 +- .../.lfs/visualize_occupancy_turbo.png.tar.gz | 4 +- dimos/agents/mcp/mcp_client.py | 5 + dimos/agents/mcp/mcp_server.py | 12 + dimos/agents/mcp/test_generated_image.py | 56 + dimos/agents/mcp/test_mcp_client.py | 5 +- dimos/agents/skills/navigation.py | 35 +- dimos/agents/skills/observe_skill.py | 3 +- .../skills/test_generated_navigation.py | 70 ++ .../skills/test_generated_person_follow.py | 81 ++ dimos/agents/skills/test_navigation.py | 4 +- dimos/agents/skills/test_observe_skill.py | 18 +- .../skills/test_unitree_skill_container.py | 17 +- .../skills/visual_servoing/bbox_navigation.py | 17 +- .../visual_servoing/detection_navigation.py | 44 +- dimos/agents/skills/visual_servoing/query.py | 3 +- .../visual_servoing/visual_servoing_2d.py | 18 +- dimos/agents/test_generated_vlm.py | 80 ++ dimos/agents/testing/vlm_stream_tester.py | 2 +- dimos/agents/vlm_agent.py | 10 +- dimos/agents/vlm_agent_spec.py | 3 +- dimos/cli/bake/test_bake_e2e.py | 36 +- dimos/cli/bake/test_cli.py | 2 +- dimos/cli/bake/test_codegen.py | 8 +- dimos/cli/bake/test_discovery.py | 10 +- dimos/cli/bake/test_graph.py | 14 +- dimos/cli/hardware/g1.py | 2 +- dimos/cli/spy/test_spy.py | 31 +- dimos/cli/test_topic.py | 20 +- dimos/cli/topic.py | 78 +- dimos/control/_control_test_helpers.py | 4 +- dimos/control/benchmarking/benchmark.py | 68 +- dimos/control/benchmarking/paths.py | 59 +- dimos/control/benchmarking/score.py | 21 +- dimos/control/benchmarking/scoring.py | 17 +- dimos/control/benchmarking/test_benchmark.py | 34 +- .../control/benchmarking/velocity_profile.py | 6 +- dimos/control/examples/cartesian_ik_jogger.py | 23 +- dimos/control/path_following_coordinator.py | 5 +- dimos/control/task.py | 8 +- .../cartesian_ik_task/cartesian_ik_task.py | 11 +- .../test_cartesian_ik_task.py | 16 +- .../tasks/eef_twist_task/eef_twist_task.py | 23 +- .../eef_twist_task/test_eef_twist_task.py | 22 +- .../g1_groot_wbc_task/g1_groot_wbc_task.py | 3 +- .../test_g1_groot_wbc_task.py | 7 +- .../tasks/gripper_task/gripper_task.py | 2 +- .../tasks/gripper_task/test_gripper_task.py | 2 +- .../holonomic_pose_follower_task.py | 19 +- .../progress_reference.py | 7 +- .../test_holonomic_pose_follower_task.py | 26 +- .../test_progress_reference.py | 14 +- .../test_path_follower_task.py | 16 +- .../planar_base_trajectory_task.py | 37 +- .../test_planar_base_trajectory_task.py | 29 +- dimos/control/tasks/pose_target_ik.py | 25 +- .../rpp_path_follower_task.py | 29 +- .../test_rpp_path_follower_task.py | 45 +- .../tasks/teleop_ik_task/teleop_ik_task.py | 36 +- .../teleop_ik_task/test_teleop_ik_task.py | 37 +- dimos/control/tasks/test_pose_target_ik.py | 19 +- .../trajectory_task/test_trajectory_task.py | 61 +- .../tasks/trajectory_task/trajectory_task.py | 62 +- dimos/control/tasks/velocity_profiler.py | 5 +- .../tasks/velocity_task/test_velocity_task.py | 3 +- .../tasks/velocity_task/velocity_task.py | 2 +- dimos/control/teleop_coordinator.py | 9 +- dimos/control/test_control.py | 189 ++- dimos/control/test_coordinator_commands.py | 9 +- dimos/control/test_coordinator_routing.py | 57 +- .../test_path_following_coordinator.py | 2 +- dimos/control/test_per_robot_joint_states.py | 14 +- dimos/core/coordination/module_coordinator.py | 2 +- .../coordination/test_module_coordinator.py | 19 +- dimos/core/coordination/test_worker.py | 2 +- dimos/core/test_baked_host.py | 4 +- dimos/core/test_cdr_transport.py | 104 ++ dimos/core/test_core.py | 17 +- dimos/core/test_stream.py | 28 +- dimos/core/test_zenoh_transport.py | 19 +- dimos/core/testing.py | 36 +- dimos/core/transport_factory.py | 7 +- dimos/e2e_tests/cdr_replay_fixture.py | 62 + dimos/e2e_tests/test_cdr_spy.py | 51 + dimos/e2e_tests/test_cockpit_browser.py | 12 +- .../e2e_tests/test_custom_channel_browser.py | 26 +- dimos/e2e_tests/test_lcm_channel_browser.py | 35 +- .../test_manipulation_planning_groups.py | 4 +- dimos/e2e_tests/test_map_click_browser.py | 21 +- dimos/e2e_tests/test_sdk_browser.py | 11 +- dimos/evals/agents/lib/plain_recording.py | 15 +- dimos/evals/agents/question_answer.py | 30 +- dimos/evals/agents/test_no_dimos.py | 22 +- dimos/evals/agents/test_pi.py | 2 +- dimos/evals/conftest.py | 10 +- dimos/evals/environments/dimsim.py | 3 +- dimos/evals/environments/sim.py | 5 +- dimos/evals/environments/test_dimsim.py | 8 +- dimos/evals/suites/dimsim_house.py | 13 +- .../evals/suites/dimsim_pointcloud_mapping.py | 2 +- dimos/evals/test_evals.py | 22 +- dimos/evals/test_mem2_wiring.py | 34 +- dimos/evals/vqa/author.py | 3 +- dimos/evals/vqa/contracts.py | 3 +- dimos/evals/vqa/families.py | 3 +- dimos/evals/vqa/generate.py | 10 +- dimos/evals/vqa/pointcloud_frame.py | 71 +- dimos/evals/vqa/primitives/edge_tam.py | 3 +- dimos/evals/vqa/primitives/range.py | 10 +- dimos/evals/vqa/primitives/test_edge_tam.py | 110 +- dimos/evals/vqa/test_contracts.py | 94 +- dimos/evals/vqa/test_pointcloud_frame.py | 152 ++- .../isolated_python/example/contract.py | 3 +- .../isolated_python/test_end_to_end.py | 5 +- .../isolated_python/test_generated_example.py | 42 + dimos/experimental/memory/rust/Cargo.toml | 3 +- .../experimental/memory/rust/src/store/mod.rs | 4 +- .../memory/rust/src/store/sqlite.rs | 8 +- .../experimental/memory/rust_cli_recorder.py | 13 +- .../memory/test_rust_cli_recorder.py | 14 +- .../robot/bosdyn/spot/effectors/high_level.py | 82 +- .../bosdyn/spot/effectors/test_high_level.py | 65 + .../experimental/robot/bosdyn/spot/replay.py | 29 +- .../robot/bosdyn/spot/test_utils.py | 144 ++- dimos/experimental/robot/bosdyn/spot/utils.py | 92 +- dimos/experimental/security_demo/conftest.py | 16 +- .../security_demo/depth_estimator.py | 12 +- .../security_demo/test_security_module.py | 15 +- dimos/experimental/world_belief/absence.py | 17 +- dimos/experimental/world_belief/recall.py | 12 +- dimos/experimental/world_belief/scene_scan.py | 66 +- .../experimental/world_belief/test_absence.py | 63 + .../world_belief/test_generated_recall.py | 90 ++ .../experimental/world_belief/world_belief.py | 36 +- .../world_belief/worldbelief_module.py | 19 +- .../world_belief/worldbelief_recorder.py | 11 +- .../world_belief/xarm6_blueprint.py | 21 +- .../hardware/drive_trains/flowbase/config.py | 8 +- .../drive_trains/transport/adapter.py | 12 +- .../drive_trains/transport/test_adapter.py | 63 + dimos/hardware/manipulators/spec.py | 11 - .../camera/gstreamer/gstreamer_camera.py | 14 +- .../gstreamer/gstreamer_camera_test_script.py | 18 +- .../camera/gstreamer/test_generated_frames.py | 84 ++ dimos/hardware/sensors/camera/module.py | 48 +- .../sensors/camera/realsense/blueprints.py | 20 +- .../sensors/camera/realsense/camera.py | 9 +- .../sensors/camera/realsense/rust/Cargo.lock | 61 +- .../sensors/camera/realsense/rust/Cargo.toml | 26 +- .../sensors/camera/realsense/rust/flake.nix | 2 +- .../sensors/camera/realsense/rust/src/main.rs | 127 +- dimos/hardware/sensors/camera/spec.py | 7 +- dimos/hardware/sensors/camera/test_module.py | 59 + dimos/hardware/sensors/camera/test_webcam.py | 36 +- dimos/hardware/sensors/camera/webcam.py | 38 +- dimos/hardware/sensors/camera/zed/camera.py | 191 ++- dimos/hardware/sensors/camera/zed/compat.py | 2 +- .../sensors/camera/zed/test_generated_tf.py | 110 ++ .../lidar/common/point_cloud_utils.hpp | 47 +- .../sensors/lidar/fastlio2/cpp/CMakeLists.txt | 9 +- .../sensors/lidar/fastlio2/cpp/flake.lock | 18 - .../sensors/lidar/fastlio2/cpp/flake.nix | 8 +- .../sensors/lidar/fastlio2/cpp/main.cpp | 18 +- .../lidar/fastlio2/tools/pcap_to_db.py | 7 +- dimos/hardware/sensors/lidar/livox/module.py | 3 +- .../sensors/lidar/livox/rust/Cargo.toml | 6 +- .../sensors/lidar/livox/rust/src/module.rs | 35 +- .../lidar/pointlio/scripts/pcap_to_db.py | 7 +- .../sensors/lidar/test_generated_recorders.py | 93 ++ .../hardware/whole_body/transport/adapter.py | 8 +- .../whole_body/transport/test_adapter.py | 54 + .../dataprep/test_generated_images.py | 51 + dimos/imitation/policy/lerobot/module.py | 10 +- dimos/manipulation/conftest.py | 8 +- .../control/coordinator_client.py | 11 +- .../control/servo_control/README.md | 15 +- .../cartesian_motion_controller.py | 140 +-- .../test_generated_legacy_controllers.py | 106 ++ .../joint_trajectory_controller.py | 36 +- dimos/manipulation/execution_manager.py | 29 +- dimos/manipulation/grasping/grasp_gen_spec.py | 7 +- .../grasping/grasp_gen_x/module.py | 4 +- .../grasping/grasp_gen_x/test_module.py | 19 +- dimos/manipulation/grasping/grasping.py | 12 +- .../manipulation/grasping/heuristic_grasp.py | 30 +- .../grasping/test_heuristic_grasp.py | 63 +- dimos/manipulation/manipulation_spec.py | 24 +- dimos/manipulation/planning/README.md | 10 +- .../planning/groups/test_planning_groups.py | 28 +- dimos/manipulation/planning/groups/utils.py | 8 +- .../kinematics/drake_optimization_ik.py | 24 +- .../planning/kinematics/jacobian_ik.py | 50 +- .../planning/kinematics/pink_ik.py | 11 +- .../planning/kinematics/pink_solver.py | 11 +- .../planning/kinematics/pinocchio_ik.py | 10 +- .../test_drake_optimization_ik_selection.py | 21 +- .../kinematics/test_jacobian_ik_selection.py | 18 +- .../planning/kinematics/test_pink_ik.py | 194 +-- .../manipulation/planning/kinematics/utils.py | 4 +- .../monitor/test_robot_state_monitor.py | 78 ++ .../monitor/world_obstacle_monitor.py | 36 +- .../planning/planners/roboplan_planner.py | 34 +- .../planning/planners/rrt_planner.py | 10 +- .../planning/planners/selected_joint_space.py | 2 +- .../planners/test_rrt_planner_selection.py | 65 +- .../manipulation/planning/spec/joint_space.py | 4 +- dimos/manipulation/planning/spec/models.py | 11 +- dimos/manipulation/planning/spec/protocols.py | 8 +- .../planning/spec/test_joint_space.py | 2 +- .../joint_trajectory_generator.py | 28 +- .../trajectory_generator/parametrizer.py | 45 +- .../roboplan_toppra_parametrizer.py | 11 +- .../simple_parametrizer.py | 7 +- .../trajectory_generator/test_parametrizer.py | 25 +- .../test_roboplan_toppra_parametrizer.py | 5 +- .../test_simple_parametrizer.py | 7 +- .../manipulation/planning/utils/path_utils.py | 3 +- .../planning/world/drake_world.py | 46 +- .../planning/world/roboplan_model.py | 7 +- .../planning/world/roboplan_world.py | 37 +- .../world/test_drake_world_planning_groups.py | 74 +- dimos/manipulation/sdk.py | 23 +- dimos/manipulation/test_execution_manager.py | 58 +- .../test_generated_plan_materialization.py | 43 +- .../manipulation/test_manipulation_module.py | 21 +- .../test_manipulation_monitor_preview.py | 59 +- .../test_manipulation_primitives.py | 22 +- dimos/manipulation/test_manipulation_spec.py | 20 +- dimos/manipulation/test_plan_execution.py | 32 +- dimos/manipulation/test_planning_factory.py | 9 +- dimos/manipulation/test_roboplan.py | 198 ++- dimos/manipulation/test_sdk.py | 52 +- dimos/manipulation/visualization/operator.py | 32 +- .../visualization/test_factory.py | 24 +- .../visualization/test_operator.py | 33 +- dimos/manipulation/visualization/viser/gui.py | 87 +- .../manipulation/visualization/viser/scene.py | 63 +- .../manipulation/visualization/viser/state.py | 15 +- .../visualization/viser/test_gui.py | 40 +- .../viser/test_scene_grasp_proposals.py | 23 +- .../viser/test_scene_obstacles.py | 7 +- .../viser/test_viser_visualization.py | 18 +- .../viser/test_visualizer_lifecycle.py | 23 +- .../visualization/viser/visualizer.py | 22 +- dimos/mapping/cli/pose_fill.py | 20 +- dimos/mapping/cli/replay.py | 38 +- dimos/mapping/cli/replay_marker.py | 18 +- dimos/mapping/cli/test_cli.py | 21 +- dimos/mapping/cli/test_pgo_accumulate.py | 48 + dimos/mapping/cli/test_pose_fill.py | 59 +- dimos/mapping/cli/test_view.py | 40 + dimos/mapping/cli/view.py | 15 +- dimos/mapping/costmapper.py | 28 +- dimos/mapping/dim_slam/dim_slam.py | 9 +- dimos/mapping/dim_slam/rust/Cargo.lock | 61 +- dimos/mapping/dim_slam/rust/Cargo.toml | 2 +- dimos/mapping/dim_slam/rust/flake.nix | 4 + dimos/mapping/dim_slam/rust/src/main.rs | 57 +- .../mapping/dim_slam/rust/src/msg_convert.rs | 155 ++- dimos/mapping/occupancy/conftest.py | 14 +- dimos/mapping/occupancy/extrude_occupancy.py | 17 +- dimos/mapping/occupancy/gradient.py | 45 +- dimos/mapping/occupancy/inflation.py | 31 +- dimos/mapping/occupancy/operations.py | 33 +- dimos/mapping/occupancy/path_map.py | 3 +- dimos/mapping/occupancy/path_mask.py | 26 +- dimos/mapping/occupancy/path_resampling.py | 80 +- .../test_generated_visualizations.py | 86 ++ dimos/mapping/occupancy/test_gradient.py | 7 +- dimos/mapping/occupancy/test_inflation.py | 3 +- dimos/mapping/occupancy/test_operations.py | 14 +- dimos/mapping/occupancy/test_path_map.py | 3 +- .../mapping/occupancy/test_visualizations.py | 3 +- dimos/mapping/occupancy/visualizations.py | 62 +- dimos/mapping/occupancy/visualize_path.py | 33 +- dimos/mapping/odometry_hist.py | 49 +- dimos/mapping/osm/current_location_map.py | 17 +- dimos/mapping/osm/osm.py | 5 +- .../mapping/osm/test_current_location_map.py | 56 + dimos/mapping/pointclouds/demo.py | 19 +- .../pointclouds/test_generated_occupancy.py | 96 ++ dimos/mapping/pointclouds/test_occupancy.py | 49 +- .../pointclouds/tool_occupancy_speed.py | 5 +- .../mapping/ray_tracing/utils/raytrace_rrd.py | 22 +- dimos/mapping/relocalization/blueprints.py | 10 +- dimos/mapping/relocalization/go2/module.py | 15 +- dimos/mapping/relocalization/lidar/module.py | 20 +- dimos/mapping/relocalization/test_module.py | 72 +- dimos/mapping/test_generated_costmapper.py | 59 + dimos/mapping/test_generated_odometry_hist.py | 63 + dimos/mapping/test_generated_voxels.py | 76 ++ dimos/mapping/test_voxels.py | 169 +-- dimos/mapping/tool_voxels.py | 10 +- dimos/mapping/voxels/impl/o3d.py | 31 +- dimos/mapping/voxels/module.py | 3 +- dimos/mapping/voxels/test_grid.py | 48 + dimos/memory/architecture.md | 2 +- dimos/memory/cli/dataset.py | 6 +- dimos/memory/cli/render.py | 81 +- dimos/memory/cli/test_render.py | 43 +- dimos/memory/codecs/{lcm.py => cdr.py} | 14 +- dimos/memory/codecs/jpeg.py | 39 - dimos/memory/demo_data.py | 65 + dimos/memory/recording_policy.py | 24 + dimos/memory/replay_module.py | 14 +- dimos/memory/store/test_mcap.py | 231 ++-- dimos/memory/stream.py | 21 +- dimos/memory/test_replay_module.py | 43 +- dimos/memory/test_voxel_map.py | 24 +- dimos/memory/tf.py | 10 +- dimos/memory/tool_e2e.py | 91 +- dimos/memory/tool_visualizer.py | 36 +- dimos/memory/tool_voxel_map.py | 6 +- dimos/memory/transform.py | 7 +- dimos/memory/type/filter.py | 11 +- dimos/memory/type/observation.py | 27 +- dimos/memory/type/test_generated_query.py | 54 + dimos/memory/vis/space/elements.py | 20 +- dimos/memory/vis/space/geometry.py | 32 + dimos/memory/vis/space/rerun.py | 137 +- dimos/memory/vis/space/space.py | 12 +- dimos/memory/vis/space/svg.py | 68 +- .../vis/space/test_generated_geometry.py | 155 +++ dimos/memory/vis/space/test_space.py | 163 ++- dimos/memory/vis/utils.py | 7 +- dimos/models/embedding/base.py | 4 +- dimos/models/embedding/clip.py | 8 +- dimos/models/embedding/dino.py | 8 +- dimos/models/embedding/siglip.py | 9 +- dimos/models/embedding/treid.py | 8 +- dimos/models/qwen/video_query.py | 4 +- dimos/models/segmentation/edge_tam.py | 11 +- .../test_generated_edge_tam_images.py | 50 + dimos/models/segmentation/test_yoloe.py | 20 +- dimos/models/segmentation/yoloe.py | 5 +- dimos/models/vl/README.md | 10 +- dimos/models/vl/base.py | 30 +- dimos/models/vl/florence.py | 7 +- dimos/models/vl/moondream.py | 14 +- dimos/models/vl/moondream_hosted.py | 14 +- dimos/models/vl/openai.py | 22 +- dimos/models/vl/qwen.py | 24 +- dimos/models/vl/test_base.py | 5 +- dimos/models/vl/test_generated_images.py | 85 ++ dimos/models/vl/tool_base.py | 7 +- dimos/models/vl/tool_vlm.py | 17 +- dimos/msgs/camera_info.py | 169 +++ .../foxglove_msgs/test_CompressedVideo.py | 29 +- dimos/msgs/geometry.py | 374 ++++++ dimos/msgs/geometry_msgs/test_PointStamped.py | 70 +- dimos/msgs/geometry_msgs/test_Pose.py | 644 ++++------ dimos/msgs/geometry_msgs/test_PoseStamped.py | 71 +- .../geometry_msgs/test_PoseWithCovariance.py | 408 ++---- .../test_PoseWithCovarianceStamped.py | 307 ++--- dimos/msgs/geometry_msgs/test_Quaternion.py | 290 ++--- dimos/msgs/geometry_msgs/test_Transform.py | 440 +++---- dimos/msgs/geometry_msgs/test_Twist.py | 212 ++-- dimos/msgs/geometry_msgs/test_TwistStamped.py | 67 +- .../geometry_msgs/test_TwistWithCovariance.py | 443 ++----- .../test_TwistWithCovarianceStamped.py | 329 ++--- dimos/msgs/geometry_msgs/test_Vector3.py | 404 +++--- .../msgs/geometry_msgs/test_WrenchStamped.py | 168 ++- dimos/msgs/geometry_msgs/tool_publish.py | 12 +- dimos/msgs/helpers.py | 45 +- dimos/msgs/image.py | 317 +++++ .../manipulation_msgs/test_grasp_candidate.py | 30 +- dimos/msgs/nav_msgs/test_OccupancyGrid.py | 504 +++----- dimos/msgs/nav_msgs/test_Odometry.py | 254 ++-- dimos/msgs/nav_msgs/test_Path.py | 238 ++-- dimos/msgs/occupancy.py | 132 ++ dimos/msgs/pointcloud.py | 348 +++++ dimos/msgs/protocol.py | 13 +- dimos/msgs/sensor_msgs/test_CameraInfo.py | 119 +- .../msgs/sensor_msgs/test_CompressedImage.py | 174 ++- dimos/msgs/sensor_msgs/test_Joy.py | 164 +-- dimos/msgs/sensor_msgs/test_image.py | 117 +- dimos/msgs/sensor_msgs/tool_image.py | 13 +- dimos/msgs/std_msgs/test_header.py | 116 +- dimos/msgs/test_camera_info.py | 127 ++ dimos/msgs/test_generated_projection.py | 98 ++ dimos/msgs/test_generated_views.py | 232 ++++ dimos/msgs/test_geometry.py | 244 ++++ dimos/msgs/test_helpers.py | 19 +- dimos/msgs/test_msg_construction.py | 962 ++++---------- dimos/msgs/test_occupancy.py | 146 +++ dimos/msgs/test_pointcloud.py | 384 ++++++ dimos/msgs/test_pointcloud_colors.py | 65 + dimos/msgs/test_time.py | 81 ++ dimos/msgs/test_trajectory.py | 75 ++ dimos/msgs/tf2_msgs/test_TFMessage.py | 181 ++- dimos/msgs/tf2_msgs/test_TFMessage_lcmpub.py | 51 +- dimos/msgs/time.py | 70 ++ dimos/msgs/trajectory.py | 79 ++ .../vision_msgs/test_detection3d_rerun.py | 41 +- dimos/navigation/base.py | 2 +- .../holonomic_tc/test_dan_holonomic_tc.py | 34 +- .../wavefront_frontier_goal_selector.py | 139 +- .../patrolling/routers/patrol_router.py | 4 +- .../patrolling/routers/visitation_history.py | 20 +- .../experimental/patrolling/utilities.py | 12 +- .../global_planner/mls_planner/start_relay.py | 4 +- dimos/navigation/go2/loop_closure/pgo_auto.py | 195 +-- .../go2/loop_closure/utils/markers_rrd.py | 20 +- .../go2/replanning_a_star/controllers.py | 17 +- .../go2/replanning_a_star/global_planner.py | 87 +- .../go2/replanning_a_star/goal_validator.py | 62 +- .../go2/replanning_a_star/local_planner.py | 39 +- .../go2/replanning_a_star/min_cost_astar.py | 107 +- .../go2/replanning_a_star/module.py | 21 +- .../go2/replanning_a_star/module_spec.py | 3 +- .../go2/replanning_a_star/navigation_map.py | 3 +- .../go2/replanning_a_star/path_clearance.py | 24 +- .../go2/replanning_a_star/path_distancer.py | 5 +- .../go2/replanning_a_star/position_tracker.py | 5 +- .../go2/replanning_a_star/replan_limiter.py | 10 +- .../navigation/go2/replanning_a_star/spec.py | 3 +- .../movement_manager/movement_manager.py | 43 +- .../movement_manager/test_movement_manager.py | 29 +- .../test_generated_bbox_navigation.py | 73 ++ dimos/perception/detection/conftest.py | 74 +- dimos/perception/detection/detectors/base.py | 3 +- .../detection/detectors/conftest.py | 4 +- dimos/perception/detection/detectors/omdet.py | 10 +- .../detection/detectors/person/yolo.py | 6 +- .../detectors/test_bbox_detectors.py | 7 +- dimos/perception/detection/detectors/yolo.py | 6 +- dimos/perception/detection/detectors/yoloe.py | 5 +- dimos/perception/detection/module2D.py | 33 +- dimos/perception/detection/module3D.py | 62 +- dimos/perception/detection/person_tracker.py | 63 +- dimos/perception/detection/project.py | 41 +- dimos/perception/detection/reid/module.py | 33 +- .../detection/reid/test_generated_module.py | 91 ++ .../detection/type/detection2d/base.py | 4 +- .../detection/type/detection2d/bbox.py | 42 +- .../type/detection2d/imageDetections2D.py | 5 +- .../detection/type/detection2d/person.py | 5 +- .../detection/type/detection2d/point.py | 28 +- .../detection/type/detection2d/seg.py | 8 +- .../detection/type/detection2d/test_bbox.py | 4 +- .../type/detection2d/test_generated_bbox.py | 101 ++ .../detection/type/detection2d/test_person.py | 4 +- .../detection/type/detection3d/base.py | 11 +- .../detection/type/detection3d/bbox.py | 53 +- .../type/detection3d/imageDetections3D.py | 7 +- .../type/detection3d/imageDetections3DPC.py | 10 +- .../detection/type/detection3d/marker.py | 11 +- .../type/detection3d/pointcloud_filters.py | 69 +- .../detection3d/test_generated_filters.py | 68 + .../detection3d/test_generated_projection.py | 97 ++ .../detection3d/test_imageDetections3D.py | 89 +- .../detection/type/detection3d/test_marker.py | 31 +- .../detection/type/imageDetections.py | 21 +- dimos/perception/experimental/moduleDB.py | 73 +- dimos/perception/experimental/object.py | 203 +-- dimos/perception/experimental/objectDB.py | 38 +- .../experimental/object_scene_registration.py | 76 +- .../object_scene_registration_spec.py | 5 +- .../perception/experimental/object_tracker.py | 111 +- .../experimental/object_tracker_2d.py | 27 +- .../experimental/object_tracker_3d.py | 92 +- .../experimental/spatial_perception.py | 47 +- .../temporal_memory/clip_filter.py | 7 +- .../temporal_memory/entity_graph_db.py | 3 +- .../frame_window_accumulator.py | 9 +- .../temporal_memory/temporal_utils/helpers.py | 8 +- .../experimental/test_generated_tracker.py | 105 ++ .../experimental/test_object_generated.py | 96 ++ ...test_object_scene_registration_temporal.py | 136 +- .../fiducial/blueprints/desk_marker_tf.py | 55 +- .../blueprints/test_desk_marker_tf.py | 17 +- dimos/perception/fiducial/marker_detect.py | 41 +- .../marker_detection_stream_module.py | 28 +- dimos/perception/fiducial/marker_pose.py | 82 +- dimos/perception/fiducial/marker_tf_module.py | 47 +- .../perception/fiducial/marker_transformer.py | 76 +- .../fiducial/test_fixture_verification.py | 37 +- .../fiducial/test_generated_marker_detect.py | 55 + dimos/perception/fiducial/test_helpers.py | 52 +- .../test_marker_detection_stream_module.py | 105 +- dimos/perception/fiducial/test_marker_pose.py | 30 +- .../fiducial/test_marker_tf_integration.py | 41 +- .../fiducial/test_marker_tf_module.py | 83 +- .../fiducial/test_marker_transformer.py | 59 +- .../memory/test_generated_inventory.py | 73 ++ .../pubsub/benchmark/test_replay_bench.py | 44 +- .../pubsub/benchmark/tool_replay_bench.py | 57 +- dimos/protocol/pubsub/encoders.py | 42 +- dimos/protocol/pubsub/impl/jpeg_lcm.py | 26 +- dimos/protocol/pubsub/impl/lcmpubsub.py | 35 +- .../pubsub/impl/rospubsub_conversion.py | 361 +----- dimos/protocol/pubsub/impl/test_jpeg_cdr.py | 76 ++ dimos/protocol/pubsub/impl/test_lcmpubsub.py | 97 +- dimos/protocol/pubsub/impl/test_rospubsub.py | 286 +---- .../pubsub/impl/test_rospubsub_conversion.py | 133 ++ .../protocol/pubsub/impl/test_zenohpubsub.py | 36 +- .../pubsub/impl/webrtc/providers/broker.py | 4 +- .../providers/test_video_track_generated.py | 44 + .../pubsub/impl/webrtc/test_transport.py | 170 ++- .../pubsub/impl/webrtc/tool_blueprint_e2e.py | 6 +- .../pubsub/impl/webrtc/tool_broker_e2e.py | 30 +- .../pubsub/impl/webrtc/webrtcpubsub.py | 2 +- dimos/protocol/pubsub/impl/zenohpubsub.py | 48 +- dimos/protocol/pubsub/test_pattern_sub.py | 13 +- dimos/protocol/pubsub/test_spec.py | 14 +- dimos/protocol/service/test_lcmservice.py | 2 +- dimos/protocol/tf/static_tf_publisher.py | 35 +- dimos/protocol/tf/test_static_tf_publisher.py | 47 + dimos/protocol/tf/test_tf.py | 475 +++---- dimos/protocol/tf/tf.py | 131 +- dimos/robot/assembly/mid360_realsense_30.py | 9 +- dimos/robot/deeprobotics/m20/camera.py | 100 +- dimos/robot/deeprobotics/m20/connection.py | 21 +- .../deeprobotics/m20/test_camera_messages.py | 52 + .../m20/test_connection_messages.py | 53 + dimos/robot/diy/alfred/config.py | 6 +- dimos/robot/diy/alfred/effector_high_level.py | 38 +- dimos/robot/diy/alfred/mount_tf.py | 21 +- dimos/robot/diy/alfred/test_alfred.py | 2 +- .../diy/alfred/test_generated_effector.py | 76 ++ dimos/robot/diy/alfred/test_mount_tf.py | 44 + dimos/robot/drone/camera_module.py | 35 +- dimos/robot/drone/connection_module.py | 64 +- dimos/robot/drone/dji_video_stream.py | 13 +- dimos/robot/drone/drone_tracking_module.py | 31 +- dimos/robot/drone/mavlink_connection.py | 16 +- dimos/robot/drone/test_drone.py | 35 +- .../galaxea/r1pro/test_connection_messages.py | 157 +++ dimos/robot/manipulators/common/mock.py | 3 +- .../dual_openyam/test_blueprints.py | 22 +- .../dual_openyam/test_teleop_ik.py | 16 +- .../openarm/test_openarm_teleop.py | 46 +- .../manipulators/openyam/test_teleop_ik.py | 19 +- dimos/robot/unitree/b1/joystick_module.py | 33 +- dimos/robot/unitree/b1/test_cdr_forwarding.py | 87 ++ dimos/robot/unitree/b1/unitree_b1.py | 10 +- .../basic/unitree_g1_coordinator.py | 6 +- .../g1/blueprints/basic/unitree_g1_record.py | 24 +- .../perceptive/unitree_g1_detection.py | 8 +- .../primitive/test_unitree_g1_vis.py | 43 + .../g1/blueprints/primitive/unitree_g1_vis.py | 28 +- dimos/robot/unitree/g1/config.py | 6 +- dimos/robot/unitree/g1/connection.py | 2 +- dimos/robot/unitree/g1/connection_spec.py | 3 +- .../unitree/g1/debug/demo_arrow_control.py | 31 +- .../g1/debug/demo_arrow_control_cmd_vel.py | 35 +- .../g1/effectors/high_level/dds_sdk.py | 5 +- .../effectors/high_level/high_level_spec.py | 3 +- .../unitree/g1/effectors/high_level/webrtc.py | 5 +- dimos/robot/unitree/g1/g1_recorder.py | 7 +- dimos/robot/unitree/g1/g1_tf_publisher.py | 81 +- dimos/robot/unitree/g1/skill_container.py | 6 +- .../g1/test_g1_manipulation_integration.py | 3 +- .../robot/unitree/g1/test_g1_tf_publisher.py | 57 +- .../unitree/g1/test_wholebody_connection.py | 40 +- .../robot/unitree/g1/wholebody_connection.py | 40 +- .../basic/unitree_go2_coordinator.py | 4 +- .../basic/unitree_go2_holonomic_benchmark.py | 3 +- .../basic/unitree_go2_holonomic_controller.py | 11 +- .../basic/unitree_go2_rpp_benchmark.py | 3 +- .../basic/unitree_go2_rpp_controller.py | 11 +- .../smart/test_unitree_go2_markers.py | 2 +- .../blueprints/smart/unitree_go2_detection.py | 6 +- .../go2/blueprints/smart/unitree_go2_ros.py | 6 +- dimos/robot/unitree/go2/camera_calibration.py | 28 + dimos/robot/unitree/go2/connection.py | 113 +- dimos/robot/unitree/go2/dds/cli/render.py | 90 +- .../go2/dds/cli/test_render_generated.py | 106 ++ dimos/robot/unitree/go2/dds/codec.py | 11 +- dimos/robot/unitree/go2/dds/extrinsics.py | 30 +- dimos/robot/unitree/go2/dds/ros.py | 264 +--- .../unitree/go2/dds/test_generated_codecs.py | 111 ++ dimos/robot/unitree/go2/dds/test_store.py | 22 +- dimos/robot/unitree/go2/dds/video.py | 17 +- dimos/robot/unitree/go2/fleet_connection.py | 2 +- .../robot/unitree/go2/go2_mid360_recorder.py | 7 +- .../go2/go2_mid360_static_transforms.py | 14 +- .../go2/test_go2_mid360_static_transforms.py | 16 +- dimos/robot/unitree/mujoco_connection.py | 41 +- dimos/robot/unitree/test_connection.py | 116 +- .../robot/unitree/test_generated_commands.py | 87 ++ dimos/robot/unitree/test_simulation_tf.py | 87 ++ dimos/robot/unitree/testing/tool_actors.py | 2 +- dimos/robot/unitree/testing/tool_tooling.py | 4 +- dimos/robot/unitree/type/lidar.py | 60 +- dimos/robot/unitree/type/map.py | 32 +- dimos/robot/unitree/type/odometry.py | 62 +- dimos/robot/unitree/type/test_lidar.py | 54 +- dimos/robot/unitree/type/test_odometry.py | 20 +- .../unitree/type/test_pose_conversion.py | 45 + .../robot/unitree/unitree_skill_container.py | 39 +- dimos/simulation/dimsim/cdr_schemas.py | 41 + dimos/simulation/dimsim/test_cdr_schemas.py | 19 + dimos/simulation/engines/base.py | 2 +- dimos/simulation/engines/mujoco_engine.py | 8 +- .../engines/test_mujoco_sim_module.py | 83 +- dimos/simulation/habitat/nix/flake.nix | 4 +- dimos/simulation/habitat/nix/install.sh | 22 +- .../mujoco/direct_cmd_vel_explorer.py | 14 +- .../test_generated_control_consumers.py | 70 ++ dimos/spec/mapping.py | 5 +- dimos/spec/perception.py | 8 +- dimos/teleop/hosted/arm_command.py | 47 +- dimos/teleop/hosted/camera_mux.py | 49 +- dimos/teleop/hosted/go2_command.py | 61 +- dimos/teleop/hosted/hosted_stats.py | 13 +- dimos/teleop/hosted/map_compress.py | 34 +- dimos/teleop/hosted/test_arm_command.py | 52 +- dimos/teleop/hosted/test_camera_mux.py | 51 +- dimos/teleop/hosted/test_go2_command.py | 49 +- dimos/teleop/hosted/test_hosted_stats.py | 20 +- dimos/teleop/hosted/test_map_compress.py | 51 +- .../teleop/keyboard/keyboard_teleop_module.py | 16 +- .../keyboard/test_keyboard_teleop_module.py | 26 +- dimos/teleop/phone/phone_extensions.py | 9 +- dimos/teleop/phone/phone_teleop_module.py | 95 +- dimos/teleop/phone/test_phone_cdr.py | 121 ++ dimos/teleop/phone/web/static/index.html | 63 +- dimos/teleop/utils/recorder.py | 10 +- dimos/teleop/utils/report.py | 47 +- dimos/teleop/utils/teleop_transforms.py | 17 +- dimos/teleop/utils/test_report.py | 71 ++ dimos/teleop/utils/test_video_stats.py | 116 +- dimos/teleop/utils/video_stats.py | 224 +--- dimos/teleop/webxr/blueprints.py | 5 +- dimos/teleop/webxr/controller_types.py | 35 +- dimos/teleop/webxr/extensions.py | 38 +- dimos/teleop/webxr/web/static/teleop.js | 64 +- dimos/types/ros_polyfill.py | 8 +- dimos/types/test_timestamped.py | 19 +- dimos/types/timestamped.py | 48 +- dimos/utils/demo_image_encoding.py | 5 +- dimos/utils/test_transform_utils.py | 211 ++-- dimos/utils/test_twist.py | 35 + dimos/utils/testing/moment.py | 7 +- dimos/utils/testing/replay.py | 9 +- dimos/utils/testing/test_moment.py | 60 +- dimos/utils/testing/test_replay.py | 24 +- dimos/utils/testing/test_sqlite_moment.py | 75 ++ dimos/utils/transform_utils.py | 114 +- dimos/visualization/rerun/bridge.py | 69 +- dimos/visualization/rerun/init.py | 2 +- dimos/visualization/rerun/message_helpers.py | 283 +++++ .../rerun/test_detection3d_bridge.py | 44 +- .../rerun/test_generated_camera_bridge.py | 79 ++ .../rerun/test_generated_cloud_bridge.py | 118 ++ .../rerun/test_generated_entity_bridge.py | 69 + .../rerun/test_generated_navigation_bridge.py | 80 ++ .../rerun/test_generated_occupancy_bridge.py | 69 + .../rerun/test_generated_tf_bridge.py | 67 + .../visualization/rerun/test_viewer_ws_e2e.py | 19 +- .../rerun/test_websocket_server.py | 60 +- dimos/visualization/rerun/tf_tree.py | 99 ++ dimos/visualization/rerun/websocket_server.py | 34 +- dimos/web/cdr_codec.py | 74 ++ dimos/web/codecs.py | 36 +- dimos/web/lcm_codec.py | 188 --- dimos/web/relay_bridge/gen_cdr_fixtures.py | 106 ++ .../web/relay_bridge/gen_costmap_fixtures.py | 19 +- dimos/web/relay_bridge/gen_lcm_fixtures.py | 221 ---- .../web/relay_bridge/test_costmap_encoding.py | 34 +- dimos/web/relay_bridge/test_dynamic.py | 14 +- dimos/web/relay_bridge/test_map_codecs.py | 18 +- .../test_relay_bridge_authoring.py | 73 +- dimos/web/test_cdr_codec.py | 73 ++ dimos/web/test_codecs.py | 42 +- .../test_websocket_vis_module.py | 104 ++ .../manipulation/adding_a_custom_arm.md | 12 +- docs/capabilities/manipulation/python_api.md | 2 +- docs/capabilities/manipulation/xarm-grasp.md | 4 +- docs/capabilities/memory/algo_comparison.md | 46 +- .../capabilities/memory/assets/brightness.svg | 4 +- .../memory/assets/color_image.svg | 4 +- docs/capabilities/memory/assets/embedding.svg | 4 +- .../memory/assets/embedding_focused.svg | 4 +- docs/capabilities/memory/assets/grid.png | 4 +- .../memory/assets/plot_brightness_algo.svg | 4 +- .../assets/plot_brightness_algo_delta.svg | 4 +- .../memory/assets/plot_colors.svg | 4 +- .../capabilities/memory/assets/plot_named.svg | 4 +- .../memory/assets/plot_plantness.svg | 4 +- .../assets/plot_plantness_brightness.svg | 4 +- .../memory/assets/plot_plantness_gap_fill.svg | 4 +- .../memory/assets/plot_robot_data.svg | 4 +- docs/capabilities/memory/assets/speed.svg | 4 +- docs/capabilities/memory/index.md | 47 +- docs/capabilities/memory/plot.md | 51 +- docs/capabilities/navigation/deep_dive.md | 2 +- docs/capabilities/navigation/habitat.md | 2 +- docs/coding-agents/testing.md | 2 +- docs/development/large_file_management.md | 6 +- docs/requirements.md | 2 +- docs/usage/blueprints.md | 4 +- docs/usage/data_streams/advanced_streams.md | 2 +- docs/usage/data_streams/index.md | 6 +- docs/usage/data_streams/quality_filter.md | 39 +- docs/usage/data_streams/temporal_alignment.md | 31 +- docs/usage/evals.md | 20 +- docs/usage/lcm.md | 227 ++-- docs/usage/modules.md | 38 +- docs/usage/native_modules.md | 10 +- docs/usage/python-api.md | 6 +- docs/usage/sensor_streams/advanced_streams.md | 2 +- docs/usage/sensor_streams/index.md | 6 +- docs/usage/sensor_streams/quality_filter.md | 39 +- .../sensor_streams/temporal_alignment.md | 31 +- docs/usage/transforms.md | 215 ++-- docs/usage/transports/index.md | 16 +- docs/web/web_sdk.md | 17 +- examples/camera_grayscale.py | 17 +- examples/language-interop/README.md | 27 +- examples/language-interop/cpp/CMakeLists.txt | 30 +- examples/language-interop/cpp/README.md | 26 +- examples/language-interop/cpp/main.cpp | 18 +- examples/language-interop/lua/README.md | 43 +- examples/language-interop/ts/README.md | 39 +- .../message-codegen/EVIDENCE-navigation.md | 371 ++++++ examples/message-codegen/EVIDENCE-recovery.md | 285 +++++ .../message-codegen/demo_aruco_detection.py | 100 ++ .../message-codegen/demo_base_transport.py | 106 ++ .../message-codegen/demo_bbox_navigation.py | 78 ++ examples/message-codegen/demo_blueprint.py | 216 ++++ .../message-codegen/demo_camera_bridge.py | 101 ++ .../message-codegen/demo_camera_mux_cdr.py | 63 + .../message-codegen/demo_cloud_filters.py | 51 + .../message-codegen/demo_command_center.py | 59 + .../demo_coordinator_joints.py | 82 ++ .../message-codegen/demo_detection_bbox.py | 62 + .../demo_detection_projection.py | 67 + .../demo_export_detection_fixture.py | 179 +++ .../message-codegen/demo_external_native.py | 60 + examples/message-codegen/demo_frontiers.py | 67 + examples/message-codegen/demo_g1_messages.py | 65 + .../message-codegen/demo_galaxea_messages.py | 82 ++ .../message-codegen/demo_go2_command_cdr.py | 60 + .../message-codegen/demo_grasp_proposals.py | 60 + examples/message-codegen/demo_grid_path.py | 73 ++ .../message-codegen/demo_habitat_native.py | 96 ++ .../message-codegen/demo_holonomic_path.py | 103 ++ .../demo_holonomic_tracking.py | 59 + examples/message-codegen/demo_hosted_map.py | 61 + examples/message-codegen/demo_keyboard_cdr.py | 54 + examples/message-codegen/demo_lidar.cpp | 43 + examples/message-codegen/demo_live_bridge.py | 122 ++ .../message-codegen/demo_marker_messages.py | 63 + examples/message-codegen/demo_mcp_observe.py | 72 ++ .../message-codegen/demo_message_helpers.py | 60 + .../message-codegen/demo_mls_transformer.py | 156 +++ examples/message-codegen/demo_movement.py | 57 + .../message-codegen/demo_mujoco_lidar_shm.py | 80 ++ examples/message-codegen/demo_native.py | 193 +++ examples/message-codegen/demo_navigation.py | 178 +++ .../demo_navigation_recording.py | 118 ++ .../message-codegen/demo_navigation_skill.py | 60 + examples/message-codegen/demo_occupancy.py | 57 + .../message-codegen/demo_odometry_history.py | 91 ++ .../message-codegen/demo_path_following.py | 85 ++ examples/message-codegen/demo_patrol.py | 65 + .../message-codegen/demo_person_follow.py | 78 ++ examples/message-codegen/demo_phone_cdr.py | 74 ++ .../message-codegen/demo_planning_backends.py | 108 ++ .../message-codegen/demo_planning_joints.py | 75 ++ examples/message-codegen/demo_pointcloud.py | 63 + examples/message-codegen/demo_pubsub.py | 128 ++ examples/message-codegen/demo_robot_tf.py | 48 + examples/message-codegen/demo_ros_bridge.py | 94 ++ examples/message-codegen/demo_self_filter.py | 76 ++ examples/message-codegen/demo_trajectory.py | 137 ++ examples/message-codegen/demo_video_stats.py | 55 + .../message-codegen/demo_viewer_controls.py | 85 ++ examples/message-codegen/demo_visual_servo.py | 41 + examples/message-codegen/demo_vl_images.py | 50 + examples/message-codegen/demo_vlm_image.py | 54 + .../demo_whole_body_transport.py | 112 ++ .../evidence/camera-export-runtime.md | 65 + .../evidence/compressed-image-runtime.md | 35 + .../evidence/experimental-consumers.md | 44 + .../evidence/final-consumer-cutover.md | 45 + .../evidence/fixture-routing-cutover.md | 33 + .../evidence/manual-tools-replay.md | 26 + .../evidence/real-detection-fixture.md | 62 + .../evidence/recorder-graph-runtime.md | 37 + .../evidence/robot-simulation-consumers.md | 55 + .../evidence/unitree-spot-camera-runtime.md | 45 + .../external-native/CMakeLists.txt | 9 + .../message-codegen/external-native/relay.cpp | 25 + .../message-codegen/external-native/relay.rs | 65 + .../message-codegen/test_base_transport.py | 32 + examples/message-codegen/test_blueprint.py | 75 ++ examples/message-codegen/test_live_bridge.py | 42 + examples/message-codegen/test_navigation.py | 32 + examples/message-codegen/test_trajectory.py | 32 + .../test_whole_body_transport.py | 32 + examples/native-modules/cpp/CMakeLists.txt | 27 +- examples/native-modules/cpp/flake.lock | 18 - examples/native-modules/cpp/flake.nix | 11 +- examples/native-modules/cpp/src/cdr_relay.cpp | 31 + examples/native-modules/cpp/src/ping.cpp | 4 +- examples/native-modules/cpp/src/pong.cpp | 4 +- examples/native-modules/cpp_ping_pong.py | 3 +- examples/native-modules/rust/Cargo.toml | 6 +- examples/native-modules/rust/src/cdr_relay.rs | 52 + examples/native-modules/rust/src/ping.rs | 7 +- examples/native-modules/rust/src/pong.rs | 7 +- examples/native-modules/test_rust_tf.py | 99 ++ examples/simplerobot/README.md | 6 +- examples/simplerobot/simplerobot.py | 33 +- examples/simplerobot/vis.py | 21 +- misc/DimSim/README.md | 22 + misc/DimSim/cli/bridge/cdr.ts | 40 + misc/DimSim/cli/bridge/cdr_schemas.ts | 7 + misc/DimSim/cli/bridge/cdr_test.ts | 30 + misc/DimSim/cli/bridge/lidar.ts | 31 +- misc/DimSim/cli/bridge/physics.ts | 40 +- misc/DimSim/cli/bridge/server.ts | 2 +- misc/DimSim/cli/deno.json | 5 +- misc/DimSim/cli/deno.lock | 53 +- misc/DimSim/cli/vendor/lcm/types.ts | 4 +- misc/DimSim/package-lock.json | 59 + misc/DimSim/package.json | 4 +- misc/DimSim/src/bridge.ts | 35 +- misc/DimSim/src/engine.js | 8 +- misc/optimize_patrol/plot_path.py | 4 +- native/cpp/CMakeLists.txt | 19 + native/cpp/cmake/dimos_nativeConfig.cmake.in | 28 + native/cpp/include/dimos/native.hpp | 2 +- native/cpp/include/dimos/native/lcm_codec.hpp | 33 - native/cpp/include/dimos/native/module.hpp | 6 +- native/cpp/messages.nix | 26 + native/cpp/tests/CMakeLists.txt | 7 +- native/cpp/tests/test_cdr_codec.cpp | 34 + native/cpp/tests/test_lcm_codec.cpp | 57 - native/cpp/tests/test_module.cpp | 24 +- .../example/example_external/runtime.py | 5 +- .../graspgenx/graspgenx_runtime/runtime.py | 39 +- .../graspgenx_runtime/test_runtime.py | 27 +- .../python/lerobot/dimos_lerobot/runtime.py | 52 +- .../lerobot/dimos_lerobot/runtime_tests.py | 62 +- native/rust/README.md | 14 +- native/rust/dimos-module-macros/src/lib.rs | 19 +- native/rust/dimos-module/Cargo.toml | 4 +- native/rust/dimos-module/src/cdr.rs | 73 ++ native/rust/dimos-module/src/host.rs | 2 +- native/rust/dimos-module/src/lcm.rs | 8 +- native/rust/dimos-module/src/lib.rs | 5 +- native/rust/dimos-module/src/module.rs | 75 +- native/rust/dimos-module/src/tf.rs | 92 +- scripts/g1_replay.py | 5 +- scripts/install_native_messages.sh | 12 + scripts/package_native_sdk.py | 44 + scripts/test_external_native_messages.sh | 30 + scripts/test_message_user_story.sh | 23 + web/cockpit/src/App.test.tsx | 28 +- web/cockpit/src/subscriptions.test.ts | 39 +- web/sdk/fixture/demo_cdr.html | 77 ++ web/sdk/fixture/demo_cdr.ts | 59 + web/sdk/package.json | 4 + web/sdk/src/decoders/cdr.test.ts | 45 + web/sdk/src/decoders/cdr.ts | 106 ++ web/sdk/src/decoders/lcm.test.ts | 374 ------ web/sdk/src/decoders/lcm.ts | 424 ------- web/sdk/src/session.test.ts | 54 +- web/shared/fixtures/cdr_frames.json | 236 ++++ web/shared/fixtures/costmap_frames.json | 4 +- web/shared/fixtures/lcm_frames.json | 1114 ----------------- 877 files changed, 30022 insertions(+), 16191 deletions(-) create mode 100644 data/.lfs/unitree_go2_detection_cdr.tar.gz create mode 100644 dimos/agents/mcp/test_generated_image.py create mode 100644 dimos/agents/skills/test_generated_navigation.py create mode 100644 dimos/agents/skills/test_generated_person_follow.py create mode 100644 dimos/agents/test_generated_vlm.py create mode 100644 dimos/core/test_cdr_transport.py create mode 100644 dimos/e2e_tests/cdr_replay_fixture.py create mode 100644 dimos/e2e_tests/test_cdr_spy.py create mode 100644 dimos/experimental/isolated_python/test_generated_example.py create mode 100644 dimos/experimental/robot/bosdyn/spot/effectors/test_high_level.py create mode 100644 dimos/experimental/world_belief/test_absence.py create mode 100644 dimos/experimental/world_belief/test_generated_recall.py create mode 100644 dimos/hardware/drive_trains/transport/test_adapter.py create mode 100644 dimos/hardware/sensors/camera/gstreamer/test_generated_frames.py create mode 100644 dimos/hardware/sensors/camera/test_module.py create mode 100644 dimos/hardware/sensors/camera/zed/test_generated_tf.py create mode 100644 dimos/hardware/sensors/lidar/test_generated_recorders.py create mode 100644 dimos/hardware/whole_body/transport/test_adapter.py create mode 100644 dimos/imitation/dataprep/test_generated_images.py create mode 100644 dimos/manipulation/control/test_generated_legacy_controllers.py create mode 100644 dimos/manipulation/planning/monitor/test_robot_state_monitor.py create mode 100644 dimos/mapping/cli/test_pgo_accumulate.py create mode 100644 dimos/mapping/cli/test_view.py create mode 100644 dimos/mapping/occupancy/test_generated_visualizations.py create mode 100644 dimos/mapping/osm/test_current_location_map.py create mode 100644 dimos/mapping/pointclouds/test_generated_occupancy.py create mode 100644 dimos/mapping/test_generated_costmapper.py create mode 100644 dimos/mapping/test_generated_odometry_hist.py create mode 100644 dimos/mapping/test_generated_voxels.py create mode 100644 dimos/mapping/voxels/test_grid.py rename dimos/memory/codecs/{lcm.py => cdr.py} (67%) delete mode 100644 dimos/memory/codecs/jpeg.py create mode 100644 dimos/memory/demo_data.py create mode 100644 dimos/memory/recording_policy.py create mode 100644 dimos/memory/type/test_generated_query.py create mode 100644 dimos/memory/vis/space/geometry.py create mode 100644 dimos/memory/vis/space/test_generated_geometry.py create mode 100644 dimos/models/segmentation/test_generated_edge_tam_images.py create mode 100644 dimos/models/vl/test_generated_images.py create mode 100644 dimos/msgs/camera_info.py create mode 100644 dimos/msgs/geometry.py create mode 100644 dimos/msgs/image.py create mode 100644 dimos/msgs/occupancy.py create mode 100644 dimos/msgs/pointcloud.py create mode 100644 dimos/msgs/test_camera_info.py create mode 100644 dimos/msgs/test_generated_projection.py create mode 100644 dimos/msgs/test_generated_views.py create mode 100644 dimos/msgs/test_geometry.py create mode 100644 dimos/msgs/test_occupancy.py create mode 100644 dimos/msgs/test_pointcloud.py create mode 100644 dimos/msgs/test_pointcloud_colors.py create mode 100644 dimos/msgs/test_time.py create mode 100644 dimos/msgs/test_trajectory.py create mode 100644 dimos/msgs/time.py create mode 100644 dimos/msgs/trajectory.py create mode 100644 dimos/navigation/test_generated_bbox_navigation.py create mode 100644 dimos/perception/detection/reid/test_generated_module.py create mode 100644 dimos/perception/detection/type/detection2d/test_generated_bbox.py create mode 100644 dimos/perception/detection/type/detection3d/test_generated_filters.py create mode 100644 dimos/perception/detection/type/detection3d/test_generated_projection.py create mode 100644 dimos/perception/experimental/test_generated_tracker.py create mode 100644 dimos/perception/experimental/test_object_generated.py create mode 100644 dimos/perception/fiducial/test_generated_marker_detect.py create mode 100644 dimos/perception/memory/test_generated_inventory.py create mode 100644 dimos/protocol/pubsub/impl/test_jpeg_cdr.py create mode 100644 dimos/protocol/pubsub/impl/test_rospubsub_conversion.py create mode 100644 dimos/protocol/pubsub/impl/webrtc/providers/test_video_track_generated.py create mode 100644 dimos/protocol/tf/test_static_tf_publisher.py create mode 100644 dimos/robot/deeprobotics/m20/test_camera_messages.py create mode 100644 dimos/robot/deeprobotics/m20/test_connection_messages.py create mode 100644 dimos/robot/diy/alfred/test_generated_effector.py create mode 100644 dimos/robot/diy/alfred/test_mount_tf.py create mode 100644 dimos/robot/galaxea/r1pro/test_connection_messages.py create mode 100644 dimos/robot/unitree/b1/test_cdr_forwarding.py create mode 100644 dimos/robot/unitree/g1/blueprints/primitive/test_unitree_g1_vis.py create mode 100644 dimos/robot/unitree/go2/camera_calibration.py create mode 100644 dimos/robot/unitree/go2/dds/cli/test_render_generated.py create mode 100644 dimos/robot/unitree/go2/dds/test_generated_codecs.py create mode 100644 dimos/robot/unitree/test_generated_commands.py create mode 100644 dimos/robot/unitree/test_simulation_tf.py create mode 100644 dimos/robot/unitree/type/test_pose_conversion.py create mode 100644 dimos/simulation/dimsim/cdr_schemas.py create mode 100644 dimos/simulation/dimsim/test_cdr_schemas.py create mode 100644 dimos/simulation/mujoco/test_generated_control_consumers.py create mode 100644 dimos/teleop/phone/test_phone_cdr.py create mode 100644 dimos/teleop/utils/test_report.py create mode 100644 dimos/utils/test_twist.py create mode 100644 dimos/utils/testing/test_sqlite_moment.py create mode 100644 dimos/visualization/rerun/message_helpers.py create mode 100644 dimos/visualization/rerun/test_generated_camera_bridge.py create mode 100644 dimos/visualization/rerun/test_generated_cloud_bridge.py create mode 100644 dimos/visualization/rerun/test_generated_entity_bridge.py create mode 100644 dimos/visualization/rerun/test_generated_navigation_bridge.py create mode 100644 dimos/visualization/rerun/test_generated_occupancy_bridge.py create mode 100644 dimos/visualization/rerun/test_generated_tf_bridge.py create mode 100644 dimos/visualization/rerun/tf_tree.py create mode 100644 dimos/web/cdr_codec.py delete mode 100644 dimos/web/lcm_codec.py create mode 100644 dimos/web/relay_bridge/gen_cdr_fixtures.py delete mode 100644 dimos/web/relay_bridge/gen_lcm_fixtures.py create mode 100644 dimos/web/test_cdr_codec.py create mode 100644 dimos/web/websocket_vis/test_websocket_vis_module.py create mode 100644 examples/message-codegen/EVIDENCE-navigation.md create mode 100644 examples/message-codegen/EVIDENCE-recovery.md create mode 100644 examples/message-codegen/demo_aruco_detection.py create mode 100644 examples/message-codegen/demo_base_transport.py create mode 100644 examples/message-codegen/demo_bbox_navigation.py create mode 100644 examples/message-codegen/demo_blueprint.py create mode 100644 examples/message-codegen/demo_camera_bridge.py create mode 100644 examples/message-codegen/demo_camera_mux_cdr.py create mode 100644 examples/message-codegen/demo_cloud_filters.py create mode 100644 examples/message-codegen/demo_command_center.py create mode 100644 examples/message-codegen/demo_coordinator_joints.py create mode 100644 examples/message-codegen/demo_detection_bbox.py create mode 100644 examples/message-codegen/demo_detection_projection.py create mode 100644 examples/message-codegen/demo_export_detection_fixture.py create mode 100644 examples/message-codegen/demo_external_native.py create mode 100644 examples/message-codegen/demo_frontiers.py create mode 100644 examples/message-codegen/demo_g1_messages.py create mode 100644 examples/message-codegen/demo_galaxea_messages.py create mode 100644 examples/message-codegen/demo_go2_command_cdr.py create mode 100644 examples/message-codegen/demo_grasp_proposals.py create mode 100644 examples/message-codegen/demo_grid_path.py create mode 100644 examples/message-codegen/demo_habitat_native.py create mode 100644 examples/message-codegen/demo_holonomic_path.py create mode 100644 examples/message-codegen/demo_holonomic_tracking.py create mode 100644 examples/message-codegen/demo_hosted_map.py create mode 100644 examples/message-codegen/demo_keyboard_cdr.py create mode 100644 examples/message-codegen/demo_lidar.cpp create mode 100644 examples/message-codegen/demo_live_bridge.py create mode 100644 examples/message-codegen/demo_marker_messages.py create mode 100644 examples/message-codegen/demo_mcp_observe.py create mode 100644 examples/message-codegen/demo_message_helpers.py create mode 100644 examples/message-codegen/demo_mls_transformer.py create mode 100644 examples/message-codegen/demo_movement.py create mode 100644 examples/message-codegen/demo_mujoco_lidar_shm.py create mode 100644 examples/message-codegen/demo_native.py create mode 100644 examples/message-codegen/demo_navigation.py create mode 100644 examples/message-codegen/demo_navigation_recording.py create mode 100644 examples/message-codegen/demo_navigation_skill.py create mode 100644 examples/message-codegen/demo_occupancy.py create mode 100644 examples/message-codegen/demo_odometry_history.py create mode 100644 examples/message-codegen/demo_path_following.py create mode 100644 examples/message-codegen/demo_patrol.py create mode 100644 examples/message-codegen/demo_person_follow.py create mode 100644 examples/message-codegen/demo_phone_cdr.py create mode 100644 examples/message-codegen/demo_planning_backends.py create mode 100644 examples/message-codegen/demo_planning_joints.py create mode 100644 examples/message-codegen/demo_pointcloud.py create mode 100644 examples/message-codegen/demo_pubsub.py create mode 100644 examples/message-codegen/demo_robot_tf.py create mode 100644 examples/message-codegen/demo_ros_bridge.py create mode 100644 examples/message-codegen/demo_self_filter.py create mode 100644 examples/message-codegen/demo_trajectory.py create mode 100644 examples/message-codegen/demo_video_stats.py create mode 100644 examples/message-codegen/demo_viewer_controls.py create mode 100644 examples/message-codegen/demo_visual_servo.py create mode 100644 examples/message-codegen/demo_vl_images.py create mode 100644 examples/message-codegen/demo_vlm_image.py create mode 100644 examples/message-codegen/demo_whole_body_transport.py create mode 100644 examples/message-codegen/evidence/camera-export-runtime.md create mode 100644 examples/message-codegen/evidence/compressed-image-runtime.md create mode 100644 examples/message-codegen/evidence/experimental-consumers.md create mode 100644 examples/message-codegen/evidence/final-consumer-cutover.md create mode 100644 examples/message-codegen/evidence/fixture-routing-cutover.md create mode 100644 examples/message-codegen/evidence/manual-tools-replay.md create mode 100644 examples/message-codegen/evidence/real-detection-fixture.md create mode 100644 examples/message-codegen/evidence/recorder-graph-runtime.md create mode 100644 examples/message-codegen/evidence/robot-simulation-consumers.md create mode 100644 examples/message-codegen/evidence/unitree-spot-camera-runtime.md create mode 100644 examples/message-codegen/external-native/CMakeLists.txt create mode 100644 examples/message-codegen/external-native/relay.cpp create mode 100644 examples/message-codegen/external-native/relay.rs create mode 100644 examples/message-codegen/test_base_transport.py create mode 100644 examples/message-codegen/test_blueprint.py create mode 100644 examples/message-codegen/test_live_bridge.py create mode 100644 examples/message-codegen/test_navigation.py create mode 100644 examples/message-codegen/test_trajectory.py create mode 100644 examples/message-codegen/test_whole_body_transport.py create mode 100644 examples/native-modules/cpp/src/cdr_relay.cpp create mode 100644 examples/native-modules/rust/src/cdr_relay.rs create mode 100644 examples/native-modules/test_rust_tf.py create mode 100644 misc/DimSim/cli/bridge/cdr.ts create mode 100644 misc/DimSim/cli/bridge/cdr_schemas.ts create mode 100644 misc/DimSim/cli/bridge/cdr_test.ts create mode 100644 native/cpp/cmake/dimos_nativeConfig.cmake.in delete mode 100644 native/cpp/include/dimos/native/lcm_codec.hpp create mode 100644 native/cpp/messages.nix create mode 100644 native/cpp/tests/test_cdr_codec.cpp delete mode 100644 native/cpp/tests/test_lcm_codec.cpp create mode 100644 native/rust/dimos-module/src/cdr.rs create mode 100644 scripts/install_native_messages.sh create mode 100644 scripts/package_native_sdk.py create mode 100644 scripts/test_external_native_messages.sh create mode 100644 scripts/test_message_user_story.sh create mode 100644 web/sdk/fixture/demo_cdr.html create mode 100644 web/sdk/fixture/demo_cdr.ts create mode 100644 web/sdk/src/decoders/cdr.test.ts create mode 100644 web/sdk/src/decoders/cdr.ts delete mode 100644 web/sdk/src/decoders/lcm.test.ts delete mode 100644 web/sdk/src/decoders/lcm.ts create mode 100644 web/shared/fixtures/cdr_frames.json delete mode 100644 web/shared/fixtures/lcm_frames.json diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 6e117d5d33c..e19b4a87799 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -140,6 +140,12 @@ repos: bash -c ' set -euo pipefail git ls-files "*Cargo.toml" | while read -r m; do + # This consumer intentionally references generated packages absent in a clean checkout. + # Format its checked-in source directly; native CI builds and runs the full project. + if [ "$m" = examples/message-project/rust/Cargo.toml ]; then + rustfmt --edition 2024 examples/message-project/rust/src/main.rs + continue + fi cargo locate-project --manifest-path "$m" --workspace --message-format plain done | sort -u | while read -r root; do cargo fmt --manifest-path "$root" --all diff --git a/AGENTS.md b/AGENTS.md index c16f42cf93c..0d2ca2afbde 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -157,7 +157,7 @@ Autonomous subsystems. Communicate via `In[T]`/`Out[T]` typed streams. Run in fo from dimos.core.module import Module from dimos.core.stream import In, Out from dimos.core.core import rpc -from dimos.msgs.sensor_msgs import Image +from dimos_generated.sensor_msgs.msg import Image class MyModule(Module): color_image: In[Image] diff --git a/Cargo.lock b/Cargo.lock index 5b05f816169..afa496958a1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -114,7 +114,7 @@ dependencies = [ "nom 7.1.3", "num-traits", "rusticata-macros 4.1.0", - "thiserror 2.0.20", + "thiserror", "time", ] @@ -262,6 +262,20 @@ name = "bytemuck" version = "1.25.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a1f896587b6f2c069c73d2f0913e2d590c3990285cd2f0b6aa02b786b4c679c" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.4", +] [[package]] name = "byteorder" @@ -344,15 +358,6 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b0fc239e0f6cb375d2402d48afb92f76f5404fd1df208a41930ec81eda078bea" -[[package]] -name = "cmake" -version = "0.1.58" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c0f78a02292a74a88ac736019ab962ece0bc380e3f977bf72e376c5d78ff0678" -dependencies = [ - "cc", -] - [[package]] name = "combine" version = "4.6.8" @@ -610,7 +615,7 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "10d60334b3b2e7c9d91ef8150abfb6fa4c1c39ebbcf4a81c2e346aad939fee3e" dependencies = [ - "thiserror 2.0.20", + "thiserror", ] [[package]] @@ -648,13 +653,12 @@ dependencies = [ ] [[package]] -name = "dimos-lcm" +name = "dimos-generated-messages" version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" dependencies = [ - "byteorder", - "socket2 0.5.10", - "tokio", + "re_cdr", + "serde", + "serde-big-array", ] [[package]] @@ -662,9 +666,9 @@ name = "dimos-livox" version = "0.1.0" dependencies = [ "crc", + "dimos-generated-messages", "dimos-module", "etherparse", - "lcm-msgs", "pcap-parser", "serde", "serde_json", @@ -683,8 +687,8 @@ dependencies = [ "anyhow", "crossbeam-channel", "crossbeam-utils", + "dimos-generated-messages", "dimos-module", - "lcm-msgs", "lz4_flex 0.14.0", "mcap", "rayon", @@ -694,7 +698,6 @@ dependencies = [ "tempfile", "tokio", "tracing", - "turbojpeg", "validator", ] @@ -703,8 +706,8 @@ name = "dimos-mls-planner" version = "0.1.0" dependencies = [ "ahash", + "dimos-generated-messages", "dimos-module", - "lcm-msgs", "rayon", "serde", "tokio", @@ -728,9 +731,9 @@ dependencies = [ name = "dimos-module" version = "0.1.0" dependencies = [ - "dimos-lcm", + "dimos-generated-messages", + "dimos-lcm-transport", "dimos-module-macros", - "lcm-msgs", "nalgebra", "rayon", "serde", @@ -758,8 +761,8 @@ dependencies = [ name = "dimos-native-module-examples" version = "0.1.0" dependencies = [ + "dimos-generated-messages", "dimos-module", - "lcm-msgs", "serde", "tokio", "tracing", @@ -770,9 +773,9 @@ dependencies = [ name = "dimos-pointlio" version = "0.1.0" dependencies = [ + "dimos-generated-messages", "dimos-livox", "dimos-module", - "lcm-msgs", "pointlio-core", "serde", "serde_json", @@ -815,8 +818,8 @@ version = "0.1.0" dependencies = [ "ahash", "arrayvec", + "dimos-generated-messages", "dimos-module", - "lcm-msgs", "nalgebra", "rayon", "serde", @@ -1131,12 +1134,6 @@ dependencies = [ "slab", ] -[[package]] -name = "gcd" -version = "2.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d758ba1b47b00caf47f24925c0074ecb20d6dfcffe7f6d53395c0465674841a" - [[package]] name = "generic-array" version = "0.14.7" @@ -1590,7 +1587,7 @@ dependencies = [ "jni-sys", "log", "simd_cesu8", - "thiserror 2.0.20", + "thiserror", "walkdir", "windows-link", ] @@ -1698,14 +1695,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" -[[package]] -name = "lcm-msgs" -version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" -dependencies = [ - "byteorder", -] - [[package]] name = "libc" version = "0.2.189" @@ -1841,7 +1830,7 @@ dependencies = [ "num_cpus", "paste", "static_assertions", - "thiserror 2.0.20", + "thiserror", "zstd", ] @@ -2506,7 +2495,7 @@ dependencies = [ "rustc-hash", "rustls", "socket2 0.6.5", - "thiserror 2.0.20", + "thiserror", "tokio", "tracing", "web-time", @@ -2530,7 +2519,7 @@ dependencies = [ "rustls-pki-types", "rustls-platform-verifier", "slab", - "thiserror 2.0.20", + "thiserror", "tinyvec", "tracing", "web-time", @@ -2667,6 +2656,18 @@ dependencies = [ "yasna", ] +[[package]] +name = "re_cdr" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bba10c21433b597481f366b978cdf2617e56785b38cb0dcc3833c6b8ad7e061a" +dependencies = [ + "bytemuck", + "byteorder", + "serde", + "thiserror", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -2684,7 +2685,7 @@ checksum = "a4e608c6638b9c18977b00b475ac1f28d14e84b27d8d42f70e0bf1e3dec127ac" dependencies = [ "getrandom 0.2.17", "libredox", - "thiserror 2.0.20", + "thiserror", ] [[package]] @@ -2790,7 +2791,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c51c9ae4df8a7fba42103df5c621fa3c37eccf3a3c650879e90fc48b11cc192c" dependencies = [ "hashbrown 0.16.1", - "thiserror 2.0.20", + "thiserror", ] [[package]] @@ -3070,6 +3071,15 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "serde-big-array" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11fc7cc2c76d73e0f27ee52abbd64eec84d46f370c88371120433196934e4b7f" +dependencies = [ + "serde", +] + [[package]] name = "serde_core" version = "1.0.229" @@ -3500,33 +3510,13 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "thiserror" -version = "1.0.69" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" -dependencies = [ - "thiserror-impl 1.0.69", -] - [[package]] name = "thiserror" version = "2.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f" dependencies = [ - "thiserror-impl 2.0.20", -] - -[[package]] -name = "thiserror-impl" -version = "1.0.69" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.119", + "thiserror-impl", ] [[package]] @@ -3841,29 +3831,6 @@ dependencies = [ "syn 2.0.119", ] -[[package]] -name = "turbojpeg" -version = "1.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "81313fcd64aaa30ddfd59eaa122968136b2dd9cd41d9619b126cf71e67a01702" -dependencies = [ - "gcd", - "libc", - "thiserror 1.0.69", - "turbojpeg-sys", -] - -[[package]] -name = "turbojpeg-sys" -version = "1.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49b5a974bc59382d35927e3a07eef2ba9a24fe7227cefe52ec44b508e1b90f86" -dependencies = [ - "anyhow", - "cmake", - "libc", -] - [[package]] name = "twox-hash" version = "1.6.3" @@ -4420,7 +4387,7 @@ dependencies = [ "oid-registry", "ring", "rusticata-macros 4.1.0", - "thiserror 2.0.20", + "thiserror", "time", ] diff --git a/README.md b/README.md index d9bbf76ed08..18541aa8c27 100644 --- a/README.md +++ b/README.md @@ -250,8 +250,10 @@ from dimos.core.coordination.blueprints import autoconnect from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs import Twist -from dimos.msgs.sensor_msgs import Image, ImageFormat +from dimos_generated.geometry_msgs.msg import Twist +from dimos_generated.sensor_msgs.msg import Image +from dimos.msgs.image import image_from_array +from dimos.msgs.time import header_now class RobotConnection(Module): cmd_vel: In[Twist] @@ -263,10 +265,10 @@ class RobotConnection(Module): def _image_loop(self): while True: - img = Image.from_numpy( + img = image_from_array( np.zeros((120, 160, 3), np.uint8), - format=ImageFormat.RGB, - frame_id="camera_optical", + encoding="rgb8", + header=header_now(frame_id="camera_optical"), ) self.color_image.publish(img) time.sleep(0.2) @@ -296,7 +298,7 @@ A blueprint example that connects the image stream from a robot to an MCP-backed ```py skip from dimos.core.coordination.blueprints import autoconnect from dimos.core.transport import LCMTransport -from dimos.msgs.sensor_msgs import Image +from dimos_generated.sensor_msgs.msg import Image from dimos.robot.unitree.go2.connection import go2_connection from dimos.agents.mcp.mcp_client import McpClient from dimos.agents.mcp.mcp_server import McpServer @@ -344,9 +346,9 @@ uv run pytest --numprocesses=auto dimos ## Multi Language Support -Python is our glue and prototyping language, but we support many languages via LCM interop. +Python, C++ and Rust share generated `.msg` value types and CDR codecs over raw LCM or Zenoh. -Check our language interop examples: -- [C++](examples/language-interop/cpp/) -- [Lua](examples/language-interop/lua/) -- [TypeScript](examples/language-interop/ts/) +- [Add and use a message](docs/development/messages.md) +- [Three-language relay and recording demos](examples/message-codegen/README.md) +- [C++ virtual robot controller](examples/language-interop/cpp/README.md) +- [Browser CDR SDK](docs/web/web_sdk.md) diff --git a/data/.lfs/astar_corner_min_cost.png.tar.gz b/data/.lfs/astar_corner_min_cost.png.tar.gz index 35f3ffe0b68..a6c2a84d750 100644 --- a/data/.lfs/astar_corner_min_cost.png.tar.gz +++ b/data/.lfs/astar_corner_min_cost.png.tar.gz @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:42517c5f67a9f06949cb2015a345f9d6b43d22cafd50e1fefb9b5d24d8b72509 -size 5671 +oid sha256:bb08a2ae88cdd133924d3e48f796bc80d1c78c21a42a819ba9a486e1363ed0cb +size 5640 diff --git a/data/.lfs/astar_min_cost.png.tar.gz b/data/.lfs/astar_min_cost.png.tar.gz index 752a7782951..820eaab9933 100644 --- a/data/.lfs/astar_min_cost.png.tar.gz +++ b/data/.lfs/astar_min_cost.png.tar.gz @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:06b67aa0d18c291c3525e67ca3a2a9ab2530f6fe782a850872ba4c343353a20a -size 12018 +oid sha256:861d4613729fba6b43855661497dbd3c69f9f7597ea31d0c1aee9ad82666a0ab +size 11961 diff --git a/data/.lfs/gradient_simple.png.tar.gz b/data/.lfs/gradient_simple.png.tar.gz index 7232282ce42..f0833e747bf 100644 --- a/data/.lfs/gradient_simple.png.tar.gz +++ b/data/.lfs/gradient_simple.png.tar.gz @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:e418f2a6858c757cb72bd25772749a1664c97a407682d88ad2b51c4bbdcb8006 -size 11568 +oid sha256:b696ca8d86efc625fddaf7dd81fecd16cd0797e008c89b74af69fe9d7ec1e1ac +size 11503 diff --git a/data/.lfs/make_path_mask_full.png.tar.gz b/data/.lfs/make_path_mask_full.png.tar.gz index 0e9336aaea5..c28598c2644 100644 --- a/data/.lfs/make_path_mask_full.png.tar.gz +++ b/data/.lfs/make_path_mask_full.png.tar.gz @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:b772d266dffa82ccf14f13c7d8cc2443210202836883c80f016a56d4cfe2b52a -size 11213 +oid sha256:804c17a3f28ba5525d253ba89b08a7ac38bf0f62cb89b799f85b27fe431192a2 +size 11320 diff --git a/data/.lfs/make_path_mask_two_meters.png.tar.gz b/data/.lfs/make_path_mask_two_meters.png.tar.gz index 7fa9e767b8f..766886698cd 100644 --- a/data/.lfs/make_path_mask_two_meters.png.tar.gz +++ b/data/.lfs/make_path_mask_two_meters.png.tar.gz @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:da608d410f4a1afee0965abfac814bc05267bdde31b0d3a9622c39515ee4f813 -size 11395 +oid sha256:c3f8a5f9371b17229634aa51bf0fe22999d57990676dea9bc866ce3945062e92 +size 11363 diff --git a/data/.lfs/resample_path_simple.png.tar.gz b/data/.lfs/resample_path_simple.png.tar.gz index 1a8c1118d60..ff26518b329 100644 --- a/data/.lfs/resample_path_simple.png.tar.gz +++ b/data/.lfs/resample_path_simple.png.tar.gz @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:0b5c454ed6cc66cf4446ce4a246464aec27368da4902651b4ad9ed29b3ba56ec -size 118319 +oid sha256:4af4f0623cae679abf46246239558b77cead04ec024bd539d31c375c29981a06 +size 117658 diff --git a/data/.lfs/resample_path_smooth.png.tar.gz b/data/.lfs/resample_path_smooth.png.tar.gz index 80af3d3805d..af21c35e372 100644 --- a/data/.lfs/resample_path_smooth.png.tar.gz +++ b/data/.lfs/resample_path_smooth.png.tar.gz @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:6cc0dfd80bada94f2ab1bb577e2ec1734dad6894113f2fe77964bd80d886c3d3 -size 109699 +oid sha256:243ffb720247f543793243305b346e44febaff5398572a7979ed770751587302 +size 109328 diff --git a/data/.lfs/unitree_go2_detection_cdr.tar.gz b/data/.lfs/unitree_go2_detection_cdr.tar.gz new file mode 100644 index 00000000000..cfb22141ba5 --- /dev/null +++ b/data/.lfs/unitree_go2_detection_cdr.tar.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:966a7b4d542a58a7891b7c0fc4951e3ea3975c56f52cd99b0a5f0541198d2d40 +size 5823304 diff --git a/data/.lfs/visualize_occupancy_rainbow.png.tar.gz b/data/.lfs/visualize_occupancy_rainbow.png.tar.gz index 9bbd2e6ea1f..18c4da65865 100644 --- a/data/.lfs/visualize_occupancy_rainbow.png.tar.gz +++ b/data/.lfs/visualize_occupancy_rainbow.png.tar.gz @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:3dc1e3b6519f7d7ff25b16c3124ee447f02857eeb3eb20930cdab95464b1f0a3 -size 11582 +oid sha256:3513b154f81dc5c1b34a74ea517e5789389fbae05f989136393794ee8bba7047 +size 11516 diff --git a/data/.lfs/visualize_occupancy_turbo.png.tar.gz b/data/.lfs/visualize_occupancy_turbo.png.tar.gz index e2863cdae6d..249eb9f97e3 100644 --- a/data/.lfs/visualize_occupancy_turbo.png.tar.gz +++ b/data/.lfs/visualize_occupancy_turbo.png.tar.gz @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:c21874bab6ec7cd9692d2b1e67498ddfff3c832ec992e9552fee17093759b270 -size 18593 +oid sha256:7e28df6a9193dd56081b51a8af3896be6cd4c5dda1deeff44b383e05e462e10b +size 18229 diff --git a/dimos/agents/mcp/mcp_client.py b/dimos/agents/mcp/mcp_client.py index 50554e59427..62d782c2953 100644 --- a/dimos/agents/mcp/mcp_client.py +++ b/dimos/agents/mcp/mcp_client.py @@ -406,6 +406,11 @@ def _process_message( def _append_image_to_history( mcp_client: McpClient, func_name: str, uuid_: str, result: Any ) -> None: + if isinstance(result, dict) and result.get("type") == "image": + result = { + "type": "image_url", + "image_url": {"url": f"data:{result['mimeType']};base64,{result['data']}"}, + } mcp_client.add_message( HumanMessage( content=[ diff --git a/dimos/agents/mcp/mcp_server.py b/dimos/agents/mcp/mcp_server.py index 23cfa90196b..597d4754d44 100644 --- a/dimos/agents/mcp/mcp_server.py +++ b/dimos/agents/mcp/mcp_server.py @@ -14,6 +14,7 @@ from __future__ import annotations import asyncio +import base64 from collections.abc import AsyncGenerator, Callable import concurrent.futures import json @@ -22,6 +23,7 @@ from typing import TYPE_CHECKING, Any import uuid +from dimos_generated.sensor_msgs.msg import Image from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse @@ -36,6 +38,7 @@ from dimos.core.module import Module from dimos.core.rpc_client import RpcCall, RPCClient from dimos.core.transport_factory import make_transport +from dimos.msgs.image import image_to_jpeg from dimos.utils.logging_config import setup_logger from dimos.utils.safe_thread_map import safe_thread_map @@ -221,6 +224,15 @@ def _can_wait(holder: str) -> bool: duration = f"{time.monotonic() - t0:.3f}s" response = str(result)[:200] + if isinstance(result, Image): + content = { + "type": "image", + "mimeType": "image/jpeg", + "data": base64.b64encode(image_to_jpeg(result)).decode("ascii"), + } + logger.info("MCP tool done", tool=name, duration=duration, response=response) + return _jsonrpc_result(req_id, {"content": [content]}) + if hasattr(result, "agent_encode"): logger.info("MCP tool done", tool=name, duration=duration, response=response) return _jsonrpc_result(req_id, {"content": result.agent_encode()}) diff --git a/dimos/agents/mcp/test_generated_image.py b/dimos/agents/mcp/test_generated_image.py new file mode 100644 index 00000000000..3fcc5442df5 --- /dev/null +++ b/dimos/agents/mcp/test_generated_image.py @@ -0,0 +1,56 @@ +# 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. + +"""A generated CDR image survives MCP serialization and agent history conversion.""" + +import asyncio +import base64 + +import cv2 +from dimos_generated.sensor_msgs.msg import Image +import numpy as np + +from dimos.agents.mcp.mcp_client import _append_image_to_history +from dimos.agents.mcp.mcp_server import handle_request +from dimos.msgs.image import image_from_array + + +def test_generated_image_mcp_result_reaches_model_history(mocker): + pixels = np.zeros((12, 16, 3), dtype=np.uint8) + pixels[..., 0] = 255 + original = image_from_array(pixels, encoding="rgb8") + response = asyncio.run( + handle_request( + {"method": "tools/call", "id": 1, "params": {"name": "observe"}}, + [], + {"observe": lambda: Image.decode(original.encode())}, + ) + ) + assert response is not None + content = response["result"]["content"] + assert len(content) == 1 + image = content[0] + assert image["type"] == "image" + assert image["mimeType"] == "image/jpeg" + jpeg = base64.b64decode(image["data"]) + decoded = cv2.imdecode(np.frombuffer(jpeg, dtype=np.uint8), cv2.IMREAD_COLOR) + assert decoded.shape == (12, 16, 3) + assert decoded[0, 0].tolist() == [0, 0, 254] + client = mocker.Mock() + _append_image_to_history(client, "observe", "test-image", image) + message = client.add_message.call_args.args[0] + assert message.content[1] == { + "type": "image_url", + "image_url": {"url": "data:image/jpeg;base64," + image["data"]}, + } diff --git a/dimos/agents/mcp/test_mcp_client.py b/dimos/agents/mcp/test_mcp_client.py index b31637630e6..2be2fcab2e9 100644 --- a/dimos/agents/mcp/test_mcp_client.py +++ b/dimos/agents/mcp/test_mcp_client.py @@ -14,11 +14,12 @@ from typing import Any +from dimos_generated.sensor_msgs.msg import Image from langchain_core.messages import HumanMessage from dimos.agents.annotation import skill from dimos.core.module import Module -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_from_file from dimos.utils.data import get_data @@ -177,7 +178,7 @@ class Visualizer(Module): @skill def take_a_picture(self) -> Image: """Takes a picture.""" - return Image.from_file(get_data("cafe-smol.jpg")).to_rgb() + return image_from_file(get_data("cafe-smol.jpg")) def test_image(agent_setup): diff --git a/dimos/agents/skills/navigation.py b/dimos/agents/skills/navigation.py index e6053703269..ef379788be3 100644 --- a/dimos/agents/skills/navigation.py +++ b/dimos/agents/skills/navigation.py @@ -15,6 +15,8 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.sensor_msgs.msg import Image from reactivex.disposable import Disposable from dimos.agents.annotation import skill @@ -24,10 +26,8 @@ from dimos.core.module import Module from dimos.core.stream import In from dimos.models.qwen.bbox import BBox -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3, make_vector3 -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.geometry import quaternion_euler, quaternion_from_euler +from dimos.msgs.time import header_now from dimos.navigation.base import NavigationState from dimos.navigation.go2.replanning_a_star.spec import NavigationInterfaceSpec from dimos.perception.experimental.object_tracking_spec import ObjectTrackingSpec @@ -96,13 +96,13 @@ def tag_location(self, location_name: str) -> str: if not self._latest_odom: return "No odometry data received yet, cannot tag location." - position = self._latest_odom.position - rotation = self._latest_odom.orientation + position = self._latest_odom.pose.position + rotation = self._latest_odom.pose.orientation location = RobotLocation( name=location_name, position=(position.x, position.y, position.z), - rotation=(rotation.x, rotation.y, rotation.z), + rotation=quaternion_euler(rotation), ) if not self._spatial_memory.tag_location(location): @@ -158,16 +158,22 @@ def _navigate_by_tagged_location(self, query: str) -> str | None: logger.info("Found tagged location", location=robot_location) goal_pose = PoseStamped( - position=make_vector3(*robot_location.position), - orientation=Quaternion.from_euler(Vector3(*robot_location.rotation)), - frame_id="map", + header=header_now("map"), + pose=Pose( + position=Point( + x=robot_location.position[0], + y=robot_location.position[1], + z=robot_location.position[2], + ), + orientation=quaternion_from_euler(*robot_location.rotation), + ), ) return self._navigate_to(goal_pose, f"Found a tagged location called '{query}'.") def _navigate_to(self, pose: PoseStamped, message: str) -> str: logger.info( - f"Navigating to pose: ({pose.position.x:.2f}, {pose.position.y:.2f}, {pose.position.z:.2f})" + f"Navigating to pose: ({pose.pose.position.x:.2f}, {pose.pose.position.y:.2f}, {pose.pose.position.z:.2f})" ) self._navigation.set_goal(pose) @@ -281,7 +287,8 @@ def _get_goal_pose_from_result(self, result: dict[str, Any]) -> PoseStamped | No theta = first.get("rot_z", 0) return PoseStamped( - position=make_vector3(pos_x, pos_y, 0), - orientation=Quaternion.from_euler(make_vector3(0, 0, theta)), - frame_id="map", + header=header_now("map"), + pose=Pose( + position=Point(x=pos_x, y=pos_y), orientation=quaternion_from_euler(0, 0, theta) + ), ) diff --git a/dimos/agents/skills/observe_skill.py b/dimos/agents/skills/observe_skill.py index e6a1c46189a..473ad9672d6 100644 --- a/dimos/agents/skills/observe_skill.py +++ b/dimos/agents/skills/observe_skill.py @@ -12,10 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.sensor_msgs.msg import Image + from dimos.agents.annotation import skill from dimos.core.module import Module from dimos.core.stream import In -from dimos.msgs.sensor_msgs.Image import Image class ObserveSkill(Module): diff --git a/dimos/agents/skills/test_generated_navigation.py b/dimos/agents/skills/test_generated_navigation.py new file mode 100644 index 00000000000..49551742803 --- /dev/null +++ b/dimos/agents/skills/test_generated_navigation.py @@ -0,0 +1,70 @@ +# 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. + +"""Generated navigation goals and named-location orientation round trips.""" + +import math + +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +import pytest + +from dimos.agents.skills.navigation import NavigationSkillContainer +from dimos.msgs.geometry import quaternion_euler, quaternion_from_euler + + +@pytest.fixture +def skill(mocker): + value = NavigationSkillContainer() + memory = mocker.Mock() + navigator = mocker.Mock() + mocker.patch.object(value, "_spatial_memory", memory, create=True) + mocker.patch.object(value, "_navigation", navigator, create=True) + value._skill_started = True + try: + yield value, memory, navigator + finally: + value.stop() + + +def test_tagged_location_preserves_euler_orientation_through_cdr_goal(skill): + value, memory, navigator = skill + angles = (0.2, -0.3, math.pi / 2) + pose = PoseStamped( + pose=Pose(position=Point(x=1, y=2, z=3), orientation=quaternion_from_euler(*angles)) + ) + value._on_odom(PoseStamped.decode(pose.encode())) + assert "Tagged 'desk'" in value.tag_location("desk") + tagged = memory.tag_location.call_args.args[0] + assert tagged.position == (1, 2, 3) + assert tagged.rotation == pytest.approx(angles) + memory.query_tagged_location.return_value = tagged + assert "Found a tagged location" in value._navigate_by_tagged_location("desk") + goal = navigator.set_goal.call_args.args[0] + decoded = PoseStamped.decode(goal.encode()) + assert decoded.header.frame_id == "map" + assert decoded.pose.position == pose.pose.position + assert quaternion_euler(decoded.pose.orientation) == pytest.approx(angles) + + +def test_semantic_result_builds_generated_map_goal(skill): + value, _, navigator = skill + result = {"distance": 0.1, "metadata": [{"pos_x": 4, "pos_y": -2, "rot_z": math.pi / 2}]} + goal = value._get_goal_pose_from_result(result) + assert goal is not None + value._navigate_to(goal, "Found desk") + decoded = PoseStamped.decode(navigator.set_goal.call_args.args[0].encode()) + assert decoded.header.frame_id == "map" + assert decoded.pose.position == Point(x=4, y=-2) + assert quaternion_euler(decoded.pose.orientation) == pytest.approx((0, 0, math.pi / 2)) + assert value._get_goal_pose_from_result({"distance": 0.9}) is None diff --git a/dimos/agents/skills/test_generated_person_follow.py b/dimos/agents/skills/test_generated_person_follow.py new file mode 100644 index 00000000000..8122ba5ea19 --- /dev/null +++ b/dimos/agents/skills/test_generated_person_follow.py @@ -0,0 +1,81 @@ +# 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. + +from types import SimpleNamespace +from unittest.mock import MagicMock + +from dimos_generated.geometry_msgs.msg import Twist +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +import numpy as np +import pytest + +from dimos.agents.skills import person_follow +from dimos.core.global_config import GlobalConfig +from dimos.msgs.image import image_from_array +from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox +from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D + + +@pytest.mark.parametrize("lost", [False, True]) +def test_generated_person_follow_loop_stops_cleanly(monkeypatch, lost): + monkeypatch.setattr(person_follow, "create", lambda _: MagicMock()) + camera = CameraInfo(width=640, height=480, k=[500, 0, 320, 0, 500, 240, 0, 0, 1]) + module = person_follow.PersonFollowSkillContainer( + camera_info=camera, g=GlobalConfig(simulation="") + ) + image = image_from_array(np.zeros((480, 640, 3), dtype=np.uint8), encoding="rgb8") + image = Image.decode(image.encode()) + detection = Detection2DBBox( + bbox=(270, 100, 370, 300), + track_id=1, + class_id=0, + confidence=1, + name="person", + ts=0, + image=image, + ) + tracker = SimpleNamespace( + init_track=lambda **kwargs: ImageDetections2D(image, [detection]), + stop=lambda: None, + process_image=lambda frame: ImageDetections2D(frame, [] if lost else [detection]), + ) + outputs = [] + + def receive(value): + outputs.append(Twist.decode(value.encode())) + if not lost: + module._should_stop.set() + + unsubscribe = module.cmd_vel.subscribe(receive) + module._max_lost_frames = 0 + module._frequency = 1000 + try: + module._on_color_image(image) + module._tracker = tracker + status = module._follow_person("person", (270, 100, 370, 300)) + assert "Starting to follow" in status + thread = module._thread + assert thread is not None + thread.join(timeout=2) + assert not thread.is_alive() + assert outputs + assert outputs[-1] == Twist() + if lost: + assert all(value == Twist() for value in outputs) + else: + assert outputs[0].linear.x == 0.5 + assert outputs[0].angular.z == 0 + finally: + unsubscribe() + module.stop() diff --git a/dimos/agents/skills/test_navigation.py b/dimos/agents/skills/test_navigation.py index f206d63ba04..08206c2c6dc 100644 --- a/dimos/agents/skills/test_navigation.py +++ b/dimos/agents/skills/test_navigation.py @@ -14,14 +14,14 @@ from typing import Any +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import Image from langchain_core.messages import HumanMessage from dimos.agents.skills.navigation import NavigationSkillContainer from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.Image import Image from dimos.navigation.base import NavigationState from dimos.types.robot_location import RobotLocation diff --git a/dimos/agents/skills/test_observe_skill.py b/dimos/agents/skills/test_observe_skill.py index 240efc69dc4..88706d6af20 100644 --- a/dimos/agents/skills/test_observe_skill.py +++ b/dimos/agents/skills/test_observe_skill.py @@ -14,15 +14,15 @@ from collections.abc import Iterator import threading -import time +from dimos_generated.sensor_msgs.msg import Image import numpy as np import pytest from reactivex.scheduler import ThreadPoolScheduler from dimos.agents.skills.observe_skill import ObserveSkill from dimos.core.transport import LCMTransport -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_from_array, image_view @pytest.fixture @@ -34,20 +34,22 @@ def module(monkeypatch: pytest.MonkeyPatch) -> Iterator[ObserveSkill]: monkeypatch.setattr("dimos.utils.reactive.get_scheduler", lambda: scheduler) m = ObserveSkill() m.color_image.transport = LCMTransport("/test_observe/color_image", Image) - yield m - m.stop() - scheduler.executor.shutdown(wait=True) + try: + yield m + finally: + m.stop() + scheduler.executor.shutdown(wait=True) def test_observe_returns_published_frame(module: ObserveSkill) -> None: - frame = Image.from_numpy(np.zeros((8, 8, 3), dtype=np.uint8), format=ImageFormat.RGB, ts=1.0) + frame = image_from_array(np.zeros((8, 8, 3), dtype=np.uint8), encoding="rgb8") stop = threading.Event() # get_next subscribes lazily, so keep publishing until observe picks a frame up. def pump() -> None: while not stop.is_set(): module.color_image.transport.publish(frame) - time.sleep(0.05) + stop.wait(0.05) thread = threading.Thread(target=pump, daemon=True) thread.start() @@ -58,7 +60,7 @@ def pump() -> None: thread.join() assert isinstance(result, Image) - assert result.data.shape[:2] == (8, 8) + assert image_view(result).shape[:2] == (8, 8) def test_observe_without_frames_raises( diff --git a/dimos/agents/skills/test_unitree_skill_container.py b/dimos/agents/skills/test_unitree_skill_container.py index d3d2e72f4ae..278c71eb370 100644 --- a/dimos/agents/skills/test_unitree_skill_container.py +++ b/dimos/agents/skills/test_unitree_skill_container.py @@ -16,14 +16,13 @@ import math from typing import Any +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped from langchain_core.messages import HumanMessage import pytest from dimos.core.core import rpc from dimos.core.module import Module -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.geometry import quaternion_euler, quaternion_from_euler from dimos.navigation.base import NavigationState from dimos.robot.unitree.unitree_skill_container import ( _UNITREE_COMMANDS, @@ -82,23 +81,25 @@ def test_did_you_mean() -> None: def _pose(x: float, y: float, yaw_deg: float) -> PoseStamped: return PoseStamped( - position=Vector3(x, y, 0.3), - orientation=Quaternion.from_euler(Vector3(0, 0, math.radians(yaw_deg))), + pose=Pose( + position=Point(x=x, y=y, z=0.3), + orientation=quaternion_from_euler(0, 0, math.radians(yaw_deg)), + ), ) def _xy(pose: PoseStamped) -> tuple[float, float]: - return pose.position.x, pose.position.y + return pose.pose.position.x, pose.pose.position.y def _yaw_deg(pose: PoseStamped) -> float: - return math.degrees(pose.orientation.to_euler().yaw) + return math.degrees(quaternion_euler(pose.pose.orientation)[2]) def test_move_to_world_keeps_heading() -> None: goal = _goal_pose(_pose(1, 1, 90), 4, 1, None, relative=False) assert _xy(goal) == pytest.approx((4, 1)) - assert goal.position.z == pytest.approx(0.3) + assert goal.pose.position.z == pytest.approx(0.3) assert _yaw_deg(goal) == pytest.approx(90) diff --git a/dimos/agents/skills/visual_servoing/bbox_navigation.py b/dimos/agents/skills/visual_servoing/bbox_navigation.py index c5511cf7d00..de56e99ba18 100644 --- a/dimos/agents/skills/visual_servoing/bbox_navigation.py +++ b/dimos/agents/skills/visual_servoing/bbox_navigation.py @@ -14,16 +14,14 @@ import logging -from dimos_lcm.sensor_msgs import CameraInfo +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import CameraInfo +from dimos_generated.vision_msgs.msg import Detection2DArray from reactivex.disposable import Disposable from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray from dimos.utils.logging_config import setup_logger logger = setup_logger(level=logging.DEBUG) @@ -46,7 +44,7 @@ class BBoxNavigationModule(Module): @rpc def start(self) -> None: unsub = self.camera_info.subscribe( - lambda msg: setattr(self, "camera_intrinsics", [msg.K[0], msg.K[4], msg.K[2], msg.K[5]]) + lambda msg: setattr(self, "camera_intrinsics", [msg.k[0], msg.k[4], msg.k[2], msg.k[5]]) ) self.register_disposable(Disposable(unsub)) @@ -58,7 +56,7 @@ def stop(self) -> None: super().stop() def _on_detection(self, det: Detection2DArray) -> None: - if det.detections_length == 0 or not self.camera_intrinsics: + if not det.detections or not self.camera_intrinsics: return fx, fy, cx, cy = self.camera_intrinsics center_x, center_y = ( @@ -71,9 +69,8 @@ def _on_detection(self, det: Detection2DArray) -> None: self.config.goal_distance, ) goal = PoseStamped( - position=Vector3(z, -x, -y), - orientation=Quaternion(0, 0, 0, 1), - frame_id=det.header.frame_id, + header=det.header, + pose=Pose(position=Point(x=z, y=-x, z=-y), orientation=Quaternion(w=1)), ) logger.debug( f"BBox center: ({center_x:.1f}, {center_y:.1f}) → " diff --git a/dimos/agents/skills/visual_servoing/detection_navigation.py b/dimos/agents/skills/visual_servoing/detection_navigation.py index 7430a23a883..16b66c692c8 100644 --- a/dimos/agents/skills/visual_servoing/detection_navigation.py +++ b/dimos/agents/skills/visual_servoing/detection_navigation.py @@ -12,15 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos_lcm.sensor_msgs import CameraInfo as DimosLcmCameraInfo +from dimos_generated.geometry_msgs.msg import TransformStamped, Twist, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 import numpy as np -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.geometry import quaternion_euler +from dimos.msgs.pointcloud import pointcloud_xyz +from dimos.msgs.time import to_seconds from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC from dimos.protocol.tf.tf import TFLookup @@ -59,22 +57,20 @@ def compute_twist_for_detection_3d( # Get transform from world frame to camera optical frame world_to_optical = self._tf.get( - "camera_optical", pointcloud.frame_id, image.ts, time_tolerance=1.0 + "camera_optical", + pointcloud.header.frame_id, + to_seconds(image.header.stamp), + time_tolerance=1.0, ) if world_to_optical is None: logger.warning("Could not get camera transform") return None - lcm_camera_info = DimosLcmCameraInfo() - lcm_camera_info.K = self._camera_info.K - lcm_camera_info.width = self._camera_info.width - lcm_camera_info.height = self._camera_info.height - # Project to 3D using the pointcloud detection_3d = Detection3DPC.from_2d( det=detection, world_pointcloud=pointcloud, - camera_info=lcm_camera_info, + camera_info=self._camera_info, world_to_optical_transform=world_to_optical, filters=[], # Skip filtering for faster processing in follow loop ) @@ -89,7 +85,7 @@ def compute_twist_for_detection_3d( logger.warning("Could not get robot transform") return None - robot_pos = robot_transform.translation + robot_pos = robot_transform.transform.translation # Compute robust target position using front-most points target_position = self._compute_robust_target_position(detection_3d.pointcloud, robot_pos) @@ -119,7 +115,7 @@ def _compute_robust_target_position( Vector3 position representing the front of the detected object, or None if not enough valid points. """ - points, _ = pointcloud.as_numpy() + points = pointcloud_xyz(pointcloud) if len(points) < 10: return None @@ -128,7 +124,7 @@ def _compute_robust_target_position( points = points[height_mask] if len(points) < 10: # Fall back to all points if height filtering removes too many - points, _ = pointcloud.as_numpy() + points = pointcloud_xyz(pointcloud) # Compute 2D distance (XY plane) from robot to each point dx = points[:, 0] - robot_pos.x @@ -152,9 +148,11 @@ def _compute_robust_target_position( # Compute centroid of front-most points centroid = front_points.mean(axis=0) - return Vector3(centroid[0], centroid[1], centroid[2]) + return Vector3(x=float(centroid[0]), y=float(centroid[1]), z=float(centroid[2])) - def _compute_twist_from_3d(self, target_position: Vector3, robot_transform: Transform) -> Twist: + def _compute_twist_from_3d( + self, target_position: Vector3, robot_transform: TransformStamped + ) -> Twist: """Compute twist command to navigate towards a 3D target position. Args: @@ -164,7 +162,7 @@ def _compute_twist_from_3d(self, target_position: Vector3, robot_transform: Tran Returns: Twist command for the robot. """ - robot_pos = robot_transform.translation + robot_pos = robot_transform.transform.translation # Compute vector from robot to target in world frame dx = target_position.x - robot_pos.x @@ -176,7 +174,7 @@ def _compute_twist_from_3d(self, target_position: Vector3, robot_transform: Tran angle_to_target = np.arctan2(dy, dx) # Get robot's current heading from transform - robot_yaw = robot_transform.rotation.to_euler().z + robot_yaw = quaternion_euler(robot_transform.transform.rotation)[2] # Angle error (how much to turn) angle_error = angle_to_target - robot_yaw @@ -207,6 +205,6 @@ def _compute_twist_from_3d(self, target_position: Vector3, robot_transform: Tran ) return Twist( - linear=Vector3(linear_x, 0.0, 0.0), - angular=Vector3(0.0, 0.0, angular_z), + linear=Vector3(x=linear_x), + angular=Vector3(z=angular_z), ) diff --git a/dimos/agents/skills/visual_servoing/query.py b/dimos/agents/skills/visual_servoing/query.py index 16a6ca60ce8..5be1b953c55 100644 --- a/dimos/agents/skills/visual_servoing/query.py +++ b/dimos/agents/skills/visual_servoing/query.py @@ -12,9 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.sensor_msgs.msg import Image + from dimos.models.qwen.bbox import BBox from dimos.models.vl.base import VlModel -from dimos.msgs.sensor_msgs.Image import Image from dimos.utils.generic import extract_json_from_llm_response diff --git a/dimos/agents/skills/visual_servoing/visual_servoing_2d.py b/dimos/agents/skills/visual_servoing/visual_servoing_2d.py index f424b214661..fe209e7f95b 100644 --- a/dimos/agents/skills/visual_servoing/visual_servoing_2d.py +++ b/dimos/agents/skills/visual_servoing/visual_servoing_2d.py @@ -12,12 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.geometry_msgs.msg import Twist, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo import numpy as np -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo - class VisualServoing2D: """2D visual servoing controller for tracking objects using bounding boxes. @@ -91,7 +89,7 @@ def compute_twist( estimated_distance = self._estimate_distance(bbox) if estimated_distance is None: - return Twist.zero() + return Twist() # Calculate distance error (positive = too far, need to move forward) distance_error = estimated_distance - self._target_distance @@ -120,8 +118,8 @@ def compute_twist( linear_x = max(linear_x, 0.1) return Twist( - linear=Vector3(linear_x, 0.0, 0.0), - angular=Vector3(0.0, 0.0, angular_z), + linear=Vector3(x=linear_x), + angular=Vector3(z=angular_z), ) def _get_normalized_x(self, pixel_x: float) -> float: @@ -135,8 +133,8 @@ def _get_normalized_x(self, pixel_x: float) -> float: Returns: Normalized x coordinate (tan of angle from optical center) """ - fx = self._camera_info.K[0] # focal length x - cx = self._camera_info.K[2] # optical center x + fx = self._camera_info.k[0] # focal length x + cx = self._camera_info.k[2] # optical center x return (pixel_x - cx) / fx def _estimate_distance(self, bbox: tuple[float, float, float, float]) -> float | None: @@ -161,7 +159,7 @@ def _estimate_distance(self, bbox: tuple[float, float, float, float]) -> float | return None # Pinhole camera model: distance = (real_width * fx) / pixel_width - fx = self._camera_info.K[0] # focal length x in pixels + fx = self._camera_info.k[0] # focal length x in pixels estimated_distance = (self._assumed_object_width * fx) / bbox_width return estimated_distance diff --git a/dimos/agents/test_generated_vlm.py b/dimos/agents/test_generated_vlm.py new file mode 100644 index 00000000000..2a456afca18 --- /dev/null +++ b/dimos/agents/test_generated_vlm.py @@ -0,0 +1,80 @@ +# 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. + +"""Verify generated camera messages reach VLM requests through stream and RPC.""" + +import base64 + +import cv2 +from dimos_generated.sensor_msgs.msg import Image +from langchain_core.messages import AIMessage, HumanMessage +import numpy as np +import pytest + +from dimos.agents.vlm_agent import VLMAgent +from dimos.msgs.image import image_from_array + + +@pytest.fixture +def agent(mocker): + model = mocker.Mock() + model.invoke.return_value = AIMessage(content="red image") + mocker.patch("langchain.chat_models.init_chat_model", return_value=model) + value = VLMAgent() + try: + yield value, model + finally: + value.stop() + + +@pytest.mark.parametrize("use_stream", [False, True]) +def test_generated_image_reaches_model_as_jpeg(agent, use_stream): + value, model = agent + pixels = np.zeros((12, 16, 3), dtype=np.uint8) + pixels[..., 0] = 255 + original = image_from_array(pixels, encoding="rgb8") + image = Image.decode(original.encode()) + answers = [] + unsubscribe = value.answer_stream.subscribe(answers.append) + try: + if use_stream: + value._on_image(image) + value._on_query(HumanMessage(content="What color?")) + assert answers[0].content == "red image" + else: + assert value.query_image(image, "What color?") == "red image" + request = model.invoke.call_args.args[0][-1] + assert request.content[0] == {"type": "text", "text": "What color?"} + prefix, encoded = request.content[1]["image_url"]["url"].split(",", 1) + assert prefix == "data:image/jpeg;base64" + decoded = cv2.imdecode( + np.frombuffer(base64.b64decode(encoded), dtype=np.uint8), cv2.IMREAD_COLOR + ) + assert decoded.shape == (12, 16, 3) + assert decoded[0, 0].tolist() == [0, 0, 254] + assert image.encode() == original.encode() + finally: + unsubscribe() + + +def test_query_before_image_does_not_invoke_model(agent): + value, model = agent + answers = [] + unsubscribe = value.answer_stream.subscribe(answers.append) + try: + value._on_query(HumanMessage(content="What color?")) + assert answers[0].content == "No image available yet." + model.invoke.assert_not_called() + finally: + unsubscribe() diff --git a/dimos/agents/testing/vlm_stream_tester.py b/dimos/agents/testing/vlm_stream_tester.py index fecabed6ccf..d779bc16246 100644 --- a/dimos/agents/testing/vlm_stream_tester.py +++ b/dimos/agents/testing/vlm_stream_tester.py @@ -16,6 +16,7 @@ import time from typing import Any +from dimos_generated.sensor_msgs.msg import Image from langchain_core.messages import AIMessage, HumanMessage from reactivex.disposable import Disposable @@ -24,7 +25,6 @@ from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import In, Out -from dimos.msgs.sensor_msgs.Image import Image from dimos.utils.logging_config import setup_logger logger = setup_logger() diff --git a/dimos/agents/vlm_agent.py b/dimos/agents/vlm_agent.py index ceef36d93cf..30c0cffd291 100644 --- a/dimos/agents/vlm_agent.py +++ b/dimos/agents/vlm_agent.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +import base64 from typing import TYPE_CHECKING, Any +from dimos_generated.sensor_msgs.msg import Image from langchain_core.messages import AIMessage, HumanMessage, SystemMessage from reactivex.disposable import Disposable @@ -21,7 +23,7 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_to_jpeg from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -102,7 +104,11 @@ def _invoke(self, msg: HumanMessage, **kwargs: Any) -> AIMessage: def _invoke_image( self, image: Image, query: str, response_format: dict[str, Any] | None = None ) -> AIMessage: - content = [{"type": "text", "text": query}, *image.agent_encode()] + encoded = base64.b64encode(image_to_jpeg(image)).decode("ascii") + content: list[str | dict[str, Any]] = [ + {"type": "text", "text": query}, + {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encoded}"}}, + ] kwargs: dict[str, Any] = {} if response_format: kwargs["response_format"] = response_format diff --git a/dimos/agents/vlm_agent_spec.py b/dimos/agents/vlm_agent_spec.py index a1bebbfe20a..5bca2b7ddef 100644 --- a/dimos/agents/vlm_agent_spec.py +++ b/dimos/agents/vlm_agent_spec.py @@ -14,7 +14,8 @@ from typing import Any, Protocol -from dimos.msgs.sensor_msgs.Image import Image +from dimos_generated.sensor_msgs.msg import Image + from dimos.spec.utils import Spec diff --git a/dimos/cli/bake/test_bake_e2e.py b/dimos/cli/bake/test_bake_e2e.py index 7da84bd4e99..8c7c5ace192 100644 --- a/dimos/cli/bake/test_bake_e2e.py +++ b/dimos/cli/bake/test_bake_e2e.py @@ -27,14 +27,17 @@ import subprocess import time +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage import numpy as np import pytest from dimos.core.global_config import global_config from dimos.core.transport_factory import make_transport -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import time_from_nanoseconds pytestmark = pytest.mark.bake_e2e @@ -110,14 +113,16 @@ def await_listener(proc: subprocess.Popen[bytes], port: int, timeout: float = 20 raise AssertionError(f"the host never listened on {port}") -def ground_patch(ts: float) -> PointCloud2: +def ground_patch(stamp_ns: int) -> PointCloud2: """A flat floor patch below the sensor, enough for the planner to surface.""" grid = np.arange(-10, 11) * 0.1 xs, ys = np.meshgrid(grid, grid) points = np.stack([xs.ravel(), ys.ravel(), np.full(xs.size, -SENSOR_Z)], axis=1).astype( np.float32 ) - return PointCloud2.from_numpy(points, frame_id="lidar", timestamp=ts) + return pointcloud_from_xyz( + points, header=Header(stamp=time_from_nanoseconds(stamp_ns), frame_id="lidar") + ) @pytest.fixture @@ -155,7 +160,7 @@ def test_the_host_publishes_its_outputs_and_hides_the_suppressed_hop(baked, zeno transports.append(transport) lidar = make_transport("lidar", PointCloud2) - tf = make_transport("tf", Transform) + tf = make_transport("tf", TFMessage) transports += [lidar, tf] lidar.start() tf.start() @@ -163,18 +168,23 @@ def test_the_host_publishes_its_outputs_and_hides_the_suppressed_hop(baked, zeno deadline = time.monotonic() + 60 while time.monotonic() < deadline and "surface_map" not in seen: # Same stamp on both: the mapper drops a cloud it has no transform for. - ts = time.time() + stamp_ns = time.time_ns() tf.broadcast( None, - Transform( - translation=Vector3(0.0, 0.0, SENSOR_Z), - frame_id="odom", - child_frame_id="lidar", - ts=ts, + TFMessage( + transforms=[ + TransformStamped( + header=Header(stamp=time_from_nanoseconds(stamp_ns), frame_id="odom"), + child_frame_id="lidar", + transform=Transform( + translation=Vector3(z=SENSOR_Z), rotation=Quaternion(w=1.0) + ), + ) + ] ), ) time.sleep(0.05) - lidar.broadcast(None, ground_patch(ts)) + lidar.broadcast(None, ground_patch(stamp_ns)) time.sleep(0.25) assert proc.poll() is None, "the host exited before publishing anything" diff --git a/dimos/cli/bake/test_cli.py b/dimos/cli/bake/test_cli.py index 438cc02e409..2de7afd7c05 100644 --- a/dimos/cli/bake/test_cli.py +++ b/dimos/cli/bake/test_cli.py @@ -113,7 +113,7 @@ def test_emit_config_is_a_complete_stdin_blob() -> None: for section in blob["modules"].values(): assert isinstance(section["config"], dict) and section["config"] assert blob["modules"]["ray_tracing"]["topics"]["lidar"] == ( - "dimos/lidar/sensor_msgs.PointCloud2" + "dimos/lidar/sensor_msgs/msg/PointCloud2" ) assert blob["suppress"] == list(graph.suppressed_topics()) assert blob["qos"] == graph.qos() diff --git a/dimos/cli/bake/test_codegen.py b/dimos/cli/bake/test_codegen.py index ff7ed602dfe..c631fbe6031 100644 --- a/dimos/cli/bake/test_codegen.py +++ b/dimos/cli/bake/test_codegen.py @@ -64,7 +64,7 @@ def test_main_rs_lists_entries_with_their_thread_counts() -> None: main = render_main_rs("go2-nav", [MAPPER, threaded], GRAPH) assert ' ModuleEntry::new::("mapper"),' in main assert ' ModuleEntry::new::("planner").threads(2),' in main - assert ' "dimos/local_map/sensor_msgs.PointCloud2",' in main + assert ' "dimos/local_map/sensor_msgs/msg/PointCloud2",' in main assert main.startswith("// GENERATED BY dimos bake") assert 'name: "go2-nav",' in main assert 'include_str!("graph.json")' in main @@ -87,11 +87,11 @@ def test_generate_crate_writes_the_blobs_the_host_includes(tmp_path: Path) -> No assert crate == tmp_path / "build" / "dimos-bake" / "go2-nav" topics = json.loads((crate / "src" / "default_topics.json").read_text()) - assert topics["mapper"]["lidar"] == "dimos/lidar/sensor_msgs.PointCloud2" - assert topics["planner"]["local_map"] == "dimos/local_map/sensor_msgs.PointCloud2" + assert topics["mapper"]["lidar"] == "dimos/lidar/sensor_msgs/msg/PointCloud2" + assert topics["planner"]["local_map"] == "dimos/local_map/sensor_msgs/msg/PointCloud2" qos = json.loads((crate / "src" / "default_qos.json").read_text()) - assert qos["dimos/local_map/sensor_msgs.PointCloud2"]["congestion_control"] == "drop" + assert qos["dimos/local_map/sensor_msgs/msg/PointCloud2"]["congestion_control"] == "drop" graph = json.loads((crate / "src" / "graph.json").read_text()) assert graph["host"] == "go2-nav" diff --git a/dimos/cli/bake/test_discovery.py b/dimos/cli/bake/test_discovery.py index e707e9f714e..f54f3de4bc8 100644 --- a/dimos/cli/bake/test_discovery.py +++ b/dimos/cli/bake/test_discovery.py @@ -41,10 +41,10 @@ threads = 2 [package.metadata.dimos.module.demo.inputs] -lidar = "sensor_msgs.PointCloud2" +lidar = "sensor_msgs/msg/PointCloud2" [package.metadata.dimos.module.demo.outputs] -global_map = "sensor_msgs.PointCloud2" +global_map = "sensor_msgs/msg/PointCloud2" """ @@ -62,8 +62,8 @@ def test_parses_a_registered_module(tmp_path: Path) -> None: assert info.rust_path == "demo::module::Demo" assert info.python_ref == "dimos.demo.module:Demo" assert info.threads == 2 - assert info.inputs == {"lidar": "sensor_msgs.PointCloud2"} - assert info.outputs == {"global_map": "sensor_msgs.PointCloud2"} + assert info.inputs == {"lidar": "sensor_msgs/msg/PointCloud2"} + assert info.outputs == {"global_map": "sensor_msgs/msg/PointCloud2"} def test_a_crate_without_the_metadata_table_is_not_a_module(tmp_path: Path) -> None: @@ -84,7 +84,7 @@ def test_a_missing_required_key_is_an_error(tmp_path: Path) -> None: def test_a_non_string_msg_type_is_an_error(tmp_path: Path) -> None: - manifest = MANIFEST.replace('lidar = "sensor_msgs.PointCloud2"', "lidar = 3") + manifest = MANIFEST.replace('lidar = "sensor_msgs/msg/PointCloud2"', "lidar = 3") with pytest.raises(BakeError, match="lidar"): parse_manifest(write_crate(tmp_path, "dimos/demo/rust", manifest)) diff --git a/dimos/cli/bake/test_graph.py b/dimos/cli/bake/test_graph.py index 1641770470d..a3a9aa29464 100644 --- a/dimos/cli/bake/test_graph.py +++ b/dimos/cli/bake/test_graph.py @@ -22,8 +22,8 @@ from dimos.cli.bake.errors import BakeError from dimos.cli.bake.graph import Connection, Graph, build_graph, parse_remap, render -PC2 = "sensor_msgs.PointCloud2" -POSE = "geometry_msgs.PoseStamped" +PC2 = "sensor_msgs/msg/PointCloud2" +POSE = "geometry_msgs/msg/PoseStamped" def module( @@ -155,20 +155,20 @@ def test_rendered_graph_is_grouped_by_direction() -> None: "Host `host`: mapper, planner", "", "Internal connections:", - " dimos/local_map/sensor_msgs.PointCloud2 sensor_msgs.PointCloud2 [SUPPRESSED]", + " dimos/local_map/sensor_msgs/msg/PointCloud2 sensor_msgs/msg/PointCloud2 [SUPPRESSED]", " out mapper.local_map", " in planner.local_map", "", "External inputs (subscribed):", - " dimos/goal/geometry_msgs.PoseStamped geometry_msgs.PoseStamped", + " dimos/goal/geometry_msgs/msg/PoseStamped geometry_msgs/msg/PoseStamped", " in planner.goal", - " dimos/lidar/sensor_msgs.PointCloud2 sensor_msgs.PointCloud2", + " dimos/lidar/sensor_msgs/msg/PointCloud2 sensor_msgs/msg/PointCloud2", " in mapper.lidar", "", "External outputs (published):", - " dimos/path/geometry_msgs.PoseStamped geometry_msgs.PoseStamped", + " dimos/path/geometry_msgs/msg/PoseStamped geometry_msgs/msg/PoseStamped", " out planner.path", - " dimos/region_bounds/geometry_msgs.PoseStamped geometry_msgs.PoseStamped", + " dimos/region_bounds/geometry_msgs/msg/PoseStamped geometry_msgs/msg/PoseStamped", " out mapper.region_bounds", ] diff --git a/dimos/cli/hardware/g1.py b/dimos/cli/hardware/g1.py index 636c60c41e1..b41ab602078 100644 --- a/dimos/cli/hardware/g1.py +++ b/dimos/cli/hardware/g1.py @@ -19,10 +19,10 @@ import time from typing import Any, NoReturn, Protocol, TypeGuard +from dimos_generated.sensor_msgs.msg import JointState import typer from dimos.control.tasks.trajectory_task.trajectory_task import JOINT_TRAJECTORY_TASK_NAME -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.porcelain.dimos import Dimos from dimos.porcelain.module_handle import ModuleHandle from dimos.robot.unitree.g1.ready_pose import G1_READY_JOINTS, G1_READY_SPEED_SCALE diff --git a/dimos/cli/spy/test_spy.py b/dimos/cli/spy/test_spy.py index 130c393fca3..11fe0d18dea 100644 --- a/dimos/cli/spy/test_spy.py +++ b/dimos/cli/spy/test_spy.py @@ -26,6 +26,7 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import Vector3 import pytest from dimos.cli.spy.core import ( @@ -39,13 +40,12 @@ default_sources, split_type_suffix, ) -from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.protocol.pubsub.impl.lcmpubsub import LCMPubSubBase, Topic from dimos.protocol.pubsub.impl.zenohpubsub import ZenohPubSubBase from dimos.protocol.service.zenohservice import ZenohSessionPool -VEC = Vector3(1.0, 2.0, 3.0) -VEC_BYTES = VEC.lcm_encode() +VEC = Vector3(x=1.0, y=2.0, z=3.0) +VEC_BYTES = VEC.encode() # TopicStats: pure, deterministic (injected timestamps, no sleeps) @@ -106,10 +106,13 @@ def test_topic_stats_history_eviction_keeps_totals(): def test_split_type_suffix(): - assert split_type_suffix("/cmd_vel#geometry_msgs.Twist") == ("/cmd_vel", "geometry_msgs.Twist") - assert split_type_suffix("dimos/cmd_vel#geometry_msgs.Twist") == ( + assert split_type_suffix("/cmd_vel#geometry_msgs/msg/Twist") == ( + "/cmd_vel", + "geometry_msgs/msg/Twist", + ) + assert split_type_suffix("dimos/cmd_vel#geometry_msgs/msg/Twist") == ( "dimos/cmd_vel", - "geometry_msgs.Twist", + "geometry_msgs/msg/Twist", ) assert split_type_suffix("/plain") == ("/plain", None) @@ -198,7 +201,7 @@ def _assert_no_decode(monkeypatch): def boom(*a, **k): raise AssertionError("spy decoded a payload on the hot path") - monkeypatch.setattr(Vector3, "lcm_decode", staticmethod(boom)) + monkeypatch.setattr(Vector3, "decode", staticmethod(boom)) @pytest.fixture @@ -324,15 +327,15 @@ def test_transport_spy_merges_sources_and_totals(): spy.start() assert a.started and b.started - a.emit("/t#geometry_msgs.Twist", 10) - a.emit("/t#geometry_msgs.Twist", 10) - b.emit("dimos/t#geometry_msgs.Twist", 20) + a.emit("/t#geometry_msgs/msg/Twist", 10) + a.emit("/t#geometry_msgs/msg/Twist", 10) + b.emit("dimos/t#geometry_msgs/msg/Twist", 20) a.emit("/u", 5) snap = spy.snapshot() - assert snap[SpyKey("lcm", "/t#geometry_msgs.Twist")].total_msgs == 2 - assert snap[SpyKey("lcm", "/t#geometry_msgs.Twist")].total_bytes == 20 - assert snap[SpyKey("zenoh", "dimos/t#geometry_msgs.Twist")].total_msgs == 1 + assert snap[SpyKey("lcm", "/t#geometry_msgs/msg/Twist")].total_msgs == 2 + assert snap[SpyKey("lcm", "/t#geometry_msgs/msg/Twist")].total_bytes == 20 + assert snap[SpyKey("zenoh", "dimos/t#geometry_msgs/msg/Twist")].total_msgs == 1 assert snap[SpyKey("lcm", "/u")].total_bytes == 5 assert spy.totals.total_msgs == 4 assert spy.totals.total_bytes == 45 @@ -448,4 +451,4 @@ def test_fake_source_satisfies_protocol(): def test_subscribe_decoded_is_not_implemented_in_v1(): src = LCMSpySource() with pytest.raises(NotImplementedError): - src.subscribe_decoded("/x#geometry_msgs.Vector3", lambda m: None) + src.subscribe_decoded("/x#geometry_msgs/msg/Vector3", lambda m: None) diff --git a/dimos/cli/test_topic.py b/dimos/cli/test_topic.py index 1e35815ef9d..ed7340bb1da 100644 --- a/dimos/cli/test_topic.py +++ b/dimos/cli/test_topic.py @@ -12,15 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.geometry_msgs.msg import Twist, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo +from dimos_generated.std_msgs.msg import Header import pytest from dimos.cli.topic import _build_eval_context, _decode_typed_lcm_message, topic_send from dimos.core.global_config import global_config from dimos.core.transport import PubSubTransport from dimos.core.transport_factory import make_transport -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo from dimos.protocol.pubsub.impl.lcmpubsub import LCM, Topic from dimos.utils.testing.collector import CallbackCollector @@ -30,18 +30,18 @@ def test_decode_typed_lcm_message_resolves_message_submodule() -> None: width=1920, height=1080, distortion_model="plumb_bob", - frame_id="camera_optical", + header=Header(frame_id="camera_optical"), ) decoded = _decode_typed_lcm_message( - "/camera_info#sensor_msgs.CameraInfo", - msg.lcm_encode(), + "/camera_info#sensor_msgs/msg/CameraInfo", + msg.encode(), ) assert isinstance(decoded, CameraInfo) assert decoded.width == 1920 assert decoded.height == 1080 - assert decoded.frame_id == "camera_optical" + assert decoded.header.frame_id == "camera_optical" assert decoded.distortion_model == "plumb_bob" @@ -71,14 +71,14 @@ def capturing_make_transport(name: str, msg_type: type) -> PubSubTransport[objec bus = LCM(url=lcm_url) bus.start() collector = CallbackCollector(1) - bus.subscribe(Topic(topic="/test_topic_send", lcm_type=Twist), collector) + bus.subscribe(Topic(topic="/test_topic_send", msg_type=Twist), collector) try: - topic_send("/test_topic_send", "Twist(Vector3(0.5, 0, 0), Vector3(0, 0, 0))") + topic_send("/test_topic_send", "Twist(linear=Vector3(x=0.5), angular=Vector3())") collector.wait() finally: for transport in transports: transport.stop() bus.stop() - assert collector.results[0][0] == Twist(Vector3(0.5, 0, 0), Vector3(0, 0, 0)) + assert collector.results[0][0] == Twist(linear=Vector3(x=0.5), angular=Vector3()) diff --git a/dimos/cli/topic.py b/dimos/cli/topic.py index 1e5dd05dbd7..d9896ad58b2 100644 --- a/dimos/cli/topic.py +++ b/dimos/cli/topic.py @@ -15,49 +15,40 @@ from __future__ import annotations from collections.abc import Callable -import importlib -import pkgutil import re import time +from types import SimpleNamespace +from typing import cast import typer from dimos.core.global_config import global_config from dimos.core.transport import PubSubTransport from dimos.core.transport_factory import make_transport, transport_topic +from dimos.message_codegen.registry import message_types +from dimos.msgs.helpers import resolve_msg_type from dimos.protocol.pubsub.impl.lcmpubsub import LCMPubSubBase, Topic from dimos.protocol.pubsub.impl.zenohpubsub import Zenoh -_modules_to_try = [ - "dimos.msgs.geometry_msgs", - "dimos.msgs.nav_msgs", - "dimos.msgs.sensor_msgs", - "dimos.msgs.std_msgs", - "dimos.msgs.vision_msgs", - "dimos.msgs.tf2_msgs", -] - def _resolve_type(type_name: str) -> type: - for module_name in _modules_to_try: - try: - 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") + types = message_types() + if type_name in types: + return cast("type", types[type_name]) + matches = [value for name, value in types.items() if name.rsplit("/", 1)[-1] == type_name] + if len(matches) == 1: + return cast("type", matches[0]) + if matches: + raise ValueError(f"Ambiguous message type {type_name!r}; use package/msg/Type") + raise ValueError(f"Unknown installed message type {type_name!r}") def _decode_typed_lcm_message(channel: str, data: bytes) -> object: - from dimos.msgs.helpers import resolve_msg_type - - _, msg_name = channel.split("#", 1) # e.g. "nav_msgs.Odometry" + _, msg_name = channel.split("#", 1) # e.g. "nav_msgs/msg/Odometry" cls = resolve_msg_type(msg_name) if cls is None: raise ValueError(f"Could not resolve message type from channel: {channel}") - return cls.lcm_decode(data) + return cls.decode(data) def _listen_forever(listening_msg: str, on_stop: Callable[[], None] = lambda: None) -> None: @@ -72,7 +63,7 @@ def _listen_forever(listening_msg: str, on_stop: Callable[[], None] = lambda: No def topic_echo(topic: str, type_name: str | None) -> None: - # Explicit mode (legacy): backend chosen by make_transport from global_config. + # Explicit message type: backend chosen by make_transport from global_config. if type_name is not None: msg_type = _resolve_type(type_name) transport: PubSubTransport[object] = make_transport(topic, msg_type) @@ -94,7 +85,7 @@ def _topic_echo_inferred_lcm(topic: str) -> None: autoconf(check_only=True) - # Listen on /topic#pkg.Msg and decode from the msg_name suffix. + # Listen on /topic#pkg/msg/Type and decode from the msg_name suffix. bus = LCMPubSubBase() bus.start() # starts threaded handle loop @@ -107,7 +98,7 @@ def on_msg(channel: str, data: bytes) -> None: bus.l.subscribe(typed_pattern, on_msg) _listen_forever( - f"Listening on {topic} (inferring from typed LCM channels like '{topic}#pkg.Msg')... " + f"Listening on {topic} (inferring from typed LCM channels like '{topic}#pkg/msg/Type')... " "(Ctrl+C to stop)", bus.stop, ) @@ -118,37 +109,32 @@ def _topic_echo_inferred_zenoh(topic: str) -> None: bus = Zenoh() bus.start() - # Typed Zenoh keys embed the type as a trailing segment ("dimos/topic/pkg.Msg"); + # Typed Zenoh keys embed the type as trailing segments ("dimos/topic/pkg/msg/Type"); # a wildcard subscription decodes each message from that suffix. Untyped keys # don't resolve to a type and are skipped by the encoder. The ignore reflects the # pattern Topic vs the encoder's concrete-topic protocol (see lcmpubsub.py). bus.subscribe(Topic(f"{key}/**"), lambda msg, _topic: print(msg)) # type: ignore[arg-type] _listen_forever( - f"Listening on {topic} (inferring from typed Zenoh keys like '{key}/pkg.Msg')... " + f"Listening on {topic} (inferring from typed Zenoh keys like '{key}/pkg/msg/Type')... " "(Ctrl+C to stop)", bus.stop, ) def _build_eval_context() -> dict[str, object]: - # The msgs packages are namespace packages (no __init__.py), so walk their - # submodules; each message file defines a class of the same name. - eval_context: dict[str, object] = {} - for package_name in _modules_to_try: - package = importlib.import_module(package_name) - for module_info in pkgutil.iter_modules(package.__path__): - name = module_info.name - if name.startswith("test_"): - continue - try: - submodule = importlib.import_module(f"{package_name}.{name}") - except ImportError: - continue - obj = getattr(submodule, name, None) - if obj is not None: - eval_context[name] = obj - return eval_context + types = message_types() + context: dict[str, object] = {} + packages: dict[str, SimpleNamespace] = {} + short_names: dict[str, list[type]] = {} + for name, cls in types.items(): + package, _, short = name.split("/") + namespace = packages.setdefault(package, SimpleNamespace(msg=SimpleNamespace())) + setattr(namespace.msg, short, cls) + short_names.setdefault(short, []).append(cls) + context.update(packages) + context.update({name: classes[0] for name, classes in short_names.items() if len(classes) == 1}) + return context def topic_send(topic: str, message_expr: str) -> None: diff --git a/dimos/control/_control_test_helpers.py b/dimos/control/_control_test_helpers.py index 0b925b41392..b9460c95802 100644 --- a/dimos/control/_control_test_helpers.py +++ b/dimos/control/_control_test_helpers.py @@ -23,14 +23,14 @@ from typing import Any +from dimos_generated.geometry_msgs.msg import Pose, PoseStamped + from dimos.control.task import ( BaseControlTask, CoordinatorState, JointCommandOutput, ResourceClaim, ) -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped class RecordingTask(BaseControlTask): diff --git a/dimos/control/benchmarking/benchmark.py b/dimos/control/benchmarking/benchmark.py index fd22f507760..9de58d99463 100644 --- a/dimos/control/benchmarking/benchmark.py +++ b/dimos/control/benchmarking/benchmark.py @@ -51,6 +51,9 @@ import time from typing import Any, Literal +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Twist +from dimos_generated.nav_msgs.msg import Path as NavPath +from dimos_generated.std_msgs.msg import Float32, Header, Int8 import numpy as np from reactivex.disposable import Disposable @@ -68,13 +71,7 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path as NavPath -from dimos.msgs.std_msgs.Float32 import Float32 -from dimos.msgs.std_msgs.Int8 import Int8 +from dimos.msgs.geometry import quaternion_euler, quaternion_from_euler from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -119,18 +116,23 @@ def shift_path_to_start_at_pose(path: NavPath, start_pose: PoseStamped) -> NavPa at the robot's current pose, so the operator only has to roughly aim the robot. Scoring is then in the executed frame regardless of where the plant starts.""" - px0, py0 = path.poses[0].position.x, path.poses[0].position.y - pyaw0 = path.poses[0].orientation.euler[2] - sx, sy = start_pose.position.x, start_pose.position.y - dyaw = start_pose.orientation.euler[2] - pyaw0 + px0, py0 = path.poses[0].pose.position.x, path.poses[0].pose.position.y + pyaw0 = quaternion_euler(path.poses[0].pose.orientation)[2] + sx, sy = start_pose.pose.position.x, start_pose.pose.position.y + dyaw = quaternion_euler(start_pose.pose.orientation)[2] - pyaw0 cd, sd = math.cos(dyaw), math.sin(dyaw) new = [] for p in path.poses: - rx, ry = p.position.x - px0, p.position.y - py0 + rx, ry = p.pose.position.x - px0, p.pose.position.y - py0 new.append( PoseStamped( - position=Vector3(sx + rx * cd - ry * sd, sy + rx * sd + ry * cd, 0.0), - orientation=Quaternion.from_euler(Vector3(0.0, 0.0, p.orientation.euler[2] + dyaw)), + header=Header(frame_id=""), + pose=Pose( + position=Point(x=sx + rx * cd - ry * sd, y=sy + rx * sd + ry * cd, z=0.0), + orientation=quaternion_from_euler( + 0.0, 0.0, quaternion_euler(p.pose.orientation)[2] + dyaw + ), + ), ) ) return NavPath(poses=new) @@ -172,7 +174,10 @@ def from_path( velocity_threshold: float, timeout: float, ) -> RunRecording: - ref = [[p.position.x, p.position.y, p.orientation.euler[2]] for p in reference.poses] + ref = [ + [p.pose.position.x, p.pose.position.y, quaternion_euler(p.pose.orientation)[2]] + for p in reference.poses + ] return cls( robot=robot, path=path_name, @@ -197,8 +202,11 @@ def reference_path(self) -> NavPath: return NavPath( poses=[ PoseStamped( - position=Vector3(x, y, 0.0), - orientation=Quaternion.from_euler(Vector3(0.0, 0.0, yaw)), + header=Header(frame_id=""), + pose=Pose( + position=Point(x=x, y=y, z=0.0), + orientation=quaternion_from_euler(0.0, 0.0, yaw), + ), ) for x, y, yaw in self.reference ] @@ -246,12 +254,12 @@ def on_odom(self, pose: PoseStamped, now: float | None = None) -> None: else: dt = now - self._prev_t if dt > 0: - dx = pose.position.x - self._prev_pose.position.x - dy = pose.position.y - self._prev_pose.position.y - y1 = pose.orientation.euler[2] - dyaw = (y1 - self._prev_pose.orientation.euler[2] + math.pi) % ( - 2 * math.pi - ) - math.pi + dx = pose.pose.position.x - self._prev_pose.pose.position.x + dy = pose.pose.position.y - self._prev_pose.pose.position.y + y1 = quaternion_euler(pose.pose.orientation)[2] + dyaw = ( + y1 - quaternion_euler(self._prev_pose.pose.orientation)[2] + math.pi + ) % (2 * math.pi) - math.pi c, s = math.cos(y1), math.sin(y1) bx = (dx / dt) * c + (dy / dt) * s by = -(dx / dt) * s + (dy / dt) * c @@ -263,9 +271,9 @@ def on_odom(self, pose: PoseStamped, now: float | None = None) -> None: self._ticks.append( [ t_rel, - pose.position.x, - pose.position.y, - pose.orientation.euler[2], + pose.pose.position.x, + pose.pose.position.y, + quaternion_euler(pose.pose.orientation)[2], self._cmd_vx, self._cmd_vy, self._cmd_wz, @@ -326,7 +334,9 @@ def __init__( progress_frac: float = 0.7, window: int = 20, ) -> None: - self._xy = np.array([[p.position.x, p.position.y] for p in reference.poses], dtype=float) + self._xy = np.array( + [[p.pose.position.x, p.pose.position.y] for p in reference.poses], dtype=float + ) self._goal = self._xy[-1] self._n = len(self._xy) self._goal_tol = goal_tolerance @@ -544,7 +554,9 @@ def _await_completion(self, monitor: CompletionMonitor, timeout_s: float) -> tup pose = self._recorder.latest_pose() if pose is not None: lin, ang = self._recorder.body_speed() - if monitor.update(pose.position.x, pose.position.y, lin, ang, time.perf_counter()): + if monitor.update( + pose.pose.position.x, pose.pose.position.y, lin, ang, time.perf_counter() + ): return True, "goal+stop" time.sleep(0.05) return False, "timeout" diff --git a/dimos/control/benchmarking/paths.py b/dimos/control/benchmarking/paths.py index 94c688d5d27..0138ec9248e 100644 --- a/dimos/control/benchmarking/paths.py +++ b/dimos/control/benchmarking/paths.py @@ -22,13 +22,17 @@ import math +from dimos_generated.geometry_msgs.msg import ( + Point as GeoPoint, + Pose, + PoseStamped, +) +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Header + from dimos.memory.vis.space.elements import Point, Polyline, Text from dimos.memory.vis.space.space import Space -from dimos.msgs.geometry_msgs.Point import Point as GeoPoint -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import quaternion_from_euler # Plot styling constants for the trajectory renderers below. _REF_COLOR = "#cccccc" # reference path = light gray @@ -44,8 +48,10 @@ def _xy_to_path(executed_xy: list[tuple[float, float]]) -> Path: """Wrap (x, y) tuples in a nav_msgs.Path so memory Polyline can render them.""" poses = [ PoseStamped( - position=Vector3(x, y, 0.0), - orientation=Quaternion.from_euler(Vector3(0.0, 0.0, 0.0)), + header=Header(frame_id=""), + pose=Pose( + position=GeoPoint(x=x, y=y, z=0.0), orientation=quaternion_from_euler(0.0, 0.0, 0.0) + ), ) for x, y in executed_xy ] @@ -54,8 +60,10 @@ def _xy_to_path(executed_xy: list[tuple[float, float]]) -> Path: def _pose(x: float, y: float, yaw: float) -> PoseStamped: return PoseStamped( - position=Vector3(x, y, 0.0), - orientation=Quaternion.from_euler(Vector3(0.0, 0.0, yaw)), + header=Header(frame_id=""), + pose=Pose( + position=GeoPoint(x=x, y=y, z=0.0), orientation=quaternion_from_euler(0.0, 0.0, yaw) + ), ) @@ -238,9 +246,8 @@ def hold_heading(path: Path, yaw: float = 0.0) -> Path: of facing the travel direction. """ return Path( - ts=path.ts, - frame_id=path.frame_id, - poses=[_pose(p.position.x, p.position.y, yaw) for p in path.poses], + header=path.header, + poses=[_pose(p.pose.position.x, p.pose.position.y, yaw) for p in path.poses], ) @@ -438,8 +445,8 @@ def path_to_svg(path: Path, size_px: int = 400, margin_px: int = 20) -> str: sp = Space() sp.add(Polyline(msg=path, color="#000000", width=_REF_WIDTH)) - sp.add(Point(msg=path.poses[0], color=_START_COLOR, radius=_MARKER_RADIUS)) - sp.add(Point(msg=path.poses[-1], color=_END_COLOR, radius=_MARKER_RADIUS)) + sp.add(Point(msg=path.poses[0].pose.position, color=_START_COLOR, radius=_MARKER_RADIUS)) + sp.add(Point(msg=path.poses[-1].pose.position, color=_END_COLOR, radius=_MARKER_RADIUS)) return sp.to_svg() @@ -458,8 +465,8 @@ def trajectory_to_svg( sp.add(Polyline(msg=_xy_to_path(executed_xy), color=_EXE_COLOR, width=_EXE_WIDTH)) sx, sy = executed_xy[0] ex, ey = executed_xy[-1] - sp.add(Point(msg=GeoPoint(sx, sy, 0.0), color=_START_COLOR, radius=_MARKER_RADIUS)) - sp.add(Point(msg=GeoPoint(ex, ey, 0.0), color=_END_COLOR, radius=_MARKER_RADIUS)) + sp.add(Point(msg=GeoPoint(x=sx, y=sy), color=_START_COLOR, radius=_MARKER_RADIUS)) + sp.add(Point(msg=GeoPoint(x=ex, y=ey), color=_END_COLOR, radius=_MARKER_RADIUS)) return sp.to_svg() @@ -487,8 +494,8 @@ def multi_trajectory_to_svg( sp.add(Polyline(msg=reference, color=_REF_COLOR, width=_REF_WIDTH * 1.4)) # Establish bounds for legend placement (below the path) and title (above). - all_ys = [p.position.y for p in reference.poses] - all_xs = [p.position.x for p in reference.poses] + all_ys = [p.pose.position.y for p in reference.poses] + all_xs = [p.pose.position.x for p in reference.poses] for xy in cohorts.values(): all_ys.extend(y for _, y in xy) all_xs.extend(x for x, _ in xy) @@ -505,7 +512,7 @@ def multi_trajectory_to_svg( if xy: sp.add(Polyline(msg=_xy_to_path(xy), color=color, width=_EXE_WIDTH)) sx, sy = xy[0] - sp.add(Point(msg=GeoPoint(sx, sy, 0.0), color=color, radius=_MARKER_RADIUS * 0.7)) + sp.add(Point(msg=GeoPoint(x=sx, y=sy), color=color, radius=_MARKER_RADIUS * 0.7)) # Legend row (world coords below the plot). ly = y_min - 0.4 - i * 0.25 sp.add( @@ -513,12 +520,18 @@ def multi_trajectory_to_svg( msg=Path( poses=[ PoseStamped( - position=Vector3(x_min, ly, 0.0), - orientation=Quaternion.from_euler(Vector3(0, 0, 0)), + header=Header(frame_id=""), + pose=Pose( + position=GeoPoint(x=x_min, y=ly, z=0.0), + orientation=quaternion_from_euler(0, 0, 0), + ), ), PoseStamped( - position=Vector3(x_min + 0.4, ly, 0.0), - orientation=Quaternion.from_euler(Vector3(0, 0, 0)), + header=Header(frame_id=""), + pose=Pose( + position=GeoPoint(x=x_min + 0.4, y=ly, z=0.0), + orientation=quaternion_from_euler(0, 0, 0), + ), ), ] ), diff --git a/dimos/control/benchmarking/score.py b/dimos/control/benchmarking/score.py index 810fd9c587d..d17ee0036be 100644 --- a/dimos/control/benchmarking/score.py +++ b/dimos/control/benchmarking/score.py @@ -31,13 +31,14 @@ from pathlib import Path from typing import Any +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Twist, Vector3 +from dimos_generated.std_msgs.msg import Header + from dimos.control.benchmarking.benchmark import RunRecording from dimos.control.benchmarking.scoring import ExecutedTrajectory, TrajectoryTick, score_run from dimos.control.benchmarking.tuning import OperatingPoint, OperatingPointMap, invert_tolerance -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.time import time_from_seconds from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -50,11 +51,15 @@ def _executed_from_recording(rec: RunRecording) -> ExecutedTrajectory: TrajectoryTick( t=t, pose=PoseStamped( - ts=t, - position=Vector3(x, y, 0.0), - orientation=Quaternion.from_euler(Vector3(0.0, 0.0, yaw)), + header=Header(frame_id="", stamp=time_from_seconds(t)), + pose=Pose( + position=Point(x=x, y=y, z=0.0), + orientation=quaternion_from_euler(0.0, 0.0, yaw), + ), + ), + cmd_twist=Twist( + linear=Vector3(x=cvx, y=cvy, z=0.0), angular=Vector3(x=0.0, y=0.0, z=cwz) ), - cmd_twist=Twist(linear=Vector3(cvx, cvy, 0.0), angular=Vector3(0.0, 0.0, cwz)), actual_twist=Twist(), ) ) diff --git a/dimos/control/benchmarking/scoring.py b/dimos/control/benchmarking/scoring.py index a6c1b8a1ab8..5b434bff2a6 100644 --- a/dimos/control/benchmarking/scoring.py +++ b/dimos/control/benchmarking/scoring.py @@ -31,12 +31,12 @@ from dataclasses import dataclass, field import math +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist +from dimos_generated.nav_msgs.msg import Path import numpy as np from numpy.typing import NDArray -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import quaternion_euler from dimos.utils.trigonometry import angle_diff @@ -75,7 +75,7 @@ class ScoreResult: def _path_xy(path: Path) -> NDArray[np.float64]: - return np.array([[p.position.x, p.position.y] for p in path.poses], dtype=np.float64) + return np.array([[p.pose.position.x, p.pose.position.y] for p in path.poses], dtype=np.float64) def nearest_segment( @@ -146,7 +146,10 @@ def score_run(reference_path: Path, executed: ExecutedTrajectory) -> ScoreResult if len(path_xy) == 0: return ScoreResult(arrived=executed.arrived, n_ticks=len(executed.ticks)) path_yaws = np.unwrap( - np.array([p.orientation.euler[2] for p in reference_path.poses], dtype=np.float64) + np.array( + [quaternion_euler(p.pose.orientation)[2] for p in reference_path.poses], + dtype=np.float64, + ) ).astype(np.float64) cte_sq: list[float] = [] @@ -157,13 +160,13 @@ def score_run(reference_path: Path, executed: ExecutedTrajectory) -> ScoreResult ang_sq: list[float] = [] for tick in executed.ticks: - pt = np.array([tick.pose.position.x, tick.pose.position.y], dtype=np.float64) + pt = np.array([tick.pose.pose.position.x, tick.pose.pose.position.y], dtype=np.float64) seg_idx, d, t_along = nearest_segment(pt, path_xy) cte_abs.append(d) cte_sq.append(d * d) path_yaw = _reference_yaw(path_yaws, seg_idx, t_along) - he = abs(angle_diff(tick.pose.orientation.euler[2], path_yaw)) + he = abs(angle_diff(quaternion_euler(tick.pose.pose.orientation)[2], path_yaw)) he_abs.append(he) he_sq.append(he * he) diff --git a/dimos/control/benchmarking/test_benchmark.py b/dimos/control/benchmarking/test_benchmark.py index 12b9c1e2729..958d3537d90 100644 --- a/dimos/control/benchmarking/test_benchmark.py +++ b/dimos/control/benchmarking/test_benchmark.py @@ -21,6 +21,8 @@ from types import SimpleNamespace +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Twist, Vector3 +from dimos_generated.std_msgs.msg import Header import pytest from dimos.control.benchmarking.benchmark import ( @@ -44,17 +46,17 @@ DEFAULT_ARTIFACT_PATH, create_task, ) -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.geometry import quaternion_euler, quaternion_from_euler _JOINTS = ["go2/vx", "go2/vy", "go2/wz"] def _pose(x=0.0, y=0.0, yaw=0.0) -> PoseStamped: return PoseStamped( - position=Vector3(x, y, 0.0), orientation=Quaternion.from_euler(Vector3(0.0, 0.0, yaw)) + header=Header(frame_id=""), + pose=Pose( + position=Point(x=x, y=y, z=0.0), orientation=quaternion_from_euler(0.0, 0.0, yaw) + ), ) @@ -174,12 +176,12 @@ def test_battery_registry_selects_fullpose_paths(): } # The point of the battery: commanded yaw decoupled from the tangent. strafe = fullpose["strafe_left_2m"] - assert all(abs(p.orientation.euler[2]) < 1e-9 for p in strafe.poses) - assert strafe.poses[-1].position.y > 1.0 # travel is +y while yaw is 0 + assert all(abs(quaternion_euler(p.pose.orientation)[2]) < 1e-9 for p in strafe.poses) + assert strafe.poses[-1].pose.position.y > 1.0 # travel is +y while yaw is 0 # square_crab: square geometry, one held heading through all corners. crab = fullpose["square_crab"] - assert all(abs(p.orientation.euler[2]) < 1e-9 for p in crab.poses) - assert max(p.position.y for p in crab.poses) > 1.0 + assert all(abs(quaternion_euler(p.pose.orientation)[2]) < 1e-9 for p in crab.poses) + assert max(p.pose.position.y for p in crab.poses) > 1.0 # "all" = tangent-heading battery + full-pose battery. assert set(BATTERIES["all"]()) == set(path_set()) | set(fullpose) @@ -187,7 +189,7 @@ def test_battery_registry_selects_fullpose_paths(): def test_anchor_shifts_path_to_pose(): ref = straight_line(length=2.0) anchored = shift_path_to_start_at_pose(ref, _pose(5.0, 3.0, 0.0)) - p0 = anchored.poses[0].position + p0 = anchored.poses[0].pose.position assert p0.x == 5.0 and p0.y == 3.0 @@ -256,7 +258,9 @@ def test_completion_does_not_fire_prematurely_on_closed_path(): # path yet -> not complete. done = False for i in range(10): - done = mon.update(ref.poses[0].position.x, ref.poses[0].position.y, 0.0, 0.0, i * 0.1) + done = mon.update( + ref.poses[0].pose.position.x, ref.poses[0].pose.position.y, 0.0, 0.0, i * 0.1 + ) assert not done @@ -316,7 +320,9 @@ def test_end_to_end_controller_benchmark_scoring(tmp_path): out = task.compute(_state(plant.x, plant.y, plant.yaw, t=k * 0.1)) cvx, cvy, cwz = out.velocities if out is not None else (0.0, 0.0, 0.0) # Benchmark records what it sees over the transport: cmd_vel + odom. - recorder.on_cmd_vel(Twist(linear=Vector3(cvx, cvy, 0.0), angular=Vector3(0.0, 0.0, cwz))) + recorder.on_cmd_vel( + Twist(linear=Vector3(x=cvx, y=cvy, z=0.0), angular=Vector3(x=0.0, y=0.0, z=cwz)) + ) recorder.on_odom(_pose(plant.x, plant.y, plant.yaw), now=k * 0.1) lin, ang = recorder.body_speed() if monitor.update(plant.x, plant.y, lin, ang, k * 0.1): @@ -386,7 +392,9 @@ def test_end_to_end_fullpose_benchmark_scoring(tmp_path): for k in range(800): out = task.compute(_state(plant.x, plant.y, plant.yaw, t=k * 0.1)) cvx, cvy, cwz = out.velocities if out is not None else (0.0, 0.0, 0.0) - recorder.on_cmd_vel(Twist(linear=Vector3(cvx, cvy, 0.0), angular=Vector3(0.0, 0.0, cwz))) + recorder.on_cmd_vel( + Twist(linear=Vector3(x=cvx, y=cvy, z=0.0), angular=Vector3(x=0.0, y=0.0, z=cwz)) + ) recorder.on_odom(_pose(plant.x, plant.y, plant.yaw), now=k * 0.1) lin, ang = recorder.body_speed() if monitor.update(plant.x, plant.y, lin, ang, k * 0.1): diff --git a/dimos/control/benchmarking/velocity_profile.py b/dimos/control/benchmarking/velocity_profile.py index e1eec131b7b..cfa45f510b2 100644 --- a/dimos/control/benchmarking/velocity_profile.py +++ b/dimos/control/benchmarking/velocity_profile.py @@ -40,10 +40,10 @@ from dataclasses import dataclass from typing import Protocol +from dimos_generated.nav_msgs.msg import Path import numpy as np from dimos.control.tasks.velocity_profiler import VelocityProfiler -from dimos.msgs.nav_msgs.Path import Path # Go2 Rung-1 saturation envelope (mirrors runner.VX_MAX / WZ_MAX). GO2_VX_MAX = 1.0 # m/s @@ -108,7 +108,9 @@ def __init__(self, cfg: VelocityProfileConfig | None = None) -> None: def for_path(self, path: Path) -> None: """(Re)compute the speed profile for ``path``. Call on path start.""" self._profile = np.asarray(self._profiler.compute_profile(path), dtype=float) - self._pts = np.array([[p.position.x, p.position.y] for p in path.poses], dtype=float) + self._pts = np.array( + [[p.pose.position.x, p.pose.position.y] for p in path.poses], dtype=float + ) def speed_limit_at(self, x: float, y: float) -> float: """Profile speed at the nearest path index, min over the lookahead diff --git a/dimos/control/examples/cartesian_ik_jogger.py b/dimos/control/examples/cartesian_ik_jogger.py index 7e438f02391..88177241d79 100644 --- a/dimos/control/examples/cartesian_ik_jogger.py +++ b/dimos/control/examples/cartesian_ik_jogger.py @@ -37,8 +37,13 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import Point, Pose +from dimos_generated.std_msgs.msg import Header import numpy as np +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.time import time_from_seconds + try: import pygame except ImportError: @@ -125,17 +130,16 @@ def from_fk( def to_pose_stamped(self) -> Any: """Convert to PoseStamped for LCM publishing.""" - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - from dimos.msgs.geometry_msgs.Quaternion import Quaternion - from dimos.msgs.geometry_msgs.Vector3 import Vector3 + from dimos_generated.geometry_msgs.msg import PoseStamped, Vector3 - position = Vector3(self.x, self.y, self.z) - orientation = Quaternion.from_euler(Vector3(self.roll, self.pitch, self.yaw)) + position = Vector3(x=self.x, y=self.y, z=self.z) + orientation = quaternion_from_euler(self.roll, self.pitch, self.yaw) return PoseStamped( - ts=time.time(), - position=position, - orientation=orientation, + header=Header(frame_id="", stamp=time_from_seconds(time.time())), + pose=Pose( + position=Point(x=position.x, y=position.y, z=position.z), orientation=orientation + ), ) @@ -171,8 +175,9 @@ def run_jogger_ui(model_path: str | None = None, ee_joint_id: int = 6) -> None: If None, uses Piper model. ee_joint_id: End-effector joint ID in the model """ + from dimos_generated.geometry_msgs.msg import PoseStamped + from dimos.core.transport import LCMTransport - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped # Use Piper model if not specified if model_path is None: diff --git a/dimos/control/path_following_coordinator.py b/dimos/control/path_following_coordinator.py index 09772ca3b02..4b410a248c6 100644 --- a/dimos/control/path_following_coordinator.py +++ b/dimos/control/path_following_coordinator.py @@ -16,10 +16,11 @@ from __future__ import annotations +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Float32 + from dimos.control.coordinator import ControlCoordinator from dimos.core.stream import In -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.std_msgs.Float32 import Float32 class PathFollowingCoordinator(ControlCoordinator): diff --git a/dimos/control/task.py b/dimos/control/task.py index d88bf2e596d..b2215ae1f52 100644 --- a/dimos/control/task.py +++ b/dimos/control/task.py @@ -35,9 +35,8 @@ from dimos.hardware.whole_body.spec import IMUState if TYPE_CHECKING: - from dimos.msgs.geometry_msgs.Pose import Pose - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped + from dimos_generated.geometry_msgs.msg import Pose, PoseStamped, TwistStamped + from dimos.teleop.webxr.controller_types import Buttons @@ -187,6 +186,7 @@ class ControlTask(Protocol): Use state.t_now passed to compute() for all timing. Example: + >>> from dimos.msgs.trajectory import sample_trajectory >>> class MyTask: ... @property ... def name(self) -> str: @@ -204,7 +204,7 @@ class ControlTask(Protocol): ... def compute(self, state: CoordinatorState) -> JointCommandOutput | None: ... # Use state.t_now, NOT time.time()! ... t_elapsed = state.t_now - self._start_time - ... positions = self._trajectory.sample(t_elapsed) + ... positions, _ = sample_trajectory(self._trajectory, t_elapsed) ... return JointCommandOutput( ... joint_names=["left/joint1", "left/joint2"], ... positions=positions, diff --git a/dimos/control/tasks/cartesian_ik_task/cartesian_ik_task.py b/dimos/control/tasks/cartesian_ik_task/cartesian_ik_task.py index 27bc3e717ff..33784e372dc 100644 --- a/dimos/control/tasks/cartesian_ik_task/cartesian_ik_task.py +++ b/dimos/control/tasks/cartesian_ik_task/cartesian_ik_task.py @@ -17,10 +17,12 @@ from __future__ import annotations from collections.abc import Mapping +import copy import threading from typing import TYPE_CHECKING import attrs +from dimos_generated.geometry_msgs.msg import Pose, PoseStamped from dimos.control.task import CoordinatorState from dimos.control.tasks.pose_target_ik import ( @@ -31,8 +33,6 @@ PoseTargetIKTaskParams, string_tuple_converter, ) -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped if TYPE_CHECKING: from dimos.control.coordinator import TaskConfig @@ -72,12 +72,7 @@ def is_active(self) -> bool: def on_cartesian_command(self, pose: Pose | PoseStamped, t_now: float) -> bool: """Accept an absolute target pose and activate tracking.""" - target = PoseStamped( - ts=pose.ts if isinstance(pose, PoseStamped) else 0.0, - frame_id=pose.frame_id if isinstance(pose, PoseStamped) else "", - position=pose.position, - orientation=pose.orientation, - ) + target = copy.deepcopy(pose) if isinstance(pose, PoseStamped) else PoseStamped(pose=pose) with self._lock: self._target_pose = target self._last_update_time = t_now diff --git a/dimos/control/tasks/cartesian_ik_task/test_cartesian_ik_task.py b/dimos/control/tasks/cartesian_ik_task/test_cartesian_ik_task.py index 86b9f6092fe..fdf21be6b34 100644 --- a/dimos/control/tasks/cartesian_ik_task/test_cartesian_ik_task.py +++ b/dimos/control/tasks/cartesian_ik_task/test_cartesian_ik_task.py @@ -16,6 +16,9 @@ from pathlib import Path +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header import pytest from pytest_mock import MockerFixture @@ -26,9 +29,6 @@ ) from dimos.control.tasks.pose_target_ik import PinkPoseTargetSolver from dimos.manipulation.planning.spec.config import RobotModelConfig -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.robot.assets.model import RobotModel @@ -57,7 +57,9 @@ def test_cartesian_leaf_maps_absolute_pose_to_configured_frame( solver = mocker.Mock(spec=PinkPoseTargetSolver) solver.step.return_value = JointState(name=["arm/joint"], position=[0.01]) task = CartesianIKTask("cartesian", _config(), solver=solver) - target = PoseStamped(position=Vector3(0.4, 0.2, 0.1), frame_id="world") + target = PoseStamped( + header=Header(frame_id="world"), pose=Pose(position=Point(x=0.4, y=0.2, z=0.1)) + ) task.on_cartesian_command(target, t_now=2.0) output = task.compute( @@ -75,7 +77,7 @@ def test_cartesian_leaf_maps_absolute_pose_to_configured_frame( def test_cartesian_leaf_clears_after_timeout(mocker: MockerFixture) -> None: solver = mocker.Mock(spec=PinkPoseTargetSolver) task = CartesianIKTask("cartesian", _config(), solver=solver) - task.on_cartesian_command(PoseStamped(), t_now=1.0) + task.on_cartesian_command(PoseStamped(header=Header(frame_id=""), pose=Pose()), t_now=1.0) output = task.compute( CoordinatorState( @@ -99,11 +101,11 @@ def test_cartesian_clear_reseeds_command_from_feedback(mocker: MockerFixture) -> t_now=1.0, dt=0.01, ) - task.on_cartesian_command(PoseStamped(), t_now=1.0) + task.on_cartesian_command(PoseStamped(header=Header(frame_id=""), pose=Pose()), t_now=1.0) assert task.compute(state) is not None task.clear() - task.on_cartesian_command(PoseStamped(), t_now=1.1) + task.on_cartesian_command(PoseStamped(header=Header(frame_id=""), pose=Pose()), t_now=1.1) solver.step.return_value = JointState(name=["arm/joint"], position=[0.01]) assert task.compute(state) is not None diff --git a/dimos/control/tasks/eef_twist_task/eef_twist_task.py b/dimos/control/tasks/eef_twist_task/eef_twist_task.py index b62ee7f0d5a..ac9d6ff889d 100644 --- a/dimos/control/tasks/eef_twist_task/eef_twist_task.py +++ b/dimos/control/tasks/eef_twist_task/eef_twist_task.py @@ -21,6 +21,7 @@ from typing import TYPE_CHECKING import attrs +from dimos_generated.geometry_msgs.msg import PoseStamped import numpy as np import pinocchio @@ -33,14 +34,15 @@ PoseTargetIKTaskParams, string_tuple_converter, ) -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry import pose_from_matrix, pose_matrix from dimos.utils.logging_config import setup_logger -from dimos.utils.transform_utils import matrix_to_pose, pose_to_matrix, twist_to_numpy +from dimos.utils.transform_utils import twist_to_numpy if TYPE_CHECKING: + from dimos_generated.geometry_msgs.msg import TwistStamped + from dimos.control.coordinator import TaskConfig from dimos.control.hardware_interface import ConnectedHardware, ConnectedWholeBody - from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped logger = setup_logger() @@ -89,7 +91,7 @@ def on_cartesian_command(self, pose: object, t_now: float) -> bool: return False def on_ee_twist_command(self, twist: TwistStamped, t_now: float) -> bool: - values = twist_to_numpy(twist) + values = twist_to_numpy(twist.twist) if values.shape != (6,) or not np.all(np.isfinite(values)): logger.warning("EEFTwistTask rejecting invalid twist", task=self._name) return False @@ -193,16 +195,11 @@ def create_task( def _integrate_twist(pose: PoseStamped, twist: TwistStamped, dt: float) -> PoseStamped: - matrix = np.asarray(pose_to_matrix(pose), dtype=np.float64) - values = twist_to_numpy(twist) + matrix = np.asarray(pose_matrix(pose.pose), dtype=np.float64) + values = twist_to_numpy(twist.twist) matrix[:3, 3] += values[:3] * dt angular_step = values[3:] * dt if np.linalg.norm(angular_step) > 0.0: matrix[:3, :3] = pinocchio.exp3(angular_step) @ matrix[:3, :3] - integrated = matrix_to_pose(matrix) - return PoseStamped( - ts=pose.ts, - frame_id=pose.frame_id, - position=integrated.position, - orientation=integrated.orientation, - ) + integrated = pose_from_matrix(matrix) + return PoseStamped(header=pose.header, pose=integrated) diff --git a/dimos/control/tasks/eef_twist_task/test_eef_twist_task.py b/dimos/control/tasks/eef_twist_task/test_eef_twist_task.py index 864267901b9..7e69037f904 100644 --- a/dimos/control/tasks/eef_twist_task/test_eef_twist_task.py +++ b/dimos/control/tasks/eef_twist_task/test_eef_twist_task.py @@ -14,6 +14,9 @@ from pathlib import Path +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Twist, TwistStamped, Vector3 +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from pytest_mock import MockerFixture @@ -25,9 +28,6 @@ ) from dimos.control.tasks.pose_target_ik import PinkPoseTargetSolver from dimos.manipulation.planning.spec.config import RobotModelConfig -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.robot.assets.model import RobotModel @@ -41,7 +41,9 @@ def _robot_model() -> RobotModelConfig: def _solver(mocker: MockerFixture) -> PinkPoseTargetSolver: solver = mocker.Mock(spec=PinkPoseTargetSolver) solver.frame_poses.return_value = { - "tool": PoseStamped(frame_id="world", position=[0.0, 0.0, 0.0]) + "tool": PoseStamped( + header=Header(frame_id="world"), pose=Pose(position=Point(x=0.0, y=0.0, z=0.0)) + ) } solver.step.return_value = JointState( name=["arm/joint1", "arm/joint2"], @@ -92,7 +94,7 @@ def _state(t_now: float = 1.0, *, dt: float = 0.01) -> CoordinatorState: def _twist(x: float = 0.1) -> TwistStamped: - return TwistStamped(frame_id="tool", linear=[x, 0.0, 0.0], angular=[0.0, 0.0, 0.0]) + return TwistStamped(header=Header(frame_id="tool"), twist=Twist(linear=Vector3(x=x))) def test_twist_integrates_target_and_uses_shared_persistent_command( @@ -108,8 +110,8 @@ def test_twist_integrates_target_and_uses_shared_persistent_command( assert second is not None first_target = solver.step.call_args_list[0].args[0]["tool"] second_target = solver.step.call_args_list[1].args[0]["tool"] - assert first_target.position.x == pytest.approx(0.001) - assert second_target.position.x == pytest.approx(0.002) + assert first_target.pose.position.x == pytest.approx(0.001) + assert second_target.pose.position.x == pytest.approx(0.002) def test_zero_twist_holds_the_integrated_target(mocker: MockerFixture) -> None: @@ -121,7 +123,7 @@ def test_zero_twist_holds_the_integrated_target(mocker: MockerFixture) -> None: task.compute(_state(1.03)) target = solver.step.call_args.args[0]["tool"] - assert target.position.x == pytest.approx(0.001) + assert target.pose.position.x == pytest.approx(0.001) def test_stale_twist_stops_motion_without_dropping_hold(mocker: MockerFixture) -> None: @@ -133,7 +135,7 @@ def test_stale_twist_stops_motion_without_dropping_hold(mocker: MockerFixture) - assert output is not None target = solver.step.call_args.args[0]["tool"] - assert target.position.x == pytest.approx(0.001) + assert target.pose.position.x == pytest.approx(0.001) def test_estop_rejects_input_and_reanchors_after_clear(mocker: MockerFixture) -> None: @@ -165,6 +167,6 @@ def test_preemption_discards_twist_target_and_command_state(mocker: MockerFixtur def test_invalid_twist_is_rejected(mocker: MockerFixture) -> None: task, _ = _task(mocker) invalid = _twist() - invalid.linear.x = np.nan + invalid.twist.linear.x = np.nan assert not task.on_ee_twist_command(invalid, t_now=1.0) diff --git a/dimos/control/tasks/g1_groot_wbc_task/g1_groot_wbc_task.py b/dimos/control/tasks/g1_groot_wbc_task/g1_groot_wbc_task.py index deed610f770..9f925a729cd 100644 --- a/dimos/control/tasks/g1_groot_wbc_task/g1_groot_wbc_task.py +++ b/dimos/control/tasks/g1_groot_wbc_task/g1_groot_wbc_task.py @@ -48,8 +48,9 @@ from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: + from dimos_generated.geometry_msgs.msg import Twist + from dimos.hardware.whole_body.spec import WholeBodyAdapter - from dimos.msgs.geometry_msgs.Twist import Twist logger = setup_logger() diff --git a/dimos/control/tasks/g1_groot_wbc_task/test_g1_groot_wbc_task.py b/dimos/control/tasks/g1_groot_wbc_task/test_g1_groot_wbc_task.py index 95304b2a10c..6a13203b4f9 100644 --- a/dimos/control/tasks/g1_groot_wbc_task/test_g1_groot_wbc_task.py +++ b/dimos/control/tasks/g1_groot_wbc_task/test_g1_groot_wbc_task.py @@ -26,6 +26,7 @@ from typing import Any from unittest.mock import MagicMock +from dimos_generated.geometry_msgs.msg import Twist, Vector3 import numpy as np import pytest @@ -37,7 +38,6 @@ G1GrootWBCTaskConfig, ) from dimos.hardware.whole_body.spec import IMUState -from dimos.msgs.geometry_msgs.Twist import Twist class _StubSession: @@ -186,7 +186,10 @@ def test_nonzero_cmd_uses_walk_until_timeout( def test_on_twist_command_sets_velocity_command(task: G1GrootWBCTask) -> None: - task.on_twist_command(Twist(linear=[0.5, 0.25, 0.0], angular=[0.0, 0.0, 0.3]), t_now=100.0) + task.on_twist_command( + Twist(linear=Vector3(x=0.5, y=0.25, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.3)), + t_now=100.0, + ) assert list(task._cmd) == pytest.approx([0.5, 0.25, 0.3]) assert task._last_cmd_time == 100.0 diff --git a/dimos/control/tasks/gripper_task/gripper_task.py b/dimos/control/tasks/gripper_task/gripper_task.py index 2945f69d877..d41ac8acaac 100644 --- a/dimos/control/tasks/gripper_task/gripper_task.py +++ b/dimos/control/tasks/gripper_task/gripper_task.py @@ -32,7 +32,7 @@ from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: - from dimos.msgs.std_msgs.Float32 import Float32 + from dimos_generated.std_msgs.msg import Float32 logger = setup_logger() diff --git a/dimos/control/tasks/gripper_task/test_gripper_task.py b/dimos/control/tasks/gripper_task/test_gripper_task.py index 9407332121f..75f56b851f9 100644 --- a/dimos/control/tasks/gripper_task/test_gripper_task.py +++ b/dimos/control/tasks/gripper_task/test_gripper_task.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.std_msgs.msg import Float32 import pytest from pytest_mock import MockerFixture @@ -26,7 +27,6 @@ ) from dimos.hardware.manipulators.spec import ManipulatorAdapter from dimos.hardware.spec import JointLimits -from dimos.msgs.std_msgs.Float32 import Float32 def _task() -> GripperControlTask: diff --git a/dimos/control/tasks/holonomic_pose_follower_task/holonomic_pose_follower_task.py b/dimos/control/tasks/holonomic_pose_follower_task/holonomic_pose_follower_task.py index 773a5abc2bd..e953e0ff094 100644 --- a/dimos/control/tasks/holonomic_pose_follower_task/holonomic_pose_follower_task.py +++ b/dimos/control/tasks/holonomic_pose_follower_task/holonomic_pose_follower_task.py @@ -26,6 +26,10 @@ from pathlib import Path as _FsPath from typing import Any, Literal +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Float32, Header + from dimos.control.benchmarking.tuning import TuningConfig from dimos.control.task import ( BaseControlTask, @@ -42,11 +46,7 @@ from dimos.control.tasks.holonomic_pose_follower_task.progress_reference import ( ProgressPathReference, ) -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.std_msgs.Float32 import Float32 +from dimos.msgs.geometry import quaternion_from_euler from dimos.protocol.service.spec import BaseConfig from dimos.utils.logging_config import setup_logger from dimos.utils.trigonometry import angle_diff @@ -79,8 +79,11 @@ def _clamp(v: float, limit: float) -> float: def _pose_stamped(pose: tuple[float, float, float]) -> PoseStamped: return PoseStamped( - position=Vector3(pose[0], pose[1], 0.0), - orientation=Quaternion.from_euler(Vector3(0.0, 0.0, pose[2])), + header=Header(frame_id=""), + pose=Pose( + position=Point(x=pose[0], y=pose[1], z=0.0), + orientation=quaternion_from_euler(0.0, 0.0, pose[2]), + ), ) @@ -417,7 +420,7 @@ def start_path(self, path: Path, current_odom: PoseStamped) -> bool: self._reference = reference # Re-project onto the new path. _v_path deliberately carries over, so an # in-motion replan does not re-ramp from rest. - reference.advance(float(current_odom.position.x), float(current_odom.position.y)) + reference.advance(float(current_odom.pose.position.x), float(current_odom.pose.position.y)) self._last_t = None self._last_pose = None self._last_pose_t = None diff --git a/dimos/control/tasks/holonomic_pose_follower_task/progress_reference.py b/dimos/control/tasks/holonomic_pose_follower_task/progress_reference.py index c754a72daa9..0f468391ccd 100644 --- a/dimos/control/tasks/holonomic_pose_follower_task/progress_reference.py +++ b/dimos/control/tasks/holonomic_pose_follower_task/progress_reference.py @@ -32,9 +32,10 @@ from dataclasses import dataclass import math +from dimos_generated.nav_msgs.msg import Path import numpy as np -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import quaternion_euler _EPS = 1e-9 @@ -68,8 +69,8 @@ def __init__(self, path: Path, *, window_m: float = 1.0, back_m: float = 0.5) -> ys: list[float] = [] yaws: list[float] = [] for p in path.poses: - x, y = float(p.position.x), float(p.position.y) - yaw = float(p.orientation.euler[2]) + x, y = float(p.pose.position.x), float(p.pose.position.y) + yaw = float(quaternion_euler(p.pose.orientation)[2]) if xs and math.hypot(x - xs[-1], y - ys[-1]) < _EPS: # Coincident waypoint: keep the position once but adopt the # later pose's yaw (a stop-and-rotate hint survives as a yaw diff --git a/dimos/control/tasks/holonomic_pose_follower_task/test_holonomic_pose_follower_task.py b/dimos/control/tasks/holonomic_pose_follower_task/test_holonomic_pose_follower_task.py index 79d801c2940..1f8d17f1bfa 100644 --- a/dimos/control/tasks/holonomic_pose_follower_task/test_holonomic_pose_follower_task.py +++ b/dimos/control/tasks/holonomic_pose_follower_task/test_holonomic_pose_follower_task.py @@ -21,6 +21,9 @@ import math from types import SimpleNamespace +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Twist, Vector3 +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Float32, Header import pytest from dimos.control.benchmarking.paths import ( @@ -42,12 +45,7 @@ DEFAULT_ARTIFACT_PATH, create_task, ) -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.std_msgs.Float32 import Float32 +from dimos.msgs.geometry import quaternion_euler, quaternion_from_euler from dimos.utils.trigonometry import angle_diff _JOINTS = ["go2/vx", "go2/vy", "go2/wz"] @@ -56,8 +54,10 @@ def _pose(x=0.0, y=0.0, yaw=0.0): return PoseStamped( - position=Vector3(x, y, 0.0), - orientation=Quaternion.from_euler(Vector3(0.0, 0.0, yaw)), + header=Header(frame_id=""), + pose=Pose( + position=Point(x=x, y=y, z=0.0), orientation=quaternion_from_euler(0.0, 0.0, yaw) + ), ) @@ -106,7 +106,9 @@ def _run_closed_loop(task, path, max_ticks=2400, on_tick=None): TrajectoryTick( t=k * _DT, pose=_pose(plant.x, plant.y, plant.yaw), - cmd_twist=Twist(linear=Vector3(vx, vy, 0.0), angular=Vector3(0.0, 0.0, wz)), + cmd_twist=Twist( + linear=Vector3(x=vx, y=vy, z=0.0), angular=Vector3(x=0.0, y=0.0, z=wz) + ), actual_twist=Twist(), ) ) @@ -206,7 +208,7 @@ def test_strafes_holding_commanded_yaw(): task = _task() plant, executed = _run_closed_loop(task, strafe_line(length=2.0)) assert executed.arrived - max_yaw = max(abs(t.pose.orientation.euler[2]) for t in executed.ticks) + max_yaw = max(abs(quaternion_euler(t.pose.pose.orientation)[2]) for t in executed.ticks) assert max_yaw < 0.1, "robot rotated instead of strafing" assert plant.y == pytest.approx(2.0, abs=0.25) # The motion was carried by the lateral channel. @@ -221,7 +223,7 @@ def test_crab_walks_square_holding_heading(): task = _task() plant, executed = _run_closed_loop(task, path) assert executed.arrived - max_yaw = max(abs(t.pose.orientation.euler[2]) for t in executed.ticks) + max_yaw = max(abs(quaternion_euler(t.pose.pose.orientation)[2]) for t in executed.ticks) assert max_yaw < 0.1, "robot turned instead of crab-walking the square" score = score_run(path, executed) assert score.cte_rms < 0.15 @@ -284,7 +286,7 @@ def test_stops_inside_the_goal_when_the_plant_runs_hot(): break hot.step(vx, vy, wz, _DT) assert task.get_state() == "arrived" - goal = path.poses[-1].position + goal = path.poses[-1].pose.position rest_err = math.hypot(hot.x - goal.x, hot.y - goal.y) assert rest_err < task._config.goal_tolerance, f"rested {rest_err:.3f} m from goal" # Came to rest, not still coasting through. diff --git a/dimos/control/tasks/holonomic_pose_follower_task/test_progress_reference.py b/dimos/control/tasks/holonomic_pose_follower_task/test_progress_reference.py index b9aba1d71f1..d23fca46d50 100644 --- a/dimos/control/tasks/holonomic_pose_follower_task/test_progress_reference.py +++ b/dimos/control/tasks/holonomic_pose_follower_task/test_progress_reference.py @@ -19,21 +19,23 @@ import math +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Header import pytest from dimos.control.tasks.holonomic_pose_follower_task.progress_reference import ( ProgressPathReference, ) -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import quaternion_from_euler def _pose(x, y, yaw): return PoseStamped( - position=Vector3(x, y, 0.0), - orientation=Quaternion.from_euler(Vector3(0.0, 0.0, yaw)), + header=Header(frame_id=""), + pose=Pose( + position=Point(x=x, y=y, z=0.0), orientation=quaternion_from_euler(0.0, 0.0, yaw) + ), ) diff --git a/dimos/control/tasks/path_follower_task/test_path_follower_task.py b/dimos/control/tasks/path_follower_task/test_path_follower_task.py index f6c2e06298e..8e21d67dee7 100644 --- a/dimos/control/tasks/path_follower_task/test_path_follower_task.py +++ b/dimos/control/tasks/path_follower_task/test_path_follower_task.py @@ -20,6 +20,9 @@ import math +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.std_msgs.msg import Header + from dimos.control.benchmarking.paths import circle, straight_line from dimos.control.benchmarking.velocity_profile import VelocityProfileConfig from dimos.control.task import CoordinatorState, JointStateSnapshot @@ -29,9 +32,7 @@ PathFollowerTaskConfig, ) from dimos.core.global_config import global_config as _gc -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.geometry import quaternion_euler, quaternion_from_euler _JOINTS = ["go2/vx", "go2/vy", "go2/wz"] @@ -54,8 +55,13 @@ def _task(**overrides) -> PathFollowerTask: def _start_aligned(task: PathFollowerTask, path) -> None: odom = PoseStamped( - position=Vector3(path.poses[0].position.x, path.poses[0].position.y, 0.0), - orientation=Quaternion.from_euler(Vector3(0.0, 0.0, path.poses[0].orientation.euler[2])), + header=Header(frame_id=""), + pose=Pose( + position=Point(x=path.poses[0].pose.position.x, y=path.poses[0].pose.position.y, z=0.0), + orientation=quaternion_from_euler( + 0.0, 0.0, quaternion_euler(path.poses[0].pose.orientation)[2] + ), + ), ) assert task.start_path(path, odom) diff --git a/dimos/control/tasks/planar_base_trajectory_task/planar_base_trajectory_task.py b/dimos/control/tasks/planar_base_trajectory_task/planar_base_trajectory_task.py index 8636b29dc69..f1b65f8bf54 100644 --- a/dimos/control/tasks/planar_base_trajectory_task/planar_base_trajectory_task.py +++ b/dimos/control/tasks/planar_base_trajectory_task/planar_base_trajectory_task.py @@ -22,6 +22,8 @@ import threading from typing import Any +from dimos_generated.dimos_msgs.msg import TrajectoryStatus +from dimos_generated.trajectory_msgs.msg import JointTrajectory from pydantic import ConfigDict, NonNegativeFloat, PositiveFloat from dimos.control.task import ( @@ -35,8 +37,8 @@ TrajectoryExecutionResult, TrajectoryExecutionStatus, ) -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState, TrajectoryStatus +from dimos.msgs.time import duration_from_seconds, header_now, to_nanoseconds +from dimos.msgs.trajectory import TrajectoryState, sample_trajectory, trajectory_duration from dimos.protocol.service.spec import BaseConfig from dimos.utils.logging_config import setup_logger from dimos.utils.trigonometry import angle_diff @@ -156,20 +158,23 @@ def cancel(self) -> bool: def get_status(self, t_now: float | None = None) -> TrajectoryStatus: with self._lock: - duration = self._trajectory.duration if self._trajectory is not None else 0.0 + duration = ( + trajectory_duration(self._trajectory) if self._trajectory is not None else 0.0 + ) elapsed = self._elapsed if ( t_now is not None and self._start_t is not None - and self._state is TrajectoryState.EXECUTING + and self._state == TrajectoryState.EXECUTING ): elapsed = t_now - self._start_t progress = 1.0 if duration <= 0.0 else min(1.0, elapsed / duration) return TrajectoryStatus( + header=header_now(), state=self._state, progress=progress, - time_elapsed=elapsed, - time_remaining=max(0.0, duration - elapsed), + time_elapsed=duration_from_seconds(elapsed), + time_remaining=duration_from_seconds(max(0.0, duration - elapsed)), error=self._error, ) @@ -221,7 +226,7 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None: if self._start_t is None: self._start_t = state.t_now self._elapsed = state.t_now - self._start_t - reference, reference_velocity = self._trajectory.sample(self._elapsed) + reference, reference_velocity = sample_trajectory(self._trajectory, self._elapsed) position_error = math.hypot(reference[0] - pose[0], reference[1] - pose[1]) yaw_error = abs(angle_diff(reference[2], pose[2])) @@ -243,14 +248,14 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None: state.t_now, f"base is {position_error:.3f} m / {yaw_error:.3f} rad off the trajectory", ) - if self._elapsed >= self._trajectory.duration: + if self._elapsed >= trajectory_duration(self._trajectory): if ( position_error < config.position_goal_tolerance and yaw_error < config.orientation_goal_tolerance ): self._stop(state.t_now, TrajectoryState.COMPLETED, "") return self._command(0.0, 0.0, 0.0) - if self._elapsed - self._trajectory.duration > config.settle_timeout: + if self._elapsed - trajectory_duration(self._trajectory) > config.settle_timeout: return self._fail(state.t_now, "base did not reach the end of the trajectory") # Settling is pure feedback; the last point's velocity would carry it past. reference_velocity = [0.0, 0.0, 0.0] @@ -327,18 +332,22 @@ def _trajectory_problem( ) -> str: if trajectory is None or not trajectory.points: return "Base trajectory has no points" - previous_time: float | None = None + previous_time: int | None = None for point in trajectory.points: if len(point.positions) != 3 or len(point.velocities) != 3: return "Base trajectory points must carry x, y and yaw" - values = (*point.positions, *point.velocities, point.time_from_start) + try: + point_time = to_nanoseconds(point.time_from_start) + except ValueError: + return "Base trajectory contains an invalid duration" + values = (*point.positions, *point.velocities, point_time) if not all(math.isfinite(value) for value in values): return "Base trajectory contains non-finite values" - if previous_time is None and point.time_from_start != 0.0: + if previous_time is None and point_time != 0.0: return "Base trajectory must start at t=0" - if previous_time is not None and point.time_from_start <= previous_time: + if previous_time is not None and point_time <= previous_time: return "Base trajectory has non-increasing timestamps" - previous_time = point.time_from_start + previous_time = point_time vx, vy, wz = point.velocities if ( math.hypot(vx, vy) > config.max_linear + _SPEED_SLACK diff --git a/dimos/control/tasks/planar_base_trajectory_task/test_planar_base_trajectory_task.py b/dimos/control/tasks/planar_base_trajectory_task/test_planar_base_trajectory_task.py index 833a3400fb7..4ac516a1b47 100644 --- a/dimos/control/tasks/planar_base_trajectory_task/test_planar_base_trajectory_task.py +++ b/dimos/control/tasks/planar_base_trajectory_task/test_planar_base_trajectory_task.py @@ -14,6 +14,7 @@ import math +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint import pytest from dimos.control.task import CoordinatorState, JointStateSnapshot @@ -22,9 +23,8 @@ PlanarBaseTrajectoryTaskConfig, ) from dimos.control.tasks.trajectory_task.trajectory_task import TrajectoryExecutionStatus -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint -from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState +from dimos.msgs.time import duration_from_seconds, header_now +from dimos.msgs.trajectory import TrajectoryState, sample_trajectory from dimos.utils.trigonometry import angle_diff JOINTS = ["base/vx", "base/vy", "base/wz"] @@ -58,14 +58,14 @@ def step(self, command): def _line(start, goal, duration, n=60): velocity = [(g - s) / duration for s, g in zip(start, goal, strict=True)] points = [ - TrajectoryPoint( - time_from_start=duration * i / (n - 1), + JointTrajectoryPoint( + time_from_start=duration_from_seconds(duration * i / (n - 1)), positions=[s + (g - s) * i / (n - 1) for s, g in zip(start, goal, strict=True)], velocities=velocity, ) for i in range(n) ] - return JointTrajectory(points=points, joint_names=["x", "y", "yaw"]) + return JointTrajectory(header=header_now(), points=points, joint_names=["x", "y", "yaw"]) def _run(task, base, seconds, t=100.0): @@ -94,7 +94,7 @@ def test_follows_a_timed_trajectory_across_the_yaw_seam(): task.execute(plan) _, t = _run(task, base, seconds=2.0) - reference, _ = plan.sample(2.0) + reference, _ = sample_trajectory(plan, 2.0) assert math.hypot(base.x - reference[0], base.y - reference[1]) < 0.01 assert abs(angle_diff(base.yaw, reference[2])) < 0.01 @@ -113,7 +113,7 @@ def test_feedback_pulls_an_offset_base_onto_the_plan(): _run(task, base, seconds=2.0) - reference, _ = plan.sample(2.0) + reference, _ = sample_trajectory(plan, 2.0) assert math.hypot(base.x - reference[0], base.y - reference[1]) < 0.02 assert abs(angle_diff(base.yaw, reference[2])) < 0.02 @@ -141,7 +141,7 @@ def test_a_second_trajectory_starts_on_its_first_tick(): plan = _line((base.x, base.y, base.yaw), (base.x + 1.0, base.y, base.yaw), 2.0) task.execute(plan) _, t = _run(task, base, seconds=1.0, t=t) - assert abs(base.x - plan.sample(1.0)[0][0]) < 0.01 + assert abs(base.x - sample_trajectory(plan, 1.0)[0][0]) < 0.01 _run(task, base, seconds=2.0, t=t) assert task.get_status().state == TrajectoryState.COMPLETED @@ -150,11 +150,16 @@ def test_a_second_trajectory_starts_on_its_first_tick(): def test_rejects_trajectories_it_cannot_follow(): nan_velocity = _line((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), duration=2.0) - nan_velocity.points[5].velocities = [math.nan, 0.0, 0.0] + point = nan_velocity.points[5] + point.velocities = [math.nan, 0.0, 0.0] + nan_velocity.points[5] = point never_ends = _line((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), duration=2.0) - never_ends.points[-1].time_from_start = math.inf + point = never_ends.points[-1] + point.time_from_start.nanosec = 1_000_000_000 + never_ends.points[-1] = point two_columns = JointTrajectory( - points=[TrajectoryPoint(positions=[0.0, 0.0], velocities=[0.0, 0.0])] + header=header_now(), + points=[JointTrajectoryPoint(positions=[0.0, 0.0], velocities=[0.0, 0.0])], ) diagonal_too_fast = _line((0.0, 0.0, 0.0), (1.0, 1.0, 0.0), duration=1.0) diff --git a/dimos/control/tasks/pose_target_ik.py b/dimos/control/tasks/pose_target_ik.py index 22a3e38d027..f7d9ee52698 100644 --- a/dimos/control/tasks/pose_target_ik.py +++ b/dimos/control/tasks/pose_target_ik.py @@ -19,6 +19,7 @@ from abc import abstractmethod from collections import deque from collections.abc import Iterable, Mapping, Sequence +import copy from dataclasses import dataclass, field import math import threading @@ -26,6 +27,9 @@ from typing import TYPE_CHECKING import attrs +from dimos_generated.geometry_msgs.msg import Pose, PoseStamped +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header import numpy as np import pink from pydantic import Field @@ -46,12 +50,10 @@ ) from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.validation import prepare_robot_model -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.geometry import pose_from_matrix, pose_matrix from dimos.protocol.service.spec import BaseConfig from dimos.robot.assets.model import RobotModel from dimos.utils.logging_config import setup_logger -from dimos.utils.transform_utils import matrix_to_pose, pose_to_matrix if TYPE_CHECKING: from numpy.typing import NDArray @@ -290,7 +292,7 @@ def step( ) -> JointState | None: """Advance the persistent command trajectory by one bounded QP step.""" with self._command_state_lock: - command_state = JointState(self._command_state or measured_state) + command_state = copy.deepcopy(self._command_state or measured_state) command_increment_history = tuple( increment.copy() for increment in self._command_increment_history ) @@ -315,7 +317,7 @@ def step( with self._command_state_lock: if self._command_state_generation != generation: return None - self._command_state = JointState(result.command) + self._command_state = copy.deepcopy(result.command) self._command_increment_history.append(result.bounded_increment.copy()) return result.command @@ -508,14 +510,13 @@ def _frame_poses( ) positions = _seed_positions_for_mapping(seed, context.robot.mapping) q = self._q_from_dimos_positions(context.robot, positions) - base_world = pose_to_matrix(robot_model.base_pose) + base_world = pose_matrix(robot_model.base_pose.pose) poses: dict[str, PoseStamped] = {} for frame_name, frame_context in context.frames.items(): - pose = matrix_to_pose(base_world @ self._current_frame_matrix(frame_context, q)) + pose = pose_from_matrix(base_world @ self._current_frame_matrix(frame_context, q)) poses[frame_name] = PoseStamped( - frame_id=robot_model.base_link, - position=pose.position, - orientation=pose.orientation, + header=Header(frame_id=robot_model.base_link), + pose=Pose(position=pose.position, orientation=pose.orientation), ) return poses @@ -727,7 +728,9 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None: return None if result is None: return None - if result.name != list(self._joint_names) or len(result.position) != len(self._joint_names): + if list(result.name) != list(self._joint_names) or len(result.position) != len( + self._joint_names + ): logger.warning("Pink control step returned an invalid joint result", task=self._name) return None diff --git a/dimos/control/tasks/rpp_path_follower_task/rpp_path_follower_task.py b/dimos/control/tasks/rpp_path_follower_task/rpp_path_follower_task.py index 6c8dd1f7b51..6607ec70067 100644 --- a/dimos/control/tasks/rpp_path_follower_task/rpp_path_follower_task.py +++ b/dimos/control/tasks/rpp_path_follower_task/rpp_path_follower_task.py @@ -38,6 +38,9 @@ from pathlib import Path as _Path from typing import TYPE_CHECKING, Any +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.nav_msgs.msg import Path + from dimos.control.benchmarking.tuning import TuningConfig from dimos.control.benchmarking.velocity_profile import PathSpeedCap, VelocityProfileConfig from dimos.control.tasks.feedforward_gain_compensator import ( @@ -49,10 +52,7 @@ PathFollowerTaskConfig, ) from dimos.core.global_config import global_config as _gc -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import quaternion_euler, quaternion_from_euler from dimos.protocol.service.spec import BaseConfig from dimos.utils.logging_config import setup_logger @@ -87,7 +87,7 @@ def _with_tangent_headings(path: Path) -> Path: if poses is None or len(poses) < 2: return path - yaws = [p.orientation.euler[2] for p in poses] + yaws = [quaternion_euler(p.pose.orientation)[2] for p in poses] if max(yaws) - min(yaws) > _HEADING_DEGENERATE_EPS: return path # planner provided real per-pose headings; leave them. @@ -95,27 +95,28 @@ def _with_tangent_headings(path: Path) -> Path: n = len(poses) for i, p in enumerate(poses): if i < n - 1: - dx = poses[i + 1].position.x - p.position.x - dy = poses[i + 1].position.y - p.position.y + dx = poses[i + 1].pose.position.x - p.pose.position.x + dy = poses[i + 1].pose.position.y - p.pose.position.y yaw = ( math.atan2(dy, dx) if (dx * dx + dy * dy) > 1e-12 # Coincident waypoints: keep the previous heading rather than # snapping to 0. - else (new_poses[-1].orientation.euler[2] if new_poses else 0.0) + else (quaternion_euler(new_poses[-1].pose.orientation)[2] if new_poses else 0.0) ) else: # Last pose has no next segment — settle to the approach heading. - yaw = new_poses[-1].orientation.euler[2] if new_poses else 0.0 + yaw = quaternion_euler(new_poses[-1].pose.orientation)[2] if new_poses else 0.0 new_poses.append( PoseStamped( - ts=p.ts, - frame_id=p.frame_id, - position=Vector3(p.position.x, p.position.y, p.position.z), - orientation=Quaternion.from_euler(Vector3(0.0, 0.0, yaw)), + header=p.header, + pose=Pose( + position=Point(x=p.pose.position.x, y=p.pose.position.y, z=p.pose.position.z), + orientation=quaternion_from_euler(0.0, 0.0, yaw), + ), ) ) - return Path(ts=path.ts, frame_id=path.frame_id, poses=new_poses) + return Path(header=path.header, poses=new_poses) class RPPPathFollowerTask(PathFollowerTask): diff --git a/dimos/control/tasks/rpp_path_follower_task/test_rpp_path_follower_task.py b/dimos/control/tasks/rpp_path_follower_task/test_rpp_path_follower_task.py index c2ceb44fca7..d65870a14fc 100644 --- a/dimos/control/tasks/rpp_path_follower_task/test_rpp_path_follower_task.py +++ b/dimos/control/tasks/rpp_path_follower_task/test_rpp_path_follower_task.py @@ -19,6 +19,8 @@ from types import SimpleNamespace +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.std_msgs.msg import Header import pytest from dimos.control.benchmarking.paths import circle, straight_line @@ -41,9 +43,8 @@ RPPPathFollowerTask, create_task, ) -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.geometry import quaternion_euler, quaternion_from_euler +from dimos.msgs.time import time_from_seconds _JOINTS = ["go2/vx", "go2/vy", "go2/wz"] @@ -83,7 +84,10 @@ def _task(artifact_path: str, **params) -> RPPPathFollowerTask: def _odom(x=0.0, y=0.0, yaw=0.0) -> PoseStamped: return PoseStamped( - position=Vector3(x, y, 0.0), orientation=Quaternion.from_euler(Vector3(0.0, 0.0, yaw)) + header=Header(frame_id=""), + pose=Pose( + position=Point(x=x, y=y, z=0.0), orientation=quaternion_from_euler(0.0, 0.0, yaw) + ), ) @@ -232,20 +236,24 @@ def test_drives_to_arrival_forward_only(artifact_path): import math +from dimos_generated.nav_msgs.msg import Path + from dimos.control.tasks.rpp_path_follower_task.rpp_path_follower_task import ( _with_tangent_headings, ) -from dimos.msgs.nav_msgs.Path import Path def _ident_path(points): """Path with identity orientation on every pose (what MLSPlannerNative emits).""" return Path( - ts=0.0, - frame_id="odom", + header=Header(frame_id="odom"), poses=[ PoseStamped( - ts=0.0, position=Vector3(x, y, 0.0), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + header=Header(frame_id="", stamp=time_from_seconds(0.0)), + pose=Pose( + position=Point(x=x, y=y, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) for x, y in points ], @@ -255,7 +263,7 @@ def _ident_path(points): def test_degenerate_path_gets_tangent_headings(): # Straight +x then turn +y: first segments face 0 rad, the corner faces +pi/2. out = _with_tangent_headings(_ident_path([(0, 0), (1, 0), (2, 0), (2, 1), (2, 2)])) - yaws = [p.orientation.euler[2] for p in out.poses] + yaws = [quaternion_euler(p.pose.orientation)[2] for p in out.poses] assert math.isclose(yaws[0], 0.0, abs_tol=1e-6) assert math.isclose(yaws[1], 0.0, abs_tol=1e-6) assert math.isclose(yaws[2], math.pi / 2, abs_tol=1e-6) # tangent into the turn @@ -267,24 +275,25 @@ def test_degenerate_path_gets_tangent_headings(): def test_positions_unchanged_by_synthesis(): pts = [(0, 0), (1, 0), (1, 1)] out = _with_tangent_headings(_ident_path(pts)) - assert [(round(p.position.x, 6), round(p.position.y, 6)) for p in out.poses] == pts + assert [(round(p.pose.position.x, 6), round(p.pose.position.y, 6)) for p in out.poses] == pts def test_path_with_real_headings_is_untouched(): # A planner that already provides per-pose orientations must be left alone. oriented = Path( - ts=0.0, - frame_id="odom", + header=Header(frame_id="odom"), poses=[ PoseStamped( - ts=0.0, - position=Vector3(0, 0, 0), - orientation=Quaternion.from_euler(Vector3(0, 0, 0.3)), + header=Header(frame_id="", stamp=time_from_seconds(0.0)), + pose=Pose( + position=Point(x=0, y=0, z=0), orientation=quaternion_from_euler(0, 0, 0.3) + ), ), PoseStamped( - ts=0.0, - position=Vector3(1, 0, 0), - orientation=Quaternion.from_euler(Vector3(0, 0, 1.4)), + header=Header(frame_id="", stamp=time_from_seconds(0.0)), + pose=Pose( + position=Point(x=1, y=0, z=0), orientation=quaternion_from_euler(0, 0, 1.4) + ), ), ], ) diff --git a/dimos/control/tasks/teleop_ik_task/teleop_ik_task.py b/dimos/control/tasks/teleop_ik_task/teleop_ik_task.py index 0ce08bf6fed..752a2ca78cf 100644 --- a/dimos/control/tasks/teleop_ik_task/teleop_ik_task.py +++ b/dimos/control/tasks/teleop_ik_task/teleop_ik_task.py @@ -17,12 +17,15 @@ from __future__ import annotations from collections.abc import Iterable, Mapping +import copy from dataclasses import dataclass from enum import Enum import threading from typing import TYPE_CHECKING import attrs +from dimos_generated.geometry_msgs.msg import Pose, PoseStamped +from dimos_generated.std_msgs.msg import Header, UInt32 from dimos.control.task import CoordinatorState from dimos.control.tasks.pose_target_ik import ( @@ -35,8 +38,7 @@ ) from dimos.manipulation.planning.kinematics.config import PinkKinematicsConfig from dimos.manipulation.planning.spec.config import RobotModelConfig -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry import pose_from_matrix, pose_matrix from dimos.protocol.service.spec import BaseConfig from dimos.teleop.webxr.controller_types import Buttons @@ -173,12 +175,7 @@ def _on_controller_pose( ) -> bool: if hand not in self._bindings: return False - sample = PoseStamped( - ts=pose.ts if isinstance(pose, PoseStamped) else 0.0, - frame_id=pose.frame_id if isinstance(pose, PoseStamped) else "", - position=pose.position, - orientation=pose.orientation, - ) + sample = copy.deepcopy(pose) if isinstance(pose, PoseStamped) else PoseStamped(pose=pose) with self._lock: if self._session_state is _SessionState.ESTOPPED: return False @@ -187,11 +184,12 @@ def _on_controller_pose( state.last_update_time = t_now return True - def on_teleop_buttons(self, msg: Buttons, t_now: float) -> bool: + def on_teleop_buttons(self, msg: UInt32, t_now: float) -> bool: """Update the all-bound-hands deadman condition.""" + buttons = Buttons(data=msg.data) grip_by_hand = { - OperatorHand.LEFT: msg.left_grip, - OperatorHand.RIGHT: msg.right_grip, + OperatorHand.LEFT: buttons.left_grip, + OperatorHand.RIGHT: buttons.right_grip, } with self._lock: self._last_button_update_time = t_now @@ -267,11 +265,19 @@ def _frame_target_snapshot(self, state: CoordinatorState) -> FrameTargetSnapshot robot_reference = hand_state.robot_reference if current is None or controller_reference is None or robot_reference is None: return None - delta = current - controller_reference + current_matrix = pose_matrix(current.pose) + reference_matrix = pose_matrix(controller_reference.pose) + target_matrix = pose_matrix(robot_reference.pose) + target_matrix[:3, 3] += current_matrix[:3, 3] - reference_matrix[:3, 3] + target_matrix[:3, :3] = ( + current_matrix[:3, :3] @ reference_matrix[:3, :3].T @ target_matrix[:3, :3] + ) targets[binding.target_frame] = PoseStamped( - frame_id=self._teleop_config.robot_model.base_pose.frame_id, - position=robot_reference.position + delta.position, - orientation=delta.orientation * robot_reference.orientation, + header=Header( + frame_id=self._teleop_config.robot_model.base_pose.header.frame_id, + stamp=current.header.stamp, + ), + pose=pose_from_matrix(target_matrix), ) update_times.append(hand_state.last_update_time) return FrameTargetSnapshot( diff --git a/dimos/control/tasks/teleop_ik_task/test_teleop_ik_task.py b/dimos/control/tasks/teleop_ik_task/test_teleop_ik_task.py index fb88bd40486..a7bf6309e41 100644 --- a/dimos/control/tasks/teleop_ik_task/test_teleop_ik_task.py +++ b/dimos/control/tasks/teleop_ik_task/test_teleop_ik_task.py @@ -17,6 +17,9 @@ from pathlib import Path from typing import cast +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header, UInt32 import pytest from pytest_mock import MockerFixture @@ -31,10 +34,6 @@ create_task, ) from dimos.manipulation.planning.spec.config import RobotModelConfig -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.robot.assets.model import RobotModel from dimos.teleop.webxr.controller_types import Buttons @@ -69,12 +68,18 @@ def _solver(mocker: MockerFixture) -> PinkPoseTargetSolver: solver = mocker.Mock(spec=PinkPoseTargetSolver) solver.frame_poses.return_value = { "left_tool": PoseStamped( - position=Vector3(1.0, 0.0, 0.0), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + header=Header(frame_id=""), + pose=Pose( + position=Point(x=1.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), "right_tool": PoseStamped( - position=Vector3(-1.0, 0.0, 0.0), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + header=Header(frame_id=""), + pose=Pose( + position=Point(x=-1.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), } solver.step.return_value = JointState( @@ -95,11 +100,11 @@ def _buttons( *, left: bool = False, right: bool = False, -) -> Buttons: +) -> UInt32: buttons = Buttons() buttons.left_grip = left buttons.right_grip = right - return buttons + return buttons.to_message() def test_face_buttons_do_not_engage_arm_teleop(mocker: MockerFixture) -> None: @@ -111,7 +116,7 @@ def test_face_buttons_do_not_engage_arm_teleop(mocker: MockerFixture) -> None: buttons = Buttons() buttons.right_primary = True - task.on_teleop_buttons(buttons, 1.0) + task.on_teleop_buttons(buttons.to_message(), 1.0) task.on_right_cartesian_command(_pose(0.5), 1.0) assert task.compute(_state()) is None @@ -119,8 +124,10 @@ def test_face_buttons_do_not_engage_arm_teleop(mocker: MockerFixture) -> None: def _pose(x: float) -> PoseStamped: return PoseStamped( - position=Vector3(x, 0.0, 0.0), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + header=Header(frame_id=""), + pose=Pose( + position=Point(x=x, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), ) @@ -179,7 +186,7 @@ def test_single_binding_tracks_relative_controller_motion(mocker: MockerFixture) assert output is not None target = solver.step.call_args.args[0]["right_tool"] - assert target.position.x == pytest.approx(-0.8) + assert target.pose.position.x == pytest.approx(-0.8) def test_bimanual_task_requires_both_hands_and_releases_atomically( @@ -306,7 +313,7 @@ def test_bimanual_timeout_clears_both_sides_and_reengagement_recaptures( def test_stale_deadman_stops_fresh_pose_streams( mocker: MockerFixture, bindings: tuple[TeleopHandBinding, ...], - buttons: Buttons, + buttons: UInt32, ) -> None: solver = _solver(mocker) task = TeleopIKTask("teleop", _config(bindings, timeout=0.2), solver=solver) diff --git a/dimos/control/tasks/test_pose_target_ik.py b/dimos/control/tasks/test_pose_target_ik.py index 0869431f589..fac8b268b6c 100644 --- a/dimos/control/tasks/test_pose_target_ik.py +++ b/dimos/control/tasks/test_pose_target_ik.py @@ -16,6 +16,9 @@ from pathlib import Path +from dimos_generated.geometry_msgs.msg import Pose, PoseStamped +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from pytest_mock import MockerFixture @@ -31,8 +34,6 @@ ) from dimos.manipulation.planning.kinematics.pink_solver import _PinkSolverCore from dimos.manipulation.planning.spec.config import RobotModelConfig -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.robot.assets.model import RobotModel @@ -97,7 +98,9 @@ def _solver(mocker: MockerFixture, positions: list[float] | None = None) -> Pink solver.step.return_value = JointState( name=["arm/a", "arm/b"], position=positions or [0.01, -0.01] ) - solver.frame_poses.return_value = {"tool": PoseStamped(frame_id="base")} + solver.frame_poses.return_value = { + "tool": PoseStamped(header=Header(frame_id="base"), pose=Pose()) + } return solver @@ -116,7 +119,7 @@ def _snapshot( extra_joint_positions: dict[str, float] | None = None, ) -> FrameTargetSnapshot: return FrameTargetSnapshot( - targets=targets or {"tool": PoseStamped(frame_id="world")}, + targets=targets or {"tool": PoseStamped(header=Header(frame_id="world"), pose=Pose())}, last_update_time=last_update_time, extra_joint_positions=extra_joint_positions or {}, ) @@ -145,7 +148,7 @@ def advance(**kwargs: object) -> _StreamingStepResult: ) step = mocker.patch.object(solver, "_step_frame_targets", side_effect=advance) - targets = {"tool": PoseStamped()} + targets = {"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())} assert solver.step(targets, JointState(name=["arm/a", "arm/b"], position=[0.0, 0.0]), 0.01) assert solver.step( @@ -180,7 +183,7 @@ def test_pose_target_solver_reset_reseeds_from_feedback(mocker: MockerFixture) - ), ], ) - targets = {"tool": PoseStamped()} + targets = {"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())} assert solver.step(targets, JointState(name=["arm/a", "arm/b"], position=[0.0, 0.0]), 0.01) solver.reset() @@ -209,7 +212,7 @@ def reset_during_step(**_kwargs: object) -> _StreamingStepResult: mocker.patch.object(solver, "_step_frame_targets", side_effect=reset_during_step) result = solver.step( - {"tool": PoseStamped()}, + {"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())}, JointState(name=["arm/a", "arm/b"], position=[0.0, 0.0]), 0.01, ) @@ -362,7 +365,7 @@ def test_current_frame_poses_uses_live_coordinator_seed(mocker: MockerFixture) - poses = task.current_frame_poses(_state(positions={"arm/a": 0.2, "arm/b": 0.3}), ["tool"]) - assert poses == {"tool": PoseStamped(frame_id="base")} + assert poses == {"tool": PoseStamped(header=Header(frame_id="base"), pose=Pose())} seed = solver.frame_poses.call_args.args[0] assert seed.name == ["arm/a", "arm/b"] assert seed.position == [0.2, 0.3] diff --git a/dimos/control/tasks/trajectory_task/test_trajectory_task.py b/dimos/control/tasks/trajectory_task/test_trajectory_task.py index cbef3f88cf9..cf126d5e73f 100644 --- a/dimos/control/tasks/trajectory_task/test_trajectory_task.py +++ b/dimos/control/tasks/trajectory_task/test_trajectory_task.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.builtin_interfaces.msg import Duration +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint import pytest from dimos.control.task import ( @@ -23,8 +25,7 @@ JointTrajectoryTaskConfig, TrajectoryExecutionStatus, ) -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +from dimos.msgs.time import duration_from_seconds, header_now @pytest.mark.parametrize("single_point", [False, True]) @@ -36,7 +37,9 @@ def test_jtt_does_not_pull_joint_back_after_another_task_moves_it(single_point, ) first_target = 1.0 if preempted else 0.1 first = JointTrajectory( - joint_names=["joint"], points=[TrajectoryPoint(positions=[first_target])] + header=header_now(), + joint_names=["joint"], + points=[JointTrajectoryPoint(positions=[first_target])], ) assert task.execute(first, {}).status is TrajectoryExecutionStatus.ACCEPTED output = task.compute(state) @@ -51,15 +54,16 @@ def test_jtt_does_not_pull_joint_back_after_another_task_moves_it(single_point, joints=JointStateSnapshot(joint_positions={"joint": -1.0}), t_now=1.0, dt=0.01 ) assert task.compute(state) is None - points = [TrajectoryPoint(positions=[-1.2])] + points = [JointTrajectoryPoint(positions=[-1.2])] if not single_point: points = [ - TrajectoryPoint(positions=[-1.0]), - TrajectoryPoint(positions=[-1.2], time_from_start=0.1), + JointTrajectoryPoint(positions=[-1.0]), + JointTrajectoryPoint(positions=[-1.2], time_from_start=duration_from_seconds(0.1)), ] assert ( task.execute( - JointTrajectory(joint_names=["joint"], points=points), state.joints.joint_positions + JointTrajectory(header=header_now(), joint_names=["joint"], points=points), + state.joints.joint_positions, ).status is TrajectoryExecutionStatus.ACCEPTED ) @@ -84,7 +88,9 @@ def test_completed_joint_reanchors_while_other_joint_keeps_command_continuity(): task = JointTrajectoryTask(JointTrajectoryTaskConfig(joint_names=["finished", "running"])) measured = JointStateSnapshot(joint_positions={"finished": 0.0, "running": 0.0}) initial = JointTrajectory( - joint_names=["finished", "running"], points=[TrajectoryPoint(positions=[0.1, 1.0])] + header=header_now(), + joint_names=["finished", "running"], + points=[JointTrajectoryPoint(positions=[0.1, 1.0])], ) assert task.execute(initial, {}).status is TrajectoryExecutionStatus.ACCEPTED output = task.compute(CoordinatorState(joints=measured, t_now=0.1, dt=0.1)) @@ -92,10 +98,11 @@ def test_completed_joint_reanchors_while_other_joint_keeps_command_continuity(): assert output.positions == pytest.approx([0.1, 0.1]) measured.joint_positions["finished"] = -1.0 replacement = JointTrajectory( + header=header_now(), joint_names=["finished"], points=[ - TrajectoryPoint(positions=[-1.0]), - TrajectoryPoint(positions=[0.1], time_from_start=2.0), + JointTrajectoryPoint(positions=[-1.0]), + JointTrajectoryPoint(positions=[0.1], time_from_start=duration_from_seconds(2.0)), ], ) assert ( @@ -116,7 +123,9 @@ def test_completed_joint_reanchors_while_other_joint_keeps_command_continuity(): ) def test_completed_trajectory_does_not_bypass_start_validation(positions, expected): task = JointTrajectoryTask(JointTrajectoryTaskConfig(joint_names=["joint"])) - initial = JointTrajectory(joint_names=["joint"], points=[TrajectoryPoint(positions=[0.1])]) + initial = JointTrajectory( + header=header_now(), joint_names=["joint"], points=[JointTrajectoryPoint(positions=[0.1])] + ) assert task.execute(initial, {}).status is TrajectoryExecutionStatus.ACCEPTED task.compute( CoordinatorState( @@ -124,10 +133,36 @@ def test_completed_trajectory_does_not_bypass_start_validation(positions, expect ) ) trajectory = JointTrajectory( + header=header_now(), joint_names=["joint"], points=[ - TrajectoryPoint(positions=[0.1]), - TrajectoryPoint(positions=[1.0], time_from_start=2.0), + JointTrajectoryPoint(positions=[0.1]), + JointTrajectoryPoint(positions=[1.0], time_from_start=duration_from_seconds(2.0)), ], ) assert task.execute(trajectory, positions).status is expected + + +@pytest.mark.parametrize("nanosec,accepted", [(0, True), (1_000_000_000, False)]) +def test_jtt_validates_normalized_durations_without_float_ordering(nanosec, accepted): + task = JointTrajectoryTask(JointTrajectoryTaskConfig(joint_names=["joint"])) + # Far enough out that float seconds cannot distinguish adjacent nanoseconds. + trajectory = JointTrajectory( + joint_names=["joint"], + points=[ + JointTrajectoryPoint(positions=[0.0]), + JointTrajectoryPoint( + positions=[0.1], time_from_start=Duration(sec=100_000_000, nanosec=nanosec) + ), + JointTrajectoryPoint( + positions=[0.2], time_from_start=Duration(sec=100_000_000, nanosec=1) + ), + ], + ) + result = task.execute(trajectory, {"joint": 0.0}) + expected = ( + TrajectoryExecutionStatus.ACCEPTED + if accepted + else TrajectoryExecutionStatus.INVALID_TRAJECTORY + ) + assert result.status == expected diff --git a/dimos/control/tasks/trajectory_task/trajectory_task.py b/dimos/control/tasks/trajectory_task/trajectory_task.py index 7a20538788e..50697d2da2a 100644 --- a/dimos/control/tasks/trajectory_task/trajectory_task.py +++ b/dimos/control/tasks/trajectory_task/trajectory_task.py @@ -28,6 +28,9 @@ import math from typing import TYPE_CHECKING, Annotated, Any +from dimos_generated.dimos_msgs.msg import TrajectoryStatus +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint from pydantic import BeforeValidator, ConfigDict, Field from pydantic.dataclasses import dataclass as pydantic_dataclass @@ -38,10 +41,8 @@ JointCommandOutput, ResourceClaim, ) -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint -from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState, TrajectoryStatus +from dimos.msgs.time import duration_from_seconds, header_now, to_nanoseconds +from dimos.msgs.trajectory import TrajectoryState, sample_trajectory, trajectory_duration from dimos.protocol.service.spec import BaseConfig from dimos.utils.logging_config import setup_logger @@ -241,8 +242,9 @@ def on_joint_command(self, msg: JointState, t_now: float) -> bool: if not selected: return False trajectory = JointTrajectory( + header=header_now(), joint_names=[name for name, _ in selected], - points=[TrajectoryPoint(positions=[position for _, position in selected])], + points=[JointTrajectoryPoint(positions=[position for _, position in selected])], ) return self.execute(trajectory, {}).status is TrajectoryExecutionStatus.ACCEPTED @@ -274,7 +276,7 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None: self._pending_start = False elapsed = max(0.0, state.t_now - run.start_time) self._last_elapsed = max(self._last_elapsed, elapsed) - desired = run.trajectory.sample(elapsed)[0][index] + desired = sample_trajectory(run.trajectory, elapsed)[0][index] current = self._commanded_positions.get(joint_name) if current is None: current = state.joints.get_position(joint_name) @@ -287,7 +289,7 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None: self._commanded_positions[joint_name] = commanded final_position = run.trajectory.points[-1].positions[index] - nominal_complete = elapsed >= run.trajectory.duration + nominal_complete = elapsed >= trajectory_duration(run.trajectory) reached = math.isclose(commanded, final_position, abs_tol=1e-9) if nominal_complete and reached: del self._motions[joint_name] @@ -355,9 +357,9 @@ def _validate_trajectory(self, trajectory: JointTrajectory) -> bool: logger.warning("Empty trajectory for %s", self._name) return False width = len(joint_names) - previous_time: float | None = None + previous_time: int | None = None for index, point in enumerate(trajectory.points): - if len(point.positions) != width or len(point.velocities) != width: + if len(point.positions) != width or len(point.velocities) not in (0, width): logger.warning("Trajectory point %d for %s has invalid width", index, self._name) return False if not all(math.isfinite(value) for value in point.positions): @@ -370,17 +372,19 @@ def _validate_trajectory(self, trajectory: JointTrajectory) -> bool: "Trajectory point %d for %s has non-finite velocities", index, self._name ) return False - if not math.isfinite(point.time_from_start): - logger.warning("Trajectory point %d for %s has non-finite time", index, self._name) + try: + point_time = to_nanoseconds(point.time_from_start) + except ValueError: + logger.warning("Trajectory point %d for %s has invalid duration", index, self._name) return False - if index == 0 and point.time_from_start != 0.0: + if index == 0 and point_time != 0.0: logger.warning("Trajectory for %s must start at t=0", self._name) return False - if previous_time is not None and point.time_from_start <= previous_time: + if previous_time is not None and point_time <= previous_time: logger.warning("Trajectory for %s has non-increasing timestamps", self._name) return False - previous_time = point.time_from_start - if len(trajectory.points) > 1 and trajectory.duration <= 0.0: + previous_time = point_time + if len(trajectory.points) > 1 and trajectory_duration(trajectory) <= 0.0: logger.warning("Trajectory for %s has nonpositive duration", self._name) return False return True @@ -448,14 +452,16 @@ def execute( trajectory = JointTrajectory( joint_names=list(trajectory.joint_names), points=[ - TrajectoryPoint( + JointTrajectoryPoint( time_from_start=trajectory.points[0].time_from_start, positions=first_positions, - velocities=list(trajectory.points[0].velocities), + velocities=trajectory.points[0].velocities, + accelerations=trajectory.points[0].accelerations, + effort=trajectory.points[0].effort, ), - *trajectory.points[1:], + *list(trajectory.points)[1:], ], - timestamp=trajectory.timestamp, + header=trajectory.header, ) run = _TrajectoryRun(trajectory) @@ -464,14 +470,14 @@ def execute( if len(trajectory.points) > 1 and joint_name not in self._commanded_positions: self._commanded_positions[joint_name] = current_positions[joint_name] self._trajectory = trajectory - self._last_duration = trajectory.duration + self._last_duration = trajectory_duration(trajectory) self._last_elapsed = 0.0 self._pending_start = True # Start time set on first compute() self._state = TrajectoryState.EXECUTING logger.info( f"Executing trajectory on {self._name}: " - f"{len(trajectory.points)} points, duration={trajectory.duration:.3f}s" + f"{len(trajectory.points)} points, duration={trajectory_duration(trajectory):.3f}s" ) return TrajectoryExecutionResult(TrajectoryExecutionStatus.ACCEPTED) @@ -528,7 +534,7 @@ def _active_run_status(self, t_now: float) -> tuple[float, float, float]: remaining_times: list[float] = [] for run in runs: elapsed = 0.0 if run.start_time is None else max(0.0, t_now - run.start_time) - duration = run.trajectory.duration + duration = trajectory_duration(run.trajectory) progress = 0.0 if duration <= 0.0 else min(1.0, elapsed / duration) progresses.append(progress) elapsed_times.append(elapsed) @@ -543,17 +549,21 @@ def get_status(self, t_now: float) -> TrajectoryStatus: if self._state == TrajectoryState.EXECUTING: progress, elapsed, remaining = self._active_run_status(t_now) return TrajectoryStatus( + header=header_now(), state=self._state, progress=progress, - time_elapsed=elapsed, - time_remaining=remaining, + time_elapsed=duration_from_seconds(elapsed), + time_remaining=duration_from_seconds(remaining), ) completed = self._state == TrajectoryState.COMPLETED return TrajectoryStatus( + header=header_now(), state=self._state, progress=1.0 if completed else 0.0, - time_elapsed=self._last_duration if completed else self._last_elapsed, - time_remaining=0.0, + time_elapsed=duration_from_seconds( + self._last_duration if completed else self._last_elapsed + ), + time_remaining=duration_from_seconds(0.0), ) diff --git a/dimos/control/tasks/velocity_profiler.py b/dimos/control/tasks/velocity_profiler.py index 00d481bb6de..51d00519e63 100644 --- a/dimos/control/tasks/velocity_profiler.py +++ b/dimos/control/tasks/velocity_profiler.py @@ -23,11 +23,10 @@ from __future__ import annotations +from dimos_generated.nav_msgs.msg import Path import numpy as np from numpy.typing import NDArray -from dimos.msgs.nav_msgs.Path import Path - class VelocityProfiler: """Compute an optimal speed profile along a path.""" @@ -60,7 +59,7 @@ def compute_profile(self, path: Path) -> NDArray[np.float64]: if len(path.poses) < 2: return np.array([self._min_speed]) - pts = np.array([[p.position.x, p.position.y] for p in path.poses]) + pts = np.array([[p.pose.position.x, p.pose.position.y] for p in path.poses]) curvatures = self._compute_curvatures(pts) max_speeds = self._curvature_speed_limits(curvatures) velocities = self._acceleration_pass(pts, max_speeds, forward=True) diff --git a/dimos/control/tasks/velocity_task/test_velocity_task.py b/dimos/control/tasks/velocity_task/test_velocity_task.py index 800595d0a6b..1a4fbaa9bdf 100644 --- a/dimos/control/tasks/velocity_task/test_velocity_task.py +++ b/dimos/control/tasks/velocity_task/test_velocity_task.py @@ -16,12 +16,13 @@ from __future__ import annotations +from dimos_generated.sensor_msgs.msg import JointState + from dimos.control.task import CoordinatorState, JointStateSnapshot from dimos.control.tasks.velocity_task.velocity_task import ( JointVelocityTask, JointVelocityTaskConfig, ) -from dimos.msgs.sensor_msgs.JointState import JointState def _task() -> JointVelocityTask: diff --git a/dimos/control/tasks/velocity_task/velocity_task.py b/dimos/control/tasks/velocity_task/velocity_task.py index e9538cd56c1..d88071be47d 100644 --- a/dimos/control/tasks/velocity_task/velocity_task.py +++ b/dimos/control/tasks/velocity_task/velocity_task.py @@ -38,7 +38,7 @@ from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: - from dimos.msgs.sensor_msgs.JointState import JointState + from dimos_generated.sensor_msgs.msg import JointState logger = setup_logger() diff --git a/dimos/control/teleop_coordinator.py b/dimos/control/teleop_coordinator.py index af8baf2a694..1052d8fba20 100644 --- a/dimos/control/teleop_coordinator.py +++ b/dimos/control/teleop_coordinator.py @@ -14,12 +14,11 @@ """Control coordinator carrying spatial poses and keyboard end-effector twists.""" +from dimos_generated.geometry_msgs.msg import PoseStamped, TwistStamped +from dimos_generated.std_msgs.msg import Float32, UInt32 + from dimos.control.coordinator import ControlCoordinator from dimos.core.stream import In -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped -from dimos.msgs.std_msgs.Float32 import Float32 -from dimos.teleop.webxr.controller_types import Buttons class TeleopControlCoordinator(ControlCoordinator): @@ -29,5 +28,5 @@ class TeleopControlCoordinator(ControlCoordinator): right_cartesian_command: In[PoseStamped] left_gripper_command: In[Float32] right_gripper_command: In[Float32] - teleop_buttons: In[Buttons] + teleop_buttons: In[UInt32] ee_twist_command: In[TwistStamped] diff --git a/dimos/control/test_control.py b/dimos/control/test_control.py index bff9528cf8f..2bda2c44147 100644 --- a/dimos/control/test_control.py +++ b/dimos/control/test_control.py @@ -23,6 +23,9 @@ from typing import Any from unittest.mock import MagicMock +from dimos_generated.geometry_msgs.msg import Twist, TwistStamped, Vector3 +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint import pytest from dimos.control._control_test_helpers import RecordingTask @@ -63,12 +66,8 @@ from dimos.hardware.manipulators.spec import ManipulatorAdapter from dimos.hardware.spec import JointLimits from dimos.hardware.whole_body.spec import MotorState, WholeBodyAdapter -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint -from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState +from dimos.msgs.time import duration_from_seconds, header_now, to_seconds +from dimos.msgs.trajectory import TrajectoryState @pytest.fixture @@ -115,17 +114,18 @@ def trajectory_task(): def simple_trajectory(): """Create a simple 2-point trajectory.""" return JointTrajectory( + header=header_now(), joint_names=["arm/joint1", "arm/joint2", "arm/joint3"], points=[ - TrajectoryPoint( + JointTrajectoryPoint( positions=[0.0, 0.0, 0.0], velocities=[0.0, 0.0, 0.0], - time_from_start=0.0, + time_from_start=duration_from_seconds(0.0), ), - TrajectoryPoint( + JointTrajectoryPoint( positions=[1.0, 0.5, 0.25], velocities=[0.0, 0.0, 0.0], - time_from_start=1.0, + time_from_start=duration_from_seconds(1.0), ), ], ) @@ -384,7 +384,7 @@ def test_map_twist_to_base_joints_routes_planar_twist_via_joint_command( dispatch = mocker.patch.object(coordinator, "_dispatch") coordinator._map_twist_to_base_joints( - Twist(linear=[1.0, 2.0, 0.0], angular=[0.0, 0.0, 3.0]) + Twist(linear=Vector3(x=1.0, y=2.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=3.0)) ) stream, joint_state = dispatch.call_args.args @@ -531,7 +531,11 @@ def test_removing_trajectory_task_allows_replacement(self, make_coordinator): def test_execute_and_cancel_without_trajectory_task_are_semantic(self, make_coordinator): coordinator = make_coordinator() - execute_result = coordinator.execute_trajectory(JointTrajectory()) + execute_result = coordinator.execute_trajectory( + JointTrajectory( + header=header_now(), + ) + ) cancel_result = coordinator.cancel_trajectory() assert execute_result.status is TrajectoryExecutionStatus.NO_TRAJECTORY_TASK @@ -624,21 +628,30 @@ def test_status_snapshot_is_non_destructive(self, trajectory_task, simple_trajec trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=10.0, dt=0.01)) active = trajectory_task.get_status(10.25) - assert active.state is TrajectoryState.EXECUTING + assert active.state == TrajectoryState.EXECUTING assert active.progress == pytest.approx(0.25) trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=11.5, dt=0.01)) terminal = trajectory_task.get_status(11.5) - assert terminal.state is TrajectoryState.COMPLETED + assert terminal.state == TrajectoryState.COMPLETED assert terminal.progress == pytest.approx(1.0) - assert trajectory_task.get_status(11.6).state is TrajectoryState.COMPLETED + assert trajectory_task.get_status(11.6).state == TrajectoryState.COMPLETED def test_execute_partial_subset_and_claims_full_configuration(self, trajectory_task): trajectory = JointTrajectory( + header=header_now(), joint_names=["arm/joint2", "arm/joint3"], points=[ - TrajectoryPoint(positions=[0.0, 0.0], velocities=[0.0, 0.0], time_from_start=0.0), - TrajectoryPoint(positions=[0.5, 1.0], velocities=[0.0, 0.0], time_from_start=1.0), + JointTrajectoryPoint( + positions=[0.0, 0.0], + velocities=[0.0, 0.0], + time_from_start=duration_from_seconds(0.0), + ), + JointTrajectoryPoint( + positions=[0.5, 1.0], + velocities=[0.0, 0.0], + time_from_start=duration_from_seconds(1.0), + ), ], ) @@ -684,41 +697,82 @@ def test_execute_accepts_start_position_at_tolerance(self, trajectory_task, simp "trajectory", [ JointTrajectory( + header=header_now(), joint_names=[], - points=[TrajectoryPoint(time_from_start=0.0, positions=[], velocities=[])], + points=[ + JointTrajectoryPoint( + time_from_start=duration_from_seconds(0.0), positions=[], velocities=[] + ) + ], ), JointTrajectory( + header=header_now(), joint_names=["arm/joint1", "arm/joint1"], points=[ - TrajectoryPoint( - time_from_start=0.0, positions=[0.0, 0.0], velocities=[0.0, 0.0] + JointTrajectoryPoint( + time_from_start=duration_from_seconds(0.0), + positions=[0.0, 0.0], + velocities=[0.0, 0.0], ) ], ), JointTrajectory( + header=header_now(), joint_names=["arm/missing"], - points=[TrajectoryPoint(time_from_start=0.0, positions=[0.0], velocities=[0.0])], + points=[ + JointTrajectoryPoint( + time_from_start=duration_from_seconds(0.0), + positions=[0.0], + velocities=[0.0], + ) + ], ), - JointTrajectory(joint_names=["arm/joint1"], points=[]), + JointTrajectory(header=header_now(), joint_names=["arm/joint1"], points=[]), JointTrajectory( + header=header_now(), joint_names=["arm/joint1"], - points=[TrajectoryPoint(time_from_start=0.0, positions=[], velocities=[0.0])], + points=[ + JointTrajectoryPoint( + time_from_start=duration_from_seconds(0.0), positions=[], velocities=[0.0] + ) + ], ), JointTrajectory( + header=header_now(), joint_names=["arm/joint1"], points=[ - TrajectoryPoint(time_from_start=0.0, positions=[float("nan")], velocities=[0.0]) + JointTrajectoryPoint( + time_from_start=duration_from_seconds(0.0), + positions=[float("nan")], + velocities=[0.0], + ) ], ), JointTrajectory( + header=header_now(), joint_names=["arm/joint1"], - points=[TrajectoryPoint(time_from_start=0.1, positions=[0.0], velocities=[0.0])], + points=[ + JointTrajectoryPoint( + time_from_start=duration_from_seconds(0.1), + positions=[0.0], + velocities=[0.0], + ) + ], ), JointTrajectory( + header=header_now(), joint_names=["arm/joint1"], points=[ - TrajectoryPoint(time_from_start=0.0, positions=[0.0], velocities=[0.0]), - TrajectoryPoint(time_from_start=0.0, positions=[1.0], velocities=[0.0]), + JointTrajectoryPoint( + time_from_start=duration_from_seconds(0.0), + positions=[0.0], + velocities=[0.0], + ), + JointTrajectoryPoint( + time_from_start=duration_from_seconds(0.0), + positions=[1.0], + velocities=[0.0], + ), ], ), ], @@ -737,10 +791,15 @@ def test_invalid_partial_inputs_reject_before_state_changes(self, trajectory_tas def test_compute_emits_active_subset_and_retains_final_target(self, trajectory_task): trajectory = JointTrajectory( + header=header_now(), joint_names=["arm/joint2"], points=[ - TrajectoryPoint(positions=[0.0], velocities=[0.0], time_from_start=0.0), - TrajectoryPoint(positions=[1.0], velocities=[0.0], time_from_start=1.0), + JointTrajectoryPoint( + positions=[0.0], velocities=[0.0], time_from_start=duration_from_seconds(0.0) + ), + JointTrajectoryPoint( + positions=[1.0], velocities=[0.0], time_from_start=duration_from_seconds(1.0) + ), ], ) assert ( @@ -764,17 +823,27 @@ def test_compute_emits_active_subset_and_retains_final_target(self, trajectory_t def test_replacement_reset_and_cancel_clear_active_subset(self, trajectory_task): first = JointTrajectory( + header=header_now(), joint_names=["arm/joint1"], points=[ - TrajectoryPoint(positions=[0.0], velocities=[0.0], time_from_start=0.0), - TrajectoryPoint(positions=[1.0], velocities=[0.0], time_from_start=1.0), + JointTrajectoryPoint( + positions=[0.0], velocities=[0.0], time_from_start=duration_from_seconds(0.0) + ), + JointTrajectoryPoint( + positions=[1.0], velocities=[0.0], time_from_start=duration_from_seconds(1.0) + ), ], ) second = JointTrajectory( + header=header_now(), joint_names=["arm/joint3"], points=[ - TrajectoryPoint(positions=[2.0], velocities=[0.0], time_from_start=0.0), - TrajectoryPoint(positions=[3.0], velocities=[0.0], time_from_start=1.0), + JointTrajectoryPoint( + positions=[2.0], velocities=[0.0], time_from_start=duration_from_seconds(0.0) + ), + JointTrajectoryPoint( + positions=[3.0], velocities=[0.0], time_from_start=duration_from_seconds(1.0) + ), ], ) assert ( @@ -858,8 +927,9 @@ def test_one_point_stream_target_is_velocity_bounded(self): ) ) target = JointTrajectory( + header=header_now(), joint_names=["arm/joint1"], - points=[TrajectoryPoint(positions=[1.0])], + points=[JointTrajectoryPoint(positions=[1.0])], ) state = JointStateSnapshot(joint_positions={"arm/joint1": 0.0}) @@ -901,16 +971,19 @@ def test_replacement_anchors_at_bounded_command_and_preserves_other_joints(self) ) state = JointStateSnapshot(joint_positions={"arm/joint1": 0.0, "arm/joint2": 0.0}) first = JointTrajectory( + header=header_now(), joint_names=["arm/joint1"], - points=[TrajectoryPoint(positions=[1.0])], + points=[JointTrajectoryPoint(positions=[1.0])], ) other = JointTrajectory( + header=header_now(), joint_names=["arm/joint2"], - points=[TrajectoryPoint(positions=[-1.0])], + points=[JointTrajectoryPoint(positions=[-1.0])], ) replacement = JointTrajectory( + header=header_now(), joint_names=["arm/joint1"], - points=[TrajectoryPoint(positions=[-1.0])], + points=[JointTrajectoryPoint(positions=[-1.0])], ) task.execute(first, {}) @@ -933,21 +1006,25 @@ def test_partial_replacement_status_includes_untouched_joint_run(self): ) state = JointStateSnapshot(joint_positions={"arm/joint1": 0.0, "arm/joint2": 0.0}) long_trajectory = JointTrajectory( + header=header_now(), joint_names=["arm/joint1", "arm/joint2"], points=[ - TrajectoryPoint(positions=[0.0, 0.0], velocities=[0.0, 0.0]), - TrajectoryPoint( + JointTrajectoryPoint(positions=[0.0, 0.0], velocities=[0.0, 0.0]), + JointTrajectoryPoint( positions=[10.0, 10.0], velocities=[0.0, 0.0], - time_from_start=10.0, + time_from_start=duration_from_seconds(10.0), ), ], ) replacement = JointTrajectory( + header=header_now(), joint_names=["arm/joint1"], points=[ - TrajectoryPoint(positions=[2.5], velocities=[0.0]), - TrajectoryPoint(positions=[3.5], velocities=[0.0], time_from_start=1.0), + JointTrajectoryPoint(positions=[2.5], velocities=[0.0]), + JointTrajectoryPoint( + positions=[3.5], velocities=[0.0], time_from_start=duration_from_seconds(1.0) + ), ], ) @@ -963,10 +1040,10 @@ def test_partial_replacement_status_includes_untouched_joint_run(self): status = task.get_status(13.5) - assert status.state is TrajectoryState.EXECUTING + assert status.state == TrajectoryState.EXECUTING assert status.progress == pytest.approx(0.35) - assert status.time_elapsed == pytest.approx(3.5) - assert status.time_remaining == pytest.approx(6.5) + assert to_seconds(status.time_elapsed) == pytest.approx(3.5) + assert to_seconds(status.time_remaining) == pytest.approx(6.5) def test_cancel_trajectory(self, trajectory_task, simple_trajectory): trajectory_task.execute(simple_trajectory, trajectory_start_positions(simple_trajectory)) @@ -1191,10 +1268,19 @@ def test_partial_trajectory_and_gripper_command_share_hardware_write(self, mocke limits=[(0.0, 1.0)], ) trajectory = JointTrajectory( + header=header_now(), joint_names=["arm/joint1", "arm/joint2"], points=[ - TrajectoryPoint(positions=[0.0, 0.0], velocities=[0.0, 0.0], time_from_start=0.0), - TrajectoryPoint(positions=[0.5, 0.5], velocities=[0.0, 0.0], time_from_start=1.0), + JointTrajectoryPoint( + positions=[0.0, 0.0], + velocities=[0.0, 0.0], + time_from_start=duration_from_seconds(0.0), + ), + JointTrajectoryPoint( + positions=[0.5, 0.5], + velocities=[0.0, 0.0], + time_from_start=duration_from_seconds(1.0), + ), ], ) tick_loop = TickLoop( @@ -1332,17 +1418,18 @@ def test_full_trajectory_execution(self, mock_adapter, wait_until): ) trajectory = JointTrajectory( + header=header_now(), joint_names=[f"arm/joint{i + 1}" for i in range(6)], points=[ - TrajectoryPoint( + JointTrajectoryPoint( positions=[0.0] * 6, velocities=[0.0] * 6, - time_from_start=0.0, + time_from_start=duration_from_seconds(0.0), ), - TrajectoryPoint( + JointTrajectoryPoint( positions=[0.5] * 6, velocities=[0.0] * 6, - time_from_start=0.5, + time_from_start=duration_from_seconds(0.5), ), ], ) diff --git a/dimos/control/test_coordinator_commands.py b/dimos/control/test_coordinator_commands.py index edae7d196f7..ec21eee326b 100644 --- a/dimos/control/test_coordinator_commands.py +++ b/dimos/control/test_coordinator_commands.py @@ -31,6 +31,7 @@ from collections.abc import Iterator from typing import Any +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint import pytest import dimos.control.coordinator as coord_mod @@ -41,8 +42,7 @@ JointCommandOutput, ResourceClaim, ) -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +from dimos.msgs.time import duration_from_seconds, header_now ARM_JOINTS = frozenset({"arm/joint1", "arm/joint2"}) @@ -126,10 +126,11 @@ def record_time(self, t_now: float | None = None) -> float | None: def _trajectory() -> JointTrajectory: return JointTrajectory( + header=header_now(), joint_names=["arm/joint1", "arm/joint2"], points=[ - TrajectoryPoint(time_from_start=0.0, positions=[0.0, 0.0]), - TrajectoryPoint(time_from_start=1.0, positions=[0.1, 0.2]), + JointTrajectoryPoint(time_from_start=duration_from_seconds(0.0), positions=[0.0, 0.0]), + JointTrajectoryPoint(time_from_start=duration_from_seconds(1.0), positions=[0.1, 0.2]), ], ) diff --git a/dimos/control/test_coordinator_routing.py b/dimos/control/test_coordinator_routing.py index fd9eb830559..a2f7e9d6eda 100644 --- a/dimos/control/test_coordinator_routing.py +++ b/dimos/control/test_coordinator_routing.py @@ -28,6 +28,9 @@ import threading from typing import Any +from dimos_generated.geometry_msgs.msg import Pose, PoseStamped, Twist, TwistStamped, Vector3 +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header import pytest from dimos.control._control_test_helpers import RecordingTask @@ -47,10 +50,6 @@ from dimos.control.teleop_coordinator import TeleopControlCoordinator from dimos.core.stream import In from dimos.hardware.drive_trains.registry import twist_base_adapter_registry -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.teleop.webxr.controller_types import Buttons ARM_JOINTS = ["arm/joint1", "arm/joint2"] @@ -255,14 +254,14 @@ def test_dual_arm_ports_isolate_left_from_right(self, make_coordinator): ) coordinator.start() - taps["left_cartesian"].emit(PoseStamped()) + taps["left_cartesian"].emit(PoseStamped(header=Header(frame_id=""), pose=Pose())) left = coordinator.get_task("cartesian_left") right = coordinator.get_task("cartesian_right") assert len(left.cartesian_calls) == 1 assert right.cartesian_calls == [] - taps["right_cartesian"].emit(PoseStamped()) + taps["right_cartesian"].emit(PoseStamped(header=Header(frame_id=""), pose=Pose())) assert len(left.cartesian_calls) == 1 assert len(right.cartesian_calls) == 1 @@ -275,7 +274,7 @@ def test_cartesian_name_match_needs_no_stream_bind(self, make_coordinator): ) coordinator.start() - taps["cartesian_command"].emit(PoseStamped()) + taps["cartesian_command"].emit(PoseStamped(header=Header(frame_id=""), pose=Pose())) calls = coordinator.get_task("cart").cartesian_calls assert len(calls) == 1 @@ -293,12 +292,14 @@ def test_frame_id_is_not_consulted(self, make_coordinator, mocker): coordinator.start() warn = mocker.patch.object(coord_mod.logger, "warning") - taps["cartesian_command"].emit(PoseStamped(frame_id="some_other_task")) - taps["cartesian_command"].emit(PoseStamped(frame_id="")) + taps["cartesian_command"].emit( + PoseStamped(header=Header(frame_id="some_other_task"), pose=Pose()) + ) + taps["cartesian_command"].emit(PoseStamped(header=Header(frame_id=""), pose=Pose())) calls = coordinator.get_task("cart").cartesian_calls assert len(calls) == 2 - assert calls[0][0].frame_id == "some_other_task" + assert calls[0][0].header.frame_id == "some_other_task" assert not warn.called @@ -311,7 +312,7 @@ def test_buttons_reach_teleop_task(self, make_coordinator): ) coordinator.start() - taps["teleop_buttons"].emit(Buttons()) + taps["teleop_buttons"].emit(Buttons().to_message()) assert len(coordinator.get_task("teleop1").buttons_calls) == 1 @@ -341,7 +342,9 @@ def test_base_twist_maps_to_virtual_joint_velocities(self, make_coordinator): ) coordinator.start() - taps["twist_command"].emit(Twist(linear=[1.0, 2.0, 0.0], angular=[0.0, 0.0, 3.0])) + taps["twist_command"].emit( + Twist(linear=Vector3(x=1.0, y=2.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=3.0)) + ) assert coordinator.get_task("basevel")._velocities == [1.0, 2.0, 3.0] @@ -360,7 +363,9 @@ def test_base_twist_both_maps_joints_and_fans_out(self, make_coordinator): coordinator.add_task(capable, task_type="g1_groot_wbc") coordinator.start() - taps["twist_command"].emit(Twist(linear=[1.0, 2.0, 0.0], angular=[0.0, 0.0, 3.0])) + taps["twist_command"].emit( + Twist(linear=Vector3(x=1.0, y=2.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=3.0)) + ) assert coordinator.get_task("basevel")._velocities == [1.0, 2.0, 3.0] assert len(capable.velocity_commands) == 1 @@ -400,7 +405,9 @@ def test_base_suffix_subset_maps_only_declared_joints(self, make_coordinator): ) coordinator.start() - taps["twist_command"].emit(Twist(linear=[1.0, 2.0, 0.0], angular=[0.0, 0.0, 3.0])) + taps["twist_command"].emit( + Twist(linear=Vector3(x=1.0, y=2.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=3.0)) + ) assert coordinator.get_task("basevel")._velocities == [1.0, 3.0] @@ -420,7 +427,9 @@ def test_base_full_6d_twist_maps_all_axes(self, make_coordinator): ) coordinator.start() - taps["twist_command"].emit(Twist(linear=[1.0, 2.0, 3.0], angular=[4.0, 5.0, 6.0])) + taps["twist_command"].emit( + Twist(linear=Vector3(x=1.0, y=2.0, z=3.0), angular=Vector3(x=4.0, y=5.0, z=6.0)) + ) assert coordinator.get_task("basevel")._velocities == [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] @@ -447,7 +456,9 @@ def test_twist_mapping_ignores_non_base_hardware(self, make_coordinator): ) coordinator.start() - taps["twist_command"].emit(Twist(linear=[1.0, 2.0, 0.0], angular=[0.0, 0.0, 3.0])) + taps["twist_command"].emit( + Twist(linear=Vector3(x=1.0, y=2.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=3.0)) + ) assert coordinator.get_task("basevel")._velocities == [1.0, 2.0, 3.0] assert coordinator.get_task("armvel")._velocities is None @@ -459,7 +470,9 @@ def test_twist_reaches_declaring_velocity_task_without_base(self, make_coordinat coordinator.start() assert taps["twist_command"].subscribed - taps["twist_command"].emit(Twist(linear=[1.0, 2.0, 0.0], angular=[0.0, 0.0, 3.0])) + taps["twist_command"].emit( + Twist(linear=Vector3(x=1.0, y=2.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=3.0)) + ) assert len(task.velocity_commands) == 1 vx, vy, wz, t_now = task.velocity_commands[0] @@ -476,7 +489,7 @@ def test_card_routed_twist_delivers_raw_msg_to_on_twist_command(self, make_coord coordinator.add_task(task, task_type="g1_groot_wbc") coordinator.start() - msg = Twist(linear=[1.0, 2.0, 0.0], angular=[0.0, 0.0, 3.0]) + msg = Twist(linear=Vector3(x=1.0, y=2.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=3.0)) taps["twist_command"].emit(msg) assert len(task.twist_msgs) == 1 @@ -490,7 +503,9 @@ def test_bare_set_velocity_command_stub_gets_nothing(self, make_coordinator): coordinator.start() assert taps["twist_command"].subscribed # the base keeps the stream alive - taps["twist_command"].emit(Twist(linear=[1.0, 2.0, 0.0], angular=[0.0, 0.0, 3.0])) + taps["twist_command"].emit( + Twist(linear=Vector3(x=1.0, y=2.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=3.0)) + ) assert bare.velocity_commands == [] @@ -515,7 +530,7 @@ def test_twist_delivery_concurrent_with_remove_hardware(self, make_coordinator): coordinator.start() assert taps["twist_command"].subscribed - msg = Twist(linear=[1.0, 0.0, 0.0], angular=[0.0, 0.0, 0.5]) + msg = Twist(linear=Vector3(x=1.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.5)) stop = threading.Event() def pump() -> None: @@ -653,7 +668,7 @@ def test_buttons_skip_card_less_tasks(self, make_coordinator): coordinator.add_task(cardless) coordinator.start() - taps["teleop_buttons"].emit(Buttons()) + taps["teleop_buttons"].emit(Buttons().to_message()) assert len(coordinator.get_task("teleop1").buttons_calls) == 1 assert cardless.buttons_calls == [] diff --git a/dimos/control/test_path_following_coordinator.py b/dimos/control/test_path_following_coordinator.py index 34962572cb8..a2bf5112858 100644 --- a/dimos/control/test_path_following_coordinator.py +++ b/dimos/control/test_path_following_coordinator.py @@ -19,6 +19,7 @@ from typing import Any +from dimos_generated.std_msgs.msg import Float32 import pytest from dimos.control.benchmarking.paths import straight_rotate @@ -27,7 +28,6 @@ from dimos.control.path_following_coordinator import PathFollowingCoordinator from dimos.control.tasks.registry import control_task_registry from dimos.control.tick_loop import TickLoop -from dimos.msgs.std_msgs.Float32 import Float32 JOINTS = make_twist_base_joints("go2") diff --git a/dimos/control/test_per_robot_joint_states.py b/dimos/control/test_per_robot_joint_states.py index bfec8d9f6f7..ca68cc8de15 100644 --- a/dimos/control/test_per_robot_joint_states.py +++ b/dimos/control/test_per_robot_joint_states.py @@ -21,6 +21,7 @@ from typing import Any from unittest.mock import MagicMock +from dimos_generated.sensor_msgs.msg import JointState import pytest from dimos.control.components import HardwareComponent, HardwareType, make_joints @@ -31,7 +32,6 @@ from dimos.core.stream import In, Out from dimos.hardware.manipulators.registry import adapter_registry from dimos.hardware.manipulators.spec import ManipulatorAdapter -from dimos.msgs.sensor_msgs.JointState import JointState LEFT_JOINTS = make_joints("left_arm", 2) RIGHT_JOINTS = make_joints("right_arm", 3) @@ -188,12 +188,12 @@ def test_each_port_gets_only_its_own_robot(self, make_coordinator, wait_until): left_msg = left.latest() assert list(left_msg.name) == LEFT_JOINTS - assert left_msg.frame_id == "left_arm" + assert left_msg.header.frame_id == "left_arm" assert list(left_msg.position) == LEFT_POSITIONS right_msg = right.latest() assert list(right_msg.name) == RIGHT_JOINTS - assert right_msg.frame_id == "right_arm" + assert right_msg.header.frame_id == "right_arm" assert list(right_msg.position) == RIGHT_POSITIONS assert set(merged.latest().name) == {*LEFT_JOINTS, *RIGHT_JOINTS} @@ -255,16 +255,16 @@ def test_message_carries_only_that_robots_values(self): assert list(left_msg.position) == [0.1, 0.2] assert list(left_msg.velocity) == [1.1, 1.2] assert list(left_msg.effort) == [2.1, 2.2] - assert left_msg.frame_id == "left_arm" + assert left_msg.header.frame_id == "left_arm" right_msg = by_id["right_arm"] assert list(right_msg.position) == [0.3, 0.4, 0.5] assert list(right_msg.velocity) == [1.3, 1.4, 1.5] assert list(right_msg.effort) == [2.3, 2.4, 2.5] - assert right_msg.frame_id == "right_arm" + assert right_msg.header.frame_id == "right_arm" - assert left_msg.ts == merged[0].ts - assert right_msg.ts == merged[0].ts + assert left_msg.header.stamp == merged[0].header.stamp + assert right_msg.header.stamp == merged[0].header.stamp def test_hardware_is_read_once_per_tick(self): published: list[tuple[str, JointState]] = [] diff --git a/dimos/core/coordination/module_coordinator.py b/dimos/core/coordination/module_coordinator.py index ca0d4f86390..f7d1624a687 100644 --- a/dimos/core/coordination/module_coordinator.py +++ b/dimos/core/coordination/module_coordinator.py @@ -787,7 +787,7 @@ def _coerce_transport_to_backend(transport: Transport[Any]) -> Transport[Any]: if is_pickled: raw, msg_type = transport.topic, None else: - raw, msg_type = transport.topic.topic, transport.topic.lcm_type + raw, msg_type = transport.topic.topic, transport.topic.msg_type # Strip the Zenoh 'dimos/' namespace (if present) back to the logical name. # The factory re-applies the right prefix for the target backend. logical = raw[len("dimos/") :] if raw.startswith("dimos/") else raw diff --git a/dimos/core/coordination/test_module_coordinator.py b/dimos/core/coordination/test_module_coordinator.py index a570ad7d1f2..d7b374c1426 100644 --- a/dimos/core/coordination/test_module_coordinator.py +++ b/dimos/core/coordination/test_module_coordinator.py @@ -18,6 +18,10 @@ from types import MappingProxyType from typing import Any, Protocol +from dimos_generated.geometry_msgs.msg import TransformStamped +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage from pydantic import BaseModel import pytest @@ -48,9 +52,6 @@ from dimos.core.stream import IO, In, Out, Stream from dimos.core.transport import CloudflareTransport, PubSubTransport from dimos.core.transport_factory import transport_topic -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.tf2_msgs.TFMessage import TFMessage import dimos.robot.get_all_blueprints as resolver from dimos.spec.utils import Spec @@ -1163,7 +1164,11 @@ class IoTfPublisher(Module): @rpc def send(self, child: str) -> None: - self.tf.publish(TFMessage(Transform(frame_id="world", child_frame_id=child))) + self.tf.publish( + TFMessage( + transforms=[TransformStamped(header=Header(frame_id="world"), child_frame_id=child)] + ) + ) class IoTfEcho(Module): @@ -1181,7 +1186,11 @@ async def handle_tf(self, msg: TFMessage) -> None: @rpc def send(self, child: str) -> None: - self.tf.publish(TFMessage(Transform(frame_id="world", child_frame_id=child))) + self.tf.publish( + TFMessage( + transforms=[TransformStamped(header=Header(frame_id="world"), child_frame_id=child)] + ) + ) @rpc def seen(self) -> list[str]: diff --git a/dimos/core/coordination/test_worker.py b/dimos/core/coordination/test_worker.py index 73167e2b46b..34db8548d83 100644 --- a/dimos/core/coordination/test_worker.py +++ b/dimos/core/coordination/test_worker.py @@ -16,6 +16,7 @@ import time from typing import TYPE_CHECKING +from dimos_generated.geometry_msgs.msg import Vector3 import pytest from dimos.agents.annotation import skill @@ -24,7 +25,6 @@ from dimos.core.global_config import GlobalConfig, global_config from dimos.core.module import Module from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.Vector3 import Vector3 if TYPE_CHECKING: from dimos.core.resource_monitor.stats import WorkerStats diff --git a/dimos/core/test_baked_host.py b/dimos/core/test_baked_host.py index 119cb1b4045..ec8eb8379ae 100644 --- a/dimos/core/test_baked_host.py +++ b/dimos/core/test_baked_host.py @@ -20,13 +20,13 @@ import pickle from typing import get_args, get_origin, get_type_hints +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import PointCloud2 import pytest from dimos.core.baked_host import BakedHost, baked_host from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.stream import IO, In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 class MapperConfig(NativeModuleConfig): diff --git a/dimos/core/test_cdr_transport.py b/dimos/core/test_cdr_transport.py new file mode 100644 index 00000000000..ed175046001 --- /dev/null +++ b/dimos/core/test_cdr_transport.py @@ -0,0 +1,104 @@ +# 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. + +"""Generated CDR crosses actual transport implementations, including large images.""" + +from contextlib import ExitStack +import threading +import uuid + +from dimos_generated.geometry_msgs.msg import Point +from dimos_generated.sensor_msgs.msg import Image +import numpy as np +import pytest + +from dimos.core.transport import LCMTransport, SHMTransport, ZenohTransport +from dimos.protocol.service.zenohservice import ZenohSessionPool + + +@pytest.fixture(params=["lcm", "zenoh", "shm"]) +def transport_pair(request, lcm_url): + with ExitStack() as stack: + if request.param == "zenoh": + pool = ZenohSessionPool() + stack.callback(pool.close_all) + + def make(message_type): + topic = f"dimos/cdr/{uuid.uuid4().hex[:8]}" + transports = [] + for _ in range(2): + if request.param == "lcm": + transport = LCMTransport(topic, message_type, url=lcm_url) + elif request.param == "zenoh": + transport = ZenohTransport( + topic, + message_type, + session_pool=pool, + scouting=False, + multicast=False, + gossip=False, + connect=[], + ) + else: + transport = SHMTransport(topic, message_type, prefer="cpu") + stack.callback(transport.stop) + transport.start() + transports.append(transport) + return transports + + yield make + + +def test_generated_message_round_trip(transport_pair, retry_until): + publisher, subscriber = transport_pair(Point) + received = [] + ready = threading.Event() + + def on_message(message): + received.append((message.x, message.y, message.z)) + ready.set() + + subscriber.subscribe(on_message) + retry_until(ready, lambda: publisher.broadcast(None, Point(x=1.25, y=-2.5, z=3))) + + assert received[0] == (1.25, -2.5, 3.0) + + +def test_large_generated_image_preserves_pixels_and_header(transport_pair, retry_until): + publisher, subscriber = transport_pair(Image) + pixels = np.arange(640 * 480 * 3, dtype=np.uint8) + message = Image(width=640, height=480, encoding="rgb8", step=1920, data=pixels) + message.header.frame_id = "camera" + message.header.stamp.sec = 1_700_000_000 + message.header.stamp.nanosec = 123_456_789 + received = [] + ready = threading.Event() + + def on_message(value): + received.append(value) + ready.set() + + subscriber.subscribe(on_message) + retry_until(ready, lambda: publisher.broadcast(None, message), timeout=5) + + decoded = received[0] + assert (decoded.width, decoded.height, decoded.step, decoded.encoding) == ( + 640, + 480, + 1920, + "rgb8", + ) + assert decoded.header.frame_id == "camera" + assert (decoded.header.stamp.sec, decoded.header.stamp.nanosec) == (1_700_000_000, 123_456_789) + np.testing.assert_array_equal(decoded.data.view(), pixels) diff --git a/dimos/core/test_core.py b/dimos/core/test_core.py index 139d1b7ec91..048d0a668c9 100644 --- a/dimos/core/test_core.py +++ b/dimos/core/test_core.py @@ -14,6 +14,8 @@ import time +from dimos_generated.geometry_msgs.msg import PoseStamped, Vector3 +from dimos_generated.sensor_msgs.msg import PointCloud2 import pytest from reactivex.disposable import Disposable @@ -22,16 +24,14 @@ from dimos.core.stream import In, Out from dimos.core.testing import MockRobotClient from dimos.core.transport import LCMTransport, pLCMTransport -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.robot.unitree.type.odometry import Odometry +from dimos.types.timestamped import TimestampedData class Navigation(Module): mov: Out[Vector3] lidar: In[PointCloud2] target_position: In[Vector3] - odometry: In[Odometry] + odometry: In[TimestampedData[PoseStamped]] odom_msg_count = 0 lidar_msg_count = 0 @@ -43,8 +43,9 @@ def navigate_to(self, target: Vector3) -> bool: ... def start(self) -> None: def _odom(msg) -> None: self.odom_msg_count += 1 - print("RCV:", (time.perf_counter() - msg.pubtime) * 1000, msg) - self.mov.publish(msg.position) + print("RCV:", (time.perf_counter() - msg.ts) * 1000, msg) + position = msg.value.pose.position + self.mov.publish(Vector3(x=position.x, y=position.y, z=position.z)) unsub = self.odometry.subscribe(_odom) self.register_disposable(Disposable(unsub)) @@ -111,11 +112,11 @@ def test_basic_deployment(dimos) -> None: robot.lidar.transport = LCMTransport("/lidar", PointCloud2) # odometry & mov using just a pickle over LCM - robot.odometry.transport = pLCMTransport("/odom") + robot.timed_odometry.transport = pLCMTransport("/odom") nav.mov.transport = pLCMTransport("/mov") nav.lidar.connect(robot.lidar) - nav.odometry.connect(robot.odometry) + nav.odometry.connect(robot.timed_odometry) robot.mov.connect(nav.mov) robot.start() diff --git a/dimos/core/test_stream.py b/dimos/core/test_stream.py index 9fab86480b8..062aef1b595 100644 --- a/dimos/core/test_stream.py +++ b/dimos/core/test_stream.py @@ -17,6 +17,8 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import PointCloud2 import pytest from dimos.core.core import rpc @@ -24,13 +26,11 @@ from dimos.core.stream import In from dimos.core.testing import MockRobotClient from dimos.core.transport import LCMTransport, pLCMTransport -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.robot.unitree.type.odometry import Odometry class SubscriberBase(Module): - sub1_msgs: list[Odometry] - sub2_msgs: list[Odometry] + sub1_msgs: list[PoseStamped] + sub2_msgs: list[PoseStamped] def __init__(self, **kwargs: Any) -> None: self.sub1_msgs = [] @@ -75,7 +75,7 @@ def sub2_msgs_len(self) -> int: class ClassicSubscriber(SubscriberBase): - odom: In[Odometry] + odom: In[PoseStamped] unsub: Callable[[], None] | None = None unsub2: Callable[[], None] | None = None @@ -98,7 +98,7 @@ def unsub_all(self) -> None: class RXPYSubscriber(SubscriberBase): - odom: In[Odometry] + odom: In[PoseStamped] unsub: Callable[[], None] | None = None unsub2: Callable[[], None] | None = None @@ -142,7 +142,7 @@ class SpyLCMTransport(LCMTransport): active_subscribers: int = 0 def __reduce__(self): - return (SpyLCMTransport, (self.topic.topic, self.topic.lcm_type)) + return (SpyLCMTransport, (self.topic.topic, self.topic.msg_type)) def __init__(self, topic: str, type: type, **kwargs) -> None: super().__init__(topic, type, **kwargs) @@ -174,7 +174,7 @@ def test_subscription(dimos, subscriber_class) -> None: robot = dimos.deploy(MockRobotClient) robot.lidar.transport = SpyLCMTransport("/lidar", PointCloud2) - robot.odometry.transport = SpyLCMTransport("/odom", Odometry) + robot.odometry.transport = SpyLCMTransport("/odom", PoseStamped) robot.mov.transport = pLCMTransport("/mov") subscriber = dimos.deploy(subscriber_class) @@ -215,7 +215,7 @@ def test_get_next(dimos) -> None: robot = dimos.deploy(MockRobotClient) robot.lidar.transport = SpyLCMTransport("/lidar", PointCloud2) - robot.odometry.transport = SpyLCMTransport("/odom", Odometry) + robot.odometry.transport = SpyLCMTransport("/odom", PoseStamped) robot.mov.transport = pLCMTransport("/mov") subscriber = dimos.deploy(RXPYSubscriber) @@ -226,14 +226,14 @@ def test_get_next(dimos) -> None: odom = subscriber.get_next() - assert isinstance(odom, Odometry) + assert isinstance(odom, PoseStamped) assert subscriber.active_subscribers() == 0 time.sleep(0.2) next_odom = subscriber.get_next() - assert isinstance(next_odom, Odometry) + assert isinstance(next_odom, PoseStamped) assert subscriber.active_subscribers() == 0 assert next_odom != odom @@ -247,7 +247,7 @@ def test_hot_getter(dimos) -> None: robot = dimos.deploy(MockRobotClient) robot.lidar.transport = SpyLCMTransport("/lidar", PointCloud2) - robot.odometry.transport = SpyLCMTransport("/odom", Odometry) + robot.odometry.transport = SpyLCMTransport("/odom", PoseStamped) robot.mov.transport = pLCMTransport("/mov") subscriber = dimos.deploy(RXPYSubscriber) @@ -261,7 +261,7 @@ def test_hot_getter(dimos) -> None: odom = subscriber.get_hot() subscriber.stop_hot_getter() - assert isinstance(odom, Odometry) + assert isinstance(odom, PoseStamped) time.sleep(0.3) # there are no subs @@ -272,7 +272,7 @@ def test_hot_getter(dimos) -> None: time.sleep(0.3) next_odom = subscriber.get_hot() - assert isinstance(next_odom, Odometry) + assert isinstance(next_odom, PoseStamped) assert next_odom != odom subscriber.stop_hot_getter() diff --git a/dimos/core/test_zenoh_transport.py b/dimos/core/test_zenoh_transport.py index 31c1532292c..c47fe6e772f 100644 --- a/dimos/core/test_zenoh_transport.py +++ b/dimos/core/test_zenoh_transport.py @@ -19,6 +19,7 @@ from types import SimpleNamespace from typing import cast +from dimos_generated.sensor_msgs.msg import Image import numpy as np from pydantic import ValidationError import pytest @@ -39,7 +40,6 @@ pLCMTransport, pZenohTransport, ) -from dimos.msgs.sensor_msgs.Image import Image from dimos.protocol.pubsub.impl.zenohpubsub import ( QOS_LATEST_WINS, QOS_NEVER_DROP, @@ -49,15 +49,22 @@ class TypedMsg: - """A fake typed message with lcm_encode for testing.""" + """A generated-message codec contract for transport selection tests.""" + + msg_name = "test_msgs/msg/TypedMsg" + schema = "" @staticmethod - def lcm_encode() -> bytes: + def encode() -> bytes: return b"" + @staticmethod + def decode(data: bytes): + return TypedMsg() + class UntypedMsg: - """A message without lcm_encode. Triggers pickle transport.""" + """A Python-only message without a generated codec contract.""" pass @@ -189,7 +196,7 @@ def test_zenoh_transport_broadcast_and_subscribe(retry_until, session_pool, coll t.start() t.subscribe(collector.callback) - test_img = Image(np.zeros((2, 2, 3), dtype=np.uint8)) + test_img = Image(width=2, height=2, step=6, encoding="rgb8", data=np.zeros(12, dtype=np.uint8)) retry_until(collector.event, lambda: t.broadcast(None, test_img)) assert isinstance(collector.received[0], Image) t.stop() @@ -228,7 +235,7 @@ def test_zenoh_transport_pickle_preserves_topic_qos() -> None: t = ZenohTransport(ZenohTopic("dimos/camera/color", Image, qos=QOS_LATEST_WINS)) t2 = pickle.loads(pickle.dumps(t)) assert type(t2) is ZenohTransport - assert t2.topic == t.topic # topic, lcm_type and qos all round-trip + assert t2.topic == t.topic # topic, msg_type and qos all round-trip def test_pzenoh_transport_pickle_preserves_topic_qos() -> None: diff --git a/dimos/core/testing.py b/dimos/core/testing.py index ebf4aa93430..3fc4a9d34f5 100644 --- a/dimos/core/testing.py +++ b/dimos/core/testing.py @@ -16,19 +16,33 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import PoseStamped, Vector3 +from dimos_generated.sensor_msgs.msg import PointCloud2 + from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.robot.unitree.type.lidar import pointcloud2_from_webrtc_lidar -from dimos.robot.unitree.type.odometry import Odometry +from dimos.msgs.time import to_seconds +from dimos.robot.unitree.type.lidar import RawLidarMsg, pointcloud2_from_webrtc_lidar +from dimos.robot.unitree.type.odometry import RawOdometryMessage, pose_from_webrtc_odometry +from dimos.types.timestamped import TimestampedData from dimos.utils.testing.replay import SensorReplay +def _replay_lidar(raw: RawLidarMsg) -> TimestampedData[PointCloud2]: + cloud = pointcloud2_from_webrtc_lidar(raw) + return TimestampedData(cloud, to_seconds(cloud.header.stamp)) + + +def _replay_odometry(raw: RawOdometryMessage) -> TimestampedData[PoseStamped]: + pose = pose_from_webrtc_odometry(raw) + return TimestampedData(pose, to_seconds(pose.header.stamp)) + + class MockRobotClient(Module): - odometry: Out[Odometry] + odometry: Out[PoseStamped] + timed_odometry: Out[TimestampedData[PoseStamped]] lidar: Out[PointCloud2] mov: In[Vector3] @@ -59,8 +73,8 @@ def stop(self) -> None: super().stop() def odomloop(self) -> None: - odomdata = SensorReplay("raw_odometry_rotate_walk", autocast=Odometry.from_msg) - lidardata = SensorReplay("office_lidar", autocast=pointcloud2_from_webrtc_lidar) + odomdata = SensorReplay("raw_odometry_rotate_walk", autocast=_replay_odometry) + lidardata = SensorReplay("office_lidar", autocast=_replay_lidar) lidariter = lidardata.iterate() self._stop_event.clear() @@ -69,9 +83,11 @@ def odomloop(self) -> None: if self._stop_event.is_set(): return print(odom) - odom.pubtime = time.perf_counter() - self.odometry.publish(odom) + # Benchmark publication time belongs to the Python-object stream, + # keeping the generated pose's exact source header untouched. + self.timed_odometry.publish(TimestampedData(odom.value, time.perf_counter())) + self.odometry.publish(odom.value) lidarmsg = next(lidariter) - self.lidar.publish(lidarmsg) + self.lidar.publish(lidarmsg.value) time.sleep(0.1) diff --git a/dimos/core/transport_factory.py b/dimos/core/transport_factory.py index b86b271682a..074df0a92ae 100644 --- a/dimos/core/transport_factory.py +++ b/dimos/core/transport_factory.py @@ -27,6 +27,7 @@ pLCMTransport, pZenohTransport, ) +from dimos.msgs.protocol import DimosMsg from dimos.protocol.pubsub.impl.zenohpubsub import ( QOS_LATEST_WINS, QOS_NEVER_DROP, @@ -58,7 +59,7 @@ def transport_topic(name: str, g: GlobalConfig = global_config) -> str: # High-rate sensor streams: drop stale frames under congestion, never stall the # publisher. Matched by message type since that is what makes them high-rate. -_LATEST_WINS_TYPES = ("sensor_msgs.Image", "sensor_msgs.PointCloud2") +_LATEST_WINS_TYPES = ("sensor_msgs/msg/Image", "sensor_msgs/msg/PointCloud2") # Low-rate channels where a drop loses something that never comes back: a whole # turn of agent/human conversation, a one-shot robot action verb, a push-to-talk # chunk where one gap discards the whole utterance, or a region of a seeded map @@ -100,7 +101,7 @@ def make_transport( """Construct the active-backend pub/sub transport for a logical channel. A pickled (self-describing) transport is used when no `msg_type` is given or - the type has no `lcm_encode`. Otherwise a typed transport is used. + the type does not implement the generated message contract. Otherwise CDR is used. A channel name alone doesn't fully define a topic: backends have per-topic settings (Zenoh publisher QoS, for one) that live on their Topic objects. @@ -110,7 +111,7 @@ def make_transport( transport map. LCM (UDP multicast) has no per-topic settings. """ - use_pickled = msg_type is None or getattr(msg_type, "lcm_encode", None) is None + use_pickled = msg_type is None or not isinstance(msg_type, DimosMsg) topic = transport_topic(name, g) if g.transport == "zenoh": ztopic = ZenohTopic( diff --git a/dimos/e2e_tests/cdr_replay_fixture.py b/dimos/e2e_tests/cdr_replay_fixture.py new file mode 100644 index 00000000000..cb953a948cd --- /dev/null +++ b/dimos/e2e_tests/cdr_replay_fixture.py @@ -0,0 +1,62 @@ +# 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. + +"""Deterministic generated-message replay for browser integration tests.""" + +import math +from pathlib import Path + +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header +import numpy as np + +from dimos.memory.codecs.cdr import CdrCodec +from dimos.memory.codecs.lz4 import Lz4Codec +from dimos.memory.store.sqlite import SqliteStore +from dimos.msgs.image import image_from_array +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import time_from_nanoseconds + + +def write_go2_cdr_replay(path: Path, *, duration_s: int = 180) -> None: + """Record varying camera pixels, world poses and lidar as typed CDR streams.""" + base_ns = 1_700_000_000_123_456_789 + points = np.array([[1, 0, 0], [1, 1, 0], [2, 0, 0]], dtype=np.float32) + with SqliteStore(path=str(path)) as store: + video = store.stream("color_image", Image, codec=Lz4Codec(CdrCodec(Image))) + odom = store.stream("odom", PoseStamped) + lidar = store.stream("lidar", PointCloud2) + for index in range(duration_s * 4): + stamp_ns = base_ns + index * 250_000_000 + stamp = time_from_nanoseconds(stamp_ns) + angle = index / 40 + header = Header(stamp=stamp, frame_id="world") + pose = PoseStamped( + header=header, + pose=Pose( + position=Point(x=math.cos(angle), y=math.sin(angle)), + orientation=Quaternion(z=math.sin(angle / 2), w=math.cos(angle / 2)), + ), + ) + pixels = np.full((240, 320, 3), index % 256, dtype=np.uint8) + pixels[:, (index * 7) % 320] = (255, 0, 127) + image = image_from_array( + pixels, encoding="rgb8", header=Header(stamp=stamp, frame_id="camera_optical") + ) + cloud = pointcloud_from_xyz(points, header=Header(stamp=stamp, frame_id="base_link")) + receipt = stamp_ns / 1e9 + video.append(image, ts=receipt) + odom.append(pose, ts=receipt) + lidar.append(cloud, ts=receipt) diff --git a/dimos/e2e_tests/test_cdr_spy.py b/dimos/e2e_tests/test_cdr_spy.py new file mode 100644 index 00000000000..4c72e0b163d --- /dev/null +++ b/dimos/e2e_tests/test_cdr_spy.py @@ -0,0 +1,51 @@ +# 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. + +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.std_msgs.msg import Header + +from dimos.core.global_config import global_config +from dimos.e2e_tests.lcm_spy import LcmSpy + + +def test_spy_publishes_generated_cdr_on_qualified_channel(mocker, monkeypatch): + monkeypatch.setattr(global_config, "transport", "lcm") + bus = mocker.patch("dimos.e2e_tests.lcm_spy.LCMPubSubBase").return_value + spy = LcmSpy() + message = PoseStamped(header=Header(frame_id="map"), pose=Pose(position=Point(x=2, y=3))) + spy.publish("/goal_request#geometry_msgs/msg/PoseStamped", message) + topic, payload = bus.publish.call_args.args + assert str(topic) == "/goal_request#geometry_msgs/msg/PoseStamped" + decoded = PoseStamped.decode(payload) + assert (decoded.pose.position.x, decoded.pose.position.y, decoded.header.frame_id) == ( + 2, + 3, + "map", + ) + + +def test_spy_odometry_predicate_uses_nested_pose(mocker, monkeypatch): + monkeypatch.setattr(global_config, "transport", "lcm") + mocker.patch("dimos.e2e_tests.lcm_spy.LCMPubSubBase") + spy = LcmSpy() + wait = mocker.patch.object(spy, "wait_for_message_result") + spy.wait_until_odom_position(2, 3, threshold=0.5) + topic, message_type, predicate, _, timeout = wait.call_args.args + assert (topic, message_type, timeout) == ( + "/odom#geometry_msgs/msg/PoseStamped", + PoseStamped, + 60, + ) + assert predicate(PoseStamped(pose=Pose(position=Point(x=2, y=3)))) + assert not predicate(PoseStamped(pose=Pose(position=Point(x=3, y=3)))) diff --git a/dimos/e2e_tests/test_cockpit_browser.py b/dimos/e2e_tests/test_cockpit_browser.py index 14e0414af6a..f92de13ef7b 100644 --- a/dimos/e2e_tests/test_cockpit_browser.py +++ b/dimos/e2e_tests/test_cockpit_browser.py @@ -35,6 +35,7 @@ """ from collections.abc import Callable, Iterator +from pathlib import Path import pytest @@ -57,7 +58,7 @@ @pytest.fixture -def start_go2_replay() -> Iterator[Callable[[], DimosCliCall]]: +def start_go2_replay(cdr_go2_replay_db: Path) -> Iterator[Callable[[], DimosCliCall]]: calls: list[DimosCliCall] = [] def start() -> DimosCliCall: @@ -65,7 +66,14 @@ def start() -> DimosCliCall: call.simulator = None # --viewer none: no rerun viewer spawn (headless CI); the replay # connection means no MuJoCo and no hardware. - call.global_args = ["--robot-ip", "fake", "--viewer", "none"] + call.global_args = [ + "--robot-ip", + "fake", + "--viewer", + "none", + "--replay-db", + str(cdr_go2_replay_db), + ] call.extra_env = {"RELAYBRIDGEMODULE__OPEN_BROWSER": "false"} call.demo_args = ["run", "unitree-go2-basic", "--local-relay"] call.start() diff --git a/dimos/e2e_tests/test_custom_channel_browser.py b/dimos/e2e_tests/test_custom_channel_browser.py index 029726696ca..2ff83eb292a 100644 --- a/dimos/e2e_tests/test_custom_channel_browser.py +++ b/dimos/e2e_tests/test_custom_channel_browser.py @@ -31,10 +31,12 @@ from collections.abc import Callable import struct +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Header import pytest -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.time import time_from_seconds from dimos.web.cockpit import Channel, cockpit from dimos.web.codecs import EncodedPayload, web_encoder from dimos.web.relay_bridge.locate import find_web_dir @@ -48,12 +50,20 @@ TOPIC = "/custom_path_e2e/nav_path" PATH_MSG = Path( - ts=42.5, - frame_id="world", + header=Header(stamp=time_from_seconds(42.5), frame_id="world"), poses=[ - PoseStamped(ts=42.5, position=[1.5, -2.5, 0.0], orientation=[0.0, 0.0, 0.0, 1.0]), - PoseStamped(ts=42.5, position=[2.0, 0.25, 0.0], orientation=[0.0, 0.0, 0.0, 1.0]), - PoseStamped(ts=42.5, position=[3.5, 4.0, 0.0], orientation=[0.0, 0.0, 0.0, 1.0]), + PoseStamped( + header=Header(stamp=time_from_seconds(42.5)), + pose=Pose(position=Point(x=1.5, y=-2.5), orientation=Quaternion(w=1)), + ), + PoseStamped( + header=Header(stamp=time_from_seconds(42.5)), + pose=Pose(position=Point(x=2.0, y=0.25), orientation=Quaternion(w=1)), + ), + PoseStamped( + header=Header(stamp=time_from_seconds(42.5)), + pose=Pose(position=Point(x=3.5, y=4.0), orientation=Quaternion(w=1)), + ), ], ) @@ -62,7 +72,7 @@ # pairs, meta.n = point count; web/examples/custom-path decodes it. @web_encoder("path.points.v1") def encode_path_points(msg: Path) -> EncodedPayload: - payload = b"".join(struct.pack(" relay -> schema-driven browser decoder, end to end. - -Marked web_browser: excluded from the default suite (needs the -`browser-tests` dependency group, Playwright browsers, and built web dists). -Locally: -`uv run --group browser-tests pytest -m web_browser dimos/e2e_tests/test_lcm_channel_browser.py`. +"""Generated CDR over LCM browser E2E. + +A typed channel advertises its complete ROS2 schema. The browser decodes the +Python CDR payload from the manifest alone; LCM remains the raw transport. +Requires the browser-tests group, installed browsers and built web dists. """ from collections.abc import Callable from pathlib import Path +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.std_msgs.msg import Header import pytest -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.time import time_from_seconds from dimos.web.cockpit import Channel, cockpit pytest.importorskip("playwright") @@ -44,10 +38,11 @@ TOPIC = "/lcm_channel_e2e/lcm_pose" POSE = PoseStamped( - ts=42.5, frame_id="map", position=[1.5, -2.5, 0.25], orientation=[0.0, 0.0, 0.0, 1.0] + header=Header(stamp=time_from_seconds(42.5), frame_id="map"), + pose=Pose(position=Point(x=1.5, y=-2.5, z=0.25), orientation=Quaternion(w=1)), ) -# A page whose SDK and relay live on another origin, reading the decoded LCM +# A page whose SDK and relay live on another origin, reading the decoded CDR # value's fields straight off the store snapshot. SDK_PAGE = """ @@ -60,7 +55,7 @@ const v = snapshot.slot?.value; if (v === undefined || v === null) return; document.querySelector("#pose").textContent = - `x=${v.pose.position.x} frame=${v.header.frame_id} nsec=${v.header.stamp.nsec}`; + `x=${v.pose.position.x} frame=${v.header.frame_id} nanosec=${v.header.stamp.nanosec}`; }); @@ -83,11 +78,11 @@ def test_cockpit_channel_table_decodes_the_lcm_channel( chromium_page.get_by_test_id("view-channels").click() expect(chromium_page.get_by_test_id("ch-lcm_pose-seq")).not_to_have_text("-", timeout=120_000) value = chromium_page.get_by_test_id("ch-lcm_pose-value") - # The preview is the schema-driven decoder's: LCM field names, wire order. + # The preview is the schema-driven decoder's: generated ROS2 fields, decoded from the advertised schema. expect(value).to_contain_text("position: {x: 1.5, y: -2.5, z: 0.25}") expect(value).to_contain_text('frame_id: "map"') expect(chromium_page.get_by_test_id("ch-lcm_pose-decode-error")).to_have_count(0) - expect(chromium_page.get_by_text("geometry_msgs.PoseStamped.lcm.v1")).to_be_visible() + expect(chromium_page.get_by_text("geometry_msgs/msg/PoseStamped.cdr.v1")).to_be_visible() def test_sdk_page_reads_the_decoded_lcm_fields( @@ -97,5 +92,5 @@ def test_sdk_page_reads_the_decoded_lcm_fields( page.write_text(SDK_PAGE % {"url": cockpit_url}) chromium_page.goto(page.as_uri()) expect(chromium_page.locator("#pose")).to_have_text( - "x=1.5 frame=map nsec=500000000", timeout=120_000 + "x=1.5 frame=map nanosec=500000000", timeout=120_000 ) diff --git a/dimos/e2e_tests/test_manipulation_planning_groups.py b/dimos/e2e_tests/test_manipulation_planning_groups.py index a263c8da752..004e0018575 100644 --- a/dimos/e2e_tests/test_manipulation_planning_groups.py +++ b/dimos/e2e_tests/test_manipulation_planning_groups.py @@ -23,6 +23,7 @@ from collections.abc import Callable import time +from dimos_generated.sensor_msgs.msg import JointState import pytest from dimos.control.coordinator import ControlCoordinator @@ -36,11 +37,10 @@ ManipulationSnapshot, PlanningGroupInfo, ) -from dimos.msgs.sensor_msgs.JointState import JointState pytestmark = [pytest.mark.self_hosted] -JOINT_STATE_TOPIC = "/coordinator_joint_state#sensor_msgs.JointState" +JOINT_STATE_TOPIC = "/coordinator_joint_state#sensor_msgs/msg/JointState" BLUEPRINT = "openarm-planner-coordinator" LEFT_GROUP_ID = "left_arm" RIGHT_GROUP_ID = "right_arm" diff --git a/dimos/e2e_tests/test_map_click_browser.py b/dimos/e2e_tests/test_map_click_browser.py index 2152eb80e9c..0695f848b6f 100644 --- a/dimos/e2e_tests/test_map_click_browser.py +++ b/dimos/e2e_tests/test_map_click_browser.py @@ -18,13 +18,14 @@ import threading from typing import NamedTuple +from dimos_generated.geometry_msgs.msg import Point, PointStamped, Pose, PoseStamped, Quaternion +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Header import pytest from reactivex.disposable import Disposable from dimos.core.transport import pLCMTransport -from dimos.msgs.geometry_msgs.PointStamped import PointStamped -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.time import time_from_seconds from dimos.web.cockpit import Map2D, cockpit from dimos.web.relay_bridge.e2e_support import stop_module from dimos.web.relay_bridge.gen_costmap_fixtures import grid_msg @@ -50,9 +51,13 @@ class MapBridge(NamedTuple): def _path(*xy: tuple[float, float]) -> Path: poses = [ - PoseStamped(ts=1.0, position=[x, y, 0.0], orientation=[0.0, 0.0, 0.0, 1.0]) for x, y in xy + PoseStamped( + header=Header(stamp=time_from_seconds(1.0)), + pose=Pose(position=Point(x=x, y=y), orientation=Quaternion(w=1)), + ) + for x, y in xy ] - return Path(ts=1.0, frame_id="world", poses=poses) + return Path(header=Header(stamp=time_from_seconds(1.0), frame_id="world"), poses=poses) @pytest.fixture(scope="module") @@ -125,9 +130,9 @@ def test_click_to_goal_and_cancel(map_bridge: MapBridge, chromium_page: Page) -> ) chromium_page.click(CANVAS) assert wait_until(lambda: len(clicks) == 1, timeout=15.0) - assert clicks[0].x == pytest.approx(0.75, abs=0.01) - assert clicks[0].y == pytest.approx(0.0, abs=0.01) - assert clicks[0].frame_id == "world" + assert clicks[0].point.x == pytest.approx(0.75, abs=0.01) + assert clicks[0].point.y == pytest.approx(0.0, abs=0.01) + assert clicks[0].header.frame_id == "world" cancel = chromium_page.get_by_test_id("map2d-global_costmap-cancel") expect(cancel).to_have_count(0) diff --git a/dimos/e2e_tests/test_sdk_browser.py b/dimos/e2e_tests/test_sdk_browser.py index 419ac776804..a5095781486 100644 --- a/dimos/e2e_tests/test_sdk_browser.py +++ b/dimos/e2e_tests/test_sdk_browser.py @@ -77,12 +77,19 @@ @pytest.fixture(scope="module") -def go2_replay_serve_dir() -> Iterator[DimosCliCall]: +def go2_replay_serve_dir(cdr_go2_replay_db: Path) -> Iterator[DimosCliCall]: call = DimosCliCall() call.simulator = None # --viewer none: no rerun viewer spawn (headless CI); the replay # connection means no MuJoCo and no hardware. - call.global_args = ["--robot-ip", "fake", "--viewer", "none"] + call.global_args = [ + "--robot-ip", + "fake", + "--viewer", + "none", + "--replay-db", + str(cdr_go2_replay_db), + ] call.extra_env = {"RELAYBRIDGEMODULE__OPEN_BROWSER": "false"} call.demo_args = [ "run", diff --git a/dimos/evals/agents/lib/plain_recording.py b/dimos/evals/agents/lib/plain_recording.py index 4b2abc83d4a..fffd61c3210 100644 --- a/dimos/evals/agents/lib/plain_recording.py +++ b/dimos/evals/agents/lib/plain_recording.py @@ -22,10 +22,11 @@ from pathlib import Path from typing import TYPE_CHECKING, Any +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 import numpy as np -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.image import image_to_rgb +from dimos.msgs.pointcloud import pointcloud_rgb, pointcloud_xyz if TYPE_CHECKING: from dimos.memory.stream import Stream @@ -48,10 +49,14 @@ def plain_recording(streams: Sequence[Stream[Any, Any]], directory: Path) -> Pat record: dict[str, Any] = { "stream": stream.name, "timestamp": obs.ts, - "frame_id": getattr(data, "frame_id", None), + "frame_id": data.header.frame_id + if isinstance(data, (Image, PointCloud2)) + else None, } if isinstance(data, PointCloud2): - points, colors = data.as_numpy() + points = pointcloud_xyz(data) + rgb = pointcloud_rgb(data) + colors = None if rgb is None else rgb.astype(np.float64) / 255.0 values = points if colors is None else np.column_stack((points, colors)) path = directory / f"{stem}.csv" np.savetxt( @@ -65,7 +70,7 @@ def plain_recording(streams: Sequence[Stream[Any, Any]], directory: Path) -> Pat record["format"] = "XYZ in source coordinates; optional RGB in [0,1]" elif isinstance(data, Image): path = directory / f"{stem}.png" - PILImage.fromarray(data.to_rgb().data).save(path) + PILImage.fromarray(image_to_rgb(data)).save(path) record["format"] = "RGB PNG" elif data is None or isinstance(data, (bool, int, float, str, list, dict)): path = directory / f"{stem}.json" diff --git a/dimos/evals/agents/question_answer.py b/dimos/evals/agents/question_answer.py index d0d372fb33b..b2c9d0b895e 100644 --- a/dimos/evals/agents/question_answer.py +++ b/dimos/evals/agents/question_answer.py @@ -12,14 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""One model call over the encoded recording. The only place in -:mod:`dimos.evals` that calls ``agent_encode()`` — the surface under test.""" +"""One model call over generated image content and explicit custom agent encoders.""" from __future__ import annotations +import base64 import json from typing import TYPE_CHECKING, Any +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import Image from pydantic import Field from dimos.evals.agents.lib.single_call import ( @@ -28,15 +30,35 @@ SingleCallAgentConfig, ) from dimos.evals.types import RunningEnvironment +from dimos.msgs.image import image_to_jpeg if TYPE_CHECKING: from dimos.memory.type.observation import Observation def _observation_blocks(obs: Observation[Any], stamp: str) -> Blocks: - """One observation as ``agent_encode()`` renders it. ``str(data)`` where a - type has no encoder: an encoding gap the eval surfaces by design.""" + """Render generated images as JPEG blocks and custom values through their encoders.""" data = obs.data + if isinstance(data, Image): + jpeg = base64.b64encode(image_to_jpeg(data, quality=80)).decode("ascii") + return [ + {"type": "text", "text": stamp}, + {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{jpeg}"}}, + ] + if isinstance(data, PoseStamped): + position = data.pose.position + orientation = data.pose.orientation + return [ + { + "type": "text", + "text": ( + f"{stamp} frame={data.header.frame_id!r} " + f"position=({position.x:.3f}, {position.y:.3f}, {position.z:.3f}) " + f"quaternion=({orientation.x:.3f}, {orientation.y:.3f}, " + f"{orientation.z:.3f}, {orientation.w:.3f})" + ), + } + ] encoded = data.agent_encode() if hasattr(data, "agent_encode") else None if isinstance(encoded, list): # e.g. Image -> image_url blocks return [{"type": "text", "text": stamp}, *encoded] diff --git a/dimos/evals/agents/test_no_dimos.py b/dimos/evals/agents/test_no_dimos.py index c864966fe25..f2371d0b2a1 100644 --- a/dimos/evals/agents/test_no_dimos.py +++ b/dimos/evals/agents/test_no_dimos.py @@ -18,6 +18,8 @@ import os from pathlib import Path +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header import numpy as np from PIL import Image as PILImage import pytest @@ -30,8 +32,9 @@ from dimos.evals.constants import NO_DIMOS_KEYWORDS, RAW_README from dimos.evals.types import RunningEnvironment from dimos.memory.store.sqlite import SqliteStore -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.image import image_from_array, image_to_rgb +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import time_from_seconds def test_no_dimos_defaults_the_excluded_keywords_but_keeps_explicit_ones() -> None: @@ -82,12 +85,19 @@ def test_recorded_observations_are_exported_as_plain_files(tmp_path: Path) -> No points = np.array([[1.125, -2.5, 3.0], [4.0, 5.0, 6.0]]) with SqliteStore(path=tmp_path / "source.db") as store: lidar = store.stream("lidar", PointCloud2) - lidar.append(PointCloud2.from_numpy(points, frame_id="world", timestamp=1.0), ts=1.0) - lidar.append(PointCloud2.from_numpy(points * 2, timestamp=2.0), ts=2.0) + lidar.append( + pointcloud_from_xyz( + points, header=Header(frame_id="world", stamp=time_from_seconds(1.0)) + ), + ts=1.0, + ) + lidar.append( + pointcloud_from_xyz(points * 2, header=Header(stamp=time_from_seconds(2.0))), ts=2.0 + ) images = store.stream("camera", Image) - images.append(Image.from_numpy(pixels, format=ImageFormat.RGB), ts=3.0) + images.append(image_from_array(pixels, encoding="rgb8"), ts=3.0) store.stream("secret", str).append("grader-only") - stored_pixels = images.first().data.to_rgb().data + stored_pixels = image_to_rgb(images.first().data) manifest = plain_recording((lidar.limit(1), images), tmp_path / "input") records = json.loads(manifest.read_text())["observations"] assert [record["stream"] for record in records] == ["lidar", "camera"] diff --git a/dimos/evals/agents/test_pi.py b/dimos/evals/agents/test_pi.py index 763445ac4e1..709bbdbd270 100644 --- a/dimos/evals/agents/test_pi.py +++ b/dimos/evals/agents/test_pi.py @@ -90,7 +90,7 @@ def test_recording_export_contains_only_selected_data(dataset: str, tmp_path: Pa with SqliteStore(path=str(exported), must_exist=True) as copy: assert copy.list_streams() == ["odom"] - assert [(obs.ts, obs.data.position.x) for obs in copy.streams.odom] == [ + assert [(obs.ts, obs.data.pose.position.x) for obs in copy.streams.odom] == [ (1000.0, 0.0), (1001.0, 1.0), ] diff --git a/dimos/evals/conftest.py b/dimos/evals/conftest.py index 22a7446c4bf..8ea5e741760 100644 --- a/dimos/evals/conftest.py +++ b/dimos/evals/conftest.py @@ -14,12 +14,11 @@ from pathlib import Path +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.std_msgs.msg import Header import pytest from dimos.memory.store.sqlite import SqliteStore -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import make_vector3 @pytest.fixture @@ -30,9 +29,8 @@ def dataset(tmp_path: Path) -> str: stream = store.stream("odom", PoseStamped) for i in range(5): pose = PoseStamped( - position=make_vector3(float(i), 0.0, 0.0), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="world", + header=Header(frame_id="world"), + pose=Pose(position=Point(x=float(i)), orientation=Quaternion(w=1.0)), ) stream.append(pose, ts=1000.0 + i) return str(path) diff --git a/dimos/evals/environments/dimsim.py b/dimos/evals/environments/dimsim.py index ab4581a0b40..db8c45ce2a6 100644 --- a/dimos/evals/environments/dimsim.py +++ b/dimos/evals/environments/dimsim.py @@ -23,9 +23,10 @@ from dimos.evals.environments.sim import Sim, SimConfig if TYPE_CHECKING: + from dimos_generated.geometry_msgs.msg import PoseStamped + from dimos.e2e_tests.dimos_cli_call import DimosCliCall from dimos.memory.store.base import Store - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped class DimSimEnvironmentConfig(SimConfig): diff --git a/dimos/evals/environments/sim.py b/dimos/evals/environments/sim.py index 0a4a2887c85..824e545f18e 100644 --- a/dimos/evals/environments/sim.py +++ b/dimos/evals/environments/sim.py @@ -34,9 +34,10 @@ from dimos.evals.types import RunningEnvironment if TYPE_CHECKING: + from dimos_generated.geometry_msgs.msg import PoseStamped + from dimos.evals.agents.base import Agent from dimos.memory.store.base import Store - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped class SimConfig(EnvironmentConfig): @@ -191,7 +192,7 @@ def settle(self, budget_s: float) -> None: pose = self.latest_pose(self._recording) except LookupError: return - position = pose.position + position = pose.pose.position if ( anchor is None or math.hypot(position.x - anchor.x, position.y - anchor.y) > self.config.at_rest_m diff --git a/dimos/evals/environments/test_dimsim.py b/dimos/evals/environments/test_dimsim.py index 3a851b068e4..3ed0818edf7 100644 --- a/dimos/evals/environments/test_dimsim.py +++ b/dimos/evals/environments/test_dimsim.py @@ -14,12 +14,14 @@ from unittest.mock import Mock +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.std_msgs.msg import Header import pytest from dimos.e2e_tests.dimos_cli_call import DimosCliCall from dimos.evals.environments.dimsim import DimSimEnvironment from dimos.memory.store.memory import MemoryStore -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.time import time_from_seconds, to_seconds def test_dimsim_launch_setup_and_pose(mocker): @@ -37,9 +39,9 @@ def test_dimsim_launch_setup_and_pose(mocker): with MemoryStore() as store: with pytest.raises(LookupError): env.latest_pose(store) - pose = PoseStamped(ts=123, frame_id="world") + pose = PoseStamped(header=Header(stamp=time_from_seconds(123), frame_id="world")) store.stream("odom", PoseStamped).append(pose) - assert env.latest_pose(store).ts == 123 + assert to_seconds(env.latest_pose(store).header.stamp) == 123 finally: env.stop() client.stop.assert_called_once() diff --git a/dimos/evals/suites/dimsim_house.py b/dimos/evals/suites/dimsim_house.py index 0a6f8b79566..c043d6da94f 100644 --- a/dimos/evals/suites/dimsim_house.py +++ b/dimos/evals/suites/dimsim_house.py @@ -28,17 +28,19 @@ from __future__ import annotations from collections.abc import Callable +import math from typing import TYPE_CHECKING +from dimos_generated.geometry_msgs.msg import Vector3 + from dimos.evals.environments.dimsim import DimSimEnvironment from dimos.evals.scorers import ramp from dimos.evals.types import EvalCase, Outcome, Suite, recording -from dimos.msgs.geometry_msgs.Vector3 import Vector3 if TYPE_CHECKING: from dimos.e2e_tests.dim_sim_client import DimSimClient -BED = Vector3(-3.567, -1.332, 0.0) +BED = Vector3(x=-3.567, y=-1.332) _HOUSE_TOUR = [ (3.881, 4.803), @@ -69,8 +71,9 @@ def _explore_house(sim: DimSimClient) -> None: import time + from dimos_generated.geometry_msgs.msg import PoseStamped + from dimos.core.transport import LCMTransport - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.simulation.mujoco.direct_cmd_vel_explorer import DirectCmdVelExplorer # dimsim spawns the robot well after MCP is up — wait for odom before driving @@ -98,10 +101,10 @@ def _ended_near(target: Vector3) -> Callable[[Outcome], float]: def grade(o: Outcome) -> float: store = recording(o) try: - p = store.streams.odom.last().data.position + p = store.streams.odom.last().data.pose.position finally: store.stop() - d = Vector3(p.x - target.x, p.y - target.y, 0.0).length() + d = math.hypot(p.x - target.x, p.y - target.y) return ramp(max(0.0, d - 2.0), band=2.0) # full credit inside the e2e 2m threshold return grade diff --git a/dimos/evals/suites/dimsim_pointcloud_mapping.py b/dimos/evals/suites/dimsim_pointcloud_mapping.py index 30522d1828f..8cb30210118 100644 --- a/dimos/evals/suites/dimsim_pointcloud_mapping.py +++ b/dimos/evals/suites/dimsim_pointcloud_mapping.py @@ -61,7 +61,7 @@ def grade(o: Outcome) -> float: count = 0.0 store = recording(o) try: - path = [(p.data.position.x, p.data.position.y) for p in store.streams.odom] + path = [(p.data.pose.position.x, p.data.pose.position.y) for p in store.streams.odom] finally: store.stop() visited = sum( diff --git a/dimos/evals/test_evals.py b/dimos/evals/test_evals.py index 764ee8534a0..251777c1a86 100644 --- a/dimos/evals/test_evals.py +++ b/dimos/evals/test_evals.py @@ -29,9 +29,12 @@ from types import SimpleNamespace from typing import Any +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.std_msgs.msg import Header from langchain_core.language_models.fake_chat_models import FakeListChatModel from langchain_core.messages import AIMessage, HumanMessage, ToolMessage import numpy as np +from PIL import Image as PILImage from pydantic import ValidationError import pytest from pytest_mock import MockerFixture @@ -41,7 +44,7 @@ from dimos.evals.agents.blind import BLIND_BLOCK, Blind from dimos.evals.agents.lib.trajectory_builder import TrajectoryBuilder from dimos.evals.agents.mcp_client_adapter import McpClientAdapter -from dimos.evals.agents.question_answer import QuestionAnswer +from dimos.evals.agents.question_answer import QuestionAnswer, _observation_blocks from dimos.evals.cli import load_agent from dimos.evals.constants import RAW_README from dimos.evals.environments.base import Environment @@ -78,17 +81,12 @@ ) from dimos.memory.store.memory import MemoryStore from dimos.memory.store.sqlite import SqliteStore -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import make_vector3 -from dimos.msgs.sensor_msgs.Image import Image def _pose(x: float, y: float) -> PoseStamped: return PoseStamped( - position=make_vector3(x, y, 0.0), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="world", + header=Header(frame_id="world"), + pose=Pose(position=Point(x=x, y=y), orientation=Quaternion(w=1.0)), ) @@ -199,7 +197,7 @@ def test_dataset_start_hands_out_the_selection(dataset: str) -> None: assert odom.name == "odom" observations = list(odom) assert [o.ts for o in observations] == [1000.0, 1001.0] - assert observations[0].data.position.x == 0.0 + assert observations[0].data.pose.position.x == 0.0 assert running.mcp_url == "" and not env.has_robot assert running.artifacts["recording"] == Path(dataset) finally: @@ -316,14 +314,14 @@ def test_sim_launches_base_blueprints_and_agent_modules_in_order( def test_image_file_environment(tmp_path: Path) -> None: path = tmp_path / "frame.png" - Image.from_numpy(np.full((8, 8, 3), 200, dtype=np.uint8)).save(str(path)) + PILImage.fromarray(np.full((8, 8, 3), 200, dtype=np.uint8)).save(path) env = ImageFile(path) env.preflight(QuestionAnswer()) running = env.start(()) try: (image_stream,) = running.streams (obs,) = list(image_stream) - assert obs.data.agent_encode()[0]["type"] == "image_url" + assert _observation_blocks(obs, "[t=0.0s]")[1]["type"] == "image_url" assert running.artifacts == {"image": path} finally: env.stop() @@ -705,7 +703,7 @@ def test_recording_helper_opens_the_artifact(dataset: str) -> None: with recording( Outcome(trajectory=_trajectory("", Path()), artifacts={"recording": Path(dataset)}) ) as store: - assert store.streams.odom.last().data.position.x == 4.0 + assert store.streams.odom.last().data.pose.position.x == 4.0 def test_count_rooms_grader_scores_reply_and_coverage(tmp_path: Path) -> None: diff --git a/dimos/evals/test_mem2_wiring.py b/dimos/evals/test_mem2_wiring.py index 1b4fe9c8cd9..39cfaf9ac5a 100644 --- a/dimos/evals/test_mem2_wiring.py +++ b/dimos/evals/test_mem2_wiring.py @@ -26,6 +26,9 @@ from pathlib import Path +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header from langchain_core.language_models.fake_chat_models import FakeListChatModel import numpy as np import pytest @@ -45,17 +48,14 @@ recording, ) from dimos.memory.store.sqlite import SqliteStore -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import make_vector3 -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_from_array +from dimos.msgs.time import time_from_seconds def _pose(x: float, y: float) -> PoseStamped: return PoseStamped( - position=make_vector3(x, y, 0.0), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="world", + header=Header(frame_id="world"), + pose=Pose(position=Point(x=x, y=y), orientation=Quaternion(w=1.0)), ) @@ -74,7 +74,14 @@ def test_selected_streams_reach_the_prompt( frame = np.full((16, 16, 3), 200, dtype=np.uint8) images = store.stream("color_image", Image) for i in range(3): - images.append(Image.from_numpy(frame, frame_id="cam", ts=1000.0 + i), ts=1000.0 + i) + images.append( + image_from_array( + frame, + encoding="rgb8", + header=Header(frame_id="cam", stamp=time_from_seconds(1000.0 + i)), + ), + ts=1000.0 + i, + ) case = EvalCase( id="wiring", @@ -99,7 +106,8 @@ def test_selected_streams_reach_the_prompt( # the actual observation data crossed from memory into the prompt: assert len(image_blocks) == 2, "both selected image observations should be encoded" assert image_blocks[0]["image_url"]["url"].startswith("data:image/jpeg;base64,") - assert str(_pose(19.0, 2.5)) in text, "last odom pose must reach the prompt" + assert "position=(19.000, 2.500, 0.000)" in text, "last odom pose must reach the prompt" + assert "frame='world'" in text and "quaternion=(0.000, 0.000, 0.000, 1.000)" in text assert case.inputs in text @@ -132,7 +140,11 @@ def test_frames_per_stream_downsamples_not_truncates( for block in content if block["type"] == "text" and block["text"].startswith("[t=") ] - assert stamped == [f"[t={float(i):.1f}s] {_pose(float(i), 0.0)}" for i in positions] + assert stamped == [ + f"[t={float(i):.1f}s] frame='world' position=({float(i):.3f}, 0.000, 0.000) " + "quaternion=(0.000, 0.000, 0.000, 1.000)" + for i in positions + ] def test_grader_reads_the_history_the_environment_recorded( @@ -152,7 +164,7 @@ def run( def grade(outcome: Outcome) -> float: with recording(outcome) as rec: - poses = [obs.data.position.x for obs in rec.streams.odom] + poses = [obs.data.pose.position.x for obs in rec.streams.odom] assert poses == [5.0, 2.0, 1.0, 0.0] return ramp(abs(poses[-1]), band=2.0) diff --git a/dimos/evals/vqa/author.py b/dimos/evals/vqa/author.py index c114b4edcc1..295db4d4366 100644 --- a/dimos/evals/vqa/author.py +++ b/dimos/evals/vqa/author.py @@ -25,8 +25,9 @@ from dimos.evals.vqa.contracts import FamilySpec, QuestionProposal if TYPE_CHECKING: + from dimos_generated.sensor_msgs.msg import Image + from dimos.models.vl.base import VlModel - from dimos.msgs.sensor_msgs.Image import Image class QuestionAuthor(Protocol): diff --git a/dimos/evals/vqa/contracts.py b/dimos/evals/vqa/contracts.py index a5df7351e25..8dbc93e87d6 100644 --- a/dimos/evals/vqa/contracts.py +++ b/dimos/evals/vqa/contracts.py @@ -29,10 +29,11 @@ ) if TYPE_CHECKING: + from dimos_generated.sensor_msgs.msg import Image + from dimos.evals.vqa.pointcloud_frame import PointCloudFrame from dimos.evals.vqa.primitives.edge_tam import ObjectMaskEvidence from dimos.evals.vqa.primitives.range import ObjectRangeEvidence - from dimos.msgs.sensor_msgs.Image import Image NonEmptyString = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)] FamilyName = Literal[ diff --git a/dimos/evals/vqa/families.py b/dimos/evals/vqa/families.py index f706da2f8d1..6ab8d3c9fc1 100644 --- a/dimos/evals/vqa/families.py +++ b/dimos/evals/vqa/families.py @@ -31,9 +31,10 @@ ) if TYPE_CHECKING: + from dimos_generated.sensor_msgs.msg import Image + from dimos.evals.vqa.pointcloud_frame import PointCloudFrame from dimos.models.vl.base import VlModel - from dimos.msgs.sensor_msgs.Image import Image PRESENCE_FAMILY = FamilySpec( diff --git a/dimos/evals/vqa/generate.py b/dimos/evals/vqa/generate.py index 1c47d8c8004..359152fabbf 100644 --- a/dimos/evals/vqa/generate.py +++ b/dimos/evals/vqa/generate.py @@ -24,6 +24,8 @@ import tempfile from typing import TYPE_CHECKING +from dimos_generated.sensor_msgs.msg import Image +from PIL import Image as PILImage from pydantic import BaseModel, ConfigDict, Field, model_validator from dimos.constants import STATE_DIR @@ -49,7 +51,8 @@ ) from dimos.evals.vqa.primitives.edge_tam import EdgeTAMObjectMaskPipeline from dimos.evals.vqa.primitives.range import LidarRangeEstimator -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_to_rgb +from dimos.msgs.time import to_seconds if TYPE_CHECKING: from dimos.models.vl.base import VlModel @@ -395,8 +398,7 @@ def _write_frame(output: Path, frame: _GeneratedFrame) -> None: image_name = f"assets/frame-{frame.index:06d}.png" image_path = output / image_name - if not frame.image.save(str(image_path)): - raise RuntimeError(f"failed to write VQA image: {image_path}") + PILImage.fromarray(image_to_rgb(frame.image)).save(image_path) case_rows = [case.model_dump(mode="json") for case in frame.cases] label_rows = [label.model_dump(mode="json") for label in frame.labels] @@ -408,7 +410,7 @@ def _write_frame(output: Path, frame: _GeneratedFrame) -> None: { "frame_index": frame.index, "image": image_name, - "timestamp": frame.image.ts, + "timestamp": to_seconds(frame.image.header.stamp), "question_count": len(frame.cases), "rejected_question_count": sum( row.get("status") == "rejected" for row in frame.audit_rows diff --git a/dimos/evals/vqa/pointcloud_frame.py b/dimos/evals/vqa/pointcloud_frame.py index af7b355c4c2..9ed1527838b 100644 --- a/dimos/evals/vqa/pointcloud_frame.py +++ b/dimos/evals/vqa/pointcloud_frame.py @@ -21,15 +21,17 @@ from types import TracebackType from typing import TYPE_CHECKING, Any, cast +from dimos_generated.geometry_msgs.msg import PoseStamped, TransformStamped +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header import numpy as np from dimos.memory.cli.dataset import open_dataset from dimos.memory.store.base import Store from dimos.memory.tf import StreamTF -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.camera_info import camera_info_from_intrinsics +from dimos.msgs.geometry import pose_from_matrix, transform_from_pose, transform_matrix +from dimos.msgs.image import image_from_array, image_view if TYPE_CHECKING: from dimos.memory.stream import Stream @@ -45,7 +47,7 @@ class PointCloudFrame: image: Image pointcloud: PointCloud2 camera_info: CameraInfo - pointcloud_to_camera: Transform + pointcloud_to_camera: TransformStamped image_observation_timestamp: float pointcloud_observation_timestamp: float calibration_source: str @@ -184,8 +186,7 @@ def load(self, frame_index: int) -> PointCloudFrame: image_query = self._images.offset(frame_index).limit(1) image_obs = image_query.first() lidar_obs = _align_one(image_query, self._lidar, self._tolerance_s) - source_image = image_obs.data.copy() - source_image.ts = image_obs.ts + source_image = image_obs.data recorded_camera_info = self._recorded_camera_info recorded_tf = self._recorded_tf @@ -196,19 +197,22 @@ def load(self, frame_index: int) -> PointCloudFrame: pointcloud_to_camera = _recorded_pointcloud_to_camera( image_obs, lidar_obs, - camera_info.frame_id, + camera_info.header.frame_id, recorded_tf, self._tolerance_s, ) rectification = np.eye(4, dtype=np.float64) - rectification[:3, :3] = np.asarray(source_camera_info.R, dtype=np.float64).reshape(3, 3) - pointcloud_to_camera = Transform.from_matrix( - rectification @ pointcloud_to_camera.to_matrix(), - ts=image_obs.ts, - frame_id=camera_info.frame_id, - child_frame_id=lidar_obs.data.frame_id, + rectification[:3, :3] = np.asarray(source_camera_info.r, dtype=np.float64).reshape(3, 3) + pointcloud_to_camera = transform_from_pose( + PoseStamped( + header=Header(stamp=image.header.stamp, frame_id=camera_info.header.frame_id), + pose=pose_from_matrix( + rectification @ transform_matrix(pointcloud_to_camera.transform) + ), + ), + child_frame_id=lidar_obs.data.header.frame_id, ) - if not np.all(np.isfinite(pointcloud_to_camera.to_matrix())): + if not np.all(np.isfinite(transform_matrix(pointcloud_to_camera.transform))): raise PointCloudFrameUnavailableError( "recorded point-cloud-to-camera transform must be finite" ) @@ -231,8 +235,7 @@ def load_image(self, frame_index: int) -> Image: if self._images is None: raise RuntimeError("point-cloud frame loader must be started before loading frames") observation = self._images.offset(frame_index).limit(1).first() - image = observation.data.copy() - image.ts = observation.ts + image = observation.data recorded_camera_info = self._recorded_camera_info if recorded_camera_info is None: raise RuntimeError("camera calibration was not initialized") @@ -262,30 +265,32 @@ def rectify(self, image: Image, source: CameraInfo) -> tuple[Image, CameraInfo]: image.width, image.height, source.distortion_model, - tuple(source.K), - tuple(source.D), - tuple(source.R), - tuple(source.P), + tuple(source.k), + tuple(source.d), + tuple(source.r), + tuple(source.p), ) maps = self._maps.get(key) if maps is None: maps = _rectification_maps(image, source) self._maps[key] = maps map_x, map_y, output_matrix = maps - data = cv2.remap(image.data, map_x, map_y, interpolation=cv2.INTER_LINEAR) - frame_id = source.frame_id or image.frame_id + data = cv2.remap(image_view(image), map_x, map_y, interpolation=cv2.INTER_LINEAR) + frame_id = source.header.frame_id or image.header.frame_id if not frame_id: raise ValueError("camera calibration requires a camera frame_id") - camera_info = CameraInfo.from_intrinsics( + camera_info = camera_info_from_intrinsics( output_matrix[0, 0], output_matrix[1, 1], output_matrix[0, 2], output_matrix[1, 2], image.width, image.height, - frame_id=frame_id, - ).with_ts(image.ts) - return Image(data=data, format=image.format, frame_id=frame_id, ts=image.ts), camera_info + header=Header(stamp=image.header.stamp, frame_id=frame_id), + ) + return image_from_array( + data, encoding=image.encoding, header=camera_info.header + ), camera_info def _align_one( @@ -317,8 +322,8 @@ def _recorded_pointcloud_to_camera( camera_frame: str, tf: TFLookup, tolerance_s: float, -) -> Transform: - cloud_frame = lidar_obs.data.frame_id +) -> TransformStamped: + cloud_frame = lidar_obs.data.header.frame_id if not cloud_frame: raise PointCloudFrameUnavailableError("recorded point cloud requires a frame_id") camera_from_cloud = tf.get( @@ -341,15 +346,15 @@ def _rectification_maps( # OpenCV is intentionally lazy because importing it loads a large native library. import cv2 - intrinsics = np.asarray(source.K, dtype=np.float64) + intrinsics = np.asarray(source.k, dtype=np.float64) if intrinsics.size != 9: raise ValueError("camera intrinsics must contain nine values") matrix = intrinsics.reshape(3, 3) if not np.all(np.isfinite(matrix)) or matrix[0, 0] <= 0 or matrix[1, 1] <= 0: raise ValueError("camera intrinsics must be finite with positive focal lengths") - distortion = np.asarray(source.D, dtype=np.float64) - rectification = np.asarray(source.R, dtype=np.float64).reshape(3, 3) - projection = np.asarray(source.P, dtype=np.float64).reshape(3, 4)[:, :3] + distortion = np.asarray(source.d, dtype=np.float64) + rectification = np.asarray(source.r, dtype=np.float64).reshape(3, 3) + projection = np.asarray(source.p, dtype=np.float64).reshape(3, 4)[:, :3] output_matrix = projection if projection[0, 0] > 0 and projection[1, 1] > 0 else matrix size = (image.width, image.height) model = source.distortion_model.strip().lower() diff --git a/dimos/evals/vqa/primitives/edge_tam.py b/dimos/evals/vqa/primitives/edge_tam.py index dfb0f7ac831..d1e59a2f557 100644 --- a/dimos/evals/vqa/primitives/edge_tam.py +++ b/dimos/evals/vqa/primitives/edge_tam.py @@ -26,9 +26,10 @@ from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D if TYPE_CHECKING: + from dimos_generated.sensor_msgs.msg import Image + from dimos.models.segmentation.edge_tam import EdgeTAMImageSegmenterCompatible from dimos.models.vl.base import VlModel - from dimos.msgs.sensor_msgs.Image import Image @dataclass(frozen=True) diff --git a/dimos/evals/vqa/primitives/range.py b/dimos/evals/vqa/primitives/range.py index 7677d6180e8..b13393f6d4f 100644 --- a/dimos/evals/vqa/primitives/range.py +++ b/dimos/evals/vqa/primitives/range.py @@ -23,6 +23,8 @@ from pydantic import BaseModel, ConfigDict, Field from dimos.evals.vqa.contracts import InsufficientEvidenceError, NonEmptyString +from dimos.msgs.geometry import transform_matrix +from dimos.msgs.pointcloud import pointcloud_xyz from dimos.perception.detection.type.detection2d.bbox import Bbox if TYPE_CHECKING: @@ -41,13 +43,15 @@ class _FrameProjection: def from_frame(cls, frame: PointCloudFrame) -> _FrameProjection: """Transform and project valid points into image coordinates.""" camera_info = frame.camera_info - intrinsics = np.asarray(camera_info.K, dtype=np.float64).reshape(3, 3) + intrinsics = np.asarray(camera_info.k, dtype=np.float64).reshape(3, 3) - raw_points, _ = frame.pointcloud.as_numpy() + raw_points = pointcloud_xyz(frame.pointcloud) points = np.asarray(raw_points, dtype=np.float64) points = points[np.all(np.isfinite(points), axis=1)] - transform = np.asarray(frame.pointcloud_to_camera.to_matrix(), dtype=np.float64) + transform = np.asarray( + transform_matrix(frame.pointcloud_to_camera.transform), dtype=np.float64 + ) homogeneous = np.column_stack((points, np.ones(len(points), dtype=np.float64))) camera_points = (transform @ homogeneous.T).T[:, :3] in_front = np.all(np.isfinite(camera_points), axis=1) & (camera_points[:, 2] > 0) diff --git a/dimos/evals/vqa/primitives/test_edge_tam.py b/dimos/evals/vqa/primitives/test_edge_tam.py index a29eeab2175..48fcabe3c42 100644 --- a/dimos/evals/vqa/primitives/test_edge_tam.py +++ b/dimos/evals/vqa/primitives/test_edge_tam.py @@ -14,38 +14,46 @@ from __future__ import annotations -from typing import TYPE_CHECKING, cast +from collections import Counter +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from dimos.evals.vqa.contracts import InsufficientEvidenceError from dimos.evals.vqa.pointcloud_frame import PointCloudFrame +from dimos.evals.vqa.primitives import range as range_primitives from dimos.evals.vqa.primitives.edge_tam import EdgeTAMObjectMaskPipeline from dimos.evals.vqa.primitives.range import LidarRangeEstimator from dimos.models.vl.base import VlModel -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.camera_info import camera_info_from_intrinsics +from dimos.msgs.image import image_from_array +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.perception.detection.type.detection2d.bbox import Bbox, Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.detection.type.detection2d.seg import Detection2DSeg -if TYPE_CHECKING: - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + +def _cloud(points: np.ndarray) -> PointCloud2: + return pointcloud_from_xyz( + points, header=Header(frame_id="lidar", stamp=time_from_seconds(9.98)) + ) -class _PointCloud: - def __init__(self, points: np.ndarray) -> None: - self._points = points - self.read_count = 0 - self.frame_id = "lidar" - self.ts = 9.98 +@pytest.fixture +def cloud_reads(monkeypatch: pytest.MonkeyPatch) -> Counter[int]: + reads: Counter[int] = Counter() + original = range_primitives.pointcloud_xyz - def as_numpy(self) -> tuple[np.ndarray, None]: - self.read_count += 1 - return self._points, None + def counted(cloud: PointCloud2) -> np.ndarray: + reads[id(cloud)] += 1 + return original(cloud) + + monkeypatch.setattr(range_primitives, "pointcloud_xyz", counted) + return reads class _TestVlModel(VlModel): @@ -73,7 +81,7 @@ def query_detections( class_id=0, confidence=1.0, name=query, - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, ) for index, box in enumerate(self._boxes) @@ -98,7 +106,7 @@ def query_detections( class_id=0, confidence=1.0, name=query, - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, ) ], @@ -138,7 +146,7 @@ def segment( def _frame( - pointcloud: _PointCloud, + pointcloud: PointCloud2, *, index: int = 0, transform: Transform | None = None, @@ -147,13 +155,21 @@ def _frame( cx: float = 5.0, cy: float = 5.0, ) -> PointCloudFrame: - image = Image(data=np.zeros((10, 10, 3), dtype=np.uint8), ts=10.0) + image = image_from_array( + np.zeros((10, 10, 3), dtype=np.uint8), + encoding="rgb8", + header=Header(stamp=time_from_seconds(10)), + ) return PointCloudFrame( index=index, image=image, - pointcloud=cast("PointCloud2", pointcloud), - camera_info=CameraInfo.from_intrinsics(fx, fy, cx, cy, 10, 10), - pointcloud_to_camera=transform or Transform.identity(), + pointcloud=pointcloud, + camera_info=camera_info_from_intrinsics(fx, fy, cx, cy, 10, 10, header=Header()), + pointcloud_to_camera=TransformStamped( + header=Header(frame_id="optical"), + child_frame_id="lidar", + transform=transform or Transform(rotation=Quaternion(w=1)), + ), image_observation_timestamp=10.0, pointcloud_observation_timestamp=9.98, calibration_source="synthetic-calibration", @@ -176,7 +192,11 @@ def _range_estimator( def test_mask_estimator_batches_and_caches_without_pointcloud() -> None: - image = Image(data=np.zeros((10, 10, 3), dtype=np.uint8), ts=10.0) + image = image_from_array( + np.zeros((10, 10, 3), dtype=np.uint8), + encoding="rgb8", + header=Header(stamp=time_from_seconds(10)), + ) left_mask = np.zeros((10, 10), dtype=np.uint8) right_mask = np.zeros((10, 10), dtype=np.uint8) left_mask[:, :4] = 1 @@ -200,14 +220,18 @@ def test_mask_estimator_batches_and_caches_without_pointcloud() -> None: assert cached is left assert segmenter.call_count == 1 - next_image = Image(data=np.zeros((10, 10, 3), dtype=np.uint8), ts=11.0) + next_image = image_from_array( + np.zeros((10, 10, 3), dtype=np.uint8), + encoding="rgb8", + header=Header(stamp=time_from_seconds(11)), + ) estimator.estimate(next_image, "left person") assert segmenter.call_count == 2 def test_projects_transformed_points_and_rejects_invalid_projections() -> None: - pointcloud = _PointCloud( + pointcloud = _cloud( np.array( [ [0.0, 0.0, 1.0], @@ -218,7 +242,7 @@ def test_projects_transformed_points_and_rejects_invalid_projections() -> None: ) frame = _frame( pointcloud, - transform=Transform(translation=Vector3(1.0, 0.0, 0.0)), + transform=Transform(translation=Vector3(x=1.0)), cx=0.0, cy=0.0, ) @@ -234,7 +258,7 @@ def test_projects_transformed_points_and_rejects_invalid_projections() -> None: def test_mask_selection_uses_nearest_camera_z_point_per_pixel() -> None: - pointcloud = _PointCloud( + pointcloud = _cloud( np.array( [ [0.0, 0.0, 1.0], @@ -256,18 +280,18 @@ def test_mask_selection_uses_nearest_camera_z_point_per_pixel() -> None: @pytest.mark.parametrize("boxes", [[], [(0.0, 0.0, 4.0, 4.0), (5.0, 5.0, 9.0, 9.0)]]) -def test_requires_exactly_one_valid_detection(boxes: list[Bbox]) -> None: - cloud = _PointCloud(np.array([[0.0, 0.0, 1.0]])) +def test_requires_exactly_one_valid_detection(boxes: list[Bbox], cloud_reads: Counter[int]) -> None: + cloud = _cloud(np.array([[0.0, 0.0, 1.0]])) estimator = _range_estimator(_Detector(boxes), _Segmenter(_full_mask())) with pytest.raises(InsufficientEvidenceError, match="exactly one valid detected"): estimator.estimate(_frame(cloud), "cup") - assert cloud.read_count == 0 + assert cloud_reads[id(cloud)] == 0 def test_rejects_too_few_mask_supporting_points() -> None: - pointcloud = _PointCloud( + pointcloud = _cloud( np.array([[-2.0, 0.0, 1.0], [-1.5, 0.0, 1.0], [-1.0, 0.0, 1.0], [-0.5, 0.0, 1.0]]) ) estimator = _range_estimator(_Detector([(0.0, 0.0, 10.0, 10.0)]), _Segmenter(_full_mask())) @@ -284,13 +308,13 @@ def test_rejects_segmentation_mask_with_wrong_dimensions() -> None: ) with pytest.raises(InsufficientEvidenceError, match="one valid segmentation mask"): - estimator.estimate(_frame(_PointCloud(np.array([[0.0, 0.0, 1.0]]))), "bottle") + estimator.estimate(_frame(_cloud(np.array([[0.0, 0.0, 1.0]]))), "bottle") def test_returns_median_euclidean_range_and_auditable_quartiles() -> None: points = np.array([[u * z, 0.0, z] for u, z in zip(range(1, 6), [1, 2, 3, 4, 20], strict=True)]) expected_ranges = np.linalg.norm(points, axis=1) - frame = _frame(_PointCloud(points), fx=1.0, fy=1.0, cx=0.0, cy=0.0) + frame = _frame(_cloud(points), fx=1.0, fy=1.0, cx=0.0, cy=0.0) evidence = _range_estimator( _Detector([(0.0, 0.0, 9.0, 9.0)]), _Segmenter(_full_mask()) @@ -306,10 +330,12 @@ def test_returns_median_euclidean_range_and_auditable_quartiles() -> None: assert evidence.model_dump(mode="json")["mask_bbox_xyxy"] == [0.0, 0.0, 9.0, 9.0] -def test_projection_cache_reuses_across_objects_only_for_the_same_explicit_frame() -> None: +def test_projection_cache_reuses_across_objects_only_for_the_same_explicit_frame( + cloud_reads: Counter[int], +) -> None: points = np.array([[float(2 * x), 0.0, 2.0] for x in range(5)]) - first_cloud = _PointCloud(points) - second_cloud = _PointCloud(points) + first_cloud = _cloud(points) + second_cloud = _cloud(points) first_frame = _frame(first_cloud, index=3, fx=1.0, fy=1.0, cx=0.0, cy=0.0) second_frame = _frame(second_cloud, index=3, fx=1.0, fy=1.0, cx=0.0, cy=0.0) estimator = _range_estimator(_Detector([(0.0, 0.0, 9.0, 9.0)]), _Segmenter(_full_mask())) @@ -318,12 +344,12 @@ def test_projection_cache_reuses_across_objects_only_for_the_same_explicit_frame estimator.estimate(first_frame, "crate") estimator.estimate(second_frame, "box") - assert first_cloud.read_count == 1 - assert second_cloud.read_count == 1 + assert cloud_reads[id(first_cloud)] == 1 + assert cloud_reads[id(second_cloud)] == 1 -def test_estimate_many_batches_masks_and_reuses_one_projection() -> None: - cloud = _PointCloud(np.array([[-1.5, 0.0, 1.0], [3.0, 0.0, 3.0]])) +def test_estimate_many_batches_masks_and_reuses_one_projection(cloud_reads: Counter[int]) -> None: + cloud = _cloud(np.array([[-1.5, 0.0, 1.0], [3.0, 0.0, 3.0]])) left_mask = np.zeros((10, 10), dtype=np.uint8) right_mask = np.zeros((10, 10), dtype=np.uint8) left_mask[5, 2] = 1 @@ -346,4 +372,4 @@ def test_estimate_many_batches_masks_and_reuses_one_projection() -> None: assert right.object_name == "table" assert left.camera_range_m < right.camera_range_m assert segmenter.call_count == 1 - assert cloud.read_count == 1 + assert cloud_reads[id(cloud)] == 1 diff --git a/dimos/evals/vqa/test_contracts.py b/dimos/evals/vqa/test_contracts.py index e0ad255ed22..79fbd9807b5 100644 --- a/dimos/evals/vqa/test_contracts.py +++ b/dimos/evals/vqa/test_contracts.py @@ -18,7 +18,10 @@ import json from pathlib import Path +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header import numpy as np +from PIL import Image as PILImage from pydantic import ValidationError import pytest @@ -52,11 +55,16 @@ from dimos.evals.vqa.primitives.range import ObjectRangeEvidence from dimos.evals.vqa.suite import load_suite from dimos.models.vl.base import VlModel -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_from_array, image_view +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.perception.detection.type.detection2d.bbox import Bbox, Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D +def _image(pixels: np.ndarray, *, ts: float = 0.0) -> Image: + return image_from_array(pixels, encoding="rgb8", header=Header(stamp=time_from_seconds(ts))) + + class _TestVlModel(VlModel): def __init__(self) -> None: pass @@ -148,7 +156,7 @@ def propose(self, image: Image, families: Sequence[FamilySpec]) -> Sequence[Ques class _EmptyThenAuthor: def propose(self, image: Image, families: Sequence[FamilySpec]) -> Sequence[QuestionProposal]: - if image.ts == 1.0: + if to_seconds(image.header.stamp) == 1.0: return () return (QuestionProposal(family="presence", object_names=("chair",)),) @@ -178,7 +186,7 @@ def query_detections( class_id=-1, confidence=1.0, name=query, - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, ) ) @@ -201,7 +209,7 @@ def query_detections( class_id=-1, confidence=1.0, name=query, - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, ) for index, box in enumerate(self._boxes) @@ -310,7 +318,7 @@ def estimate(self, image: Image, object_name: str) -> ObjectMaskEvidence: class_id=-1, confidence=1.0, name=object_name, - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, ) return ObjectMaskEvidence( @@ -439,7 +447,7 @@ def test_closest_object_proposal_requires_two_to_five_distinct_references() -> N def test_object_distance_requires_range_evidence() -> None: - image = Image.from_numpy(np.zeros((4, 4, 3), dtype=np.uint8)) + image = _image(np.zeros((4, 4, 3), dtype=np.uint8)) proposal = QuestionProposal(family="object_distance", object_names=("chair",)) with pytest.raises(InsufficientEvidenceError, match="image-aligned point-cloud"): @@ -456,7 +464,7 @@ def test_object_distance_requires_range_evidence() -> None: ), ) def test_object_distance_buckets_range_evidence(range_m: float, expected: str) -> None: - image = Image.from_numpy(np.zeros((4, 4, 3), dtype=np.uint8)) + image = _image(np.zeros((4, 4, 3), dtype=np.uint8)) proposal = QuestionProposal(family="object_distance", object_names=("chair",)) answer = answer_question( @@ -473,7 +481,7 @@ def test_object_distance_buckets_range_evidence(range_m: float, expected: str) - def test_object_distance_rejects_evidence_crossing_bucket_boundary() -> None: - image = Image.from_numpy(np.zeros((4, 4, 3), dtype=np.uint8)) + image = _image(np.zeros((4, 4, 3), dtype=np.uint8)) proposal = QuestionProposal(family="object_distance", object_names=("chair",)) with pytest.raises(InsufficientEvidenceError, match="uncertainty crosses"): @@ -503,7 +511,7 @@ def test_closest_object_compares_multiple_named_references( quartiles: tuple[tuple[float, float, float], ...], expected: str, ) -> None: - image = Image.from_numpy(np.zeros((4, 10, 3), dtype=np.uint8)) + image = _image(np.zeros((4, 10, 3), dtype=np.uint8)) proposal = QuestionProposal( family="closest_object", object_names=("left person", "right person", "chair"), @@ -524,7 +532,7 @@ def test_closest_object_compares_multiple_named_references( def test_closest_object_rejects_overlapping_range_intervals() -> None: - image = Image.from_numpy(np.zeros((4, 10, 3), dtype=np.uint8)) + image = _image(np.zeros((4, 10, 3), dtype=np.uint8)) proposal = QuestionProposal( family="closest_object", object_names=("left person", "right person", "chair"), @@ -553,7 +561,7 @@ def test_closest_object_rejects_reordered_range_results() -> None: with pytest.raises(ValueError, match="requested object order"): answer_question( proposal, - Image.from_numpy(np.zeros((4, 10, 3), dtype=np.uint8)), + _image(np.zeros((4, 10, 3), dtype=np.uint8)), _Detector(present=True), object(), _ReorderedClosestRangeEstimator((1.0, 1.1, 1.2), (2.0, 2.1, 2.2)), @@ -586,7 +594,7 @@ def test_family_answer_requires_case_insensitive_unique_choices() -> None: def test_presence_family_derives_answer_from_detector_evidence() -> None: - image = Image.from_numpy(np.zeros((4, 4, 3), dtype=np.uint8)) + image = _image(np.zeros((4, 4, 3), dtype=np.uint8)) proposal = QuestionProposal(family="presence", object_names=("chair",)) present = answer_question(proposal, image, _Detector(present=True)) @@ -610,7 +618,7 @@ def test_presence_family_derives_answer_from_detector_evidence() -> None: ), ) def test_horizontal_direction_uses_detection_center(box: Bbox, expected: str) -> None: - image = Image.from_numpy(np.zeros((60, 90, 3), dtype=np.uint8)) + image = _image(np.zeros((60, 90, 3), dtype=np.uint8)) proposal = QuestionProposal(family="horizontal_direction", object_names=("robot",)) answer = answer_question(proposal, image, _BoxesDetector((box,))) @@ -621,7 +629,7 @@ def test_horizontal_direction_uses_detection_center(box: Bbox, expected: str) -> def test_horizontal_direction_rejects_ambiguous_instances() -> None: - image = Image.from_numpy(np.zeros((60, 90, 3), dtype=np.uint8)) + image = _image(np.zeros((60, 90, 3), dtype=np.uint8)) proposal = QuestionProposal(family="horizontal_direction", object_names=("chair",)) detector = _BoxesDetector( ( @@ -639,7 +647,7 @@ def test_horizontal_direction_rejects_ambiguous_instances() -> None: ((1, "one"), (2, "two"), (3, "three"), (4, "four or more"), (5, "four or more")), ) def test_object_count_buckets_detected_instances(count: int, expected: str) -> None: - image = Image.from_numpy(np.zeros((60, 90, 3), dtype=np.uint8)) + image = _image(np.zeros((60, 90, 3), dtype=np.uint8)) proposal = QuestionProposal(family="object_count", object_names=("box",)) boxes = tuple((float(index), 0.0, float(index + 1), 1.0) for index in range(count)) @@ -652,7 +660,7 @@ def test_object_count_buckets_detected_instances(count: int, expected: str) -> N def test_object_count_requires_at_least_one_detection() -> None: - image = Image.from_numpy(np.zeros((60, 90, 3), dtype=np.uint8)) + image = _image(np.zeros((60, 90, 3), dtype=np.uint8)) proposal = QuestionProposal(family="object_count", object_names=("box",)) with pytest.raises(InsufficientEvidenceError, match="to count"): @@ -672,7 +680,7 @@ def test_image_coverage_chooses_scheme_with_largest_boundary_margin( boundaries: list[float], margin: float, ) -> None: - image = Image.from_numpy(np.zeros((10, 10, 3), dtype=np.uint8)) + image = _image(np.zeros((10, 10, 3), dtype=np.uint8)) proposal = QuestionProposal(family="image_coverage", object_names=("chair",)) answer = answer_question( @@ -694,7 +702,7 @@ def test_image_coverage_requires_mask_evidence() -> None: with pytest.raises(InsufficientEvidenceError, match="segmentation-mask"): answer_question( proposal, - Image.from_numpy(np.zeros((10, 10, 3), dtype=np.uint8)), + _image(np.zeros((10, 10, 3), dtype=np.uint8)), _Detector(present=True), ) @@ -710,7 +718,7 @@ def test_largest_visible_area_requires_distinct_references() -> None: def test_largest_visible_area_accepts_winner_at_twenty_percent_margin() -> None: - image = Image.from_numpy(np.zeros((10, 10, 3), dtype=np.uint8)) + image = _image(np.zeros((10, 10, 3), dtype=np.uint8)) proposal = QuestionProposal( family="largest_visible_area", object_names=("chair", "table", "box"), @@ -729,7 +737,7 @@ def test_largest_visible_area_accepts_winner_at_twenty_percent_margin() -> None: def test_largest_visible_area_rejects_close_mask_areas() -> None: - image = Image.from_numpy(np.zeros((10, 10, 3), dtype=np.uint8)) + image = _image(np.zeros((10, 10, 3), dtype=np.uint8)) proposal = QuestionProposal( family="largest_visible_area", object_names=("chair", "table"), @@ -745,7 +753,7 @@ def test_largest_visible_area_rejects_close_mask_areas() -> None: def test_openai_author_parses_constrained_proposals() -> None: - image = Image.from_numpy(np.zeros((4, 4, 3), dtype=np.uint8)) + image = _image(np.zeros((4, 4, 3), dtype=np.uint8)) author = OpenAIQuestionAuthor(_QuestionModel()) proposals = author.propose(image, AVAILABLE_FAMILIES) @@ -768,7 +776,7 @@ def test_openai_author_parses_constrained_proposals() -> None: def test_openai_author_keeps_valid_items_from_partly_invalid_response() -> None: - image = Image.from_numpy(np.zeros((4, 4, 3), dtype=np.uint8)) + image = _image(np.zeros((4, 4, 3), dtype=np.uint8)) author = OpenAIQuestionAuthor(_PartlyInvalidQuestionModel()) proposals = author.propose(image, AVAILABLE_FAMILIES) @@ -777,7 +785,7 @@ def test_openai_author_keeps_valid_items_from_partly_invalid_response() -> None: def test_openai_author_skips_unavailable_family() -> None: - image = Image.from_numpy(np.zeros((4, 4, 3), dtype=np.uint8)) + image = _image(np.zeros((4, 4, 3), dtype=np.uint8)) author = OpenAIQuestionAuthor(_UnavailableFamilyQuestionModel()) proposals = author.propose(image, (AVAILABLE_FAMILIES[0],)) @@ -941,7 +949,7 @@ def test_suite_reports_empty_cases_before_duplicate_labels(tmp_path: Path) -> No def test_generate_and_evaluate_one_image(tmp_path: Path) -> None: output = tmp_path / "nested" / "vqa" - image = Image.from_numpy(np.arange(48, dtype=np.uint8).reshape(4, 4, 3), ts=12.5) + image = _image(np.arange(48, dtype=np.uint8).reshape(4, 4, 3), ts=12.5) request = GenerationRequest(dataset="recording.db", image_index=4, output=output) result = generate_frames_dataset( @@ -954,7 +962,7 @@ def test_generate_and_evaluate_one_image(tmp_path: Path) -> None: assert result.cases[0].id == "frame-000004-chair-presence" image_path = output / "assets" / "frame-000004.png" assert image_path.is_file() - assert np.array_equal(Image.from_file(image_path).data, image.data) + assert np.array_equal(np.asarray(PILImage.open(image_path)), image_view(image)) assert (output / "cases.jsonl").is_file() assert (output / "labels.jsonl").is_file() assert (output / "audit" / "frame-000004" / "ground_truth.json").is_file() @@ -978,7 +986,7 @@ def test_generation_rejects_nonempty_output_without_modifying_it(tmp_path: Path) existing = output / "existing.txt" existing.write_text("keep me") request = GenerationRequest(dataset="recording.db", image_index=4, output=output) - image = Image.from_numpy(np.zeros((4, 4, 3), dtype=np.uint8)) + image = _image(np.zeros((4, 4, 3), dtype=np.uint8)) with pytest.raises(FileExistsError, match="output directory is not empty"): generate_frames_dataset( @@ -994,7 +1002,7 @@ def test_generation_rejects_nonempty_output_without_modifying_it(tmp_path: Path) def test_generation_audits_invalid_family_proposals_and_continues(tmp_path: Path) -> None: output = tmp_path / "vqa" - image = Image.from_numpy(np.zeros((4, 4, 3), dtype=np.uint8)) + image = _image(np.zeros((4, 4, 3), dtype=np.uint8)) request = GenerationRequest(dataset="recording.db", image_index=4, output=output) result = generate_frames_dataset( @@ -1012,7 +1020,7 @@ def test_generation_audits_invalid_family_proposals_and_continues(tmp_path: Path def test_generate_distance_case_from_pointcloud_frame(tmp_path: Path) -> None: output = tmp_path / "vqa" - image = Image.from_numpy(np.zeros((4, 4, 3), dtype=np.uint8), ts=12.5) + image = _image(np.zeros((4, 4, 3), dtype=np.uint8), ts=12.5) request = GenerationRequest(dataset="recording.db", image_index=4, output=output) result = generate_frames_dataset( @@ -1036,7 +1044,7 @@ def test_generate_distance_case_from_pointcloud_frame(tmp_path: Path) -> None: def test_generate_closest_object_case_from_pointcloud_frame(tmp_path: Path) -> None: output = tmp_path / "vqa" - image = Image.from_numpy(np.zeros((4, 10, 3), dtype=np.uint8), ts=12.5) + image = _image(np.zeros((4, 10, 3), dtype=np.uint8), ts=12.5) request = GenerationRequest(dataset="recording.db", image_index=4, output=output) result = generate_frames_dataset( @@ -1065,7 +1073,7 @@ def test_generate_closest_object_case_from_pointcloud_frame(tmp_path: Path) -> N def test_generate_image_coverage_case_without_pointcloud(tmp_path: Path) -> None: output = tmp_path / "vqa" - image = Image.from_numpy(np.zeros((10, 10, 3), dtype=np.uint8), ts=12.5) + image = _image(np.zeros((10, 10, 3), dtype=np.uint8), ts=12.5) result = generate_frames_dataset( GenerationRequest(dataset="recording.db", image_index=4, output=output), @@ -1086,7 +1094,7 @@ def test_generation_deduplicates_reordered_closest_object_references(tmp_path: P ( GenerationFrame( 4, - Image.from_numpy(np.zeros((4, 10, 3), dtype=np.uint8)), + _image(np.zeros((4, 10, 3), dtype=np.uint8)), object(), ), ), @@ -1108,7 +1116,7 @@ def test_generation_deduplicates_reordered_closest_object_references(tmp_path: P def test_generation_deduplicates_reordered_largest_area_references(tmp_path: Path) -> None: result = generate_frames_dataset( GenerationRequest(dataset="recording.db", image_index=4, output=tmp_path / "vqa"), - (GenerationFrame(4, Image.from_numpy(np.zeros((10, 10, 3), dtype=np.uint8))),), + (GenerationFrame(4, _image(np.zeros((10, 10, 3), dtype=np.uint8))),), _DuplicateLargestVisibleAreaAuthor(), _Detector(present=True), mask_estimator=_MaskEstimator({"chair": 60, "table": 40, "box": 10}), @@ -1120,7 +1128,7 @@ def test_generation_deduplicates_reordered_largest_area_references(tmp_path: Pat def test_frame_without_pointcloud_does_not_expose_distance_family(tmp_path: Path) -> None: output = tmp_path / "vqa" - image = Image.from_numpy(np.zeros((4, 4, 3), dtype=np.uint8), ts=12.5) + image = _image(np.zeros((4, 4, 3), dtype=np.uint8), ts=12.5) request = GenerationRequest(dataset="recording.db", image_index=4, output=output) author = _RecordingAuthor() @@ -1141,7 +1149,7 @@ def test_frame_without_pointcloud_exposes_mask_families_when_masks_are_available tmp_path: Path, ) -> None: output = tmp_path / "vqa" - image = Image.from_numpy(np.zeros((10, 10, 3), dtype=np.uint8), ts=12.5) + image = _image(np.zeros((10, 10, 3), dtype=np.uint8), ts=12.5) author = _RecordingAuthor() generate_frames_dataset( @@ -1171,8 +1179,8 @@ def test_generate_multiple_images_aggregates_frame_artifacts(tmp_path: Path) -> output=output, ) frames = ( - GenerationFrame(1, Image.from_numpy(np.zeros((4, 4, 3), dtype=np.uint8), ts=1.0)), - GenerationFrame(3, Image.from_numpy(np.zeros((4, 4, 3), dtype=np.uint8), ts=3.0)), + GenerationFrame(1, _image(np.zeros((4, 4, 3), dtype=np.uint8), ts=1.0)), + GenerationFrame(3, _image(np.zeros((4, 4, 3), dtype=np.uint8), ts=3.0)), ) result = generate_frames_dataset( @@ -1209,8 +1217,8 @@ def test_empty_frame_does_not_count_as_rejected_question(tmp_path: Path) -> None output = tmp_path / "vqa" request = GenerationRequest(dataset="recording.db", start=1, stop=3, output=output) frames = ( - GenerationFrame(1, Image.from_numpy(np.zeros((4, 4, 3), dtype=np.uint8), ts=1.0)), - GenerationFrame(2, Image.from_numpy(np.zeros((4, 4, 3), dtype=np.uint8), ts=2.0)), + GenerationFrame(1, _image(np.zeros((4, 4, 3), dtype=np.uint8), ts=1.0)), + GenerationFrame(2, _image(np.zeros((4, 4, 3), dtype=np.uint8), ts=2.0)), ) generate_frames_dataset( @@ -1228,7 +1236,7 @@ def test_empty_frame_does_not_count_as_rejected_question(tmp_path: Path) -> None def test_generation_keeps_answered_questions_and_audits_rejections(tmp_path: Path) -> None: output = tmp_path / "vqa" - image = Image.from_numpy(np.zeros((60, 90, 3), dtype=np.uint8), ts=12.5) + image = _image(np.zeros((60, 90, 3), dtype=np.uint8), ts=12.5) request = GenerationRequest(dataset="recording.db", image_index=4, output=output) detector = _BoxesDetector( ( @@ -1252,7 +1260,7 @@ def test_generation_keeps_answered_questions_and_audits_rejections(tmp_path: Pat def test_generation_propagates_detector_failures(tmp_path: Path) -> None: output = tmp_path / "vqa" - image = Image.from_numpy(np.zeros((4, 4, 3), dtype=np.uint8)) + image = _image(np.zeros((4, 4, 3), dtype=np.uint8)) request = GenerationRequest(dataset="recording.db", image_index=4, output=output) with pytest.raises(ValueError, match="detector broke"): @@ -1270,8 +1278,8 @@ def test_later_frame_failure_does_not_publish_partial_dataset(tmp_path: Path) -> output = tmp_path / "vqa" request = GenerationRequest(dataset="recording.db", start=1, stop=3, output=output) frames = ( - GenerationFrame(1, Image.from_numpy(np.zeros((4, 4, 3), dtype=np.uint8), ts=1.0)), - GenerationFrame(2, Image.from_numpy(np.zeros((4, 4, 3), dtype=np.uint8), ts=2.0)), + GenerationFrame(1, _image(np.zeros((4, 4, 3), dtype=np.uint8), ts=1.0)), + GenerationFrame(2, _image(np.zeros((4, 4, 3), dtype=np.uint8), ts=2.0)), ) with pytest.raises(ValueError, match="detector broke"): @@ -1287,7 +1295,7 @@ def test_later_frame_failure_does_not_publish_partial_dataset(tmp_path: Path) -> def test_generation_deduplicates_proposals_and_suffixes_id_collisions(tmp_path: Path) -> None: output = tmp_path / "vqa" - image = Image.from_numpy(np.zeros((4, 4, 3), dtype=np.uint8)) + image = _image(np.zeros((4, 4, 3), dtype=np.uint8)) request = GenerationRequest(dataset="recording.db", image_index=4, output=output) result = generate_frames_dataset( diff --git a/dimos/evals/vqa/test_pointcloud_frame.py b/dimos/evals/vqa/test_pointcloud_frame.py index bdb262edb31..119af072cff 100644 --- a/dimos/evals/vqa/test_pointcloud_frame.py +++ b/dimos/evals/vqa/test_pointcloud_frame.py @@ -14,6 +14,10 @@ from pathlib import Path +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage import numpy as np import pytest @@ -25,12 +29,43 @@ ) from dimos.memory.store.memory import MemoryStore from dimos.memory.store.sqlite import SqliteStore -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.camera_info import camera_info_from_intrinsics +from dimos.msgs.geometry import inverse_transform, transform_matrix +from dimos.msgs.image import image_from_array, image_view +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import time_from_seconds, to_seconds + + +def _image(data, *, frame_id="", ts=0.0): + return image_from_array( + data, encoding="rgb8", header=Header(frame_id=frame_id, stamp=time_from_seconds(ts)) + ) + + +def _cloud(data, *, frame_id="", timestamp=0.0): + return pointcloud_from_xyz( + data, header=Header(frame_id=frame_id, stamp=time_from_seconds(timestamp)) + ) + + +def _calibration(fx, fy, cx, cy, width, height, *, frame_id=""): + result = camera_info_from_intrinsics( + max(fx, 1.0), max(fy, 1.0), cx, cy, width, height, header=Header(frame_id=frame_id) + ) + result.k[0] = fx + result.k[4] = fy + return result + + +def _transform(*, translation=None, frame_id="", child_frame_id="", ts=0.0): + return TransformStamped( + header=Header(frame_id=frame_id, stamp=time_from_seconds(ts)), + child_frame_id=child_frame_id, + transform=Transform( + translation=translation if translation is not None else Vector3(), + rotation=Quaternion(w=1.0), + ), + ) def _populate_required_streams( @@ -40,18 +75,18 @@ def _populate_required_streams( include_empty_pointlio: bool = False, resolvable_tf: bool = True, ) -> None: - image = Image.from_numpy( + image = _image( np.zeros((2, 2, 3), dtype=np.uint8), frame_id="camera_optical", ts=10.0, ) - cloud = PointCloud2.from_numpy( + cloud = _cloud( np.array([[0.0, 0.0, 1.0]]), frame_id="world", timestamp=10.0, ) - camera_info = CameraInfo.from_intrinsics(1.0, 1.0, 0.0, 0.0, 2, 2, frame_id="camera_optical") - transform = Transform( + camera_info = _calibration(1.0, 1.0, 0.0, 0.0, 2, 2, frame_id="camera_optical") + transform = _transform( frame_id="world" if resolvable_tf else "unrelated_parent", child_frame_id="camera_optical" if resolvable_tf else "unrelated_child", ts=10.0, @@ -68,7 +103,7 @@ def _populate_required_streams( if empty != "camera_info": camera_infos.append(camera_info, ts=10.0) if empty != "tf": - transforms.append(TFMessage(transform), ts=10.0) + transforms.append(TFMessage(transforms=[transform]), ts=10.0) if empty != "lidar": lidar.append(cloud, ts=10.0) @@ -85,16 +120,16 @@ def test_align_one_rejects_observation_outside_tolerance() -> None: def test_rectifier_rejects_invalid_camera_intrinsics() -> None: - image = Image.from_numpy(np.zeros((2, 2, 3), dtype=np.uint8)) - camera_info = CameraInfo.from_intrinsics(0.0, 1.0, 0.0, 0.0, 2, 2) + image = _image(np.zeros((2, 2, 3), dtype=np.uint8)) + camera_info = _calibration(0.0, 1.0, 0.0, 0.0, 2, 2) with pytest.raises(ValueError, match="positive focal lengths"): _ImageRectifier().rectify(image, camera_info) def test_rectifier_rejects_calibration_dimension_mismatch() -> None: - image = Image.from_numpy(np.zeros((2, 2, 3), dtype=np.uint8)) - camera_info = CameraInfo.from_intrinsics(1.0, 1.0, 0.0, 0.0, 1, 2) + image = _image(np.zeros((2, 2, 3), dtype=np.uint8)) + camera_info = _calibration(1.0, 1.0, 0.0, 0.0, 1, 2) with pytest.raises(ValueError, match="do not match image dimensions"): _ImageRectifier().rectify(image, camera_info) @@ -137,7 +172,7 @@ def test_uses_nonempty_lidar_when_pointlio_stream_is_empty(tmp_path: Path) -> No _populate_required_streams(store, include_empty_pointlio=True) with PointCloudFrameLoader(dataset) as loader: - assert loader.load(0).pointcloud.frame_id == "world" + assert loader.load(0).pointcloud.header.frame_id == "world" def test_unresolvable_tf_still_allows_image_only_fallback( @@ -153,28 +188,28 @@ def test_unresolvable_tf_still_allows_image_only_fallback( with loader: with pytest.raises(PointCloudFrameUnavailableError, match="cannot resolve"): loader.load(0) - assert loader.load_image(0).shape == (2, 2, 3) + assert image_view(loader.load_image(0)).shape == (2, 2, 3) def test_load_uses_recorded_camera_info_and_tf( monkeypatch: pytest.MonkeyPatch, tmp_path: Path ) -> None: dataset = tmp_path / "recording.db" - image = Image.from_numpy( + image = _image( np.zeros((2, 2, 3), dtype=np.uint8), frame_id="camera_optical", ts=10.0, ) - cloud = PointCloud2.from_numpy( + cloud = _cloud( np.array([[0.0, 0.0, 1.0]]), frame_id="world", timestamp=10.0, ) - camera_info = CameraInfo.from_intrinsics( + camera_info = _calibration( 1.0, 1.0, 0.0, 0.0, image.width, image.height, frame_id="camera_optical" ) - world_from_camera = Transform( - translation=Vector3(1.0, 2.0, 0.5), + world_from_camera = _transform( + translation=Vector3(x=1.0, y=2.0, z=0.5), frame_id="world", child_frame_id="camera_optical", ts=10.0, @@ -184,14 +219,16 @@ def test_load_uses_recorded_camera_info_and_tf( store.stream("color_image", Image).append(image, ts=10.0) store.stream("lidar", PointCloud2).append(cloud, ts=10.0) store.stream("camera_info", CameraInfo).append(camera_info, ts=9.0) - store.stream("tf", TFMessage).append(TFMessage(world_from_camera), ts=10.0) + store.stream("tf", TFMessage).append(TFMessage(transforms=[world_from_camera]), ts=10.0) loader = PointCloudFrameLoader(dataset) calibrations: list[CameraInfo] = [] def rectify(source: Image, calibration: CameraInfo) -> tuple[Image, CameraInfo]: calibrations.append(calibration) - return source, camera_info.with_ts(source.ts) + result = CameraInfo.decode(camera_info.encode()) + result.header.stamp = source.header.stamp + return source, result monkeypatch.setattr( loader._rectifier, @@ -203,46 +240,45 @@ def rectify(source: Image, calibration: CameraInfo) -> tuple[Image, CameraInfo]: assert calibrations == [camera_info] assert frame.calibration_source == "recorded" - assert frame.camera_info.ts == image.ts - assert np.allclose(frame.pointcloud_to_camera.to_matrix(), (-world_from_camera).to_matrix()) + assert frame.camera_info.header.stamp == image.header.stamp + assert np.allclose( + transform_matrix(frame.pointcloud_to_camera.transform), + transform_matrix(inverse_transform(world_from_camera).transform), + ) def test_load_resolves_tf_through_non_world_root( monkeypatch: pytest.MonkeyPatch, tmp_path: Path ) -> None: dataset = tmp_path / "recording.db" - image = Image.from_numpy( - np.zeros((2, 2, 3), dtype=np.uint8), frame_id="camera_optical", ts=10.0 - ) - cloud = PointCloud2.from_numpy(np.array([[0.0, 0.0, 1.0]]), frame_id="lidar", timestamp=10.0) - camera_info = CameraInfo.from_intrinsics(1.0, 1.0, 0.0, 0.0, 2, 2, frame_id="camera_optical") + image = _image(np.zeros((2, 2, 3), dtype=np.uint8), frame_id="camera_optical", ts=10.0) + cloud = _cloud(np.array([[0.0, 0.0, 1.0]]), frame_id="lidar", timestamp=10.0) + camera_info = _calibration(1.0, 1.0, 0.0, 0.0, 2, 2, frame_id="camera_optical") transforms = ( - Transform(frame_id="map", child_frame_id="camera_optical", ts=10.0), - Transform(frame_id="map", child_frame_id="lidar", ts=10.0), + _transform(frame_id="map", child_frame_id="camera_optical", ts=10.0), + _transform(frame_id="map", child_frame_id="lidar", ts=10.0), ) with SqliteStore(path=dataset) as store: store.stream("color_image", Image).append(image, ts=10.0) store.stream("lidar", PointCloud2).append(cloud, ts=10.0) store.stream("camera_info", CameraInfo).append(camera_info, ts=10.0) - store.stream("tf", TFMessage).append(TFMessage(*transforms), ts=10.0) + store.stream("tf", TFMessage).append(TFMessage(transforms=[*transforms]), ts=10.0) loader = PointCloudFrameLoader(dataset) monkeypatch.setattr(loader._rectifier, "rectify", lambda source, info: (source, info)) with loader: frame = loader.load(0) - assert np.allclose(frame.pointcloud_to_camera.to_matrix(), np.eye(4)) + assert np.allclose(transform_matrix(frame.pointcloud_to_camera.transform), np.eye(4)) def test_load_applies_camera_rectification_to_pointcloud_transform( monkeypatch: pytest.MonkeyPatch, tmp_path: Path ) -> None: dataset = tmp_path / "recording.db" - image = Image.from_numpy( - np.zeros((2, 2, 3), dtype=np.uint8), frame_id="camera_optical", ts=10.0 - ) - cloud = PointCloud2.from_numpy(np.array([[0.0, 0.0, 1.0]]), frame_id="lidar", timestamp=10.0) - camera_info = CameraInfo.from_intrinsics(1.0, 1.0, 0.0, 0.0, 2, 2, frame_id="camera_optical") + image = _image(np.zeros((2, 2, 3), dtype=np.uint8), frame_id="camera_optical", ts=10.0) + cloud = _cloud(np.array([[0.0, 0.0, 1.0]]), frame_id="lidar", timestamp=10.0) + camera_info = _calibration(1.0, 1.0, 0.0, 0.0, 2, 2, frame_id="camera_optical") rectification = np.array( [ [0.0, -1.0, 0.0], @@ -250,13 +286,15 @@ def test_load_applies_camera_rectification_to_pointcloud_transform( [0.0, 0.0, 1.0], ] ) - camera_info.R = rectification.ravel().tolist() + camera_info.r = rectification.ravel().tolist() with SqliteStore(path=dataset) as store: store.stream("color_image", Image).append(image, ts=10.0) store.stream("lidar", PointCloud2).append(cloud, ts=10.0) store.stream("camera_info", CameraInfo).append(camera_info, ts=10.0) store.stream("tf", TFMessage).append( - TFMessage(Transform(frame_id="camera_optical", child_frame_id="lidar", ts=10.0)), + TFMessage( + transforms=[_transform(frame_id="camera_optical", child_frame_id="lidar", ts=10.0)] + ), ts=10.0, ) @@ -265,7 +303,9 @@ def test_load_applies_camera_rectification_to_pointcloud_transform( with loader: frame = loader.load(0) - assert np.allclose(frame.pointcloud_to_camera.to_matrix()[:3, :3], rectification) + assert np.allclose( + transform_matrix(frame.pointcloud_to_camera.transform)[:3, :3], rectification + ) def test_load_selects_camera_info_for_each_image_timestamp( @@ -273,16 +313,16 @@ def test_load_selects_camera_info_for_each_image_timestamp( ) -> None: dataset = tmp_path / "recording.db" images = ( - Image.from_numpy(np.zeros((2, 2, 3), dtype=np.uint8), ts=10.0), - Image.from_numpy(np.zeros((2, 2, 3), dtype=np.uint8), ts=20.0), + _image(np.zeros((2, 2, 3), dtype=np.uint8), ts=10.0), + _image(np.zeros((2, 2, 3), dtype=np.uint8), ts=20.0), ) camera_infos = ( - CameraInfo.from_intrinsics(10.0, 10.0, 0.0, 0.0, 2, 2, frame_id="camera_early"), - CameraInfo.from_intrinsics(100.0, 100.0, 0.0, 0.0, 2, 2, frame_id="camera_late"), + _calibration(10.0, 10.0, 0.0, 0.0, 2, 2, frame_id="camera_early"), + _calibration(100.0, 100.0, 0.0, 0.0, 2, 2, frame_id="camera_late"), ) camera_transforms = ( - Transform(frame_id="world", child_frame_id="camera_early", ts=10.0), - Transform(frame_id="world", child_frame_id="camera_late", ts=20.0), + _transform(frame_id="world", child_frame_id="camera_early", ts=10.0), + _transform(frame_id="world", child_frame_id="camera_late", ts=20.0), ) with SqliteStore(path=dataset) as store: @@ -293,7 +333,7 @@ def test_load_selects_camera_info_for_each_image_timestamp( for image, timestamp in zip(images, (10.0, 20.0), strict=True): image_stream.append(image, ts=timestamp) lidar_stream.append( - PointCloud2.from_numpy( + _cloud( np.array([[0.0, 0.0, 1.0]]), frame_id="world", timestamp=timestamp, @@ -303,14 +343,18 @@ def test_load_selects_camera_info_for_each_image_timestamp( camera_info_stream.append(camera_infos[0], ts=9.0) camera_info_stream.append(camera_infos[1], ts=19.0) for transform in camera_transforms: - tf_stream.append(TFMessage(transform), ts=transform.ts) + tf_stream.append( + TFMessage(transforms=[transform]), ts=to_seconds(transform.header.stamp) + ) loader = PointCloudFrameLoader(dataset) calibrations: list[CameraInfo] = [] def rectify(source: Image, calibration: CameraInfo) -> tuple[Image, CameraInfo]: calibrations.append(calibration) - return source, calibration.with_ts(source.ts) + result = CameraInfo.decode(calibration.encode()) + result.header.stamp = source.header.stamp + return source, result monkeypatch.setattr(loader._rectifier, "rectify", rectify) with loader: @@ -318,5 +362,5 @@ def rectify(source: Image, calibration: CameraInfo) -> tuple[Image, CameraInfo]: late = loader.load(1) assert calibrations == list(camera_infos) - assert early.camera_info.frame_id == "camera_early" - assert late.camera_info.frame_id == "camera_late" + assert early.camera_info.header.frame_id == "camera_early" + assert late.camera_info.header.frame_id == "camera_late" diff --git a/dimos/experimental/isolated_python/example/contract.py b/dimos/experimental/isolated_python/example/contract.py index b8b4a664011..95ed29bf901 100644 --- a/dimos/experimental/isolated_python/example/contract.py +++ b/dimos/experimental/isolated_python/example/contract.py @@ -16,6 +16,8 @@ from typing import Protocol +from dimos_generated.std_msgs.msg import Int32 + from dimos.agents.annotation import skill from dimos.core.core import rpc from dimos.core.stream import In, Out @@ -23,7 +25,6 @@ IsolatedPythonModule, IsolatedPythonModuleConfig, ) -from dimos.msgs.std_msgs.Int32 import Int32 from dimos.spec.utils import Spec diff --git a/dimos/experimental/isolated_python/test_end_to_end.py b/dimos/experimental/isolated_python/test_end_to_end.py index 602b0e5213a..0ccf7d889bd 100644 --- a/dimos/experimental/isolated_python/test_end_to_end.py +++ b/dimos/experimental/isolated_python/test_end_to_end.py @@ -16,6 +16,8 @@ import subprocess import sys +from dimos_generated.std_msgs.msg import Int32 + from dimos.core.coordination.blueprints import autoconnect from dimos.core.coordination.module_coordinator import ModuleCoordinator from dimos.core.core import rpc @@ -23,7 +25,6 @@ from dimos.core.stream import In, Out from dimos.experimental.isolated_python.example.contract import ExampleExternal from dimos.experimental.isolated_python.example.support import Offset -from dimos.msgs.std_msgs.Int32 import Int32 from dimos.utils.testing.waiting import wait_until @@ -32,7 +33,7 @@ class Producer(Module): @rpc def publish(self, value: int) -> None: - self.value.publish(Int32(value)) + self.value.publish(Int32(data=value)) class Consumer(Module): diff --git a/dimos/experimental/isolated_python/test_generated_example.py b/dimos/experimental/isolated_python/test_generated_example.py new file mode 100644 index 00000000000..53c59e8a3eb --- /dev/null +++ b/dimos/experimental/isolated_python/test_generated_example.py @@ -0,0 +1,42 @@ +# 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. + +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path + +from dimos_generated.std_msgs.msg import Int32 + +from dimos.experimental.isolated_python.bootstrap import load_class +from dimos.experimental.isolated_python.example.contract import ExampleExternal + + +def test_external_example_uses_generated_contract_and_cdr(monkeypatch): + project = Path(__file__).parents[3] / "native/python/example" + monkeypatch.syspath_prepend(str(project)) + runtime_type = load_class("example_external.runtime:ExampleExternalRuntime") + runtime = runtime_type(_isolated_python_runtime=True, initial_multiplier=3) + received = [] + monkeypatch.setattr(runtime.doubled, "publish", received.append) + try: + streams = {s.name: s.type for s in ExampleExternal.blueprint().blueprints[0].streams} + assert streams["value"] is Int32 and streams["doubled"] is Int32 + runtime._publish_doubled(Int32.decode(Int32(data=-7).encode())) + assert len(received) == 1 + assert type(received[0]) is Int32 + assert Int32.decode(received[0].encode()).data == -21 + assert runtime.get_multiplier() == 3 + finally: + runtime.stop() diff --git a/dimos/experimental/memory/rust/Cargo.toml b/dimos/experimental/memory/rust/Cargo.toml index c0655dc8a4e..2a7322024bb 100644 --- a/dimos/experimental/memory/rust/Cargo.toml +++ b/dimos/experimental/memory/rust/Cargo.toml @@ -10,7 +10,7 @@ anyhow = "1" crossbeam-channel = "0.5" crossbeam-utils = "0.8" dimos-module = { path = "../../../../native/rust/dimos-module" } -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +dimos-generated-messages = { path = "../../../message_codegen" } lz4_flex = "0.14.0" mcap = { version = "0.25.0", default-features = false, features = ["zstd"] } rayon = "1" @@ -19,7 +19,6 @@ serde = { version = "1", features = ["derive"] } serde_json = "1" tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal"] } tracing = "0.1" -turbojpeg = { version = "1.5.1", default-features = false, features = ["cmake"] } validator = { version = "0.21", features = ["derive"] } [dev-dependencies] diff --git a/dimos/experimental/memory/rust/src/store/mod.rs b/dimos/experimental/memory/rust/src/store/mod.rs index 78c08a53e28..d0371887b64 100644 --- a/dimos/experimental/memory/rust/src/store/mod.rs +++ b/dimos/experimental/memory/rust/src/store/mod.rs @@ -41,8 +41,8 @@ impl RecordingStoreConfig { /// One ordered observation ready for durable storage. pub struct Observation { pub stream: Arc, - pub source_ts: f64, - pub reception_ts: f64, + pub source_ts: i64, + pub reception_ts: i64, pub data: Vec, } diff --git a/dimos/experimental/memory/rust/src/store/sqlite.rs b/dimos/experimental/memory/rust/src/store/sqlite.rs index 813d0451c0b..059ab59bbd9 100644 --- a/dimos/experimental/memory/rust/src/store/sqlite.rs +++ b/dimos/experimental/memory/rust/src/store/sqlite.rs @@ -90,13 +90,15 @@ fn insert_observation(connection: &Connection, observation: &Observation) -> Res if observation.stream.is_tf() { connection.execute( &format!(r#"INSERT INTO "{name}" (ts) VALUES (?1)"#), - params![observation.source_ts], + params![observation.source_ts as f64 / 1_000_000_000.0], )?; } else { - let tags = serde_json::json!({"reception_ts": observation.reception_ts}).to_string(); + let tags = + serde_json::json!({"reception_ts": observation.reception_ts as f64 / 1_000_000_000.0}) + .to_string(); connection.execute( &format!(r#"INSERT INTO "{name}" (ts, tags) VALUES (?1, jsonb(?2))"#), - params![observation.source_ts, tags], + params![observation.source_ts as f64 / 1_000_000_000.0, tags], )?; } let id = connection.last_insert_rowid(); diff --git a/dimos/experimental/memory/rust_cli_recorder.py b/dimos/experimental/memory/rust_cli_recorder.py index 2214635c4e2..3aad51177d6 100644 --- a/dimos/experimental/memory/rust_cli_recorder.py +++ b/dimos/experimental/memory/rust_cli_recorder.py @@ -33,7 +33,6 @@ from dimos.experimental.memory.rust_recorder import RustRecorderConfig, RustStreamSpec from dimos.memory.store.sqlite import SqliteStore from dimos.memory.tap import matching, recording_dir -from dimos.msgs.sensor_msgs.Image import Image from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -103,9 +102,11 @@ def make_plan(transports: dict[tuple[str, type], Any]) -> RustRecordingPlan: for index, ((name, payload_type), transport) in enumerate(transports.items()): if name not in selected: continue - if not hasattr(payload_type, "lcm_encode") or not hasattr(payload_type, "lcm_decode"): + if not all( + hasattr(payload_type, member) for member in ("encode", "decode", "msg_name", "schema") + ): logger.info( - "--record: stream is not a DimOS LCM message type; skipped", + "--record: stream is not a generated CDR message type; skipped", stream=name, payload_type=str(payload_type), ) @@ -127,11 +128,11 @@ def make_plan(transports: dict[tuple[str, type], Any]) -> RustRecordingPlan: backends.add(backend) port = f"stream_{index}" streams.append( - RustStreamSpec( + RustStreamSpec.from_type( port=port, name=name, - payload_type=f"{payload_type.__module__}.{payload_type.__qualname__}", - codec="jpeg" if issubclass(payload_type, Image) else "lcm", + payload_type=payload_type, + codec="cdr", ) ) topics[port] = transport.channel diff --git a/dimos/experimental/memory/test_rust_cli_recorder.py b/dimos/experimental/memory/test_rust_cli_recorder.py index 63dd5c1f789..fd12739be26 100644 --- a/dimos/experimental/memory/test_rust_cli_recorder.py +++ b/dimos/experimental/memory/test_rust_cli_recorder.py @@ -21,6 +21,8 @@ from types import SimpleNamespace from typing import Any +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import Image import pytest from pytest_mock import MockerFixture @@ -30,8 +32,6 @@ from dimos.experimental.memory.rust_cli_recorder import RustRecordingPlan, RustRecordingSession from dimos.experimental.memory.rust_recorder import RustStreamSpec from dimos.memory.store.sqlite import SqliteStore -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.Image import Image from dimos.protocol.pubsub.impl.zenohpubsub import Topic as ZenohTopic @@ -67,8 +67,8 @@ def test_plan_uses_actual_lcm_channels_and_memory_codecs(tmp_path: Path) -> None assert plan.backend == "lcm" assert plan.topics == {"stream_0": odom.channel, "stream_1": camera.channel} assert [(stream.name, stream.codec) for stream in plan.streams] == [ - ("odom", "lcm"), - ("color_image", "jpeg"), + ("odom", "cdr"), + ("color_image", "cdr"), ] assert plan.path == tmp_path / "memory.db" @@ -180,11 +180,11 @@ def _plan(tmp_path: Path) -> RustRecordingPlan: backend="lcm", topics={"stream_0": "/odom"}, streams=[ - RustStreamSpec( + RustStreamSpec.from_type( port="stream_0", name="odom", - payload_type="dimos.msgs.geometry_msgs.PoseStamped.PoseStamped", - codec="lcm", + payload_type=PoseStamped, + codec="cdr", ) ], payload_types={"odom": PoseStamped}, diff --git a/dimos/experimental/robot/bosdyn/spot/effectors/high_level.py b/dimos/experimental/robot/bosdyn/spot/effectors/high_level.py index b713e1cdd1d..a465b40124c 100644 --- a/dimos/experimental/robot/bosdyn/spot/effectors/high_level.py +++ b/dimos/experimental/robot/bosdyn/spot/effectors/high_level.py @@ -43,6 +43,21 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseWithCovariance, + Quaternion, + TransformStamped, + Twist, + TwistWithCovariance, + Vector3, +) +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage + from dimos.agents.annotation import skill from dimos.core.core import rpc from dimos.core.stream import In, Out @@ -71,15 +86,8 @@ rotate_camera_info_quarter_turns, rotate_image_quarter_turns, ) -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.geometry import transform_from_odometry +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.protocol.tf.static_tf_publisher import StaticTfPublisher, StaticTfPublisherConfig from dimos.utils.logging_config import setup_logger @@ -168,7 +176,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: # never reaches a half-initialised SDK. self._ready = asyncio.Event() - def transforms(self) -> list[Transform]: + def transforms(self) -> list[TransformStamped]: """Static base_link -> camera-optical extrinsics parsed from the URDF. `StaticTfPublisher` republishes these on a fixed interval; the moving @@ -410,10 +418,12 @@ async def _poll_images(self) -> None: image = decode_image(response, frame_id, time_converter) if image is None: continue - if last_published_ts.get(source_name) == image.ts: + if last_published_ts.get(source_name) == to_seconds(image.header.stamp): continue - last_published_ts[source_name] = image.ts - camera_info = camera_info_from_response(response, frame_id, image.ts) + last_published_ts[source_name] = to_seconds(image.header.stamp) + camera_info = camera_info_from_response( + response, frame_id, to_seconds(image.header.stamp) + ) if quarter_turns: image = rotate_image_quarter_turns(image, quarter_turns) if camera_info is not None: @@ -466,44 +476,26 @@ async def _poll_odom(self) -> None: def _publish_odom(self, vision_tform_body: Any, velocity: Any, ts: float) -> None: pose = Pose( - position=[vision_tform_body.x, vision_tform_body.y, vision_tform_body.z], - orientation=[ - vision_tform_body.rot.x, - vision_tform_body.rot.y, - vision_tform_body.rot.z, - vision_tform_body.rot.w, - ], + position=Point(x=vision_tform_body.x, y=vision_tform_body.y, z=vision_tform_body.z), + orientation=Quaternion( + x=vision_tform_body.rot.x, + y=vision_tform_body.rot.y, + z=vision_tform_body.rot.z, + w=vision_tform_body.rot.w, + ), ) twist = Twist( - linear=[velocity.linear.x, velocity.linear.y, velocity.linear.z], - angular=[velocity.angular.x, velocity.angular.y, velocity.angular.z], + linear=Vector3(x=velocity.linear.x, y=velocity.linear.y, z=velocity.linear.z), + angular=Vector3(x=velocity.angular.x, y=velocity.angular.y, z=velocity.angular.z), ) odometry = Odometry( - ts=ts, - frame_id=self.config.odom_frame_id, + header=Header(stamp=time_from_seconds(ts), frame_id=self.config.odom_frame_id), child_frame_id=self.config.base_frame_id, - pose=pose, - twist=twist, + pose=PoseWithCovariance(pose=pose), + twist=TwistWithCovariance(twist=twist), ) self.odometry.publish(odometry) - self.tf.publish( - TFMessage( - Transform( - translation=Vector3( - vision_tform_body.x, vision_tform_body.y, vision_tform_body.z - ), - rotation=Quaternion( - vision_tform_body.rot.x, - vision_tform_body.rot.y, - vision_tform_body.rot.z, - vision_tform_body.rot.w, - ), - frame_id=self.config.odom_frame_id, - child_frame_id=self.config.base_frame_id, - ts=ts, - ) - ) - ) + self.tf.publish(TFMessage(transforms=[transform_from_odometry(odometry)])) @rpc async def move(self, twist: Twist, duration: float = 0.0) -> bool: @@ -569,7 +561,7 @@ async def move_velocity( yaw: Rotational velocity (rad/s). duration: Seconds to move. 0 uses one `cmd_vel_timeout` window. """ - twist = Twist(linear=Vector3(x, y, 0), angular=Vector3(0, 0, yaw)) + twist = Twist(linear=Vector3(x=x, y=y), angular=Vector3(z=yaw)) if await self.move(twist, duration=duration): return f"Moving with velocity=({x}, {y}, {yaw}) for {duration} seconds" return f"Failed to move with velocity=({x}, {y}, {yaw})" diff --git a/dimos/experimental/robot/bosdyn/spot/effectors/test_high_level.py b/dimos/experimental/robot/bosdyn/spot/effectors/test_high_level.py new file mode 100644 index 00000000000..fe8905fbc63 --- /dev/null +++ b/dimos/experimental/robot/bosdyn/spot/effectors/test_high_level.py @@ -0,0 +1,65 @@ +# 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. + +"""The Spot SDK boundary publishes generated odometry and its matching TF edge.""" + +from types import SimpleNamespace + +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.tf2_msgs.msg import TFMessage +import pytest + +from dimos.experimental.robot.bosdyn.spot.effectors.high_level import SpotHighLevel + + +@pytest.fixture +def module(): + instance = SpotHighLevel(odom_frame_id="vision", base_frame_id="body") + try: + yield instance + finally: + instance.stop() + + +def test_odometry_and_tf_preserve_sdk_pose_velocity_and_source_time(module, mocker): + odometry_publish = mocker.patch.object(module.odometry, "publish") + tf_publish = mocker.patch.object(module.tf, "publish") + pose = SimpleNamespace(x=1.0, y=2.0, z=3.0, rot=SimpleNamespace(x=0, y=0, z=0.6, w=0.8)) + velocity = SimpleNamespace( + linear=SimpleNamespace(x=4, y=5, z=6), angular=SimpleNamespace(x=7, y=8, z=9) + ) + + module._publish_odom(pose, velocity, -0.5) + + odometry_publish.assert_called_once() + tf_publish.assert_called_once() + odometry = Odometry.decode(odometry_publish.call_args.args[0].encode()) + tf = TFMessage.decode(tf_publish.call_args.args[0].encode()) + assert odometry.header.frame_id == "vision" + assert odometry.child_frame_id == "body" + assert (odometry.header.stamp.sec, odometry.header.stamp.nanosec) == (-1, 500000000) + assert ( + odometry.pose.pose.position.x, + odometry.pose.pose.position.y, + odometry.pose.pose.position.z, + ) == (1, 2, 3) + assert odometry.pose.pose.orientation.z == 0.6 + assert odometry.twist.twist.linear.y == 5 + assert odometry.twist.twist.angular.z == 9 + assert len(tf.transforms) == 1 + edge = tf.transforms[0] + assert edge.header == odometry.header + assert edge.child_frame_id == odometry.child_frame_id + assert edge.transform.translation.x == odometry.pose.pose.position.x + assert edge.transform.rotation == odometry.pose.pose.orientation diff --git a/dimos/experimental/robot/bosdyn/spot/replay.py b/dimos/experimental/robot/bosdyn/spot/replay.py index e3da31702b8..14d1e45b92f 100644 --- a/dimos/experimental/robot/bosdyn/spot/replay.py +++ b/dimos/experimental/robot/bosdyn/spot/replay.py @@ -27,9 +27,13 @@ from __future__ import annotations from collections.abc import AsyncIterator -import math from pathlib import Path +from dimos_generated.geometry_msgs.msg import TransformStamped +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.tf2_msgs.msg import TFMessage + from dimos.constants import RECORDINGS_DIR from dimos.core.module import Module, ModuleConfig from dimos.core.stream import Out @@ -37,15 +41,9 @@ CAMERA_STREAM_SUFFIXES, FRONT_CAMERA_ROTATE_UPRIGHT, ) +from dimos.experimental.robot.bosdyn.spot.utils import roll_optical_frame from dimos.memory.replay import resolve_db_path from dimos.memory.store.sqlite import SqliteStore -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -122,23 +120,18 @@ def _resolve_db_path(self) -> Path: def _republish_tf(self, message: TFMessage) -> None: self.tf.publish( - TFMessage(*(self._roll_optical_frame(transform) for transform in message.transforms)) + TFMessage( + transforms=[self._roll_optical_frame(transform) for transform in message.transforms] + ) ) - def _roll_optical_frame(self, transform: Transform) -> Transform: + def _roll_optical_frame(self, transform: TransformStamped) -> TransformStamped: if not self.config.roll_front_frames: return transform turns = _OPTICAL_FRAME_ROLL_TURNS.get(transform.child_frame_id) if not turns: return transform - roll = Quaternion.from_euler(Vector3(0.0, 0.0, turns * math.pi / 2)) - return Transform( - translation=transform.translation, - rotation=transform.rotation * roll, - frame_id=transform.frame_id, - child_frame_id=transform.child_frame_id, - ts=transform.ts, - ) + return roll_optical_frame(transform, turns) async def main(self) -> AsyncIterator[None]: db_path = self._resolve_db_path() diff --git a/dimos/experimental/robot/bosdyn/spot/test_utils.py b/dimos/experimental/robot/bosdyn/spot/test_utils.py index 2be357127bc..4acbe37c943 100644 --- a/dimos/experimental/robot/bosdyn/spot/test_utils.py +++ b/dimos/experimental/robot/bosdyn/spot/test_utils.py @@ -12,9 +12,30 @@ # See the License for the specific language governing permissions and # limitations under the License. +import math from pathlib import Path +import sys +from types import ModuleType, SimpleNamespace -from dimos.experimental.robot.bosdyn.spot.utils import camera_mount_transforms +import cv2 +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest + +from dimos.experimental.robot.bosdyn.spot.utils import ( + camera_info_from_response, + camera_mount_transforms, + decode_image, + roll_optical_frame, + rotate_camera_info_quarter_turns, + rotate_image_quarter_turns, +) +from dimos.msgs.camera_info import camera_info_from_intrinsics +from dimos.msgs.geometry import yaw +from dimos.msgs.image import image_from_array, image_view def test_camera_mount_transforms_uses_loaded_robot_topology(tmp_path: Path) -> None: @@ -42,6 +63,123 @@ def test_camera_mount_transforms_uses_loaded_robot_topology(tmp_path: Path) -> N transforms = camera_mount_transforms(urdf, "body", ["camera_optical"]) assert len(transforms) == 1 - assert transforms[0].frame_id == "body" + assert transforms[0].header.frame_id == "body" assert transforms[0].child_frame_id == "camera_optical" - assert transforms[0].translation.to_list() == [1.0, 2.0, 0.0] + translation = transforms[0].transform.translation + assert (translation.x, translation.y, translation.z) == (1.0, 2.0, 0.0) + + +@pytest.mark.parametrize("turns", [-1, 0, 1]) +def test_optical_roll_preserves_mount_and_exact_stamp(turns): + edge = TransformStamped( + header=Header(frame_id="body", stamp=Time(sec=1700000000, nanosec=123456789)), + child_frame_id="optical", + transform=Transform(translation=Vector3(x=1, y=2, z=3)), + ) + rolled = roll_optical_frame(edge, turns) + assert rolled.header == edge.header + assert rolled.child_frame_id == edge.child_frame_id + assert rolled.transform.translation == edge.transform.translation + assert yaw(rolled.transform.rotation) == pytest.approx(turns * math.pi / 2) + assert edge.transform.rotation.w == 1 + + +@pytest.fixture +def image_sdk(monkeypatch): + constants = SimpleNamespace( + FORMAT_JPEG=1, + FORMAT_RAW=2, + PIXEL_FORMAT_GREYSCALE_U8=1, + PIXEL_FORMAT_RGB_U8=3, + PIXEL_FORMAT_RGBA_U8=4, + PIXEL_FORMAT_DEPTH_U16=5, + PIXEL_FORMAT_GREYSCALE_U16=6, + ) + root = ModuleType("bosdyn") + root.__path__ = [] + api = ModuleType("bosdyn.api") + api.image_pb2 = SimpleNamespace(Image=constants) + monkeypatch.setitem(sys.modules, "bosdyn", root) + monkeypatch.setitem(sys.modules, "bosdyn.api", api) + return constants + + +@pytest.mark.parametrize( + "encoding,pixel_format,dtype,channels", + [ + ("mono8", 1, np.uint8, 1), + ("rgb8", 3, np.uint8, 3), + ("rgba8", 4, np.uint8, 4), + ("16UC1", 5, np.uint16, 1), + ("mono16", 6, np.uint16, 1), + ], +) +def test_sdk_raw_camera_boundary_emits_generated_pixels( + image_sdk, encoding, pixel_format, dtype, channels +): + shape = (2, 3) if channels == 1 else (2, 3, channels) + pixels = np.arange(np.prod(shape), dtype=dtype).reshape(shape) + shot = SimpleNamespace( + format=image_sdk.FORMAT_RAW, + pixel_format=pixel_format, + rows=2, + cols=3, + data=pixels.tobytes(), + ) + response = SimpleNamespace(shot=SimpleNamespace(image=shot, acquisition_time="robot-clock")) + clock = SimpleNamespace(local_seconds_from_robot_timestamp=lambda stamp: 1700000000.25) + value = decode_image(response, "camera-optical", clock) + assert type(value) is Image and value.encoding == encoding + decoded = Image.decode(value.encode()) + assert decoded.header.frame_id == "camera-optical" + assert decoded.header.stamp == Time(sec=1700000000, nanosec=250000000) + np.testing.assert_array_equal(image_view(decoded), pixels) + + +def test_sdk_jpeg_camera_boundary_decodes_generated_image(image_sdk): + pixels = np.full((8, 8, 3), [20, 80, 140], dtype=np.uint8) + ok, jpeg = cv2.imencode(".jpg", pixels) + assert ok + shot = SimpleNamespace(format=image_sdk.FORMAT_JPEG, pixel_format=3, data=jpeg.tobytes()) + response = SimpleNamespace(shot=SimpleNamespace(image=shot, acquisition_time="robot-clock")) + clock = SimpleNamespace(local_seconds_from_robot_timestamp=lambda stamp: 1700000000.25) + value = decode_image(response, "camera-optical", clock) + assert type(value) is Image and value.encoding == "bgr8" + np.testing.assert_allclose(image_view(Image.decode(value.encode())), pixels, atol=2) + + +@pytest.mark.parametrize("turns", [-1, 0, 1, 2, 4]) +def test_generated_camera_rotation_preserves_depth_header_and_intrinsics(turns): + header = Header(frame_id="optical", stamp=Time(sec=-1, nanosec=123456789)) + pixels = np.arange(6, dtype=np.uint16).reshape(2, 3) * 1000 + image = image_from_array(pixels, encoding="16UC1", header=header) + info = camera_info_from_intrinsics(10, 20, 1, 0.5, 3, 2, header=header) + before_image, before_info = image.encode(), info.encode() + rotated = rotate_image_quarter_turns(image, turns) + calibration = rotate_camera_info_quarter_turns(info, turns) + np.testing.assert_array_equal(image_view(rotated), np.rot90(pixels, k=turns)) + assert rotated.header == calibration.header == header + assert (calibration.height, calibration.width) == image_view(rotated).shape + expected_fx, expected_fy, cx, cy, width, height = 10, 20, 1, 0.5, 3, 2 + for _ in range(turns % 4): + expected_fx, expected_fy = expected_fy, expected_fx + cx, cy = cy, width - 1 - cx + width, height = height, width + assert list(calibration.k) == [expected_fx, 0, cx, 0, expected_fy, cy, 0, 0, 1] + assert image.encode() == before_image and info.encode() == before_info + + +def test_sdk_calibration_boundary_emits_plain_generated_value(): + source = SimpleNamespace( + HasField=lambda name: name == "pinhole", + cols=6, + rows=4, + pinhole=SimpleNamespace( + intrinsics=SimpleNamespace( + focal_length=SimpleNamespace(x=10, y=20), principal_point=SimpleNamespace(x=3, y=2) + ) + ), + ) + value = camera_info_from_response(SimpleNamespace(source=source), "optical", 1700000000.25) + assert list(value.k) == [10, 0, 3, 0, 20, 2, 0, 0, 1] + assert value.header.stamp == Time(sec=1700000000, nanosec=250000000) diff --git a/dimos/experimental/robot/bosdyn/spot/utils.py b/dimos/experimental/robot/bosdyn/spot/utils.py index d5e6f77fab4..30e5d7c80b3 100644 --- a/dimos/experimental/robot/bosdyn/spot/utils.py +++ b/dimos/experimental/robot/bosdyn/spot/utils.py @@ -20,13 +20,15 @@ from pathlib import Path from typing import Any +from dimos_generated.geometry_msgs.msg import Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.std_msgs.msg import Header import numpy as np -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.camera_info import camera_info_from_intrinsics +from dimos.msgs.geometry import compose_transforms, quaternion_from_euler +from dimos.msgs.image import image_from_array, image_view +from dimos.msgs.time import time_from_seconds from dimos.robot.assets.model import JointDescription, RobotModel from dimos.utils.logging_config import setup_logger @@ -56,8 +58,12 @@ def decode_image(response: Any, frame_id: str, time_converter: Any) -> Image | N if decoded is None: logger.error(f"Failed to decode JPEG image from {frame_id}") return None - image_format = ImageFormat.GRAY if decoded.ndim == 2 else ImageFormat.BGR - return Image.from_numpy(decoded, format=image_format, frame_id=frame_id, ts=ts) + image_format = "mono8" if decoded.ndim == 2 else "bgr8" + return image_from_array( + decoded, + encoding=image_format, + header=Header(frame_id=frame_id, stamp=time_from_seconds(ts)), + ) if shot.format != image_pb2.Image.FORMAT_RAW: logger.error(f"Unsupported Spot image encoding {shot.format} from {frame_id}") @@ -74,21 +80,27 @@ def decode_image(response: Any, frame_id: str, time_converter: Any) -> Image | N if channels == 1 else array.reshape(shot.rows, shot.cols, channels) ) - return Image.from_numpy(array, format=image_format, frame_id=frame_id, ts=ts) + return image_from_array( + array, encoding=image_format, header=Header(frame_id=frame_id, stamp=time_from_seconds(ts)) + ) -def joint_to_transform(joint: JointDescription) -> Transform: - return Transform( - translation=Vector3(*joint.origin_xyz), - rotation=Quaternion.from_euler(Vector3(*joint.origin_rpy)), - frame_id=joint.parent_link, +def joint_to_transform(joint: JointDescription) -> TransformStamped: + return TransformStamped( + header=Header(frame_id=joint.parent_link), child_frame_id=joint.child_link, + transform=Transform( + translation=Vector3( + x=joint.origin_xyz[0], y=joint.origin_xyz[1], z=joint.origin_xyz[2] + ), + rotation=quaternion_from_euler(*joint.origin_rpy), + ), ) def camera_mount_transforms( urdf_path: str | Path, base_frame_id: str, optical_frames: list[str] -) -> list[Transform]: +) -> list[TransformStamped]: """Compose each base_frame_id -> optical_frame extrinsic from the URDF's fixed joints. Walks the fixed-joint chain (base_link -> body -> {pos}_camera -> optical) up @@ -98,7 +110,7 @@ def camera_mount_transforms( model = RobotModel.from_file(urdf_path).load() urdf_root = model.root_link joint_by_child = {joint.child_link: joint for joint in model.joints} - transforms: list[Transform] = [] + transforms: list[TransformStamped] = [] for optical_frame in optical_frames: chain: list[JointDescription] = [] current = optical_frame @@ -112,13 +124,13 @@ def camera_mount_transforms( edges = [joint_to_transform(joint) for joint in reversed(chain)] composed = edges[0] for edge in edges[1:]: - composed = composed + edge - composed.frame_id = base_frame_id + composed = compose_transforms(composed, edge) + composed.header.frame_id = base_frame_id transforms.append(composed) return transforms -def roll_optical_frame(transform: Transform, quarter_turns: int) -> Transform: +def roll_optical_frame(transform: TransformStamped, quarter_turns: int) -> TransformStamped: """Roll a camera's optical frame `quarter_turns` * 90° about its viewing (z) axis. Pairs with `rotate_image_quarter_turns`: rotating the image alone leaves the 3D @@ -127,20 +139,18 @@ def roll_optical_frame(transform: Transform, quarter_turns: int) -> Transform: """ if not quarter_turns: return transform - roll = Quaternion.from_euler(Vector3(0.0, 0.0, quarter_turns * math.pi / 2)) - return Transform( - translation=transform.translation, - rotation=transform.rotation * roll, - frame_id=transform.frame_id, + roll = TransformStamped( + header=Header(frame_id=transform.child_frame_id), child_frame_id=transform.child_frame_id, - ts=transform.ts, + transform=Transform(rotation=quaternion_from_euler(0, 0, quarter_turns * math.pi / 2)), ) + return compose_transforms(transform, roll) def rotate_image_quarter_turns(image: Image, quarter_turns: int) -> Image: """Rotate an Image by `quarter_turns` * 90° CCW (negative for CW).""" - rotated = np.rot90(image.data, k=quarter_turns) - return Image.from_numpy(rotated, format=image.format, frame_id=image.frame_id, ts=image.ts) + rotated = np.rot90(image_view(image), k=quarter_turns) + return image_from_array(rotated, encoding=image.encoding, header=image.header) def rotate_camera_info_quarter_turns(info: CameraInfo, quarter_turns: int) -> CameraInfo: @@ -149,15 +159,15 @@ def rotate_camera_info_quarter_turns(info: CameraInfo, quarter_turns: int) -> Ca Each CCW quarter turn swaps the focal lengths and remaps the principal point so the intrinsics stay consistent with the rotated pixel grid (width/height swap). """ - fx, fy, cx, cy = info.K[0], info.K[4], info.K[2], info.K[5] + fx, fy, cx, cy = info.k[0], info.k[4], info.k[2], info.k[5] width, height = info.width, info.height for _ in range(quarter_turns % 4): fx, fy = fy, fx cx, cy = cy, (width - 1) - cx width, height = height, width - return CameraInfo.from_intrinsics( - fx=fx, fy=fy, cx=cx, cy=cy, width=width, height=height, frame_id=info.frame_id - ).with_ts(info.ts) + return camera_info_from_intrinsics( + fx=fx, fy=fy, cx=cx, cy=cy, width=width, height=height, header=info.header + ) def camera_info_from_response(response: Any, source_name: str, ts: float) -> CameraInfo | None: @@ -166,30 +176,30 @@ def camera_info_from_response(response: Any, source_name: str, ts: float) -> Cam if not source.HasField("pinhole"): return None intrinsics = source.pinhole.intrinsics - info = CameraInfo.from_intrinsics( + info = camera_info_from_intrinsics( fx=intrinsics.focal_length.x, fy=intrinsics.focal_length.y, cx=intrinsics.principal_point.x, cy=intrinsics.principal_point.y, width=source.cols, height=source.rows, - frame_id=source_name, + header=Header(frame_id=source_name, stamp=time_from_seconds(ts)), ) - return info.with_ts(ts) + return info def clamp(value: float, low: float, high: float) -> float: return max(low, min(high, value)) -def raw_layout(pixel_format: int) -> tuple[Any, int, ImageFormat]: +def raw_layout(pixel_format: int) -> tuple[Any, int, str]: from bosdyn.api import image_pb2 # type: ignore[import-not-found] - layouts: dict[int, tuple[Any, int, ImageFormat]] = { - image_pb2.Image.PIXEL_FORMAT_GREYSCALE_U8: (np.uint8, 1, ImageFormat.GRAY), - image_pb2.Image.PIXEL_FORMAT_GREYSCALE_U16: (np.uint16, 1, ImageFormat.GRAY16), - image_pb2.Image.PIXEL_FORMAT_DEPTH_U16: (np.uint16, 1, ImageFormat.DEPTH16), - image_pb2.Image.PIXEL_FORMAT_RGB_U8: (np.uint8, 3, ImageFormat.RGB), - image_pb2.Image.PIXEL_FORMAT_RGBA_U8: (np.uint8, 4, ImageFormat.RGBA), + layouts: dict[int, tuple[Any, int, str]] = { + image_pb2.Image.PIXEL_FORMAT_GREYSCALE_U8: (np.uint8, 1, "mono8"), + image_pb2.Image.PIXEL_FORMAT_GREYSCALE_U16: (np.uint16, 1, "mono16"), + image_pb2.Image.PIXEL_FORMAT_DEPTH_U16: (np.uint16, 1, "16UC1"), + image_pb2.Image.PIXEL_FORMAT_RGB_U8: (np.uint8, 3, "rgb8"), + image_pb2.Image.PIXEL_FORMAT_RGBA_U8: (np.uint8, 4, "rgba8"), } - return layouts.get(pixel_format, (None, 0, ImageFormat.GRAY)) + return layouts.get(pixel_format, (None, 0, "mono8")) diff --git a/dimos/experimental/security_demo/conftest.py b/dimos/experimental/security_demo/conftest.py index 364a90e6d32..6e7cc79d682 100644 --- a/dimos/experimental/security_demo/conftest.py +++ b/dimos/experimental/security_demo/conftest.py @@ -14,11 +14,13 @@ from __future__ import annotations +from dimos_generated.sensor_msgs.msg import CameraInfo +import numpy as np +from PIL import Image as PILImage import pytest from dimos.experimental.security_demo.security_module import SecurityModule -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_from_array from dimos.perception.detection.detectors.yolo import Yolo2DDetector from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.utils.data import get_data @@ -33,12 +35,18 @@ def yolo_detector(): @pytest.fixture(scope="session") def person_image(): - return Image.from_file(get_data("security_detection.png")) + return image_from_array( + np.asarray(PILImage.open(get_data("security_detection.png")).convert("RGB")), + encoding="rgb8", + ) @pytest.fixture(scope="session") def empty_image(): - return Image.from_file(get_data("security_no_detection.png")) + return image_from_array( + np.asarray(PILImage.open(get_data("security_no_detection.png")).convert("RGB")), + encoding="rgb8", + ) @pytest.fixture() diff --git a/dimos/experimental/security_demo/depth_estimator.py b/dimos/experimental/security_demo/depth_estimator.py index 588057d7112..3fc069f1fcb 100644 --- a/dimos/experimental/security_demo/depth_estimator.py +++ b/dimos/experimental/security_demo/depth_estimator.py @@ -17,12 +17,13 @@ from collections.abc import Callable import threading +from dimos_generated.sensor_msgs.msg import Image import numpy as np from PIL import Image as PILImage import torch from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_from_array, image_to_rgb _DEPTH_MODEL_NAME = "depth-anything/Depth-Anything-V2-Small-hf" _DEPTH_MAX_WIDTH = 640 @@ -86,8 +87,7 @@ def _loop(self) -> None: self._process(image) def _process(self, image: Image) -> None: - rgb = image.to_rgb() - pil_image = PILImage.fromarray(rgb.data) + pil_image = PILImage.fromarray(image_to_rgb(image)) if pil_image.width > _DEPTH_MAX_WIDTH: scale = _DEPTH_MAX_WIDTH / pil_image.width new_h = int(pil_image.height * scale) @@ -105,8 +105,4 @@ def _process(self, image: Image) -> None: ).squeeze() depth_np = depth.cpu().numpy().astype(np.float32) - self._publish( - Image.from_numpy( - depth_np, format=ImageFormat.DEPTH, frame_id=image.frame_id, ts=image.ts - ) - ) + self._publish(image_from_array(depth_np, encoding="32FC1", header=image.header)) diff --git a/dimos/experimental/security_demo/test_security_module.py b/dimos/experimental/security_demo/test_security_module.py index 307595402ef..fa308ff475a 100644 --- a/dimos/experimental/security_demo/test_security_module.py +++ b/dimos/experimental/security_demo/test_security_module.py @@ -16,10 +16,9 @@ import threading +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Twist, Vector3 import pytest -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D @@ -50,7 +49,9 @@ def test_patrol_step_transitions_to_following_on_detection( module._state = "PATROLLING" module._has_active_goal = True module._latest_image = person_image - module._latest_pose = PoseStamped(position=[1, 2, 0], orientation=[0, 0, 0, 1]) + module._latest_pose = PoseStamped( + pose=Pose(position=Point(x=1, y=2), orientation=Quaternion(w=1)) + ) det = make_detection() module._detector.process_image.return_value = ImageDetections2D( @@ -78,7 +79,7 @@ def test_patrol_step_requests_goal_when_no_active_goal(security_module): module._has_active_goal = False module._latest_image = None # causes early return after goal logic - goal = PoseStamped(position=[5, 5, 0], orientation=[0, 0, 0, 1]) + goal = PoseStamped(pose=Pose(position=Point(x=5, y=5), orientation=Quaternion(w=1))) module._router.next_goal.return_value = goal module._patrol_step() @@ -100,7 +101,7 @@ def test_follow_step_publishes_twist_when_tracking( image=person_image, detections=[det] ) - twist = Twist(linear=[0.3, 0, 0], angular=[0, 0, 0.1]) + twist = Twist(linear=Vector3(x=0.3), angular=Vector3(z=0.1)) module._visual_servo.compute_twist.return_value = twist mocker.patch("dimos.experimental.security_demo.security_module.time.sleep") @@ -125,7 +126,7 @@ def test_follow_step_transitions_to_patrolling_on_person_lost(security_module, p # Zero twist should be published to stop the robot published_twist = module.cmd_vel.publish.call_args[0][0] - assert published_twist.is_zero() + assert published_twist == Twist() module._speak_skill.speak.assert_called_with( "Lost sight of intruder, resuming patrol", blocking=False @@ -157,7 +158,7 @@ def patrol_side_effect(): # Verify zero twist published on shutdown last_twist = module.cmd_vel.publish.call_args[0][0] - assert last_twist.is_zero() + assert last_twist == Twist() # Verify state transitions: PATROLLING then IDLE state_values = [call.args[0].data for call in module.security_state.publish.call_args_list] diff --git a/dimos/experimental/world_belief/absence.py b/dimos/experimental/world_belief/absence.py index 8003d671346..2802ea3ced6 100644 --- a/dimos/experimental/world_belief/absence.py +++ b/dimos/experimental/world_belief/absence.py @@ -19,11 +19,12 @@ import numpy as np if TYPE_CHECKING: + from dimos_generated.geometry_msgs.msg import TransformStamped, Vector3 + from dimos_generated.sensor_msgs.msg import CameraInfo from numpy.typing import NDArray - from dimos.msgs.geometry_msgs.Transform import Transform - from dimos.msgs.geometry_msgs.Vector3 import Vector3 - from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo +from dimos.msgs.camera_info import intrinsic_matrix +from dimos.msgs.geometry import inverse_transform, transform_matrix OUT_OF_VIEW = "out_of_view" PRESENT = "present" @@ -38,20 +39,20 @@ def classify_visibility( center_world: Vector3, camera_info: CameraInfo, - world_from_camera: Transform, + world_from_camera: TransformStamped, depth_m: NDArray[np.float32], *, near_extent_m: float = 0.0, ) -> str: """Classify a world point as PRESENT, ABSENT, OCCLUDED, or OUT_OF_VIEW. ``near_extent_m`` widens the expected PRESENT depth band.""" - cam_from_world = world_from_camera.inverse() - p = cam_from_world.rotation.rotate_vector(center_world) + cam_from_world.translation - x, y, z = p.x, p.y, p.z + cam_from_world = inverse_transform(world_from_camera) + matrix = transform_matrix(cam_from_world.transform) + x, y, z = (matrix @ np.array([center_world.x, center_world.y, center_world.z, 1.0]))[:3] if z <= _NEAR_M: # Too close to camera; sensor is unreliable return OUT_OF_VIEW - k = camera_info.get_K_matrix() + k = intrinsic_matrix(camera_info) u = k[0, 0] * x / z + k[0, 2] v = k[1, 1] * y / z + k[1, 2] h, w = depth_m.shape[:2] diff --git a/dimos/experimental/world_belief/recall.py b/dimos/experimental/world_belief/recall.py index 7d989b0c6bb..450a4a5482e 100644 --- a/dimos/experimental/world_belief/recall.py +++ b/dimos/experimental/world_belief/recall.py @@ -23,8 +23,10 @@ import re from typing import TYPE_CHECKING, Any, cast +from dimos_generated.sensor_msgs.msg import Image + from dimos.memory.embed import EmbedImages -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_brightness, image_resize_to_fit, image_sharpness from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -86,14 +88,14 @@ def best_in_new_buckets(upstream: Iterator[Any]) -> Iterator[Any]: if bucket in indexed_buckets: continue indexed_buckets.add(bucket) - yield max(observations, key=lambda obs: obs.data.sharpness) + yield max(observations, key=lambda obs: image_sharpness(obs.data)) pipeline = ( # Skip pose-less frames — they can't answer "where" (world-frame frames need no pose). src_stream.filter( - lambda obs: obs.pose is not None or (obs.data.frame_id or "") in ("", "world") + lambda obs: obs.pose is not None or (obs.data.header.frame_id or "") in ("", "world") ) - .filter(lambda obs: obs.data.brightness > 0.1) + .filter(lambda obs: image_brightness(obs.data) > 0.1) .transform(best_in_new_buckets) .transform(EmbedImages(model)) ) @@ -102,7 +104,7 @@ def best_in_new_buckets(upstream: Iterator[Any]) -> Iterator[Any]: embedded = cast("EmbeddedObservation[Image]", obs) payload = obs.data if index_store is not None: - payload, _scale = payload.resize_to_fit(thumbnail_px, thumbnail_px) + payload, _scale = image_resize_to_fit(payload, thumbnail_px, thumbnail_px) index.append( payload, ts=obs.ts, diff --git a/dimos/experimental/world_belief/scene_scan.py b/dimos/experimental/world_belief/scene_scan.py index f386e32e145..9b38aafacb8 100644 --- a/dimos/experimental/world_belief/scene_scan.py +++ b/dimos/experimental/world_belief/scene_scan.py @@ -18,14 +18,18 @@ from dataclasses import dataclass import itertools +import math from typing import TYPE_CHECKING, Any +from dimos_generated.geometry_msgs.msg import PoseStamped, TransformStamped +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.std_msgs.msg import Header import numpy as np from dimos.experimental.world_belief.world_belief import WorldBelief -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.geometry import point_distance, quaternion_angle, transform_from_pose +from dimos.msgs.image import image_from_array, image_to_rgb, image_view +from dimos.msgs.time import time_from_seconds from dimos.perception.experimental.object import Object from dimos.utils.logging_config import setup_logger @@ -156,10 +160,7 @@ def warmup(self) -> None: set_prompts = getattr(detector, "set_prompts", None) if self._text_prompts and callable(set_prompts): set_prompts(text=list(self._text_prompts)) - dummy = Image( - data=np.zeros((64, 64, 3), dtype=np.uint8), - format=ImageFormat.RGB, - ) + dummy = image_from_array(np.zeros((64, 64, 3), dtype=np.uint8), encoding="rgb8") detector_warmup = getattr(detector, "warmup", None) if callable(detector_warmup): detector_warmup(dummy) @@ -237,7 +238,9 @@ def scan( last_source_ts = obs_ts if last_source_ts is None else max(last_source_ts, obs_ts) if obs.pose is None: if poseless_in_target is None: - frame_id = obs.data.frame_id or "" # one decode, first pose-less frame only + frame_id = ( + obs.data.header.frame_id or "" + ) # one decode, first pose-less frame only poseless_in_target = frame_id in ("", self._target_frame) if not poseless_in_target: skipped += 1 @@ -325,8 +328,8 @@ def _select_frames(self, frames: list[Any]) -> list[Any]: p0 is not None and p1 is not None and ( - p0.position.distance(p1.position) / dt > _TRANS_SPEED - or p0.orientation.angle_to(p1.orientation) / dt > _ROT_SPEED + point_distance(p0.position, p1.position) / dt > _TRANS_SPEED + or quaternion_angle(p0.orientation, p1.orientation) / dt > _ROT_SPEED ) ) period = 1.0 / (_MOVING_HZ if moving else self._stationary_hz) @@ -344,7 +347,7 @@ def _detect_frame( depth_stream: Any, depth_tolerance: float, detector: Detector, - ) -> tuple[list[Object], Any, Transform | None] | None: + ) -> tuple[list[Object], Any, TransformStamped | None] | None: """Detect + NMS + 2D→3D lift + embed one frame; None if no depth within tolerance.""" color_img: Image = obs.data depth_img = self._nearest_depth(depth_stream, obs.ts, depth_tolerance) @@ -352,12 +355,19 @@ def _detect_frame( return None process = getattr(detector, "predict_image", detector.process_image) detections: ImageDetections2D[Any] = process(color_img) - frame_id = color_img.frame_id or "" + frame_id = color_img.header.frame_id or "" in_target = frame_id in ("", self._target_frame) - camera_transform = None if in_target else Transform.from_pose(frame_id, obs.pose) - if camera_transform is not None: - camera_transform.frame_id = self._target_frame - camera_transform.ts = float(obs.ts) + camera_transform = None + if not in_target: + if obs.pose is None: + return None + camera_transform = transform_from_pose( + PoseStamped( + header=Header(frame_id=self._target_frame, stamp=time_from_seconds(obs.ts)), + pose=obs.pose, + ), + child_frame_id=frame_id, + ) objects = Object.from_2d_to_list( detections_2d=detections, color_image=color_img, @@ -366,7 +376,9 @@ def _detect_frame( camera_transform=camera_transform, ) for obj in objects: - obj.frame_id = obj.pose.frame_id = obj.pointcloud.frame_id = self._target_frame + obj.frame_id = obj.pose.header.frame_id = obj.pointcloud.header.frame_id = ( + self._target_frame + ) objects = [o for o in objects if float(o.confidence) >= self._detector_conf] # Remove overlapping prompts and same-depth contained regions; a contained object # at a different depth survives. @@ -377,7 +389,10 @@ def _detect_frame( _bbox_iou(o.bbox, k.bbox) > _NMS_IOU or ( _bbox_containment(o.bbox, k.bbox) >= _NMS_CONTAINMENT - and o.center.distance(k.center) <= _SUBREGION_COLOCATED_M + and math.dist( + (o.center.x, o.center.y, o.center.z), (k.center.x, k.center.y, k.center.z) + ) + <= _SUBREGION_COLOCATED_M ) for k in kept ): @@ -390,7 +405,7 @@ def _detect_frame( x1 <= 3 or y1 <= 3 or x2 >= color_img.width - 3 or y2 >= color_img.height - 3 ) self._attach_embeddings(kept, color_img) - depth_arr = np.asarray(depth_img.to_opencv(), dtype=np.float32) + depth_arr = np.asarray(image_view(depth_img), dtype=np.float32) return kept, depth_arr, camera_transform def _attach_embeddings(self, objects: list[Object], color_img: Image) -> None: @@ -401,11 +416,12 @@ def _attach_embeddings(self, objects: list[Object], color_img: Image) -> None: crops: list[Image] = [] idxs: list[int] = [] for i, obj in enumerate(objects): - crop = obj.cropped_image(padding=0).to_rgb() + crop_source = obj.cropped_image(padding=0) + crop = image_from_array( + image_to_rgb(crop_source), encoding="rgb8", header=color_img.header + ) if crop.width < 2 or crop.height < 2: continue - crop.frame_id = color_img.frame_id - crop.ts = color_img.ts crops.append(crop) idxs.append(i) if not crops: @@ -427,9 +443,9 @@ def _nearest_depth(depth_stream: Any, ts: float, tolerance: float) -> Image | No return None depth_obs = min(candidates, key=lambda o: abs(float(o.ts) - ts)) raw: Image = depth_obs.data - cv = raw.to_opencv() - if raw.format == ImageFormat.DEPTH16: + cv = image_view(raw) + if raw.encoding == "16UC1": cv = cv.astype("float32") / 1000.0 elif cv.dtype != "float32": cv = cv.astype("float32") - return Image(data=cv, format=ImageFormat.DEPTH, frame_id=raw.frame_id, ts=raw.ts) + return image_from_array(cv, encoding="32FC1", header=raw.header) diff --git a/dimos/experimental/world_belief/test_absence.py b/dimos/experimental/world_belief/test_absence.py new file mode 100644 index 00000000000..0f7172da35d --- /dev/null +++ b/dimos/experimental/world_belief/test_absence.py @@ -0,0 +1,63 @@ +# 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. + +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest + +from dimos.experimental.world_belief.absence import ( + ABSENT, + OCCLUDED, + OUT_OF_VIEW, + PRESENT, + classify_visibility, +) +from dimos.msgs.camera_info import camera_info_from_intrinsics + + +@pytest.mark.parametrize( + "depth,expected", [(3.0, ABSENT), (1.0, PRESENT), (0.5, OCCLUDED), (0.0, OCCLUDED)] +) +def test_visibility_inverts_generated_camera_transform(depth, expected): + calibration = camera_info_from_intrinsics( + 100, 100, 16, 16, 32, 32, header=Header(frame_id="camera") + ) + world_from_camera = TransformStamped( + header=Header(frame_id="world"), + child_frame_id="camera", + transform=Transform(translation=Vector3(x=10), rotation=Quaternion(w=1)), + ) + # In world coordinates x=10; after inverse transform it is on the optical axis. + center = Vector3(x=10, z=1) + depth_m = np.full((32, 32), depth, dtype=np.float32) + assert classify_visibility(center, calibration, world_from_camera, depth_m) == expected + assert ( + classify_visibility(Vector3(x=10, z=-1), calibration, world_from_camera, depth_m) + == OUT_OF_VIEW + ) + + +def test_visibility_requires_all_valid_depth_samples_to_show_absence(): + calibration = camera_info_from_intrinsics( + 100, 100, 16, 16, 32, 32, header=Header(frame_id="camera") + ) + identity = TransformStamped( + header=Header(frame_id="world"), + child_frame_id="camera", + transform=Transform(rotation=Quaternion(w=1)), + ) + depth = np.full((32, 32), 3, dtype=np.float32) + depth[16, 16] = 1 + assert classify_visibility(Vector3(z=1), calibration, identity, depth) == PRESENT diff --git a/dimos/experimental/world_belief/test_generated_recall.py b/dimos/experimental/world_belief/test_generated_recall.py new file mode 100644 index 00000000000..9b696b79c4b --- /dev/null +++ b/dimos/experimental/world_belief/test_generated_recall.py @@ -0,0 +1,90 @@ +# 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. + +from dimos_generated.geometry_msgs.msg import Pose, Quaternion +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest + +from dimos.experimental.world_belief.recall import build_frame_clip_index, index_stream_name +from dimos.memory.store.memory import MemoryStore +from dimos.msgs.image import image_from_array, image_resize_to_fit, image_view +from dimos.msgs.time import time_from_seconds + + +class _EmbeddingModel: + def __init__(self) -> None: + self.seen: list[Image] = [] + + def embed(self, *images: Image) -> list[np.ndarray]: + self.seen.extend(images) + return [np.array([1.0, 0.0], dtype=np.float32) for _ in images] + + +def test_index_selects_sharp_generated_frame_and_preserves_stamp_and_thumbnail() -> None: + source, destination = MemoryStore(), MemoryStore() + try: + frames = source.stream("color_image", Image) + sharp = np.tile(np.array([[64, 255], [255, 64]], dtype=np.uint8), (4, 8)) + sharp = np.repeat(sharp[:, :, None], 3, axis=2) + image = image_from_array( + sharp, encoding="rgb8", header=Header(stamp=time_from_seconds(0.2), frame_id="camera") + ) + identity = Pose(orientation=Quaternion(w=1.0)) + frames.append( + image_from_array(np.full_like(sharp, 150), encoding="rgb8", header=image.header), + ts=0.1, + pose=identity, + ) + frames.append(image, ts=0.2, pose=identity) + # A bright pose-less camera frame and a dark posed frame are ineligible. + frames.append(image, ts=1.1) + frames.append( + image_from_array(np.zeros_like(sharp), encoding="rgb8", header=image.header), + ts=2.1, + pose=identity, + ) + model = _EmbeddingModel() + assert ( + build_frame_clip_index( + source, model=model, index_store=destination, thumbnail_px=4, source_tag="test" + ) + == 1 + ) + assert len(model.seen) == 1 + assert np.array_equal(image_view(model.seen[0]), sharp) + result = destination.stream( + index_stream_name("openai/clip-vit-base-patch32"), Image + ).first() + assert (result.data.width, result.data.height) == (4, 2) + assert result.data.header == image.header + assert result.ts == 0.2 + assert result.tags["rec"] == "test" + assert ( + build_frame_clip_index( + source, model=model, index_store=destination, thumbnail_px=4, source_tag="test" + ) + == 0 + ) + finally: + source.stop() + destination.stop() + + +def test_generated_resize_keeps_small_images_and_rejects_invalid_dimensions() -> None: + image = image_from_array(np.ones((2, 3, 3), dtype=np.uint8), encoding="bgr8") + assert image_resize_to_fit(image, 5, 5) == (image, 1.0) + with pytest.raises(ValueError, match="positive"): + image_resize_to_fit(image, 0, 5) diff --git a/dimos/experimental/world_belief/world_belief.py b/dimos/experimental/world_belief/world_belief.py index dcbe3cf1e36..30d980069a3 100644 --- a/dimos/experimental/world_belief/world_belief.py +++ b/dimos/experimental/world_belief/world_belief.py @@ -22,6 +22,7 @@ from typing import TYPE_CHECKING, Any import uuid +from dimos_generated.geometry_msgs.msg import Vector3 import numpy as np from scipy.optimize import linear_sum_assignment @@ -36,16 +37,16 @@ normalize_embedding, ) from dimos.models.embedding.base import Embedding -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.geometry import pose_from_transform, pose_matrix +from dimos.msgs.pointcloud import pointcloud_xyz from dimos.protocol.service.spec import BaseConfig from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: + from dimos_generated.geometry_msgs.msg import Pose, TransformStamped + from dimos_generated.sensor_msgs.msg import CameraInfo from numpy.typing import NDArray - from dimos.msgs.geometry_msgs.Pose import Pose - from dimos.msgs.geometry_msgs.Transform import Transform - from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo from dimos.perception.experimental.object import Object logger = setup_logger() @@ -89,7 +90,10 @@ class WorldBeliefConfig(BaseConfig): def camera_viewpoint(pose: Pose) -> Viewpoint: """Return camera position and forward axis for viewpoint diversity.""" - return pose.position, pose.orientation.rotate_vector(Vector3.unit_z()) + forward = pose_matrix(pose)[:3, 2] + return Vector3(x=pose.position.x, y=pose.position.y, z=pose.position.z), Vector3( + x=forward[0], y=forward[1], z=forward[2] + ) def is_novel_viewpoint( @@ -103,8 +107,10 @@ def is_novel_viewpoint( pos, fwd = viewpoint cos_thresh = math.cos(math.radians(angle_deg)) for prev_pos, prev_fwd in existing: - near_in_space = pos.distance(prev_pos) < min_baseline_m - near_in_angle = fwd.dot(prev_fwd) > cos_thresh + near_in_space = ( + math.dist((pos.x, pos.y, pos.z), (prev_pos.x, prev_pos.y, prev_pos.z)) < min_baseline_m + ) + near_in_angle = (fwd.x * prev_fwd.x + fwd.y * prev_fwd.y + fwd.z * prev_fwd.z) > cos_thresh if near_in_space and near_in_angle: return False return True @@ -287,7 +293,7 @@ def observe( objects: list[Object], *, frame_ts: float, - camera_transform: Transform | None, + camera_transform: TransformStamped | None, camera_info: CameraInfo, depth_m: NDArray[np.float32], ) -> None: @@ -298,7 +304,9 @@ def observe( previous_ts = self._now # Without a camera pose, fold identity but not viewpoint or absence evidence. viewpoint = ( - camera_viewpoint(camera_transform.to_pose()) if camera_transform is not None else None + camera_viewpoint(pose_from_transform(camera_transform).pose) + if camera_transform is not None + else None ) for obj in objects: # Normalize the center before association and publication. @@ -509,8 +517,7 @@ def _hit( track.obj.pointcloud = obj.pointcloud track.obj.ts = obj.ts track.obj.frame_id = obj.frame_id - track.obj.pose.ts = obj.pose.ts - track.obj.pose.frame_id = obj.pose.frame_id + track.obj.pose.header = obj.pose.header track.obj.image = obj.image track.obj.mask = obj.mask track.obj.bbox = obj.bbox @@ -557,12 +564,13 @@ def _hit( def _observation_center(obj: Object) -> Vector3: """Return the median world-space surface point, falling back to ``obj.center``.""" try: - pts = np.asarray(obj.pointcloud.pointcloud.points) + pts = pointcloud_xyz(obj.pointcloud) except Exception: return obj.center if pts.ndim != 2 or len(pts) < 8: return obj.center - return Vector3(np.median(pts, axis=0)) + center = np.median(pts, axis=0) + return Vector3(x=center[0], y=center[1], z=center[2]) def _geometry_suspect(self, track: _Track, obj: Object) -> bool: """Reject partial or >2x median-size geometry without folding it into the reference.""" @@ -581,7 +589,7 @@ def _vote_absence( self, matched_tracks: set[int], camera_info: CameraInfo, - world_from_camera: Transform, + world_from_camera: TransformStamped, depth_m: NDArray[np.float32], ) -> None: for track in self._tracks.values(): diff --git a/dimos/experimental/world_belief/worldbelief_module.py b/dimos/experimental/world_belief/worldbelief_module.py index a84fbd16d96..8d637d5d47c 100644 --- a/dimos/experimental/world_belief/worldbelief_module.py +++ b/dimos/experimental/world_belief/worldbelief_module.py @@ -23,6 +23,9 @@ import threading from typing import TYPE_CHECKING, Any, ClassVar +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.vision_msgs.msg import Detection3DArray from pydantic import Field from dimos.agents.annotation import skill @@ -41,9 +44,8 @@ WorldBeliefRecorderSpec, ) from dimos.memory.module import MemoryModuleConfig -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.msgs.pointcloud import voxel_downsample_cloud +from dimos.msgs.time import time_from_seconds from dimos.perception.experimental.object import Object from dimos.utils.logging_config import setup_logger @@ -306,9 +308,8 @@ def scan(self, prompt: list[str] | None = None, window: float = 60.0) -> SkillRe self.detections_3d.publish( to_detection3d_array(present, frame_id=frame_id, ts=result.as_of_ts) ) - cloud = aggregate_pointclouds(present).voxel_downsample(0.005) - cloud.frame_id = frame_id - cloud.ts = result.as_of_ts + cloud = voxel_downsample_cloud(aggregate_pointclouds(present), 0.005) + cloud.header = Header(frame_id=frame_id, stamp=time_from_seconds(result.as_of_ts)) self.pointcloud.publish(cloud) self.objects.publish(present) return SkillResult( @@ -387,7 +388,11 @@ def recall(self, text: str, k: int = 20) -> dict[str, Any] | None: "when_ts": round(float(hit.ts), 3), "where_camera": None if pose is None - else [round(pose.x, 3), round(pose.y, 3), round(pose.z, 3)], + else [ + round(pose.pose.position.x, 3), + round(pose.pose.position.y, 3), + round(pose.pose.position.z, 3), + ], "where_object": None if obj_center is None else [ diff --git a/dimos/experimental/world_belief/worldbelief_recorder.py b/dimos/experimental/world_belief/worldbelief_recorder.py index 276b0897d1c..cfe344623a8 100644 --- a/dimos/experimental/world_belief/worldbelief_recorder.py +++ b/dimos/experimental/world_belief/worldbelief_recorder.py @@ -18,14 +18,13 @@ from pathlib import Path from typing import Any, Protocol +from dimos_generated.geometry_msgs.msg import Pose, Quaternion +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, JointState + from dimos.constants import STATE_DIR from dimos.core.core import rpc from dimos.core.stream import In from dimos.memory.module import Recorder, RecorderConfig, pose_setter_for -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.spec.utils import Spec from dimos.utils.logging_config import setup_logger @@ -85,9 +84,9 @@ def recording_path(self) -> str: def _prepare_streams(self) -> None: super()._prepare_streams() depth = self.config.stream_remapping.get("depth_image", "depth_image") - self.store.stream(depth, Image, codec="lz4+lcm") + self.store.stream(depth, Image, codec="lz4+cdr") @pose_setter_for("coordinator_joint_state") async def _proprio_pose(self, msg: Any) -> Any: """Use an identity pose for proprioceptive joint-state records.""" - return Transform.identity().to_pose() + return Pose(orientation=Quaternion(w=1.0)) diff --git a/dimos/experimental/world_belief/xarm6_blueprint.py b/dimos/experimental/world_belief/xarm6_blueprint.py index 8292a164684..543bb16ef40 100644 --- a/dimos/experimental/world_belief/xarm6_blueprint.py +++ b/dimos/experimental/world_belief/xarm6_blueprint.py @@ -17,7 +17,9 @@ from __future__ import annotations from functools import partial -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any + +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 from dimos.agents.mcp.mcp_server import McpServer from dimos.constants import STATE_DIR @@ -30,19 +32,19 @@ ) from dimos.hardware.sensors.camera.realsense.camera import RealSenseCamera from dimos.manipulation.manipulation_module import ManipulationModule -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.robot.manipulators.common.blueprints import coordinator, trajectory_task from dimos.robot.manipulators.xarm.config import make_xarm6_model_config, xarm6_hardware from dimos.visualization.rerun.bridge import RerunBridgeModule +from dimos.visualization.rerun.message_helpers import camera_pinhole if TYPE_CHECKING: import rerun.blueprint as rrb -XARM6_WORLDBELIEF_CAMERA_TRANSFORM = Transform( - translation=Vector3(x=0.06693724, y=-0.0309563, z=0.00691482), - rotation=Quaternion(0.70513398, 0.00535696, 0.70897578, -0.01052180), +XARM6_WORLDBELIEF_CAMERA_TRANSFORM = TransformStamped( + transform=Transform( + translation=Vector3(x=0.06693724, y=-0.0309563, z=0.00691482), + rotation=Quaternion(x=0.70513398, y=0.00535696, z=0.70897578, w=-0.01052180), + ) ) @@ -59,10 +61,7 @@ def _topic_to_entity(topic: Any) -> str: def _camera_info_to_rerun(msg: Any, image_topic: str) -> list[tuple[str, Any]]: - return cast( - "list[tuple[str, Any]]", - msg.to_rerun(image_topic=image_topic, optical_frame=getattr(msg, "frame_id", None)), - ) + return [(image_topic, camera_pinhole(msg))] def _rerun_blueprint() -> rrb.Blueprint: diff --git a/dimos/hardware/drive_trains/flowbase/config.py b/dimos/hardware/drive_trains/flowbase/config.py index 494247d6c74..a002be9596d 100644 --- a/dimos/hardware/drive_trains/flowbase/config.py +++ b/dimos/hardware/drive_trains/flowbase/config.py @@ -14,9 +14,7 @@ """FlowBase platform constants.""" -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion -# Mid-360 mount pose on the FlowBase (position + orientation) in the base frame. -FLOWBASE_MID360_MOUNT = Pose(0.20, -0.20, 0.10, *Quaternion.from_euler(Vector3(0, 0, 0))) +# Mid-360 mount pose on the FlowBase in the base frame. +FLOWBASE_MID360_MOUNT = Pose(position=Point(x=0.20, y=-0.20, z=0.10), orientation=Quaternion(w=1.0)) diff --git a/dimos/hardware/drive_trains/transport/adapter.py b/dimos/hardware/drive_trains/transport/adapter.py index 921cd90d378..77b9dcd457e 100644 --- a/dimos/hardware/drive_trains/transport/adapter.py +++ b/dimos/hardware/drive_trains/transport/adapter.py @@ -23,10 +23,10 @@ import threading from typing import Any +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist, Vector3 + from dimos.core.transport_factory import make_transport -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.geometry import yaw from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -141,7 +141,11 @@ def read_enabled(self) -> bool: def _on_odom(self, msg: PoseStamped) -> None: with self._lock: - self._latest_odom = [msg.x, msg.y, msg.yaw] + self._latest_odom = [ + msg.pose.position.x, + msg.pose.position.y, + yaw(msg.pose.orientation), + ] def transport_lcm_factory(**kwargs: Any) -> TransportTwistAdapter: diff --git a/dimos/hardware/drive_trains/transport/test_adapter.py b/dimos/hardware/drive_trains/transport/test_adapter.py new file mode 100644 index 00000000000..eb2f7d8905d --- /dev/null +++ b/dimos/hardware/drive_trains/transport/test_adapter.py @@ -0,0 +1,63 @@ +# 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. + +import math + +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Twist, Vector3 +import pytest + +from dimos.hardware.drive_trains.transport.adapter import TransportTwistAdapter + + +def test_nested_pose_becomes_planar_odometry_without_sharing_state(): + adapter = TransportTwistAdapter() + assert adapter.read_odometry() is None + adapter._on_odom( + PoseStamped( + pose=Pose( + position=Point(x=2.0, y=-1.0), + orientation=Quaternion(z=math.sqrt(0.5), w=math.sqrt(0.5)), + ) + ) + ) + odometry = adapter.read_odometry() + assert odometry == pytest.approx([2.0, -1.0, math.pi / 2]) + odometry[0] = 99.0 + assert adapter.read_odometry()[0] == 2.0 + + +@pytest.mark.parametrize( + "dof,values,expected", + [ + (1, [0.1], Twist(linear=Vector3(x=0.1))), + (2, [0.1, 0.2], Twist(linear=Vector3(x=0.1, y=0.2))), + (3, [0.1, 0.2, 0.3], Twist(linear=Vector3(x=0.1, y=0.2), angular=Vector3(z=0.3))), + ], +) +def test_generated_commands_preserve_enable_and_stop_contract( + monkeypatch, mocker, dof, values, expected +): + adapter = TransportTwistAdapter(dof=dof) + transport = mocker.Mock() + monkeypatch.setattr(adapter, "_cmd_vel_transport", transport) + assert not adapter.write_velocities(values) + transport.publish.assert_not_called() + adapter.write_enable(True) + assert adapter.write_velocities(values) + message = transport.publish.call_args.args[0] + assert Twist.decode(message.encode()) == expected + assert adapter.read_velocities() == values + adapter.write_enable(False) + assert transport.publish.call_args.args[0] == Twist() + assert adapter.read_velocities() == [0.0] * dof diff --git a/dimos/hardware/manipulators/spec.py b/dimos/hardware/manipulators/spec.py index 4ca30e43a38..60deec651b4 100644 --- a/dimos/hardware/manipulators/spec.py +++ b/dimos/hardware/manipulators/spec.py @@ -27,9 +27,6 @@ from typing import Protocol, runtime_checkable from dimos.hardware.spec import JointLimits as _JointLimits -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 class DriverStatus(Enum): @@ -65,14 +62,6 @@ class ManipulatorInfo: serial_number: str | None = None -def default_base_transform() -> Transform: - """Default identity transform for arm mounting.""" - return Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - ) - - @runtime_checkable class ManipulatorAdapter(Protocol): """Protocol for hardware-specific IO. diff --git a/dimos/hardware/sensors/camera/gstreamer/gstreamer_camera.py b/dimos/hardware/sensors/camera/gstreamer/gstreamer_camera.py index 055d50f9602..4bf7bbf239b 100644 --- a/dimos/hardware/sensors/camera/gstreamer/gstreamer_camera.py +++ b/dimos/hardware/sensors/camera/gstreamer/gstreamer_camera.py @@ -20,13 +20,16 @@ import time from typing import Any +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header import numpy as np from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import Out -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_from_array +from dimos.msgs.time import time_from_seconds from dimos.utils.logging_config import setup_logger # Add system path for gi module if needed @@ -279,11 +282,10 @@ def _on_new_sample(self, appsink): # type: ignore[no-untyped-def] image_array = data.reshape((height, width, 3)) # Create an Image message with the absolute timestamp - image_msg = Image( - data=image_array.copy(), # Make a copy to ensure data persistence - format=ImageFormat.BGR, - frame_id=self.frame_id, - ts=timestamp, + image_msg = image_from_array( + image_array, + encoding="bgr8", + header=Header(frame_id=self.frame_id, stamp=time_from_seconds(timestamp)), ) # Publish the image diff --git a/dimos/hardware/sensors/camera/gstreamer/gstreamer_camera_test_script.py b/dimos/hardware/sensors/camera/gstreamer/gstreamer_camera_test_script.py index 7bef41e5e1c..5c7babb8ece 100755 --- a/dimos/hardware/sensors/camera/gstreamer/gstreamer_camera_test_script.py +++ b/dimos/hardware/sensors/camera/gstreamer/gstreamer_camera_test_script.py @@ -18,10 +18,12 @@ import logging import time +from dimos_generated.sensor_msgs.msg import Image + from dimos.core.coordination.module_coordinator import ModuleCoordinator from dimos.core.transport import LCMTransport from dimos.hardware.sensors.camera.gstreamer.gstreamer_camera import GstreamerCameraModule -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.time import to_seconds from dimos.protocol.service.lcmservice import autoconf from dimos.utils.logging_config import setup_logger @@ -81,25 +83,27 @@ def main() -> None: # Counter for received frames frame_count = [0] last_log_time = [time.time()] - first_timestamp = [None] + first_timestamp: list[float | None] = [None] - def on_frame(msg) -> None: # type: ignore[no-untyped-def] + def on_frame(msg: Image) -> None: frame_count[0] += 1 current_time = time.time() # Capture first timestamp to show absolute timestamps are preserved if first_timestamp[0] is None: - first_timestamp[0] = msg.ts - logger.info(f"First frame absolute timestamp: {msg.ts:.6f}") + first_timestamp[0] = to_seconds(msg.header.stamp) + logger.info(f"First frame absolute timestamp: {to_seconds(msg.header.stamp):.6f}") # Log stats every 2 seconds if current_time - last_log_time[0] >= 2.0: fps = frame_count[0] / (current_time - last_log_time[0]) - timestamp_delta = msg.ts - first_timestamp[0] + initial_stamp = first_timestamp[0] + assert initial_stamp is not None + timestamp_delta = to_seconds(msg.header.stamp) - initial_stamp logger.info( f"Received {frame_count[0]} frames - FPS: {fps:.1f} - " f"Resolution: {msg.width}x{msg.height} - " - f"Timestamp: {msg.ts:.3f} (delta: {timestamp_delta:.3f}s)" + f"Timestamp: {to_seconds(msg.header.stamp):.3f} (delta: {timestamp_delta:.3f}s)" ) frame_count[0] = 0 last_log_time[0] = current_time diff --git a/dimos/hardware/sensors/camera/gstreamer/test_generated_frames.py b/dimos/hardware/sensors/camera/gstreamer/test_generated_frames.py new file mode 100644 index 00000000000..cc413b20089 --- /dev/null +++ b/dimos/hardware/sensors/camera/gstreamer/test_generated_frames.py @@ -0,0 +1,84 @@ +# 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. + +import importlib.util +from pathlib import Path +import sys +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock + +from dimos_generated.sensor_msgs.msg import Image +import numpy as np +import pytest + +from dimos.msgs.image import image_view +from dimos.msgs.time import to_nanoseconds + + +@pytest.fixture +def camera_module(monkeypatch): + gi = ModuleType("gi") + gi.require_version = Mock() + repository = ModuleType("gi.repository") + repository.GLib = Mock() + repository.Gst = SimpleNamespace( + init=Mock(), + FlowReturn=SimpleNamespace(OK=0, ERROR=-1), + MapFlags=SimpleNamespace(READ=1), + CLOCK_TIME_NONE=-1, + ) + monkeypatch.setitem(sys.modules, "gi", gi) + monkeypatch.setitem(sys.modules, "gi.repository", repository) + spec = importlib.util.spec_from_file_location( + "_gstreamer_generated_test", Path(__file__).with_name("gstreamer_camera.py") + ) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + return module + + +def test_generated_bgr_frame_owns_pixels_after_sdk_buffer_unmap(camera_module, monkeypatch): + camera = camera_module.GstreamerCameraModule(frame_id="camera_test") + camera.running = True + pixels = np.arange(18, dtype=np.uint8).reshape(2, 3, 3) + raw = bytearray(pixels.tobytes()) + mapped = SimpleNamespace(data=raw) + buffer = Mock(pts=1700000000125000000) + buffer.map.return_value = (True, mapped) + buffer.unmap.side_effect = lambda info: raw.__setitem__(slice(None), bytes(18)) + caps = Mock() + caps.get_structure.return_value.get_value.side_effect = {"width": 3, "height": 2}.__getitem__ + sample = Mock() + sample.get_buffer.return_value = buffer + sample.get_caps.return_value = caps + sink = Mock() + sink.emit.return_value = sample + received = [] + monkeypatch.setattr(camera.video, "publish", received.append) + try: + assert camera._on_new_sample(sink) == 0 + assert len(received) == 1 and type(received[0]) is Image + value = Image.decode(received[0].encode()) + assert value.header.frame_id == "camera_test" + assert value.encoding == "bgr8" and value.step == 9 + assert to_nanoseconds(value.header.stamp) == 1700000000125000000 + np.testing.assert_array_equal(image_view(value), pixels) + buffer.unmap.assert_called_once_with(mapped) + # Existing invalid/pre-2000 timestamp rejection occurs before mapping. + buffer.pts = 1 + assert camera._on_new_sample(sink) == 0 and len(received) == 1 + assert buffer.map.call_count == 1 + finally: + camera.stop() diff --git a/dimos/hardware/sensors/camera/module.py b/dimos/hardware/sensors/camera/module.py index 83020165da7..2ef9856b843 100644 --- a/dimos/hardware/sensors/camera/module.py +++ b/dimos/hardware/sensors/camera/module.py @@ -13,8 +13,13 @@ # limitations under the License. from collections.abc import Callable +from copy import copy import time +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage from pydantic import Field import reactivex as rx @@ -25,28 +30,24 @@ from dimos.core.stream import Out from dimos.hardware.sensors.camera.spec import CameraHardware from dimos.hardware.sensors.camera.webcam import Webcam -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image, sharpness_barrier -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.image import image_sharpness +from dimos.msgs.time import time_from_nanoseconds from dimos.spec import perception +from dimos.utils.reactive import quality_barrier from dimos.visualization.vis_module import vis_module -def default_transform() -> Transform: - return Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="base_link", +def default_transform() -> TransformStamped: + return TransformStamped( + header=Header(frame_id="base_link"), child_frame_id="camera_link", + transform=Transform(rotation=Quaternion(w=1.0)), ) class CameraModuleConfig(ModuleConfig): frame_id: str = "camera_link" - transform: Transform | None = Field(default_factory=default_transform) + transform: TransformStamped | None = Field(default_factory=default_transform) hardware: Callable[[], CameraHardware] | CameraHardware = Webcam frequency: float = 0.0 # Hz, 0 means no limit @@ -71,7 +72,7 @@ def start(self) -> None: stream = self.hardware.image_stream() if self.config.frequency > 0: - stream = stream.pipe(sharpness_barrier(self.config.frequency)) + stream = stream.pipe(quality_barrier(image_sharpness, self.config.frequency)) self.register_disposable( stream.subscribe(self.color_image.publish), @@ -82,24 +83,21 @@ def start(self) -> None: ) def publish_metadata(self) -> None: - camera_info = self.hardware.camera_info.with_ts(time.time()) + camera_info = copy(self.hardware.camera_info) + camera_info.header.stamp = time_from_nanoseconds(time.time_ns()) self.camera_info.publish(camera_info) if not self.config.transform: return - camera_link = self.config.transform - camera_link.ts = camera_info.ts - - camera_optical = Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(-0.5, 0.5, -0.5, 0.5), - frame_id="camera_link", - child_frame_id="camera_optical", - ts=camera_link.ts, + camera_link = copy(self.config.transform) + camera_link.header.stamp = camera_info.header.stamp + camera_optical = TransformStamped( + header=Header(frame_id=camera_link.child_frame_id, stamp=camera_info.header.stamp), + child_frame_id=camera_info.header.frame_id, + transform=Transform(rotation=Quaternion(x=-0.5, y=0.5, z=-0.5, w=0.5)), ) - - self.tf.publish(TFMessage(camera_link, camera_optical)) + self.tf.publish(TFMessage(transforms=[camera_link, camera_optical])) @rpc def stop(self) -> None: diff --git a/dimos/hardware/sensors/camera/realsense/blueprints.py b/dimos/hardware/sensors/camera/realsense/blueprints.py index 5f4c8e38648..c7c449b58fd 100644 --- a/dimos/hardware/sensors/camera/realsense/blueprints.py +++ b/dimos/hardware/sensors/camera/realsense/blueprints.py @@ -25,14 +25,15 @@ from typing import Any +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header + from dimos.core.coordination.blueprints import autoconnect from dimos.core.global_config import global_config from dimos.hardware.sensors.camera.realsense.camera import RealSenseCamera -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.protocol.tf.static_tf_publisher import StaticTfPublisher +from dimos.visualization.rerun.message_helpers import cloud_archetype from dimos.visualization.vis_module import vis_module # Dot radius in screen-space UI points. @@ -46,13 +47,12 @@ class RealSenseMountTf(StaticTfPublisher): path from world to it and draws nothing. """ - def transforms(self) -> list[Transform]: + def transforms(self) -> list[TransformStamped]: return [ - Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="world", + TransformStamped( + header=Header(frame_id="world"), child_frame_id="camera_link", + transform=Transform(rotation=Quaternion(w=1.0)), ) ] @@ -78,7 +78,7 @@ def _rerun_blueprint() -> Any: def _cloud(cloud: PointCloud2) -> Any: # The default "spheres" mode is sized for sparse lidar; a dense RGBD cloud # reads better as flat dots that keep their size as you zoom. - return cloud.to_rerun(mode="points", ui_radius=_CLOUD_RADIUS_UI) + return cloud_archetype(cloud, mode="points", ui_radius=_CLOUD_RADIUS_UI) _vis = vis_module( diff --git a/dimos/hardware/sensors/camera/realsense/camera.py b/dimos/hardware/sensors/camera/realsense/camera.py index edc80b80f4b..3afd870ba36 100644 --- a/dimos/hardware/sensors/camera/realsense/camera.py +++ b/dimos/hardware/sensors/camera/realsense/camera.py @@ -19,17 +19,14 @@ from typing import TYPE_CHECKING, Any +from dimos_generated.dimos_msgs.msg import ImuInfo +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, Imu, PointCloud2 +from dimos_generated.tf2_msgs.msg import TFMessage from pydantic import Field from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.stream import Out from dimos.hardware.sensors.camera.spec import DepthCameraConfig -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.Imu import Imu -from dimos.msgs.sensor_msgs.ImuInfo import ImuInfo -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.spec import perception diff --git a/dimos/hardware/sensors/camera/realsense/rust/Cargo.lock b/dimos/hardware/sensors/camera/realsense/rust/Cargo.lock index 1a7cdf31bcc..2ab58459a48 100644 --- a/dimos/hardware/sensors/camera/realsense/rust/Cargo.lock +++ b/dimos/hardware/sensors/camera/realsense/rust/Cargo.lock @@ -220,6 +220,20 @@ name = "bytemuck" version = "1.25.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a1f896587b6f2c069c73d2f0913e2d590c3990285cd2f0b6aa02b786b4c679c" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.4", +] [[package]] name = "byteorder" @@ -540,9 +554,17 @@ dependencies = [ ] [[package]] -name = "dimos-lcm" +name = "dimos-generated-messages" +version = "0.1.0" +dependencies = [ + "re_cdr", + "serde", + "serde-big-array", +] + +[[package]] +name = "dimos-lcm-transport" version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" dependencies = [ "byteorder", "socket2 0.5.10", @@ -553,9 +575,9 @@ dependencies = [ name = "dimos-module" version = "0.1.0" dependencies = [ - "dimos-lcm", + "dimos-generated-messages", + "dimos-lcm-transport", "dimos-module-macros", - "lcm-msgs", "nalgebra", "rayon", "serde", @@ -582,8 +604,8 @@ dependencies = [ name = "dimos-realsense" version = "0.1.0" dependencies = [ + "dimos-generated-messages", "dimos-module", - "lcm-msgs", "pkg-config", "realsense-rust", "realsense-sys", @@ -1355,14 +1377,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" -[[package]] -name = "lcm-msgs" -version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" -dependencies = [ - "byteorder", -] - [[package]] name = "libc" version = "0.2.189" @@ -2154,6 +2168,18 @@ dependencies = [ "yasna", ] +[[package]] +name = "re_cdr" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bba10c21433b597481f366b978cdf2617e56785b38cb0dcc3833c6b8ad7e061a" +dependencies = [ + "bytemuck", + "byteorder", + "serde", + "thiserror 2.0.20", +] + [[package]] name = "realsense-rust" version = "1.3.0" @@ -2532,6 +2558,15 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "serde-big-array" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11fc7cc2c76d73e0f27ee52abbd64eec84d46f370c88371120433196934e4b7f" +dependencies = [ + "serde", +] + [[package]] name = "serde_core" version = "1.0.229" diff --git a/dimos/hardware/sensors/camera/realsense/rust/Cargo.toml b/dimos/hardware/sensors/camera/realsense/rust/Cargo.toml index db7e2d03df8..4f66e1f7922 100644 --- a/dimos/hardware/sensors/camera/realsense/rust/Cargo.toml +++ b/dimos/hardware/sensors/camera/realsense/rust/Cargo.toml @@ -19,24 +19,24 @@ path = "dimos_realsense::RealSense" python = "dimos.hardware.sensors.camera.realsense.camera:RealSenseCamera" [package.metadata.dimos.module.realsense.inputs] -tf = "tf2_msgs.TFMessage" +tf = "tf2_msgs/msg/TFMessage" [package.metadata.dimos.module.realsense.outputs] -color_image = "sensor_msgs.Image" -depth_image = "sensor_msgs.Image" -infrared_left = "sensor_msgs.Image" -infrared_right = "sensor_msgs.Image" -imu = "sensor_msgs.Imu" -imu_info = "sensor_msgs.ImuInfo" -pointcloud = "sensor_msgs.PointCloud2" -camera_info = "sensor_msgs.CameraInfo" -depth_camera_info = "sensor_msgs.CameraInfo" -infrared_left_camera_info = "sensor_msgs.CameraInfo" -infrared_right_camera_info = "sensor_msgs.CameraInfo" +color_image = "sensor_msgs/msg/Image" +depth_image = "sensor_msgs/msg/Image" +infrared_left = "sensor_msgs/msg/Image" +infrared_right = "sensor_msgs/msg/Image" +imu = "sensor_msgs/msg/Imu" +imu_info = "dimos_msgs/msg/ImuInfo" +pointcloud = "sensor_msgs/msg/PointCloud2" +camera_info = "sensor_msgs/msg/CameraInfo" +depth_camera_info = "sensor_msgs/msg/CameraInfo" +infrared_left_camera_info = "sensor_msgs/msg/CameraInfo" +infrared_right_camera_info = "sensor_msgs/msg/CameraInfo" [dependencies] dimos-module = { path = "../../../../../../native/rust/dimos-module" } -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +dimos-generated-messages = { path = "../../../../../message_codegen" } realsense-rust = "1.3" realsense-sys = "2.56" tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal"] } diff --git a/dimos/hardware/sensors/camera/realsense/rust/flake.nix b/dimos/hardware/sensors/camera/realsense/rust/flake.nix index 3fa255d6046..5bdf5bfd3a7 100644 --- a/dimos/hardware/sensors/camera/realsense/rust/flake.nix +++ b/dimos/hardware/sensors/camera/realsense/rust/flake.nix @@ -11,7 +11,7 @@ # Also works when NativeModule builds from a regular Python environment. # Use the Nix compiler/linker so libc matches the camera libraries. devShells = forAll (pkgs: { - default = pkgs.mkShell { packages = [ pkgs.cargo pkgs.rustc pkgs.clippy pkgs.librealsense pkgs.pkg-config ]; }; + default = pkgs.mkShell { packages = [ pkgs.python3 pkgs.cargo pkgs.rustc pkgs.clippy pkgs.librealsense pkgs.pkg-config ]; }; }); }; } diff --git a/dimos/hardware/sensors/camera/realsense/rust/src/main.rs b/dimos/hardware/sensors/camera/realsense/rust/src/main.rs index 36de88262f3..a6d22933cf4 100644 --- a/dimos/hardware/sensors/camera/realsense/rust/src/main.rs +++ b/dimos/hardware/sensors/camera/realsense/rust/src/main.rs @@ -18,20 +18,26 @@ use std::collections::{BTreeSet, HashMap, HashSet, VecDeque}; use std::ffi::{CStr, CString}; -use std::io::Write; use std::os::raw::c_void; use std::ptr::{self, NonNull}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{mpsc, Arc, Mutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use dimos_generated_messages::builtin_interfaces::msg::Time; +use dimos_generated_messages::dimos_msgs::msg::ImuInfo; +use dimos_generated_messages::geometry_msgs::msg::{ + Quaternion as QuaternionMsg, Vector3 as Vector3Msg, +}; +use dimos_generated_messages::sensor_msgs::msg::{ + CameraInfo, Image, Imu, PointCloud2, PointField, RegionOfInterest, +}; +use dimos_generated_messages::std_msgs::msg::Header; +use dimos_module::cdr; use dimos_module::nalgebra::{ Isometry3, Matrix3, Quaternion, Rotation3, Translation3, UnitQuaternion, Vector3, }; use dimos_module::{native_config, run_with_transport, Module, Output, Tf, Transform}; -use lcm_msgs::geometry_msgs::{Quaternion as QuaternionMsg, Vector3 as Vector3Msg}; -use lcm_msgs::sensor_msgs::{CameraInfo, Image, Imu, PointCloud2, PointField, RegionOfInterest}; -use lcm_msgs::std_msgs::{Header, Time}; use realsense_rust::base::{Rs2Extrinsics, Rs2Intrinsics}; use realsense_rust::config::Config as RsConfig; use realsense_rust::context::Context; @@ -56,9 +62,6 @@ const IMU_CLOCK_WINDOW_SECONDS: i64 = 30; const FRAME_TIMEOUT: Duration = Duration::from_secs(1); // Body (REP-103) from optical axes, x y z w. const OPTICAL_ROTATION: [f64; 4] = [-0.5, 0.5, -0.5, 0.5]; -// lcm-gen fingerprint of sensor_msgs.ImuInfo; it has no generated bindings. -const IMU_INFO_FINGERPRINT: u64 = 0xe6aa5563f2a33280; - // ---- config, field for field with RealSenseCameraConfig ---- #[derive(Debug, Clone, Deserialize, Serialize)] @@ -170,36 +173,6 @@ impl Config { } } -// ---- sensor_msgs.ImuInfo, hand-encoded ---- - -#[derive(Clone)] -struct ImuInfo { - header: Header, - gyro_noise_density: f64, - gyro_random_walk: f64, - accel_noise_density: f64, - accel_random_walk: f64, - frequency: f64, -} - -impl ImuInfo { - fn encode(&self) -> Vec { - let mut buf = Vec::with_capacity(8 + self.header.encoded_size() + 40); - buf.write_all(&IMU_INFO_FINGERPRINT.to_be_bytes()).unwrap(); - self.header.encode_one(&mut buf).unwrap(); - for v in [ - self.gyro_noise_density, - self.gyro_random_walk, - self.accel_noise_density, - self.accel_random_walk, - self.frequency, - ] { - buf.write_all(&v.to_be_bytes()).unwrap(); - } - buf - } -} - // ---- module ---- /// What the capture thread learns at start and the tickers publish from. @@ -250,27 +223,27 @@ struct Outs { #[derive(Module)] #[module(name = "realsense", setup = start, teardown = stop)] struct RealSense { - #[output(encode = Image::encode)] + #[output(encode = cdr::encode)] color_image: Output, - #[output(encode = Image::encode)] + #[output(encode = cdr::encode)] depth_image: Output, - #[output(encode = Image::encode)] + #[output(encode = cdr::encode)] infrared_left: Output, - #[output(encode = Image::encode)] + #[output(encode = cdr::encode)] infrared_right: Output, - #[output(encode = Imu::encode)] + #[output(encode = cdr::encode)] imu: Output, - #[output(encode = ImuInfo::encode)] + #[output(encode = cdr::encode)] imu_info: Output, - #[output(encode = PointCloud2::encode)] + #[output(encode = cdr::encode)] pointcloud: Output, - #[output(encode = CameraInfo::encode)] + #[output(encode = cdr::encode)] camera_info: Output, - #[output(encode = CameraInfo::encode)] + #[output(encode = cdr::encode)] depth_camera_info: Output, - #[output(encode = CameraInfo::encode)] + #[output(encode = cdr::encode)] infrared_left_camera_info: Output, - #[output(encode = CameraInfo::encode)] + #[output(encode = cdr::encode)] infrared_right_camera_info: Output, #[tf] tf: Tf, @@ -360,10 +333,9 @@ fn unix_now() -> f64 { fn header(frame_id: &str, ts: f64) -> Header { let sec = ts.trunc(); Header { - seq: 0, stamp: Time { sec: sec as i32, - nsec: ((ts - sec) * 1e9) as i32, + nanosec: ((ts - sec) * 1e9) as u32, }, frame_id: frame_id.to_string(), } @@ -390,11 +362,11 @@ fn image( }; Image { header: header(frame_id, ts), - height: frame.height() as i32, - width: frame.width() as i32, + height: frame.height() as u32, + width: frame.width() as u32, encoding: encoding.to_string(), is_bigendian: 0, - step: (frame.width() * bytes_per_pixel) as i32, + step: (frame.width() * bytes_per_pixel) as u32, data: data.to_vec(), } } @@ -411,13 +383,13 @@ fn camera_info(intrinsics: &Rs2Intrinsics, frame_id: &str) -> CameraInfo { }; CameraInfo { header: header(frame_id, 0.0), - height: intrinsics.height() as i32, - width: intrinsics.width() as i32, + height: intrinsics.height() as u32, + width: intrinsics.width() as u32, distortion_model: model.to_string(), - D: intrinsics.0.coeffs.iter().map(|&c| c as f64).collect(), - K: [fx, 0.0, cx, 0.0, fy, cy, 0.0, 0.0, 1.0], - R: [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0], - P: [fx, 0.0, cx, 0.0, 0.0, fy, cy, 0.0, 0.0, 0.0, 1.0, 0.0], + d: intrinsics.0.coeffs.iter().map(|&c| c as f64).collect(), + k: [fx, 0.0, cx, 0.0, fy, cy, 0.0, 0.0, 1.0], + r: [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0], + p: [fx, 0.0, cx, 0.0, 0.0, fy, cy, 0.0, 0.0, 0.0, 1.0, 0.0], binning_x: 0, binning_y: 0, roi: RegionOfInterest { @@ -677,7 +649,7 @@ fn capture_thread(cfg: &Config, shared: &Shared, handle: &Handle, outs: &Outs) { // The pair is rectified, so the whole of their offset is along x. if let (Some(info), Ok(ext)) = (infra2_info.as_mut(), p2.extrinsics(p1)) { let baseline = (ext.translation()[0] as f64).abs(); - info.P[3] = -info.P[0] * baseline; + info.p[3] = -info.p[0] * baseline; } } } @@ -937,10 +909,10 @@ fn pointcloud_thread(cfg: &Config, shared: &Shared, handle: &Handle, outs: &Outs } fn empty_cloud(header: &Header) -> PointCloud2 { - let field = |name: &str, offset: i32| PointField { + let field = |name: &str, offset: u32| PointField { name: name.to_string(), offset, - datatype: PointField::FLOAT32 as u8, + datatype: PointField::FLOAT32, count: 1, }; PointCloud2 { @@ -968,8 +940,8 @@ fn rgbd_to_cloud( stride: usize, depth_trunc_m: f32, ) -> PointCloud2 { - let (fx, fy) = (info.K[0] as f32, info.K[4] as f32); - let (cx, cy) = (info.K[2] as f32, info.K[5] as f32); + let (fx, fy) = (info.k[0] as f32, info.k[4] as f32); + let (cx, cy) = (info.k[2] as f32, info.k[5] as f32); let (w, h) = (rgbd.depth.width as usize, rgbd.depth.height as usize); let rgb = &rgbd.color.data; if rgb.len() != w * h * 3 || rgbd.depth.data.len() != w * h * 2 { @@ -1020,13 +992,13 @@ fn rgbd_to_cloud( ((c[0] * 255.0) as u32) << 16 | ((c[1] * 255.0) as u32) << 8 | (c[2] * 255.0) as u32; data.extend_from_slice(&packed.to_le_bytes()); } - let field = |name: &str, offset: i32| PointField { + let field = |name: &str, offset: u32| PointField { name: name.to_string(), offset, - datatype: PointField::FLOAT32 as u8, + datatype: PointField::FLOAT32, count: 1, }; - let n = points.len() as i32; + let n = points.len() as u32; PointCloud2 { header: rgbd.depth.header.clone(), height: if n == 0 { 0 } else { 1 }, @@ -1406,10 +1378,9 @@ mod tests { accel_random_walk: 4.0, frequency: 5.0, }; - let bytes = msg.encode(); - assert_eq!(&bytes[..8], &IMU_INFO_FINGERPRINT.to_be_bytes()); - assert_eq!(bytes.len(), 8 + msg.header.encoded_size() + 40); - assert_eq!(&bytes[bytes.len() - 8..], &5.0f64.to_be_bytes()); + let bytes = cdr::encode(&msg).unwrap(); + assert_eq!(&bytes[..4], &[0, 1, 0, 0]); + assert_eq!(cdr::decode::(&bytes).unwrap(), msg); } } @@ -1427,20 +1398,20 @@ mod bench { let rgbd = Rgbd { color: Image { header: header("f", 0.0), - height: h as i32, - width: w as i32, + height: h as u32, + width: w as u32, encoding: "rgb8".into(), is_bigendian: 0, - step: (w * 3) as i32, + step: (w * 3) as u32, data: vec![128; w * h * 3], }, depth: Image { header: header("f", 0.0), - height: h as i32, - width: w as i32, + height: h as u32, + width: w as u32, encoding: "16UC1".into(), is_bigendian: 0, - step: (w * 2) as i32, + step: (w * 2) as u32, data: depth, }, }; @@ -1463,7 +1434,7 @@ mod bench { let cloud = rgbd_to_cloud(&rgbd, &info, 0.001, stride, 5.0); let build = t.elapsed(); let t = std::time::Instant::now(); - let bytes = cloud.encode(); + let bytes = cdr::encode(&cloud).unwrap(); eprintln!( "stride={stride} points={} build={:?} encode={:?} bytes={}", cloud.width, diff --git a/dimos/hardware/sensors/camera/spec.py b/dimos/hardware/sensors/camera/spec.py index ee6d2c9ee3b..613d912b439 100644 --- a/dimos/hardware/sensors/camera/spec.py +++ b/dimos/hardware/sensors/camera/spec.py @@ -14,14 +14,13 @@ from abc import ABC, abstractmethod +from dimos_generated.geometry_msgs.msg import Quaternion +from dimos_generated.sensor_msgs.msg import CameraInfo, Image from reactivex.observable import Observable -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image from dimos.protocol.service.spec import BaseConfig, Configurable -OPTICAL_ROTATION = Quaternion(-0.5, 0.5, -0.5, 0.5) +OPTICAL_ROTATION = Quaternion(x=-0.5, y=0.5, z=-0.5, w=0.5) class CameraConfig(BaseConfig): diff --git a/dimos/hardware/sensors/camera/test_module.py b/dimos/hardware/sensors/camera/test_module.py new file mode 100644 index 00000000000..ec1c5b30519 --- /dev/null +++ b/dimos/hardware/sensors/camera/test_module.py @@ -0,0 +1,59 @@ +# 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. + +"""Camera metadata uses the same generated header and frame chain as its images.""" + +from dimos_generated.sensor_msgs.msg import CameraInfo +from dimos_generated.tf2_msgs.msg import TFMessage +import numpy as np +import pytest + +from dimos.hardware.sensors.camera.module import CameraModule +from dimos.hardware.sensors.camera.webcam import Webcam +from dimos.msgs.geometry import transform_matrix + + +@pytest.fixture +def module(): + camera = CameraModule() + camera.hardware = Webcam(frame_id_prefix="robot") + try: + yield camera + finally: + camera.stop() + + +def test_metadata_preserves_calibration_and_does_not_mutate_mount_template(module, mocker): + info_publish = mocker.patch.object(module.camera_info, "publish") + tf_publish = mocker.patch.object(module.tf, "publish") + mocker.patch( + "dimos.hardware.sensors.camera.module.time.time_ns", return_value=1700000000123456789 + ) + original = module.config.transform.encode() + + module.publish_metadata() + + info_publish.assert_called_once() + tf_publish.assert_called_once() + info = CameraInfo.decode(info_publish.call_args.args[0].encode()) + transforms = TFMessage.decode(tf_publish.call_args.args[0].encode()).transforms + assert (info.header.stamp.sec, info.header.stamp.nanosec) == (1700000000, 123456789) + assert len(transforms) == 2 + assert [(edge.header.frame_id, edge.child_frame_id) for edge in transforms] == [ + ("base_link", "camera_link"), + ("camera_link", "robot/camera_optical"), + ] + assert all(edge.header.stamp == info.header.stamp for edge in transforms) + assert module.config.transform.encode() == original + np.testing.assert_allclose(transform_matrix(transforms[0].transform), np.eye(4)) diff --git a/dimos/hardware/sensors/camera/test_webcam.py b/dimos/hardware/sensors/camera/test_webcam.py index 6e390d078c7..44264d4af3e 100644 --- a/dimos/hardware/sensors/camera/test_webcam.py +++ b/dimos/hardware/sensors/camera/test_webcam.py @@ -12,27 +12,53 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest + from dimos.hardware.sensors.camera.webcam import Webcam -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo +from dimos.msgs.camera_info import camera_info_from_fov +from dimos.msgs.image import image_view def test_webcam_without_intrinsics_reports_a_nominal_pinhole() -> None: info = Webcam(width=1280, height=720).camera_info assert (info.width, info.height) == (1280, 720) - assert info.K[0] > 0 - assert info.frame_id == "camera_optical" + assert info.k[0] > 0 + assert info.header.frame_id == "camera_optical" def test_webcam_keeps_configured_intrinsics() -> None: - configured = CameraInfo.from_fov(90.0, 640, 480) + configured = camera_info_from_fov(90.0, 640, 480, header=Header()) assert Webcam(camera_info=configured).camera_info is configured def test_webcam_with_size_but_no_focal_length_gets_a_nominal_pinhole() -> None: info = Webcam(camera_info=CameraInfo(width=1280, height=720)).camera_info - assert info.K[0] > 0 + assert info.k[0] > 0 def test_webcam_stereo_slice_halves_the_nominal_pinhole_width() -> None: info = Webcam(width=1280, height=720, stereo_slice="left").camera_info assert (info.width, info.height) == (640, 720) + + +@pytest.mark.parametrize("side,start", [("left", 0), ("right", 2)]) +def test_capture_converts_bgr_and_crops_before_cdr_encoding(side, start, mocker): + camera = Webcam(stereo_slice=side, frame_id_prefix="robot") + pixels = np.arange(2 * 4 * 3, dtype=np.uint8).reshape(2, 4, 3) + capture = mocker.Mock() + capture.read.return_value = (True, pixels) + camera._capture = capture + mocker.patch("dimos.msgs.time.time.time_ns", return_value=1700000000123456789) + try: + message = Image.decode(camera.capture_frame().encode()) + finally: + camera.stop() + assert message.encoding == "rgb8" + assert (message.width, message.height, message.step) == (2, 2, 6) + assert message.header.frame_id == "robot/camera_optical" + assert (message.header.stamp.sec, message.header.stamp.nanosec) == (1700000000, 123456789) + np.testing.assert_array_equal(image_view(message), pixels[:, start : start + 2, ::-1]) + capture.release.assert_called_once() diff --git a/dimos/hardware/sensors/camera/webcam.py b/dimos/hardware/sensors/camera/webcam.py index 8523f356b6d..a64f6081b9d 100644 --- a/dimos/hardware/sensors/camera/webcam.py +++ b/dimos/hardware/sensors/camera/webcam.py @@ -18,13 +18,16 @@ import time from typing import Annotated, Any, Literal +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.std_msgs.msg import Header from pydantic import BeforeValidator, Field from reactivex import create from reactivex.observable import Observable from dimos.hardware.sensors.camera.spec import CameraConfig, CameraHardware -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.camera_info import camera_info_from_fov +from dimos.msgs.image import image_from_array +from dimos.msgs.time import header_now from dimos.utils.reactive import backpressure @@ -142,22 +145,13 @@ def capture_frame(self) -> Image: # Convert BGR to RGB (OpenCV uses BGR by default) frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - # Create Image message - # Using Image.from_numpy() since it's designed for numpy arrays - # Setting format to RGB since we converted from BGR->RGB above - image = Image.from_numpy( - frame_rgb, - format=ImageFormat.RGB, # We converted to RGB above - frame_id=self._frame("camera_optical"), # Standard frame ID for camera images - ts=time.time(), # Current timestamp - ) - if self.config.stereo_slice in ("left", "right"): - half_width = image.width // 2 - if self.config.stereo_slice == "left": - image = image.crop(0, 0, half_width, image.height) - else: - image = image.crop(half_width, 0, half_width, image.height) + half_width = frame_rgb.shape[1] // 2 + start = 0 if self.config.stereo_slice == "left" else half_width + frame_rgb = frame_rgb[:, start : start + half_width] + image = image_from_array( + frame_rgb, encoding="rgb8", header=header_now(self._frame("camera_optical")) + ) self._emitted_size = (image.width, image.height) return image @@ -190,7 +184,7 @@ def _capture_loop(self) -> None: @property def camera_info(self) -> CameraInfo: info = self.config.camera_info - if info.width and info.height and info.K[0] > 0 and info.K[4] > 0: + if info.width and info.height and info.k[0] > 0 and info.k[4] > 0: return info # No intrinsics configured: a nominal pinhole so the image still renders. # Sized from the frames actually emitted (the stereo slice halves them). @@ -199,8 +193,12 @@ def camera_info(self) -> CameraInfo: else: width = self.config.width // 2 if self.config.stereo_slice else self.config.width height = self.config.height - return CameraInfo.from_fov( - 60.0, width, height, axis="horizontal", frame_id=self._frame("camera_optical") + return camera_info_from_fov( + 60.0, + width, + height, + axis="horizontal", + header=Header(frame_id=self._frame("camera_optical")), ) def emit(self, image: Image) -> None: ... diff --git a/dimos/hardware/sensors/camera/zed/camera.py b/dimos/hardware/sensors/camera/zed/camera.py index fd777339dfc..b240f7adcad 100644 --- a/dimos/hardware/sensors/camera/zed/camera.py +++ b/dimos/hardware/sensors/camera/zed/camera.py @@ -18,6 +18,10 @@ import threading import time +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage from pydantic import Field import pyzed.sl as sl import reactivex as rx @@ -33,13 +37,11 @@ DepthCameraConfig, DepthCameraHardware, ) -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.camera_info import camera_info_with_stamp +from dimos.msgs.geometry import compose_transforms, inverse_transform +from dimos.msgs.image import image_from_array +from dimos.msgs.pointcloud import pointcloud_from_rgbd, voxel_downsample_cloud +from dimos.msgs.time import time_from_seconds from dimos.spec import perception from dimos.utils.reactive import backpressure @@ -47,8 +49,8 @@ def default_base_transform() -> Transform: """Default identity transform for camera mounting.""" return Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), + translation=Vector3(), + rotation=Quaternion(w=1.0), ) @@ -131,11 +133,13 @@ def __init__(self, *args, **kwargs) -> None: # type: ignore[no-untyped-def] def _publish_camera_info(self) -> None: ts = time.time() if self._color_camera_info: - self._color_camera_info.ts = ts - self.camera_info.publish(self._color_camera_info) + self.camera_info.publish( + camera_info_with_stamp(self._color_camera_info, time_from_seconds(ts)) + ) if self._depth_camera_info: - self._depth_camera_info.ts = ts - self.depth_camera_info.publish(self._depth_camera_info) + self.depth_camera_info.publish( + camera_info_with_stamp(self._depth_camera_info, time_from_seconds(ts)) + ) @rpc def start(self) -> None: @@ -232,10 +236,11 @@ def _intrinsics_to_camera_info( height=self._stream_height, width=self._stream_width, distortion_model="plumb_bob", - D=D, - K=K, - P=P, - frame_id=frame_id, + d=D, + k=K, + r=[1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0], + p=P, + header=Header(frame_id=frame_id), ) def _get_extrinsics(self) -> None: @@ -251,15 +256,16 @@ def _extrinsics_to_transform( frame_id: str, child_frame_id: str, ts: float, - ) -> Transform: + ) -> TransformStamped: translation = extrinsics.get_translation().get() - quat = extrinsics.get_orientation().get() # [x, y, z, w] - return Transform( - translation=Vector3(*translation), - rotation=Quaternion(quat[0], quat[1], quat[2], quat[3]), - frame_id=frame_id, + quat = extrinsics.get_orientation().get() + return TransformStamped( + header=Header(frame_id=frame_id, stamp=time_from_seconds(ts)), child_frame_id=child_frame_id, - ts=ts, + transform=Transform( + translation=Vector3(x=translation[0], y=translation[1], z=translation[2]), + rotation=Quaternion(x=quat[0], y=quat[1], z=quat[2], w=quat[3]), + ), ) def _enable_tracking(self) -> None: @@ -301,11 +307,10 @@ def _capture_loop(self) -> None: if color_data.ndim == 3 and color_data.shape[2] == 4: color_data = color_data[:, :, :3] color_data = cv2.cvtColor(color_data, cv2.COLOR_BGR2RGB) - color_img = Image( - data=color_data, - format=ImageFormat.RGB, - frame_id=self._color_optical_frame, - ts=ts, + color_img = image_from_array( + color_data, + encoding="rgb8", + header=Header(frame_id=self._color_optical_frame, stamp=time_from_seconds(ts)), ) self.color_image.publish(color_img) @@ -320,11 +325,10 @@ def _capture_loop(self) -> None: if self.config.align_depth_to_color else self._depth_optical_frame ) - depth_img = Image( - data=depth_data, - format=ImageFormat.DEPTH, - frame_id=depth_frame_id, - ts=ts, + depth_img = image_from_array( + depth_data, + encoding="32FC1", + header=Header(frame_id=depth_frame_id, stamp=time_from_seconds(ts)), ) self.depth_image.publish(depth_img) @@ -335,97 +339,62 @@ def _capture_loop(self) -> None: self._publish_tf(ts) - def _tracking_transform(self, ts: float) -> Transform | None: + def _tracking_transform(self, ts: float) -> TransformStamped | None: if not self._tracking_enabled or self._zed is None or self._pose is None: return None state = self._zed.get_position(self._pose, sl.REFERENCE_FRAME.WORLD) if state != sl.POSITIONAL_TRACKING_STATE.OK: return None - translation = self._pose.get_translation().get().tolist() - rotation = self._pose.get_orientation().get().tolist() - world_to_camera = Transform( - translation=Vector3(*translation), - rotation=Quaternion(*rotation), - frame_id=self.config.world_frame, - child_frame_id=self._camera_link, - ts=ts, + world_to_camera = self._extrinsics_to_transform( + self._pose, self.config.world_frame, self._camera_link, ts ) if self.config.base_transform is None: return world_to_camera - - base_to_camera = Transform( - translation=self.config.base_transform.translation, - rotation=self.config.base_transform.rotation, - frame_id=self.config.base_frame_id, + base_to_camera = TransformStamped( + header=Header(frame_id=self.config.base_frame_id, stamp=time_from_seconds(ts)), child_frame_id=self._camera_link, - ts=ts, + transform=self.config.base_transform, ) - camera_to_base = base_to_camera.inverse() - world_to_base = world_to_camera + camera_to_base - world_to_base.frame_id = self.config.world_frame - world_to_base.child_frame_id = self.config.base_frame_id - world_to_base.ts = ts - return world_to_base + return compose_transforms(world_to_camera, inverse_transform(base_to_camera)) def _publish_tf(self, ts: float) -> None: - transforms = [] - + transforms: list[TransformStamped] = [] + stamp = time_from_seconds(ts) if self.config.base_transform is not None: - base_to_camera = Transform( - translation=self.config.base_transform.translation, - rotation=self.config.base_transform.rotation, - frame_id=self.config.base_frame_id, - child_frame_id=self._camera_link, - ts=ts, + transforms.append( + TransformStamped( + header=Header(frame_id=self.config.base_frame_id, stamp=stamp), + child_frame_id=self._camera_link, + transform=self.config.base_transform, + ) ) - transforms.append(base_to_camera) - - # camera_imu_transform is IMU -> left_camera (coordinate transform), - # we need to invert to get the pose of left camera in camera_link frame - camera_link_to_depth = self._extrinsics_to_transform( - self._camera_link_to_color_extrinsics, - self._camera_link, - self._depth_frame, - ts, - ).inverse() - camera_link_to_depth.frame_id = self._camera_link - camera_link_to_depth.child_frame_id = self._depth_frame - transforms.append(camera_link_to_depth) - - depth_to_depth_optical = Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=OPTICAL_ROTATION, - frame_id=self._depth_frame, - child_frame_id=self._depth_optical_frame, - ts=ts, - ) - transforms.append(depth_to_depth_optical) - - color_tf = self._extrinsics_to_transform( - self._camera_link_to_color_extrinsics, - self._camera_link, - self._color_frame, - ts, - ).inverse() - color_tf.frame_id = self._camera_link - color_tf.child_frame_id = self._color_frame - transforms.append(color_tf) - - color_to_color_optical = Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=OPTICAL_ROTATION, - frame_id=self._color_frame, - child_frame_id=self._color_optical_frame, - ts=ts, - ) - transforms.append(color_to_color_optical) + # SDK extrinsics are the IMU-to-left-camera coordinate transform; + # retain the inverse used by the existing camera-link pose convention. + for frame, optical_frame in ( + (self._depth_frame, self._depth_optical_frame), + (self._color_frame, self._color_optical_frame), + ): + extrinsics = inverse_transform( + self._extrinsics_to_transform( + self._camera_link_to_color_extrinsics, self._camera_link, frame, ts + ) + ) + extrinsics.header.frame_id = self._camera_link + extrinsics.child_frame_id = frame + transforms.append(extrinsics) + transforms.append( + TransformStamped( + header=Header(frame_id=frame, stamp=stamp), + child_frame_id=optical_frame, + transform=Transform(rotation=OPTICAL_ROTATION), + ) + ) tracking_tf = self._tracking_transform(ts) if tracking_tf is not None: transforms.append(tracking_tf) - - self.tf.publish(TFMessage(*transforms)) + self.tf.publish(TFMessage(transforms=transforms)) def _generate_pointcloud(self) -> None: with self._pointcloud_lock: @@ -436,13 +405,13 @@ def _generate_pointcloud(self) -> None: return try: - pcd = PointCloud2.from_rgbd( - color_image=color_img, - depth_image=depth_img, - camera_info=self._color_camera_info, + pcd = pointcloud_from_rgbd( + color=color_img, + depth=depth_img, + calibration=self._color_camera_info, depth_scale=self._depth_scale, ) - pcd = pcd.voxel_downsample(0.005) + pcd = voxel_downsample_cloud(pcd, 0.005) self.pointcloud.publish(pcd) except Exception as e: print(f"Pointcloud generation error: {e}") diff --git a/dimos/hardware/sensors/camera/zed/compat.py b/dimos/hardware/sensors/camera/zed/compat.py index f7e78fc5d62..2ee92c2341f 100644 --- a/dimos/hardware/sensors/camera/zed/compat.py +++ b/dimos/hardware/sensors/camera/zed/compat.py @@ -16,7 +16,7 @@ from pathlib import Path -from dimos.msgs.sensor_msgs.CameraInfo import CalibrationProvider +from dimos.msgs.camera_info import CalibrationProvider try: import pyzed.sl # noqa: F401 diff --git a/dimos/hardware/sensors/camera/zed/test_generated_tf.py b/dimos/hardware/sensors/camera/zed/test_generated_tf.py new file mode 100644 index 00000000000..84ea16dd3ae --- /dev/null +++ b/dimos/hardware/sensors/camera/zed/test_generated_tf.py @@ -0,0 +1,110 @@ +# 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. + +from enum import Enum +import importlib.util +from pathlib import Path +import sys +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock + +from dimos_generated.geometry_msgs.msg import Transform, Vector3 +from dimos_generated.tf2_msgs.msg import TFMessage +import numpy as np +import pytest + +from dimos.msgs.geometry import transform_matrix +from dimos.msgs.time import time_from_seconds + + +@pytest.fixture +def camera_module(monkeypatch): + # Load the real camera class with only the optional SDK boundary mocked. + # Neither Camera() nor start() is invoked. + sl = ModuleType("pyzed.sl") + sl.DEPTH_MODE = Enum("DEPTH_MODE", {"NEURAL": 1}) + sl.REFERENCE_FRAME = SimpleNamespace(WORLD=1) + sl.POSITIONAL_TRACKING_STATE = SimpleNamespace(OK=1) + package = ModuleType("pyzed") + package.sl = sl + monkeypatch.setitem(sys.modules, "pyzed", package) + monkeypatch.setitem(sys.modules, "pyzed.sl", sl) + spec = importlib.util.spec_from_file_location( + "_zed_generated_tf_test", Path(__file__).with_name("camera.py") + ) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + return module + + +def sdk_pose(translation): + return SimpleNamespace( + get_translation=lambda: SimpleNamespace(get=lambda: np.asarray(translation)), + get_orientation=lambda: SimpleNamespace(get=lambda: np.asarray([0.0, 0.0, 0.0, 1.0])), + ) + + +@pytest.mark.parametrize("mount", [None, Transform(translation=Vector3(x=0.25))]) +def test_tracking_uses_generated_frame_chain(camera_module, mount): + camera = camera_module.ZEDCamera(base_transform=mount) + camera._zed = Mock() + camera._zed.get_position.return_value = 1 + camera._pose = sdk_pose([1.0, 2.0, 3.0]) + camera._tracking_enabled = True + try: + result = camera._tracking_transform(1700000000.125) + assert result.header.frame_id == "world" + assert result.header.stamp == time_from_seconds(1700000000.125) + assert result.child_frame_id == ("camera_link" if mount is None else "base_link") + assert result.transform.translation.x == (1.0 if mount is None else 0.75) + assert result.transform.translation.y == 2.0 + camera._zed.get_position.return_value = 0 + assert camera._tracking_transform(2.0) is None + finally: + camera._zed = None + camera.stop() + + +def test_generated_tf_preserves_mount_extrinsics_and_optical_frames(camera_module, monkeypatch): + camera = camera_module.ZEDCamera() + camera._camera_link_to_color_extrinsics = sdk_pose([0.1, 0.2, 0.3]) + received = [] + monkeypatch.setattr(camera.tf, "publish", received.append) + try: + camera._publish_tf(1700000000.125) + message = TFMessage.decode(received[0].encode()) + assert len(message.transforms) == 5 + frames = [(t.header.frame_id, t.child_frame_id) for t in message.transforms] + assert frames == [ + ("base_link", "camera_link"), + ("camera_link", "camera_depth_frame"), + ("camera_depth_frame", "camera_depth_optical_frame"), + ("camera_link", "camera_color_frame"), + ("camera_color_frame", "camera_color_optical_frame"), + ] + assert all(t.header.stamp == time_from_seconds(1700000000.125) for t in message.transforms) + np.testing.assert_array_equal(transform_matrix(message.transforms[0].transform), np.eye(4)) + for index in (1, 3): + np.testing.assert_allclose( + transform_matrix(message.transforms[index].transform)[:3, 3], [-0.1, -0.2, -0.3] + ) + for index in (2, 4): + np.testing.assert_allclose( + transform_matrix(message.transforms[index].transform)[:3, :3], + [[0.0, 0.0, 1.0], [-1.0, 0.0, 0.0], [0.0, -1.0, 0.0]], + atol=1e-15, + ) + finally: + camera.stop() diff --git a/dimos/hardware/sensors/lidar/common/point_cloud_utils.hpp b/dimos/hardware/sensors/lidar/common/point_cloud_utils.hpp index 1a9d43d2139..f9b95a3e4d3 100644 --- a/dimos/hardware/sensors/lidar/common/point_cloud_utils.hpp +++ b/dimos/hardware/sensors/lidar/common/point_cloud_utils.hpp @@ -5,29 +5,32 @@ #pragma once -#include +#include +#include +#include #include #include +#include -#include "sensor_msgs/PointCloud2.hpp" -#include "sensor_msgs/PointField.hpp" -#include "std_msgs/Header.hpp" -#include "std_msgs/Time.hpp" namespace dimos { -inline std_msgs::Time time_from_seconds(double t) { - std_msgs::Time ts; - ts.sec = static_cast(t); - ts.nsec = static_cast((t - ts.sec) * 1e9); +inline builtin_interfaces::msg::Time time_from_seconds(double t) { + if (!std::isfinite(t)) throw std::invalid_argument("timestamp must be finite"); + double seconds = std::floor(t); + auto nanoseconds = std::llround((t - seconds) * 1e9); + if (nanoseconds == 1000000000) { ++seconds; nanoseconds = 0; } + if (seconds < std::numeric_limits::min() || seconds > std::numeric_limits::max()) + throw std::invalid_argument("timestamp exceeds ROS Time range"); + builtin_interfaces::msg::Time ts; + ts.sec = static_cast(seconds); + ts.nanosec = static_cast(nanoseconds); return ts; } -// Stamped Header with a process-wide auto-incrementing sequence number. -inline std_msgs::Header make_header(const std::string& frame_id, double ts) { - static std::atomic seq{0}; - std_msgs::Header h; - h.seq = seq.fetch_add(1, std::memory_order_relaxed); +// ROS Header carries the source timestamp and coordinate frame. +inline std_msgs::msg::Header make_header(const std::string& frame_id, double ts) { + std_msgs::msg::Header h; h.stamp = time_from_seconds(ts); h.frame_id = frame_id; return h; @@ -39,22 +42,23 @@ inline constexpr int32_t kXyziPointStep = kXyziFieldCount * sizeof(float); // Empty XYZI PointCloud2 sized for num_points. The caller fills each point // through xyzi_point() below. -inline sensor_msgs::PointCloud2 make_xyzi_cloud(const std::string& frame_id, double ts, +inline sensor_msgs::msg::PointCloud2 make_xyzi_cloud(const std::string& frame_id, double ts, int num_points) { - sensor_msgs::PointCloud2 pc; + if (num_points < 0 || static_cast(num_points) * kXyziPointStep > std::numeric_limits::max()) + throw std::invalid_argument("point cloud row exceeds ROS uint32 dimensions"); + sensor_msgs::msg::PointCloud2 pc; pc.header = make_header(frame_id, ts); pc.height = 1; pc.width = num_points; pc.is_bigendian = 0; pc.is_dense = 1; - pc.fields_length = kXyziFieldCount; pc.fields.resize(kXyziFieldCount); auto make_field = [](const std::string& name, int32_t index) { - sensor_msgs::PointField f; + sensor_msgs::msg::PointField f; f.name = name; f.offset = index * static_cast(sizeof(float)); - f.datatype = sensor_msgs::PointField::FLOAT32; + f.datatype = sensor_msgs::msg::PointField::FLOAT32; f.count = 1; return f; }; @@ -65,13 +69,12 @@ inline sensor_msgs::PointCloud2 make_xyzi_cloud(const std::string& frame_id, dou pc.point_step = kXyziPointStep; pc.row_step = pc.point_step * num_points; - pc.data_length = pc.row_step; - pc.data.resize(pc.data_length); + pc.data.resize(pc.row_step); return pc; } // The x/y/z/intensity slot of point i in a cloud from make_xyzi_cloud. -inline float* xyzi_point(sensor_msgs::PointCloud2& pc, int i) { +inline float* xyzi_point(sensor_msgs::msg::PointCloud2& pc, int i) { return reinterpret_cast(pc.data.data() + i * kXyziPointStep); } diff --git a/dimos/hardware/sensors/lidar/fastlio2/cpp/CMakeLists.txt b/dimos/hardware/sensors/lidar/fastlio2/cpp/CMakeLists.txt index 91e96cb082d..5b953318420 100644 --- a/dimos/hardware/sensors/lidar/fastlio2/cpp/CMakeLists.txt +++ b/dimos/hardware/sensors/lidar/fastlio2/cpp/CMakeLists.txt @@ -39,13 +39,7 @@ if(NOT FASTLIO_DIR) set(FASTLIO_DIR ${CMAKE_CURRENT_SOURCE_DIR}/fast-lio-non-ros) endif() -# dimos-lcm C++ message headers -FetchContent_Declare(dimos_lcm - GIT_REPOSITORY https://github.com/dimensionalOS/dimos-lcm.git - GIT_TAG main - GIT_SHALLOW TRUE -) -FetchContent_MakeAvailable(dimos_lcm) + # The header-only dimos native C++ SDK. Overridable so nix can point at the # flake input instead of the in-repo relative path. @@ -87,7 +81,6 @@ endif() target_include_directories(fastlio2_native PRIVATE ${FASTLIO_DIR}/include ${FASTLIO_DIR}/src - ${dimos_lcm_SOURCE_DIR}/generated/cpp_lcm_msgs ${LCM_INCLUDE_DIRS} ${EIGEN3_INCLUDE_DIR} ${PCL_INCLUDE_DIRS} diff --git a/dimos/hardware/sensors/lidar/fastlio2/cpp/flake.lock b/dimos/hardware/sensors/lidar/fastlio2/cpp/flake.lock index 7adfee54000..e520f5b284a 100644 --- a/dimos/hardware/sensors/lidar/fastlio2/cpp/flake.lock +++ b/dimos/hardware/sensors/lidar/fastlio2/cpp/flake.lock @@ -1,22 +1,5 @@ { "nodes": { - "dimos-lcm": { - "flake": false, - "locked": { - "lastModified": 1769774949, - "narHash": "sha256-icRK7jerqNlwK1WZBrnIP04I2WozzFqTD7qsmnPxQuo=", - "owner": "dimensionalOS", - "repo": "dimos-lcm", - "rev": "0aa72b7b1bd3a65f50f5c03485ee9b728df56afe", - "type": "github" - }, - "original": { - "owner": "dimensionalOS", - "ref": "main", - "repo": "dimos-lcm", - "type": "github" - } - }, "fast-lio": { "flake": false, "locked": { @@ -163,7 +146,6 @@ }, "root": { "inputs": { - "dimos-lcm": "dimos-lcm", "fast-lio": "fast-lio", "flake-utils": "flake-utils", "lcm-extended": "lcm-extended", diff --git a/dimos/hardware/sensors/lidar/fastlio2/cpp/flake.nix b/dimos/hardware/sensors/lidar/fastlio2/cpp/flake.nix index 4ae6a9d0798..34ecabe9f0c 100644 --- a/dimos/hardware/sensors/lidar/fastlio2/cpp/flake.nix +++ b/dimos/hardware/sensors/lidar/fastlio2/cpp/flake.nix @@ -8,10 +8,6 @@ livox-sdk.url = "path:../../livox/cpp"; livox-sdk.inputs.nixpkgs.follows = "nixpkgs"; livox-sdk.inputs.flake-utils.follows = "flake-utils"; - dimos-lcm = { - url = "github:dimensionalOS/dimos-lcm/main"; - flake = false; - }; # Standalone Boost.PFR, consumed by the SDK via a FetchContent source override. pfr = { url = "github:apolukhin/pfr_non_boost/2.3.2"; @@ -30,7 +26,7 @@ }; }; - outputs = { self, nixpkgs, zenoh, flake-utils, livox-sdk, dimos-lcm, pfr, fast-lio, lcm-extended, ... }: + outputs = { self, nixpkgs, zenoh, flake-utils, livox-sdk, pfr, fast-lio, lcm-extended, ... }: flake-utils.lib.eachDefaultSystem (system: let # Overlay fixes for darwin-broken nixpkgs recipes in our transitive @@ -83,6 +79,7 @@ nativeBuildInputs = [ pkgs.cmake pkgs.pkg-config ]; buildInputs = [ + (import ../../../../../../native/cpp/messages.nix { inherit pkgs; }) livox-sdk2 lcm pkgs.glib @@ -97,7 +94,6 @@ cmakeFlags = [ "-DCMAKE_POLICY_VERSION_MINIMUM=3.5" - "-DFETCHCONTENT_SOURCE_DIR_DIMOS_LCM=${dimos-lcm}" "-DFETCHCONTENT_SOURCE_DIR_PFR=${pfr}" "-DFASTLIO_DIR=${fast-lio}" "-DLIVOX_COMMON_DIR=${livox-common}" diff --git a/dimos/hardware/sensors/lidar/fastlio2/cpp/main.cpp b/dimos/hardware/sensors/lidar/fastlio2/cpp/main.cpp index e45adcf03ba..b8c89ab074f 100644 --- a/dimos/hardware/sensors/lidar/fastlio2/cpp/main.cpp +++ b/dimos/hardware/sensors/lidar/fastlio2/cpp/main.cpp @@ -21,10 +21,8 @@ #include "point_cloud_utils.hpp" #include "dimos/native.hpp" +#include -#include "nav_msgs/Odometry.hpp" -#include "sensor_msgs/PointCloud2.hpp" -#include "std_msgs/Header.hpp" // FAST-LIO (header-only core, compiled sources linked via CMake) #include "fast_lio.hpp" @@ -108,8 +106,8 @@ class FastLio2 : public Module { public: void build(Builder& builder, Config& config) override { cfg_ = config.parse(); - lidar_ = builder.output("lidar"); - odometry_ = builder.output("odometry"); + lidar_ = builder.output("lidar"); + odometry_ = builder.output("odometry"); frame_interval_ = std::chrono::microseconds(static_cast(1e6 / cfg_.frequency)); @@ -371,7 +369,7 @@ class FastLio2 : public Module { void publish_pointcloud(const PointCloudXYZI::Ptr& cloud, double ts) { int num_points = static_cast(cloud->size()); - sensor_msgs::PointCloud2 pc = make_xyzi_cloud(cfg_.sensor_frame_id, ts, num_points); + sensor_msgs::msg::PointCloud2 pc = make_xyzi_cloud(cfg_.sensor_frame_id, ts, num_points); for (int i = 0; i < num_points; ++i) { float* dst = xyzi_point(pc, i); @@ -386,7 +384,7 @@ class FastLio2 : public Module { // Publish odometry as frame_id (fixed) -> sensor_frame_id (moving sensor). void publish_odometry(const custom_messages::Odometry& odom, double ts) { - nav_msgs::Odometry msg; + nav_msgs::msg::Odometry msg; msg.header = make_header(cfg_.frame_id, ts); msg.child_frame_id = cfg_.sensor_frame_id; @@ -409,14 +407,14 @@ class FastLio2 : public Module { msg.twist.twist.angular.x = 0; msg.twist.twist.angular.y = 0; msg.twist.twist.angular.z = 0; - std::memset(msg.twist.covariance, 0, sizeof(msg.twist.covariance)); + msg.twist.covariance.fill(0); odometry_.publish(msg); } FastLio2Config cfg_; - Output lidar_; - Output odometry_; + Output lidar_; + Output odometry_; std::unique_ptr fast_lio_; // All four come from config, set in build(). diff --git a/dimos/hardware/sensors/lidar/fastlio2/tools/pcap_to_db.py b/dimos/hardware/sensors/lidar/fastlio2/tools/pcap_to_db.py index 0c42313508b..c2bddba6b96 100644 --- a/dimos/hardware/sensors/lidar/fastlio2/tools/pcap_to_db.py +++ b/dimos/hardware/sensors/lidar/fastlio2/tools/pcap_to_db.py @@ -148,12 +148,13 @@ def _write_rrd(db_path: Path, odom_stream: str, lidar_stream: str, voxel: float) FastLio2 publishes its cloud in the sensor/body frame, so each frame is transformed to world by its pose here, then voxel-deduped. Best-effort: any failure is non-fatal to the recording. Returns the .rrd path, or None.""" + from dimos_generated.nav_msgs.msg import Odometry + from dimos_generated.sensor_msgs.msg import PointCloud2 import numpy as np import rerun as rr from dimos.memory.store.sqlite import SqliteStore - from dimos.msgs.nav_msgs.Odometry import Odometry - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + from dimos.msgs.pointcloud import pointcloud_xyz from dimos.visualization.rerun.init import rerun_init store = SqliteStore(path=str(db_path), must_exist=True) @@ -188,7 +189,7 @@ def _write_rrd(db_path: Path, odom_stream: str, lidar_stream: str, voxel: float) j = int(np.argmin(np.abs(ots - lid.ts))) if abs(ots[j] - lid.ts) > _POSE_MATCH_TOL: continue - pts = np.asarray(lid.data.as_numpy()[0])[:, :3].astype(np.float64) + pts = pointcloud_xyz(lid.data).astype(np.float64) if pts.shape[0] == 0: continue world = pts @ _quat_to_rot(*oquat[j]).T + opos[j] diff --git a/dimos/hardware/sensors/lidar/livox/module.py b/dimos/hardware/sensors/lidar/livox/module.py index 5f6279f1f1e..9f28e91953a 100644 --- a/dimos/hardware/sensors/lidar/livox/module.py +++ b/dimos/hardware/sensors/lidar/livox/module.py @@ -30,6 +30,7 @@ import sys from typing import TYPE_CHECKING, Any, Literal +from dimos_generated.sensor_msgs.msg import Imu, PointCloud2 from pydantic import Field, field_validator from dimos.core.core import rpc @@ -47,8 +48,6 @@ SDK_POINT_DATA_PORT, SDK_PUSH_MSG_PORT, ) -from dimos.msgs.sensor_msgs.Imu import Imu -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.spec import perception diff --git a/dimos/hardware/sensors/lidar/livox/rust/Cargo.toml b/dimos/hardware/sensors/lidar/livox/rust/Cargo.toml index 61c75c6aa6d..4f2a56f4cac 100644 --- a/dimos/hardware/sensors/lidar/livox/rust/Cargo.toml +++ b/dimos/hardware/sensors/lidar/livox/rust/Cargo.toml @@ -17,14 +17,14 @@ path = "dimos_livox::module::Mid360" python = "dimos.hardware.sensors.lidar.livox.module:Mid360" [package.metadata.dimos.module.mid360.outputs] -lidar = "sensor_msgs.PointCloud2" -imu = "sensor_msgs.Imu" +lidar = "sensor_msgs/msg/PointCloud2" +imu = "sensor_msgs/msg/Imu" [dependencies] crc = "3" dimos-module = { path = "../../../../../../native/rust/dimos-module" } etherparse = "0.21.0" -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +dimos-generated-messages = { path = "../../../../../message_codegen" } pcap-parser = "0.17.0" serde = { version = "1", features = ["derive"] } serde_json = "1" diff --git a/dimos/hardware/sensors/lidar/livox/rust/src/module.rs b/dimos/hardware/sensors/lidar/livox/rust/src/module.rs index 15502ed40dd..097c84621c4 100644 --- a/dimos/hardware/sensors/lidar/livox/rust/src/module.rs +++ b/dimos/hardware/sensors/lidar/livox/rust/src/module.rs @@ -18,10 +18,12 @@ use crate::live::{LiveConfig, LiveSource, Ports}; use crate::pcap::PcapSource; use crate::pipeline::{Frame, ImuRecord, PacketSource, Pipeline}; use crate::wire::{DataPacket, DataType}; +use dimos_generated_messages::builtin_interfaces::msg::Time; +use dimos_generated_messages::geometry_msgs::msg::{Quaternion, Vector3}; +use dimos_generated_messages::sensor_msgs::msg::{Imu, PointCloud2, PointField}; +use dimos_generated_messages::std_msgs::msg::Header; +use dimos_module::cdr; use dimos_module::{native_config, Module, Output}; -use lcm_msgs::geometry_msgs::{Quaternion, Vector3}; -use lcm_msgs::sensor_msgs::{Imu, PointCloud2, PointField}; -use lcm_msgs::std_msgs::{Header, Time}; use serde::{Deserialize, Serialize}; use std::net::Ipv4Addr; use std::sync::atomic::{AtomicBool, Ordering}; @@ -98,10 +100,10 @@ enum PointFormat { } impl PointFormat { - fn point_step(self) -> i32 { + fn point_step(self) -> u32 { match self { - PointFormat::Full => (OFFSET_FULL_LINE + 1) as i32, - PointFormat::Minimal | PointFormat::Legacy => (OFFSET_FOURTH + 4) as i32, + PointFormat::Full => (OFFSET_FULL_LINE + 1) as u32, + PointFormat::Minimal | PointFormat::Legacy => (OFFSET_FOURTH + 4) as u32, } } } @@ -124,10 +126,10 @@ impl Config { #[derive(Module)] #[module(name = "mid360", setup = start, handle = wait, teardown = stop)] pub struct Mid360 { - #[output(encode = PointCloud2::encode)] + #[output(encode = cdr::encode)] lidar: Output, - #[output(encode = Imu::encode)] + #[output(encode = cdr::encode)] imu: Output, #[config] @@ -277,10 +279,9 @@ fn run_pipeline( fn header(frame_id: &str, ts_ns: u64) -> Header { Header { - seq: 0, stamp: Time { sec: (ts_ns / 1_000_000_000) as i32, - nsec: (ts_ns % 1_000_000_000) as i32, + nanosec: (ts_ns % 1_000_000_000) as u32, }, frame_id: frame_id.to_string(), } @@ -324,13 +325,13 @@ const OFFSET_FULL_LINE: usize = 21; fn cloud_fields(format: PointFormat) -> Vec { let make_field = |name: &str, offset: usize, datatype: u8| PointField { name: name.into(), - offset: offset as i32, + offset: offset as u32, datatype, count: 1, }; - let f32t = PointField::FLOAT32 as u8; - let u32t = PointField::UINT32 as u8; - let u8t = PointField::UINT8 as u8; + let f32t = PointField::FLOAT32; + let u32t = PointField::UINT32; + let u8t = PointField::UINT8; let mut fields = vec![ make_field("x", OFFSET_X, f32t), make_field("y", OFFSET_Y, f32t), @@ -374,7 +375,7 @@ fn cloud_message(format: PointFormat, frame_id: &str, frame: &Frame) -> PointClo } } - let point_count = frame.points.len() as i32; + let point_count = frame.points.len() as u32; PointCloud2 { header: header(frame_id, frame.start_ns), height: 1, @@ -420,7 +421,7 @@ mod tests { assert_eq!(cloud.width, 2); assert_eq!(cloud.header.frame_id, "lidar_link"); assert_eq!(cloud.header.stamp.sec, 1_700_000_000); - assert_eq!(cloud.header.stamp.nsec, 123_000_000); + assert_eq!(cloud.header.stamp.nanosec, 123_000_000); let names: Vec<&str> = cloud.fields.iter().map(|f| f.name.as_str()).collect(); assert_eq!( names, @@ -471,7 +472,7 @@ mod tests { }; let msg = imu_message("imu_link", &record); assert_eq!(msg.header.stamp.sec, 5); - assert_eq!(msg.header.stamp.nsec, 500_000_000); + assert_eq!(msg.header.stamp.nanosec, 500_000_000); assert_eq!(msg.orientation_covariance[0], -1.0); assert_eq!(msg.orientation.w, 1.0); assert_eq!(msg.angular_velocity.z, 0.3); diff --git a/dimos/hardware/sensors/lidar/pointlio/scripts/pcap_to_db.py b/dimos/hardware/sensors/lidar/pointlio/scripts/pcap_to_db.py index dd6e66f745d..69d57f44aed 100644 --- a/dimos/hardware/sensors/lidar/pointlio/scripts/pcap_to_db.py +++ b/dimos/hardware/sensors/lidar/pointlio/scripts/pcap_to_db.py @@ -229,12 +229,13 @@ def _write_rrd(db_path: Path, odom_stream: str, lidar_stream: str, voxel: float) Point-LIO publishes its cloud in the sensor/body frame, so each frame is transformed to world by its pose here, then voxel-deduped. Best-effort: any failure is non-fatal to the recording. Returns the .rrd path, or None.""" + from dimos_generated.nav_msgs.msg import Odometry + from dimos_generated.sensor_msgs.msg import PointCloud2 import numpy as np import rerun as rr from dimos.memory.store.sqlite import SqliteStore - from dimos.msgs.nav_msgs.Odometry import Odometry - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + from dimos.msgs.pointcloud import pointcloud_xyz from dimos.visualization.rerun.init import rerun_init store = SqliteStore(path=str(db_path), must_exist=True) @@ -269,7 +270,7 @@ def _write_rrd(db_path: Path, odom_stream: str, lidar_stream: str, voxel: float) j = int(np.argmin(np.abs(ots - lid.ts))) if abs(ots[j] - lid.ts) > _POSE_MATCH_TOL: continue - pts = np.asarray(lid.data.as_numpy()[0])[:, :3].astype(np.float64) + pts = pointcloud_xyz(lid.data).astype(np.float64) if pts.shape[0] == 0: continue world = pts @ _quat_to_rot(*oquat[j]).T + opos[j] diff --git a/dimos/hardware/sensors/lidar/test_generated_recorders.py b/dimos/hardware/sensors/lidar/test_generated_recorders.py new file mode 100644 index 00000000000..09137c3ff48 --- /dev/null +++ b/dimos/hardware/sensors/lidar/test_generated_recorders.py @@ -0,0 +1,93 @@ +# 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. + +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseWithCovariance, Quaternion +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest + +from dimos.hardware.sensors.camera.realsense.blueprints import RealSenseMountTf, _cloud +from dimos.hardware.sensors.lidar.fastlio2.recorder import FastLio2Recorder +from dimos.hardware.sensors.lidar.pointlio.recorder import PointlioRecorder +from dimos.mapping.relocalization.blueprints import RecordingPlayer, _fine_points +from dimos.msgs.geometry import transform_matrix +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.robot.assembly.mid360_realsense_30 import Mid360RealsenseRecorder +from dimos.robot.unitree.g1.g1_recorder import G1Recorder, G1RecorderConfig + + +@pytest.mark.asyncio +@pytest.mark.parametrize("recorder_type", [FastLio2Recorder, PointlioRecorder]) +async def test_generated_recorders_anchor_clouds_to_last_odometry(recorder_type, tmp_path): + recorder = recorder_type(db_path=tmp_path / "unused.db") + try: + cloud = pointcloud_from_xyz(np.array([[1, 2, 3]], dtype=np.float32), header=Header()) + assert await recorder._lidar_pose(cloud) is None + pose = Pose(position=Point(x=2, y=3), orientation=Quaternion(w=1)) + odometry = Odometry(header=Header(frame_id="world"), pose=PoseWithCovariance(pose=pose)) + assert await recorder._odom_pose(odometry) == pose + assert await recorder._lidar_pose(cloud) == pose + assert not (tmp_path / "unused.db").exists() + finally: + recorder.stop() + + +def test_recorder_and_replay_ports_use_generated_message_identities(): + for recorder_type in [FastLio2Recorder, PointlioRecorder]: + streams = { + stream.name: stream.type for stream in recorder_type.blueprint().blueprints[0].streams + } + prefix = "fastlio" if recorder_type is FastLio2Recorder else "pointlio" + assert streams[f"{prefix}_odometry"] is Odometry + assert streams[f"{prefix}_lidar"] is PointCloud2 + streams = { + stream.name: stream.type + for stream in Mid360RealsenseRecorder.blueprint().blueprints[0].streams + } + assert streams["color_image"] is Image + assert streams["realsense_depth_image"] is Image + assert streams["realsense_pointcloud"] is PointCloud2 + assert streams["realsense_camera_info"] is CameraInfo + streams = { + stream.name: stream.type for stream in RecordingPlayer.blueprint().blueprints[0].streams + } + assert streams["lidar"] is PointCloud2 + + +def test_realsense_mount_is_an_identity_and_generated_cloud_views_are_typed(): + mount = RealSenseMountTf() + try: + edge = mount.transforms()[0] + assert edge.header.frame_id == "world" and edge.child_frame_id == "camera_link" + np.testing.assert_allclose(transform_matrix(edge.transform), np.eye(4)) + finally: + mount.stop() + cloud = pointcloud_from_xyz(np.array([[1, 2, 3]], dtype=np.float32), header=Header()) + before = cloud.encode() + assert _cloud(cloud).positions is not None + assert _fine_points(cloud).positions is not None + assert cloud.encode() == before + + +def test_g1_depth_recorder_declares_generated_lossless_cdr(): + config = G1RecorderConfig() + assert config.stream_codecs["realsense_depth_image"] == "lz4+cdr" + streams = {s.name: s.type for s in G1Recorder.blueprint().blueprints[0].streams} + assert streams["realsense_depth_image"] is Image + assert streams["realsense_camera_info"] is CameraInfo diff --git a/dimos/hardware/whole_body/transport/adapter.py b/dimos/hardware/whole_body/transport/adapter.py index 2ad2b223b84..2b496609a1f 100644 --- a/dimos/hardware/whole_body/transport/adapter.py +++ b/dimos/hardware/whole_body/transport/adapter.py @@ -24,12 +24,13 @@ import threading from typing import Any +from dimos_generated.dimos_msgs.msg import MotorCommandArray +from dimos_generated.sensor_msgs.msg import Imu, JointState + from dimos.core.transport_factory import make_transport from dimos.hardware.spec import JointLimits from dimos.hardware.whole_body.spec import IMUState, MotorCommand, MotorState -from dimos.msgs.sensor_msgs.Imu import Imu -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.sensor_msgs.MotorCommandArray import MotorCommandArray +from dimos.msgs.time import header_now from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -135,6 +136,7 @@ def write_motor_commands(self, commands: list[MotorCommand]) -> bool: return False msg = MotorCommandArray( + header=header_now(), q=[c.q for c in commands], dq=[c.dq for c in commands], kp=[c.kp for c in commands], diff --git a/dimos/hardware/whole_body/transport/test_adapter.py b/dimos/hardware/whole_body/transport/test_adapter.py new file mode 100644 index 00000000000..b779185731b --- /dev/null +++ b/dimos/hardware/whole_body/transport/test_adapter.py @@ -0,0 +1,54 @@ +# 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. + +from dimos_generated.geometry_msgs.msg import Quaternion, Vector3 +from dimos_generated.sensor_msgs.msg import Imu, JointState +import pytest + +from dimos.hardware.whole_body.spec import IMUState, MotorCommand, MotorState +from dimos.hardware.whole_body.transport.adapter import TransportWholeBodyAdapter + + +def test_unconnected_adapter_does_not_send_commands(): + adapter = TransportWholeBodyAdapter(dof=2) + assert not adapter.write_motor_commands([MotorCommand(), MotorCommand()]) + assert not adapter.has_motor_states() + assert adapter.read_motor_states() == [MotorState(), MotorState()] + + +@pytest.mark.parametrize("field", ["position", "velocity", "effort"]) +def test_short_feedback_preserves_last_complete_snapshot(field): + adapter = TransportWholeBodyAdapter(dof=2) + message = JointState(position=[1.0, 2.0], velocity=[3.0, 4.0], effort=[5.0, 6.0]) + adapter._on_motor_states(message) + setattr(message, field, [0.0]) + adapter._on_motor_states(message) + assert adapter.read_motor_states() == [ + MotorState(q=1.0, dq=3.0, tau=5.0), + MotorState(q=2.0, dq=4.0, tau=6.0), + ] + + +def test_generated_imu_converts_xyzw_to_hardware_wxyz(): + adapter = TransportWholeBodyAdapter(dof=2) + adapter._on_imu( + Imu( + orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.4), + angular_velocity=Vector3(x=1.0, y=2.0, z=3.0), + linear_acceleration=Vector3(x=4.0, y=5.0, z=6.0), + ) + ) + assert adapter.read_imu() == IMUState( + quaternion=(0.4, 0.1, 0.2, 0.3), gyroscope=(1.0, 2.0, 3.0), accelerometer=(4.0, 5.0, 6.0) + ) diff --git a/dimos/imitation/dataprep/test_generated_images.py b/dimos/imitation/dataprep/test_generated_images.py new file mode 100644 index 00000000000..7f4e9a4f72e --- /dev/null +++ b/dimos/imitation/dataprep/test_generated_images.py @@ -0,0 +1,51 @@ +# 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. + +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest + +from dimos.imitation.dataprep.core import StreamField, is_image_array, resolve_field +from dimos.msgs.image import compressed_image_from_image, image_from_array + + +@pytest.mark.parametrize("compressed", [False, True]) +def test_generated_camera_field_is_a_pixel_array_after_cdr(compressed): + pixels = np.full((8, 10, 3), [30, 90, 150], dtype=np.uint8) + image = image_from_array(pixels, encoding="bgr8", header=Header(frame_id="camera")) + message = compressed_image_from_image(image, quality=100) if compressed else image + message = type(message).decode(message.encode()) + before = message.encode() + array = resolve_field(message, StreamField(stream="camera")) + assert array.shape == pixels.shape and array.dtype == pixels.dtype + assert is_image_array(array) + np.testing.assert_allclose(array, pixels, atol=2 if compressed else 0) + assert message.encode() == before + + +def test_generated_mono16_camera_respects_row_padding_and_endian(): + pixels = np.array([[1, 65535], [256, 4096]], dtype=">u2") + padded = np.zeros((2, 6), dtype=np.uint8) + padded[:, :4] = pixels.view(np.uint8).reshape(2, 4) + message = Image( + width=2, height=2, encoding="mono16", is_bigendian=1, step=6, data=padded.ravel() + ) + array = resolve_field({"camera": message}, StreamField(stream="dict", field="camera")) + assert is_image_array(array) + np.testing.assert_array_equal(array, pixels) + assert array.strides == (6, 2) diff --git a/dimos/imitation/policy/lerobot/module.py b/dimos/imitation/policy/lerobot/module.py index 9f23719c5e6..5e6978c1b1a 100644 --- a/dimos/imitation/policy/lerobot/module.py +++ b/dimos/imitation/policy/lerobot/module.py @@ -19,6 +19,9 @@ from pathlib import Path from typing import Protocol, TypedDict +from dimos_generated.sensor_msgs.msg import Image, JointState +from dimos_generated.std_msgs.msg import UInt32 +from dimos_generated.trajectory_msgs.msg import JointTrajectory from pydantic import Field, field_validator from dimos.control.tasks.trajectory_task.trajectory_task import ( @@ -31,9 +34,6 @@ IsolatedPythonModule, IsolatedPythonModuleConfig, ) -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory from dimos.spec.utils import Spec from dimos.teleop.webxr.controller_types import BUTTON_ALIASES, Buttons @@ -119,8 +119,8 @@ class LeRobotPolicyModule(IsolatedPythonModule): color_image: In[Image] coordinator_joint_state: In[JointState] - button_pressed: In[Buttons] - teleop_buttons: In[Buttons] + button_pressed: In[UInt32] + teleop_buttons: In[UInt32] _control: PolicyControlSpec diff --git a/dimos/manipulation/conftest.py b/dimos/manipulation/conftest.py index dff9c169ab6..706bbec57fb 100644 --- a/dimos/manipulation/conftest.py +++ b/dimos/manipulation/conftest.py @@ -19,6 +19,7 @@ from typing import Any, Protocol, cast from unittest.mock import DEFAULT, MagicMock +from dimos_generated.dimos_msgs.msg import TrajectoryStatus import pytest from pytest_mock import MockerFixture @@ -32,7 +33,8 @@ from dimos.manipulation.manipulation_module import ManipulationModule from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig -from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState, TrajectoryStatus +from dimos.msgs.time import header_now +from dimos.msgs.trajectory import TrajectoryState from dimos.robot.assets.model import RobotModel @@ -55,7 +57,9 @@ def invoke(task: str, method: str, args: dict | None = None): return TrajectoryCancellationResult(TrajectoryCancellationStatus.ALREADY_STOPPED) return DEFAULT - coordinator.task_invoke.return_value = TrajectoryStatus(state=TrajectoryState.IDLE) + coordinator.task_invoke.return_value = TrajectoryStatus( + header=header_now(), state=TrajectoryState.IDLE + ) coordinator.task_invoke.side_effect = invoke return coordinator diff --git a/dimos/manipulation/control/coordinator_client.py b/dimos/manipulation/control/coordinator_client.py index 93a22be8f0d..fc118d6dd36 100644 --- a/dimos/manipulation/control/coordinator_client.py +++ b/dimos/manipulation/control/coordinator_client.py @@ -58,9 +58,10 @@ from dimos.manipulation.planning.trajectory_generator.joint_trajectory_generator import ( JointTrajectoryGenerator, ) +from dimos.msgs.trajectory import sample_trajectory, trajectory_duration if TYPE_CHECKING: - from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory + from dimos_generated.trajectory_msgs.msg import JointTrajectory class CoordinatorClient: @@ -294,7 +295,7 @@ def preview_trajectory(trajectory: JointTrajectory, joint_names: list[str]) -> N print("\n" + "=" * 70) print("GENERATED TRAJECTORY") print("=" * 70) - print(f"Duration: {trajectory.duration:.3f}s") + print(f"Duration: {trajectory_duration(trajectory):.3f}s") print(f"Points: {len(trajectory.points)}") print("-" * 70) print(f"{'Time':>6} | {header_str} (degrees)") @@ -302,8 +303,8 @@ def preview_trajectory(trajectory: JointTrajectory, joint_names: list[str]) -> N num_samples = min(10, max(len(trajectory.points) // 10, 5)) for i in range(num_samples + 1): - t = (i / num_samples) * trajectory.duration - q_ref, _ = trajectory.sample(t) + t = (i / num_samples) * trajectory_duration(trajectory) + q_ref, _ = sample_trajectory(trajectory, t) q_deg = [f"{math.degrees(q):7.1f}" for q in q_ref] print(f"{t:6.2f} | {' '.join(q_deg)}") @@ -432,7 +433,7 @@ def run(self) -> None: if result.status is TrajectoryExecutionStatus.ACCEPTED: print( "Trajectory accepted " - f"(expected duration: {self._generated_trajectory.duration:.2f}s)" + f"(expected duration: {trajectory_duration(self._generated_trajectory):.2f}s)" ) else: print(f"Failed to start trajectory: {result.message or result.status.name}") diff --git a/dimos/manipulation/control/servo_control/README.md b/dimos/manipulation/control/servo_control/README.md index d2bec661156..33197bea94d 100644 --- a/dimos/manipulation/control/servo_control/README.md +++ b/dimos/manipulation/control/servo_control/README.md @@ -117,7 +117,7 @@ print("Target reached!") ### Using Quaternions ```python -from dimos.msgs.geometry_msgs import Quaternion +from dimos_generated.geometry_msgs.msg import Quaternion # Create quaternion (identity rotation) quat = Quaternion(x=0, y=0, z=0, w=1) @@ -131,19 +131,20 @@ controller.set_target_pose( ### Using PoseStamped Messages ```python -from dimos.msgs.geometry_msgs import PoseStamped +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos.msgs.geometry import vector_array, quaternion_array +from dimos.msgs.time import header_now # Create target pose target = PoseStamped( - frame_id="world", - position=[0.3, 0.2, 0.5], - orientation=[0, 0, 0, 1] # quaternion + header=header_now(frame_id="world"), + pose=Pose(position=Point(x=0.3,y=0.2,z=0.5)), ) # Option 1: Via RPC controller.set_target_pose( - position=list(target.position), - orientation=list(target.orientation) + position=vector_array(target.pose.position).tolist(), + orientation=quaternion_array(target.pose.orientation).tolist() ) # Option 2: Via topic (if connected) diff --git a/dimos/manipulation/control/servo_control/cartesian_motion_controller.py b/dimos/manipulation/control/servo_control/cartesian_motion_controller.py index c0d858a0e40..457cd4e97a0 100644 --- a/dimos/manipulation/control/servo_control/cartesian_motion_controller.py +++ b/dimos/manipulation/control/servo_control/cartesian_motion_controller.py @@ -31,19 +31,19 @@ import time from typing import Any +from dimos_generated.dimos_msgs.msg import JointCommand, RobotState +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Twist, Vector3 +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header + from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out from dimos.manipulation.control.arm_driver_spec import ArmDriverSpec -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.JointCommand import JointCommand -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.sensor_msgs.RobotState import RobotState +from dimos.manipulation.planning.utils.kinematics_utils import compute_pose_error +from dimos.msgs.geometry import pose_matrix, quaternion_euler, quaternion_from_euler +from dimos.msgs.time import time_from_seconds from dimos.utils.logging_config import setup_logger from dimos.utils.simple_controller import PIDController @@ -257,12 +257,18 @@ def set_target_pose( # Detect if orientation is euler (3 elements) or quaternion (4 elements) if len(orientation) == 3: # Convert euler to quaternion using Pose's built-in conversion - euler_angles = Vector3(orientation[0], orientation[1], orientation[2]) - quat = Quaternion.from_euler(euler_angles) + euler_angles = Vector3(x=orientation[0], y=orientation[1], z=orientation[2]) + quat = quaternion_from_euler(euler_angles.x, euler_angles.y, euler_angles.z) orientation = [quat.x, quat.y, quat.z, quat.w] target = PoseStamped( - ts=time.time(), frame_id=frame_id, position=position, orientation=orientation + header=Header(frame_id=frame_id, stamp=time_from_seconds(time.time())), + pose=Pose( + position=Point(x=position[0], y=position[1], z=position[2]), + orientation=Quaternion( + x=orientation[0], y=orientation[1], z=orientation[2], w=orientation[3] + ), + ), ) with self._target_lock: @@ -309,7 +315,7 @@ def is_converged(self) -> bool: if not target_pose or not current_pose: return False - pos_error, ori_error = self._compute_pose_error(current_pose, target_pose) + pos_error, ori_error = self._compute_pose_error(current_pose, target_pose.pose) return ( pos_error < self.config.position_tolerance and ori_error < self.config.orientation_tolerance @@ -397,20 +403,21 @@ def _control_loop(self) -> None: current_pose_list[2] / 1000.0, ] euler_angles = Vector3( - current_pose_list[3], current_pose_list[4], current_pose_list[5] + x=current_pose_list[3], y=current_pose_list[4], z=current_pose_list[5] ) - quat = Quaternion.from_euler(euler_angles) + quat = quaternion_from_euler(euler_angles.x, euler_angles.y, euler_angles.z) self._current_tcp_pose = Pose( - position=position_m, - orientation=[quat.x, quat.y, quat.z, quat.w], + position=Point(x=position_m[0], y=position_m[1], z=position_m[2]), + orientation=Quaternion(x=quat.x, y=quat.y, z=quat.z, w=quat.w), ) # Publish current pose for target setters to use current_pose_stamped = PoseStamped( - ts=current_time, - frame_id="world", - position=position_m, - orientation=[quat.x, quat.y, quat.z, quat.w], + header=Header(frame_id="world", stamp=time_from_seconds(current_time)), + pose=Pose( + position=Point(x=position_m[0], y=position_m[1], z=position_m[2]), + orientation=Quaternion(x=quat.x, y=quat.y, z=quat.z, w=quat.w), + ), ) self.current_pose.publish(current_pose_stamped) else: @@ -439,7 +446,7 @@ def _control_loop(self) -> None: # Compute Cartesian error pos_error_mag, ori_error_mag = self._compute_pose_error( - self._current_tcp_pose, target_pose + self._current_tcp_pose, target_pose.pose ) # Log error periodically (every 1 second) @@ -447,7 +454,7 @@ def _control_loop(self) -> None: self._last_error_log_time = 0.0 if current_time - self._last_error_log_time > 1.0: logger.info( - f"Curr=[{self._current_tcp_pose.x:.3f},{self._current_tcp_pose.y:.3f},{self._current_tcp_pose.z:.3f}]m Tgt=[{target_pose.x:.3f},{target_pose.y:.3f},{target_pose.z:.3f}]m Err={pos_error_mag * 1000:.1f}mm" + f"Curr=[{self._current_tcp_pose.position.x:.3f},{self._current_tcp_pose.position.y:.3f},{self._current_tcp_pose.position.z:.3f}]m Tgt=[{target_pose.pose.position.x:.3f},{target_pose.pose.position.y:.3f},{target_pose.pose.position.z:.3f}]m Err={pos_error_mag * 1000:.1f}mm" ) self._last_error_log_time = current_time @@ -486,7 +493,7 @@ def _control_loop(self) -> None: # Generate Cartesian velocity command cartesian_twist = self._compute_cartesian_velocity( - self._current_tcp_pose, target_pose, dt + self._current_tcp_pose, target_pose.pose, dt ) # Publish debug twist @@ -505,12 +512,12 @@ def _control_loop(self) -> None: # Convert Pose to xArm format: [x, y, z, roll, pitch, yaw] # Note: xArm IK expects position in mm, so convert from m to mm next_pose_list = [ - next_pose.x * 1000.0, # m to mm - next_pose.y * 1000.0, # m to mm - next_pose.z * 1000.0, # m to mm - next_pose.roll, - next_pose.pitch, - next_pose.yaw, + next_pose.position.x * 1000.0, # m to mm + next_pose.position.y * 1000.0, # m to mm + next_pose.position.z * 1000.0, # m to mm + quaternion_euler(next_pose.orientation)[0], + quaternion_euler(next_pose.orientation)[1], + quaternion_euler(next_pose.orientation)[2], ] logger.debug( @@ -541,7 +548,7 @@ def _control_loop(self) -> None: # Publish joint command joint_cmd = JointCommand( - timestamp=current_time, + header=Header(stamp=time_from_seconds(current_time)), positions=list(target_joints), ) @@ -573,24 +580,7 @@ def _compute_pose_error(self, current_pose: Pose, target_pose: Pose) -> tuple[fl Returns: Tuple of (position_error_magnitude, orientation_error_magnitude) """ - # Position error (Euclidean distance) - pos_error = Vector3( - target_pose.x - current_pose.x, - target_pose.y - current_pose.y, - target_pose.z - current_pose.z, - ) - pos_error_mag = math.sqrt(pos_error.x**2 + pos_error.y**2 + pos_error.z**2) - - # Orientation error (angle between quaternions) - # q_error = q_current^-1 * q_target - q_current_inv = current_pose.orientation.conjugate() - q_error = q_current_inv * target_pose.orientation - - # Extract angle from axis-angle representation - # For quaternion [x, y, z, w], angle = 2 * acos(w) - ori_error_mag = 2 * math.acos(min(1.0, abs(q_error.w))) - - return pos_error_mag, ori_error_mag + return compute_pose_error(pose_matrix(current_pose), pose_matrix(target_pose)) def _compute_cartesian_velocity( self, current_pose: Pose, target_pose: Pose, dt: float @@ -607,9 +597,9 @@ def _compute_cartesian_velocity( Twist message with linear and angular velocities """ # Position error - error_x = target_pose.x - current_pose.x - error_y = target_pose.y - current_pose.y - error_z = target_pose.z - current_pose.z + error_x = target_pose.position.x - current_pose.position.x + error_y = target_pose.position.y - current_pose.position.y + error_z = target_pose.position.z - current_pose.position.z # Compute linear velocities via PID vel_x = self._pid_x.update(error_x, dt) # type: ignore[no-untyped-call] @@ -618,9 +608,18 @@ def _compute_cartesian_velocity( # Orientation error (convert to euler for simpler PID) # This is an approximation; axis-angle would be more accurate - error_roll = self._normalize_angle(target_pose.roll - current_pose.roll) - error_pitch = self._normalize_angle(target_pose.pitch - current_pose.pitch) - error_yaw = self._normalize_angle(target_pose.yaw - current_pose.yaw) + error_roll = self._normalize_angle( + quaternion_euler(target_pose.orientation)[0] + - quaternion_euler(current_pose.orientation)[0] + ) + error_pitch = self._normalize_angle( + quaternion_euler(target_pose.orientation)[1] + - quaternion_euler(current_pose.orientation)[1] + ) + error_yaw = self._normalize_angle( + quaternion_euler(target_pose.orientation)[2] + - quaternion_euler(current_pose.orientation)[2] + ) # Compute angular velocities via PID omega_x = self._pid_roll.update(error_roll, dt) # type: ignore[no-untyped-call] @@ -628,7 +627,8 @@ def _compute_cartesian_velocity( omega_z = self._pid_yaw.update(error_yaw, dt) # type: ignore[no-untyped-call] return Twist( - linear=Vector3(vel_x, vel_y, vel_z), angular=Vector3(omega_x, omega_y, omega_z) + linear=Vector3(x=vel_x, y=vel_y, z=vel_z), + angular=Vector3(x=omega_x, y=omega_y, z=omega_z), ) def _integrate_velocity(self, current_pose: Pose, velocity: Twist, dt: float) -> Pose: @@ -645,27 +645,27 @@ def _integrate_velocity(self, current_pose: Pose, velocity: Twist, dt: float) -> """ # Integrate position (simple Euler integration) next_position = Vector3( - current_pose.x + velocity.linear.x * dt, - current_pose.y + velocity.linear.y * dt, - current_pose.z + velocity.linear.z * dt, + x=current_pose.position.x + velocity.linear.x * dt, + y=current_pose.position.y + velocity.linear.y * dt, + z=current_pose.position.z + velocity.linear.z * dt, ) # Integrate orientation (simple euler integration - good for small dt) - next_roll = current_pose.roll + velocity.angular.x * dt - next_pitch = current_pose.pitch + velocity.angular.y * dt - next_yaw = current_pose.yaw + velocity.angular.z * dt + next_roll = quaternion_euler(current_pose.orientation)[0] + velocity.angular.x * dt + next_pitch = quaternion_euler(current_pose.orientation)[1] + velocity.angular.y * dt + next_yaw = quaternion_euler(current_pose.orientation)[2] + velocity.angular.z * dt - euler_angles = Vector3(next_roll, next_pitch, next_yaw) - next_orientation = Quaternion.from_euler(euler_angles) + euler_angles = Vector3(x=next_roll, y=next_pitch, z=next_yaw) + next_orientation = quaternion_from_euler(euler_angles.x, euler_angles.y, euler_angles.z) return Pose( - position=next_position, - orientation=[ - next_orientation.x, - next_orientation.y, - next_orientation.z, - next_orientation.w, - ], + position=Point(x=next_position.x, y=next_position.y, z=next_position.z), + orientation=Quaternion( + x=next_orientation.x, + y=next_orientation.y, + z=next_orientation.z, + w=next_orientation.w, + ), ) @staticmethod diff --git a/dimos/manipulation/control/test_generated_legacy_controllers.py b/dimos/manipulation/control/test_generated_legacy_controllers.py new file mode 100644 index 00000000000..75dc43288b7 --- /dev/null +++ b/dimos/manipulation/control/test_generated_legacy_controllers.py @@ -0,0 +1,106 @@ +# 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. + +from types import SimpleNamespace +from unittest.mock import MagicMock + +from dimos_generated.dimos_msgs.msg import JointCommand +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import pytest + +from dimos.manipulation.control.servo_control.cartesian_motion_controller import ( + CartesianMotionController, +) +from dimos.manipulation.control.trajectory_controller.joint_trajectory_controller import ( + JointTrajectoryController, +) +from dimos.msgs.time import duration_from_seconds +from dimos.msgs.trajectory import TrajectoryState + + +def test_trajectory_loop_publishes_generated_sample_with_source_time( + monkeypatch: pytest.MonkeyPatch, +) -> None: + controller = JointTrajectoryController() + controller._trajectory = JointTrajectory( + joint_names=["a", "b"], + points=[ + JointTrajectoryPoint(positions=[0.0, 0.0], time_from_start=duration_from_seconds(0)), + JointTrajectoryPoint(positions=[1.0, 2.0], time_from_start=duration_from_seconds(1)), + ], + ) + controller._state = TrajectoryState.EXECUTING + controller._start_time = 100.0 + controller.joint_position_command = MagicMock() + controller.joint_position_command.publish.side_effect = ( + lambda message: controller._stop_event.set() + ) + monkeypatch.setattr( + "dimos.manipulation.control.trajectory_controller.joint_trajectory_controller.time", + SimpleNamespace(time=lambda: 100.5, sleep=lambda seconds: None), + ) + monkeypatch.setattr("dimos.msgs.time.time.time_ns", lambda: 100500000000) + try: + controller._execution_loop() + controller.joint_position_command.publish.assert_called_once() + message = JointCommand.decode( + controller.joint_position_command.publish.call_args.args[0].encode() + ) + assert list(message.positions) == [0.5, 1.0] + assert (message.header.stamp.sec, message.header.stamp.nanosec) == (100, 500000000) + finally: + controller.stop() + + +def test_cartesian_generated_command_preserves_joint_count_and_pid_clamp( + monkeypatch: pytest.MonkeyPatch, +) -> None: + controller = CartesianMotionController() + driver = MagicMock() + driver.get_forward_kinematics.return_value = (0, [0.0] * 6) + driver.get_inverse_kinematics.return_value = (0, [0.1, 0.2, 0.3]) + controller._arm_driver = driver + controller._latest_joint_state = JointState(position=[0.0, 0.0]) + controller._target_pose_ = PoseStamped( + pose=Pose(position=Point(x=0.1), orientation=Quaternion(w=1.0)) + ) + controller._is_tracking = True + controller._last_target_time = 100.0 + controller.current_pose = MagicMock() + controller.cartesian_velocity = MagicMock() + controller.joint_position_command = MagicMock() + controller.joint_position_command.publish.side_effect = ( + lambda message: controller._stop_event.set() + ) + monkeypatch.setattr( + "dimos.manipulation.control.servo_control.cartesian_motion_controller.time", + SimpleNamespace(time=lambda: 100.5), + ) + try: + controller._control_loop() + controller.joint_position_command.publish.assert_called_once() + message = JointCommand.decode( + controller.joint_position_command.publish.call_args.args[0].encode() + ) + assert list(message.positions) == [0.1, 0.2] + assert (message.header.stamp.sec, message.header.stamp.nanosec) == (100, 500000000) + velocity = controller.cartesian_velocity.publish.call_args.args[0] + assert velocity.linear.x == pytest.approx(controller.config.max_linear_velocity) + # FK/IK only invoked on the fake driver; no driver actuation interface exists here. + assert driver.get_forward_kinematics.call_count == 1 + assert driver.get_inverse_kinematics.call_count == 1 + finally: + controller.stop() diff --git a/dimos/manipulation/control/trajectory_controller/joint_trajectory_controller.py b/dimos/manipulation/control/trajectory_controller/joint_trajectory_controller.py index 5676c741393..52f6c30b923 100644 --- a/dimos/manipulation/control/trajectory_controller/joint_trajectory_controller.py +++ b/dimos/manipulation/control/trajectory_controller/joint_trajectory_controller.py @@ -33,15 +33,16 @@ import time from typing import Any +from dimos_generated.dimos_msgs.msg import JointCommand, RobotState, TrajectoryStatus +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.trajectory_msgs.msg import JointTrajectory + from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.sensor_msgs.JointCommand import JointCommand -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.sensor_msgs.RobotState import RobotState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState, TrajectoryStatus +from dimos.msgs.time import duration_from_seconds, header_now +from dimos.msgs.trajectory import TrajectoryState, sample_trajectory, trajectory_duration from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -178,7 +179,7 @@ def execute_trajectory(self, trajectory: JointTrajectory) -> bool: return False # Validate trajectory - if trajectory is None or trajectory.duration <= 0: + if trajectory is None or trajectory_duration(trajectory) <= 0: logger.warning("Invalid trajectory: None or zero duration") return False @@ -198,7 +199,7 @@ def execute_trajectory(self, trajectory: JointTrajectory) -> bool: logger.info( f"Executing trajectory: {len(trajectory.points)} points, " - f"duration={trajectory.duration:.3f}s" + f"duration={trajectory_duration(trajectory):.3f}s" ) return True @@ -255,18 +256,19 @@ def get_status(self) -> TrajectoryStatus: if self._trajectory is not None and self._state == TrajectoryState.EXECUTING: time_elapsed = time.time() - self._start_time - time_remaining = max(0.0, self._trajectory.duration - time_elapsed) + time_remaining = max(0.0, trajectory_duration(self._trajectory) - time_elapsed) progress = ( - min(1.0, time_elapsed / self._trajectory.duration) - if self._trajectory.duration > 0 + min(1.0, time_elapsed / trajectory_duration(self._trajectory)) + if trajectory_duration(self._trajectory) > 0 else 1.0 ) return TrajectoryStatus( + header=header_now(), state=self._state, progress=progress, - time_elapsed=time_elapsed, - time_remaining=time_remaining, + time_elapsed=duration_from_seconds(time_elapsed), + time_remaining=duration_from_seconds(time_remaining), error=self._error_message, ) @@ -281,7 +283,7 @@ def _on_robot_state(self, msg: RobotState) -> None: def _on_trajectory(self, msg: JointTrajectory) -> None: """Callback when trajectory is received via topic.""" logger.info( - f"Received trajectory via topic: {len(msg.points)} points, duration={msg.duration:.3f}s" + f"Received trajectory via topic: {len(msg.points)} points, duration={trajectory_duration(msg):.3f}s" ) self.execute_trajectory(msg) @@ -313,17 +315,17 @@ def _execution_loop(self) -> None: if self._trajectory is None: self._state = TrajectoryState.FAULT logger.error("Trajectory is None during execution") - elif t >= self._trajectory.duration: + elif t >= trajectory_duration(self._trajectory): self._state = TrajectoryState.COMPLETED logger.info( - f"Trajectory completed: duration={self._trajectory.duration:.3f}s" + f"Trajectory completed: duration={trajectory_duration(self._trajectory):.3f}s" ) else: # Sample trajectory - q_ref, _qd_ref = self._trajectory.sample(t) + q_ref, _qd_ref = sample_trajectory(self._trajectory, t) # Create and publish command (outside lock would be better but simpler here) - cmd = JointCommand(positions=q_ref, timestamp=time.time()) + cmd = JointCommand(positions=q_ref, header=header_now()) # Publish - must release lock first for thread safety trajectory_active = True diff --git a/dimos/manipulation/execution_manager.py b/dimos/manipulation/execution_manager.py index 78f37be4119..d9535409222 100644 --- a/dimos/manipulation/execution_manager.py +++ b/dimos/manipulation/execution_manager.py @@ -21,6 +21,9 @@ import threading import time +from dimos_generated.dimos_msgs.msg import TrajectoryStatus +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint + from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.control.coordinator import ControlCoordinator from dimos.control.tasks.trajectory_task.trajectory_task import ( @@ -32,9 +35,7 @@ ) from dimos.manipulation.manipulation_spec import ExecutionResult, ExecutionStatus from dimos.manipulation.planning.spec.models import GeneratedPlan -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint -from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState, TrajectoryStatus +from dimos.msgs.trajectory import TrajectoryState from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -249,7 +250,7 @@ def cancel(self, timeout: float = 1.0) -> ExecutionResult: ExecutionStatus.FAULT, }: return latest - if status.state is TrajectoryState.IDLE: + if status.state == TrajectoryState.IDLE: result = ExecutionResult(ExecutionStatus.NO_EXECUTION, cancellation.message) self._store(result, active=False) return result @@ -306,7 +307,7 @@ def _combine(self, statuses: dict[str, TrajectoryStatus]) -> ExecutionResult: other for other, other_status in statuses.items() if other != task - and other_status.state is TrajectoryState.EXECUTING + and other_status.state == TrajectoryState.EXECUTING and not self._cancel_task(other) ] if failed: @@ -317,7 +318,7 @@ def _combine(self, statuses: dict[str, TrajectoryStatus]) -> ExecutionResult: ) mapped = self._result_from_status(status).status return ExecutionResult(mapped, f"{task}: {status.error}", trajectory_status=primary) - if all(status.state is TrajectoryState.COMPLETED for status in statuses.values()): + if all(status.state == TrajectoryState.COMPLETED for status in statuses.values()): return ExecutionResult(ExecutionStatus.COMPLETED, trajectory_status=primary) return ExecutionResult(ExecutionStatus.EXECUTING, trajectory_status=primary) @@ -409,7 +410,7 @@ def _result_from_status(status: TrajectoryStatus) -> ExecutionResult: TrajectoryState.COMPLETED: ExecutionStatus.COMPLETED, TrajectoryState.ABORTED: ExecutionStatus.ABORTED, TrajectoryState.FAULT: ExecutionStatus.FAULT, - }[status.state] + }[TrajectoryState(status.state)] return ExecutionResult(mapped, status.error, trajectory_status=status) def _store(self, result: ExecutionResult, *, active: bool, run_id: int | None = None) -> None: @@ -428,7 +429,7 @@ def _prepare_trajectory(self, plan: GeneratedPlan) -> list[tuple[str, JointTraje if not plan.is_success(): raise _PlanRejectedError("Generated plan status is not successful") - names = plan.trajectory.joint_names + names = list(plan.trajectory.joint_names) unknown = [name for name in names if name not in self._joint_names] if unknown: raise _PlanRejectedError(f"Generated trajectory has unknown joints: {unknown}") @@ -476,12 +477,18 @@ def _columns(trajectory: JointTrajectory, columns: list[int]) -> JointTrajectory return JointTrajectory( joint_names=[trajectory.joint_names[index] for index in columns], points=[ - TrajectoryPoint( + JointTrajectoryPoint( time_from_start=point.time_from_start, positions=[point.positions[index] for index in columns], - velocities=[point.velocities[index] for index in columns], + velocities=[point.velocities[index] for index in columns] + if point.velocities + else [], + accelerations=[point.accelerations[index] for index in columns] + if point.accelerations + else [], + effort=[point.effort[index] for index in columns] if point.effort else [], ) for point in trajectory.points ], - timestamp=trajectory.timestamp, + header=trajectory.header, ) diff --git a/dimos/manipulation/grasping/grasp_gen_spec.py b/dimos/manipulation/grasping/grasp_gen_spec.py index d3071d76d8e..2ac410fd0a4 100644 --- a/dimos/manipulation/grasping/grasp_gen_spec.py +++ b/dimos/manipulation/grasping/grasp_gen_spec.py @@ -14,9 +14,10 @@ from typing import Protocol -from dimos.msgs.geometry_msgs.PoseArray import PoseArray -from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos_generated.dimos_msgs.msg import GraspCandidateArray +from dimos_generated.geometry_msgs.msg import PoseArray +from dimos_generated.sensor_msgs.msg import PointCloud2 + from dimos.spec.utils import Spec diff --git a/dimos/manipulation/grasping/grasp_gen_x/module.py b/dimos/manipulation/grasping/grasp_gen_x/module.py index 7a608e7a90d..4e73f119615 100644 --- a/dimos/manipulation/grasping/grasp_gen_x/module.py +++ b/dimos/manipulation/grasping/grasp_gen_x/module.py @@ -18,6 +18,8 @@ from typing import Annotated, Literal, TypeAlias +from dimos_generated.dimos_msgs.msg import GraspCandidateArray +from dimos_generated.sensor_msgs.msg import PointCloud2 import numpy as np from pydantic import Field, FiniteFloat, field_validator @@ -27,8 +29,6 @@ IsolatedPythonModuleConfig, ) from dimos.manipulation.grasping.grasp_gen_spec import GraspGenSpec -from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.protocol.service.spec import BaseConfig GRASPGENX_MODEL_REPO = "adithyamurali/GraspGenXModel" diff --git a/dimos/manipulation/grasping/grasp_gen_x/test_module.py b/dimos/manipulation/grasping/grasp_gen_x/test_module.py index 55d7e4c301f..edcbdf69d0d 100644 --- a/dimos/manipulation/grasping/grasp_gen_x/test_module.py +++ b/dimos/manipulation/grasping/grasp_gen_x/test_module.py @@ -23,15 +23,15 @@ import sys from typing import Any +from dimos_generated.dimos_msgs.msg import GraspCandidate, GraspCandidateArray +from dimos_generated.geometry_msgs.msg import Point, Pose +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from dimos.manipulation.grasping.grasp_gen_spec import GraspGenSpec, LegacyGraspGenSpec from dimos.manipulation.grasping.grasp_gen_x.module import GraspGenXConfig -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.manipulation_msgs.GraspCandidate import GraspCandidate -from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray -from dimos.msgs.std_msgs.Header import Header +from dimos.msgs.time import time_from_seconds def config(**overrides: Any) -> GraspGenXConfig: @@ -68,15 +68,20 @@ def test_public_adapter_import_does_not_load_optional_runtime() -> None: def test_messages_round_trip_empty_and_score() -> None: - value = GraspCandidateArray(Header(3.0, "camera"), [GraspCandidate(Pose(1, 2, 3), 0.25)]) + value = GraspCandidateArray( + header=Header(stamp=time_from_seconds(3.0), frame_id="camera"), + candidates=[GraspCandidate(pose=Pose(position=Point(x=1, y=2)), score=0.25)], + ) decoded = GraspCandidateArray.decode(value.encode()) assert decoded.header.frame_id == "camera" - assert decoded.header.timestamp == pytest.approx(3.0) + assert decoded.header.stamp == time_from_seconds(3.0) assert decoded.candidates[0].score == pytest.approx(0.25) assert ( GraspCandidateArray.decode( - GraspCandidateArray(Header(3.0, "camera"), []).encode() + GraspCandidateArray( + header=Header(stamp=time_from_seconds(3.0), frame_id="camera"), candidates=[] + ).encode() ).candidates == [] ) diff --git a/dimos/manipulation/grasping/grasping.py b/dimos/manipulation/grasping/grasping.py index 9fd1cc9178a..2e0683efc97 100644 --- a/dimos/manipulation/grasping/grasping.py +++ b/dimos/manipulation/grasping/grasping.py @@ -19,20 +19,22 @@ from __future__ import annotations +import math from typing import TYPE_CHECKING +from dimos_generated.geometry_msgs.msg import PoseArray + from dimos.agents.annotation import skill from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import Out from dimos.manipulation.grasping.grasp_gen_spec import LegacyGraspGenSpec -from dimos.msgs.geometry_msgs.PoseArray import PoseArray +from dimos.msgs.geometry import quaternion_euler from dimos.perception.experimental.object_scene_registration_spec import ObjectSceneRegistrationSpec from dimos.utils.logging_config import setup_logger -from dimos.utils.transform_utils import quaternion_to_euler if TYPE_CHECKING: - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + from dimos_generated.sensor_msgs.msg import PointCloud2 logger = setup_logger() @@ -129,9 +131,9 @@ def _format_grasp_result(self, grasps: PoseArray, object_name: str) -> str: """Format grasp result for agent/human consumption.""" best = grasps.poses[0] pos = best.position - rpy = quaternion_to_euler(best.orientation, degrees=True) + rpy = tuple(math.degrees(value) for value in quaternion_euler(best.orientation)) return ( f"Generated {len(grasps.poses)}" f"Best grasp: pos=({pos.x:.4f}, {pos.y:.4f}, {pos.z:.4f}), " - f"rpy=({rpy.x:.1f}, {rpy.y:.1f}, {rpy.z:.1f}) degrees" + f"rpy=({rpy[0]:.1f}, {rpy[1]:.1f}, {rpy[2]:.1f}) degrees" ) diff --git a/dimos/manipulation/grasping/heuristic_grasp.py b/dimos/manipulation/grasping/heuristic_grasp.py index aff5442661e..74fec3113db 100644 --- a/dimos/manipulation/grasping/heuristic_grasp.py +++ b/dimos/manipulation/grasping/heuristic_grasp.py @@ -18,19 +18,18 @@ import math +from dimos_generated.dimos_msgs.msg import GraspCandidate, GraspCandidateArray +from dimos_generated.geometry_msgs.msg import Point, Pose +from dimos_generated.sensor_msgs.msg import PointCloud2 import numpy as np from numpy.typing import NDArray from dimos.core.core import rpc from dimos.core.module import Module from dimos.manipulation.grasping.grasp_gen_spec import GraspGenSpec -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.manipulation_msgs.GraspCandidate import GraspCandidate -from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.std_msgs.Header import Header +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.pointcloud import pointcloud_xyz +from dimos.msgs.time import to_nanoseconds class HeuristicGraspModule(Module, GraspGenSpec): @@ -41,11 +40,10 @@ class HeuristicGraspModule(Module, GraspGenSpec): @rpc def propose_grasps(self, object_pointcloud: PointCloud2) -> GraspCandidateArray: - if object_pointcloud.ts is None or not math.isfinite(float(object_pointcloud.ts)): - raise ValueError("object pointcloud must have a finite timestamp") - if not object_pointcloud.frame_id: + to_nanoseconds(object_pointcloud.header.stamp) + if not object_pointcloud.header.frame_id: raise ValueError("object pointcloud frame_id must not be empty") - points = object_pointcloud.points_f32() + points = pointcloud_xyz(object_pointcloud).astype(np.float32) if points.ndim != 2 or points.shape[1] != 3 or len(points) < 3: raise ValueError("object pointcloud must contain at least three XYZ points") if not np.all(np.isfinite(points)): @@ -55,12 +53,14 @@ def propose_grasps(self, object_pointcloud: PointCloud2) -> GraspCandidateArray: center_xy = np.median(xy, axis=0) low_z, high_z = np.quantile(points[:, 2], [0.05, 0.95]) pose = Pose( - Vector3(float(center_xy[0]), float(center_xy[1]), float((low_z + high_z) / 2.0)), - Quaternion.from_euler(Vector3(-math.pi, 0.0, self._narrow_axis_yaw(xy))), + position=Point( + x=float(center_xy[0]), y=float(center_xy[1]), z=float((low_z + high_z) / 2.0) + ), + orientation=quaternion_from_euler(-math.pi, 0.0, self._narrow_axis_yaw(xy)), ) return GraspCandidateArray( - Header(float(object_pointcloud.ts), object_pointcloud.frame_id), - [GraspCandidate(pose, score=1.0)], + header=object_pointcloud.header, + candidates=[GraspCandidate(pose=pose, score=1.0)], ) @staticmethod diff --git a/dimos/manipulation/grasping/test_heuristic_grasp.py b/dimos/manipulation/grasping/test_heuristic_grasp.py index 9862b615c44..5fa2c160067 100644 --- a/dimos/manipulation/grasping/test_heuristic_grasp.py +++ b/dimos/manipulation/grasping/test_heuristic_grasp.py @@ -15,20 +15,25 @@ from collections.abc import Iterator import math +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from dimos.manipulation.grasping.grasp_gen_spec import GraspGenSpec from dimos.manipulation.grasping.heuristic_grasp import HeuristicGraspModule -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.geometry import pose_matrix +from dimos.msgs.pointcloud import pointcloud_from_xyz from dimos.spec.utils import spec_annotation_compliance def _cloud( - points: np.ndarray, *, frame_id: str = "world", timestamp: float | None = 1.0 + points: np.ndarray, *, frame_id: str = "world", stamp: Time | None = None ) -> PointCloud2: - return PointCloud2.from_numpy(points.astype(np.float32), frame_id=frame_id, timestamp=timestamp) + return pointcloud_from_xyz( + points.astype(np.float32), header=Header(frame_id=frame_id, stamp=stamp or Time(sec=1)) + ) @pytest.fixture @@ -58,13 +63,13 @@ def test_heuristic_grasp_proposes_centered_top_down_pose(module: HeuristicGraspM assert len(proposals.candidates) == 1 assert proposals.header.frame_id == "world" - assert proposals.header.timestamp == pytest.approx(1.0) + assert proposals.header.stamp == Time(sec=1) pose = proposals.candidates[0].pose assert pose.position.x == pytest.approx(0.0) assert pose.position.y == pytest.approx(0.0) assert pose.position.z == pytest.approx(0.15) - approach = pose.orientation.rotate_vector(Vector3(0.0, 0.0, 1.0)) - assert approach.z == pytest.approx(-1.0) + approach = pose_matrix(pose)[:3, 2] + assert approach[2] == pytest.approx(-1.0) assert proposals.candidates[0].score == pytest.approx(1.0) @@ -82,10 +87,10 @@ def test_heuristic_grasp_aligns_jaw_axis_with_narrow_axis(module: HeuristicGrasp ) ) - jaw_axis = proposals.candidates[0].pose.orientation.rotate_vector(Vector3(0.0, 1.0, 0.0)) - assert abs(jaw_axis.x) == pytest.approx(1.0) - assert jaw_axis.y == pytest.approx(0.0, abs=1e-6) - assert jaw_axis.z == pytest.approx(0.0, abs=1e-6) + jaw_axis = pose_matrix(proposals.candidates[0].pose)[:3, 1] + assert abs(jaw_axis[0]) == pytest.approx(1.0) + assert jaw_axis[1] == pytest.approx(0.0, abs=1e-6) + assert jaw_axis[2] == pytest.approx(0.0, abs=1e-6) def test_heuristic_grasp_canonicalizes_pca_eigenvector_sign( @@ -104,27 +109,23 @@ def test_heuristic_grasp_canonicalizes_pca_eigenvector_sign( @pytest.mark.parametrize( - "points, frame_id, timestamp, error", + "points, frame_id, stamp, error", [ - ( - np.asarray([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]), - "world", - 1.0, - "at least three", - ), - (np.asarray([[0.0, 0.0, math.nan]] * 3), "world", 1.0, "finite"), - (np.zeros((3, 3)), "", 1.0, "frame_id"), - (np.zeros((3, 3)), "world", None, "timestamp"), - (np.zeros((3, 3)), "world", math.nan, "finite timestamp"), - (np.zeros((3, 3)), "world", math.inf, "finite timestamp"), + (np.zeros((2, 3)), "world", Time(), "at least three"), + (np.asarray([[0.0, 0.0, math.nan]] * 3), "world", Time(), "finite"), + (np.zeros((3, 3)), "", Time(), "frame_id"), + (np.zeros((3, 3)), "world", Time(nanosec=1000000000), "nanosec"), ], ) -def test_heuristic_grasp_rejects_invalid_pointclouds( - module: HeuristicGraspModule, - points: np.ndarray, - frame_id: str, - timestamp: float | None, - error: str, -) -> None: +def test_heuristic_grasp_rejects_invalid_pointclouds(module, points, frame_id, stamp, error): with pytest.raises(ValueError, match=error): - module.propose_grasps(_cloud(points, frame_id=frame_id, timestamp=timestamp)) + module.propose_grasps(_cloud(points, frame_id=frame_id, stamp=stamp)) + + +@pytest.mark.parametrize("stamp", [Time(), Time(sec=1700000000, nanosec=123456789)]) +def test_heuristic_grasp_preserves_exact_header_without_aliasing(module, stamp): + cloud = _cloud(np.zeros((3, 3)), stamp=stamp) + proposals = module.propose_grasps(cloud) + assert proposals.header == cloud.header + proposals.header.frame_id = "changed" + assert cloud.header.frame_id == "world" diff --git a/dimos/manipulation/manipulation_spec.py b/dimos/manipulation/manipulation_spec.py index 9c5e0307dcd..c77ebc0a1d6 100644 --- a/dimos/manipulation/manipulation_spec.py +++ b/dimos/manipulation/manipulation_spec.py @@ -21,12 +21,13 @@ from enum import Enum, auto from typing import Protocol +from dimos_generated.dimos_msgs.msg import GraspCandidateArray, TrajectoryStatus +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import JointState + from dimos.control.tasks.trajectory_task.trajectory_task import TrajectoryExecutionResult from dimos.manipulation.planning.spec.models import GeneratedPlan, PlanningGroupID -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryStatus +from dimos.msgs.trajectory import trajectory_duration from dimos.spec.utils import Spec @@ -108,11 +109,20 @@ class PlanningGroupState: joint_presets: Mapping[str, JointState] = field(default_factory=dict) def __repr__(self) -> str: + joints = ( + None + if self.joints is None + else dict(zip(self.joints.name, self.joints.position, strict=True)) + ) + presets = { + name: dict(zip(state.name, state.position, strict=True)) + for name, state in self.joint_presets.items() + } return ( - f"PlanningGroupState(joints={self.joints!r}, " + f"PlanningGroupState(joints={joints!r}, " f"end_effector_pose={self.end_effector_pose!r}, " f"gripper_position={self.gripper_position!r}, " - f"joint_presets={dict(self.joint_presets)!r})" + f"joint_presets={presets!r})" ) @@ -153,7 +163,7 @@ def __repr__(self) -> str: return f"PlanResult({self.status.name}, message={self.message!r})" return ( f"PlanResult({self.status.name}, groups={self.plan.group_ids!r}, " - f"waypoints={len(self.plan.path)}, duration={self.plan.trajectory.duration:.3g}s, " + f"waypoints={len(self.plan.path)}, duration={trajectory_duration(self.plan.trajectory):.3g}s, " f"path_length={self.plan.path_length:.3g}, planning_time={self.plan.planning_time:.3g}s, " f"iterations={self.plan.iterations}, message={self.message!r})" ) diff --git a/dimos/manipulation/planning/README.md b/dimos/manipulation/planning/README.md index 3d997bebbb9..7b4182eb448 100644 --- a/dimos/manipulation/planning/README.md +++ b/dimos/manipulation/planning/README.md @@ -72,15 +72,15 @@ manual motion and explicit planning, preview, and execution through `arm.rpc`. from dimos.manipulation import ManipulationModule from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec import RobotModelConfig -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.JointState import JointState +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.geometry_msgs.msg import Quaternion +from dimos_generated.geometry_msgs.msg import Vector3 +from dimos_generated.sensor_msgs.msg import JointState from dimos.robot.assets.model import RobotModel config = RobotModelConfig( model=RobotModel.from_file("/path/to/xarm7.urdf"), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + base_pose=PoseStamped(), joint_names=["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"], base_link="link_base", planning_groups=[ diff --git a/dimos/manipulation/planning/groups/test_planning_groups.py b/dimos/manipulation/planning/groups/test_planning_groups.py index 84567f296b2..c93b5f8f331 100644 --- a/dimos/manipulation/planning/groups/test_planning_groups.py +++ b/dimos/manipulation/planning/groups/test_planning_groups.py @@ -17,6 +17,9 @@ from dataclasses import dataclass from pathlib import Path +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header import pytest from dimos.manipulation.planning.groups.discovery import ( @@ -34,7 +37,6 @@ normalize_joint_target, ) from dimos.manipulation.planning.spec.config import RobotModelConfig -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.robot.assets.model import JointDescription, RobotModel @@ -165,8 +167,8 @@ def test_normalize_joint_target_accepts_exact_or_unnamed_target() -> None: JointState(name=["left/j1", "left/j2"], position=[1.0, 2.0]), ) unnamed = normalize_joint_target(group, JointState(position=[3.0, 4.0])) - assert named.name == ["left/j1", "left/j2"] - assert unnamed.name == ["left/j1", "left/j2"] + assert list(named.name) == ["left/j1", "left/j2"] + assert list(unnamed.name) == ["left/j1", "left/j2"] with pytest.raises(ValueError, match="missing joints"): normalize_joint_target( group, @@ -177,10 +179,26 @@ def test_normalize_joint_target_accepts_exact_or_unnamed_target() -> None: def test_state_projection_requires_exact_canonical_names() -> None: state = JointState(name=["right/j1", "left/j2", "left/j1"], position=[3.0, 2.0, 1.0]) projected = filter_joint_state_to_selected_joints(state, ("left/j1", "right/j1")) - assert projected.name == ["left/j1", "right/j1"] - assert projected.position == [1.0, 3.0] + assert list(projected.name) == ["left/j1", "right/j1"] + assert list(projected.position) == [1.0, 3.0] assert joint_state_to_ordered_positions( state, joint_names=("left/j1", "left/j2", "right/j1") ).tolist() == [1.0, 2.0, 3.0] with pytest.raises(ValueError, match="missing"): filter_joint_state_to_selected_joints(state, ("missing",)) + + +def test_projection_and_target_normalization_copy_source_header(): + group = PlanningGroup("arm", ("arm/a", "arm/b"), "base", "tool") + source = JointState( + header=Header(stamp=Time(sec=-1, nanosec=999999999), frame_id="robot"), + name=["arm/b", "arm/a"], + position=[2.0, 1.0], + ) + projected = filter_joint_state_to_selected_joints(source, ["arm/a"]) + normalized = normalize_joint_target(group, source) + assert projected.header == normalized.header == source.header + assert list(projected.position) == [1.0] + assert list(normalized.position) == [1.0, 2.0] + source.header.frame_id = "changed" + assert projected.header.frame_id == normalized.header.frame_id == "robot" diff --git a/dimos/manipulation/planning/groups/utils.py b/dimos/manipulation/planning/groups/utils.py index cb59f885157..a92768aff10 100644 --- a/dimos/manipulation/planning/groups/utils.py +++ b/dimos/manipulation/planning/groups/utils.py @@ -16,11 +16,11 @@ from collections.abc import Sequence +from dimos_generated.sensor_msgs.msg import JointState import numpy as np from numpy.typing import NDArray from dimos.manipulation.planning.groups.models import PlanningGroup -from dimos.msgs.sensor_msgs.JointState import JointState def filter_joint_state_to_selected_joints( @@ -32,6 +32,7 @@ def filter_joint_state_to_selected_joints( if missing: raise ValueError(f"Joint state is missing selected joints: {missing}") return JointState( + header=joint_state.header, name=list(joint_names), position=[float(positions_by_name[name]) for name in joint_names], ) @@ -45,7 +46,9 @@ def normalize_joint_target(group: PlanningGroup, target: JointState) -> JointSta f"Target for '{group.id}' has {len(target.position)} positions, " f"expected {len(group.joint_names)}" ) - return JointState(name=list(group.joint_names), position=list(target.position)) + return JointState( + header=target.header, name=list(group.joint_names), position=list(target.position) + ) if len(target.name) != len(target.position): raise ValueError( f"Target for '{group.id}' has {len(target.name)} names but " @@ -59,6 +62,7 @@ def normalize_joint_target(group: PlanningGroup, target: JointState) -> JointSta if extra: raise ValueError(f"Target for '{group.id}' has extra joints: {sorted(extra)}") return JointState( + header=target.header, name=list(group.joint_names), position=[float(positions[name]) for name in group.joint_names], ) diff --git a/dimos/manipulation/planning/kinematics/drake_optimization_ik.py b/dimos/manipulation/planning/kinematics/drake_optimization_ik.py index e65afde96aa..80a4f9ee35b 100644 --- a/dimos/manipulation/planning/kinematics/drake_optimization_ik.py +++ b/dimos/manipulation/planning/kinematics/drake_optimization_ik.py @@ -19,6 +19,8 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, cast +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import JointState import numpy as np from dimos.manipulation.planning.groups.models import PlanningGroup @@ -31,9 +33,7 @@ from dimos.manipulation.planning.spec.models import IKResult from dimos.manipulation.planning.spec.protocols import WorldSpec from dimos.manipulation.planning.utils.kinematics_utils import compute_pose_error -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.geometry import pose_matrix from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -98,11 +98,8 @@ def solve( "DrakeOptimizationIK requires exactly one pose-targetable planning group for legacy solve()", ) - # Convert PoseStamped to 4x4 matrix via Transform - target_matrix = Transform( - translation=target_pose.position, - rotation=target_pose.orientation, - ).to_matrix() + # Convert PoseStamped to 4x4 matrix via TransformStamped + target_matrix = pose_matrix(target_pose.pose) # Get joint limits lower_limits, upper_limits = world.get_prepared_model().joint_space.position_limits() @@ -113,7 +110,7 @@ def solve( seed = world.get_joint_state(ctx) # Extract joint names and seed positions - joint_names = seed.name + joint_names = list(seed.name) seed_positions = np.array(seed.position, dtype=np.float64) # Target transform @@ -201,10 +198,7 @@ def solve_pose_targets( ) lower_limits, upper_limits = world.get_prepared_model().joint_space.position_limits() - target_matrix = Transform( - translation=request.target_pose.position, - rotation=request.target_pose.orientation, - ).to_matrix() + target_matrix = pose_matrix(request.target_pose.pose) target_transform = RigidTransform(target_matrix) locked_positions = { index: float(request.seed_positions[index]) @@ -328,7 +322,7 @@ def _solve_single( joint_solution = np.clip(joint_solution, lower_limits, upper_limits) # Compute actual error using FK - solution_state = JointState({"name": joint_names, "position": joint_solution.tolist()}) + solution_state = JointState(name=joint_names, position=joint_solution.tolist()) with world.scratch_context() as ctx: world.set_joint_state(ctx, solution_state) actual_matrix = world.get_link_pose(ctx, target_frame_name) @@ -356,7 +350,7 @@ def _create_success_result( ) -> IKResult: return IKResult( status=IKStatus.SUCCESS, - joint_state=JointState({"name": joint_names, "position": joint_positions.tolist()}), + joint_state=JointState(name=joint_names, position=joint_positions.tolist()), position_error=position_error, orientation_error=orientation_error, iterations=iterations, diff --git a/dimos/manipulation/planning/kinematics/jacobian_ik.py b/dimos/manipulation/planning/kinematics/jacobian_ik.py index 07029aeca99..fd37e8b44bb 100644 --- a/dimos/manipulation/planning/kinematics/jacobian_ik.py +++ b/dimos/manipulation/planning/kinematics/jacobian_ik.py @@ -42,17 +42,14 @@ compute_pose_error, damped_pseudoinverse, ) +from dimos.msgs.geometry import pose_matrix from dimos.utils.logging_config import setup_logger -from dimos.utils.transform_utils import pose_to_matrix if TYPE_CHECKING: from numpy.typing import NDArray -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.JointState import JointState +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist, Vector3 +from dimos_generated.sensor_msgs.msg import JointState logger = setup_logger() @@ -136,7 +133,7 @@ def solve( seed = world.get_joint_state(ctx) # Extract joint names for creating random seeds - joint_names = seed.name + joint_names = list(seed.name) best_result: IKResult | None = None best_error = float("inf") @@ -148,9 +145,7 @@ def solve( else: # Random seed within joint limits random_positions = np.random.uniform(lower_limits, upper_limits) - current_seed = JointState( - {"name": joint_names, "position": random_positions.tolist()} - ) + current_seed = JointState(name=joint_names, position=random_positions.tolist()) # Solve iterative IK result = self.solve_iterative( @@ -215,9 +210,7 @@ def solve_pose_targets( if request is None: return _create_failure_result(IKStatus.NO_SOLUTION, "Invalid pose target request") - full_seed = JointState( - {"name": request.joint_names, "position": request.seed_positions.tolist()} - ) + full_seed = JointState(name=request.joint_names, position=request.seed_positions.tolist()) result = self.solve_iterative( world=world, target_pose=request.target_pose, @@ -236,9 +229,7 @@ def solve_pose_targets( full_positions[request.group_indices] = np.asarray( result.joint_state.position, dtype=np.float64 ) - full_state = JointState( - {"name": request.joint_names, "position": full_positions.tolist()} - ) + full_state = JointState(name=request.joint_names, position=full_positions.tolist()) if not world.check_config_collision_free(full_state): return _create_failure_result(IKStatus.COLLISION, "IK solution is in collision") return result @@ -271,12 +262,9 @@ def solve_iterative( IKResult with solution or failure status """ # Convert to internal representation - target_matrix = Transform( - translation=target_pose.position, - rotation=target_pose.orientation, - ).to_matrix() + target_matrix = pose_matrix(target_pose.pose) current_joints = np.array(seed.position, dtype=np.float64) - joint_names = seed.name + joint_names = list(seed.name) active_joint_indices = active_joint_indices or list(range(len(joint_names))) result_joint_names = list(group.joint_names) if group is not None else joint_names @@ -286,15 +274,13 @@ def solve_iterative( for iteration in range(max_iterations): with world.scratch_context() as ctx: # Set current position (convert to JointState for API) - current_state = JointState( - {"name": joint_names, "position": current_joints.tolist()} - ) + current_state = JointState(name=joint_names, position=current_joints.tolist()) world.set_joint_state(ctx, current_state) if group is None: - current_pose = pose_to_matrix(world.get_ee_pose(ctx)) + current_pose = pose_matrix(world.get_ee_pose(ctx).pose) else: - current_pose = pose_to_matrix(world.get_group_ee_pose(ctx, group.id)) + current_pose = pose_matrix(world.get_group_ee_pose(ctx, group.id).pose) # Compute error pos_error, ori_error = compute_pose_error(current_pose, target_matrix) @@ -345,12 +331,12 @@ def solve_iterative( # Compute final error with world.scratch_context() as ctx: - final_state = JointState({"name": joint_names, "position": current_joints.tolist()}) + final_state = JointState(name=joint_names, position=current_joints.tolist()) world.set_joint_state(ctx, final_state) if group is None: - final_pose = pose_to_matrix(world.get_ee_pose(ctx)) + final_pose = pose_matrix(world.get_ee_pose(ctx).pose) else: - final_pose = pose_to_matrix(world.get_group_ee_pose(ctx, group.id)) + final_pose = pose_matrix(world.get_group_ee_pose(ctx, group.id).pose) pos_error, ori_error = compute_pose_error(final_pose, target_matrix) return _create_failure_result( @@ -416,7 +402,7 @@ def solve_differential( if max_ratio > 1.0: q_dot = q_dot / max_ratio - return JointState({"name": joint_names, "velocity": q_dot.tolist()}) + return JointState(name=joint_names, velocity=q_dot.tolist()) def solve_differential_position_only( self, @@ -462,7 +448,7 @@ def solve_differential_position_only( # Compute joint velocities q_dot = J_pinv @ vel_array - return JointState({"name": joint_names, "velocity": q_dot.tolist()}) + return JointState(name=joint_names, velocity=q_dot.tolist()) # Result Helpers @@ -478,7 +464,7 @@ def _create_success_result( """Create a successful IK result.""" return IKResult( status=IKStatus.SUCCESS, - joint_state=JointState({"name": joint_names, "position": joint_positions.tolist()}), + joint_state=JointState(name=joint_names, position=joint_positions.tolist()), position_error=position_error, orientation_error=orientation_error, iterations=iterations, diff --git a/dimos/manipulation/planning/kinematics/pink_ik.py b/dimos/manipulation/planning/kinematics/pink_ik.py index c182ec2b66c..667a562cefc 100644 --- a/dimos/manipulation/planning/kinematics/pink_ik.py +++ b/dimos/manipulation/planning/kinematics/pink_ik.py @@ -20,6 +20,8 @@ from dataclasses import replace from typing import TYPE_CHECKING +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import JointState import numpy as np import pink from pink.exceptions import NoSolutionFound @@ -41,8 +43,6 @@ from dimos.manipulation.planning.spec.models import IKResult from dimos.manipulation.planning.spec.protocols import WorldSpec from dimos.manipulation.planning.utils.kinematics_utils import compute_pose_error -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.JointState import JointState if TYPE_CHECKING: from numpy.typing import NDArray @@ -408,7 +408,7 @@ def _success( ) -> IKResult: return IKResult( status=IKStatus.SUCCESS, - joint_state=JointState({"name": joint_names, "position": joint_positions.tolist()}), + joint_state=JointState(name=joint_names, position=joint_positions.tolist()), position_error=position_error, orientation_error=orientation_error, iterations=iterations, @@ -430,10 +430,7 @@ def _selected_result(result: IKResult, selected_names: Sequence[str]) -> IKResul return IKResult( status=result.status, joint_state=JointState( - { - "name": list(selected_names), - "position": [float(positions[name]) for name in selected_names], - } + name=list(selected_names), position=[float(positions[name]) for name in selected_names] ), position_error=result.position_error, orientation_error=result.orientation_error, diff --git a/dimos/manipulation/planning/kinematics/pink_solver.py b/dimos/manipulation/planning/kinematics/pink_solver.py index fb21dad90dd..bb8f01a0997 100644 --- a/dimos/manipulation/planning/kinematics/pink_solver.py +++ b/dimos/manipulation/planning/kinematics/pink_solver.py @@ -33,15 +33,16 @@ ) raise ImportError(msg) from exc +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import JointState + from dimos.manipulation.planning.kinematics.config import PinkKinematicsConfig from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.joint_space import CoordinateTopology, JointSpace from dimos.manipulation.planning.spec.validation import PreparedRobotModel from dimos.manipulation.planning.utils.mesh_utils import prepare_urdf_for_drake -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.geometry import pose_matrix from dimos.utils.logging_config import setup_logger -from dimos.utils.transform_utils import pose_to_matrix if TYPE_CHECKING: from numpy.typing import NDArray @@ -314,8 +315,8 @@ def _current_frame_matrix( def _target_in_model_frame( self, config: RobotModelConfig, target_pose: PoseStamped ) -> NDArray[np.float64]: - target_world = pose_to_matrix(target_pose) - base_world = pose_to_matrix(config.base_pose) + target_world = pose_matrix(target_pose.pose) + base_world = pose_matrix(config.base_pose.pose) target_model: NDArray[np.float64] = np.asarray( np.linalg.inv(base_world) @ target_world, dtype=np.float64 ) diff --git a/dimos/manipulation/planning/kinematics/pinocchio_ik.py b/dimos/manipulation/planning/kinematics/pinocchio_ik.py index 79eb831eef1..cb8d64526c5 100644 --- a/dimos/manipulation/planning/kinematics/pinocchio_ik.py +++ b/dimos/manipulation/planning/kinematics/pinocchio_ik.py @@ -35,6 +35,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Any +from dimos_generated.geometry_msgs.msg import PoseStamped import numpy as np from numpy.linalg import norm, solve import pinocchio @@ -42,11 +43,9 @@ from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: + from dimos_generated.geometry_msgs.msg import Pose from numpy.typing import NDArray - from dimos.msgs.geometry_msgs.Pose import Pose - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - logger = setup_logger() @@ -211,8 +210,9 @@ def forward_kinematics(self, joint_positions: NDArray[np.floating[Any]]) -> pino def pose_to_se3(pose: Pose | PoseStamped) -> pinocchio.SE3: """Convert Pose or PoseStamped to pinocchio SE3""" - position = np.array([pose.x, pose.y, pose.z]) - quat = pose.orientation + value = pose.pose if isinstance(pose, PoseStamped) else pose + position = np.array([value.position.x, value.position.y, value.position.z]) + quat = value.orientation rotation = pinocchio.Quaternion(quat.w, quat.x, quat.y, quat.z).toRotationMatrix() return pinocchio.SE3(rotation, position) diff --git a/dimos/manipulation/planning/kinematics/test_drake_optimization_ik_selection.py b/dimos/manipulation/planning/kinematics/test_drake_optimization_ik_selection.py index 51133160860..c1a4b6354ed 100644 --- a/dimos/manipulation/planning/kinematics/test_drake_optimization_ik_selection.py +++ b/dimos/manipulation/planning/kinematics/test_drake_optimization_ik_selection.py @@ -17,6 +17,9 @@ from contextlib import contextmanager from pathlib import Path +from dimos_generated.geometry_msgs.msg import Pose, PoseStamped +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header import numpy as np from dimos.manipulation.planning.groups.models import PlanningGroup @@ -31,8 +34,6 @@ ) from dimos.manipulation.planning.spec.models import IKResult from dimos.manipulation.planning.spec.validation import PreparedRobotModel -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.robot.assets.model import LoadedRobotModel, RobotModel @@ -43,10 +44,8 @@ def __init__(self) -> None: joint_names=["arm/base", "arm/shoulder", "arm/elbow", "arm/wrist"], ) self.current_state = JointState( - { - "name": ["arm/base", "arm/shoulder", "arm/elbow", "arm/wrist"], - "position": [1.0, 2.0, 3.0, 4.0], - } + name=["arm/base", "arm/shoulder", "arm/elbow", "arm/wrist"], + position=[1.0, 2.0, 3.0, 4.0], ) self.collision_checked_state: JointState | None = None self.prepared = PreparedRobotModel( @@ -96,10 +95,8 @@ def fake_solve_single(self, **kwargs) -> IKResult: return IKResult( status=IKStatus.SUCCESS, joint_state=JointState( - { - "name": ["arm/base", "arm/shoulder", "arm/elbow", "arm/wrist"], - "position": [10.0, 20.0, 30.0, 40.0], - } + name=["arm/base", "arm/shoulder", "arm/elbow", "arm/wrist"], + position=[10.0, 20.0, 30.0, 40.0], ), position_error=0.0, orientation_error=0.0, @@ -110,8 +107,8 @@ def fake_solve_single(self, **kwargs) -> IKResult: result = DrakeOptimizationIK().solve_pose_targets( world=world, # type: ignore[arg-type] - pose_targets={group: PoseStamped()}, - seed=JointState({"name": ["arm/shoulder"], "position": [22.0]}), + pose_targets={group: PoseStamped(header=Header(frame_id=""), pose=Pose())}, + seed=JointState(name=["arm/shoulder"], position=[22.0]), check_collision=False, max_attempts=1, ) diff --git a/dimos/manipulation/planning/kinematics/test_jacobian_ik_selection.py b/dimos/manipulation/planning/kinematics/test_jacobian_ik_selection.py index 8f8b000da8d..e75e7551b61 100644 --- a/dimos/manipulation/planning/kinematics/test_jacobian_ik_selection.py +++ b/dimos/manipulation/planning/kinematics/test_jacobian_ik_selection.py @@ -19,6 +19,9 @@ from contextlib import nullcontext from pathlib import Path +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header import numpy as np from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupDefinition @@ -31,15 +34,16 @@ JointSpace, ) from dimos.manipulation.planning.spec.validation import PreparedRobotModel -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.robot.assets.model import LoadedRobotModel, RobotModel def _pose(x: float = 0.0) -> PoseStamped: - return PoseStamped(position=Vector3(x, 0.0, 0.0), orientation=Quaternion(0.0, 0.0, 0.0, 1.0)) + return PoseStamped( + header=Header(frame_id=""), + pose=Pose( + position=Point(x=x, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + ) def _group(tip_link: str | None = "tool") -> PlanningGroup: @@ -93,7 +97,7 @@ def scratch_context(self) -> nullcontext[None]: def get_joint_state(self, ctx: object) -> JointState: return JointState( - {"name": ["arm/joint_a", "arm/joint_b", "arm/gripper"], "position": [0.0, 0.0, 0.9]} + name=["arm/joint_a", "arm/joint_b", "arm/gripper"], position=[0.0, 0.0, 0.9] ) def set_joint_state(self, ctx: object, joint_state: JointState) -> None: @@ -117,7 +121,7 @@ def test_solve_pose_targets_filters_to_group_and_uses_group_world_methods() -> N world=world, pose_targets={_group(): _pose()}, seed=JointState( - {"name": ["arm/joint_a", "arm/joint_b", "arm/gripper"], "position": [0.0, 0.0, 0.9]} + name=["arm/joint_a", "arm/joint_b", "arm/gripper"], position=[0.0, 0.0, 0.9] ), max_attempts=1, ) diff --git a/dimos/manipulation/planning/kinematics/test_pink_ik.py b/dimos/manipulation/planning/kinematics/test_pink_ik.py index 467f7e469f9..b2370c81485 100644 --- a/dimos/manipulation/planning/kinematics/test_pink_ik.py +++ b/dimos/manipulation/planning/kinematics/test_pink_ik.py @@ -24,6 +24,9 @@ from types import MappingProxyType, ModuleType, SimpleNamespace from typing import Any, cast +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header import numpy as np from pink.exceptions import NoSolutionFound import pytest @@ -64,12 +67,8 @@ ) from dimos.manipulation.planning.spec.models import IKResult from dimos.manipulation.planning.spec.validation import PreparedRobotModel, prepare_robot_model -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.geometry import pose_from_matrix from dimos.robot.assets.model import LoadedRobotModel, PlanarBaseDefinition, RobotModel -from dimos.utils.transform_utils import matrix_to_pose _TRACKING_ERROR_RAD = np.deg2rad(10.0) @@ -396,7 +395,10 @@ def _install_fake_modules(mocker: MockerFixture, converge: bool = True) -> _Fake def _robot_config() -> RobotModelConfig: return RobotModelConfig( model=RobotModel.from_file(Path("/tmp/fake.urdf")), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0)), + base_pose=PoseStamped( + header=Header(frame_id=""), + pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), + ), joint_names=["joint_a", "joint_b", "joint_c"], base_link="base", planning_groups=[ @@ -548,7 +550,7 @@ def scratch_context(self) -> nullcontext[None]: def get_joint_state(self, ctx: object) -> JointState: self.joint_state_calls += 1 - return JointState({"name": ["joint_b", "joint_c", "joint_a"], "position": [0.0, 0.0, 0.0]}) + return JointState(name=["joint_b", "joint_c", "joint_a"], position=[0.0, 0.0, 0.0]) def get_joint_limits(self) -> tuple[np.ndarray, np.ndarray]: return np.array([-1.0, -1.0, -1.0]), np.array([1.0, 1.0, 1.0]) @@ -775,7 +777,7 @@ def apply_increment(**kwargs: Any) -> None: commands.append( ik.step_frame_targets( robot_model=_robot_config(), - frame_targets={"tool": PoseStamped()}, + frame_targets={"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())}, controlled_joints=["joint_a"], command_state=commands[-1], measured_state=initial, @@ -816,7 +818,7 @@ def apply_increment(**kwargs: Any) -> None: commands.append( ik.step_frame_targets( robot_model=_robot_config(), - frame_targets={"tool": PoseStamped()}, + frame_targets={"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())}, controlled_joints=["joint_a"], command_state=commands[-1], measured_state=initial, @@ -853,7 +855,7 @@ def apply_increment(**kwargs: Any) -> None: initial = JointState(name=["joint_a"], position=[0.0]) first = ik.step_frame_targets( robot_model=_robot_config(), - frame_targets={"tool": PoseStamped()}, + frame_targets={"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())}, controlled_joints=["joint_a"], command_state=initial, measured_state=initial, @@ -865,7 +867,7 @@ def apply_increment(**kwargs: Any) -> None: second = ik.step_frame_targets( robot_model=_robot_config(), - frame_targets={"tool": PoseStamped()}, + frame_targets={"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())}, controlled_joints=["joint_a"], command_state=first, measured_state=JointState(name=["joint_a"], position=[-0.1]), @@ -875,8 +877,8 @@ def apply_increment(**kwargs: Any) -> None: joint_command_filter_cutoff_hz=5.0, ) - assert first.position == pytest.approx([0.2]) - assert second.position == pytest.approx([0.05]) + assert list(first.position) == pytest.approx([0.2]) + assert list(second.position) == pytest.approx([0.05]) def test_step_frame_targets_preserves_controlled_joint_order( @@ -893,8 +895,11 @@ def test_step_frame_targets_preserves_controlled_joint_order( robot_model=_robot_config(), frame_targets={ "tool": PoseStamped( - position=Vector3(0.1, 0.2, 0.3), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + header=Header(frame_id=""), + pose=Pose( + position=Point(x=0.1, y=0.2, z=0.3), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) }, controlled_joints=["joint_c", "joint_a"], @@ -905,7 +910,7 @@ def test_step_frame_targets_preserves_controlled_joint_order( ) assert result.name == ["joint_c", "joint_a"] - assert result.position == pytest.approx([0.1, 0.08]) + assert list(result.position) == pytest.approx([0.1, 0.08]) def test_step_frame_targets_builds_both_frame_tasks_with_tuning( @@ -932,12 +937,15 @@ def test_step_frame_targets_builds_both_frame_tasks_with_tuning( robot_model=_robot_config(), frame_targets={ "tool": PoseStamped( - position=Vector3(0.1, 0.2, 0.3), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + header=Header(frame_id=""), + pose=Pose( + position=Point(x=0.1, y=0.2, z=0.3), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), "base": PoseStamped( - position=Vector3(), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + header=Header(frame_id=""), + pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), ), }, controlled_joints=["joint_a", "joint_b", "joint_c"], @@ -1009,12 +1017,16 @@ def test_streaming_reuses_task_stack_across_steps( mocker.patch.object(ik, "_get_control_context", return_value=context) create_tasks = mocker.spy(ik, "_create_tasks") first_target = PoseStamped( - position=Vector3(0.1, 0.2, 0.3), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + header=Header(frame_id=""), + pose=Pose( + position=Point(x=0.1, y=0.2, z=0.3), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), ) second_target = PoseStamped( - position=Vector3(0.3, 0.2, 0.1), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + header=Header(frame_id=""), + pose=Pose( + position=Point(x=0.3, y=0.2, z=0.1), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), ) ik.step_frame_targets( @@ -1053,8 +1065,11 @@ def test_task_hooks_receive_read_only_stack_and_successful_velocity( _robot_config(), { "tool": PoseStamped( - position=Vector3(0.1, 0.2, 0.3), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + header=Header(frame_id=""), + pose=Pose( + position=Point(x=0.1, y=0.2, z=0.3), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) }, ["joint_a", "joint_b", "joint_c"], @@ -1082,7 +1097,7 @@ def test_after_solve_hook_is_not_called_when_solver_raises( with pytest.raises(RuntimeError, match="no solution"): ik.step_frame_targets( _robot_config(), - {"tool": PoseStamped()}, + {"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())}, ["joint_a", "joint_b", "joint_c"], JointState( name=["joint_a", "joint_b", "joint_c"], @@ -1136,7 +1151,7 @@ def step_outside_limits(**kwargs: Any) -> None: result = ik.step_frame_targets( robot_model=_robot_config(), - frame_targets={"tool": PoseStamped()}, + frame_targets={"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())}, controlled_joints=["joint_a", "joint_b", "joint_c"], command_state=JointState( name=["joint_a", "joint_b", "joint_c"], @@ -1149,7 +1164,7 @@ def step_outside_limits(**kwargs: Any) -> None: max_command_tracking_error_rad=_TRACKING_ERROR_RAD, ) - assert result.position == pytest.approx([-0.9999, 0.1, _TRACKING_ERROR_RAD]) + assert list(result.position) == pytest.approx([-0.9999, 0.1, _TRACKING_ERROR_RAD]) def test_step_frame_targets_rejects_feedback_beyond_tolerance( @@ -1166,7 +1181,7 @@ def test_step_frame_targets_rejects_feedback_beyond_tolerance( with pytest.raises(PinkJointLimitError, match="joint_a.*lower limit"): ik.step_frame_targets( robot_model=_robot_config(), - frame_targets={"tool": PoseStamped()}, + frame_targets={"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())}, controlled_joints=["joint_a"], command_state=JointState(name=["joint_a"], position=[0.0]), measured_state=JointState(name=["joint_a"], position=[-1.0011]), @@ -1194,14 +1209,14 @@ def step_unbounded(**kwargs: Any) -> None: result = ik.step_frame_targets( robot_model=_robot_config(), - frame_targets={"tool": PoseStamped()}, + frame_targets={"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())}, controlled_joints=["joint_b"], command_state=JointState(name=["joint_b"], position=[5.0]), measured_state=JointState(name=["joint_b"], position=[5.0]), max_command_tracking_error_rad=_TRACKING_ERROR_RAD, ) - assert result.position == pytest.approx([5.1]) + assert list(result.position) == pytest.approx([5.1]) def test_validate_frame_targets_rejects_margin_wider_than_joint_range( @@ -1241,7 +1256,7 @@ def test_step_frame_targets_rejects_unknown_frame(mocker: MockerFixture, tmp_pat with pytest.raises(ValueError, match="missing_frame"): _StreamingTestPinkIK(PinkIKConfig()).step_frame_targets( robot_model=config, - frame_targets={"missing_frame": PoseStamped()}, + frame_targets={"missing_frame": PoseStamped(header=Header(frame_id=""), pose=Pose())}, controlled_joints=config.joint_names, command_state=JointState(name=config.joint_names, position=[0.0, 0.0, 0.0]), measured_state=JointState(name=config.joint_names, position=[0.0, 0.0, 0.0]), @@ -1274,7 +1289,7 @@ def test_solve_targets_returns_successful_ik_result(mocker: MockerFixture) -> No assert result.status == IKStatus.SUCCESS assert result.joint_state is not None assert result.joint_state.name == ["joint_a", "joint_b", "joint_c"] - assert result.joint_state.position == pytest.approx([0.2, 0.1, 0.3]) + assert list(result.joint_state.position) == pytest.approx([0.2, 0.1, 0.3]) def test_planning_uses_named_stack_and_task_lifecycle_hooks( @@ -1385,11 +1400,16 @@ def is_collision_free(self, ctx: object) -> bool: [0.05003819, 0.15888552, 0.11027428, -0.10991712, -0.07993349, 0.14942138, -0.19789388] ) target_q = ik._q_from_dimos_positions(context, target_positions) - target_pose = matrix_to_pose(ik._current_frame_matrix(context, target_q)) + target_pose = pose_from_matrix(ik._current_frame_matrix(context, target_q)) result = ik.solve_pose_targets( cast("Any", World()), - {group: PoseStamped(position=target_pose.position, orientation=target_pose.orientation)}, + { + group: PoseStamped( + header=Header(frame_id=""), + pose=Pose(position=target_pose.position, orientation=target_pose.orientation), + ) + }, seed=seed, check_collision=False, max_attempts=1, @@ -1410,8 +1430,11 @@ def test_solve_rejects_collision_candidate(mocker: MockerFixture) -> None: result = ik.solve( world=cast("Any", _FakeWorld(collision_free=False)), target_pose=PoseStamped( - position=Vector3(0.1, 0.0, 0.0), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + header=Header(frame_id=""), + pose=Pose( + position=Point(x=0.1, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), check_collision=True, max_attempts=1, @@ -1452,8 +1475,11 @@ def fake_solve_targets(**_: object) -> IKResult: result = ik.solve( world=cast("Any", _FakeWorld(collision_free=True)), target_pose=PoseStamped( - position=Vector3(0.1, 0.0, 0.0), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + header=Header(frame_id=""), + pose=Pose( + position=Point(x=0.1, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), check_collision=True, max_attempts=2, @@ -1520,7 +1546,7 @@ def test_solve_pose_targets_uses_group_tip_and_filters_group_joints( return_value=IKResult( status=IKStatus.SUCCESS, joint_state=JointState( - {"name": ["joint_a", "joint_b", "joint_c"], "position": [0.1, 0.2, 0.3]} + name=["joint_a", "joint_b", "joint_c"], position=[0.1, 0.2, 0.3] ), ), ) @@ -1530,10 +1556,11 @@ def test_solve_pose_targets_uses_group_tip_and_filters_group_joints( world=cast("Any", world), pose_targets={ world.groups["manipulator"]: PoseStamped( - position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + header=Header(frame_id=""), + pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), ) }, - seed=JointState({"name": ["joint_a", "joint_b", "joint_c"], "position": [0.0, 0.0, 0.0]}), + seed=JointState(name=["joint_a", "joint_b", "joint_c"], position=[0.0, 0.0, 0.0]), max_attempts=1, ) @@ -1553,7 +1580,8 @@ def test_solve_pose_targets_rejects_group_without_tip(mocker: MockerFixture) -> world=cast("Any", world), pose_targets={ world.groups["no_tip"]: PoseStamped( - position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + header=Header(frame_id=""), + pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), ) }, ) @@ -1571,7 +1599,7 @@ def test_solve_pose_targets_partial_seed_reads_world_state(mocker: MockerFixture return_value=IKResult( status=IKStatus.SUCCESS, joint_state=JointState( - {"name": ["joint_a", "joint_b", "joint_c"], "position": [0.1, 0.2, 0.3]} + name=["joint_a", "joint_b", "joint_c"], position=[0.1, 0.2, 0.3] ), ), ) @@ -1581,10 +1609,11 @@ def test_solve_pose_targets_partial_seed_reads_world_state(mocker: MockerFixture world=cast("Any", world), pose_targets={ world.groups["manipulator"]: PoseStamped( - position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + header=Header(frame_id=""), + pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), ) }, - seed=JointState({"name": ["joint_a"], "position": [0.0]}), + seed=JointState(name=["joint_a"], position=[0.0]), max_attempts=1, ) @@ -1602,7 +1631,7 @@ def test_solve_pose_targets_multi_target_uses_multi_frame_solve(mocker: MockerFi return_value=IKResult( status=IKStatus.SUCCESS, joint_state=JointState( - {"name": ["joint_a", "joint_b", "joint_c"], "position": [0.1, 0.2, 0.3]} + name=["joint_a", "joint_b", "joint_c"], position=[0.1, 0.2, 0.3] ), position_error=0.0, orientation_error=0.0, @@ -1613,13 +1642,15 @@ def test_solve_pose_targets_multi_target_uses_multi_frame_solve(mocker: MockerFi world=cast("Any", world), pose_targets={ world.groups["manipulator"]: PoseStamped( - position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + header=Header(frame_id=""), + pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), ), world.groups["wrist"]: PoseStamped( - position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + header=Header(frame_id=""), + pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), ), }, - seed=JointState({"name": ["joint_a", "joint_b", "joint_c"], "position": [0.0, 0.0, 0.0]}), + seed=JointState(name=["joint_a", "joint_b", "joint_c"], position=[0.0, 0.0, 0.0]), max_attempts=1, ) @@ -1642,7 +1673,7 @@ def test_solve_pose_targets_checks_multi_group_solution_together( return_value=IKResult( status=IKStatus.SUCCESS, joint_state=JointState( - {"name": ["joint_a", "joint_b", "joint_c"], "position": [0.1, 0.2, 0.3]} + name=["joint_a", "joint_b", "joint_c"], position=[0.1, 0.2, 0.3] ), ), ) @@ -1651,18 +1682,15 @@ def test_solve_pose_targets_checks_multi_group_solution_together( world=cast("Any", world), pose_targets={ world.groups["manipulator"]: PoseStamped( - position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + header=Header(frame_id=""), + pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), ), world.groups["wrist"]: PoseStamped( - position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + header=Header(frame_id=""), + pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), ), }, - seed=JointState( - { - "name": ["joint_a", "joint_b", "joint_c"], - "position": [0.0, 0.0, 0.0], - } - ), + seed=JointState(name=["joint_a", "joint_b", "joint_c"], position=[0.0, 0.0, 0.0]), max_attempts=1, ) @@ -1681,7 +1709,7 @@ def test_solve_pose_targets_auxiliary_only_retains_seed_selection_order( world=cast("Any", world), pose_targets={}, auxiliary_groups=[world.groups["no_tip"], world.groups["manipulator"]], - seed=JointState({"name": ["joint_a", "joint_b", "joint_c"], "position": [0.1, 0.2, 0.3]}), + seed=JointState(name=["joint_a", "joint_b", "joint_c"], position=[0.1, 0.2, 0.3]), ) assert result.status == IKStatus.SUCCESS @@ -1691,7 +1719,7 @@ def test_solve_pose_targets_auxiliary_only_retains_seed_selection_order( def _solved_joint_state() -> JointState: - return JointState({"name": ["joint_a", "joint_b", "joint_c"], "position": [0.1, 0.2, 0.3]}) + return JointState(name=["joint_a", "joint_b", "joint_c"], position=[0.1, 0.2, 0.3]) def _qp_infeasible() -> NoSolutionFound: @@ -1726,8 +1754,11 @@ def fake_solve_targets(**_: object) -> IKResult: result = ik.solve( world=cast("Any", _FakeWorld(collision_free=True)), target_pose=PoseStamped( - position=Vector3(0.1, 0.0, 0.0), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + header=Header(frame_id=""), + pose=Pose( + position=Point(x=0.1, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), check_collision=True, max_attempts=2, @@ -1752,8 +1783,11 @@ def fake_solve_targets(**_: object) -> IKResult: result = ik.solve( world=cast("Any", _FakeWorld(collision_free=True)), target_pose=PoseStamped( - position=Vector3(0.1, 0.0, 0.0), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + header=Header(frame_id=""), + pose=Pose( + position=Point(x=0.1, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), check_collision=True, max_attempts=3, @@ -1774,8 +1808,11 @@ def test_solve_does_not_retry_unexpected_exception(mocker: MockerFixture) -> Non result = ik.solve( world=cast("Any", _FakeWorld(collision_free=True)), target_pose=PoseStamped( - position=Vector3(0.1, 0.0, 0.0), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + header=Header(frame_id=""), + pose=Pose( + position=Point(x=0.1, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), check_collision=True, max_attempts=3, @@ -1796,8 +1833,11 @@ def test_solve_mapping_value_error_fails_without_retrying(mocker: MockerFixture) result = ik.solve( world=cast("Any", _FakeWorld(collision_free=True)), target_pose=PoseStamped( - position=Vector3(0.1, 0.0, 0.0), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + header=Header(frame_id=""), + pose=Pose( + position=Point(x=0.1, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), check_collision=True, max_attempts=5, @@ -1809,7 +1849,7 @@ def test_solve_mapping_value_error_fails_without_retrying(mocker: MockerFixture) def _pose_targets_seed() -> JointState: - return JointState({"name": ["joint_a", "joint_b", "joint_c"], "position": [0.0, 0.0, 0.0]}) + return JointState(name=["joint_a", "joint_b", "joint_c"], position=[0.0, 0.0, 0.0]) def test_solve_pose_targets_retries_after_no_solution_found(mocker: MockerFixture) -> None: @@ -1838,7 +1878,8 @@ def fake_solve_targets(**_: object) -> IKResult: world=cast("Any", world), pose_targets={ world.groups["manipulator"]: PoseStamped( - position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + header=Header(frame_id=""), + pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), ) }, seed=_pose_targets_seed(), @@ -1866,7 +1907,8 @@ def fake_solve_targets(**_: object) -> IKResult: world=cast("Any", world), pose_targets={ world.groups["manipulator"]: PoseStamped( - position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + header=Header(frame_id=""), + pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), ) }, seed=_pose_targets_seed(), @@ -1892,7 +1934,8 @@ def test_solve_pose_targets_does_not_retry_unexpected_exception( world=cast("Any", world), pose_targets={ world.groups["manipulator"]: PoseStamped( - position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + header=Header(frame_id=""), + pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), ) }, seed=_pose_targets_seed(), @@ -1918,7 +1961,8 @@ def test_solve_pose_targets_mapping_value_error_fails_without_retrying( world=cast("Any", world), pose_targets={ world.groups["manipulator"]: PoseStamped( - position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + header=Header(frame_id=""), + pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), ) }, seed=_pose_targets_seed(), diff --git a/dimos/manipulation/planning/kinematics/utils.py b/dimos/manipulation/planning/kinematics/utils.py index 8fff0b12b6d..d32e40c02d4 100644 --- a/dimos/manipulation/planning/kinematics/utils.py +++ b/dimos/manipulation/planning/kinematics/utils.py @@ -17,6 +17,8 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import JointState import numpy as np from numpy.typing import NDArray @@ -25,8 +27,6 @@ from dimos.manipulation.planning.spec.enums import IKStatus from dimos.manipulation.planning.spec.models import IKResult from dimos.manipulation.planning.spec.protocols import WorldSpec -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.JointState import JointState @dataclass(frozen=True) diff --git a/dimos/manipulation/planning/monitor/test_robot_state_monitor.py b/dimos/manipulation/planning/monitor/test_robot_state_monitor.py new file mode 100644 index 00000000000..f22d99c99d4 --- /dev/null +++ b/dimos/manipulation/planning/monitor/test_robot_state_monitor.py @@ -0,0 +1,78 @@ +# 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. + +import threading + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header +import pytest + +from dimos.manipulation.planning.monitor.robot_state_monitor import RobotStateMonitor +from dimos.manipulation.planning.spec.protocols import WorldSpec + + +@pytest.fixture() +def monitor(mocker): + world = mocker.Mock(spec=WorldSpec) + monitor = RobotStateMonitor(world, threading.RLock(), ["arm/a", "arm/b"]) + monitor.start() + try: + yield monitor, world + finally: + monitor.stop() + + +def test_generated_feedback_reorders_joints_and_preserves_source_header(monitor): + monitor, world = monitor + received = [] + monitor.add_state_callback(received.append) + source = JointState( + header=Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id="robot"), + name=["arm/b", "other", "arm/a"], + position=[2.0, 9.0, 1.0], + velocity=[0.2, 0.9, 0.1], + ) + monitor.on_joint_state(JointState.decode(source.encode())) + assert monitor.get_current_positions().tolist() == [1.0, 2.0] + assert monitor.get_current_velocities().tolist() == [0.1, 0.2] + expected = JointState( + header=source.header, name=["arm/a", "arm/b"], position=[1.0, 2.0], velocity=[0.1, 0.2] + ) + world.sync_from_joint_state.assert_called_once_with(expected) + assert received == [expected] + snapshot = monitor.get_current_joint_state() + assert snapshot == expected + snapshot.header.stamp.nanosec = 0 + snapshot.position[0] = 99 + received[0].position[0] = 42 + assert monitor.get_current_joint_state() == expected + source.header.stamp.nanosec = 0 + assert received[0].header.stamp.nanosec == 123456789 + + +def test_missing_joint_does_not_replace_valid_state(monitor): + monitor, world = monitor + monitor.on_joint_state(JointState(name=["arm/a", "arm/b"], position=[1.0, 2.0])) + monitor.on_joint_state(JointState(name=["arm/a"], position=[9.0])) + assert monitor.get_current_positions().tolist() == [1.0, 2.0] + assert world.sync_from_joint_state.call_count == 1 + + +def test_stopped_monitor_does_not_sync(monitor): + monitor, world = monitor + monitor.stop() + monitor.on_joint_state(JointState(name=["arm/a", "arm/b"], position=[1.0, 2.0])) + assert monitor.get_current_positions() is None + world.sync_from_joint_state.assert_not_called() diff --git a/dimos/manipulation/planning/monitor/world_obstacle_monitor.py b/dimos/manipulation/planning/monitor/world_obstacle_monitor.py index 41ff2b2fbcc..2f5f1a29da9 100644 --- a/dimos/manipulation/planning/monitor/world_obstacle_monitor.py +++ b/dimos/manipulation/planning/monitor/world_obstacle_monitor.py @@ -30,22 +30,23 @@ import time from typing import TYPE_CHECKING, Any +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion + from dimos.manipulation.planning.spec.enums import ObstacleType from dimos.manipulation.planning.spec.models import ( DEFAULT_OBSTACLE_RGBA, CollisionObjectMessage, Obstacle, ) -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.pointcloud import pointcloud_xyz from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: from collections.abc import Callable + from dimos_generated.vision_msgs.msg import Detection3D + from dimos.manipulation.planning.monitor.world_monitor import WorldMonitor - from dimos.msgs.vision_msgs.Detection3D import Detection3D from dimos.perception.experimental.object import Object logger = setup_logger() @@ -56,7 +57,7 @@ class WorldObstacleMonitor: This class handles updates from: - Explicit collision objects (CollisionObjectMessage) - - Perception detections (Detection3D from dimos.msgs.vision_msgs) + - Perception detections (generated Detection3D) ## Thread Safety @@ -267,7 +268,7 @@ def on_detections(self, detections: list[Detection3D]) -> None: - Removes obstacles for detections that are no longer present Args: - detections: List of Detection3D messages from dimos.msgs.vision_msgs + detections: List of generated Detection3D messages """ if not self._running: return @@ -313,11 +314,7 @@ def on_detections(self, detections: list[Detection3D]) -> None: def _detection3d_to_pose(self, detection: Detection3D) -> PoseStamped: """Convert Detection3D bbox.center to PoseStamped.""" - center = detection.bbox.center - return PoseStamped( - position=center.position, - orientation=center.orientation, - ) + return PoseStamped(header=detection.header, pose=detection.bbox.center) def _detection_to_obstacle(self, detection: Detection3D) -> Obstacle: """Convert Detection3D to Obstacle.""" @@ -644,8 +641,8 @@ def _object_to_obstacle(self, obj: object) -> Obstacle: pointcloud_to_convex_hull_obj, ) - points, _ = obj.pointcloud.as_numpy() - if points is not None and points.shape[0] >= 4: + points = pointcloud_xyz(obj.pointcloud) + if points.shape[0] >= 4: # Keyed on the object's stable unique id: rescans overwrite # in place, and no two objects share a file. hull = pointcloud_to_convex_hull_obj(points, cache_key=name) @@ -658,10 +655,13 @@ def _object_to_obstacle(self, obj: object) -> Obstacle: name=name, obstacle_type=ObstacleType.MESH, pose=PoseStamped( - ts=obj.pose.ts, - frame_id=obj.pose.frame_id, - position=Vector3(hull.centroid), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + header=obj.pose.header, + pose=Pose( + position=Point( + x=hull.centroid[0], y=hull.centroid[1], z=hull.centroid[2] + ), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), color=(0.2, 0.8, 0.2, 0.6), mesh_path=hull.path, @@ -673,7 +673,7 @@ def _object_to_obstacle(self, obj: object) -> Obstacle: return Obstacle( name=name, obstacle_type=ObstacleType.BOX, - pose=obj.pose or PoseStamped(position=obj.center), + pose=obj.pose, dimensions=(float(obj.size.x), float(obj.size.y), float(obj.size.z)), color=(0.2, 0.8, 0.2, 0.6), ) diff --git a/dimos/manipulation/planning/planners/roboplan_planner.py b/dimos/manipulation/planning/planners/roboplan_planner.py index 2c006997830..ecdd485859b 100644 --- a/dimos/manipulation/planning/planners/roboplan_planner.py +++ b/dimos/manipulation/planning/planners/roboplan_planner.py @@ -23,6 +23,8 @@ import numpy as np +from dimos.msgs.geometry import transform_matrix + try: import roboplan.cartesian_planning as roboplan_cartesian import roboplan.core as roboplan_core @@ -33,6 +35,9 @@ "Install the manipulation extra before selecting the roboplan backend." ) from exc +from dimos_generated.geometry_msgs.msg import PoseStamped, TransformStamped +from dimos_generated.sensor_msgs.msg import JointState + from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupSelection from dimos.manipulation.planning.planners.roboplan_config import ( RoboPlanCartesianPathConfig, @@ -55,11 +60,8 @@ RoboPlanContext, RoboPlanWorld, ) -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.geometry import pose_matrix from dimos.utils.logging_config import setup_logger -from dimos.utils.transform_utils import pose_to_matrix if TYPE_CHECKING: from numpy.typing import NDArray @@ -374,8 +376,8 @@ def _validate_cartesian_request( waypoint_type = ( PoseStamped if isinstance(target[0], PoseStamped) - else Transform - if isinstance(target[0], Transform) + else TransformStamped + if isinstance(target[0], TransformStamped) else None ) if waypoint_type is None or any( @@ -385,10 +387,10 @@ def _validate_cartesian_request( status=PlanningStatus.INVALID_GOAL, message=( f"Cartesian target for '{group_id}' must contain only PoseStamped " - "waypoints or only Transform waypoints" + "waypoints or only TransformStamped waypoints" ), ) - if any(waypoint.frame_id != "world" for waypoint in target): + if any(waypoint.header.frame_id != "world" for waypoint in target): return PlanningResult( status=PlanningStatus.UNSUPPORTED, message="Cartesian planning supports only world-frame waypoints", @@ -420,7 +422,7 @@ def _build_cartesian_path( if group.tip_link is None: raise ValueError(f"Planning group '{group.id}' has no TCP tip link") start_pose = self._world.get_group_ee_pose(ctx, group.id) - start_matrix = np.asarray(pose_to_matrix(start_pose), dtype=np.float64) + start_matrix = np.asarray(pose_matrix(start_pose.pose), dtype=np.float64) target_matrices = [ self._resolve_cartesian_waypoint(start_matrix, waypoint) for waypoint in target ] @@ -436,12 +438,12 @@ def _build_cartesian_path( def _resolve_cartesian_waypoint( self, start_matrix: NDArray[np.float64], - waypoint: PoseStamped | Transform, + waypoint: PoseStamped | TransformStamped, ) -> NDArray[np.float64]: if isinstance(waypoint, PoseStamped): - return np.asarray(pose_to_matrix(waypoint), dtype=np.float64) + return np.asarray(pose_matrix(waypoint.pose), dtype=np.float64) target_matrix = start_matrix.copy() - delta_matrix = np.asarray(waypoint.to_matrix(), dtype=np.float64) + delta_matrix = np.asarray(transform_matrix(waypoint.transform), dtype=np.float64) target_matrix[:3, 3] += delta_matrix[:3, 3] target_matrix[:3, :3] = delta_matrix[:3, :3] @ start_matrix[:3, :3] return target_matrix @@ -695,15 +697,15 @@ def _validate_shortcut_path( if not shortened_path: raise ValueError("RoboPlan path shortcutter returned an empty path") if not np.allclose( - shortened_path[0].position, - original_path[0].position, + np.asarray(shortened_path[0].position, dtype=np.float64), + np.asarray(original_path[0].position, dtype=np.float64), atol=1e-9, rtol=0.0, ): raise ValueError("RoboPlan path shortcutter changed the start configuration") if not np.allclose( - shortened_path[-1].position, - original_path[-1].position, + np.asarray(shortened_path[-1].position, dtype=np.float64), + np.asarray(original_path[-1].position, dtype=np.float64), atol=1e-9, rtol=0.0, ): diff --git a/dimos/manipulation/planning/planners/rrt_planner.py b/dimos/manipulation/planning/planners/rrt_planner.py index 796a1a6b164..8be423a13eb 100644 --- a/dimos/manipulation/planning/planners/rrt_planner.py +++ b/dimos/manipulation/planning/planners/rrt_planner.py @@ -25,6 +25,7 @@ import time from typing import TYPE_CHECKING +from dimos_generated.sensor_msgs.msg import JointState import numpy as np from dimos.manipulation.planning.groups.models import PlanningGroupSelection @@ -43,7 +44,6 @@ ) from dimos.manipulation.planning.spec.protocols import WorldSpec from dimos.manipulation.planning.utils.path_utils import compute_path_length -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -104,7 +104,7 @@ def plan_joint_path( # Extract positions as numpy arrays for internal computation q_start = np.array(start.position, dtype=np.float64) q_goal = np.array(goal.position, dtype=np.float64) - joint_names = start.name # Store for converting back to JointState + joint_names = list(start.name) # Store for converting back to JointState error = self._validate_inputs(world, start, goal) if error is not None: @@ -471,8 +471,8 @@ def _extend_tree( new_config = nearest.config + step_size * (diff / dist) # Check validity of edge using context-free method - start_state = JointState({"name": joint_names, "position": nearest.config.tolist()}) - end_state = JointState({"name": joint_names, "position": new_config.tolist()}) + start_state = JointState(name=joint_names, position=nearest.config.tolist()) + end_state = JointState(name=joint_names, position=new_config.tolist()) if world.check_edge_collision_free(start_state, end_state, self._collision_step_size): new_node = TreeNode(config=new_config, parent=nearest) nearest.children.append(new_node) @@ -519,7 +519,7 @@ def _extract_path( full_path_arrays = start_path + list(reversed(goal_path)) # Convert to list of JointState - return [JointState({"name": joint_names, "position": q.tolist()}) for q in full_path_arrays] + return [JointState(name=joint_names, position=q.tolist()) for q in full_path_arrays] def _simplify_path( self, diff --git a/dimos/manipulation/planning/planners/selected_joint_space.py b/dimos/manipulation/planning/planners/selected_joint_space.py index 2a1ac19bfa9..ff792ae0f84 100644 --- a/dimos/manipulation/planning/planners/selected_joint_space.py +++ b/dimos/manipulation/planning/planners/selected_joint_space.py @@ -18,6 +18,7 @@ from dataclasses import dataclass +from dimos_generated.sensor_msgs.msg import JointState import numpy as np from numpy.typing import NDArray @@ -25,7 +26,6 @@ from dimos.manipulation.planning.spec.joint_space import JointSpace from dimos.manipulation.planning.spec.models import JointPath from dimos.manipulation.planning.spec.protocols import WorldSpec -from dimos.msgs.sensor_msgs.JointState import JointState @dataclass(frozen=True) diff --git a/dimos/manipulation/planning/planners/test_rrt_planner_selection.py b/dimos/manipulation/planning/planners/test_rrt_planner_selection.py index 64ffb146ab9..fe5c85c0d92 100644 --- a/dimos/manipulation/planning/planners/test_rrt_planner_selection.py +++ b/dimos/manipulation/planning/planners/test_rrt_planner_selection.py @@ -20,6 +20,17 @@ import math from pathlib import Path +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from pytest_mock import MockerFixture @@ -40,16 +51,14 @@ JointSpace, ) from dimos.manipulation.planning.spec.validation import PreparedRobotModel -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.robot.assets.model import LoadedRobotModel, PlanarBaseDefinition, RobotModel def _pose() -> PoseStamped: - return PoseStamped(position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0)) + return PoseStamped( + header=Header(frame_id=""), + pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), + ) def _group(name: str, joints: tuple[str, ...]) -> PlanningGroup: @@ -103,9 +112,7 @@ def scratch_context(self) -> nullcontext[None]: return nullcontext(None) def get_joint_state(self, ctx: object) -> JointState: - return JointState( - {"name": ["arm/joint_a", "arm/joint_b", "arm/gripper"], "position": self.current} - ) + return JointState(name=["arm/joint_a", "arm/joint_b", "arm/gripper"], position=self.current) def get_joint_limits(self) -> tuple[np.ndarray, np.ndarray]: return np.array([-1.0, -1.0, -1.0]), np.array([1.0, 1.0, 1.0]) @@ -197,14 +204,14 @@ def check_config_collision_free(self, joint_state: JointState) -> bool: ("start", "goal", "expected_start", "expected_goal"), [ ( - JointState({"position": [0.1, 0.2]}), - JointState({"position": [0.3, 0.4]}), + JointState(position=[0.1, 0.2]), + JointState(position=[0.3, 0.4]), [0.1, 0.2], [0.3, 0.4], ), ( - JointState({"name": ["arm/joint_b", "arm/joint_a"], "position": [0.2, 0.1]}), - JointState({"name": ["arm/joint_b", "arm/joint_a"], "position": [0.4, 0.3]}), + JointState(name=["arm/joint_b", "arm/joint_a"], position=[0.2, 0.1]), + JointState(name=["arm/joint_b", "arm/joint_a"], position=[0.4, 0.3]), [0.1, 0.2], [0.3, 0.4], ), @@ -229,22 +236,20 @@ def test_plan_selected_joint_path_normalizes_target_forms( ("start", "goal", "status", "message"), [ ( - JointState({"name": ["arm/joint_a"], "position": [0.0]}), - JointState({"position": [0.0, 0.0]}), + JointState(name=["arm/joint_a"], position=[0.0]), + JointState(position=[0.0, 0.0]), PlanningStatus.INVALID_START, "missing", ), ( - JointState({"position": [0.0, 0.0]}), - JointState( - {"name": ["arm/joint_a", "arm/joint_b", "arm/extra"], "position": [0.0, 0.0, 0.0]} - ), + JointState(position=[0.0, 0.0]), + JointState(name=["arm/joint_a", "arm/joint_b", "arm/extra"], position=[0.0, 0.0, 0.0]), PlanningStatus.INVALID_GOAL, "extra", ), ( - JointState({"name": ["arm/joint_a", "joint_b"], "position": [0.0, 0.0]}), - JointState({"position": [0.0, 0.0]}), + JointState(name=["arm/joint_a", "joint_b"], position=[0.0, 0.0]), + JointState(position=[0.0, 0.0]), PlanningStatus.INVALID_START, "missing", ), @@ -270,8 +275,8 @@ def test_plan_selected_joint_path_rejects_noncanonical_names() -> None: result = RRTConnectPlanner().plan_selected_joint_path( _World(), selection, - JointState({"name": ["joint_a", "gripper"], "position": [0.0, 0.0]}), - JointState({"position": [0.1, 0.2]}), + JointState(name=["joint_a", "gripper"], position=[0.0, 0.0]), + JointState(position=[0.1, 0.2]), ) assert result.status == PlanningStatus.INVALID_START @@ -285,8 +290,8 @@ def test_plan_selected_joint_path_direct_edge_projects_full_state_with_unselecte result = RRTConnectPlanner().plan_selected_joint_path( world, PlanningGroupSelection.from_groups((group,)), - JointState({"position": [0.1, 0.2]}), - JointState({"position": [0.3, 0.4]}), + JointState(position=[0.1, 0.2]), + JointState(position=[0.3, 0.4]), ) assert result.status == PlanningStatus.SUCCESS @@ -305,11 +310,15 @@ def test_plan_cartesian_path_is_explicitly_unsupported() -> None: result = RRTConnectPlanner().plan_cartesian_path( _World(), selection, - JointState({"position": [0.0, 0.0]}), + JointState(position=[0.0, 0.0]), { group.id: ( - Transform.identity(), - Transform(translation=Vector3(0.1, 0.0, 0.0)), + TransformStamped(header=Header(frame_id="world"), child_frame_id=""), + TransformStamped( + header=Header(frame_id="world"), + transform=Transform(translation=Vector3(x=0.1, y=0.0, z=0.0)), + child_frame_id="", + ), ) }, RoboPlanCartesianPathConfig(), diff --git a/dimos/manipulation/planning/spec/joint_space.py b/dimos/manipulation/planning/spec/joint_space.py index 98cd89f3d3c..f1c0374ab4e 100644 --- a/dimos/manipulation/planning/spec/joint_space.py +++ b/dimos/manipulation/planning/spec/joint_space.py @@ -23,13 +23,13 @@ import math from typing import Annotated, Literal +from dimos_generated.sensor_msgs.msg import JointState import numpy as np from numpy.typing import NDArray from pydantic import ConfigDict, Field, model_validator from pydantic.dataclasses import dataclass as pydantic_dataclass from typing_extensions import Self -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.utils.trigonometry import angle_diff @@ -137,7 +137,7 @@ def normalize_positions( def from_joint_state(self, state: JointState) -> NDArray[np.float64]: """Extract and validate this space's coordinates from a named state.""" if not state.name: - return self.normalize_positions(state.position) + return self.normalize_positions(np.asarray(state.position, dtype=np.float64)) if len(state.name) != len(set(state.name)): raise ValueError("Joint state contains duplicate coordinate names") positions = dict(zip(state.name, state.position, strict=True)) diff --git a/dimos/manipulation/planning/spec/models.py b/dimos/manipulation/planning/spec/models.py index 125b5939fb7..0ae43a458bb 100644 --- a/dimos/manipulation/planning/spec/models.py +++ b/dimos/manipulation/planning/spec/models.py @@ -28,15 +28,14 @@ ) if TYPE_CHECKING: + from dimos_generated.geometry_msgs.msg import PoseStamped, TransformStamped + from dimos_generated.sensor_msgs.msg import JointState + from dimos_generated.trajectory_msgs.msg import JointTrajectory import numpy as np from numpy.typing import NDArray from dimos.manipulation.planning.groups.models import PlanningGroup from dimos.manipulation.planning.spec.validation import PreparedRobotModel - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - from dimos.msgs.geometry_msgs.Transform import Transform - from dimos.msgs.sensor_msgs.JointState import JointState - from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory PlanningGroupID: TypeAlias = str @@ -49,10 +48,10 @@ """List of joint states forming a path (each waypoint has names + positions)""" -CartesianWaypoint: TypeAlias = "PoseStamped | Transform" +CartesianWaypoint: TypeAlias = "PoseStamped | TransformStamped" """One absolute TCP pose or relative rigid displacement from the planning start.""" -CartesianTarget: TypeAlias = "Sequence[PoseStamped] | Sequence[Transform]" +CartesianTarget: TypeAlias = "Sequence[PoseStamped] | Sequence[TransformStamped]" """Ordered homogeneous Cartesian waypoints for one planning group.""" diff --git a/dimos/manipulation/planning/spec/protocols.py b/dimos/manipulation/planning/spec/protocols.py index 4891ef23305..bf80e9b6a03 100644 --- a/dimos/manipulation/planning/spec/protocols.py +++ b/dimos/manipulation/planning/spec/protocols.py @@ -26,6 +26,10 @@ if TYPE_CHECKING: from contextlib import AbstractContextManager + from dimos_generated.dimos_msgs.msg import GraspCandidateArray + from dimos_generated.geometry_msgs.msg import PoseStamped + from dimos_generated.sensor_msgs.msg import JointState + from dimos_generated.trajectory_msgs.msg import JointTrajectory import numpy as np from numpy.typing import NDArray @@ -42,10 +46,6 @@ VisualizationStateFrame, ) from dimos.manipulation.planning.spec.validation import PreparedRobotModel - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray - from dimos.msgs.sensor_msgs.JointState import JointState - from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory @runtime_checkable diff --git a/dimos/manipulation/planning/spec/test_joint_space.py b/dimos/manipulation/planning/spec/test_joint_space.py index 8bad5f17f83..4b3c5563f81 100644 --- a/dimos/manipulation/planning/spec/test_joint_space.py +++ b/dimos/manipulation/planning/spec/test_joint_space.py @@ -18,13 +18,13 @@ from pathlib import Path from typing import Any +from dimos_generated.sensor_msgs.msg import JointState import numpy as np import pytest from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.joint_space import CoordinateTopology, JointCoordinate from dimos.manipulation.planning.spec.validation import prepare_robot_model -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.robot.assets.model import RobotModel diff --git a/dimos/manipulation/planning/trajectory_generator/joint_trajectory_generator.py b/dimos/manipulation/planning/trajectory_generator/joint_trajectory_generator.py index e635ebbd03c..81226240591 100644 --- a/dimos/manipulation/planning/trajectory_generator/joint_trajectory_generator.py +++ b/dimos/manipulation/planning/trajectory_generator/joint_trajectory_generator.py @@ -32,8 +32,10 @@ import math -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint + +from dimos.msgs.time import duration_from_seconds, header_now, to_seconds +from dimos.msgs.trajectory import trajectory_duration class JointTrajectoryGenerator: @@ -112,13 +114,13 @@ def generate(self, waypoints: list[list[float]]) -> JointTrajectory: if not waypoints or len(waypoints) < 2: raise ValueError("Need at least 2 waypoints") - all_points: list[TrajectoryPoint] = [] + all_points: list[JointTrajectoryPoint] = [] current_time = 0.0 # Add first waypoint all_points.append( - TrajectoryPoint( - time_from_start=0.0, + JointTrajectoryPoint( + time_from_start=duration_from_seconds(0.0), positions=list(waypoints[0]), velocities=[0.0] * self.num_joints, ) @@ -136,14 +138,14 @@ def generate(self, waypoints: list[list[float]]) -> JointTrajectory: all_points.extend(segment_points[1:]) current_time += segment_duration - return JointTrajectory(points=all_points) + return JointTrajectory(header=header_now(), points=all_points) def _generate_segment( self, start: list[float], end: list[float], start_time: float, - ) -> tuple[list[TrajectoryPoint], float]: + ) -> tuple[list[JointTrajectoryPoint], float]: """ Generate trajectory points for a single segment using trapezoidal profile. @@ -153,7 +155,7 @@ def _generate_segment( start_time: Time offset for this segment Returns: - Tuple of (list of TrajectoryPoints, segment duration) + Tuple of (list of JointTrajectoryPoints, segment duration) """ # Calculate displacement for each joint displacements = [end[j] - start[j] for j in range(self.num_joints)] @@ -172,7 +174,7 @@ def _generate_segment( segment_duration = max(segment_duration, 0.01) # Generate points along the segment - points: list[TrajectoryPoint] = [] + points: list[JointTrajectoryPoint] = [] for i in range(self.points_per_segment + 1): # Normalized time [0, 1] @@ -204,8 +206,8 @@ def _generate_segment( velocities.append(vel) points.append( - TrajectoryPoint( - time_from_start=t, + JointTrajectoryPoint( + time_from_start=duration_from_seconds(t), positions=positions, velocities=velocities, ) @@ -429,7 +431,7 @@ def preview(self, trajectory: JointTrajectory) -> str: lines = [ "Trajectory Preview", "=" * 60, - f"Duration: {trajectory.duration:.3f}s", + f"Duration: {trajectory_duration(trajectory):.3f}s", f"Points: {len(trajectory.points)}", "", "Waypoints (time -> positions):", @@ -447,7 +449,7 @@ def preview(self, trajectory: JointTrajectory) -> str: pt = trajectory.points[i] pos_str = ", ".join(f"{p:+.3f}" for p in pt.positions) vel_str = ", ".join(f"{v:+.3f}" for v in pt.velocities) - lines.append(f" t={pt.time_from_start:6.3f}s: pos=[{pos_str}]") + lines.append(f" t={to_seconds(pt.time_from_start):6.3f}s: pos=[{pos_str}]") lines.append(f" vel=[{vel_str}]") lines.append("-" * 60) diff --git a/dimos/manipulation/planning/trajectory_generator/parametrizer.py b/dimos/manipulation/planning/trajectory_generator/parametrizer.py index 3497b29e94d..6cc100c94ba 100644 --- a/dimos/manipulation/planning/trajectory_generator/parametrizer.py +++ b/dimos/manipulation/planning/trajectory_generator/parametrizer.py @@ -16,14 +16,16 @@ from abc import ABC, abstractmethod from collections.abc import Sequence +import copy import math +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint + from dimos.manipulation.planning.groups.models import PlanningGroupSelection from dimos.manipulation.planning.spec.models import GeneratedPlan, PlanningResult from dimos.manipulation.planning.spec.protocols import WorldSpec -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +from dimos.msgs.time import duration_from_seconds, header_now, to_nanoseconds _TRAJECTORY_POSITION_TOLERANCE = 1e-6 @@ -49,7 +51,7 @@ def materialize_plan( f"Cannot materialize unsuccessful planning result: {result.status.name}" ) - path = [JointState(state) for state in result.path] + path = [copy.deepcopy(state) for state in result.path] self._validate_selected_path(path, selection.joint_names) if result.timestamps is None: trajectory = self._parametrize_path(world, selection, tuple(path), speed_scale) @@ -58,8 +60,8 @@ def materialize_plan( self._validate_trajectory( trajectory, selection.joint_names, - expected_start=path[0].position, - expected_goal=path[-1].position, + expected_start=list(path[0].position), + expected_goal=list(path[-1].position), ) return GeneratedPlan( @@ -127,7 +129,7 @@ def _timed_trajectory( ) -> JointTrajectory: if len(timestamps) != len(path): raise TrajectoryParametrizationError("Planner must return one timestamp per waypoint") - points: list[TrajectoryPoint] = [] + points: list[JointTrajectoryPoint] = [] for waypoint_index, (state, timestamp) in enumerate(zip(path, timestamps, strict=True)): velocities = list(state.velocity) if len(velocities) != len(selection.joint_names): @@ -135,13 +137,14 @@ def _timed_trajectory( f"Waypoint {waypoint_index} velocity dimension mismatch" ) points.append( - TrajectoryPoint( - time_from_start=float(timestamp), + JointTrajectoryPoint( + time_from_start=duration_from_seconds(float(timestamp)), positions=list(state.position), velocities=velocities, ) ) return JointTrajectory( + header=header_now(), joint_names=list(selection.joint_names), points=points, ) @@ -162,36 +165,38 @@ def _validate_trajectory( ) if not trajectory.points: raise TrajectoryParametrizationError("Generated trajectory has no points") - previous_time: float | None = None + previous_time: int | None = None for point_index, point in enumerate(trajectory.points): if len(point.positions) != len(expected) or len(point.velocities) != len(expected): raise TrajectoryParametrizationError( f"Generated point {point_index} dimension mismatch" ) cls._assert_finite_sequence( - point.positions, + list(point.positions), f"Generated point {point_index} positions", ) cls._assert_finite_sequence( - point.velocities, + list(point.velocities), f"Generated point {point_index} velocities", ) - if not math.isfinite(point.time_from_start): + try: + point_time = to_nanoseconds(point.time_from_start) + except ValueError as exc: raise TrajectoryParametrizationError( - f"Generated point {point_index} time is non-finite" - ) - if point_index == 0 and point.time_from_start != 0.0: + f"Generated point {point_index} duration is invalid" + ) from exc + if point_index == 0 and point_time != 0.0: raise TrajectoryParametrizationError("Generated trajectory must start at time 0") - if previous_time is not None and point.time_from_start <= previous_time: + if previous_time is not None and point_time <= previous_time: raise TrajectoryParametrizationError( "Generated trajectory times must be strictly increasing" ) - previous_time = point.time_from_start - if not cls._positions_close(trajectory.points[0].positions, expected_start): + previous_time = point_time + if not cls._positions_close(list(trajectory.points[0].positions), expected_start): raise TrajectoryParametrizationError( "Generated trajectory does not preserve the path start" ) - if not cls._positions_close(trajectory.points[-1].positions, expected_goal): + if not cls._positions_close(list(trajectory.points[-1].positions), expected_goal): raise TrajectoryParametrizationError( "Generated trajectory does not preserve the path goal" ) diff --git a/dimos/manipulation/planning/trajectory_generator/roboplan_toppra_parametrizer.py b/dimos/manipulation/planning/trajectory_generator/roboplan_toppra_parametrizer.py index b41e16975e3..061f9d57496 100644 --- a/dimos/manipulation/planning/trajectory_generator/roboplan_toppra_parametrizer.py +++ b/dimos/manipulation/planning/trajectory_generator/roboplan_toppra_parametrizer.py @@ -19,6 +19,8 @@ import sys from typing import Any +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint import numpy as np import roboplan.core as roboplan_core import roboplan.toppra as roboplan_toppra @@ -35,9 +37,7 @@ ) from dimos.manipulation.planning.world.roboplan_model import RoboPlanGroup, RoboPlanModel from dimos.manipulation.planning.world.roboplan_world import RoboPlanWorld -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +from dimos.msgs.time import duration_from_seconds, header_now @dataclass(frozen=True) @@ -231,8 +231,8 @@ def _canonical_result( [joint_space.normalize_positions(position) for position in canonical_positions] ) points = [ - TrajectoryPoint( - time_from_start=time, + JointTrajectoryPoint( + time_from_start=duration_from_seconds(time), positions=list(position), velocities=[float(velocity[index]) for index in output_indices], ) @@ -244,6 +244,7 @@ def _canonical_result( ) ] return JointTrajectory( + header=header_now(), joint_names=list(selection.joint_names), points=points, ) diff --git a/dimos/manipulation/planning/trajectory_generator/simple_parametrizer.py b/dimos/manipulation/planning/trajectory_generator/simple_parametrizer.py index 4f06d783c27..b1f2218ee67 100644 --- a/dimos/manipulation/planning/trajectory_generator/simple_parametrizer.py +++ b/dimos/manipulation/planning/trajectory_generator/simple_parametrizer.py @@ -16,6 +16,9 @@ import math +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.trajectory_msgs.msg import JointTrajectory + from dimos.manipulation.planning.groups.models import PlanningGroupSelection from dimos.manipulation.planning.spec.protocols import WorldSpec from dimos.manipulation.planning.trajectory_generator.config import ( @@ -28,8 +31,6 @@ BaseTrajectoryParametrizer, TrajectoryParametrizationError, ) -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory class SimpleTrapezoidParametrizer(BaseTrajectoryParametrizer): @@ -71,7 +72,7 @@ def _parametrize_path( return JointTrajectory( joint_names=list(selection.joint_names), points=generated.points, - timestamp=generated.timestamp, + header=generated.header, ) @staticmethod diff --git a/dimos/manipulation/planning/trajectory_generator/test_parametrizer.py b/dimos/manipulation/planning/trajectory_generator/test_parametrizer.py index cb303cde22a..7c53e113f86 100644 --- a/dimos/manipulation/planning/trajectory_generator/test_parametrizer.py +++ b/dimos/manipulation/planning/trajectory_generator/test_parametrizer.py @@ -16,6 +16,8 @@ from unittest.mock import MagicMock +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint import pytest from dimos.manipulation.planning.groups.models import ( @@ -29,9 +31,7 @@ BaseTrajectoryParametrizer, TrajectoryParametrizationError, ) -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +from dimos.msgs.time import duration_from_seconds, header_now, to_seconds class _FixedParametrizer(BaseTrajectoryParametrizer): @@ -73,15 +73,16 @@ def _path() -> list[JointState]: def _output() -> JointTrajectory: return JointTrajectory( + header=header_now(), joint_names=["arm/a", "arm/b"], points=[ - TrajectoryPoint( - time_from_start=0.0, + JointTrajectoryPoint( + time_from_start=duration_from_seconds(0.0), positions=[0.0, 0.0], velocities=[0.0, 0.0], ), - TrajectoryPoint( - time_from_start=0.5, + JointTrajectoryPoint( + time_from_start=duration_from_seconds(0.5), positions=[0.4, 0.0], velocities=[0.0, 0.0], ), @@ -143,7 +144,7 @@ def test_timed_planner_result_bypasses_backend_path_conversion() -> None: ) assert parametrizer.calls == [] - assert [point.time_from_start for point in plan.trajectory.points] == [ + assert [to_seconds(point.time_from_start) for point in plan.trajectory.points] == [ 0.0, 0.75, ] @@ -168,7 +169,9 @@ def test_timed_planner_result_requires_velocity_for_each_joint() -> None: def test_rejects_backend_trajectory_with_nonincreasing_time() -> None: output = _output() - output.points[-1].time_from_start = 0.0 + point = output.points[-1] + point.time_from_start = duration_from_seconds(0.0) + output.points[-1] = point parametrizer = _FixedParametrizer(output) with pytest.raises(TrajectoryParametrizationError, match="strictly increasing"): @@ -181,7 +184,9 @@ def test_rejects_backend_trajectory_with_nonincreasing_time() -> None: def test_rejects_backend_trajectory_that_changes_path_goal() -> None: output = _output() - output.points[-1].positions = [0.3, 0.0] + point = output.points[-1] + point.positions = [0.3, 0.0] + output.points[-1] = point parametrizer = _FixedParametrizer(output) with pytest.raises(TrajectoryParametrizationError, match="path goal"): diff --git a/dimos/manipulation/planning/trajectory_generator/test_roboplan_toppra_parametrizer.py b/dimos/manipulation/planning/trajectory_generator/test_roboplan_toppra_parametrizer.py index 52b77ea4325..2359a76e8c7 100644 --- a/dimos/manipulation/planning/trajectory_generator/test_roboplan_toppra_parametrizer.py +++ b/dimos/manipulation/planning/trajectory_generator/test_roboplan_toppra_parametrizer.py @@ -27,6 +27,8 @@ pytest.importorskip("roboplan.toppra") +from dimos_generated.sensor_msgs.msg import JointState + from dimos.manipulation.planning.groups.models import ( PlanningGroup, PlanningGroupDefinition, @@ -53,7 +55,6 @@ ) from dimos.manipulation.planning.world.roboplan_model import RoboPlanGroup, RoboPlanModel from dimos.manipulation.planning.world.roboplan_world import RoboPlanWorld -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.robot.assets.model import PlanarBaseDefinition, RobotModel pytestmark = pytest.mark.self_hosted @@ -319,7 +320,7 @@ def test_roboplan_toppra_parametrizes_unbounded_planar_base(tmp_path: Path) -> N world, selection, result ) - positions = [point.positions for point in plan.trajectory.points] + positions = [list(point.positions) for point in plan.trajectory.points] assert positions[0] == pytest.approx(start) assert positions[-1] == pytest.approx(goal) assert max(abs(current[2] - previous[2]) for previous, current in pairwise(positions)) < math.pi diff --git a/dimos/manipulation/planning/trajectory_generator/test_simple_parametrizer.py b/dimos/manipulation/planning/trajectory_generator/test_simple_parametrizer.py index ad82cf9510b..376bb2dca19 100644 --- a/dimos/manipulation/planning/trajectory_generator/test_simple_parametrizer.py +++ b/dimos/manipulation/planning/trajectory_generator/test_simple_parametrizer.py @@ -17,6 +17,9 @@ from pathlib import Path from unittest.mock import MagicMock +from dimos_generated.geometry_msgs.msg import Pose, PoseStamped +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header import pytest from pytest_mock import MockerFixture @@ -43,8 +46,6 @@ from dimos.manipulation.planning.trajectory_generator.simple_parametrizer import ( SimpleTrapezoidParametrizer, ) -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.robot.assets.model import LoadedRobotModel, RobotModel @@ -64,7 +65,7 @@ def _selection() -> PlanningGroupSelection: def _world(*, velocity: float = 2.0, acceleration: float = 6.0) -> WorldSpec: config = RobotModelConfig( model=RobotModel.from_file(Path("/robot.urdf")), - base_pose=PoseStamped(), + base_pose=PoseStamped(header=Header(frame_id=""), pose=Pose()), joint_names=["arm/a", "arm/b"], base_link="base", ) diff --git a/dimos/manipulation/planning/utils/path_utils.py b/dimos/manipulation/planning/utils/path_utils.py index 072e8cc7a42..290c00949aa 100644 --- a/dimos/manipulation/planning/utils/path_utils.py +++ b/dimos/manipulation/planning/utils/path_utils.py @@ -30,10 +30,9 @@ from typing import TYPE_CHECKING +from dimos_generated.sensor_msgs.msg import JointState import numpy as np -from dimos.msgs.sensor_msgs.JointState import JointState - if TYPE_CHECKING: from numpy.typing import NDArray diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index 84181a117c6..0e685aaa7af 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -25,6 +25,8 @@ from typing import TYPE_CHECKING, Any import xml.etree.ElementTree as ET +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion +from dimos_generated.std_msgs.msg import Header import numpy as np from dimos.manipulation.planning.groups.identifiers import assert_valid_group_id @@ -39,6 +41,9 @@ validate_obstacle, ) from dimos.manipulation.planning.utils.mesh_utils import prepare_urdf_for_drake +from dimos.msgs.geometry import pose_matrix +from dimos.msgs.time import to_seconds +from dimos.msgs.trajectory import trajectory_duration from dimos.robot.assets.model import LoadedRobotModel from dimos.utils.logging_config import setup_logger @@ -47,11 +52,10 @@ from numpy.typing import NDArray -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos_generated.dimos_msgs.msg import GraspCandidateArray +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.trajectory_msgs.msg import JointTrajectory if TYPE_CHECKING: from dimos.manipulation.planning.spec.models import ( @@ -537,11 +541,8 @@ def _add_obstacle_to_meshcat(self, obstacle: Obstacle, obstacle_id: str) -> None def _pose_to_rigid_transform(self, pose: PoseStamped) -> Any: """Convert PoseStamped to Drake RigidTransform.""" - pose_matrix = Transform( - translation=pose.position, - rotation=pose.orientation, - ).to_matrix() - return RigidTransform(pose_matrix) + matrix = pose_matrix(pose.pose) + return RigidTransform(matrix) def _create_shape(self, obstacle: Obstacle) -> Any: """Create Drake shape from obstacle specification.""" @@ -559,11 +560,8 @@ def _create_shape(self, obstacle: Obstacle) -> Any: raise ValueError(f"Unsupported obstacle type: {obstacle.obstacle_type}") def _validate_obstacle(self, obstacle: Obstacle, *, allow_empty_name: bool = False) -> None: - pose_matrix = Transform( - translation=obstacle.pose.position, - rotation=obstacle.pose.orientation, - ).to_matrix() - validate_obstacle(obstacle, pose_matrix, allow_empty_name=allow_empty_name) + matrix = pose_matrix(obstacle.pose.pose) + validate_obstacle(obstacle, matrix, allow_empty_name=allow_empty_name) def _remove_obstacle_geometry(self, obstacle_data: _ObstacleData) -> None: self._scene_graph.RemoveGeometry( @@ -994,9 +992,13 @@ def _get_group_ee_pose(self, ctx: Context, group_id: PlanningGroupID) -> PoseSta quat = X_WE.rotation().ToQuaternion() # Drake returns [w, x, y, z] return PoseStamped( - frame_id="world", - position=[float(pos[0]), float(pos[1]), float(pos[2])], - orientation=[float(quat.x()), float(quat.y()), float(quat.z()), float(quat.w())], + header=Header(frame_id="world"), + pose=Pose( + position=Point(x=float(pos[0]), y=float(pos[1]), z=float(pos[2])), + orientation=Quaternion( + x=float(quat.x()), y=float(quat.y()), z=float(quat.z()), w=float(quat.w()) + ), + ), ) def get_link_pose(self, ctx: Context, link_name: str) -> NDArray[np.float64]: @@ -1179,9 +1181,9 @@ def animate_trajectory( trajectory_indices = self._trajectory_indices(trajectory) playback_scale = 1.0 if duration is not None: - if duration <= 0.0 or trajectory.duration <= 0.0: + if duration <= 0.0 or trajectory_duration(trajectory) <= 0.0: raise ValueError("preview duration must be positive") - playback_scale = duration / trajectory.duration + playback_scale = duration / trajectory_duration(trajectory) with self._lock: assert self._plant_context is not None assert self._live_context is not None @@ -1199,7 +1201,7 @@ def animate_trajectory( ) try: - previous_time = trajectory.points[0].time_from_start + previous_time = to_seconds(trajectory.points[0].time_from_start) for frame_index, point in enumerate(trajectory.points): with self._lock: if self._preview_animation_generation != generation: @@ -1213,7 +1215,7 @@ def animate_trajectory( self._set_preview_positions(self._plant_context, positions) self._publish_visualization() if frame_index < len(trajectory.points) - 1: - next_time = trajectory.points[frame_index + 1].time_from_start + next_time = to_seconds(trajectory.points[frame_index + 1].time_from_start) time.sleep((next_time - previous_time) * playback_scale) previous_time = next_time finally: diff --git a/dimos/manipulation/planning/world/roboplan_model.py b/dimos/manipulation/planning/world/roboplan_model.py index 62277f1b866..4513f14c3a3 100644 --- a/dimos/manipulation/planning/world/roboplan_model.py +++ b/dimos/manipulation/planning/world/roboplan_model.py @@ -25,6 +25,7 @@ import xml.etree.ElementTree as ET from xml.sax.saxutils import escape +from dimos_generated.geometry_msgs.msg import PoseStamped import numpy as np import yaml @@ -34,8 +35,8 @@ from dimos.manipulation.planning.spec.models import PlanningGroupID from dimos.manipulation.planning.spec.validation import PreparedRobotModel from dimos.manipulation.planning.utils.mesh_utils import prepare_urdf_for_drake +from dimos.msgs.geometry import pose_matrix from dimos.robot.assets.model import LoadedRobotModel -from dimos.utils.transform_utils import pose_to_matrix ROBOPLAN_WORLD_FRAME = "dimos_world" @@ -353,8 +354,8 @@ def _apply_collision_exclusions(scene: Any, srdf: str) -> None: scene.setCollisions(first, second, False) -def _pose_attributes(pose: Any) -> dict[str, str]: - matrix = pose_to_matrix(pose) +def _pose_attributes(pose: PoseStamped) -> dict[str, str]: + matrix = pose_matrix(pose.pose) sy = float(np.hypot(matrix[0, 0], matrix[1, 0])) if sy > 1e-9: rpy = ( diff --git a/dimos/manipulation/planning/world/roboplan_world.py b/dimos/manipulation/planning/world/roboplan_world.py index 16fa2b633f7..ee8789f4e05 100644 --- a/dimos/manipulation/planning/world/roboplan_world.py +++ b/dimos/manipulation/planning/world/roboplan_world.py @@ -29,6 +29,8 @@ from threading import RLock from typing import TYPE_CHECKING, Any +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion +from dimos_generated.std_msgs.msg import Header import numpy as np try: @@ -39,6 +41,9 @@ "Install the manipulation extra before selecting the roboplan backend." ) from exc +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import JointState + from dimos.manipulation.planning.groups.models import PlanningGroup from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry from dimos.manipulation.planning.groups.utils import joint_state_to_ordered_positions @@ -59,10 +64,8 @@ RoboPlanModel, build_roboplan_model, ) -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.geometry import pose_from_matrix, pose_matrix from dimos.utils.logging_config import setup_logger -from dimos.utils.transform_utils import matrix_to_pose, pose_to_matrix if TYPE_CHECKING: from collections.abc import Generator @@ -202,7 +205,7 @@ def update_obstacle_pose(self, obstacle_id: str, pose: PoseStamped) -> bool: with self._lock: self._require_finalized() replacement_pose = deepcopy(pose) - matrix = pose_to_matrix(replacement_pose) + matrix = pose_matrix(replacement_pose.pose) if not np.isfinite(matrix).all(): raise ValueError("Obstacle pose must contain only finite values") if obstacle_id not in self._obstacles: @@ -354,16 +357,18 @@ def get_group_ee_pose(self, ctx: RoboPlanContext, group_id: PlanningGroupID) -> if group.tip_link is None: raise ValueError(f"Planning group '{group_id}' has no tip link") mat = self.get_link_pose(ctx, group.tip_link) - pose = matrix_to_pose(mat) + pose = pose_from_matrix(mat) return PoseStamped( - frame_id="world", - position=[pose.position.x, pose.position.y, pose.position.z], - orientation=[ - pose.orientation.x, - pose.orientation.y, - pose.orientation.z, - pose.orientation.w, - ], + header=Header(frame_id="world"), + pose=Pose( + position=Point(x=pose.position.x, y=pose.position.y, z=pose.position.z), + orientation=Quaternion( + x=pose.orientation.x, + y=pose.orientation.y, + z=pose.orientation.z, + w=pose.orientation.w, + ), + ), ) def get_link_pose(self, ctx: RoboPlanContext, link_name: str) -> NDArray[np.float64]: @@ -417,7 +422,7 @@ def get_group_jacobian( def _validate_model_config(self, config: RobotModelConfig) -> None: if not config.joint_names: raise ValueError("RoboPlanWorld requires explicit joint_names") - if config.base_pose.frame_id not in ("", "world"): + if config.base_pose.header.frame_id not in ("", "world"): raise ValueError("RoboPlanWorld base_pose frame_id must be empty or 'world'") def _validate_planning_group_config(self, config: RobotModelConfig) -> None: @@ -553,7 +558,7 @@ def _call_path_collision_checker( def _add_obstacle_to_scene(self, obstacle: Obstacle, obstacle_id: str) -> None: scene = self._require_scene() - matrix = pose_to_matrix(obstacle.pose) + matrix = pose_matrix(obstacle.pose.pose) color = np.asarray(obstacle.color, dtype=np.float64) if obstacle.obstacle_type == ObstacleType.BOX: self._require_dimensions(obstacle, 3) @@ -608,7 +613,7 @@ def _add_obstacle_to_scene(self, obstacle: Obstacle, obstacle_id: str) -> None: def _validate_obstacle(self, obstacle: Obstacle, *, allow_empty_name: bool = False) -> None: validate_obstacle( - obstacle, pose_to_matrix(obstacle.pose), allow_empty_name=allow_empty_name + obstacle, pose_matrix(obstacle.pose.pose), allow_empty_name=allow_empty_name ) def _require_dimensions(self, obstacle: Obstacle, n_dims: int) -> None: diff --git a/dimos/manipulation/planning/world/test_drake_world_planning_groups.py b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py index 0f377b3b415..1693fa99a38 100644 --- a/dimos/manipulation/planning/world/test_drake_world_planning_groups.py +++ b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py @@ -17,6 +17,10 @@ from dataclasses import replace from pathlib import Path +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint import numpy as np import pytest @@ -26,10 +30,7 @@ from dimos.manipulation.planning.spec.models import Obstacle from dimos.manipulation.planning.spec.validation import prepare_robot_model from dimos.manipulation.planning.world.drake_world import DRAKE_AVAILABLE, DrakeWorld -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +from dimos.msgs.time import duration_from_seconds, header_now from dimos.robot.assets.model import PlanarBaseDefinition, RobotModel requires_drake = pytest.mark.skipif( @@ -40,10 +41,19 @@ def _trajectory(names: list[str], first: list[float], second: list[float]) -> JointTrajectory: return JointTrajectory( + header=header_now(), joint_names=names, points=[ - TrajectoryPoint(time_from_start=0.0, positions=first, velocities=[0.0] * len(names)), - TrajectoryPoint(time_from_start=2.0, positions=second, velocities=[0.0] * len(names)), + JointTrajectoryPoint( + time_from_start=duration_from_seconds(0.0), + positions=first, + velocities=[0.0] * len(names), + ), + JointTrajectoryPoint( + time_from_start=duration_from_seconds(2.0), + positions=second, + velocities=[0.0] * len(names), + ), ], ) @@ -123,7 +133,10 @@ def _config( ) -> RobotModelConfig: return RobotModelConfig( model=RobotModel.from_file(path).with_default_joint_acceleration_limit(2.0), - base_pose=PoseStamped(position=[0, 0, 0], orientation=[0, 0, 0, 1]), + base_pose=PoseStamped( + header=Header(frame_id=""), + pose=Pose(position=Point(x=0, y=0, z=0), orientation=Quaternion(x=0, y=0, z=0, w=1)), + ), joint_names=joints or ["joint1", "joint2"], base_link="base_link", planning_groups=groups, @@ -214,7 +227,10 @@ def test_drake_obstacle_ids_are_world_owned_and_invalid_insertions_are_rejected( obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, - pose=PoseStamped(position=[2, 0, 0], orientation=[0, 0, 0, 1]), + pose=PoseStamped( + header=Header(frame_id=""), + pose=Pose(position=Point(x=2, y=0, z=0), orientation=Quaternion(x=0, y=0, z=0, w=1)), + ), dimensions=(0.1, 0.1, 0.1), ) unnamed = replace(obstacle, name="") @@ -240,10 +256,15 @@ def test_drake_obstacle_ids_are_world_owned_and_invalid_insertions_are_rejected( assert world.remove_obstacle("missing") is False assert world.update_obstacle(replace(obstacle, name="missing")) is False assert world.update_obstacle_pose("missing", obstacle.pose) is False - moved_pose = PoseStamped(position=[0.0, 0.0, 0.0], orientation=[0.0, 0.0, 0.0, 1.0]) + moved_pose = PoseStamped( + header=Header(frame_id=""), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + ) assert world.update_obstacle_pose("box", moved_pose) assert world._obstacles["box"].geometry_id != original_geometry_id - assert world.get_obstacles()[0].pose.position.x == pytest.approx(0.0) + assert world.get_obstacles()[0].pose.pose.position.x == pytest.approx(0.0) assert not world.check_config_collision_free(joint_state) replacement = replace( @@ -272,7 +293,10 @@ def test_drake_obstacle_replacement_failure_invalidates_world( obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, - pose=PoseStamped(position=[0, 0, 0], orientation=[0, 0, 0, 1]), + pose=PoseStamped( + header=Header(frame_id=""), + pose=Pose(position=Point(x=0, y=0, z=0), orientation=Quaternion(x=0, y=0, z=0, w=1)), + ), dimensions=(0.1, 0.1, 0.1), ) world.add_obstacle(obstacle) @@ -296,13 +320,13 @@ def test_drake_group_fk_uses_tip_link_and_unique_pose_group(tmp_path: Path) -> N _load(world, _config(urdf, [_arm_group("joint1", "joint2")])) world.finalize() ctx = world.get_live_context() - world.set_joint_state(ctx, JointState({"name": ["joint1", "joint2"], "position": [0.0, 0.0]})) + world.set_joint_state(ctx, JointState(name=["joint1", "joint2"], position=[0.0, 0.0])) group_pose = world.get_group_ee_pose(ctx, "arm") default_pose = world.get_ee_pose(ctx) - assert group_pose.position.x == pytest.approx(2.0) - assert default_pose.position.x == pytest.approx(group_pose.position.x) + assert group_pose.pose.position.x == pytest.approx(2.0) + assert default_pose.pose.position.x == pytest.approx(group_pose.pose.position.x) assert world.get_jacobian(ctx).shape == (6, 2) @@ -317,7 +341,12 @@ def test_drake_applies_config_base_pose_when_urdf_has_world_base_joint( world, RobotModelConfig( model=RobotModel.from_file(urdf).with_default_joint_acceleration_limit(2.0), - base_pose=PoseStamped(position=[0, 0.5, 0], orientation=[0, 0, 0, 1]), + base_pose=PoseStamped( + header=Header(frame_id=""), + pose=Pose( + position=Point(x=0, y=0.5, z=0), orientation=Quaternion(x=0, y=0, z=0, w=1) + ), + ), joint_names=["joint1", "joint2"], base_link="base_link", planning_groups=[_arm_group("joint1", "joint2")], @@ -349,7 +378,12 @@ def test_drake_planar_base_coordinates_move_original_robot_root(tmp_path: Path) .with_default_joint_acceleration_limit(2.0) .with_planar_base(planar_base) ), - base_pose=PoseStamped(position=[0, 0, 0.5], orientation=[0, 0, 0, 1]), + base_pose=PoseStamped( + header=Header(frame_id=""), + pose=Pose( + position=Point(x=0, y=0, z=0.5), orientation=Quaternion(x=0, y=0, z=0, w=1) + ), + ), joint_names=joint_names, base_link=planar_base.root_link, planning_groups=[ @@ -373,9 +407,9 @@ def test_drake_planar_base_coordinates_move_original_robot_root(tmp_path: Path) tool_pose = world.get_group_ee_pose(context, "mobile_arm") np.testing.assert_allclose(original_root_pose[:3, 3], [1.0, 2.0, 0.5], atol=1e-8) - assert tool_pose.position.x == pytest.approx(1.0) - assert tool_pose.position.y == pytest.approx(4.0) - assert tool_pose.position.z == pytest.approx(0.5) + assert tool_pose.pose.position.x == pytest.approx(1.0) + assert tool_pose.pose.position.y == pytest.approx(4.0) + assert tool_pose.pose.position.z == pytest.approx(0.5) @requires_drake @@ -395,7 +429,7 @@ def test_drake_group_jacobian_shape_and_group_local_order(tmp_path: Path) -> Non ) world.finalize() ctx = world.get_live_context() - world.set_joint_state(ctx, JointState({"name": ["joint1", "joint2"], "position": [0.0, 0.0]})) + world.set_joint_state(ctx, JointState(name=["joint1", "joint2"], position=[0.0, 0.0])) forward_jacobian = world.get_group_jacobian(ctx, "wrist_forward") reverse_jacobian = world.get_group_jacobian(ctx, "wrist_reverse") diff --git a/dimos/manipulation/sdk.py b/dimos/manipulation/sdk.py index bfc8814155e..ae31e9b4b92 100644 --- a/dimos/manipulation/sdk.py +++ b/dimos/manipulation/sdk.py @@ -19,6 +19,9 @@ import math from typing import TYPE_CHECKING +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header import numpy as np from numpy.typing import ArrayLike, NDArray @@ -34,9 +37,6 @@ ) from dimos.manipulation.planning.groups.utils import joint_state_to_ordered_positions from dimos.manipulation.planning.spec.models import PlanningGroupID -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.sensor_msgs.JointState import JointState if TYPE_CHECKING: from dimos.porcelain.dimos import Dimos @@ -165,14 +165,17 @@ def move_pose( This targets an endpoint; use ``move_linear`` for a straight translation. """ xyz = _vector(position, 3, "position") - rotation = ( - self.pose().orientation - if orientation is None - else Quaternion(_vector(orientation, 4, "orientation")) - ) - if math.hypot(*rotation.to_tuple()) == 0.0: + if orientation is None: + rotation = self.pose().pose.orientation + else: + x, y, z, w = _vector(orientation, 4, "orientation") + rotation = Quaternion(x=x, y=y, z=z, w=w) + if math.hypot(rotation.x, rotation.y, rotation.z, rotation.w) == 0.0: raise ValueError("orientation must have nonzero norm") - target = PoseStamped(frame_id="world", position=xyz, orientation=rotation) + target = PoseStamped( + header=Header(frame_id="world"), + pose=Pose(position=Point(x=xyz[0], y=xyz[1], z=xyz[2]), orientation=rotation), + ) plan = self.rpc.plan_to_poses({self.info.id: target}, speed_scale=speed_scale) return self._execute("move_pose", plan, timeout) diff --git a/dimos/manipulation/test_execution_manager.py b/dimos/manipulation/test_execution_manager.py index 013cb178110..42570bb7652 100644 --- a/dimos/manipulation/test_execution_manager.py +++ b/dimos/manipulation/test_execution_manager.py @@ -16,6 +16,11 @@ from unittest.mock import MagicMock +from dimos_generated.builtin_interfaces.msg import Duration, Time +from dimos_generated.dimos_msgs.msg import TrajectoryStatus +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint import pytest from dimos.control.coordinator import ControlCoordinator @@ -30,10 +35,8 @@ from dimos.manipulation.manipulation_spec import ExecutionStatus from dimos.manipulation.planning.spec.enums import PlanningStatus from dimos.manipulation.planning.spec.models import GeneratedPlan -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint -from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState, TrajectoryStatus +from dimos.msgs.time import duration_from_seconds, header_now, to_seconds +from dimos.msgs.trajectory import TrajectoryState def _plan( @@ -41,12 +44,16 @@ def _plan( status: PlanningStatus = PlanningStatus.SUCCESS, ) -> GeneratedPlan: points = [ - TrajectoryPoint(positions=[0.0] * len(names), time_from_start=0.0), - TrajectoryPoint(positions=[1.0] * len(names), time_from_start=1.0), + JointTrajectoryPoint( + positions=[0.0] * len(names), time_from_start=duration_from_seconds(0.0) + ), + JointTrajectoryPoint( + positions=[1.0] * len(names), time_from_start=duration_from_seconds(1.0) + ), ] return GeneratedPlan( group_ids=("both_arms",), - trajectory=JointTrajectory(joint_names=list(names), points=points), + trajectory=JointTrajectory(header=header_now(), joint_names=list(names), points=points), path=[JointState(name=list(names), position=point.positions) for point in points], status=status, ) @@ -62,7 +69,7 @@ def invoke(task: str, method: str, args: dict | None = None): return TrajectoryExecutionResult(TrajectoryExecutionStatus.ACCEPTED) if method == "cancel": return TrajectoryCancellationResult(TrajectoryCancellationStatus.ALREADY_STOPPED) - return TrajectoryStatus(state=TrajectoryState.IDLE) + return TrajectoryStatus(header=header_now(), state=TrajectoryState.IDLE) coordinator.task_invoke.side_effect = invoke return coordinator @@ -169,7 +176,9 @@ def _task_invoke(self, task, method, args=None): return self.cancellations.get( task, TrajectoryCancellationResult(TrajectoryCancellationStatus.CANCELLED) ) - return TrajectoryStatus(state=self.states[task], error=self.errors.get(task, "")) + return TrajectoryStatus( + header=header_now(), state=self.states[task], error=self.errors.get(task, "") + ) def manager(self) -> PlanExecutionManager: return PlanExecutionManager( @@ -190,7 +199,10 @@ def test_whole_body_plan_splits_into_joint_and_base_columns() -> None: assert robot.dispatched[JOINT_TRAJECTORY_TASK_NAME].joint_names == ["left/j1"] assert robot.dispatched["base_traj"].joint_names == list(BASE) - assert [p.time_from_start for p in robot.dispatched["base_traj"].points] == [0.0, 1.0] + assert [to_seconds(p.time_from_start) for p in robot.dispatched["base_traj"].points] == [ + 0.0, + 1.0, + ] robot.states = dict.fromkeys(robot.states, TrajectoryState.COMPLETED) assert manager.wait().status is ExecutionStatus.COMPLETED @@ -220,7 +232,7 @@ def test_a_refused_part_cancels_the_parts_already_dispatched() -> None: assert result.status is ExecutionStatus.REJECTED assert "too fast" in result.message - assert robot.states[JOINT_TRAJECTORY_TASK_NAME] is TrajectoryState.ABORTED + assert robot.states[JOINT_TRAJECTORY_TASK_NAME] == TrajectoryState.ABORTED def test_a_refused_part_is_uncertain_when_the_others_cannot_be_confirmed_stopped() -> None: @@ -255,3 +267,27 @@ def test_a_poll_from_a_finished_run_does_not_disturb_the_next_one() -> None: assert manager.status is ExecutionStatus.ACCEPTED assert set(robot.states.values()) == {TrajectoryState.EXECUTING} + + +def test_split_cdr_plan_preserves_header_durations_and_optional_point_fields(): + robot = _WholeBody() + manager = robot.manager() + plan = _plan(("base/yaw", "left/j1", "base/x", "base/y")) + plan.trajectory.header = Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id="robot") + points = list(plan.trajectory.points) + for point in points: + point.velocities = [0.1, 0.2, 0.3, 0.4] + point.accelerations = [1.0, 2.0, 3.0, 4.0] + point.effort = [5.0, 6.0, 7.0, 8.0] + points[-1].time_from_start = Duration(sec=1, nanosec=9) + plan.trajectory.points = points + try: + assert manager.execute(plan, blocking=False).status == ExecutionStatus.ACCEPTED + part = JointTrajectory.decode(robot.dispatched["base_traj"].encode()) + assert part.header == plan.trajectory.header + assert part.points[-1].time_from_start == Duration(sec=1, nanosec=9) + assert list(part.points[-1].velocities) == [0.3, 0.4, 0.1] + assert list(part.points[-1].accelerations) == [3.0, 4.0, 1.0] + assert list(part.points[-1].effort) == [7.0, 8.0, 5.0] + finally: + manager.cancel() diff --git a/dimos/manipulation/test_generated_plan_materialization.py b/dimos/manipulation/test_generated_plan_materialization.py index 4d405769d24..3b5becd08cb 100644 --- a/dimos/manipulation/test_generated_plan_materialization.py +++ b/dimos/manipulation/test_generated_plan_materialization.py @@ -17,6 +17,18 @@ from pathlib import Path from unittest.mock import MagicMock +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint import pytest from dimos.manipulation.manipulation_module import ManipulationState @@ -38,13 +50,7 @@ from dimos.manipulation.planning.trajectory_generator.simple_parametrizer import ( SimpleTrapezoidParametrizer, ) -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +from dimos.msgs.time import duration_from_seconds, header_now, to_seconds from dimos.robot.assets.model import LoadedRobotModel, RobotModel @@ -66,21 +72,24 @@ def __init__( def generate(self, waypoints: list[list[float]]) -> JointTrajectory: RecordingGenerator.calls.append(waypoints) return JointTrajectory( + header=header_now(), points=[ - TrajectoryPoint( - time_from_start=float(index), + JointTrajectoryPoint( + time_from_start=duration_from_seconds(float(index)), positions=list(point), velocities=[0.0] * self.num_joints, ) for index, point in enumerate(waypoints) - ] + ], ) def _model() -> RobotModelConfig: return RobotModelConfig( model=RobotModel.from_file(Path("/robot.urdf")), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + base_pose=PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + ), joint_names=["left/a", "left/b", "right/c"], base_link="base", planning_groups=[ @@ -162,7 +171,7 @@ def test_materializes_once_with_reordered_groups_heterogeneous_limits_and_distin assert module._last_plan is not None assert module._last_plan.path is not module._last_plan.trajectory.points assert module._last_plan.trajectory.joint_names == names - assert module._last_plan.trajectory.points[-1].time_from_start == 1.0 + assert to_seconds(module._last_plan.trajectory.points[-1].time_from_start) == 1.0 def test_cartesian_plan_preserves_planner_timestamps_and_velocities(monkeypatch, module_factory): @@ -184,8 +193,12 @@ def test_cartesian_plan_preserves_planner_timestamps_and_velocities(monkeypatch, plan = module.generate_cartesian_plan( { "left_arm": ( - Transform.identity(), - Transform(translation=Vector3(0.01, 0.0, 0.0)), + TransformStamped(header=Header(frame_id="world"), child_frame_id=""), + TransformStamped( + header=Header(frame_id="world"), + transform=Transform(translation=Vector3(x=0.01, y=0.0, z=0.0)), + child_frame_id="", + ), ) }, RoboPlanCartesianPathConfig( @@ -196,7 +209,7 @@ def test_cartesian_plan_preserves_planner_timestamps_and_velocities(monkeypatch, ) assert plan is not None - assert [point.time_from_start for point in plan.trajectory.points] == [0.0, 0.25] + assert [to_seconds(point.time_from_start) for point in plan.trajectory.points] == [0.0, 0.25] assert [point.velocities for point in plan.trajectory.points] == [ [0.0, 0.0], [0.4, 0.2], diff --git a/dimos/manipulation/test_manipulation_module.py b/dimos/manipulation/test_manipulation_module.py index 448d74ec487..15c15fa915b 100644 --- a/dimos/manipulation/test_manipulation_module.py +++ b/dimos/manipulation/test_manipulation_module.py @@ -24,6 +24,9 @@ import importlib.util from unittest.mock import DEFAULT, MagicMock +from dimos_generated.dimos_msgs.msg import TrajectoryStatus +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion +from dimos_generated.sensor_msgs.msg import JointState import pytest from dimos.control.coordinator import ControlCoordinator @@ -40,11 +43,8 @@ from dimos.manipulation.manipulation_spec import ExecutionStatus from dimos.manipulation.planning.planners.config import RRTConnectPlannerConfig from dimos.manipulation.planning.spec.config import RobotModelConfig -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState, TrajectoryStatus +from dimos.msgs.time import header_now +from dimos.msgs.trajectory import TrajectoryState, trajectory_duration from dimos.robot.manipulators.xarm.config import make_xarm7_model_config pytestmark = pytest.mark.self_hosted @@ -96,7 +96,9 @@ def invoke(task: str, method: str, args: dict | None = None): return TrajectoryCancellationResult(TrajectoryCancellationStatus.ALREADY_STOPPED) return DEFAULT - coordinator.task_invoke.return_value = TrajectoryStatus(state=TrajectoryState.COMPLETED) + coordinator.task_invoke.return_value = TrajectoryStatus( + header=header_now(), state=TrajectoryState.COMPLETED + ) coordinator.task_invoke.side_effect = invoke mod = ManipulationModule( model=xarm7_config, @@ -157,7 +159,7 @@ def test_plan_to_joints(self, module, joint_state_zeros): assert module._last_plan is not None assert len(module._last_plan.trajectory.points) > 1 - assert module._last_plan.trajectory.duration > 0 + assert trajectory_duration(module._last_plan.trajectory) > 0 assert module._last_plan.group_ids == ("manipulator",) def test_plan_to_explicit_joint_target(self, module, joint_state_zeros): @@ -177,10 +179,7 @@ def test_add_and_remove_obstacle(self, module, joint_state_zeros): """Test adding and removing obstacles.""" module._on_joint_state(joint_state_zeros) - pose = Pose( - position=Vector3(0.5, 0.0, 0.3), - orientation=Quaternion(), # default is identity (w=1) - ) + pose = Pose(position=Point(x=0.5, y=0.0, z=0.3), orientation=Quaternion()) obstacle_id = module.add_obstacle("test_box", pose, "box", [0.1, 0.1, 0.1]) assert obstacle_id != "" diff --git a/dimos/manipulation/test_manipulation_monitor_preview.py b/dimos/manipulation/test_manipulation_monitor_preview.py index 6328ad55282..cbe008cf333 100644 --- a/dimos/manipulation/test_manipulation_monitor_preview.py +++ b/dimos/manipulation/test_manipulation_monitor_preview.py @@ -19,6 +19,11 @@ from pathlib import Path from unittest.mock import MagicMock +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint import pytest from dimos.manipulation.manipulation_module import ManipulationModule @@ -30,12 +35,7 @@ from dimos.manipulation.planning.spec.enums import PlanningStatus from dimos.manipulation.planning.spec.models import GeneratedPlan from dimos.manipulation.planning.spec.protocols import VisualizationSpec -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +from dimos.msgs.time import duration_from_seconds, header_now from dimos.robot.assets.model import RobotModel @@ -44,7 +44,9 @@ def canonical_model_config() -> RobotModelConfig: """Create a model whose joint names match coordinator-facing names.""" return RobotModelConfig( model=RobotModel.from_file(Path("/path/to/robot.urdf")), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + base_pose=PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + ), joint_names=["left/joint1", "left/joint2", "left/joint3"], base_link="link_base", planning_groups=[ @@ -61,7 +63,9 @@ def canonical_model_config() -> RobotModelConfig: def _one_joint_config(name: str = "arm") -> RobotModelConfig: return RobotModelConfig( model=RobotModel.from_file(Path("/path")), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + base_pose=PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + ), joint_names=["j0"], base_link="base_link", planning_groups=[ @@ -86,10 +90,11 @@ def _install_generated_plan( ) module._last_plan = GeneratedPlan( trajectory=JointTrajectory( + header=header_now(), joint_names=config.joint_names, points=[ - TrajectoryPoint( - time_from_start=float(index), + JointTrajectoryPoint( + time_from_start=duration_from_seconds(float(index)), positions=list(point), velocities=[0.0 for _ in config.joint_names], ) @@ -127,9 +132,12 @@ def _make_path(*points: list[float]) -> list[JointState]: def _make_trajectory(*points: tuple[float, list[float]]) -> JointTrajectory: joint_names = [f"j{i}" for i in range(len(points[0][1]))] if points else [] return JointTrajectory( + header=header_now(), joint_names=joint_names, points=[ - TrajectoryPoint(time_from_start=time_from_start, positions=positions) + JointTrajectoryPoint( + time_from_start=duration_from_seconds(time_from_start), positions=positions + ) for time_from_start, positions in points ], ) @@ -187,6 +195,7 @@ def test_routes_canonical_state_in_model_order(self, canonical_model_config, mod module.config.model = canonical_model_config msg = JointState( + header=Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id="robot"), name=["left/joint3", "unrelated", "left/joint1", "left/joint2"], position=[0.3, 9.0, 0.1, 0.2], velocity=[3.0, 9.0, 1.0, 2.0], @@ -197,6 +206,9 @@ def test_routes_canonical_state_in_model_order(self, canonical_model_config, mod assert state.name == canonical_model_config.joint_names assert state.position == [0.1, 0.2, 0.3] assert state.velocity == [1.0, 2.0, 3.0] + assert state.header == msg.header + msg.header.frame_id = "changed" + assert state.header.frame_id == "robot" def test_aliases_coordinator_names_onto_model_joints( self, canonical_model_config, module_factory @@ -280,7 +292,9 @@ def test_visualization_routing_and_stop_all_monitors(self): monitor.cancel_preview_animation() path = _make_path([1.0], [2.0], [3.0]) plan = GeneratedPlan( - trajectory=JointTrajectory(), + trajectory=JointTrajectory( + header=header_now(), + ), group_ids=("manipulator",), status=PlanningStatus.SUCCESS, path=path, @@ -303,7 +317,12 @@ def test_visualization_none_is_noop(self): assert monitor.get_visualization_url() is None monitor.update_visualization_state() monitor.cancel_preview_animation() - monitor.animate_trajectory(JointTrajectory(), 1.0) + monitor.animate_trajectory( + JointTrajectory( + header=header_now(), + ), + 1.0, + ) monitor.start_visualization_thread() assert monitor._viz_thread is None @@ -320,7 +339,13 @@ def test_preview_without_a_plan_returns_rejected(self, module_factory): def test_clear_planned_path_invalidates_before_dismissing_preview(self, module_factory): module = module_factory() - plan = GeneratedPlan(trajectory=JointTrajectory(), group_ids=("manipulator",), path=[]) + plan = GeneratedPlan( + trajectory=JointTrajectory( + header=header_now(), + ), + group_ids=("manipulator",), + path=[], + ) module._last_plan = plan module._world_monitor = MagicMock() plan_during_dismissal: list[GeneratedPlan | None] = [] @@ -337,7 +362,11 @@ def test_clear_planned_path_invalidates_before_dismissing_preview(self, module_f def test_clear_planned_path_clears_without_a_world_monitor(self, module_factory): module = module_factory() module._last_plan = GeneratedPlan( - trajectory=JointTrajectory(), group_ids=("manipulator",), path=[] + trajectory=JointTrajectory( + header=header_now(), + ), + group_ids=("manipulator",), + path=[], ) assert module.clear_planned_path().succeeded diff --git a/dimos/manipulation/test_manipulation_primitives.py b/dimos/manipulation/test_manipulation_primitives.py index e276b3d3446..7c4d15dc7be 100644 --- a/dimos/manipulation/test_manipulation_primitives.py +++ b/dimos/manipulation/test_manipulation_primitives.py @@ -18,6 +18,10 @@ from pathlib import Path from unittest.mock import MagicMock +from dimos_generated.geometry_msgs.msg import TransformStamped +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint import pytest from pytest_mock import MockerFixture @@ -36,10 +40,7 @@ from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import PlanningStatus from dimos.manipulation.planning.spec.models import GeneratedPlan -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +from dimos.msgs.time import duration_from_seconds, header_now from dimos.robot.assets.model import RobotModel from dimos.robot.manipulators.openyam.config import make_openyam_model_config @@ -69,10 +70,11 @@ def _plan() -> GeneratedPlan: group_ids=("tool",), status=PlanningStatus.SUCCESS, trajectory=JointTrajectory( + header=header_now(), joint_names=names, points=[ - TrajectoryPoint(positions=[0.0], time_from_start=0.0), - TrajectoryPoint(positions=[0.1], time_from_start=1.0), + JointTrajectoryPoint(positions=[0.0], time_from_start=duration_from_seconds(0.0)), + JointTrajectoryPoint(positions=[0.1], time_from_start=duration_from_seconds(1.0)), ], ), ) @@ -103,10 +105,10 @@ def test_move_linear_uses_world_relative_target_and_default_speed( assert result.succeeded targets, config = generate.call_args.args start, relative = targets["tool"] - assert start == Transform.identity() - assert relative.translation.x == pytest.approx(0.02) - assert relative.translation.y == pytest.approx(0.0) - assert relative.translation.z == pytest.approx(-0.01) + assert start == TransformStamped(header=Header(frame_id="world"), child_frame_id="") + assert relative.transform.translation.x == pytest.approx(0.02) + assert relative.transform.translation.y == pytest.approx(0.0) + assert relative.transform.translation.z == pytest.approx(-0.01) assert generate.call_args.kwargs["check_collision"] is False assert generate.call_args.kwargs["speed_scale"] == pytest.approx(0.5) execute.assert_called_once_with(blocking=False, timeout=None, plan_id=generated.plan_id) diff --git a/dimos/manipulation/test_manipulation_spec.py b/dimos/manipulation/test_manipulation_spec.py index 915cb70a842..a9c383835d9 100644 --- a/dimos/manipulation/test_manipulation_spec.py +++ b/dimos/manipulation/test_manipulation_spec.py @@ -20,6 +20,9 @@ from types import UnionType from typing import get_args, get_origin, get_type_hints +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint + from dimos.manipulation.manipulation_module import ManipulationModule import dimos.manipulation.manipulation_spec as manipulation_spec from dimos.manipulation.manipulation_spec import ( @@ -37,9 +40,7 @@ ) from dimos.manipulation.planning.spec.enums import PlanningStatus from dimos.manipulation.planning.spec.models import GeneratedPlan -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +from dimos.msgs.time import duration_from_seconds, header_now from dimos.spec.utils import spec_annotation_compliance @@ -124,10 +125,17 @@ def test_plan_result_repr_summarizes_trajectory_without_dumping_points() -> None plan = GeneratedPlan( group_ids=("arm/manipulator",), trajectory=JointTrajectory( + header=header_now(), joint_names=["arm/j0"], points=[ - TrajectoryPoint(positions=[0.0], velocities=[0.0], time_from_start=0.0), - TrajectoryPoint(positions=[123.456], velocities=[0.0], time_from_start=1.0), + JointTrajectoryPoint( + positions=[0.0], velocities=[0.0], time_from_start=duration_from_seconds(0.0) + ), + JointTrajectoryPoint( + positions=[123.456], + velocities=[0.0], + time_from_start=duration_from_seconds(1.0), + ), ], ), path=[ @@ -145,7 +153,7 @@ def test_plan_result_repr_summarizes_trajectory_without_dumping_points() -> None assert "waypoints=2" in text assert "path_length=123" in text assert "iterations=7" in text - assert "TrajectoryPoint" not in text + assert "JointTrajectoryPoint" not in text assert "positions" not in text diff --git a/dimos/manipulation/test_plan_execution.py b/dimos/manipulation/test_plan_execution.py index cbd20a1116d..0676e7e494d 100644 --- a/dimos/manipulation/test_plan_execution.py +++ b/dimos/manipulation/test_plan_execution.py @@ -16,6 +16,9 @@ from unittest.mock import DEFAULT, MagicMock +from dimos_generated.dimos_msgs.msg import TrajectoryStatus +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint import pytest from dimos.control.coordinator import ControlCoordinator @@ -30,27 +33,26 @@ from dimos.manipulation.planning.spec.enums import PlanningStatus from dimos.manipulation.planning.spec.models import GeneratedPlan from dimos.manipulation.visualization.operator import ManipulationOperator -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint -from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState, TrajectoryStatus +from dimos.msgs.time import duration_from_seconds, header_now +from dimos.msgs.trajectory import TrajectoryState from dimos.robot.assets.model import PlanarBaseDefinition def _plan(final_position: float = 1.0, joint_name: str = "arm/j0") -> GeneratedPlan: names = [joint_name] trajectory = JointTrajectory( + header=header_now(), joint_names=names, points=[ - TrajectoryPoint( + JointTrajectoryPoint( positions=[0.0], velocities=[0.0], - time_from_start=0.0, + time_from_start=duration_from_seconds(0.0), ), - TrajectoryPoint( + JointTrajectoryPoint( positions=[final_position], velocities=[0.0], - time_from_start=1.0, + time_from_start=duration_from_seconds(1.0), ), ], ) @@ -149,7 +151,7 @@ def test_execute_sends_planar_base_columns_to_the_base_task(module_factory) -> N coordinator.task_invoke.side_effect = lambda task, method, args: ( TrajectoryExecutionResult(TrajectoryExecutionStatus.ACCEPTED) if method == "execute" - else TrajectoryStatus(state=state) + else TrajectoryStatus(header=header_now(), state=state) ) module = _module_with_coordinator(coordinator, module_factory) planar_base = PlanarBaseDefinition( @@ -165,12 +167,12 @@ def test_execute_sends_planar_base_columns_to_the_base_task(module_factory) -> N module._initialize_execution() names = ["arm/j0", *planar_base.joint_names] points = [ - TrajectoryPoint(positions=[0.0] * 4, time_from_start=0.0), - TrajectoryPoint(positions=[1.0] * 4, time_from_start=1.0), + JointTrajectoryPoint(positions=[0.0] * 4, time_from_start=duration_from_seconds(0.0)), + JointTrajectoryPoint(positions=[1.0] * 4, time_from_start=duration_from_seconds(1.0)), ] module._last_plan = GeneratedPlan( group_ids=("manipulator",), - trajectory=JointTrajectory(joint_names=names, points=points), + trajectory=JointTrajectory(header=header_now(), joint_names=names, points=points), path=[JointState(name=names, position=point.positions) for point in points], status=PlanningStatus.SUCCESS, ) @@ -249,10 +251,12 @@ def read_status(): return module.get_state().operation_status.name assert operator.execute(_plan()) is True - coordinator.task_invoke.return_value = TrajectoryStatus(state=TrajectoryState.EXECUTING) + coordinator.task_invoke.return_value = TrajectoryStatus( + header=header_now(), state=TrajectoryState.EXECUTING + ) assert read_status() == "EXECUTING" - coordinator.task_invoke.return_value = TrajectoryStatus(state=terminal) + coordinator.task_invoke.return_value = TrajectoryStatus(header=header_now(), state=terminal) assert read_status() == operation assert module.get_state().execution_status.name == terminal.name coordinator.task_invoke.reset_mock() diff --git a/dimos/manipulation/test_planning_factory.py b/dimos/manipulation/test_planning_factory.py index d8985706595..0b2f3376079 100644 --- a/dimos/manipulation/test_planning_factory.py +++ b/dimos/manipulation/test_planning_factory.py @@ -23,6 +23,8 @@ from typing import Any from unittest.mock import ANY +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.std_msgs.msg import Header from pydantic import ValidationError import pytest from pytest_mock import MockerFixture @@ -50,9 +52,6 @@ SimpleTrapezoidParametrizationConfig, TrajectoryParametrizationConfig, ) -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.robot.assets.model import RobotModel @@ -75,7 +74,9 @@ def _make(**kwargs: Any) -> ManipulationModule: def robot_config() -> RobotModelConfig: return RobotModelConfig( model=RobotModel.from_file(Path("/path/to/robot.urdf")), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), # type: ignore[call-arg] + base_pose=PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + ), # type: ignore[call-arg] joint_names=["joint1", "joint2"], planning_groups=[ PlanningGroupDefinition( diff --git a/dimos/manipulation/test_roboplan.py b/dimos/manipulation/test_roboplan.py index b19dd48081c..61d2e8d0642 100644 --- a/dimos/manipulation/test_roboplan.py +++ b/dimos/manipulation/test_roboplan.py @@ -26,6 +26,17 @@ from typing import Any, ClassVar import xml.etree.ElementTree as ET +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from pytest_mock import MockerFixture @@ -45,13 +56,8 @@ from dimos.manipulation.planning.spec.enums import ObstacleType, PlanningStatus from dimos.manipulation.planning.spec.models import Obstacle from dimos.manipulation.planning.spec.validation import MAX_OCTREE_POINTS, prepare_robot_model -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.geometry import pose_matrix, quaternion_from_euler from dimos.robot.assets.model import RobotModel -from dimos.utils.transform_utils import pose_to_matrix class FakeJointConfiguration: @@ -514,7 +520,9 @@ def robot_config(tmp_path: Path) -> RobotModelConfig: ) return RobotModelConfig( model=RobotModel.from_file(model_path).with_default_joint_acceleration_limit(2.0), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), # type: ignore[call-arg] + base_pose=PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + ), # type: ignore[call-arg] joint_names=["joint1", "joint2"], base_link="base", planning_groups=[ @@ -601,16 +609,14 @@ def _selection( ) -def _relative_target(*waypoints: Transform) -> tuple[Transform, ...]: - return (Transform.identity(), *waypoints) +def _relative_target(*waypoints: TransformStamped) -> tuple[TransformStamped, ...]: + return (TransformStamped(header=Header(frame_id="world"), child_frame_id=""), *waypoints) def _absolute_target(*waypoints: PoseStamped) -> tuple[PoseStamped, ...]: return ( PoseStamped( - frame_id="world", - position=Vector3(), - orientation=Quaternion(), + header=Header(frame_id="world"), pose=Pose(position=Point(), orientation=Quaternion()) ), *waypoints, ) @@ -724,22 +730,27 @@ def test_obstacle_mutation_updates_scene_and_stored_pose( obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, - pose=PoseStamped(position=Vector3(), orientation=Quaternion()), # type: ignore[call-arg] + pose=PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + ), # type: ignore[call-arg] dimensions=(0.1, 0.2, 0.3), ) assert world.add_obstacle(obstacle) == "box" assert add_box.call_args.args[2] == "dimos_world" assert "box" in world._scene.geometry - updated_pose = PoseStamped(position=Vector3(1, 0, 0), orientation=Quaternion()) # type: ignore[call-arg] + updated_pose = PoseStamped( + header=Header(frame_id=""), + pose=Pose(position=Point(x=1, y=0, z=0), orientation=Quaternion()), + ) # type: ignore[call-arg] assert world.update_obstacle_pose( "box", updated_pose, ) np.testing.assert_allclose( - pose_to_matrix(world.get_obstacles()[0].pose), - pose_to_matrix(updated_pose), + pose_matrix(world.get_obstacles()[0].pose.pose), + pose_matrix(updated_pose.pose), ) - np.testing.assert_allclose(world._scene.geometry["box"], pose_to_matrix(updated_pose)) + np.testing.assert_allclose(world._scene.geometry["box"], pose_matrix(updated_pose.pose)) assert world.add_obstacle(obstacle) is None assert world.remove_obstacle("box") assert world.get_obstacles() == [] @@ -754,7 +765,9 @@ def test_octree_obstacle_reaches_the_scene_as_occupied_cells( obstacle = Obstacle( name="voxel-map", obstacle_type=ObstacleType.OCTREE, - pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + pose=PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + ), points=((0.0, 0.0, 0.0), (0.05, 0.0, 0.0)), octree_resolution=0.05, ) @@ -782,7 +795,9 @@ def test_octree_obstacle_survives_the_deepcopy_and_equality_the_world_does( obstacle = Obstacle( name="voxel-map", obstacle_type=ObstacleType.OCTREE, - pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + pose=PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + ), points=((0.0, 0.0, 0.0),), octree_resolution=0.05, ) @@ -804,7 +819,9 @@ def octree(**kwargs: object) -> Obstacle: base: dict[str, object] = { "name": "voxel-map", "obstacle_type": ObstacleType.OCTREE, - "pose": PoseStamped(position=Vector3(), orientation=Quaternion()), + "pose": PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + ), "points": ((0.0, 0.0, 0.0),), "octree_resolution": 0.05, } @@ -833,7 +850,9 @@ def test_obstacle_operations_require_finalization( obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, - pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + pose=PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + ), dimensions=(0.1, 0.2, 0.3), ) @@ -859,7 +878,9 @@ def test_failed_obstacle_add_rolls_back_and_can_be_retried( obstacle = Obstacle( name="retryable", obstacle_type=ObstacleType.BOX, - pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + pose=PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + ), dimensions=(0.1, 0.2, 0.3), ) add_box = mocker.patch.object( @@ -887,7 +908,9 @@ def test_concurrent_remove_waits_for_obstacle_add( obstacle = Obstacle( name="concurrent", obstacle_type=ObstacleType.BOX, - pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + pose=PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + ), dimensions=(0.1, 0.2, 0.3), ) backend_add_started = threading.Event() @@ -935,7 +958,9 @@ def test_obstacle_ids_are_world_owned_and_invalid_insertions_are_rejected( unnamed = Obstacle( name="", obstacle_type=ObstacleType.BOX, - pose=PoseStamped(position=Vector3(), orientation=Quaternion()), # type: ignore[call-arg] + pose=PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + ), # type: ignore[call-arg] dimensions=(0.1, 0.1, 0.1), ) named = replace(unnamed, name="world-owned") @@ -955,7 +980,9 @@ def test_complete_update_rejects_invalid_obstacle_values( valid = Obstacle( name="shape", obstacle_type=ObstacleType.BOX, - pose=PoseStamped(position=Vector3(), orientation=Quaternion()), # type: ignore[call-arg] + pose=PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + ), # type: ignore[call-arg] dimensions=(0.1, 0.2, 0.3), ) invalid_obstacles = [ @@ -970,11 +997,14 @@ def test_complete_update_rejects_invalid_obstacle_values( replace( valid, pose=PoseStamped( - position=[np.nan, 0.0, 0.0], - orientation=[0.0, 0.0, 0.0, 1.0], + header=Header(frame_id=""), + pose=Pose( + position=Point(x=np.nan, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), ), - "finite values", + "Position components must be finite", ), ] @@ -992,7 +1022,9 @@ def test_complete_replacement_and_defensive_obstacle_snapshots( original = Obstacle( name="shape", obstacle_type=ObstacleType.BOX, - pose=PoseStamped(position=Vector3(), orientation=Quaternion()), # type: ignore[call-arg] + pose=PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + ), # type: ignore[call-arg] dimensions=(0.1, 0.2, 0.3), color=(1.0, 0.0, 0.0, 1.0), ) @@ -1003,7 +1035,10 @@ def test_complete_replacement_and_defensive_obstacle_snapshots( replacement = Obstacle( name="shape", obstacle_type=ObstacleType.SPHERE, - pose=PoseStamped(position=Vector3(1, 2, 3), orientation=Quaternion()), # type: ignore[call-arg] + pose=PoseStamped( + header=Header(frame_id=""), + pose=Pose(position=Point(x=1, y=2, z=3), orientation=Quaternion()), + ), # type: ignore[call-arg] dimensions=(0.4,), color=(0.0, 1.0, 0.0, 0.5), ) @@ -1016,9 +1051,9 @@ def test_complete_replacement_and_defensive_obstacle_snapshots( replacement.dimensions = (2.0,) stored.dimensions = (3.0,) - stored.pose.position.x = 99.0 + stored.pose.pose.position.x = 99.0 assert world.get_obstacles()[0].dimensions == (0.4,) - assert world.get_obstacles()[0].pose.position.x == pytest.approx(1.0) + assert world.get_obstacles()[0].pose.pose.position.x == pytest.approx(1.0) assert world.update_obstacle(replace(replacement, name="missing")) is False @@ -1031,7 +1066,9 @@ def test_collision_query_blocks_during_obstacle_replacement( obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, - pose=PoseStamped(position=Vector3(), orientation=Quaternion()), # type: ignore[call-arg] + pose=PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + ), # type: ignore[call-arg] dimensions=(0.1, 0.2, 0.3), ) world.add_obstacle(obstacle) @@ -1076,7 +1113,9 @@ def test_obstacle_replacement_blocks_during_collision_query( obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, - pose=PoseStamped(position=Vector3(), orientation=Quaternion()), # type: ignore[call-arg] + pose=PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + ), # type: ignore[call-arg] dimensions=(0.1, 0.2, 0.3), ) world.add_obstacle(obstacle) @@ -1121,7 +1160,9 @@ def test_native_update_failure_invalidates_world( obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, - pose=PoseStamped(position=Vector3(), orientation=Quaternion()), # type: ignore[call-arg] + pose=PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + ), # type: ignore[call-arg] dimensions=(0.1, 0.2, 0.3), ) world.add_obstacle(obstacle) @@ -1150,7 +1191,9 @@ def test_native_pose_update_failure_invalidates_world( obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, - pose=PoseStamped(position=Vector3(), orientation=Quaternion()), # type: ignore[call-arg] + pose=PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + ), # type: ignore[call-arg] dimensions=(0.1, 0.2, 0.3), ) world.add_obstacle(obstacle) @@ -1214,7 +1257,9 @@ def test_generic_planner_allows_update_between_collision_checks( obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, - pose=PoseStamped(position=Vector3(), orientation=Quaternion()), # type: ignore[call-arg] + pose=PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + ), # type: ignore[call-arg] dimensions=(0.1, 0.2, 0.3), ) world.add_obstacle(obstacle) @@ -1253,7 +1298,7 @@ def test_fk_jacobian_and_explicit_min_distance_unsupported( world.set_joint_state(ctx, JointState(name=["joint1", "joint2"], position=[0.25, 0.5])) pose = world.get_ee_pose(ctx) - assert pose.position.x == pytest.approx(0.75) + assert pose.pose.position.x == pytest.approx(0.75) assert world.get_jacobian(ctx).shape == (6, 2) with pytest.raises(NotImplementedError, match="get_min_distance"): world.get_min_distance(ctx) @@ -1301,14 +1346,14 @@ def fake_jacobian( ctx = world.get_live_context() world.set_joint_state( ctx, - JointState({"name": ["joint1", "joint2", "joint3"], "position": [1.0, 2.0, 3.0]}), + JointState(name=["joint1", "joint2", "joint3"], position=[1.0, 2.0, 3.0]), ) pose = world.get_group_ee_pose(ctx, "wrist") jacobian = world.get_group_jacobian(ctx, "wrist") assert fk_frames == ["tcp"] - assert pose.position.x == pytest.approx(6.0) + assert pose.pose.position.x == pytest.approx(6.0) np.testing.assert_allclose(jacobian, np.arange(18, dtype=np.float64).reshape(6, 3)[:, [2, 0]]) @@ -1616,7 +1661,9 @@ def test_native_planning_blocks_obstacle_replacement( obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, - pose=PoseStamped(position=Vector3(), orientation=Quaternion()), # type: ignore[call-arg] + pose=PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + ), # type: ignore[call-arg] dimensions=(0.1, 0.2, 0.3), ) world.add_obstacle(obstacle) @@ -1720,8 +1767,8 @@ def capture_scene_start( ) assert result.status == PlanningStatus.SUCCESS - assert result.path[0].position == pytest.approx(start.position) - assert observed_scene_start[:2] == pytest.approx(start.position) + assert list(result.path[0].position) == pytest.approx(list(start.position)) + assert observed_scene_start[:2] == pytest.approx(list(start.position)) def test_native_selected_planner_supports_non_overlapping_multi_group_selection( @@ -1747,7 +1794,7 @@ def test_native_selected_planner_supports_non_overlapping_multi_group_selection( assert result.status == PlanningStatus.SUCCESS assert result.path[-1].name == ["joint1", "joint2"] - assert result.path[-1].position == pytest.approx([0.1, 0.1]) + assert list(result.path[-1].position) == pytest.approx([0.1, 0.1]) def test_cartesian_planner_returns_timed_canonical_joint_states_and_options( @@ -1781,9 +1828,13 @@ def test_cartesian_planner_returns_timed_canonical_joint_states_and_options( JointState(name=list(selection.joint_names), position=[0.0, 0.0]), { "manipulator": _relative_target( - Transform( - translation=Vector3(0.1, 0.0, 0.0), - rotation=Quaternion.from_euler(Vector3(0.0, 0.0, np.pi / 2.0)), + TransformStamped( + header=Header(frame_id="world"), + transform=Transform( + translation=Vector3(x=0.1, y=0.0, z=0.0), + rotation=quaternion_from_euler(0.0, 0.0, np.pi / 2.0), + ), + child_frame_id="", ) ) }, @@ -1796,8 +1847,8 @@ def test_cartesian_planner_returns_timed_canonical_joint_states_and_options( assert result.status == PlanningStatus.SUCCESS assert result.timestamps == [0.0, 0.02, 0.04] assert [state.name for state in result.path] == [list(selection.joint_names)] * 3 - assert result.path[-1].position == pytest.approx([0.1, 0.1]) - assert result.path[1].velocity == pytest.approx([0.5, 0.5]) + assert list(result.path[-1].position) == pytest.approx([0.1, 0.1]) + assert list(result.path[1].velocity) == pytest.approx([0.5, 0.5]) planner = FakeCartesianPathPlanner.instances[-1] assert planner.options.speed_mode == FakeCartesianSpeedMode.TimeOptimal for field_name, expected in option_overrides.items(): @@ -1830,8 +1881,16 @@ def test_cartesian_zero_rotation_preserves_start_orientation( JointState(name=list(selection.joint_names), position=[0.0, 0.0]), { "manipulator": _relative_target( - Transform(translation=Vector3(0.05, 0.02, 0.0)), - Transform(translation=Vector3(0.1, 0.0, 0.0)), + TransformStamped( + header=Header(frame_id="world"), + transform=Transform(translation=Vector3(x=0.05, y=0.02, z=0.0)), + child_frame_id="", + ), + TransformStamped( + header=Header(frame_id="world"), + transform=Transform(translation=Vector3(x=0.1, y=0.0, z=0.0)), + child_frame_id="", + ), ) }, RoboPlanCartesianPathConfig(), @@ -1858,12 +1917,20 @@ def test_cartesian_uses_explicit_start_after_live_state_advances( world, selection, start, - {"manipulator": _relative_target(Transform(translation=Vector3(0.1, 0.0, 0.0)))}, + { + "manipulator": _relative_target( + TransformStamped( + header=Header(frame_id="world"), + transform=Transform(translation=Vector3(x=0.1, y=0.0, z=0.0)), + child_frame_id="", + ) + ) + }, RoboPlanCartesianPathConfig(), ) assert result.status == PlanningStatus.SUCCESS - assert result.path[0].position == pytest.approx(start.position) + assert list(result.path[0].position) == pytest.approx(list(start.position)) start_pose = FakeCartesianPathPlanner.instances[-1].paths[0].tforms[0][0] assert start_pose[0, 3] == pytest.approx(sum(start.position)) @@ -1886,7 +1953,15 @@ def test_cartesian_rejects_official_planner_failure( world, selection, JointState(name=list(selection.joint_names), position=[0.0, 0.0]), - {"manipulator": _relative_target(Transform(translation=Vector3(0.1, 0.0, 0.0)))}, + { + "manipulator": _relative_target( + TransformStamped( + header=Header(frame_id="world"), + transform=Transform(translation=Vector3(x=0.1, y=0.0, z=0.0)), + child_frame_id="", + ) + ) + }, RoboPlanCartesianPathConfig(), ) @@ -1938,7 +2013,15 @@ def collides_only_mid_edge(scene: FakeScene, q: np.ndarray) -> bool: world, selection, JointState(name=list(selection.joint_names), position=[0.0, 0.0]), - {"manipulator": _relative_target(Transform(translation=Vector3(0.05, 0.0, 0.0)))}, + { + "manipulator": _relative_target( + TransformStamped( + header=Header(frame_id="world"), + transform=Transform(translation=Vector3(x=0.05, y=0.0, z=0.0)), + child_frame_id="", + ) + ) + }, RoboPlanCartesianPathConfig(), ) @@ -2062,8 +2145,9 @@ def test_composed_model_fills_only_missing_acceleration_limits( def test_base_pose_is_written_to_composed_model( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: - robot_config.base_pose = PoseStamped( # type: ignore[call-arg] - position=Vector3(1, 0, 0), orientation=Quaternion() + robot_config.base_pose = PoseStamped( + header=Header(frame_id=""), + pose=Pose(position=Point(x=1, y=0, z=0), orientation=Quaternion()), ) world = _make_world(fake_roboplan, robot_config) diff --git a/dimos/manipulation/test_sdk.py b/dimos/manipulation/test_sdk.py index 7b31097f7d4..cc5b1bc51d5 100644 --- a/dimos/manipulation/test_sdk.py +++ b/dimos/manipulation/test_sdk.py @@ -15,6 +15,10 @@ import atexit from dataclasses import replace +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header +from dimos_generated.trajectory_msgs.msg import JointTrajectory from IPython.core.completer import provisionalcompleter from IPython.core.interactiveshell import InteractiveShell import numpy as np @@ -38,9 +42,7 @@ ) from dimos.manipulation.planning.spec.models import GeneratedPlan from dimos.manipulation.sdk import Arm, MotionError -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.time import header_now from dimos.porcelain.dimos import Dimos @@ -60,7 +62,11 @@ def rpc(mocker): "arm": PlanningGroupState( joints=JointState(name=["j0", "j1"], position=[0.1, 0.2]), end_effector_pose=PoseStamped( - frame_id="world", position=[0.4, 0.0, 0.3], orientation=[0.0, 1.0, 0.0, 0.0] + header=Header(frame_id="world"), + pose=Pose( + position=Point(x=0.4, y=0.0, z=0.3), + orientation=Quaternion(x=0.0, y=1.0, z=0.0, w=0.0), + ), ), gripper_position=0.5, joint_presets={ @@ -71,10 +77,22 @@ def rpc(mocker): }, ) proxy.plan_to_joints.return_value = PlanResult( - PlanStatus.SUCCEEDED, plan=GeneratedPlan(("arm",), JointTrajectory()) + PlanStatus.SUCCEEDED, + plan=GeneratedPlan( + ("arm",), + JointTrajectory( + header=header_now(), + ), + ), ) proxy.plan_to_poses.return_value = PlanResult( - PlanStatus.SUCCEEDED, plan=GeneratedPlan(("arm",), JointTrajectory()) + PlanStatus.SUCCEEDED, + plan=GeneratedPlan( + ("arm",), + JointTrajectory( + header=header_now(), + ), + ), ) proxy.execute.return_value = ExecutionResult(ExecutionStatus.COMPLETED) proxy.move_linear.return_value = MoveResult( @@ -173,7 +191,7 @@ def test_explicit_group_and_module_selection(app, rpc): def test_joints_are_fresh_arrays_in_declared_order(arm, rpc): state = rpc.get_state.return_value.groups["arm"] - state.joints.name[:] = ["j1", "j0"] + state.joints.name = ["j1", "j0"] positions = arm.joints() np.testing.assert_array_equal(positions, [0.2, 0.1]) @@ -226,9 +244,16 @@ def test_pose_preserves_current_orientation(arm, rpc): arm.move_pose([0.5, 0.1, 0.4], speed_scale=0.3, timeout=20.0) target = rpc.plan_to_poses.call_args.args[0]["arm"] - assert target.frame_id == "world" - np.testing.assert_array_equal(target.position.to_numpy(), [0.5, 0.1, 0.4]) - assert target.orientation.to_tuple() == (0.0, 1.0, 0.0, 0.0) + assert target.header.frame_id == "world" + np.testing.assert_array_equal( + [target.pose.position.x, target.pose.position.y, target.pose.position.z], [0.5, 0.1, 0.4] + ) + assert ( + target.pose.orientation.x, + target.pose.orientation.y, + target.pose.orientation.z, + target.pose.orientation.w, + ) == (0.0, 1.0, 0.0, 0.0) assert rpc.plan_to_poses.call_args.kwargs == {"speed_scale": 0.3} rpc.execute.assert_called_once_with( blocking=True, timeout=20.0, plan_id=rpc.plan_to_poses.return_value.plan.plan_id @@ -239,7 +264,12 @@ def test_explicit_orientation_needs_no_state_read(arm, rpc): arm.move_pose([0.4, 0.0, 0.3], orientation=(0.0, 0.0, 0.0, 1.0)) target = rpc.plan_to_poses.call_args.args[0]["arm"] - assert target.orientation.to_tuple() == (0.0, 0.0, 0.0, 1.0) + assert ( + target.pose.orientation.x, + target.pose.orientation.y, + target.pose.orientation.z, + target.pose.orientation.w, + ) == (0.0, 0.0, 0.0, 1.0) rpc.get_state.assert_not_called() diff --git a/dimos/manipulation/visualization/operator.py b/dimos/manipulation/visualization/operator.py index 55649ce6676..19d832bea25 100644 --- a/dimos/manipulation/visualization/operator.py +++ b/dimos/manipulation/visualization/operator.py @@ -17,15 +17,17 @@ from __future__ import annotations from collections.abc import Mapping, Sequence +import copy from dataclasses import dataclass, field import math from typing import TYPE_CHECKING +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import JointState + from dimos.manipulation.planning.groups.models import PlanningGroup from dimos.manipulation.planning.planners.config import CartesianPathConfig from dimos.manipulation.planning.spec.models import GeneratedPlan, PlanningGroupID -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.JointState import JointState if TYPE_CHECKING: from dimos.manipulation.manipulation_module import ManipulationModule @@ -111,7 +113,7 @@ def set_motion_speed(self, speed_scale: float) -> bool: def get_init_joints(self) -> JointState | None: """Return the operator-authoritative initialization state.""" init = self._module.get_init_joints() - return None if init is None else JointState(init) + return None if init is None else copy.deepcopy(init) def evaluate_joint_target(self, request: JointTargetRequest) -> TargetEvaluationResult: """Validate and evaluate a canonical joint target.""" @@ -122,7 +124,7 @@ def evaluate_joint_target(self, request: JointTargetRequest) -> TargetEvaluation complete = self._complete_states(groups, request.target) if complete is None: return self._invalid(request.group_ids, "Incomplete model target state") - return self._evaluate_complete_target(groups, JointState(request.target), complete) + return self._evaluate_complete_target(groups, copy.deepcopy(request.target), complete) def evaluate_pose_target(self, request: PoseTargetRequest) -> TargetEvaluationResult: """Validate and evaluate explicit world-frame pose targets.""" @@ -132,7 +134,7 @@ def evaluate_pose_target(self, request: PoseTargetRequest) -> TargetEvaluationRe ik = self._module.inverse_kinematics( pose_targets=dict(request.pose_targets), auxiliary_group_ids=request.auxiliary_group_ids, - seed=JointState(request.seed) if request.seed is not None else None, + seed=copy.deepcopy(request.seed) if request.seed is not None else None, check_collision=True, ) if not ik.is_success() or ik.joint_state is None: @@ -155,12 +157,9 @@ def plan_to_joints(self, request: JointTargetRequest) -> GeneratedPlan | None: assert groups is not None targets = { group.id: JointState( - { - "name": list(group.joint_names), - "position": list( - request.target.position[offset : offset + len(group.joint_names)] - ), - } + name=list(group.joint_names), + header=request.target.header, + position=list(request.target.position)[offset : offset + len(group.joint_names)], ) for group, offset in self._group_offsets(groups) } @@ -260,9 +259,9 @@ def _validate_pose_request( group_ids, f"Planning group '{group.id}' has no tip_link" ) for group_id, pose in request.pose_targets.items(): - if pose.frame_id != "world": + if pose.header.frame_id != "world": return group_ids, self._invalid( - group_ids, f"Unsupported pose frame for '{group_id}': {pose.frame_id}" + group_ids, f"Unsupported pose frame for '{group_id}': {pose.header.frame_id}" ) if not self._pose_is_finite(pose): return group_ids, self._invalid( @@ -315,7 +314,7 @@ def _complete_states( if value is None: return None positions.append(value) - return JointState({"name": list(config.joint_names), "position": positions}) + return JointState(name=list(config.joint_names), position=positions) def _evaluate_complete_target( self, @@ -350,7 +349,7 @@ def _evaluate_complete_target( message="Target is collision-free" if valid else "Target is infeasible", collision_free=valid, group_ids=tuple(group.id for group in groups), - target_joints=JointState(target), + target_joints=copy.deepcopy(target), group_diagnostics=diagnostics, group_poses=poses, ) @@ -370,5 +369,6 @@ def _group_offsets(groups: Sequence[PlanningGroup]) -> tuple[tuple[PlanningGroup @staticmethod def _pose_is_finite(pose: PoseStamped) -> bool: - values = [*pose.position, *pose.orientation] + point, rotation = pose.pose.position, pose.pose.orientation + values = [point.x, point.y, point.z, rotation.x, rotation.y, rotation.z, rotation.w] return len(values) == 7 and all(math.isfinite(float(value)) for value in values) diff --git a/dimos/manipulation/visualization/test_factory.py b/dimos/manipulation/visualization/test_factory.py index bedfdcce871..d1b770e120e 100644 --- a/dimos/manipulation/visualization/test_factory.py +++ b/dimos/manipulation/visualization/test_factory.py @@ -18,6 +18,11 @@ from pathlib import Path from unittest.mock import MagicMock +from dimos_generated.dimos_msgs.msg import GraspCandidateArray +from dimos_generated.geometry_msgs.msg import Pose, PoseStamped +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header +from dimos_generated.trajectory_msgs.msg import JointTrajectory import numpy as np from numpy.typing import NDArray from pydantic import ValidationError @@ -43,10 +48,7 @@ ) from dimos.manipulation.visualization.factory import create_manipulation_visualization from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.time import header_now from dimos.robot.assets.model import LoadedRobotModel, RobotModel @@ -97,7 +99,7 @@ def load_model(self, model: PreparedRobotModel) -> None: def get_model_config(self) -> RobotModelConfig: return RobotModelConfig( model=RobotModel.from_file(Path("fake.urdf")), - base_pose=PoseStamped(), + base_pose=PoseStamped(header=Header(frame_id=""), pose=Pose()), joint_names=["joint1"], planning_groups=[ PlanningGroupDefinition( @@ -160,7 +162,7 @@ def set_joint_state(self, ctx: object, joint_state: JointState) -> None: return None def get_joint_state(self, ctx: object) -> JointState: - return JointState({}) + return JointState() def is_collision_free(self, ctx: object) -> bool: return True @@ -180,7 +182,7 @@ def check_edge_collision_free( return True def get_ee_pose(self, ctx: object) -> PoseStamped: - return PoseStamped() + return PoseStamped(header=Header(frame_id=""), pose=Pose()) def get_link_pose(self, ctx: object, link_name: str) -> NDArray[np.float64]: return np.eye(4, dtype=np.float64) @@ -189,7 +191,7 @@ def get_jacobian(self, ctx: object) -> NDArray[np.float64]: return np.zeros((6, 0), dtype=np.float64) def get_group_ee_pose(self, ctx: object, group_id: str) -> PoseStamped: - return PoseStamped() + return PoseStamped(header=Header(frame_id=""), pose=Pose()) def get_group_jacobian(self, ctx: object, group_id: str) -> NDArray[np.float64]: return np.zeros((6, 0), dtype=np.float64) @@ -323,11 +325,11 @@ def test_create_visualization_meshcat_accepts_structural_world() -> None: PlanningSceneInfo(model=fake_world.get_prepared_model()), operator=object() ) frame = VisualizationStateFrame(joint_state=None) - trajectory = JointTrajectory(joint_names=["arm/j1"], points=[]) + trajectory = JointTrajectory(header=header_now(), joint_names=["arm/j1"], points=[]) obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, - pose=PoseStamped(), + pose=PoseStamped(header=Header(frame_id=""), pose=Pose()), dimensions=(1.0, 1.0, 1.0), ) visualization.initialize(session) @@ -394,7 +396,7 @@ def test_drake_meshcat_visualization_lifecycle_is_noop_without_meshcat() -> None obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, - pose=PoseStamped(), + pose=PoseStamped(header=Header(frame_id=""), pose=Pose()), dimensions=(1.0, 1.0, 1.0), ) world.add_vis_obstacle("box", obstacle) diff --git a/dimos/manipulation/visualization/test_operator.py b/dimos/manipulation/visualization/test_operator.py index 2943a237890..8e9cfbc483c 100644 --- a/dimos/manipulation/visualization/test_operator.py +++ b/dimos/manipulation/visualization/test_operator.py @@ -18,6 +18,11 @@ from types import SimpleNamespace from unittest.mock import MagicMock +from dimos_generated.geometry_msgs.msg import Pose, PoseStamped +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint + from dimos.manipulation.manipulation_spec import ( CommandResult, CommandStatus, @@ -37,10 +42,7 @@ ManipulationOperator, PoseTargetRequest, ) -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +from dimos.msgs.time import duration_from_seconds, header_now from dimos.robot.assets.model import RobotModel @@ -58,8 +60,11 @@ def _config() -> RobotModelConfig: def _plan() -> GeneratedPlan: trajectory = JointTrajectory( + header=header_now(), joint_names=["left/j1", "left/j2"], - points=[TrajectoryPoint(positions=[0.4, 0.5], time_from_start=1.0)], + points=[ + JointTrajectoryPoint(positions=[0.4, 0.5], time_from_start=duration_from_seconds(1.0)) + ], ) return GeneratedPlan( group_ids=("left_arm",), @@ -96,7 +101,9 @@ def _operator() -> tuple[ManipulationOperator, MagicMock, MagicMock]: name=config.joint_names, position=[0.0, 0.0, 0.0] ) monitor.is_state_valid.return_value = True - monitor.get_group_ee_pose.return_value = PoseStamped(frame_id="world") + monitor.get_group_ee_pose.return_value = PoseStamped( + header=Header(frame_id="world"), pose=Pose() + ) return ManipulationOperator(module, monitor), module, monitor @@ -138,7 +145,7 @@ def test_joint_evaluation_rejects_noncanonical_unknown_and_overlapping_selection def test_pose_evaluation_routes_group_id_and_canonical_seed() -> None: operator, module, _ = _operator() - pose = PoseStamped(frame_id="world") + pose = PoseStamped(header=Header(frame_id="world"), pose=Pose()) seed = JointState(name=["left/j1", "left/j2"], position=[0.0, 0.0]) result = operator.evaluate_pose_target(PoseTargetRequest({"left_arm": pose}, seed=seed)) assert result.success @@ -149,9 +156,11 @@ def test_pose_evaluation_routes_group_id_and_canonical_seed() -> None: def test_pose_evaluation_rejects_non_world_frame_and_ambiguous_local_seed() -> None: operator, _, _ = _operator() - bad_frame = PoseTargetRequest({"left_arm": PoseStamped(frame_id="camera")}) + bad_frame = PoseTargetRequest( + {"left_arm": PoseStamped(header=Header(frame_id="camera"), pose=Pose())} + ) bad_seed = PoseTargetRequest( - {"left_arm": PoseStamped(frame_id="world")}, + {"left_arm": PoseStamped(header=Header(frame_id="world"), pose=Pose())}, seed=JointState(name=["j1", "j2"], position=[0.0, 0.0]), ) assert not operator.evaluate_pose_target(bad_frame).success @@ -163,7 +172,9 @@ def test_planning_and_actions_return_exact_generated_plan() -> None: joint_request = JointTargetRequest( ("left_arm",), JointState(name=["left/j1", "left/j2"], position=[0.1, 0.2]) ) - pose_request = PoseTargetRequest({"left_arm": PoseStamped(frame_id="world")}) + pose_request = PoseTargetRequest( + {"left_arm": PoseStamped(header=Header(frame_id="world"), pose=Pose())} + ) assert ( operator.plan_to_joints(joint_request) is module.generate_plan_to_joint_targets.return_value ) @@ -176,7 +187,7 @@ def test_planning_and_actions_return_exact_generated_plan() -> None: def test_cartesian_planning_uses_current_group_pose() -> None: operator, module, monitor = _operator() - target = PoseStamped(frame_id="world") + target = PoseStamped(header=Header(frame_id="world"), pose=Pose()) config = RoboPlanCartesianPathConfig() result = operator.plan_cartesian(CartesianTargetRequest({"left_arm": target}, config)) assert result is module.generate_cartesian_plan.return_value diff --git a/dimos/manipulation/visualization/viser/gui.py b/dimos/manipulation/visualization/viser/gui.py index c1e04aa57bf..b977c3f83fc 100644 --- a/dimos/manipulation/visualization/viser/gui.py +++ b/dimos/manipulation/visualization/viser/gui.py @@ -15,10 +15,15 @@ from __future__ import annotations from collections.abc import Callable, Mapping, Sequence +import copy import math from threading import RLock from typing import Any, TypeAlias, cast +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header + from dimos.manipulation.planning.groups.models import PlanningGroup from dimos.manipulation.planning.planners.roboplan_config import RoboPlanCartesianPathConfig from dimos.manipulation.planning.spec.config import RobotModelConfig @@ -51,9 +56,6 @@ TargetEvaluationWorker, TargetStatus, ) -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.utils.logging_config import setup_logger from dimos.utils.trigonometry import angle_diff @@ -96,7 +98,7 @@ def group_display_name(group: PlanningGroup) -> str: def _copy_joint_state(state: JointState | None) -> JointState | None: - return None if state is None else JointState(state) + return None if state is None else copy.deepcopy(state) ROBOT_DISPLAY_LABELS = tuple(mode.value.title() for mode in RobotDisplayMode) @@ -204,10 +206,7 @@ def _current_target_for_group(self, group: PlanningGroup) -> dict[PlanningGroupI return {} return { group.id: JointState( - { - "name": list(group.joint_names), - "position": [values[name] for name in group.joint_names], - } + name=list(group.joint_names), position=[values[name] for name in group.joint_names] ) } @@ -234,21 +233,16 @@ def evaluate_joint_target_set( names.extend(str(name) for name in target.name) positions.extend(float(value) for value in target.position) return self.operator.evaluate_joint_target( - JointTargetRequest(tuple(group_ids), JointState({"name": names, "position": positions})) + JointTargetRequest(tuple(group_ids), JointState(name=names, position=positions)) ) def evaluate_pose_target_set( self, - pose_targets: Mapping[PlanningGroupID, Pose], + pose_targets: Mapping[PlanningGroupID, PoseStamped], auxiliary_group_ids: Sequence[PlanningGroupID] = (), seed: JointState | None = None, ) -> TargetEvaluationResult: - stamped = { - group_id: PoseStamped( - frame_id="world", position=pose.position, orientation=pose.orientation - ) - for group_id, pose in pose_targets.items() - } + stamped = dict(pose_targets) return self.operator.evaluate_pose_target( PoseTargetRequest(stamped, tuple(auxiliary_group_ids), _copy_joint_state(seed)) ) @@ -271,24 +265,17 @@ def plan_to_selected_joints( names.extend(str(name) for name in target.name) positions.extend(float(value) for value in target.position) plan = self.operator.plan_to_joints( - JointTargetRequest(tuple(group_ids), JointState({"name": names, "position": positions})) + JointTargetRequest(tuple(group_ids), JointState(name=names, position=positions)) ) self.state.plan_state.plan = plan return plan is not None def plan_cartesian( self, - pose_targets: Mapping[PlanningGroupID, Pose], + pose_targets: Mapping[PlanningGroupID, PoseStamped], auxiliary_group_ids: Sequence[PlanningGroupID], ) -> bool: - stamped = { - group_id: PoseStamped( - frame_id="world", - position=pose.position, - orientation=pose.orientation, - ) - for group_id, pose in pose_targets.items() - } + stamped = dict(pose_targets) plan = self.operator.plan_cartesian( CartesianTargetRequest( stamped, @@ -701,10 +688,8 @@ def _initialize_selected_group_targets(self) -> None: if any(str(name) not in values for name in group.joint_names): continue self.state.group_joint_targets[group_id] = JointState( - { - "name": list(group.joint_names), - "position": [float(values[str(name)]) for name in group.joint_names], - } + name=list(group.joint_names), + position=[float(values[str(name)]) for name in group.joint_names], ) if group.has_pose_target and group_id not in self.state.pose_targets: pose = self.get_group_ee_pose(group_id) @@ -735,11 +720,9 @@ def _refresh_target_joints_from_groups(self) -> None: if target is not None: names.extend(str(name) for name in target.name) positions.extend(float(value) for value in target.position) - self.state.target_joints = ( - JointState({"name": names, "position": positions}) if names else None - ) + self.state.target_joints = JointState(name=names, position=positions) if names else None - def _active_pose_targets(self) -> dict[PlanningGroupID, Pose]: + def _active_pose_targets(self) -> dict[PlanningGroupID, PoseStamped]: return { group_id: self.state.pose_targets[group_id] for group_id in self.state.selected_group_ids @@ -806,7 +789,7 @@ def _apply_preset(self, preset: str) -> None: ) return positions = [float(values[str(name)]) for name in group.joint_names] - targets[group_id] = JointState({"name": list(group.joint_names), "position": positions}) + targets[group_id] = JointState(name=list(group.joint_names), position=positions) slider_values.append((group_id, group.joint_names, positions)) self.state.group_joint_targets.update(targets) with self._joint_controls_lock: @@ -854,9 +837,7 @@ def _target_set_from_sliders(self) -> dict[PlanningGroupID, JointState] | None: self._set_error(f"Missing target slider for {group_id}/{joint_name}") return None positions.append(float(handle.value)) - targets[group_id] = JointState( - {"name": list(group.joint_names), "position": positions} - ) + targets[group_id] = JointState(name=list(group.joint_names), position=positions) return targets def _on_joint_slider_update(self, _group_id: PlanningGroupID, _joint_name: str) -> None: @@ -893,7 +874,7 @@ def _on_transform_update( joints=( None if self.state.target_joints is None - else JointState(self.state.target_joints) + else copy.deepcopy(self.state.target_joints) ), pose_targets=dict(self._active_pose_targets()), ) @@ -926,7 +907,7 @@ def _move_joint_target_visuals(self, targets: Mapping[PlanningGroupID, JointStat state = self._target_ghost_state(targets) if state is not None: config = self.get_model_config() - self.scene.set_target_joints(config.joint_names, state.position) + self.scene.set_target_joints(config.joint_names, list(state.position)) def _target_ghost_state( self, targets: Mapping[PlanningGroupID, JointState] @@ -944,10 +925,7 @@ def _target_ghost_state( if not all(name in values for name in config.joint_names): return None return JointState( - { - "name": list(config.joint_names), - "position": [values[name] for name in config.joint_names], - } + name=list(config.joint_names), position=[values[name] for name in config.joint_names] ) def _sync_target_ghost_visibility(self) -> None: @@ -992,12 +970,12 @@ def _apply_target_evaluation_result( ) self.state.error = "" if success and collision_free else self.state.feasibility.message if result.target_joints is not None: - self.state.target_joints = JointState(result.target_joints) + self.state.target_joints = copy.deepcopy(result.target_joints) self._split_target_joints_by_group(result.target_joints) self.state.group_poses = { str(group_id): pose for group_id, pose in result.group_poses.items() - if isinstance(pose, Pose) + if isinstance(pose, PoseStamped) } if request.source == "joints": self._sync_pose_targets_from_group_poses() @@ -1025,10 +1003,8 @@ def _split_target_joints_by_group(self, target_joints: JointState) -> None: if group is None or any(str(name) not in positions for name in group.joint_names): continue self.state.group_joint_targets[group_id] = JointState( - { - "name": list(group.joint_names), - "position": [positions[str(name)] for name in group.joint_names], - } + name=list(group.joint_names), + position=[positions[str(name)] for name in group.joint_names], ) def _sync_pose_targets_from_group_poses(self) -> None: @@ -1118,7 +1094,7 @@ def _submit_plan(self) -> None: selection_epoch = self.state.selection_epoch target_sequence_id = self.state.latest_sequence_id joint_targets: dict[PlanningGroupID, JointState] | None = None - pose_targets: dict[PlanningGroupID, Pose] = {} + pose_targets: dict[PlanningGroupID, PoseStamped] = {} auxiliary_group_ids: tuple[PlanningGroupID, ...] = () if planning_mode == PlanningMode.CARTESIAN_SPACE: pose_targets = self._active_pose_targets() @@ -1413,10 +1389,15 @@ def _set_visible(self, key: str, visible: bool) -> None: def _set_optional_handle_attr(handle: object, attr: str, value: object) -> None: setattr(handle, attr, value) - def _pose_from_transform_target(self, target: TransformControlsHandle) -> Pose | None: + def _pose_from_transform_target(self, target: TransformControlsHandle) -> PoseStamped | None: px, py, pz = (float(value) for value in target.position) qw, qx, qy, qz = (float(value) for value in target.wxyz) - return Pose({"position": [px, py, pz], "orientation": [qx, qy, qz, qw]}) + return PoseStamped( + header=Header(frame_id="world"), + pose=Pose( + position=Point(x=px, y=py, z=pz), orientation=Quaternion(x=qx, y=qy, z=qz, w=qw) + ), + ) def _feasibility_status( self, result: TargetEvaluationResult, success: bool, collision_free: bool diff --git a/dimos/manipulation/visualization/viser/scene.py b/dimos/manipulation/visualization/viser/scene.py index 5591d1a263a..ecab0f6c0ae 100644 --- a/dimos/manipulation/visualization/viser/scene.py +++ b/dimos/manipulation/visualization/viser/scene.py @@ -25,6 +25,9 @@ from typing import Any, Protocol, TypeAlias, cast import xml.etree.ElementTree as ET +from dimos_generated.dimos_msgs.msg import GraspCandidate, GraspCandidateArray +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import JointState import numpy as np from numpy.typing import NDArray import trimesh @@ -43,11 +46,7 @@ VISER_INSTALL_HINT, VISER_URDF_INSTALL_HINT, ) -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.manipulation_msgs.GraspCandidate import GraspCandidate -from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray -from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.geometry import pose_matrix from dimos.robot.assets.model import LoadedRobotModel from dimos.utils.logging_config import setup_logger @@ -330,7 +329,7 @@ def show_grasp_proposals(self, candidates: GraspCandidateArray) -> None: if self._closed: return self.clear_grasp_proposals() - drawn = list(candidates.candidates[:GRASP_PROPOSAL_DRAW_LIMIT]) + drawn = list(candidates.candidates)[:GRASP_PROPOSAL_DRAW_LIMIT] if not drawn: return emphasis = min(GRASP_PROPOSAL_EMPHASIS_COUNT, len(drawn)) @@ -396,7 +395,7 @@ def _grasp_proposal_segments( ) for rank in range(start, stop): pose = candidates[rank].pose - rotation = np.asarray(pose.orientation.to_rotation_matrix(), dtype=np.float32) + rotation = np.asarray(pose_matrix(pose)[:3, :3], dtype=np.float32) origin = np.asarray( [pose.position.x, pose.position.y, pose.position.z], dtype=np.float32 ) @@ -448,12 +447,12 @@ def _obstacle_pose( ) -> tuple[tuple[float, float, float], tuple[float, float, float, float]]: pose = obstacle.pose return ( - (float(pose.position.x), float(pose.position.y), float(pose.position.z)), + (float(pose.pose.position.x), float(pose.pose.position.y), float(pose.pose.position.z)), ( - float(pose.orientation.w), - float(pose.orientation.x), - float(pose.orientation.y), - float(pose.orientation.z), + float(pose.pose.orientation.w), + float(pose.pose.orientation.x), + float(pose.pose.orientation.y), + float(pose.pose.orientation.z), ), ) @@ -661,14 +660,14 @@ def update_current_model(self, joint_state: JointState | None) -> None: return self._ensure_robot_urdfs(config) current = self._urdfs.get("current") - self.set_urdf_joints(current, config.joint_names, joint_state.position) + self.set_urdf_joints(current, config.joint_names, list(joint_state.position)) if self._target_tracks_current: - self._set_target_joints(config.joint_names, joint_state.position) + self._set_target_joints(config.joint_names, list(joint_state.position)) self._set_target_visibility(self._target_active) self.set_urdf_joints( self._collision_fallback_urdf, config.joint_names, - joint_state.position, + list(joint_state.position), ) def cancel_preview_animation(self) -> None: @@ -730,20 +729,20 @@ def _set_preview_ghost_joints( ghost = self._urdfs.get("preview") self.set_urdf_joints(ghost, joint_names, joints) - def set_target_pose(self, control_id: str, pose: Pose | None) -> None: + def set_target_pose(self, control_id: str, pose: PoseStamped | None) -> None: handle = self._handles.get(f"{control_id}:ee_control") if handle is None or pose is None: return handle.position = ( - float(pose.position.x), - float(pose.position.y), - float(pose.position.z), + float(pose.pose.position.x), + float(pose.pose.position.y), + float(pose.pose.position.z), ) handle.wxyz = ( - float(pose.orientation.w), - float(pose.orientation.x), - float(pose.orientation.y), - float(pose.orientation.z), + float(pose.pose.orientation.w), + float(pose.pose.orientation.x), + float(pose.pose.orientation.y), + float(pose.pose.orientation.z), ) def set_target_control_visual_state(self, control_id: str, feasible: bool) -> None: @@ -990,23 +989,21 @@ def _ensure_base_pose_frame(self, kind: str, config: RobotModelConfig) -> None: frame_name, show_axes=False, position=( - float(pose.position.x), - float(pose.position.y), - float(pose.position.z), + float(pose.pose.position.x), + float(pose.pose.position.y), + float(pose.pose.position.z), ), wxyz=( - float(pose.orientation.w), - float(pose.orientation.x), - float(pose.orientation.y), - float(pose.orientation.z), + float(pose.pose.orientation.w), + float(pose.pose.orientation.x), + float(pose.pose.orientation.y), + float(pose.pose.orientation.z), ), ) @staticmethod def _has_non_identity_base_pose(config: RobotModelConfig) -> bool: - pose = getattr(config, "base_pose", None) - if pose is None: - return False + pose = config.base_pose.pose return any( abs(value) > 1e-12 for value in ( diff --git a/dimos/manipulation/visualization/viser/state.py b/dimos/manipulation/visualization/viser/state.py index 90d43037dfe..f836b17fcc2 100644 --- a/dimos/manipulation/visualization/viser/state.py +++ b/dimos/manipulation/visualization/viser/state.py @@ -21,10 +21,11 @@ import threading from typing import Literal +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import JointState + from dimos.manipulation.planning.spec.models import GeneratedPlan, PlanningGroupID from dimos.manipulation.visualization.operator import TargetEvaluationResult -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -107,17 +108,17 @@ class PanelState: selected_group_ids: tuple[PlanningGroupID, ...] = () planning_mode: PlanningMode = PlanningMode.JOINT_SPACE selection_epoch: int = 0 - pose_targets: dict[PlanningGroupID, Pose] = field(default_factory=dict) + pose_targets: dict[PlanningGroupID, PoseStamped] = field(default_factory=dict) group_joint_targets: dict[PlanningGroupID, JointState] = field(default_factory=dict) target_joints: JointState | None = None - group_poses: dict[PlanningGroupID, Pose] = field(default_factory=dict) + group_poses: dict[PlanningGroupID, PoseStamped] = field(default_factory=dict) runtime: PanelRuntime = PanelRuntime.STOPPED backend_status: BackendConnectionStatus = BackendConnectionStatus.DISCONNECTED target_status: TargetStatus = TargetStatus.EMPTY action_status: ActionStatus = ActionStatus.IDLE manipulation_state: str = "DISCONNECTED" current_joints: list[float] | None = None - cartesian_target: Pose | None = None + cartesian_target: PoseStamped | None = None feasibility: FeasibilityState = field(default_factory=FeasibilityState) latest_sequence_id: int = 0 plan_state: PanelPlanState = field(default_factory=PanelPlanState) @@ -207,10 +208,10 @@ class TargetEvaluationRequest: source: PreviewSource selection_epoch: int = 0 group_ids: tuple[PlanningGroupID, ...] = () - pose: Pose | None = None + pose: PoseStamped | None = None joints: JointState | None = None auxiliary_group_ids: tuple[PlanningGroupID, ...] = () - pose_targets: dict[PlanningGroupID, Pose] = field(default_factory=dict) + pose_targets: dict[PlanningGroupID, PoseStamped] = field(default_factory=dict) joint_targets: dict[PlanningGroupID, JointState] = field(default_factory=dict) diff --git a/dimos/manipulation/visualization/viser/test_gui.py b/dimos/manipulation/visualization/viser/test_gui.py index d7b8b3ebcad..b89c7e42583 100644 --- a/dimos/manipulation/visualization/viser/test_gui.py +++ b/dimos/manipulation/visualization/viser/test_gui.py @@ -23,8 +23,14 @@ import pytest +from dimos.msgs.time import duration_from_seconds, header_now + pytest.importorskip("viser", reason="Viser optional dependency is not installed") +from dimos_generated.dimos_msgs.msg import TrajectoryStatus +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint + from dimos.control.tasks.trajectory_task.trajectory_task import ( TrajectoryExecutionResult, TrajectoryExecutionStatus, @@ -61,10 +67,7 @@ TargetEvaluationWorker, TargetStatus, ) -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint -from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState, TrajectoryStatus +from dimos.msgs.trajectory import TrajectoryState from dimos.robot.assets.model import LoadedRobotModel, PlanarBaseDefinition, RobotModel @@ -264,10 +267,15 @@ def executable_gui(monkeypatch, mocker): plan=GeneratedPlan( group_ids=gui.state.selected_group_ids, trajectory=JointTrajectory( + header=header_now(), joint_names=["arm/j0"], points=[ - TrajectoryPoint(positions=[0.0], time_from_start=0.0), - TrajectoryPoint(positions=[1.0], time_from_start=1.0), + JointTrajectoryPoint( + positions=[0.0], time_from_start=duration_from_seconds(0.0) + ), + JointTrajectoryPoint( + positions=[1.0], time_from_start=duration_from_seconds(1.0) + ), ], ), path=[JointState(name=["arm/j0"], position=[value]) for value in (0.0, 1.0)], @@ -353,7 +361,7 @@ def test_gui_completion_enables_next_plan_without_cancel(executable_gui, module_ def invoke(task, method, args=None): if method == "execute": return TrajectoryExecutionResult(TrajectoryExecutionStatus.ACCEPTED) - return TrajectoryStatus(state=reported["state"]) + return TrajectoryStatus(header=header_now(), state=reported["state"]) mocker.patch.object(module._control_coordinator, "task_invoke", side_effect=invoke) cancel = mocker.spy(module, "cancel") @@ -606,8 +614,10 @@ def test_gui_preview_enters_previewing_before_worker_runs( gui.state.plan_state.target_sequence_id = gui.state.latest_sequence_id gui.state.plan_state.plan = GeneratedPlan( group_ids=gui.state.selected_group_ids, - trajectory=JointTrajectory(), - path=[JointState({"name": [], "position": []})], + trajectory=JointTrajectory( + header=header_now(), + ), + path=[JointState(name=[], position=[])], ) assert gui.state.can_execute() is True @@ -643,8 +653,10 @@ def test_gui_selection_change_clears_invalidated_preview( gui.state.plan_state.target_sequence_id = gui.state.latest_sequence_id gui.state.plan_state.plan = GeneratedPlan( group_ids=gui.state.selected_group_ids, - trajectory=JointTrajectory(), - path=[JointState({"name": [], "position": []})], + trajectory=JointTrajectory( + header=header_now(), + ), + path=[JointState(name=[], position=[])], ) gui._submit_preview() @@ -679,8 +691,10 @@ def test_gui_selection_change_ignores_invalidated_preview_error( gui.state.plan_state.target_sequence_id = gui.state.latest_sequence_id gui.state.plan_state.plan = GeneratedPlan( group_ids=gui.state.selected_group_ids, - trajectory=JointTrajectory(), - path=[JointState({"name": [], "position": []})], + trajectory=JointTrajectory( + header=header_now(), + ), + path=[JointState(name=[], position=[])], ) gui._submit_preview() diff --git a/dimos/manipulation/visualization/viser/test_scene_grasp_proposals.py b/dimos/manipulation/visualization/viser/test_scene_grasp_proposals.py index 88ec1369799..f81f7de9131 100644 --- a/dimos/manipulation/visualization/viser/test_scene_grasp_proposals.py +++ b/dimos/manipulation/visualization/viser/test_scene_grasp_proposals.py @@ -17,22 +17,22 @@ import inspect from unittest.mock import MagicMock +from dimos_generated.geometry_msgs.msg import Point import numpy as np import pytest pytest.importorskip("viser", reason="Viser optional dependency is not installed") +from dimos_generated.dimos_msgs.msg import GraspCandidate, GraspCandidateArray +from dimos_generated.geometry_msgs.msg import Pose, Quaternion +from dimos_generated.std_msgs.msg import Header + from dimos.manipulation.visualization.viser.scene import ( GRASP_PROPOSAL_DRAW_LIMIT, GRASP_PROPOSAL_EMPHASIS_COUNT, ViserManipulationScene, ) -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.manipulation_msgs.GraspCandidate import GraspCandidate -from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray -from dimos.msgs.std_msgs.Header import Header +from dimos.msgs.time import time_from_seconds def _scene() -> ViserManipulationScene: @@ -41,11 +41,14 @@ def _scene() -> ViserManipulationScene: def _proposals(count: int) -> GraspCandidateArray: return GraspCandidateArray( - Header(1.0, "world"), - [ + header=Header(stamp=time_from_seconds(1.0), frame_id="world"), + candidates=[ GraspCandidate( - Pose(Vector3(0.1 * index, 0.0, 0.2), Quaternion(0.0, 0.0, 0.0, 1.0)), - 1.0 - 0.01 * index, + pose=Pose( + position=Point(x=0.1 * index, y=0.0, z=0.2), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + score=1.0 - 0.01 * index, ) for index in range(count) ], diff --git a/dimos/manipulation/visualization/viser/test_scene_obstacles.py b/dimos/manipulation/visualization/viser/test_scene_obstacles.py index fd88770c5d9..02bda7e1916 100644 --- a/dimos/manipulation/visualization/viser/test_scene_obstacles.py +++ b/dimos/manipulation/visualization/viser/test_scene_obstacles.py @@ -16,15 +16,18 @@ from unittest.mock import MagicMock +from dimos_generated.geometry_msgs.msg import Pose +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest pytest.importorskip("viser", reason="Viser optional dependency is not installed") +from dimos_generated.geometry_msgs.msg import PoseStamped + from dimos.manipulation.planning.spec.enums import ObstacleType from dimos.manipulation.planning.spec.models import Obstacle from dimos.manipulation.visualization.viser.scene import ViserManipulationScene -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped def _scene() -> ViserManipulationScene: @@ -35,7 +38,7 @@ def _octree(points: np.ndarray, resolution: float = 0.025) -> Obstacle: return Obstacle( name="mapping/voxel-map", obstacle_type=ObstacleType.OCTREE, - pose=PoseStamped(frame_id="world"), + pose=PoseStamped(header=Header(frame_id="world"), pose=Pose()), points=points, octree_resolution=resolution, ) diff --git a/dimos/manipulation/visualization/viser/test_viser_visualization.py b/dimos/manipulation/visualization/viser/test_viser_visualization.py index 30d96f2b385..e7f57cdd812 100644 --- a/dimos/manipulation/visualization/viser/test_viser_visualization.py +++ b/dimos/manipulation/visualization/viser/test_viser_visualization.py @@ -19,8 +19,13 @@ import pytest +from dimos.msgs.time import duration_from_seconds, header_now + pytest.importorskip("viser", reason="Viser optional dependency is not installed") +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint + from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.models import ( @@ -34,9 +39,6 @@ scaled_frame_delays, ) from dimos.manipulation.visualization.viser.visualizer import ViserManipulationVisualizer -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint from dimos.robot.assets.model import RobotModel @@ -69,8 +71,9 @@ def test_visualizer_builds_full_model_preview_from_selected_canonical_joints() - visualizer._model_config = _model() visualizer._current_state = JointState(name=["left/j1", "right/j1"], position=[0.1, 0.2]) trajectory = JointTrajectory( + header=header_now(), joint_names=["right/j1"], - points=[TrajectoryPoint(positions=[0.8], time_from_start=1.0)], + points=[JointTrajectoryPoint(positions=[0.8], time_from_start=duration_from_seconds(1.0))], ) preview = visualizer._raw_preview_animation(trajectory) assert preview == PreviewAnimation(("left/j1", "right/j1"), (PreviewFrame(1.0, (0.1, 0.8)),)) @@ -82,8 +85,13 @@ def test_visualizer_rejects_unknown_or_duplicate_trajectory_joints() -> None: visualizer._current_state = JointState(name=["left/j1", "right/j1"], position=[0.1, 0.2]) for names in (["unknown"], ["left/j1", "left/j1"]): trajectory = JointTrajectory( + header=header_now(), joint_names=names, - points=[TrajectoryPoint(positions=[0.0] * len(names), time_from_start=1.0)], + points=[ + JointTrajectoryPoint( + positions=[0.0] * len(names), time_from_start=duration_from_seconds(1.0) + ) + ], ) assert visualizer._raw_preview_animation(trajectory) is None diff --git a/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py b/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py index 44631f6bcfb..4d12c4409c6 100644 --- a/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py +++ b/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py @@ -17,10 +17,18 @@ from pathlib import Path from types import SimpleNamespace +from dimos_generated.geometry_msgs.msg import Pose +from dimos_generated.std_msgs.msg import Header import pytest +from dimos.msgs.time import header_now + pytest.importorskip("viser", reason="Viser optional dependency is not installed") +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.trajectory_msgs.msg import JointTrajectory + from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.joint_space import JointSpace @@ -38,9 +46,6 @@ from dimos.manipulation.visualization.viser.runtime import ViserRuntime from dimos.manipulation.visualization.viser.scene import RobotDisplayMode, ViserManipulationScene from dimos.manipulation.visualization.viser.visualizer import ViserManipulationVisualizer -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory from dimos.robot.assets.model import LoadedRobotModel, RobotModel @@ -86,7 +91,7 @@ def stop(self) -> None: def fake_robot_config(name: str) -> RobotModelConfig: return RobotModelConfig( model=RobotModel.from_file(Path(f"{name}.urdf")), - base_pose=PoseStamped(), + base_pose=PoseStamped(header=Header(frame_id=""), pose=Pose()), joint_names=[], planning_groups=[ PlanningGroupDefinition( @@ -424,7 +429,7 @@ def test_visualizer_publish_preview_and_close_paths( monkeypatch: pytest.MonkeyPatch, ) -> None: calls: list[tuple[str, str]] = [] - current = JointState({"name": ["joint1"], "position": [0.5]}) + current = JointState(name=["joint1"], position=[0.5]) class FakeRuntime: url = "http://localhost:8095" @@ -476,7 +481,9 @@ def close(self) -> None: visualizer.cancel_preview_animation() visualizer.update_state(VisualizationStateFrame(current)) visualizer.cancel_preview_animation() - visualizer.animate_trajectory(JointTrajectory(joint_names=["arm/joint1"]), duration=1.5) + visualizer.animate_trajectory( + JointTrajectory(header=header_now(), joint_names=["arm/joint1"]), duration=1.5 + ) visualizer.close() visualizer.update_state(VisualizationStateFrame(current)) @@ -541,7 +548,7 @@ def __init__( config = fake_robot_config("arm") config.model = RobotModel.from_file(fixed_world_root) - config.base_pose.position.x = 1.0 + config.base_pose.pose.position.x = 1.0 def prepare(description: object, **kwargs: object) -> object: prepared.append(kwargs) @@ -595,7 +602,7 @@ def test_selected_display_mode_survives_primary_recreation_and_joint_updates( scene.register_model(prepared) current = scene._urdfs["current"] - scene.update_current_model(JointState({"name": ["joint1"], "position": [0.75]})) + scene.update_current_model(JointState(name=["joint1"], position=[0.75])) assert current is not old_current assert scene.robot_display_mode == mode diff --git a/dimos/manipulation/visualization/viser/visualizer.py b/dimos/manipulation/visualization/viser/visualizer.py index bfe364bf040..3fc8b0f1643 100644 --- a/dimos/manipulation/visualization/viser/visualizer.py +++ b/dimos/manipulation/visualization/viser/visualizer.py @@ -15,8 +15,13 @@ from __future__ import annotations from contextlib import suppress +import copy from typing import TYPE_CHECKING +from dimos_generated.dimos_msgs.msg import GraspCandidateArray +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.trajectory_msgs.msg import JointTrajectory + from dimos.manipulation.visualization.viser.animation import ( PreviewAnimation, PreviewFrame, @@ -30,9 +35,8 @@ ) from dimos.manipulation.visualization.viser.scene import ViserManipulationScene from dimos.manipulation.visualization.viser.theme import apply_dimos_theme -from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.time import to_seconds +from dimos.msgs.trajectory import trajectory_duration from dimos.utils.logging_config import setup_logger try: @@ -47,6 +51,8 @@ raise ModuleNotFoundError(VISER_URDF_INSTALL_HINT) from e if TYPE_CHECKING: + from dimos_generated.geometry_msgs.msg import PoseStamped + from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.models import ( Obstacle, @@ -54,7 +60,6 @@ VisualizationSession, VisualizationStateFrame, ) - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped logger = setup_logger() @@ -233,7 +238,7 @@ def update_state(self, frame: VisualizationStateFrame) -> None: if self._scene is None: return if frame.joint_state is not None: - self._current_state = JointState(frame.joint_state) + self._current_state = copy.deepcopy(frame.joint_state) self._scene.update_current_model(frame.joint_state) if self._gui is not None: self._gui.refresh() @@ -249,7 +254,10 @@ def animate_trajectory( preview = self._raw_preview_animation(trajectory) if preview is not None: self._scene.animate_preview( - preview, duration if duration is not None else max(float(trajectory.duration), 0.0) + preview, + duration + if duration is not None + else max(float(trajectory_duration(trajectory)), 0.0), ) def cancel_preview_animation(self) -> None: @@ -280,7 +288,7 @@ def _raw_preview_animation(self, trajectory: JointTrajectory) -> PreviewAnimatio for point in trajectory.points: selected = {name: float(point.positions[index]) for name, index in indices.items()} positions = [float(selected.get(name, baseline[name])) for name in config.joint_names] - frames.append(PreviewFrame(float(point.time_from_start), tuple(positions))) + frames.append(PreviewFrame(float(to_seconds(point.time_from_start)), tuple(positions))) return PreviewAnimation(tuple(config.joint_names), tuple(frames)) if frames else None @staticmethod diff --git a/dimos/mapping/cli/pose_fill.py b/dimos/mapping/cli/pose_fill.py index 38465ff59fb..d7fb3f1cb98 100644 --- a/dimos/mapping/cli/pose_fill.py +++ b/dimos/mapping/cli/pose_fill.py @@ -37,10 +37,11 @@ # Heavy dimos imports are deferred (TYPE_CHECKING / inside functions) so that # `dimos map --help` stays fast. See test_cli_startup.py. if TYPE_CHECKING: + from dimos_generated.geometry_msgs.msg import TransformStamped + from dimos.memory.backend import Backend from dimos.memory.stream import Stream from dimos.memory.type.observation import Observation - from dimos.msgs.geometry_msgs.Transform import Transform def pose_fill( @@ -48,7 +49,7 @@ def pose_fill( pose_stream: Stream[Any], *, tolerance: float = 0.1, - mount: Transform | None = None, + mount: TransformStamped | None = None, ) -> Stream[Any]: """Re-pose each observation in *stream* from the nearest entry in *pose_stream*. @@ -64,10 +65,10 @@ def pose_fill( camera_optical`` so a base-frame odometry source yields an optical-frame image pose. """ + from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion + from dimos.memory.type.observation import _to_tuple - from dimos.msgs.geometry_msgs.Quaternion import Quaternion - from dimos.msgs.geometry_msgs.Transform import Transform - from dimos.msgs.geometry_msgs.Vector3 import Vector3 + from dimos.msgs.geometry import pose_from_matrix, pose_matrix, transform_matrix def _fill(pair_obs: Observation[Any]) -> Observation[Any]: primary, secondary = cast( @@ -76,8 +77,13 @@ def _fill(pair_obs: Observation[Any]) -> Observation[Any]: if mount is None: return primary.with_pose(secondary.data) x, y, z, qx, qy, qz, qw = cast("tuple[float, ...]", _to_tuple(secondary.data)) - world_base = Transform(translation=Vector3(x, y, z), rotation=Quaternion(qx, qy, qz, qw)) - return primary.with_pose(world_base + mount) + world_base = Pose( + position=Point(x=x, y=y, z=z), + orientation=Quaternion(x=qx, y=qy, z=qz, w=qw), + ) + return primary.with_pose( + pose_from_matrix(pose_matrix(world_base) @ transform_matrix(mount.transform)) + ) return stream.align(pose_stream, tolerance=tolerance).map(_fill) diff --git a/dimos/mapping/cli/replay.py b/dimos/mapping/cli/replay.py index 9f30e89cfd2..6a93394cbe8 100644 --- a/dimos/mapping/cli/replay.py +++ b/dimos/mapping/cli/replay.py @@ -40,8 +40,10 @@ from dimos.memory.utils.progress import progress if TYPE_CHECKING: + from dimos_generated.sensor_msgs.msg import PointCloud2 + from dimos.memory.stream import Stream - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + TIMELINE = "ts" @@ -64,6 +66,8 @@ def _log_clouds( """ import rerun as rr + from dimos.visualization.rerun.message_helpers import cloud_archetype + n = total if total is not None else stream.count() with progress(n, label) as bar: for obs in stream: @@ -71,7 +75,9 @@ def _log_clouds( rr.set_time(TIMELINE, timestamp=obs.ts) rr.log( entity, - obs.data.to_rerun(voxel_size=voxel, mode=point_mode, bottom_cutoff=bottom_cutoff), + cloud_archetype( + obs.data, voxel_size=voxel, mode=point_mode, bottom_cutoff=bottom_cutoff + ), ) @@ -174,22 +180,27 @@ def main( ), ) -> None: """Dump a recording to .rrd (lidar clouds + camera frames) and open it in rerun.""" + from dimos_generated.geometry_msgs.msg import PoseStamped + from dimos_generated.nav_msgs.msg import Odometry + from dimos_generated.sensor_msgs.msg import Image, PointCloud2 import rerun as rr from dimos.mapping.voxels.module import VoxelMapTransformer from dimos.memory.cli.dataset import open_store, resolve_dataset, stream_payload_types from dimos.memory.transform import throttle - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - from dimos.msgs.nav_msgs.Odometry import Odometry - from dimos.msgs.sensor_msgs.Image import Image - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2, register_colormap_annotation - from dimos.robot.unitree.go2.connection import _camera_info_static + from dimos.robot.unitree.go2.camera_calibration import front_camera_calibration + from dimos.visualization.rerun.message_helpers import ( + camera_pinhole, + cloud_archetype, + image_archetype, + register_colormap_annotation, + ) src_path = resolve_dataset(dataset) store = open_store(src_path) if out is None: out = Path.cwd() / f"{src_path.stem}.rrd" - cam_info = _camera_info_static() + cam_info = front_camera_calibration() with store: # Resolve which streams to voxelize: all PointCloud2 streams, or the @@ -205,7 +216,7 @@ def main( # Static pinhole on the camera entity; per-frame Transform3D goes on the # same entity. Image is the child so it projects through the pinhole. - pinhole = cam_info.to_rerun() + pinhole = camera_pinhole(cam_info) assert not isinstance(pinhole, list) rr.log("world/camera", pinhole, static=True) @@ -268,8 +279,11 @@ def clipped(name: str, ptype: type[Any]) -> Stream[Any]: ).last() rr.log( f"world/{name}_map", - final.data.to_rerun( - voxel_size=voxel / 4, mode=point_mode, bottom_cutoff=bottom_cutoff + cloud_archetype( + final.data, + voxel_size=voxel / 4, + mode=point_mode, + bottom_cutoff=bottom_cutoff, ), static=True, ) @@ -339,7 +353,7 @@ def clipped(name: str, ptype: type[Any]) -> Stream[Any]: translation=[x, y, z], quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw]) ), ) - rr.log("world/camera/image", img_obs.data.to_rerun()) + rr.log("world/camera/image", image_archetype(img_obs.data)) print(f"wrote {out}") if no_gui: diff --git a/dimos/mapping/cli/replay_marker.py b/dimos/mapping/cli/replay_marker.py index d2e72a2bc47..3525ea7c69f 100644 --- a/dimos/mapping/cli/replay_marker.py +++ b/dimos/mapping/cli/replay_marker.py @@ -34,6 +34,7 @@ if TYPE_CHECKING: from dimos.memory.stream import Stream + TIMELINE = "ts" @@ -62,28 +63,29 @@ def main( ), ) -> None: """Dump an AprilTag detection replay to .rrd and open it in rerun.""" + from dimos_generated.sensor_msgs.msg import Image, PointCloud2 import rerun as rr from dimos.memory.cli.dataset import open_store, resolve_dataset from dimos.memory.transform import QualityWindow, SpeedLimit from dimos.memory.vis.color import Color - from dimos.msgs.sensor_msgs.Image import Image - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + from dimos.msgs.image import image_sharpness from dimos.perception.fiducial.marker_transformer import DetectMarkers - from dimos.robot.unitree.go2.connection import _camera_info_static + from dimos.robot.unitree.go2.camera_calibration import front_camera_calibration + from dimos.visualization.rerun.message_helpers import camera_pinhole, image_archetype db_path = resolve_dataset(dataset) store = open_store(db_path) if out is None: out = Path.cwd() / f"{db_path.stem}.rrd" - cam_info = _camera_info_static() + cam_info = front_camera_calibration() rr.init("dimos markers", recording_id=db_path.stem) rr.save(str(out)) # Static pinhole on the camera entity; per-frame Transform3D goes on the # same entity. Image is the child so it projects through the pinhole. - pinhole = cam_info.to_rerun() + pinhole = camera_pinhole(cam_info) assert not isinstance(pinhole, list) rr.log("world/camera", pinhole, static=True) @@ -116,7 +118,7 @@ def main( translation=[x, y, z], quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw]) ), ) - rr.log("world/camera/image", img_obs.data.to_rerun()) + rr.log("world/camera/image", image_archetype(img_obs.data)) # Clear any prior detection box at this ts; pass 3 will overwrite # this with the actual detection iff one fires for this frame. rr.log( @@ -129,7 +131,7 @@ def main( # Pass 3: marker detections (filtered same way as `dimos map`). xf = DetectMarkers(camera_info=cam_info, marker_length_m=marker_size) pipeline: Stream[Any] = color_image.transform( - QualityWindow(lambda img: img.sharpness, window=quality_window) + QualityWindow(image_sharpness, window=quality_window) ) if marker_max_speed > 0: pipeline = pipeline.transform( @@ -195,7 +197,7 @@ def main( smoothing_window=smoothing_window, ) pipeline_tracked: Stream[Any] = color_image.transform( - QualityWindow(lambda img: img.sharpness, window=quality_window) + QualityWindow(image_sharpness, window=quality_window) ) if marker_max_speed > 0: pipeline_tracked = pipeline_tracked.transform( diff --git a/dimos/mapping/cli/test_cli.py b/dimos/mapping/cli/test_cli.py index 53674915c6d..12fe2a9ae95 100644 --- a/dimos/mapping/cli/test_cli.py +++ b/dimos/mapping/cli/test_cli.py @@ -151,8 +151,9 @@ def test_pose_fill_snippet(dataset: str, tmp_path: Path) -> None: assert res.returncode == 0, res.stderr assert out.exists() + from dimos_generated.sensor_msgs.msg import PointCloud2 + from dimos.memory.store.sqlite import SqliteStore - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 store = SqliteStore(path=str(out), must_exist=True) try: @@ -223,21 +224,3 @@ def test_global_snippet(dataset: str, tmp_path: Path) -> None: ) assert res.returncode == 0, res.stderr assert out.exists() and out.stat().st_size > 0 - - -def test_view_pc2(tmp_path: Path) -> None: - import numpy as np - - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 - - pc2 = tmp_path / "map.pc2.lcm" - cloud = PointCloud2.from_numpy( - np.random.default_rng(0).random((256, 3)), frame_id="world", timestamp=0.0 - ) - pc2.write_bytes(cloud.lcm_encode()) - - out = tmp_path / "view.rrd" - res = _run("view", str(pc2), "--out", str(out)) - assert res.returncode == 0, res.stderr - assert "256 points" in res.stdout - assert out.exists() and out.stat().st_size > 0 diff --git a/dimos/mapping/cli/test_pgo_accumulate.py b/dimos/mapping/cli/test_pgo_accumulate.py new file mode 100644 index 00000000000..0c25f3e57bf --- /dev/null +++ b/dimos/mapping/cli/test_pgo_accumulate.py @@ -0,0 +1,48 @@ +# 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. + +from dimos_generated.geometry_msgs.msg import Transform, TransformStamped, Vector3 +from dimos_generated.std_msgs.msg import Header +import numpy as np + +from dimos.mapping.cli.map import _accumulate +from dimos.mapping.loop_closure.pgo import PoseGraph +from dimos.memory.store.memory import MemoryStore +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.msgs.time import time_from_nanoseconds + + +def test_pgo_accumulate_applies_generated_correction_before_mapping(mocker): + graph = mocker.Mock(spec=PoseGraph) + graph.correction_at.return_value = TransformStamped( + header=Header(frame_id="world_corrected"), + child_frame_id="world", + transform=Transform(translation=Vector3(x=2.0)), + ) + mapper = mocker.Mock(side_effect=lambda observations: observations) + mocker.patch("dimos.mapping.voxels.module.VoxelMapTransformer", return_value=mapper) + stamp = time_from_nanoseconds(42_123_456_789) + cloud = pointcloud_from_xyz( + np.array([[1.0, 2.0, 3.0]]), header=Header(stamp=stamp, frame_id="world") + ) + with MemoryStore() as store: + stream = store.stream("cloud", type(cloud)) + stream.append(cloud, ts=42.0, pose=(0.0, 0.0, 0.0)) + result = _accumulate(stream, voxel=0.1, block_count=10, device="CPU:0", graph=graph) + assert result is not None + np.testing.assert_array_equal(pointcloud_xyz(result), [[3.0, 2.0, 3.0]]) + assert result.header.frame_id == "world_corrected" + assert result.header.stamp.sec == 42 + assert result.header.stamp.nanosec == 123_456_789 + graph.correction_at.assert_called_once_with(42.0) diff --git a/dimos/mapping/cli/test_pose_fill.py b/dimos/mapping/cli/test_pose_fill.py index 2953d9e13e6..4db1cab4c57 100644 --- a/dimos/mapping/cli/test_pose_fill.py +++ b/dimos/mapping/cli/test_pose_fill.py @@ -20,20 +20,29 @@ from __future__ import annotations +import math + +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +import pytest + from dimos.mapping.cli.pose_fill import pose_fill from dimos.memory.store.memory import MemoryStore -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 def test_pose_fill_attaches_nearest_pose() -> None: """Each target obs gets the nearest pose-source pose; payload is preserved.""" with MemoryStore() as store: target = store.stream("image", str) - poses = store.stream("odom", Transform) + poses = store.stream("odom", Pose) target.append("img0", ts=0.0) - poses.append(Transform(translation=Vector3(1.0, 0.0, 0.0)), ts=0.001) + poses.append(Pose(position=Point(x=1.0), orientation=Quaternion(w=1)), ts=0.001) out = pose_fill(target, poses, tolerance=0.05).to_list() @@ -47,11 +56,13 @@ def test_pose_fill_mount_composes_static_child_transform() -> None: """`mount` composes ``world_base + mount`` onto each attached pose.""" with MemoryStore() as store: target = store.stream("image", str) - poses = store.stream("odom", Transform) + poses = store.stream("odom", Pose) target.append("img0", ts=0.0) # Base pose 1m forward in x; mount offsets 1m in y (identity rotations). - poses.append(Transform(translation=Vector3(1.0, 0.0, 0.0)), ts=0.001) - mount = Transform(translation=Vector3(0.0, 1.0, 0.0), rotation=Quaternion()) + poses.append(Pose(position=Point(x=1.0), orientation=Quaternion(w=1)), ts=0.001) + mount = TransformStamped( + transform=Transform(translation=Vector3(y=1), rotation=Quaternion(w=1)) + ) out = pose_fill(target, poses, tolerance=0.05, mount=mount).to_list() @@ -59,3 +70,35 @@ def test_pose_fill_mount_composes_static_child_transform() -> None: assert out[0].pose_tuple is not None # Identity rotations → composed translation is the component-wise sum. assert out[0].pose_tuple[:3] == (1.0, 1.0, 0.0) + + +def test_pose_fill_rotates_mount_offset_in_base_axes() -> None: + """A base yaw rotates the local camera offset before the world translation.""" + with MemoryStore() as store: + target = store.stream("image", str) + poses = store.stream("odom", Pose) + target.append("camera frame", ts=5.0) + poses.append( + Pose( + position=Point(x=2), + orientation=Quaternion(z=math.sqrt(0.5), w=math.sqrt(0.5)), + ), + ts=5.001, + ) + mount = TransformStamped( + transform=Transform(translation=Vector3(x=1), rotation=Quaternion(w=1)) + ) + result = pose_fill(target, poses, tolerance=0.05, mount=mount).to_list() + assert len(result) == 1 + assert result[0].data == "camera frame" + assert result[0].ts == 5.0 + assert result[0].pose_tuple == pytest.approx((2, 1, 0, 0, 0, math.sqrt(0.5), math.sqrt(0.5))) + + +def test_pose_fill_drops_unmatched_frame_without_loading_payload() -> None: + with MemoryStore() as store: + target = store.stream("image", str) + poses = store.stream("odom", Pose) + target.append("unmatched", ts=1.0) + poses.append(Pose(position=Point(x=1), orientation=Quaternion(w=1)), ts=3.0) + assert pose_fill(target, poses, tolerance=0.05).to_list() == [] diff --git a/dimos/mapping/cli/test_view.py b/dimos/mapping/cli/test_view.py new file mode 100644 index 00000000000..d7ddb137033 --- /dev/null +++ b/dimos/mapping/cli/test_view.py @@ -0,0 +1,40 @@ +# 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. + +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path + +from dimos_generated.std_msgs.msg import Header +import numpy as np +import rerun as rr + +from dimos.mapping.cli.view import main +from dimos.msgs.pointcloud import pointcloud_from_xyz + + +def test_view_cdr_cloud_writes_offline_recording(tmp_path: Path, capsys): + pc2 = tmp_path / "map.pc2.cdr" + cloud = pointcloud_from_xyz( + np.random.default_rng(0).random((256, 3)), header=Header(frame_id="world") + ) + pc2.write_bytes(cloud.encode()) + out = tmp_path / "view.rrd" + try: + main(pc2, voxel=0.05, bottom_cutoff=None, out=out) + finally: + rr.disconnect() + assert "256 points in 'world'" in capsys.readouterr().out + assert out.exists() and out.stat().st_size > 0 diff --git a/dimos/mapping/cli/view.py b/dimos/mapping/cli/view.py index a76eccc173d..60c330ad25d 100644 --- a/dimos/mapping/cli/view.py +++ b/dimos/mapping/cli/view.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""`dimos map view`: open a `.pc2.lcm` cloud - `dimos map global --export` - in rerun.""" +"""`dimos map view`: open a `.pc2.cdr` cloud - `dimos map global --export` - in rerun.""" from __future__ import annotations @@ -22,7 +22,7 @@ def main( - path: Path = typer.Argument(..., help="Point cloud to view (.pc2.lcm)"), + path: Path = typer.Argument(..., help="Point cloud to view (.pc2.cdr)"), voxel: float = typer.Option(0.05, "--voxel", help="Rendered point size (m)"), bottom_cutoff: float | None = typer.Option( None, @@ -33,14 +33,15 @@ def main( None, "--out", help="Write a .rrd instead of opening the viewer" ), ) -> None: - """Load an LCM-encoded PointCloud2 and display it in rerun.""" + """Load a CDR-encoded PointCloud2 and display it in rerun.""" + from dimos_generated.sensor_msgs.msg import PointCloud2 import rerun as rr - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.visualization.rerun.init import rerun_init + from dimos.visualization.rerun.message_helpers import cloud_archetype - cloud = PointCloud2.lcm_decode(path.read_bytes()) - print(f"{path}: {len(cloud.pointcloud.points)} points in {cloud.frame_id!r}") + cloud = PointCloud2.decode(path.read_bytes()) + print(f"{path}: {cloud.width * cloud.height} points in {cloud.header.frame_id!r}") rerun_init("dimos map view") if out is not None: @@ -49,7 +50,7 @@ def main( rr.spawn() rr.log( f"world/{path.name.split('.')[0]}/pointcloud", - cloud.to_rerun(voxel_size=voxel / 2, bottom_cutoff=bottom_cutoff), + cloud_archetype(cloud, ui_radius=voxel / 2, bottom_cutoff=bottom_cutoff), static=True, ) if out is not None: diff --git a/dimos/mapping/costmapper.py b/dimos/mapping/costmapper.py index 31b94b164f5..42043063b7c 100644 --- a/dimos/mapping/costmapper.py +++ b/dimos/mapping/costmapper.py @@ -15,6 +15,8 @@ from dataclasses import asdict import time +from dimos_generated.nav_msgs.msg import OccupancyGrid +from dimos_generated.sensor_msgs.msg import PointCloud2 import numpy as np from pydantic import Field from reactivex import combine_latest, operators as ops @@ -27,8 +29,8 @@ HeightCostConfig, OccupancyConfig, ) -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.geometry import pose_matrix +from dimos.msgs.occupancy import occupancy_view from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -95,20 +97,24 @@ def _calculate_costmap(self, msg: PointCloud2) -> OccupancyGrid: def _apply_initial_safe_radius(self, grid: OccupancyGrid) -> None: radius_meters = self.config.initial_safe_radius_meters - if radius_meters <= 0 or grid.grid.size == 0: + if radius_meters <= 0: + return + cells = occupancy_view(grid).copy() + if cells.size == 0: return - resolution = grid.resolution - origin_x = grid.origin.position.x - origin_y = grid.origin.position.y - - rows, columns = np.ogrid[: grid.grid.shape[0], : grid.grid.shape[1]] - cell_world_x = columns * resolution + origin_x - cell_world_y = rows * resolution + origin_y + resolution = grid.info.resolution + matrix = pose_matrix(grid.info.origin) + rows, columns = np.ogrid[: cells.shape[0], : cells.shape[1]] + local_x = columns * resolution + local_y = rows * resolution + cell_world_x = matrix[0, 0] * local_x + matrix[0, 1] * local_y + matrix[0, 3] + cell_world_y = matrix[1, 0] * local_x + matrix[1, 1] * local_y + matrix[1, 3] distance_squared_meters = cell_world_x**2 + cell_world_y**2 # Half-cell tolerance: a cell counts as inside if any part of it overlaps # the disc. Avoids floating-point boundary flakiness from radius/resolution. effective_radius_meters = radius_meters + resolution * 0.5 safe_mask = distance_squared_meters <= effective_radius_meters**2 - grid.grid[safe_mask] = 0 + cells[safe_mask] = 0 + grid.data = cells.ravel() diff --git a/dimos/mapping/dim_slam/dim_slam.py b/dimos/mapping/dim_slam/dim_slam.py index 59fa76996ac..920d8b0f79e 100644 --- a/dimos/mapping/dim_slam/dim_slam.py +++ b/dimos/mapping/dim_slam/dim_slam.py @@ -19,16 +19,13 @@ from typing import Literal +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, Imu, PointCloud2 +from dimos_generated.tf2_msgs.msg import TFMessage from pydantic import BaseModel, Field, model_validator from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.stream import IO, In, Out -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.Imu import Imu -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.utils.logging_config import setup_logger from dimos.utils.nvidia_env import driver_env, sdk_variant diff --git a/dimos/mapping/dim_slam/rust/Cargo.lock b/dimos/mapping/dim_slam/rust/Cargo.lock index 325616621aa..37dafef1184 100644 --- a/dimos/mapping/dim_slam/rust/Cargo.lock +++ b/dimos/mapping/dim_slam/rust/Cargo.lock @@ -220,6 +220,20 @@ name = "bytemuck" version = "1.25.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a1f896587b6f2c069c73d2f0913e2d590c3990285cd2f0b6aa02b786b4c679c" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.4", +] [[package]] name = "byteorder" @@ -556,8 +570,8 @@ version = "0.1.0" dependencies = [ "cu_vslam_rs", "dim_slam", + "dimos-generated-messages", "dimos-module", - "lcm-msgs", "serde", "tokio", "tracing", @@ -576,9 +590,17 @@ dependencies = [ ] [[package]] -name = "dimos-lcm" +name = "dimos-generated-messages" +version = "0.1.0" +dependencies = [ + "re_cdr", + "serde", + "serde-big-array", +] + +[[package]] +name = "dimos-lcm-transport" version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" dependencies = [ "byteorder", "socket2 0.5.10", @@ -589,9 +611,9 @@ dependencies = [ name = "dimos-module" version = "0.1.0" dependencies = [ - "dimos-lcm", + "dimos-generated-messages", + "dimos-lcm-transport", "dimos-module-macros", - "lcm-msgs", "nalgebra", "rayon", "serde", @@ -1385,14 +1407,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" -[[package]] -name = "lcm-msgs" -version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" -dependencies = [ - "byteorder", -] - [[package]] name = "libc" version = "0.2.189" @@ -2167,6 +2181,18 @@ dependencies = [ "yasna", ] +[[package]] +name = "re_cdr" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bba10c21433b597481f366b978cdf2617e56785b38cb0dcc3833c6b8ad7e061a" +dependencies = [ + "bytemuck", + "byteorder", + "serde", + "thiserror", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -2523,6 +2549,15 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "serde-big-array" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11fc7cc2c76d73e0f27ee52abbd64eec84d46f370c88371120433196934e4b7f" +dependencies = [ + "serde", +] + [[package]] name = "serde_core" version = "1.0.229" diff --git a/dimos/mapping/dim_slam/rust/Cargo.toml b/dimos/mapping/dim_slam/rust/Cargo.toml index 5af8d516d21..2807007898c 100644 --- a/dimos/mapping/dim_slam/rust/Cargo.toml +++ b/dimos/mapping/dim_slam/rust/Cargo.toml @@ -16,7 +16,7 @@ dim_slam = { git = "https://github.com/dimensionalOS/dimSLAM", tag = "v0.4.1" } # and build.rs needs it for the rpath. cu_vslam_rs = "0.2.0" dimos-module = { path = "../../../../native/rust/dimos-module" } -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +dimos-generated-messages = { path = "../../../message_codegen" } tokio = { version = "1", features = ["rt-multi-thread", "macros"] } tracing = "0.1" # native_config expands to ::serde and ::validator derives, so both must be direct deps. diff --git a/dimos/mapping/dim_slam/rust/flake.nix b/dimos/mapping/dim_slam/rust/flake.nix index 7b44713e85d..fe2288ee502 100644 --- a/dimos/mapping/dim_slam/rust/flake.nix +++ b/dimos/mapping/dim_slam/rust/flake.nix @@ -39,6 +39,9 @@ mkdir -p $out/native/rust cp -r ${dimos-repo}/native/rust/dimos-module $out/native/rust/dimos-module cp -r ${dimos-repo}/native/rust/dimos-module-macros $out/native/rust/dimos-module-macros + cp -r ${dimos-repo}/native/rust/dimos-lcm-transport $out/native/rust/dimos-lcm-transport + mkdir -p $out/dimos/message_codegen + cp -r ${dimos-repo}/dimos/message_codegen/. $out/dimos/message_codegen/ ''; generatedCargoNix = crate2nix.tools.${system}.generatedCargoNix { @@ -52,6 +55,7 @@ inherit pkgs; buildRustCrateForPkgs = cratePkgs: cratePkgs.buildRustCrate.override { defaultCrateOverrides = cratePkgs.defaultCrateOverrides // { + dimos-generated-messages = _: { nativeBuildInputs = [cratePkgs.python3]; }; # cu_vslam_rs's build.rs compiles its shim against this SDK. cu_vslam_rs = _: { CUVSLAM_SDK_DIR = sdkPackage; }; # buildRustCrate names DEP_ vars after the crate, cargo after the diff --git a/dimos/mapping/dim_slam/rust/src/main.rs b/dimos/mapping/dim_slam/rust/src/main.rs index 23b8318f935..aa938075c39 100644 --- a/dimos/mapping/dim_slam/rust/src/main.rs +++ b/dimos/mapping/dim_slam/rust/src/main.rs @@ -17,9 +17,10 @@ mod msg_convert; use dim_slam::{ CameraConfig, DimSlamCore, DimSlamCoreConfig, ImuConfig, InitialStds, SourceConfig, }; +use dimos_generated_messages::nav_msgs::msg::Odometry; +use dimos_generated_messages::sensor_msgs::msg::{CameraInfo, Image, Imu, PointCloud2}; +use dimos_module::cdr; use dimos_module::{native_config, run_with_transport, Input, Module, Output, Tf}; -use lcm_msgs::nav_msgs::Odometry; -use lcm_msgs::sensor_msgs::{CameraInfo, Image, Imu, PointCloud2}; use msg_convert::{ tf_lookup, to_camera_model, to_estimate, to_image_frame, to_imu_sample, to_odometry_msg, to_point_cloud2, to_transform, @@ -114,25 +115,25 @@ struct DimSlamConfig { #[derive(Module)] #[module(setup = init, teardown = report)] struct DimSlam { - #[input(decode = CameraInfo::decode)] + #[input(decode = cdr::decode)] camera_info: Input, - #[input(decode = Image::decode)] + #[input(decode = cdr::decode)] image: Input, /// rgbd only; unconnected otherwise. - #[input(decode = Image::decode)] + #[input(decode = cdr::decode)] depth_image: Input, /// Only needed when depth has to be reprojected onto the rig camera. - #[input(decode = CameraInfo::decode)] + #[input(decode = cdr::decode)] depth_camera_info: Input, - #[input(decode = Imu::decode)] + #[input(decode = cdr::decode)] imu: Input, /// External odometry sources, told apart by the transform each message carries. - #[input(decode = Odometry::decode)] + #[input(decode = cdr::decode)] odom_sources: Input, - #[output(encode = Odometry::encode)] + #[output(encode = cdr::encode)] odometry: Output, /// rgbd only: range-gated depth points, in the depth frame. - #[output(encode = PointCloud2::encode)] + #[output(encode = cdr::encode)] depth_cloud: Output, #[tf] tf: Tf, @@ -199,26 +200,47 @@ impl DimSlam { } async fn handle_camera_info(&mut self, info: CameraInfo) { + let info = match to_camera_model(info) { + Ok(value) => value, + Err(error) => { + tracing::warn!(%error, "invalid info message"); + return; + } + }; let Self { slam, tf, .. } = self; slam.as_mut() .expect("setup ran") - .handle_camera_info(to_camera_model(info), &tf_lookup(tf)); + .handle_camera_info(info, &tf_lookup(tf)); } async fn handle_image(&mut self, img: Image) { + let img = match to_image_frame(img) { + Ok(value) => value, + Err(error) => { + tracing::warn!(%error, "invalid img message"); + return; + } + }; let Self { slam, tf, .. } = self; slam.as_mut() .expect("setup ran") - .handle_image(to_image_frame(img), &tf_lookup(tf)); + .handle_image(img, &tf_lookup(tf)); self.publish().await; } async fn handle_depth_image(&mut self, img: Image) { + let img = match to_image_frame(img) { + Ok(value) => value, + Err(error) => { + tracing::warn!(%error, "invalid img message"); + return; + } + }; let Self { slam, tf, .. } = self; let cloud = slam .as_mut() .expect("setup ran") - .handle_depth_image(to_image_frame(img), &tf_lookup(tf)); + .handle_depth_image(img, &tf_lookup(tf)); if let Some(cloud) = cloud { self.depth_cloud .publish(&to_point_cloud2(&cloud)) @@ -229,10 +251,17 @@ impl DimSlam { } async fn handle_depth_camera_info(&mut self, info: CameraInfo) { + let info = match to_camera_model(info) { + Ok(value) => value, + Err(error) => { + tracing::warn!(%error, "invalid info message"); + return; + } + }; self.slam .as_mut() .expect("setup ran") - .handle_depth_camera_info(to_camera_model(info)); + .handle_depth_camera_info(info); } async fn handle_imu(&mut self, msg: Imu) { diff --git a/dimos/mapping/dim_slam/rust/src/msg_convert.rs b/dimos/mapping/dim_slam/rust/src/msg_convert.rs index e0e7d8289c0..2f4b15883fc 100644 --- a/dimos/mapping/dim_slam/rust/src/msg_convert.rs +++ b/dimos/mapping/dim_slam/rust/src/msg_convert.rs @@ -12,33 +12,33 @@ // See the License for the specific language governing permissions and // limitations under the License. // -// LCM messages in and out of the dim_slam library's plain types. +// Generated ROS2 messages in and out of the dim_slam library's plain types. use dim_slam::nalgebra::{Isometry3, Matrix6, Quaternion, Translation3, UnitQuaternion, Vector3}; use dim_slam::{CameraModel, ImageFrame, ImuSample, OdometryEstimate, PointCloud, Twist}; +use dimos_generated_messages::builtin_interfaces::msg::Time; +use dimos_generated_messages::geometry_msgs::msg as geometry_msgs; +use dimos_generated_messages::nav_msgs::msg::Odometry; +use dimos_generated_messages::sensor_msgs::msg::{CameraInfo, Image, Imu, PointCloud2, PointField}; +use dimos_generated_messages::std_msgs::msg::Header; use dimos_module::{Tf, Transform}; -use lcm_msgs::geometry_msgs; -use lcm_msgs::nav_msgs::Odometry; -use lcm_msgs::sensor_msgs::{CameraInfo, Image, Imu, PointCloud2, PointField}; -use lcm_msgs::std_msgs::{Header, Time}; const NS_PER_SEC: i64 = 1_000_000_000; -const BYTES_PER_POINT: i32 = 12; +const BYTES_PER_POINT: u32 = 12; fn stamp_to_ns(header: &Header) -> i64 { - header.stamp.sec as i64 * NS_PER_SEC + header.stamp.nsec as i64 + header.stamp.sec as i64 * NS_PER_SEC + header.stamp.nanosec as i64 } fn to_stamp(timestamp_ns: i64) -> Time { Time { - sec: (timestamp_ns / NS_PER_SEC) as i32, - nsec: (timestamp_ns % NS_PER_SEC) as i32, + sec: timestamp_ns.div_euclid(NS_PER_SEC) as i32, + nanosec: timestamp_ns.rem_euclid(NS_PER_SEC) as u32, } } fn to_header(timestamp_ns: i64, frame_id: &str) -> Header { Header { - seq: 0, stamp: to_stamp(timestamp_ns), frame_id: frame_id.to_string(), } @@ -74,27 +74,69 @@ pub fn tf_lookup(tf: &Tf) -> impl Fn(&str, &str) -> Option> + '_ } } -pub fn to_image_frame(img: Image) -> ImageFrame { - ImageFrame { +pub fn to_image_frame(mut img: Image) -> Result { + let (scalar_bytes, channels) = match img.encoding.as_str() { + "mono8" => (1usize, 1), + "rgb8" | "bgr8" => (1, 3), + "mono16" | "16UC1" => (2, 1), + "32FC1" => (4, 1), + _ => return Err(format!("unsupported image encoding {}", img.encoding)), + }; + let row_bytes = img.width as usize * scalar_bytes * channels; + let step = img.step as usize; + if step < row_bytes || img.data.len() != step * img.height as usize { + return Err("image dimensions/stride do not match data".into()); + } + if img.is_bigendian != 0 && scalar_bytes > 1 && step > 0 { + for row in img.data.chunks_exact_mut(step) { + for value in row[..row_bytes].chunks_exact_mut(scalar_bytes) { + value.reverse(); + } + } + } + if img.encoding == "bgr8" && step > 0 { + for row in img.data.chunks_exact_mut(step) { + for pixel in row[..row_bytes].as_chunks_mut::<3>().0 { + pixel.swap(0, 2); + } + } + img.encoding = "rgb8".into(); + } + Ok(ImageFrame { timestamp_ns: stamp_to_ns(&img.header), frame_id: img.header.frame_id, - width: img.width, - height: img.height, + width: img + .width + .try_into() + .map_err(|_| "image width exceeds backend range")?, + height: img + .height + .try_into() + .map_err(|_| "image height exceeds backend range")?, encoding: img.encoding, - step: img.step, + step: img + .step + .try_into() + .map_err(|_| "image step exceeds backend range")?, data: img.data, - } + }) } -pub fn to_camera_model(info: CameraInfo) -> CameraModel { - CameraModel { +pub fn to_camera_model(info: CameraInfo) -> Result { + Ok(CameraModel { timestamp_ns: stamp_to_ns(&info.header), frame_id: info.header.frame_id, - width: info.width, - height: info.height, - distortion: info.D, - intrinsics: info.K, - } + width: info + .width + .try_into() + .map_err(|_| "camera width exceeds backend range")?, + height: info + .height + .try_into() + .map_err(|_| "camera height exceeds backend range")?, + distortion: info.d, + intrinsics: info.k, + }) } pub fn to_imu_sample(msg: &Imu) -> ImuSample { @@ -172,17 +214,17 @@ pub fn to_transform(estimate: &OdometryEstimate) -> Transform { ) } -fn xyz_field(name: &str, offset: i32) -> PointField { +fn xyz_field(name: &str, offset: u32) -> PointField { PointField { name: name.to_string(), offset, - datatype: PointField::FLOAT32 as u8, + datatype: PointField::FLOAT32, count: 1, } } pub fn to_point_cloud2(cloud: &PointCloud) -> PointCloud2 { - let width = cloud.points.len() as i32; + let width = cloud.points.len() as u32; let mut data = Vec::with_capacity(cloud.points.len() * BYTES_PER_POINT as usize); for point in &cloud.points { for coordinate in point { @@ -206,6 +248,67 @@ pub fn to_point_cloud2(cloud: &PointCloud) -> PointCloud2 { mod tests { use super::*; + #[test] + fn image_adapter_normalizes_depth_byte_order_and_preserves_row_padding() { + let image = Image { + width: 1, + height: 2, + step: 4, + encoding: "16UC1".into(), + is_bigendian: 1, + data: vec![0x12, 0x34, 99, 98, 0x56, 0x78, 97, 96], + ..Default::default() + }; + let frame = to_image_frame(image).unwrap(); + assert_eq!(frame.data, [0x34, 0x12, 99, 98, 0x78, 0x56, 97, 96]); + assert_eq!((frame.width, frame.height, frame.step), (1, 2, 4)); + } + + #[test] + fn image_adapter_normalizes_bgr_for_the_backend() { + let frame = to_image_frame(Image { + width: 1, + height: 1, + step: 4, + encoding: "bgr8".into(), + data: vec![1, 2, 3, 99], + ..Default::default() + }) + .unwrap(); + assert_eq!(frame.encoding, "rgb8"); + assert_eq!(frame.data, [3, 2, 1, 99]); + } + + #[test] + fn image_adapter_rejects_invalid_layout_and_backend_dimension_overflow() { + assert!(to_image_frame(Image { + width: 1, + height: 1, + step: 1, + encoding: "rgb8".into(), + data: vec![1], + ..Default::default() + }) + .is_err()); + assert!(to_camera_model(CameraInfo { + width: u32::MAX, + ..Default::default() + }) + .is_err()); + assert!(to_image_frame(Image { + encoding: "16UC1".into(), + is_bigendian: 1, + ..Default::default() + }) + .is_ok()); + } + + #[test] + fn negative_source_stamp_is_normalized() { + let stamp = to_stamp(-250000000); + assert_eq!((stamp.sec, stamp.nanosec), (-1, 750000000)); + } + #[test] fn stamp_round_trips_through_ns() { let header = Header { diff --git a/dimos/mapping/occupancy/conftest.py b/dimos/mapping/occupancy/conftest.py index f20dc1310b1..0847c3199cd 100644 --- a/dimos/mapping/occupancy/conftest.py +++ b/dimos/mapping/occupancy/conftest.py @@ -12,17 +12,27 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.geometry_msgs.msg import Pose, Quaternion +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid import numpy as np import pytest from dimos.mapping.occupancy.gradient import gradient -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid from dimos.utils.data import get_data @pytest.fixture def occupancy() -> OccupancyGrid: - return OccupancyGrid(np.load(get_data("occupancy_simple.npy"))) + cells = np.load(get_data("occupancy_simple.npy")) + return OccupancyGrid( + info=MapMetaData( + width=cells.shape[1], + height=cells.shape[0], + resolution=0.05, + origin=Pose(orientation=Quaternion(w=1)), + ), + data=cells.astype(np.int8).ravel(), + ) @pytest.fixture diff --git a/dimos/mapping/occupancy/extrude_occupancy.py b/dimos/mapping/occupancy/extrude_occupancy.py index 799319cbf6d..bf6f8d59954 100644 --- a/dimos/mapping/occupancy/extrude_occupancy.py +++ b/dimos/mapping/occupancy/extrude_occupancy.py @@ -13,10 +13,11 @@ # limitations under the License. +from dimos_generated.nav_msgs.msg import OccupancyGrid import numpy as np from numpy.typing import NDArray -from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid +from dimos.msgs.occupancy import occupancy_view # Rectangle type: (x, y, width, height) Rect = tuple[int, int, int, int] @@ -36,12 +37,10 @@ def identify_convex_shapes(occupancy_grid: OccupancyGrid) -> list[Rect]: Returns: List of rectangles as (x, y, width, height) tuples in grid coords. """ - grid = occupancy_grid.grid + grid = occupancy_view(occupancy_grid) # Create binary mask of occupied cells (treat UNKNOWN as OCCUPIED) - occupied_mask = ((grid == CostValues.OCCUPIED) | (grid == CostValues.UNKNOWN)).astype( - np.uint8 - ) * 255 + occupied_mask = ((grid == 100) | (grid == -1)).astype(np.uint8) * 255 return _decompose_to_rectangles(occupied_mask) @@ -184,9 +183,9 @@ def generate_mujoco_scene( # Get rectangles from the occupancy grid rectangles = identify_convex_shapes(occupancy_grid) - resolution = occupancy_grid.resolution - origin_x = occupancy_grid.origin.position.x - origin_y = occupancy_grid.origin.position.y + resolution = occupancy_grid.info.resolution + origin_x = occupancy_grid.info.origin.position.x + origin_y = occupancy_grid.info.origin.position.y # Build XML xml_lines = [ @@ -208,7 +207,7 @@ def generate_mujoco_scene( # Add each rectangle as a box geom for i, (gx, gy, gw, gh) in enumerate(rectangles): # Convert grid coordinates to world coordinates - # Grid origin is top-left, world origin is at occupancy_grid.origin + # Grid origin is top-left, world origin is at occupancy_grid.info.origin # gx, gy are in grid cells, need to convert to meters world_x = origin_x + (gx + gw / 2) * resolution world_y = origin_y + (gy + gh / 2) * resolution diff --git a/dimos/mapping/occupancy/gradient.py b/dimos/mapping/occupancy/gradient.py index 802b52a9faf..0478058b1c5 100644 --- a/dimos/mapping/occupancy/gradient.py +++ b/dimos/mapping/occupancy/gradient.py @@ -14,10 +14,11 @@ from typing import TYPE_CHECKING, Literal, TypeAlias, cast +from dimos_generated.nav_msgs.msg import OccupancyGrid import numpy as np from scipy import ndimage -from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid +from dimos.msgs.occupancy import occupancy_view if TYPE_CHECKING: from numpy.typing import NDArray @@ -49,19 +50,19 @@ def gradient( """ # Remember which cells are unknown - unknown_mask = occupancy_grid.grid == CostValues.UNKNOWN + unknown_mask = occupancy_view(occupancy_grid) == -1 # Create binary obstacle map # Consider cells >= threshold as obstacles (1), everything else as free (0) # Unknown cells are not considered obstacles for distance calculation - obstacle_map = (occupancy_grid.grid >= obstacle_threshold).astype(np.float32) + obstacle_map = (occupancy_view(occupancy_grid) >= obstacle_threshold).astype(np.float32) # Compute distance transform (distance to nearest obstacle in cells) # Unknown cells are treated as if they don't exist for distance calculation distance_cells = cast("NDArray[np.float64]", ndimage.distance_transform_edt(1 - obstacle_map)) # Convert to meters and clip to max distance - distance_meters = np.clip(distance_cells * occupancy_grid.resolution, 0, max_distance) + distance_meters = np.clip(distance_cells * occupancy_grid.info.resolution, 0, max_distance) # Invert and scale to 0-100 range # Far from obstacles (max_distance) -> 0 @@ -69,21 +70,17 @@ def gradient( gradient_values = (1 - distance_meters / max_distance) * 100 # Ensure obstacles are exactly 100 - gradient_values[obstacle_map > 0] = CostValues.OCCUPIED + gradient_values[obstacle_map > 0] = 100 # Convert to int8 for OccupancyGrid gradient_data = gradient_values.astype(np.int8) # Preserve unknown cells as unknown (don't apply gradient to them) - gradient_data[unknown_mask] = CostValues.UNKNOWN + gradient_data[unknown_mask] = -1 # Create new OccupancyGrid with gradient gradient_grid = OccupancyGrid( - grid=gradient_data, - resolution=occupancy_grid.resolution, - origin=occupancy_grid.origin, - frame_id=occupancy_grid.frame_id, - ts=occupancy_grid.ts, + header=occupancy_grid.header, info=occupancy_grid.info, data=gradient_data.ravel() ) return gradient_grid @@ -122,22 +119,18 @@ def voronoi_gradient( - 100: At obstacles """ # Remember which cells are unknown - unknown_mask = occupancy_grid.grid == CostValues.UNKNOWN + unknown_mask = occupancy_view(occupancy_grid) == -1 # Create binary obstacle map - obstacle_map = (occupancy_grid.grid >= obstacle_threshold).astype(np.float32) + obstacle_map = (occupancy_view(occupancy_grid) >= obstacle_threshold).astype(np.float32) # Check if there are any obstacles if not np.any(obstacle_map): # No obstacles - everything is free - gradient_data = np.zeros_like(occupancy_grid.grid, dtype=np.int8) - gradient_data[unknown_mask] = CostValues.UNKNOWN + gradient_data = np.zeros_like(occupancy_view(occupancy_grid), dtype=np.int8) + gradient_data[unknown_mask] = -1 return OccupancyGrid( - grid=gradient_data, - resolution=occupancy_grid.resolution, - origin=occupancy_grid.origin, - frame_id=occupancy_grid.frame_id, - ts=occupancy_grid.ts, + header=occupancy_grid.header, info=occupancy_grid.info, data=gradient_data.ravel() ) # Label connected obstacle regions (clusters) @@ -189,22 +182,18 @@ def voronoi_gradient( gradient_values = cost_ratio * 99 # Ensure obstacles are exactly 100 - gradient_values[obstacle_map > 0] = CostValues.OCCUPIED + gradient_values[obstacle_map > 0] = 100 # Apply max_distance clipping - cells beyond max_distance from obstacles get cost 0 - max_distance_cells = max_distance / occupancy_grid.resolution + max_distance_cells = max_distance / occupancy_grid.info.resolution gradient_values[distance_cells > max_distance_cells] = 0 # Convert to int8 gradient_data = gradient_values.astype(np.int8) # Preserve unknown cells - gradient_data[unknown_mask] = CostValues.UNKNOWN + gradient_data[unknown_mask] = -1 return OccupancyGrid( - grid=gradient_data, - resolution=occupancy_grid.resolution, - origin=occupancy_grid.origin, - frame_id=occupancy_grid.frame_id, - ts=occupancy_grid.ts, + header=occupancy_grid.header, info=occupancy_grid.info, data=gradient_data.ravel() ) diff --git a/dimos/mapping/occupancy/inflation.py b/dimos/mapping/occupancy/inflation.py index 610157ddf62..f7264039912 100644 --- a/dimos/mapping/occupancy/inflation.py +++ b/dimos/mapping/occupancy/inflation.py @@ -12,10 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +from copy import deepcopy +import math + +from dimos_generated.nav_msgs.msg import OccupancyGrid import numpy as np from scipy import ndimage -from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid +from dimos.msgs.occupancy import occupancy_extent, occupancy_view def simple_inflate(occupancy_grid: OccupancyGrid, radius: float) -> OccupancyGrid: @@ -25,29 +29,28 @@ def simple_inflate(occupancy_grid: OccupancyGrid, radius: float) -> OccupancyGri Returns: New OccupancyGrid with inflated obstacles """ + occupancy_extent(occupancy_grid) + if not math.isfinite(radius) or radius < 0: + raise ValueError("inflation radius must be finite and nonnegative") + # Convert radius to grid cells - cell_radius = int(np.ceil(radius / occupancy_grid.resolution)) + cell_radius = int(np.ceil(radius / occupancy_grid.info.resolution)) # Get grid as numpy array - grid_array = occupancy_grid.grid + grid_array = occupancy_view(occupancy_grid) # Create circular kernel for binary inflation y, x = np.ogrid[-cell_radius : cell_radius + 1, -cell_radius : cell_radius + 1] kernel = (x**2 + y**2 <= cell_radius**2).astype(np.uint8) # Find occupied cells - occupied_mask = grid_array >= CostValues.OCCUPIED + occupied_mask = grid_array >= 100 # Binary inflation inflated = ndimage.binary_dilation(occupied_mask, structure=kernel) result_grid = grid_array.copy() - result_grid[inflated] = CostValues.OCCUPIED - - # Create new OccupancyGrid with inflated data using numpy constructor - return OccupancyGrid( - grid=result_grid, - resolution=occupancy_grid.resolution, - origin=occupancy_grid.origin, - frame_id=occupancy_grid.frame_id, - ts=occupancy_grid.ts, - ) + result_grid[inflated] = 100 + + result = deepcopy(occupancy_grid) + result.data = result_grid.ravel() + return result diff --git a/dimos/mapping/occupancy/operations.py b/dimos/mapping/occupancy/operations.py index 8609ac634b3..3771bcae436 100644 --- a/dimos/mapping/occupancy/operations.py +++ b/dimos/mapping/occupancy/operations.py @@ -12,10 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.nav_msgs.msg import OccupancyGrid import numpy as np from scipy import ndimage -from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid +from dimos.msgs.occupancy import occupancy_view def smooth_occupied( @@ -33,8 +34,8 @@ def smooth_occupied( Returns: New OccupancyGrid with smoothed occupied zones """ - grid_array = occupancy_grid.grid - occupied_mask = grid_array >= CostValues.OCCUPIED + grid_array = occupancy_view(occupancy_grid) + occupied_mask = grid_array >= 100 # Count occupied neighbors for each cell (8-connectivity). kernel = np.array([[1, 1, 1], [1, 0, 1], [1, 1, 1]], dtype=np.uint8) @@ -47,14 +48,10 @@ def smooth_occupied( unsupported = occupied_mask & (neighbor_count < min_neighbors) result_grid = grid_array.copy() - result_grid[unsupported] = CostValues.FREE + result_grid[unsupported] = 0 return OccupancyGrid( - grid=result_grid, - resolution=occupancy_grid.resolution, - origin=occupancy_grid.origin, - frame_id=occupancy_grid.frame_id, - ts=occupancy_grid.ts, + header=occupancy_grid.header, info=occupancy_grid.info, data=result_grid.ravel() ) @@ -70,19 +67,13 @@ def overlay_occupied(base: OccupancyGrid, overlay: OccupancyGrid) -> OccupancyGr Returns: New OccupancyGrid with combined occupied zones """ - if base.grid.shape != overlay.grid.shape: + if occupancy_view(base).shape != occupancy_view(overlay).shape: raise ValueError( - f"Grid shapes must match: base {base.grid.shape} vs overlay {overlay.grid.shape}" + f"Grid shapes must match: base {occupancy_view(base).shape} vs overlay {occupancy_view(overlay).shape}" ) - result_grid = base.grid.copy() - overlay_occupied_mask = overlay.grid >= CostValues.OCCUPIED - result_grid[overlay_occupied_mask] = CostValues.OCCUPIED + result_grid = occupancy_view(base).copy() + overlay_occupied_mask = occupancy_view(overlay) >= 100 + result_grid[overlay_occupied_mask] = 100 - return OccupancyGrid( - grid=result_grid, - resolution=base.resolution, - origin=base.origin, - frame_id=base.frame_id, - ts=base.ts, - ) + return OccupancyGrid(header=base.header, info=base.info, data=result_grid.ravel()) diff --git a/dimos/mapping/occupancy/path_map.py b/dimos/mapping/occupancy/path_map.py index 69a1f937385..6f5a241c213 100644 --- a/dimos/mapping/occupancy/path_map.py +++ b/dimos/mapping/occupancy/path_map.py @@ -12,11 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.nav_msgs.msg import OccupancyGrid + from dimos.mapping.occupancy.gradient import GradientStrategy, gradient, voronoi_gradient from dimos.mapping.occupancy.inflation import simple_inflate from dimos.mapping.occupancy.operations import overlay_occupied, smooth_occupied from dimos.mapping.occupancy.types import NavigationStrategy -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid def make_navigation_map( diff --git a/dimos/mapping/occupancy/path_mask.py b/dimos/mapping/occupancy/path_mask.py index b71372bbf62..3fb8a4bb4e0 100644 --- a/dimos/mapping/occupancy/path_mask.py +++ b/dimos/mapping/occupancy/path_mask.py @@ -12,11 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + +from dimos_generated.nav_msgs.msg import OccupancyGrid, Path import numpy as np from numpy.typing import NDArray -from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.occupancy import occupancy_view, world_to_grid def make_path_mask( @@ -44,9 +46,9 @@ def make_path_mask( """ import cv2 - mask = np.zeros((occupancy_grid.height, occupancy_grid.width), dtype=np.uint8) + mask = np.zeros((occupancy_grid.info.height, occupancy_grid.info.width), dtype=np.uint8) - line_width_pixels = max(1, int(robot_width / occupancy_grid.resolution)) + line_width_pixels = max(1, int(robot_width / occupancy_grid.info.resolution)) poses = path.poses if len(poses) < pose_index + 2: @@ -55,8 +57,8 @@ def make_path_mask( # Draw lines between consecutive points cumulative_length = 0.0 for i in range(pose_index, len(poses) - 1): - pos1 = poses[i].position - pos2 = poses[i + 1].position + pos1 = poses[i].pose.position + pos2 = poses[i + 1].pose.position segment_length = np.sqrt( (pos2.x - pos1.x) ** 2 + (pos2.y - pos1.y) ** 2 + (pos2.z - pos1.z) ** 2 @@ -67,22 +69,22 @@ def make_path_mask( cumulative_length += segment_length - grid_pt1 = occupancy_grid.world_to_grid(pos1) - grid_pt2 = occupancy_grid.world_to_grid(pos2) + grid_pt1 = world_to_grid(occupancy_grid, pos1) + grid_pt2 = world_to_grid(occupancy_grid, pos2) - pt1 = (round(grid_pt1.x), round(grid_pt1.y)) - pt2 = (round(grid_pt2.x), round(grid_pt2.y)) + pt1 = (round(grid_pt1[0]), round(grid_pt1[1])) + pt2 = (round(grid_pt2[0]), round(grid_pt2[1])) cv2.line(mask, pt1, pt2, (255.0,), thickness=line_width_pixels) - bool_mask = mask.astype(np.bool_) + bool_mask: NDArray[np.bool_] = mask.astype(np.bool_) total_points = np.sum(bool_mask) if total_points == 0: return bool_mask - occupied_mask = occupancy_grid.grid >= CostValues.OCCUPIED + occupied_mask = occupancy_view(occupancy_grid) >= 100 occupied_in_path = bool_mask & occupied_mask occupied_count = np.sum(occupied_in_path) diff --git a/dimos/mapping/occupancy/path_resampling.py b/dimos/mapping/occupancy/path_resampling.py index 1f50fad7cdf..277600e4420 100644 --- a/dimos/mapping/occupancy/path_resampling.py +++ b/dimos/mapping/occupancy/path_resampling.py @@ -15,16 +15,13 @@ import math +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.nav_msgs.msg import Path import numpy as np from scipy.ndimage import uniform_filter1d -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import quaternion_from_euler from dimos.utils.logging_config import setup_logger -from dimos.utils.transform_utils import euler_to_quaternion logger = setup_logger() @@ -48,27 +45,30 @@ def _add_orientations_to_path(path: Path, goal_orientation: Quaternion) -> None: next_pose = path.poses[i + 1] # Calculate direction to next point - dx = next_pose.position.x - current_pose.position.x - dy = next_pose.position.y - current_pose.position.y + dx = next_pose.pose.position.x - current_pose.pose.position.x + dy = next_pose.pose.position.y - current_pose.pose.position.y # Calculate yaw angle yaw = math.atan2(dy, dx) # Convert to quaternion (roll=0, pitch=0, yaw) - orientation = euler_to_quaternion(Vector3(0, 0, yaw)) - current_pose.orientation = orientation + orientation = quaternion_from_euler(0, 0, yaw) + current_pose.pose.orientation = orientation + path.poses[i] = current_pose # Set last pose orientation - identity_quat = Quaternion(0, 0, 0, 1) + identity_quat = Quaternion(w=1) + last_pose = path.poses[-1] if goal_orientation != identity_quat: # Use the provided goal orientation if it's not the identity - path.poses[-1].orientation = goal_orientation + last_pose.pose.orientation = goal_orientation elif len(path.poses) > 1: # Use the previous pose's orientation - path.poses[-1].orientation = path.poses[-2].orientation + last_pose.pose.orientation = path.poses[-2].pose.orientation else: # Single pose with identity goal orientation - path.poses[-1].orientation = identity_quat + last_pose.pose.orientation = identity_quat + path.poses[-1] = last_pose # TODO: replace goal_pose with just goal_orientation @@ -82,7 +82,7 @@ def simple_resample_path(path: Path, goal_pose: Pose, spacing: float) -> Path: Returns: A new Path with resampled poses """ - if len(path) < 2 or spacing <= 0: + if len(path.poses) < 2 or spacing <= 0: return path resampled = [] @@ -95,8 +95,8 @@ def simple_resample_path(path: Path, goal_pose: Pose, spacing: float) -> Path: prev = path.poses[i - 1] # Calculate segment distance - dx = current.x - prev.x - dy = current.y - prev.y + dx = current.pose.position.x - prev.pose.position.x + dy = current.pose.position.y - prev.pose.position.y segment_length = (dx**2 + dy**2) ** 0.5 if segment_length < 1e-10: @@ -114,12 +114,11 @@ def simple_resample_path(path: Path, goal_pose: Pose, spacing: float) -> Path: break # Create new pose - new_x = prev.x + dir_x * dist_along - new_y = prev.y + dir_y * dist_along + new_x = prev.pose.position.x + dir_x * dist_along + new_y = prev.pose.position.y + dir_y * dist_along new_pose = PoseStamped( - frame_id=path.frame_id, - position=[new_x, new_y, 0.0], - orientation=prev.orientation, # Keep same orientation + header=path.header, + pose=Pose(position=Point(x=new_x, y=new_y), orientation=prev.pose.orientation), ) resampled.append(new_pose) @@ -133,10 +132,18 @@ def simple_resample_path(path: Path, goal_pose: Pose, spacing: float) -> Path: # Add last pose if not already there if len(path.poses) > 1: last = path.poses[-1] - if not resampled or (resampled[-1].x != last.x or resampled[-1].y != last.y): + if ( + math.hypot( + resampled[-1].pose.position.x - last.pose.position.x, + resampled[-1].pose.position.y - last.pose.position.y, + ) + <= 1e-10 + ): + resampled[-1] = last + else: resampled.append(last) - ret = Path(frame_id=path.frame_id, poses=resampled) + ret = Path(header=path.header, poses=resampled) _add_orientations_to_path(ret, goal_pose.orientation) @@ -165,21 +172,20 @@ def smooth_resample_path( """ if len(path.poses) == 1: - p = path.poses[0].position + p = path.poses[0].pose.position o = goal_pose.orientation new_pose = PoseStamped( - frame_id=path.frame_id, - position=[p.x, p.y, p.z], - orientation=[o.x, o.y, o.z, o.w], + header=path.header, + pose=Pose(position=p, orientation=o), ) - return Path(frame_id=path.frame_id, poses=[new_pose]) + return Path(header=path.header, poses=[new_pose]) - if len(path) < 2 or spacing <= 0: + if len(path.poses) < 2 or spacing <= 0: return path # Extract x, y coordinates from path - xs = np.array([p.x for p in path.poses]) - ys = np.array([p.y for p in path.poses]) + xs = np.array([p.pose.position.x for p in path.poses]) + ys = np.array([p.pose.position.y for p in path.poses]) # Remove duplicate consecutive points diffs = np.sqrt(np.diff(xs) ** 2 + np.diff(ys) ** 2) @@ -246,13 +252,15 @@ def smooth_resample_path( resampled = [] for i in range(len(sampled_x)): new_pose = PoseStamped( - frame_id=path.frame_id, - position=[float(sampled_x[i]), float(sampled_y[i]), 0.0], - orientation=Quaternion(0, 0, 0, 1), + header=path.header, + pose=Pose( + position=Point(x=float(sampled_x[i]), y=float(sampled_y[i])), + orientation=Quaternion(w=1), + ), ) resampled.append(new_pose) - ret = Path(frame_id=path.frame_id, poses=resampled) + ret = Path(header=path.header, poses=resampled) _add_orientations_to_path(ret, goal_pose.orientation) diff --git a/dimos/mapping/occupancy/test_generated_visualizations.py b/dimos/mapping/occupancy/test_generated_visualizations.py new file mode 100644 index 00000000000..a96e03144b8 --- /dev/null +++ b/dimos/mapping/occupancy/test_generated_visualizations.py @@ -0,0 +1,86 @@ +# 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. + +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +import math + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid, Path +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest + +from dimos.mapping.occupancy.visualizations import Palette, visualize_occupancy_grid +from dimos.mapping.occupancy.visualize_path import visualize_path +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.image import image_view + + +@pytest.mark.parametrize("palette", ["rainbow", "turbo"]) +def test_generated_image_preserves_header_and_colors(palette: Palette) -> None: + grid = OccupancyGrid( + header=Header(stamp=Time(sec=1, nanosec=999), frame_id="map"), + info=MapMetaData(width=3, height=1, resolution=1), + data=[-1, 0, 100], + ) + image = Image.decode(visualize_occupancy_grid(grid, palette).encode()) + assert image.header == grid.header + assert image.encoding == "bgr8" + assert image_view(image).shape == (1, 3, 3) + np.testing.assert_array_equal( + image_view(image)[0, 0], [0, 0, 0] if palette == "rainbow" else [28, 24, 34] + ) + assert not np.array_equal(image_view(image)[0, 1], image_view(image)[0, 2]) + + +@pytest.mark.parametrize("footprint", [False, True]) +def test_path_rendering_is_relative_to_rotated_grid(footprint: bool) -> None: + grid = OccupancyGrid( + info=MapMetaData( + width=10, height=10, resolution=1, origin=Pose(orientation=Quaternion(w=1)) + ), + data=[0] * 100, + ) + path = Path( + poses=[ + PoseStamped(pose=Pose(position=Point(x=x, y=3), orientation=Quaternion(w=1))) + for x in [2, 4, 6] + ] + ) + expected = ( + visualize_path(grid, path, 1, 2) + if footprint + else visualize_occupancy_grid(grid, "rainbow", path) + ) + grid.info.origin = Pose( + position=Point(x=20, y=10), orientation=quaternion_from_euler(0, 0, math.pi / 2) + ) + path.poses = [ + PoseStamped( + pose=Pose( + position=Point(x=17, y=10 + x), orientation=quaternion_from_euler(0, 0, math.pi / 2) + ) + ) + for x in [2, 4, 6] + ] + actual = ( + visualize_path(grid, path, 1, 2) + if footprint + else visualize_occupancy_grid(grid, "rainbow", path) + ) + np.testing.assert_array_equal(image_view(actual), image_view(expected)) diff --git a/dimos/mapping/occupancy/test_gradient.py b/dimos/mapping/occupancy/test_gradient.py index a097873aae7..2a328c2141e 100644 --- a/dimos/mapping/occupancy/test_gradient.py +++ b/dimos/mapping/occupancy/test_gradient.py @@ -12,18 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. +import cv2 import numpy as np import pytest from dimos.mapping.occupancy.gradient import gradient, voronoi_gradient from dimos.mapping.occupancy.visualizations import visualize_occupancy_grid -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_view from dimos.utils.data import get_data @pytest.mark.parametrize("method", ["simple", "voronoi"]) def test_gradient(occupancy, method) -> None: - expected = Image.from_file(get_data(f"gradient_{method}.png")) + expected = cv2.imread(str(get_data(f"gradient_{method}.png")), cv2.IMREAD_COLOR) match method: case "simple": @@ -34,4 +35,4 @@ def test_gradient(occupancy, method) -> None: raise ValueError(f"Unknown resampling method: {method}") actual = visualize_occupancy_grid(og, "rainbow") - np.testing.assert_array_equal(actual.data, expected.data) + np.testing.assert_array_equal(image_view(actual), expected) diff --git a/dimos/mapping/occupancy/test_inflation.py b/dimos/mapping/occupancy/test_inflation.py index a30ad413b16..03a502f05d7 100644 --- a/dimos/mapping/occupancy/test_inflation.py +++ b/dimos/mapping/occupancy/test_inflation.py @@ -19,6 +19,7 @@ from dimos.mapping.occupancy.inflation import simple_inflate from dimos.mapping.occupancy.visualizations import visualize_occupancy_grid +from dimos.msgs.image import image_view from dimos.utils.data import get_data @@ -28,4 +29,4 @@ def test_inflation(occupancy) -> None: og = simple_inflate(occupancy, 0.2) result = visualize_occupancy_grid(og, "rainbow") - np.testing.assert_array_equal(result.data, expected) + np.testing.assert_array_equal(image_view(result), expected) diff --git a/dimos/mapping/occupancy/test_operations.py b/dimos/mapping/occupancy/test_operations.py index 89332d0bdd5..56a0b91f5c0 100644 --- a/dimos/mapping/occupancy/test_operations.py +++ b/dimos/mapping/occupancy/test_operations.py @@ -14,11 +14,15 @@ # limitations under the License. +from copy import deepcopy + import cv2 import numpy as np from dimos.mapping.occupancy.operations import overlay_occupied, smooth_occupied from dimos.mapping.occupancy.visualizations import visualize_occupancy_grid +from dimos.msgs.image import image_view +from dimos.msgs.occupancy import occupancy_view from dimos.utils.data import get_data @@ -27,14 +31,16 @@ def test_smooth_occupied(occupancy) -> None: result = visualize_occupancy_grid(smooth_occupied(occupancy), "rainbow") - np.testing.assert_array_equal(result.data, expected) + np.testing.assert_array_equal(image_view(result), expected) def test_overlay_occupied(occupancy) -> None: expected = cv2.imread(get_data("overlay_occupied.png"), cv2.IMREAD_COLOR) - overlay = occupancy.copy() - overlay.grid[50:100, 50:100] = 100 + overlay = deepcopy(occupancy) + cells = occupancy_view(overlay).copy() + cells[50:100, 50:100] = 100 + overlay.data = cells.ravel() result = visualize_occupancy_grid(overlay_occupied(occupancy, overlay), "rainbow") - np.testing.assert_array_equal(result.data, expected) + np.testing.assert_array_equal(image_view(result), expected) diff --git a/dimos/mapping/occupancy/test_path_map.py b/dimos/mapping/occupancy/test_path_map.py index 8928e1ab923..46e03ea9955 100644 --- a/dimos/mapping/occupancy/test_path_map.py +++ b/dimos/mapping/occupancy/test_path_map.py @@ -20,6 +20,7 @@ from dimos.mapping.occupancy.path_map import make_navigation_map from dimos.mapping.occupancy.visualizations import visualize_occupancy_grid +from dimos.msgs.image import image_view from dimos.utils.data import get_data @@ -31,4 +32,4 @@ def test_make_navigation_map(occupancy, strategy) -> None: og = make_navigation_map(occupancy, robot_width, strategy=strategy, gradient_strategy="voronoi") result = visualize_occupancy_grid(og, "rainbow") - np.testing.assert_array_equal(result.data, expected) + np.testing.assert_array_equal(image_view(result), expected) diff --git a/dimos/mapping/occupancy/test_visualizations.py b/dimos/mapping/occupancy/test_visualizations.py index 17b2629e802..e0066c86df7 100644 --- a/dimos/mapping/occupancy/test_visualizations.py +++ b/dimos/mapping/occupancy/test_visualizations.py @@ -19,6 +19,7 @@ import pytest from dimos.mapping.occupancy.visualizations import visualize_occupancy_grid +from dimos.msgs.image import image_view from dimos.utils.data import get_data @@ -28,4 +29,4 @@ def test_visualize_occupancy_grid(occupancy_gradient, palette) -> None: result = visualize_occupancy_grid(occupancy_gradient, palette) - np.testing.assert_array_equal(result.data, expected) + np.testing.assert_array_equal(image_view(result), expected) diff --git a/dimos/mapping/occupancy/visualizations.py b/dimos/mapping/occupancy/visualizations.py index 560707ccae4..dcb2a75f574 100644 --- a/dimos/mapping/occupancy/visualizations.py +++ b/dimos/mapping/occupancy/visualizations.py @@ -12,16 +12,20 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + from functools import lru_cache from typing import Literal, TypeAlias +from dimos_generated.nav_msgs.msg import OccupancyGrid, Path +from dimos_generated.sensor_msgs.msg import Image import matplotlib.pyplot as plt import numpy as np from numpy.typing import NDArray -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.geometry import pose_matrix +from dimos.msgs.image import image_from_array +from dimos.msgs.occupancy import occupancy_extent, occupancy_view Palette: TypeAlias = Literal["rainbow", "turbo"] @@ -31,33 +35,29 @@ def visualize_occupancy_grid( ) -> Image: match palette: case "rainbow": - bgr_image = rainbow_image(occupancy_grid.grid) + bgr_image = rainbow_image(occupancy_view(occupancy_grid)) case "turbo": - bgr_image = turbo_image(occupancy_grid.grid) + bgr_image = turbo_image(occupancy_view(occupancy_grid)) case _: raise NotImplementedError() if path is not None and len(path.poses) > 0: _draw_path(occupancy_grid, bgr_image, path) - return Image( - data=bgr_image, - format=ImageFormat.BGR, - frame_id=occupancy_grid.frame_id, - ts=occupancy_grid.ts, - ) + return image_from_array(bgr_image, encoding="bgr8", header=occupancy_grid.header) def _draw_path(occupancy_grid: OccupancyGrid, bgr_image: NDArray[np.uint8], path: Path) -> None: import cv2 + occupancy_extent(occupancy_grid) + world_to_grid = np.linalg.inv(pose_matrix(occupancy_grid.info.origin)) points = [] for pose in path.poses: - grid_coord = occupancy_grid.world_to_grid([pose.x, pose.y, pose.z]) - pixel_x = int(grid_coord.x) - pixel_y = int(grid_coord.y) - - if 0 <= pixel_x < occupancy_grid.width and 0 <= pixel_y < occupancy_grid.height: + position = pose.pose.position + local = world_to_grid @ np.array([position.x, position.y, position.z, 1.0]) + pixel_x, pixel_y = np.floor(local[:2] / occupancy_grid.info.resolution).astype(int) + if 0 <= pixel_x < occupancy_grid.info.width and 0 <= pixel_y < occupancy_grid.info.height: points.append((pixel_x, pixel_y)) if len(points) > 1: @@ -140,7 +140,7 @@ def _interpolate_turbo(t: float) -> tuple[int, int, int]: def generate_rgba_texture( - grid: OccupancyGrid, + grid: NDArray[np.int8], colormap: str | None = None, opacity: float = 1.0, cost_range: tuple[int, int] | None = None, @@ -149,7 +149,7 @@ def generate_rgba_texture( """Generate RGBA texture for an occupancy grid. Args: - grid: OccupancyGrid to render. + grid: NDArray[np.int8] to render. colormap: Optional matplotlib colormap name. opacity: Blend factor (0.0 to 1.0). Blends towards background color. cost_range: Optional (min, max) cost range. Cells outside range use background. @@ -166,18 +166,18 @@ def generate_rgba_texture( bg_rgb = np.array([0, 0, 0], dtype=np.float32) if cost_range is not None: - in_range_mask = (grid.grid >= cost_range[0]) & (grid.grid <= cost_range[1]) + in_range_mask = (grid >= cost_range[0]) & (grid <= cost_range[1]) else: in_range_mask = None if colormap is not None: cmap = plt.get_cmap(colormap) - grid_float = grid.grid.astype(np.float32) + grid_float = grid.astype(np.float32) - vis = np.zeros((grid.height, grid.width, 4), dtype=np.uint8) + vis = np.zeros((*grid.shape, 4), dtype=np.uint8) - free_mask = grid.grid == 0 - occupied_mask = grid.grid > 0 + free_mask = grid == 0 + occupied_mask = grid > 0 if np.any(free_mask): fg = np.array(cmap(0.0)[:3]) * 255 @@ -193,21 +193,21 @@ def generate_rgba_texture( vis[occupied_mask, :3] = blended.astype(np.uint8) vis[occupied_mask, 3] = 255 - unknown_mask = grid.grid == -1 + unknown_mask = grid == -1 vis[unknown_mask] = 0 if in_range_mask is not None: - out_of_range = ~in_range_mask & (grid.grid != -1) + out_of_range = ~in_range_mask & (grid != -1) vis[out_of_range, :3] = bg_rgb.astype(np.uint8) vis[out_of_range, 3] = 255 return vis # Default: Foxglove-style coloring - vis = np.zeros((grid.height, grid.width, 4), dtype=np.uint8) + vis = np.zeros((*grid.shape, 4), dtype=np.uint8) - free_mask = grid.grid == 0 - occupied_mask = grid.grid > 0 + free_mask = grid == 0 + occupied_mask = grid > 0 fg_free = np.array([72, 73, 129], dtype=np.float32) blended_free = fg_free * opacity + bg_rgb * (1 - opacity) @@ -215,18 +215,18 @@ def generate_rgba_texture( vis[free_mask, 3] = 255 if np.any(occupied_mask): - costs = grid.grid[occupied_mask].astype(np.float32) + costs = grid[occupied_mask].astype(np.float32) factor = (1 - costs / 100).clip(0, 1) fg_occ = np.column_stack([72 * factor, 73 * factor, 129 * factor]) blended_occ = fg_occ * opacity + bg_rgb * (1 - opacity) vis[occupied_mask, :3] = blended_occ.astype(np.uint8) vis[occupied_mask, 3] = 255 - unknown_mask = grid.grid == -1 + unknown_mask = grid == -1 vis[unknown_mask] = 0 if in_range_mask is not None: - out_of_range = ~in_range_mask & (grid.grid != -1) + out_of_range = ~in_range_mask & (grid != -1) vis[out_of_range, :3] = bg_rgb.astype(np.uint8) vis[out_of_range, 3] = 255 diff --git a/dimos/mapping/occupancy/visualize_path.py b/dimos/mapping/occupancy/visualize_path.py index 4d3de0f0a0e..10467b3ba6b 100644 --- a/dimos/mapping/occupancy/visualize_path.py +++ b/dimos/mapping/occupancy/visualize_path.py @@ -12,12 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + +from dimos_generated.nav_msgs.msg import OccupancyGrid, Path +from dimos_generated.sensor_msgs.msg import Image import numpy as np from dimos.mapping.occupancy.visualizations import visualize_occupancy_grid -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.geometry import pose_matrix +from dimos.msgs.image import image_from_array, image_view +from dimos.msgs.occupancy import occupancy_extent def visualize_path( @@ -31,7 +35,7 @@ def visualize_path( import cv2 image = visualize_occupancy_grid(occupancy_grid, "rainbow") - bgr = image.data + bgr = image_view(image) bgr = cv2.resize( bgr, @@ -40,19 +44,21 @@ def visualize_path( ) # Convert robot dimensions from meters to grid cells, then to scaled pixels - resolution = occupancy_grid.resolution + occupancy_extent(occupancy_grid) + resolution = occupancy_grid.info.resolution + world_to_grid = np.linalg.inv(pose_matrix(occupancy_grid.info.origin)) robot_width_px = int((robot_width / resolution) * scale) robot_length_px = int((robot_length / resolution) * scale) # Draw robot rectangle at each path point for pose in path.poses: # Convert world coordinates to grid coordinates - grid_coord = occupancy_grid.world_to_grid([pose.x, pose.y, pose.z]) - cx = int(grid_coord.x * scale) - cy = int(grid_coord.y * scale) + local = world_to_grid @ pose_matrix(pose.pose) + cx = int(local[0, 3] / resolution * scale) + cy = int(local[1, 3] / resolution * scale) - # Get yaw angle from pose orientation - yaw = pose.yaw + # Get heading relative to the grid's origin orientation. + yaw = np.arctan2(local[1, 0], local[0, 0]) # Define rectangle corners centered at origin (length along x, width along y) half_length = robot_length_px / 2 @@ -81,9 +87,4 @@ def visualize_path( pts = rotated_corners.astype(np.int32).reshape((-1, 1, 2)) cv2.polylines(bgr, [pts], isClosed=True, color=(0, 0, 0), thickness=thickness) - return Image( - data=bgr, - format=ImageFormat.BGR, - frame_id=occupancy_grid.frame_id, - ts=occupancy_grid.ts, - ) + return image_from_array(bgr, encoding="bgr8", header=occupancy_grid.header) diff --git a/dimos/mapping/odometry_hist.py b/dimos/mapping/odometry_hist.py index f368a7fed1b..eb2594007ae 100644 --- a/dimos/mapping/odometry_hist.py +++ b/dimos/mapping/odometry_hist.py @@ -15,21 +15,25 @@ from __future__ import annotations from collections import deque +from copy import deepcopy import math from typing import Any +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.nav_msgs.msg import Odometry, Path from pydantic import Field from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.time import to_nanoseconds def path_at_true_height(path: Path) -> Any: """The default z lift clears a costmap a bare odometry demo has none of.""" - return path.to_rerun(z_offset=0.0, radii=0.02) + import rerun as rr + + points = [[p.pose.position.x, p.pose.position.y, p.pose.position.z] for p in path.poses] + return rr.LineStrips3D([points] if points else [], colors=[(0, 255, 128)], radii=0.02) class OdometryHistConfig(ModuleConfig): @@ -51,36 +55,37 @@ def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) self._poses: deque[PoseStamped] = deque(maxlen=self.config.max_poses) # None, not 0.0: a recording stamped from zero would suppress the first path. - self._last_publish_ts: float | None = None + self._last_publish_ns: int | None = None self._unpublished = False async def handle_odometry(self, msg: Odometry) -> None: - position = msg.pose.position + position = msg.pose.pose.position point = (position.x, position.y, position.z) - frame_id = self.config.frame_id or msg.frame_id + frame_id = self.config.frame_id or msg.header.frame_id previous = self._poses[-1] if self._poses else None if ( previous is None - or math.dist((previous.x, previous.y, previous.z), point) >= self.config.min_step_meters - ): - orientation = msg.pose.orientation - self._poses.append( - PoseStamped( - ts=msg.ts, - frame_id=frame_id, - position=list(point), - orientation=[orientation.x, orientation.y, orientation.z, orientation.w], - ) + or math.dist( + (previous.pose.position.x, previous.pose.position.y, previous.pose.position.z), + point, ) + >= self.config.min_step_meters + ): + header = deepcopy(msg.header) + header.frame_id = frame_id + self._poses.append(PoseStamped(header=header, pose=msg.pose.pose)) self._unpublished = True if not self._unpublished: return - if self._last_publish_ts is not None: - elapsed = msg.ts - self._last_publish_ts + stamp_ns = to_nanoseconds(msg.header.stamp) + if self._last_publish_ns is not None: + elapsed = stamp_ns - self._last_publish_ns # A replay restart moves the stamp backwards; publish rather than wait it out. - if 0.0 <= elapsed < self.config.min_publish_interval_seconds: + if 0 <= elapsed < round(self.config.min_publish_interval_seconds * 1_000_000_000): return - self._last_publish_ts = msg.ts + self._last_publish_ns = stamp_ns self._unpublished = False - self.odom_hist.publish(Path(ts=msg.ts, frame_id=frame_id, poses=list(self._poses))) + header = deepcopy(msg.header) + header.frame_id = frame_id + self.odom_hist.publish(Path(header=header, poses=list(self._poses))) diff --git a/dimos/mapping/osm/current_location_map.py b/dimos/mapping/osm/current_location_map.py index d573370a068..7e47d0473b8 100644 --- a/dimos/mapping/osm/current_location_map.py +++ b/dimos/mapping/osm/current_location_map.py @@ -12,12 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dataclasses import replace + +import numpy as np from PIL import Image as PILImage, ImageDraw from dimos.mapping.models import LatLon from dimos.mapping.osm.osm import MapImage, get_osm_map from dimos.mapping.osm.query import query_for_one_position, query_for_one_position_and_context from dimos.models.vl.base import VlModel +from dimos.msgs.image import image_from_array, image_to_rgb from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -70,11 +74,9 @@ def _fetch_new_map(self) -> None: self._map_image = get_osm_map(self._position, self._zoom_level, self._n_tiles) # type: ignore[arg-type] # Add position marker - import numpy as np - assert self._map_image is not None assert self._position is not None - pil_image = PILImage.fromarray(self._map_image.image.data) + pil_image = PILImage.fromarray(image_to_rgb(self._map_image.image)) draw = ImageDraw.Draw(pil_image) x, y = self._map_image.latlon_to_pixel(self._position) radius = 20 @@ -85,7 +87,12 @@ def _fetch_new_map(self) -> None: width=3, ) - self._map_image.image.data[:] = np.array(pil_image) + self._map_image = replace( + self._map_image, + image=image_from_array( + np.asarray(pil_image), encoding="rgb8", header=self._map_image.image.header + ), + ) def _position_is_too_far_off_center(self) -> bool: x, y = self._map_image.latlon_to_pixel(self._position) # type: ignore[arg-type, union-attr] @@ -108,6 +115,6 @@ def save_current_map_image(self, filepath: str = "osm_debug_map.png") -> str: self._get_current_map() # type: ignore[no-untyped-call] if self._map_image is not None: - self._map_image.image.save(filepath) + PILImage.fromarray(image_to_rgb(self._map_image.image)).save(filepath) logger.info(f"Saved OSM map image to {filepath}") return filepath diff --git a/dimos/mapping/osm/osm.py b/dimos/mapping/osm/osm.py index 505da875a8b..51350b98295 100644 --- a/dimos/mapping/osm/osm.py +++ b/dimos/mapping/osm/osm.py @@ -17,12 +17,13 @@ import io import math +from dimos_generated.sensor_msgs.msg import Image import numpy as np from PIL import Image as PILImage import requests from dimos.mapping.models import ImageCoord, LatLon -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_from_array @dataclass(frozen=True) @@ -176,7 +177,7 @@ def get_osm_map(position: LatLon, zoom_level: int = 18, n_tiles: int = 4) -> Map raise ValueError("Failed to download all tiles for the requested map.") return MapImage( - image=Image.from_numpy(np.array(output_img), format=ImageFormat.RGB), + image=image_from_array(np.array(output_img), encoding="rgb8"), position=position, zoom_level=zoom_level, n_tiles=n_tiles, diff --git a/dimos/mapping/osm/test_current_location_map.py b/dimos/mapping/osm/test_current_location_map.py new file mode 100644 index 00000000000..05e35896c24 --- /dev/null +++ b/dimos/mapping/osm/test_current_location_map.py @@ -0,0 +1,56 @@ +# 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. + +from pathlib import Path +from unittest.mock import MagicMock + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.std_msgs.msg import Header +import numpy as np +from PIL import Image as PILImage +import pytest + +from dimos.mapping.models import LatLon +from dimos.mapping.osm import current_location_map +from dimos.mapping.osm.current_location_map import CurrentLocationMap +from dimos.mapping.osm.osm import MapImage +from dimos.models.vl.base import VlModel +from dimos.msgs.image import image_from_array, image_to_rgb + + +def test_position_marker_preserves_source_stamp_color_and_cached_tile( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + position = LatLon(lat=0, lon=0) + pixels = np.full((256, 256, 3), [0, 10, 30], dtype=np.uint8) + source = image_from_array( + pixels, encoding="bgr8", header=Header(frame_id="map", stamp=Time(sec=123, nanosec=456)) + ) + tile = MapImage(source, position, 0, 1) + monkeypatch.setattr(current_location_map, "get_osm_map", lambda *_args: tile) + current = CurrentLocationMap(MagicMock(spec=VlModel)) + current.update_position(position) + current._fetch_new_map() + assert current._map_image is not None + image = current._map_image.image + assert image.header == source.header + assert image.encoding == "rgb8" + x, y = tile.latlon_to_pixel(position) + np.testing.assert_array_equal(image_to_rgb(image)[y, x], [255, 0, 0]) + np.testing.assert_array_equal(image_to_rgb(image)[0, 0], [30, 10, 0]) + np.testing.assert_array_equal(image_to_rgb(source)[y, x], [30, 10, 0]) + output = tmp_path / "map.png" + assert current.save_current_map_image(str(output)) == str(output) + with PILImage.open(output) as saved: + np.testing.assert_array_equal(np.asarray(saved), image_to_rgb(image)) diff --git a/dimos/mapping/pointclouds/demo.py b/dimos/mapping/pointclouds/demo.py index 387a32a22af..78daaf0db15 100644 --- a/dimos/mapping/pointclouds/demo.py +++ b/dimos/mapping/pointclouds/demo.py @@ -17,6 +17,9 @@ from typing import TYPE_CHECKING +from dimos_generated.nav_msgs.msg import OccupancyGrid +from dimos_generated.std_msgs.msg import Header +import numpy as np import typer from dimos.mapping.occupancy.gradient import gradient @@ -27,8 +30,8 @@ read_pointcloud, visualize, ) -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.image import image_to_bgr +from dimos.msgs.pointcloud import pointcloud_from_xyz from dimos.utils.data import get_data if TYPE_CHECKING: @@ -46,7 +49,7 @@ def _get_occupancy_grid() -> OccupancyGrid: min_height = 0.15 max_height = 0.6 occupancygrid = simple_occupancy( - PointCloud2(_get_sum_map()), + pointcloud_from_xyz(np.asarray(_get_sum_map().points), header=Header(frame_id="world")), resolution=resolution, min_height=min_height, max_height=max_height, @@ -57,15 +60,15 @@ def _get_occupancy_grid() -> OccupancyGrid: def _show_occupancy_grid(og: OccupancyGrid) -> None: import cv2 - cost_map = visualize_occupancy_grid(og, "turbo").to_opencv() - cost_map = cv2.flip(cost_map, 0) + cost_map = image_to_bgr(visualize_occupancy_grid(og, "turbo")) + flipped = cv2.flip(cost_map, 0) # Resize to make the image larger (scale by 4x) - height, width = cost_map.shape[:2] - cost_map = cv2.resize(cost_map, (width * 4, height * 4), interpolation=cv2.INTER_NEAREST) + height, width = flipped.shape[:2] + resized = cv2.resize(flipped, (width * 4, height * 4), interpolation=cv2.INTER_NEAREST) cv2.namedWindow("Occupancy Grid", cv2.WINDOW_NORMAL) - cv2.imshow("Occupancy Grid", cost_map) + cv2.imshow("Occupancy Grid", resized) cv2.waitKey(0) cv2.destroyAllWindows() diff --git a/dimos/mapping/pointclouds/test_generated_occupancy.py b/dimos/mapping/pointclouds/test_generated_occupancy.py new file mode 100644 index 00000000000..4920e8b7db5 --- /dev/null +++ b/dimos/mapping/pointclouds/test_generated_occupancy.py @@ -0,0 +1,96 @@ +# 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. + +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Callable + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.nav_msgs.msg import OccupancyGrid +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header +import numpy as np +from numpy.typing import NDArray +import pytest + +from dimos.mapping.occupancy.inflation import simple_inflate +from dimos.mapping.pointclouds.occupancy import ( + general_occupancy, + height_cost_occupancy, + simple_occupancy, +) +from dimos.msgs.occupancy import occupancy_view +from dimos.msgs.pointcloud import pointcloud_from_xyz + + +@pytest.mark.parametrize("algorithm", [general_occupancy, height_cost_occupancy, simple_occupancy]) +def test_generated_cloud_to_grid(algorithm: Callable[..., OccupancyGrid]) -> None: + header = Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id="map") + xy = np.arange(0, 2, 0.1) + x, y = np.meshgrid(xy, xy) + points = np.column_stack((x.ravel(), y.ravel(), np.zeros(x.size))) + cloud = pointcloud_from_xyz(points, header=header) + decoded = PointCloud2.decode(cloud.encode()) + grid = algorithm(decoded, resolution=0.1) + restored = OccupancyGrid.decode(grid.encode()) + assert restored.header == header + cells = occupancy_view(restored) + assert cells.shape == (39, 39) + assert np.any(cells == 0) + assert np.any(cells == -1) + assert np.all((cells >= -1) & (cells <= 100)) + override = algorithm(decoded, resolution=0.1, frame_id="world") + assert override.header.frame_id == "world" + assert override.header.stamp == header.stamp + assert decoded.header == header + + +@pytest.mark.parametrize("algorithm", [general_occupancy, height_cost_occupancy, simple_occupancy]) +@pytest.mark.parametrize("points", [np.empty((0, 3)), np.array([[np.nan, 0, 0]])]) +def test_empty_or_missing_points_remain_unknown( + algorithm: Callable[..., OccupancyGrid], + points: NDArray[np.float64], +) -> None: + cloud = pointcloud_from_xyz(points, header=Header(stamp=Time(nanosec=123), frame_id="map")) + grid = algorithm(cloud) + assert grid.header == cloud.header + assert occupancy_view(grid).tolist() == [[-1]] + assert grid.info.origin.orientation.w == 1 + + +@pytest.mark.parametrize("resolution", [0, -1, float("nan"), float("inf")]) +def test_invalid_resolution(resolution: float) -> None: + cloud = pointcloud_from_xyz(np.zeros((1, 3)), header=Header()) + with pytest.raises(ValueError, match="resolution"): + general_occupancy(cloud, resolution=resolution) + + +def test_inflation_preserves_metadata_and_source() -> None: + cloud = pointcloud_from_xyz( + np.array([[0, 0, 1.0]]), header=Header(stamp=Time(nanosec=789), frame_id="map") + ) + grid = general_occupancy(cloud, resolution=0.1) + original = occupancy_view(grid).copy() + inflated = simple_inflate(grid, 0.2) + assert inflated.header == grid.header + assert inflated.info == grid.info + assert np.count_nonzero(occupancy_view(inflated) == 100) == 13 + np.testing.assert_array_equal(occupancy_view(grid), original) + inflated.header.frame_id = "changed" + assert grid.header.frame_id == "map" + assert simple_inflate(grid, 0) == grid + for radius in [-1, float("nan"), float("inf")]: + with pytest.raises(ValueError, match="radius"): + simple_inflate(grid, radius) diff --git a/dimos/mapping/pointclouds/test_occupancy.py b/dimos/mapping/pointclouds/test_occupancy.py index 1ad3ef684d6..4f067da2e11 100644 --- a/dimos/mapping/pointclouds/test_occupancy.py +++ b/dimos/mapping/pointclouds/test_occupancy.py @@ -13,21 +13,26 @@ # See the License for the specific language governing permissions and # limitations under the License. +from pathlib import Path + import cv2 +from dimos_generated.nav_msgs.msg import OccupancyGrid +from dimos_generated.std_msgs.msg import Header import numpy as np from open3d.geometry import PointCloud import pytest from dimos.core.transport import LCMTransport +from dimos.e2e_tests.cdr_replay_fixture import write_go2_cdr_replay from dimos.mapping.occupancy.visualizations import visualize_occupancy_grid from dimos.mapping.pointclouds.occupancy import ( height_cost_occupancy, simple_occupancy, ) from dimos.mapping.pointclouds.util import read_pointcloud -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.image import image_from_file, image_view +from dimos.msgs.occupancy import occupancy_view +from dimos.msgs.pointcloud import pointcloud_from_xyz from dimos.utils.data import get_data from dimos.utils.testing.moment import OutputMoment from dimos.utils.testing.test_moment import Go2Moment @@ -53,12 +58,12 @@ def big_office() -> PointCloud: ) def test_occupancy(apartment: PointCloud, occupancy_fn, output_name: str) -> None: expected_image = cv2.imread(str(get_data(output_name)), cv2.IMREAD_GRAYSCALE) - cloud = PointCloud2.from_numpy(np.asarray(apartment.points), frame_id="map") + cloud = pointcloud_from_xyz(np.asarray(apartment.points), header=Header(frame_id="map")) occupancy_grid = occupancy_fn(cloud) # Convert grid from -1..100 to 0..101 for PNG - computed_image = (occupancy_grid.grid + 1).astype(np.uint8) + computed_image = (occupancy_view(occupancy_grid) + 1).astype(np.uint8) np.testing.assert_array_equal(computed_image, expected_image) @@ -71,14 +76,13 @@ def test_occupancy(apartment: PointCloud, occupancy_fn, output_name: str) -> Non ], ) def test_occupancy2(big_office, occupancy_fn, output_name): - expected_image = Image.from_file(get_data(output_name)) - cloud = PointCloud2.from_numpy(np.asarray(big_office.points), frame_id="") + expected_image = image_from_file(get_data(output_name)) + cloud = pointcloud_from_xyz(np.asarray(big_office.points), header=Header()) occupancy_grid = occupancy_fn(cloud) actual = visualize_occupancy_grid(occupancy_grid, "rainbow") - actual.ts = expected_image.ts - np.testing.assert_array_equal(actual, expected_image) + np.testing.assert_array_equal(image_view(actual), image_view(expected_image)) class HeightCostMoment(Go2Moment): @@ -86,8 +90,10 @@ class HeightCostMoment(Go2Moment): @pytest.fixture -def height_cost_moment(): - moment = HeightCostMoment() +def height_cost_moment(tmp_path: Path): + recording = tmp_path / "go2-cdr.db" + write_go2_cdr_replay(recording, duration_s=3) + moment = HeightCostMoment(recording) def get_moment(ts: float, publish: bool = True) -> HeightCostMoment: moment.seek(ts) @@ -103,28 +109,29 @@ def get_moment(ts: float, publish: bool = True) -> HeightCostMoment: moment.publish() return moment - yield get_moment - - moment.stop() + try: + yield get_moment + finally: + moment.stop() def test_height_cost_occupancy_from_lidar(height_cost_moment) -> None: - """Test height_cost_occupancy with real lidar data.""" + """Test height_cost_occupancy with a deterministic CDR lidar recording.""" moment = height_cost_moment(1.0) costmap = moment.costmap.value assert costmap is not None # Basic sanity checks - assert costmap.grid is not None - assert costmap.width > 0 - assert costmap.height > 0 + assert occupancy_view(costmap) is not None + assert costmap.info.width > 0 + assert costmap.info.height > 0 # Costs should be in range -1 to 100 (-1 = unknown) - assert costmap.grid.min() >= -1 - assert costmap.grid.max() <= 100 + assert occupancy_view(costmap).min() >= -1 + assert occupancy_view(costmap).max() <= 100 # Check we have some unknown, some known - known_mask = costmap.grid >= 0 + known_mask = occupancy_view(costmap) >= 0 assert known_mask.sum() > 0, "Expected some known cells" assert (~known_mask).sum() > 0, "Expected some unknown cells" diff --git a/dimos/mapping/pointclouds/tool_occupancy_speed.py b/dimos/mapping/pointclouds/tool_occupancy_speed.py index 2ae1ca66cf3..66389ba968b 100644 --- a/dimos/mapping/pointclouds/tool_occupancy_speed.py +++ b/dimos/mapping/pointclouds/tool_occupancy_speed.py @@ -14,13 +14,14 @@ import pickle +from dimos_generated.sensor_msgs.msg import PointCloud2 + from dimos.mapping.voxels.grid import VoxelGrid -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.utils.data import get_data_dir from dimos.utils.testing.replay import TimedSensorReplay -def test_build_map(): +def test_build_map() -> None: grid = VoxelGrid() replay: TimedSensorReplay[PointCloud2] = TimedSensorReplay("go2_bigoffice/lidar") diff --git a/dimos/mapping/ray_tracing/utils/raytrace_rrd.py b/dimos/mapping/ray_tracing/utils/raytrace_rrd.py index 29c2726835a..9a6cb813d2e 100644 --- a/dimos/mapping/ray_tracing/utils/raytrace_rrd.py +++ b/dimos/mapping/ray_tracing/utils/raytrace_rrd.py @@ -26,6 +26,7 @@ from pathlib import Path +from dimos_generated.sensor_msgs.msg import PointCloud2 import numpy as np from numpy.typing import NDArray import typer @@ -41,7 +42,7 @@ from dimos.memory.store.sqlite import SqliteStore from dimos.memory.tf import StreamTF from dimos.memory.vis.utils import DEFAULT_RENDER_VOXEL, default_render_voxel -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_xyz from dimos.utils.data import resolve_named_path TIMELINE = "ts" @@ -224,17 +225,26 @@ def main( for obs in lidar: t = tf.get( world_frame, - obs.data.frame_id, + obs.data.header.frame_id, time_point=obs.ts, time_tolerance=TF_MATCH_TOLERANCE_S, ) if t is None: dropped += 1 continue - x, y, z = float(t.translation.x), float(t.translation.y), float(t.translation.z) - qx, qy, qz, qw = t.rotation.x, t.rotation.y, t.rotation.z, t.rotation.w + x, y, z = ( + float(t.transform.translation.x), + float(t.transform.translation.y), + float(t.transform.translation.z), + ) + qx, qy, qz, qw = ( + t.transform.rotation.x, + t.transform.rotation.y, + t.transform.rotation.z, + t.transform.rotation.w, + ) # Sensor-frame cloud: the mapper registers it by the tf pose. - raw = obs.data.points_f32() + raw = pointcloud_xyz(obs.data).astype(np.float32) for mapper in mappers.values(): mapper.add_frame(raw, (x, y, z), (qx, qy, qz, qw)) count += 1 @@ -282,7 +292,7 @@ def main( flip = np.sum(vectors * (robot - origins), axis=1) < 0 vectors = np.where(flip[:, None], -vectors, vectors) if min_eigs is None: - lengths = np.full(len(origins), normal_scale, np.float32) + lengths: NDArray[np.float32] = np.full(len(origins), normal_scale, np.float32) else: lengths = normal_scale * _planarity_scale(min_eigs[keep]) rr.log( diff --git a/dimos/mapping/relocalization/blueprints.py b/dimos/mapping/relocalization/blueprints.py index e7acc0a469f..f62cb8712af 100644 --- a/dimos/mapping/relocalization/blueprints.py +++ b/dimos/mapping/relocalization/blueprints.py @@ -29,6 +29,9 @@ from typing import Any +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.tf2_msgs.msg import TFMessage + from dimos.core.coordination.blueprints import autoconnect from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig @@ -37,10 +40,9 @@ from dimos.mapping.relocalization.lidar.module import LocalMapRelocalization from dimos.memory.replay import ReplayStream from dimos.memory.store.sqlite import SqliteStore -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.utils.data import resolve_named_path from dimos.utils.logging_config import setup_logger +from dimos.visualization.rerun.message_helpers import cloud_archetype from dimos.visualization.vis_module import vis_module logger = setup_logger() @@ -87,9 +89,9 @@ def start(self) -> None: self.register_disposable(tf.observable().subscribe(self.tf.publish)) -def _fine_points(cloud: Any) -> Any: +def _fine_points(cloud: PointCloud2) -> Any: """The premap is millimetre-scale; draw it at that size, not the 5 cm default.""" - return cloud.to_rerun(voxel_size=0.0015) + return cloud_archetype(cloud, voxel_size=0.0015) # Off until asked for. diff --git a/dimos/mapping/relocalization/go2/module.py b/dimos/mapping/relocalization/go2/module.py index 0bec9eae10d..9dfcc3921b9 100644 --- a/dimos/mapping/relocalization/go2/module.py +++ b/dimos/mapping/relocalization/go2/module.py @@ -14,14 +14,15 @@ """Go2 relocalization: lidar relocalization plus the premap merged into the live map.""" +from dimos_generated.geometry_msgs.msg import TransformStamped +from dimos_generated.sensor_msgs.msg import PointCloud2 from reactivex import combine_latest from dimos.core.core import rpc from dimos.core.stream import In, Out from dimos.mapping.relocalization.lidar.module import LidarConfig, LidarWindowRelocalization from dimos.mapping.voxels.grid import VoxelGrid -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import concatenate_clouds, transform_cloud from dimos.utils.reactive import backpressure @@ -51,12 +52,14 @@ def start(self) -> None: ).subscribe(self._on_merge_input) ) - def _on_merge_input(self, pair: tuple[PointCloud2, Transform]) -> None: + def _on_merge_input(self, pair: tuple[PointCloud2, TransformStamped]) -> None: local, tf = pair assert self.premap is not None - premap_in_world = self.premap.transform(tf) + premap_in_world = transform_cloud(self.premap, tf) if self.config.use_carving: - grid = VoxelGrid(carve_columns=True, frame_id=local.frame_id, show_startup_log=False) + grid = VoxelGrid( + carve_columns=True, frame_id=local.header.frame_id, show_startup_log=False + ) try: grid.add_frame(premap_in_world) grid.add_frame(local) @@ -64,4 +67,4 @@ def _on_merge_input(self, pair: tuple[PointCloud2, Transform]) -> None: finally: grid.dispose() else: - self.merged_map.publish(local + premap_in_world) + self.merged_map.publish(concatenate_clouds(local, premap_in_world)) diff --git a/dimos/mapping/relocalization/lidar/module.py b/dimos/mapping/relocalization/lidar/module.py index 35ffde47898..f0ff17b0f00 100644 --- a/dimos/mapping/relocalization/lidar/module.py +++ b/dimos/mapping/relocalization/lidar/module.py @@ -16,10 +16,10 @@ from collections import deque from functools import reduce -import operator import time from typing import Any +from dimos_generated.sensor_msgs.msg import PointCloud2 from reactivex import Observable, operators as ops from dimos.core.core import rpc @@ -30,7 +30,7 @@ RelocalizeConfig, ) from dimos.mapping.relocalization.module import Config, RelocalizationModule -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import concatenate_clouds, pointcloud_to_open3d from dimos.utils.logging_config import setup_logger from dimos.utils.reactive import backpressure @@ -51,7 +51,7 @@ def window( ops.do_action(held.append), ops.throttle_first(interval), ops.filter(lambda _: len(held) >= cfg.min_frames), - ops.map(lambda _: reduce(operator.add, held)), + ops.map(lambda _: reduce(concatenate_clouds, held)), ) @@ -91,7 +91,9 @@ def start(self) -> None: return # Downsampling the premap and computing its normals and FPFH is the # pipeline's dominant cost, so it is a startup cost, not a per-fix one. - self._relocalizer = LidarRelocalizer(self.premap.pointcloud, self.config.relocalize) + self._relocalizer = LidarRelocalizer( + pointcloud_to_open3d(self.premap), self.config.relocalize + ) self.register_disposable( backpressure( self.clouds().pipe( @@ -108,20 +110,20 @@ def _maybe_log_skip(self, msg: PointCloud2) -> None: now = time.monotonic() if now - self._last_skip_log > 5.0: logger.warning( - f"relocalize skipped: n_pts={len(msg)} " + f"relocalize skipped: n_pts={msg.width * msg.height} " f"< min_local_points={self.config.min_local_points}" ) self._last_skip_log = now def _has_enough_points(self, msg: PointCloud2) -> bool: - return len(msg) >= self.config.min_local_points + return msg.width * msg.height >= self.config.min_local_points def _relocalize(self, msg: PointCloud2) -> None: assert self._relocalizer is not None t0 = time.monotonic() try: tf = self._relocalizer.relocalize( - msg.pointcloud, self.config.world_frame, self.config.map_frame + pointcloud_to_open3d(msg), self.config.world_frame, self.config.map_frame ) except Exception: logger.exception("relocalize() failed") @@ -129,11 +131,11 @@ def _relocalize(self, msg: PointCloud2) -> None: dt = time.monotonic() - t0 if tf is None: logger.info( - f"relocalize lidar: refused after {dt:.1f}s n_pts={len(msg)} " + f"relocalize lidar: refused after {dt:.1f}s n_pts={msg.width * msg.height} " f"(below fitness_threshold={self.config.relocalize.fitness_threshold})" ) return - logger.info(f"relocalize lidar: time_cost={dt:.1f}s n_pts={len(msg)}") + logger.info(f"relocalize lidar: time_cost={dt:.1f}s n_pts={msg.width * msg.height}") self.submit(tf, "lidar") diff --git a/dimos/mapping/relocalization/test_module.py b/dimos/mapping/relocalization/test_module.py index 34f44ee6d0d..3c0384c432a 100644 --- a/dimos/mapping/relocalization/test_module.py +++ b/dimos/mapping/relocalization/test_module.py @@ -15,6 +15,9 @@ from types import SimpleNamespace from typing import get_type_hints +from dimos_generated.geometry_msgs.msg import TransformStamped +from dimos_generated.std_msgs.msg import Header +from dimos_generated.vision_msgs.msg import Detection3DArray import numpy as np import pytest from scipy.spatial.transform import Rotation @@ -22,9 +25,8 @@ from dimos.core.stream import In from dimos.mapping.relocalization.lidar.module import LidarWindowRelocalization from dimos.mapping.relocalization.module import RelocalizationModule -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.msgs.geometry import inverse_transform, transform_from_matrix, transform_matrix +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_xyz @pytest.fixture @@ -53,19 +55,33 @@ def test_submit_publishes_and_checks_frames(module): """submit does not second-guess the fix; the implementation already decided.""" m = module() got = fixes(m) - tf = Transform.from_matrix(np.eye(4), frame_id="world", child_frame_id="map") + tf = TransformStamped( + header=Header(frame_id="world"), + child_frame_id="map", + transform=transform_from_matrix(np.eye(4)), + ) m.submit(tf, "x") m.submit(tf, "x") assert got == [tf, tf] # A strategy handing over the placement instead of the frame transform is # caught here rather than publishing a backwards TF. with pytest.raises(AssertionError): - m.submit(Transform.from_matrix(np.eye(4), frame_id="map", child_frame_id="world")) + m.submit( + TransformStamped( + header=Header(frame_id="map"), + child_frame_id="world", + transform=transform_from_matrix(np.eye(4)), + ) + ) def test_relocalize_once_stops_after_the_first_fix(module): """What the flag does, either way. Which one is the default is a policy call.""" - tf = Transform.from_matrix(np.eye(4), frame_id="world", child_frame_id="map") + tf = TransformStamped( + header=Header(frame_id="world"), + child_frame_id="map", + transform=transform_from_matrix(np.eye(4)), + ) once = module(relocalize_once=True) assert once.keep_relocalizing() and not once.placed @@ -86,7 +102,11 @@ def test_the_fix_goes_out_the_moment_it_is_accepted(interval): fixes, sent = Subject(), [] disposable = fix_stream(fixes, interval=interval).subscribe(sent.append) - tf = Transform.from_matrix(np.eye(4), frame_id="world", child_frame_id="map") + tf = TransformStamped( + header=Header(frame_id="world"), + child_frame_id="map", + transform=transform_from_matrix(np.eye(4)), + ) fixes.on_next(tf) assert sent == [tf] disposable.dispose() @@ -94,9 +114,9 @@ def test_the_fix_goes_out_the_moment_it_is_accepted(interval): def test_premap_defines_the_map_frame_and_waits_for_a_fix(module, tmp_path): """Loading is the base's: every strategy reads a premap and publishes it, once placed.""" - path = tmp_path / "somewhere.pc2.lcm" + path = tmp_path / "somewhere.pc2.cdr" path.write_bytes( - PointCloud2.from_numpy(np.zeros((5, 3), dtype=np.float32), timestamp=0.0).lcm_encode() + pointcloud_from_xyz(np.zeros((5, 3), dtype=np.float32), header=Header()).encode() ) m = module() published, disposables = [], [] @@ -106,11 +126,17 @@ def test_premap_defines_the_map_frame_and_waits_for_a_fix(module, tmp_path): m.register_disposable = disposables.append m._load_premap(str(path)) - assert m.premap is not None and len(m.premap) == 5 - assert m.premap.frame_id == "map" + assert m.premap is not None and m.premap.width * m.premap.height == 5 + assert m.premap.header.frame_id == "map" assert len(disposables) == 1 # the gated publish assert published == [] # ... which stays silent until a fix lands - m.submit(Transform.from_matrix(np.eye(4), frame_id="world", child_frame_id="map")) + m.submit( + TransformStamped( + header=Header(frame_id="world"), + child_frame_id="map", + transform=transform_from_matrix(np.eye(4)), + ) + ) assert published == [m.premap] # republish_loaded_map=0: once, on that fix disposables[0].dispose() @@ -137,8 +163,8 @@ def relocalizer(threshold): # Accepted: open3d places the live cloud in the map, the TF tree wants the # other direction, and relocalize() is what turns one into the other. tf = relocalizer(0.3).relocalize(None, "world", "map") - assert (tf.frame_id, tf.child_frame_id) == ("world", "map") - np.testing.assert_allclose(tf.to_matrix(), np.linalg.inv(placement), atol=1e-9) + assert (tf.header.frame_id, tf.child_frame_id) == ("world", "map") + np.testing.assert_allclose(transform_matrix(tf.transform), np.linalg.inv(placement), atol=1e-9) def test_no_config_without_naming_a_rig(): @@ -168,22 +194,28 @@ def test_the_match_runs_on_a_window_of_the_last_scans(): window(scans, cfg, interval=0.001).subscribe(matched.append) for i in range(5): - scans.on_next(PointCloud2.from_numpy(np.full((4, 3), i, dtype=np.float32), timestamp=0.0)) + scans.on_next(pointcloud_from_xyz(np.full((4, 3), i, dtype=np.float32), header=Header())) time.sleep(0.01) # clear the throttle, so every scan gets its attempt # The first scan is below min_frames; from then on the window is full. - assert [len(c) for c in matched] == [8, 12, 12, 12] + assert [c.width * c.height for c in matched] == [8, 12, 12, 12] # ... and holds the *last* three scans, not the first. - assert set(matched[-1].points_f32()[:, 0]) == {2.0, 3.0, 4.0} + assert set(pointcloud_xyz(matched[-1])[:, 0]) == {2.0, 3.0, 4.0} def test_from_matrix_inverse_matches_linalg_inv(): T = np.eye(4) T[:3, :3] = Rotation.from_euler("xyz", [0.1, -0.2, 1.3]).as_matrix() T[:3, 3] = [1.5, -2.0, 0.3] - tf = Transform.from_matrix(T, frame_id="map", child_frame_id="world").inverse() - assert (tf.frame_id, tf.child_frame_id) == ("world", "map") - np.testing.assert_allclose(tf.to_matrix(), np.linalg.inv(T), atol=1e-9) + tf = inverse_transform( + TransformStamped( + header=Header(frame_id="map"), + child_frame_id="world", + transform=transform_from_matrix(T), + ) + ) + assert (tf.header.frame_id, tf.child_frame_id) == ("world", "map") + np.testing.assert_allclose(transform_matrix(tf.transform), np.linalg.inv(T), atol=1e-9) def test_dual_strategy_merges_ports(): diff --git a/dimos/mapping/test_generated_costmapper.py b/dimos/mapping/test_generated_costmapper.py new file mode 100644 index 00000000000..1b918edf3c9 --- /dev/null +++ b/dimos/mapping/test_generated_costmapper.py @@ -0,0 +1,59 @@ +# 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. + +import math + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid +from dimos_generated.std_msgs.msg import Header +import numpy as np + +from dimos.mapping.costmapper import CostMapper +from dimos.mapping.pointclouds.occupancy import GeneralOccupancyConfig +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.occupancy import occupancy_view +from dimos.msgs.pointcloud import pointcloud_from_xyz + + +def test_cost_mapper_generates_cdr_grid_with_exact_header(request): + mapper = CostMapper(algo="general", config=GeneralOccupancyConfig(resolution=0.25)) + request.addfinalizer(mapper.stop) + header = Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id="map") + cloud = pointcloud_from_xyz(np.array([[0, 0, 0], [1, 1, 1]]), header=header) + grid = mapper._calculate_costmap(cloud) + decoded = OccupancyGrid.decode(grid.encode()) + assert decoded.header == header + assert np.any(occupancy_view(decoded) == 100) + assert np.any(occupancy_view(decoded) == 0) + + +def test_initial_safe_radius_respects_rotated_origin_and_replaces_readonly_data(request): + mapper = CostMapper(initial_safe_radius_meters=0.1) + request.addfinalizer(mapper.stop) + grid = OccupancyGrid( + info=MapMetaData( + width=3, + height=3, + resolution=1, + origin=Pose( + position=Point(x=1, y=-1), orientation=quaternion_from_euler(0, 0, math.pi / 2) + ), + ), + data=[100] * 9, + ) + mapper._apply_initial_safe_radius(grid) + expected = np.full((3, 3), 100, dtype=np.int8) + expected[1, 1] = 0 + np.testing.assert_array_equal(occupancy_view(grid), expected) diff --git a/dimos/mapping/test_generated_odometry_hist.py b/dimos/mapping/test_generated_odometry_hist.py new file mode 100644 index 00000000000..61cdae3fe40 --- /dev/null +++ b/dimos/mapping/test_generated_odometry_hist.py @@ -0,0 +1,63 @@ +# 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. + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion +from dimos_generated.nav_msgs.msg import Odometry, Path +from dimos_generated.std_msgs.msg import Header +import pytest + +from dimos.mapping.odometry_hist import OdometryHist, path_at_true_height + + +def odometry(x: float, nanosec: int, sec: int = 1700000000) -> Odometry: + message = Odometry(header=Header(stamp=Time(sec=sec, nanosec=nanosec), frame_id="odom")) + message.pose.pose = Pose(position=Point(x=x, y=2, z=3), orientation=Quaternion(w=1)) + return Odometry.decode(message.encode()) + + +@pytest.mark.asyncio +async def test_history_retains_exact_stamp_and_one_nanosecond_interval(request): + module = OdometryHist(min_publish_interval_seconds=1e-9) + request.addfinalizer(module.stop) + paths = [] + module.odom_hist.subscribe(lambda path: paths.append(Path.decode(path.encode()))) + source = odometry(1, 123456789) + await module.handle_odometry(source) + source.pose.pose.position.x = 100 + await module.handle_odometry(odometry(2, 123456790)) + assert len(paths) == 2 + assert paths[-1].header.stamp.nanosec == 123456790 + assert [p.header.stamp.nanosec for p in paths[-1].poses] == [123456789, 123456790] + assert paths[-1].poses[0].pose.position.x == 1 + assert paths[0].poses[0].pose.position.x == 1 + archetype = path_at_true_height(paths[-1]) + assert archetype.strips.as_arrow_array().to_pylist() == [[[1, 2, 3], [2, 2, 3]]] + + +@pytest.mark.asyncio +async def test_history_throttles_bounds_and_restarts_at_zero(request): + module = OdometryHist(max_poses=2, frame_id="world", min_publish_interval_seconds=0.1) + request.addfinalizer(module.stop) + paths = [] + module.odom_hist.subscribe(paths.append) + await module.handle_odometry(odometry(1, 0)) + await module.handle_odometry(odometry(2, 50000000)) + assert len(paths) == 1 + await module.handle_odometry(odometry(3, 100000000)) + assert len(paths) == 2 and len(paths[-1].poses) == 2 + assert paths[-1].header.frame_id == "world" + assert all(p.header.frame_id == "world" for p in paths[-1].poses) + await module.handle_odometry(odometry(4, 0, sec=0)) + assert len(paths) == 3 and paths[-1].header.stamp == Time() diff --git a/dimos/mapping/test_generated_voxels.py b/dimos/mapping/test_generated_voxels.py new file mode 100644 index 00000000000..6e3dc1e7722 --- /dev/null +++ b/dimos/mapping/test_generated_voxels.py @@ -0,0 +1,76 @@ +# 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. + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.std_msgs.msg import Header +import numpy as np +import open3d.core as o3c +import pytest + +from dimos.mapping.voxels.grid import VoxelGrid +from dimos.mapping.voxels.impl.o3d import O3dVoxels +from dimos.mapping.voxels.impl.packed import PackedVoxels +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_xyz + + +@pytest.mark.parametrize("backend", ["packed", "open3d"]) +@pytest.mark.parametrize("carve", [False, True]) +def test_generated_cloud_union_and_column_carving(backend, carve): + grid = ( + PackedVoxels(0.5, carve) + if backend == "packed" + else O3dVoxels(0.5, 1000, carve, o3c.Device("CPU:0")) + ) + try: + for points in [[[0.1, 0.1, 0.1], [0.1, 0.1, 1.1], [1.1, 1.1, 0.1]], [[0.1, 0.1, 0.6]]]: + grid.add_frame(pointcloud_from_xyz(np.array(points), header=Header(frame_id="map"))) + actual = {tuple(point) for point in grid.points()} + expected = {(0.25, 0.25, 0.75), (1.25, 1.25, 0.25)} + if not carve: + expected |= {(0.25, 0.25, 0.25), (0.25, 0.25, 1.25)} + assert actual == expected + grid.add_frame( + pointcloud_from_xyz(np.array([[np.nan, 0, 0], [0, np.inf, 0]]), header=Header()) + ) + assert {tuple(point) for point in grid.points()} == expected + finally: + grid.dispose() + + +def test_voxel_output_retains_exact_latest_stamp_and_invalidates_cache(): + grid = VoxelGrid(voxel_size=0.5, device="CPU:0", frame_id="map", show_startup_log=False) + try: + empty = grid.get_global_pointcloud2() + assert empty.width == 0 and empty.header.stamp == Time() + header = Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id="map") + source = pointcloud_from_xyz(np.array([[0.1, 0.2, 0.3]]), header=header) + grid.add_frame(source) + source.header.stamp.nanosec = 1 + first = grid.get_global_pointcloud2() + assert first.header.stamp == header.stamp + np.testing.assert_array_equal(pointcloud_xyz(first), [[0.25, 0.25, 0.25]]) + assert grid.get_global_pointcloud2() is first + header.stamp.nanosec += 1 + grid.add_frame(pointcloud_from_xyz(np.array([[1.1, 1.1, 1.1]]), header=header)) + second = grid.get_global_pointcloud2() + assert second is not first + assert second.header == header and second.width == 2 + assert first.width == 1 + grid.add_frame(pointcloud_from_xyz(np.empty((0, 3)), header=Header(frame_id="map"))) + assert grid.get_global_pointcloud2().header.stamp == Time() + assert grid.get_global_pointcloud2().width == 2 + finally: + grid.dispose() + with pytest.raises(RuntimeError, match="disposed"): + grid.get_global_pointcloud2() diff --git a/dimos/mapping/test_voxels.py b/dimos/mapping/test_voxels.py index 662a791b41c..26963134604 100644 --- a/dimos/mapping/test_voxels.py +++ b/dimos/mapping/test_voxels.py @@ -12,131 +12,76 @@ # See the License for the specific language governing permissions and # limitations under the License. -from collections.abc import Callable, Generator +"""Deterministic CDR replacement for the retired typed-pickle lidar fixture.""" +from collections.abc import Generator + +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest -from dimos.core.transport import LCMTransport from dimos.mapping.voxels.grid import VoxelGrid -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.utils.data import get_data -from dimos.utils.testing.legacy_pickle import LegacyPickleStore -from dimos.utils.testing.moment import OutputMoment -from dimos.utils.testing.test_moment import Go2Moment - -pytestmark = pytest.mark.self_hosted +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_xyz @pytest.fixture def grid() -> Generator[VoxelGrid, None, None]: - g = VoxelGrid() - yield g - g.dispose() - - -class Go2MapperMoment(Go2Moment): - global_map: OutputMoment[PointCloud2] = OutputMoment(LCMTransport("/global_map", PointCloud2)) - - -MomentFactory = Callable[[float, bool], Go2MapperMoment] - - -@pytest.fixture -def moment() -> Generator[MomentFactory, None, None]: - instances: list[Go2MapperMoment] = [] - - def get_moment(ts: float, publish: bool = True) -> Go2MapperMoment: - m = Go2MapperMoment() - m.seek(ts) - if publish: - m.publish() - instances.append(m) - return m - - yield get_moment - for m in instances: - m.stop() - - -@pytest.fixture -def moment1(moment: MomentFactory) -> Go2MapperMoment: - return moment(10, False) + value = VoxelGrid(device="CPU:0", show_startup_log=False) + try: + yield value + finally: + value.dispose() @pytest.fixture -def moment2(moment: MomentFactory) -> Go2MapperMoment: - return moment(85, False) +def lidar_frame() -> PointCloud2: + axis = np.arange(20) * 0.05 + 0.025 + points = np.stack(np.meshgrid(axis, axis, axis, indexing="ij"), axis=-1).reshape(-1, 3) + message = pointcloud_from_xyz(points, header=Header(frame_id="world")) + return PointCloud2.decode(message.encode()) def test_ingest_a_few(grid: VoxelGrid) -> None: - data_dir = get_data("unitree_go2_office_walk2") - lidar_store = LegacyPickleStore(f"{data_dir}/lidar") - - for i in [1, 4, 8]: - frame = lidar_store.find_closest_seek(i) - assert frame is not None - print("add", frame) - grid.add_frame(frame) - - assert len(grid.get_global_pointcloud2()) == 30136 - - -@pytest.mark.parametrize( - "voxel_size, expected_points", - [ - (0.5, 277), - (0.1, 7290), - (0.05, 28199), - ], -) -def test_roundtrip(moment1: Go2MapperMoment, voxel_size: float, expected_points: int) -> None: - lidar_frame = moment1.lidar.value - assert lidar_frame is not None - - grid = VoxelGrid(voxel_size=voxel_size) + for offset in (0.0, 1.0, 2.0): + points = np.array([[offset + 0.025, 0.025, 0.025], [offset + 0.075, 0.025, 0.025]]) + frame = pointcloud_from_xyz(points, header=Header(frame_id="world")) + grid.add_frame(PointCloud2.decode(frame.encode())) + assert grid.get_global_pointcloud2().width == 6 + np.testing.assert_allclose( + np.sort(pointcloud_xyz(grid.get_global_pointcloud2())[:, 0]), + [0.025, 0.075, 1.025, 1.075, 2.025, 2.075], + atol=1e-6, + ) + + +@pytest.mark.parametrize("voxel_size,expected_points", [(0.5, 8), (0.1, 1000), (0.05, 8000)]) +def test_roundtrip(lidar_frame: PointCloud2, voxel_size: float, expected_points: int) -> None: + grid = VoxelGrid(voxel_size=voxel_size, device="CPU:0", show_startup_log=False) + try: + grid.add_frame(lidar_frame) + first = grid.get_global_pointcloud2() + assert first.width == expected_points + if voxel_size == 0.05: + assert first.width == lidar_frame.width + np.testing.assert_allclose( + np.sort(pointcloud_xyz(first), axis=0), + np.sort(pointcloud_xyz(lidar_frame), axis=0), + atol=1e-6, + ) + grid.add_frame(PointCloud2.decode(first.encode())) + assert grid.get_global_pointcloud2().width == expected_points + np.testing.assert_array_equal( + pointcloud_xyz(grid.get_global_pointcloud2()), pointcloud_xyz(first) + ) + finally: + grid.dispose() + + +def test_roundtrip_range_preserved(grid: VoxelGrid, lidar_frame: PointCloud2) -> None: + inputs = pointcloud_xyz(lidar_frame) grid.add_frame(lidar_frame) - - global1 = grid.get_global_pointcloud2() - assert len(global1) == expected_points - - # loseless roundtrip - if voxel_size == 0.05: - assert len(global1) == len(lidar_frame) - # TODO: we want __eq__ on PointCloud2 - should actually compare - # all points in both frames - - grid.add_frame(global1) - # no new information, no global map change - assert len(grid.get_global_pointcloud2()) == len(global1) - - moment1.publish() - grid.dispose() - - -def test_roundtrip_range_preserved(grid: VoxelGrid) -> None: - """Test that input coordinate ranges are preserved in output.""" - data_dir = get_data("unitree_go2_office_walk2") - lidar_store = LegacyPickleStore(f"{data_dir}/lidar") - - frame = lidar_store.find_closest_seek(1.0) - assert frame is not None - input_pts = np.asarray(frame.pointcloud.points) - - grid.add_frame(frame) - - out_pcd = grid.get_global_pointcloud().to_legacy() - out_pts = np.asarray(out_pcd.points) - - voxel_size = grid._voxel_size - tolerance = voxel_size # Allow one voxel of difference at boundaries - - # TODO: we want __eq__ on PointCloud2 - should actually compare - # all points in both frames - - for axis, name in enumerate(["X", "Y", "Z"]): - in_min, in_max = input_pts[:, axis].min(), input_pts[:, axis].max() - out_min, out_max = out_pts[:, axis].min(), out_pts[:, axis].max() - - assert abs(in_min - out_min) < tolerance, f"{name} min mismatch: in={in_min}, out={out_min}" - assert abs(in_max - out_max) < tolerance, f"{name} max mismatch: in={in_max}, out={out_max}" + outputs = np.asarray(grid.get_global_pointcloud().to_legacy().points) + for axis in range(3): + assert abs(inputs[:, axis].min() - outputs[:, axis].min()) < grid._voxel_size + assert abs(inputs[:, axis].max() - outputs[:, axis].max()) < grid._voxel_size diff --git a/dimos/mapping/tool_voxels.py b/dimos/mapping/tool_voxels.py index 84fd48c729a..4a360fdcf11 100644 --- a/dimos/mapping/tool_voxels.py +++ b/dimos/mapping/tool_voxels.py @@ -13,14 +13,16 @@ # limitations under the License. from collections.abc import Callable, Generator +import os import time +from dimos_generated.sensor_msgs.msg import PointCloud2 import numpy as np import pytest from dimos.core.transport import LCMTransport from dimos.mapping.voxels.grid import VoxelGrid -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_xyz from dimos.utils.testing.moment import OutputMoment from dimos.utils.testing.test_moment import Go2Moment @@ -44,7 +46,7 @@ def moment() -> Generator[MomentFactory, None, None]: instances: list[Go2MapperMoment] = [] def get_moment(ts: float, publish: bool = True) -> Go2MapperMoment: - m = Go2MapperMoment() + m = Go2MapperMoment(os.environ["DIMOS_CDR_REPLAY"]) m.seek(ts) if publish: m.publish() @@ -87,8 +89,8 @@ def test_carving(grid: VoxelGrid, moment1: Go2MapperMoment, moment2: Go2MapperMo count_carving = grid.size() voxel_size = grid._voxel_size - pts1 = np.asarray(lidar_frame1.pointcloud.points) - pts2 = np.asarray(lidar_frame2.pointcloud.points) + pts1 = pointcloud_xyz(lidar_frame1) + pts2 = pointcloud_xyz(lidar_frame2) combined_vox = np.floor(np.vstack([pts1, pts2]) / voxel_size).astype(np.int64) count_additive = np.unique(combined_vox, axis=0).shape[0] diff --git a/dimos/mapping/voxels/impl/o3d.py b/dimos/mapping/voxels/impl/o3d.py index 82e80982b97..9fd3e48572d 100644 --- a/dimos/mapping/voxels/impl/o3d.py +++ b/dimos/mapping/voxels/impl/o3d.py @@ -16,33 +16,17 @@ from typing import TYPE_CHECKING +from dimos_generated.sensor_msgs.msg import PointCloud2 import numpy as np +from numpy.typing import NDArray -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_xyz if TYPE_CHECKING: import open3d as o3d # type: ignore[import-untyped] import open3d.core as o3c # type: ignore[import-untyped] -def ensure_tensor_pcd( - pcd_any: o3d.t.geometry.PointCloud | o3d.geometry.PointCloud, - device: o3c.Device, -) -> o3d.t.geometry.PointCloud: - """Convert legacy / cuda.pybind point clouds into o3d.t.geometry.PointCloud on `device`.""" - import open3d as o3d # type: ignore[import-untyped] - import open3d.core as o3c # type: ignore[import-untyped] - - if isinstance(pcd_any, o3d.t.geometry.PointCloud): - return pcd_any.to(device) - - assert isinstance(pcd_any, o3d.geometry.PointCloud), ( - "Input must be a legacy PointCloud or a tensor PointCloud" - ) - - return o3d.t.geometry.PointCloud.from_legacy(pcd_any, o3c.float32, device) - - class O3dVoxels: """CUDA voxel store: Open3D VoxelBlockGrid hashmap.""" @@ -70,11 +54,12 @@ def __init__( def add_frame(self, frame: PointCloud2) -> None: import open3d.core as o3c # type: ignore[import-untyped] - pcd = ensure_tensor_pcd(frame.pointcloud, self._dev) - if pcd.is_empty(): + values = pointcloud_xyz(frame).astype(np.float32) + values = values[np.isfinite(values).all(axis=1)] + if not len(values): return - pts = pcd.point["positions"].to(self._dev, o3c.float32) + pts = o3c.Tensor(values, device=self._dev, dtype=o3c.float32) vox = (pts / self._voxel_size).floor().to(self._key_dtype) keys_Nx3 = vox.contiguous() @@ -127,7 +112,7 @@ def _carve_and_insert(self, new_keys: o3c.Tensor) -> None: self._hashmap.activate(new_keys) - def points(self) -> np.ndarray: + def points(self) -> NDArray[np.float32]: """Voxel centers, (N, 3) float32.""" import open3d.core as o3c # type: ignore[import-untyped] diff --git a/dimos/mapping/voxels/module.py b/dimos/mapping/voxels/module.py index b273889a4c0..b67ebffb419 100644 --- a/dimos/mapping/voxels/module.py +++ b/dimos/mapping/voxels/module.py @@ -16,6 +16,8 @@ from typing import TYPE_CHECKING, Any +from dimos_generated.sensor_msgs.msg import PointCloud2 + from dimos.core.core import rpc from dimos.core.module import ModuleConfig from dimos.core.stream import In, Out @@ -23,7 +25,6 @@ from dimos.memory.module import StreamModule from dimos.memory.stream import Stream from dimos.memory.transform import Transformer -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 if TYPE_CHECKING: from collections.abc import Iterator diff --git a/dimos/mapping/voxels/test_grid.py b/dimos/mapping/voxels/test_grid.py new file mode 100644 index 00000000000..365661d437f --- /dev/null +++ b/dimos/mapping/voxels/test_grid.py @@ -0,0 +1,48 @@ +# Copyright 2025-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. + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest + +from dimos.mapping.voxels.grid import VoxelGrid +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_xyz + + +def test_cpu_accumulation_preserves_latest_exact_stamp_and_output_frame() -> None: + first = pointcloud_from_xyz( + np.array([[0.25, 0.25, 0.25]], dtype=np.float32), + header=Header(frame_id="lidar", stamp=Time(sec=1, nanosec=987654321)), + ) + second = pointcloud_from_xyz( + np.array([[1.25, 1.25, 1.25]], dtype=np.float32), + header=Header(frame_id="lidar", stamp=Time(sec=2, nanosec=123456789)), + ) + grid = VoxelGrid(voxel_size=0.5, device="CPU:0", carve_columns=False, frame_id="map") + try: + grid.add_frame(first) + grid.add_frame(second) + result = grid.get_global_pointcloud2() + decoded = type(result).decode(result.encode()) + assert decoded.header.frame_id == "map" + assert decoded.header.stamp == second.header.stamp + assert first.header.frame_id == second.header.frame_id == "lidar" + np.testing.assert_allclose( + pointcloud_xyz(decoded), [[0.25, 0.25, 0.25], [1.25, 1.25, 1.25]] + ) + finally: + grid.dispose() + with pytest.raises(RuntimeError, match="disposed"): + grid.get_global_pointcloud2() diff --git a/dimos/memory/architecture.md b/dimos/memory/architecture.md index 1e8e6316a17..5e3ce102149 100644 --- a/dimos/memory/architecture.md +++ b/dimos/memory/architecture.md @@ -17,7 +17,7 @@ Store → Stream → [filters / transforms / terminals] → Stream → [filters Supporting Systems: - BlobStore — separates large payloads from metadata. FileBlobStore (files on disk) and SqliteBlobStore (blob table per stream). Supports lazy loading. -- Codecs — codec_for() auto-selects: JpegCodec for images (TurboJPEG, ~10-20x compression), LcmCodec for DimOS messages, PickleCodec fallback. +- Codecs — codec_for() selects CdrCodec for generated messages, including raw Image and explicit CompressedImage values. Python objects use PickleCodec. SQLite can wrap CDR in LZ4; MCAP uses chunk compression. - Transformers — Transformer[T,R] ABC wrapping iterator-to-iterator. EmbedImages/EmbedText enrich observations with embeddings. QualityWindow keeps best per time window. - Backpressure Buffers — KeepLast, Bounded, DropNew, Unbounded — bridge push/pull for live mode. diff --git a/dimos/memory/cli/dataset.py b/dimos/memory/cli/dataset.py index a9e16648585..932a756d599 100644 --- a/dimos/memory/cli/dataset.py +++ b/dimos/memory/cli/dataset.py @@ -17,7 +17,7 @@ One entry point for every CLI that opens a recording. :func:`open_dataset` resolves a dataset name/path (bare names look up the cwd / repo ``data/`` dir) and picks the store by file extension: ``.db`` -> SqliteStore, ``.mcap`` -> -Go2McapStore. Use :func:`open_store` when the path is already resolved. +McapStore. Use :func:`open_store` when the path is already resolved. """ from __future__ import annotations @@ -35,9 +35,9 @@ def open_store(path: str | Path) -> Store: """Open an already-resolved dataset *path*, dispatching on its extension.""" s = str(path) if s.endswith(".mcap"): - from dimos.robot.unitree.go2.dds.store import Go2McapStore # lazy: robot-layer codecs + from dimos.memory.store.mcap import McapStore - return Go2McapStore(path=s) + return McapStore(path=s) if s.endswith(".db"): from dimos.memory.store.sqlite import SqliteStore diff --git a/dimos/memory/cli/render.py b/dimos/memory/cli/render.py index 1b6729f8023..e60664b6456 100644 --- a/dimos/memory/cli/render.py +++ b/dimos/memory/cli/render.py @@ -15,8 +15,8 @@ """Render any memory store into rerun. Generic: walks the store's streams and logs every observation whose payload -implements ``to_rerun()`` (the :class:`RerunConvertible` convention). Streams -whose payload has no ``to_rerun`` are skipped. Each stream becomes an entity +has a generated-message viewer helper or an explicit application ``to_rerun()`` +adapter. Unsupported payloads are skipped. Each stream becomes an entity path; observations share one ``time`` timeline (relative to the store's earliest observation, so streams stay aligned). CameraInfo streams are the exception: logged once as a Pinhole on their matching image entity (see @@ -25,16 +25,56 @@ from __future__ import annotations +from copy import deepcopy from pathlib import Path import shutil import subprocess from typing import TYPE_CHECKING, Any if TYPE_CHECKING: + from dimos_generated.sensor_msgs.msg import CameraInfo + from dimos.memory.store.base import Store from dimos.memory.stream import Stream from dimos.memory.type.observation import Observation - from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo + + +def _render_message(data: Any) -> Any: + """Render generated standard messages, or an application's explicit adapter.""" + from dimos_generated.geometry_msgs.msg import PointStamped, PoseStamped + from dimos_generated.nav_msgs.msg import OccupancyGrid, Odometry, Path + from dimos_generated.sensor_msgs.msg import CameraInfo, CompressedImage, Image, PointCloud2 + from dimos_generated.tf2_msgs.msg import TFMessage + from dimos_generated.vision_msgs.msg import Detection3DArray + + from dimos.visualization.rerun.message_helpers import ( + camera_pinhole, + cloud_archetype, + detection_boxes, + image_archetype, + navigation_archetype, + occupancy_mesh, + tf_archetypes, + ) + + if isinstance(data, CameraInfo): + return camera_pinhole(data) + if isinstance(data, (Image, CompressedImage)): + return image_archetype(data) + if isinstance(data, PointCloud2): + return cloud_archetype(data) + if isinstance(data, TFMessage): + return tf_archetypes(data) + if isinstance(data, (PointStamped, PoseStamped, Odometry, Path)): + return navigation_archetype(data) + if isinstance(data, OccupancyGrid): + return occupancy_mesh(data) + if isinstance(data, Detection3DArray): + return detection_boxes(data) + adapter = getattr(data, "to_rerun", None) + if callable(adapter): + return adapter() + raise TypeError(f"no viewer adapter for {type(data).__name__}") def _pair_camera_infos( @@ -46,27 +86,30 @@ def _pair_camera_infos( has one of each; the image's frame_id wins as the pinhole's parent frame. Returns (image stream -> (CameraInfo, parent frame), paired stream names). """ - from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo - from dimos.msgs.sensor_msgs.Image import Image + from dimos_generated.sensor_msgs.msg import CameraInfo, CompressedImage, Image infos = [(n, f.data) for n, _, f in renderable if isinstance(f.data, CameraInfo)] - images = [(n, f.data) for n, _, f in renderable if isinstance(f.data, Image)] + images = [(n, f.data) for n, _, f in renderable if isinstance(f.data, (Image, CompressedImage))] pinholes: dict[str, tuple[CameraInfo, str]] = {} paired: set[str] = set() for info_name, info in infos: - targets = [(n, img) for n, img in images if img.frame_id and img.frame_id == info.frame_id] + targets = [ + (n, img) + for n, img in images + if img.header.frame_id and img.header.frame_id == info.header.frame_id + ] if not targets and len(infos) == 1 and len(images) == 1: targets = [images[0]] print( - f" {info_name}: frame_id {info.frame_id!r} matches no image stream, " + f" {info_name}: frame_id {info.header.frame_id!r} matches no image stream, " f"pairing with {targets[0][0]!r} as the only image stream" ) if not targets: - print(f" {info_name}: no image stream matches frame_id {info.frame_id!r}") + print(f" {info_name}: no image stream matches frame_id {info.header.frame_id!r}") continue for target, img in targets: - pinholes[target] = (info, img.frame_id or info.frame_id) + pinholes[target] = (info, img.header.frame_id or info.header.frame_id) paired.add(info_name) return pinholes, paired @@ -81,7 +124,7 @@ def _frame_first_seen( first tf message overall (which may be an unrelated or static edge). Scans each tf stream only until every frame is found. """ - from dimos.msgs.tf2_msgs.TFMessage import TFMessage + from dimos_generated.tf2_msgs.msg import TFMessage seen: dict[str, float] = {} for _, stream, first in renderable: @@ -124,11 +167,12 @@ def render_store( (``/``) — except a stream whose name matches ``root``'s last segment, which stays at ```` itself. Returns the ``.rrd`` path. """ + from dimos_generated.tf2_msgs.msg import TFMessage import rerun as rr from dimos.memory.utils.progress import progress - from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.visualization.rerun.init import rerun_init + from dimos.visualization.rerun.message_helpers import camera_pinhole if out is None: src = getattr(store.config, "path", None) or "store" @@ -151,13 +195,10 @@ def entity(name: str) -> str: except LookupError: continue data = first.data - if not hasattr(data, "to_rerun"): - print(f" skip {name}: {type(data).__name__} has no to_rerun()") - continue try: - data.to_rerun() + _render_message(data) except Exception as e: - print(f" skip {name}: to_rerun() failed ({e})") + print(f" skip {name}: viewer conversion failed ({e})") continue renderable.append((name, stream, first)) @@ -180,7 +221,9 @@ def entity(name: str) -> str: rr.save(out) for image_name, (info, frame) in pinholes.items(): - pinhole = info.to_rerun_pinhole(optical_frame=frame) + calibration = deepcopy(info) + calibration.header.frame_id = frame + pinhole = camera_pinhole(calibration) at = frame_at.get(frame) if not has_tf: rr.log(entity(image_name), pinhole, static=True) @@ -205,7 +248,7 @@ def entity(name: str) -> str: report(obs) continue rr.set_time("time", duration=obs.ts - t0) - data = obs.data.to_rerun() + data = _render_message(obs.data) path = entity(name) if isinstance(data, list): # RerunMulti: [(subpath, archetype), ...] for sub, arch in data: diff --git a/dimos/memory/cli/test_render.py b/dimos/memory/cli/test_render.py index acb0be178b4..1b3f4540ecc 100644 --- a/dimos/memory/cli/test_render.py +++ b/dimos/memory/cli/test_render.py @@ -15,16 +15,17 @@ from pathlib import Path from types import SimpleNamespace +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage import numpy as np from dimos.memory.cli.render import _frame_first_seen, _pair_camera_infos, render_store from dimos.memory.store.memory import MemoryStore -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.camera_info import camera_info_from_intrinsics +from dimos.msgs.image import image_from_array +from dimos.msgs.time import time_from_seconds def _entry(name: str, data: object) -> tuple: # type: ignore[type-arg] @@ -32,13 +33,21 @@ def _entry(name: str, data: object) -> tuple: # type: ignore[type-arg] def _info(frame_id: str) -> CameraInfo: - return CameraInfo.from_intrinsics( - fx=500.0, fy=500.0, cx=320.0, cy=240.0, width=640, height=480, frame_id=frame_id + return camera_info_from_intrinsics( + fx=500.0, + fy=500.0, + cx=320.0, + cy=240.0, + width=640, + height=480, + header=Header(frame_id=frame_id), ) def _image(frame_id: str) -> Image: - return Image(data=np.zeros((4, 4, 3), dtype=np.uint8), frame_id=frame_id) + return image_from_array( + np.zeros((4, 4, 3), dtype=np.uint8), encoding="bgr8", header=Header(frame_id=frame_id) + ) def test_pairs_by_frame_id() -> None: @@ -50,7 +59,7 @@ def test_pairs_by_frame_id() -> None: pinholes, paired = _pair_camera_infos(renderable) assert set(pinholes) == {"color_image"} info, frame = pinholes["color_image"] - assert info.frame_id == "camera_optical" + assert info.header.frame_id == "camera_optical" assert frame == "camera_optical" assert paired == {"camera_info"} @@ -98,13 +107,13 @@ def test_no_infos() -> None: def _tf(child: str, ts: float) -> TFMessage: return TFMessage( - Transform( - translation=Vector3(1.0, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="world", - child_frame_id=child, - ts=ts, - ) + transforms=[ + TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(ts)), + child_frame_id=child, + transform=Transform(translation=Vector3(x=1), rotation=Quaternion(w=1)), + ) + ] ) diff --git a/dimos/memory/codecs/lcm.py b/dimos/memory/codecs/cdr.py similarity index 67% rename from dimos/memory/codecs/lcm.py rename to dimos/memory/codecs/cdr.py index 9dd010aa3d2..51d127d983a 100644 --- a/dimos/memory/codecs/lcm.py +++ b/dimos/memory/codecs/cdr.py @@ -20,10 +20,14 @@ from dimos.msgs.protocol import DimosMsg -class LcmCodec: - """Codec for DimosMsg types — uses lcm_encode/lcm_decode.""" +class CdrCodec: + """Store a generated message as its transport-neutral CDR bytes.""" def __init__(self, msg_type: type[DimosMsg]) -> None: + if not all( + hasattr(msg_type, member) for member in ("encode", "decode", "msg_name", "schema") + ): + raise TypeError(f"{msg_type!r} is not a generated CDR message type") self._msg_type = msg_type @property @@ -32,7 +36,9 @@ def payload_type(self) -> type[DimosMsg]: return self._msg_type def encode(self, value: DimosMsg) -> bytes: - return value.lcm_encode() + if not isinstance(value, self._msg_type): + raise TypeError(f"Expected {self._msg_type.msg_name}, got {type(value).__name__}") + return value.encode() def decode(self, data: bytes) -> DimosMsg: - return self._msg_type.lcm_decode(data) + return self._msg_type.decode(data) diff --git a/dimos/memory/codecs/jpeg.py b/dimos/memory/codecs/jpeg.py deleted file mode 100644 index 8f1bf8625bd..00000000000 --- a/dimos/memory/codecs/jpeg.py +++ /dev/null @@ -1,39 +0,0 @@ -# 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. - -from __future__ import annotations - -from dimos.msgs.sensor_msgs.Image import Image - - -class JpegCodec: - """Codec for Image types — JPEG-compressed inside an LCM Image envelope. - - Uses ``Image.lcm_jpeg_encode/decode`` which preserves ``ts``, ``frame_id``, - and all LCM header fields. Pixel data is lossy-compressed via TurboJPEG. - """ - - def __init__(self, quality: int = 50) -> None: - self._quality = quality - - @property - def payload_type(self) -> type[Image]: - """Message type decoded by this storage codec.""" - return Image - - def encode(self, value: Image) -> bytes: - return value.lcm_jpeg_encode(quality=self._quality) - - def decode(self, data: bytes) -> Image: - return Image.lcm_jpeg_decode(data) diff --git a/dimos/memory/demo_data.py b/dimos/memory/demo_data.py new file mode 100644 index 00000000000..d1abd2a84f6 --- /dev/null +++ b/dimos/memory/demo_data.py @@ -0,0 +1,65 @@ +# 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. + +"""Small deterministic CDR recordings for offline documentation examples.""" + +import math +from pathlib import Path + +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header +import numpy as np + +from dimos.memory.codecs.cdr import CdrCodec +from dimos.memory.codecs.lz4 import Lz4Codec +from dimos.memory.store.sqlite import SqliteStore +from dimos.msgs.image import image_from_array +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import time_from_nanoseconds + + +def write_demo_recording(path: Path, *, samples: int = 720) -> SqliteStore: + """Create a new image/pose/cloud recording at 4 Hz without downloads or models. + + Refuse existing files so examples cannot replace a user's recording. + Images have both exact source stamps and observation poses. Storage uses + LZ4 around generated CDR, retaining complete type/schema metadata. + """ + if path.exists(): + raise FileExistsError(path) + if samples <= 0: + raise ValueError("samples must be positive") + store = SqliteStore(path=path) + images = store.stream("color_image", Image, codec=Lz4Codec(CdrCodec(Image))) + poses = store.stream("odom", PoseStamped, codec=Lz4Codec(CdrCodec(PoseStamped))) + clouds = store.stream("lidar", PointCloud2, codec=Lz4Codec(CdrCodec(PointCloud2))) + x, y = np.meshgrid(np.linspace(-2.0, 2.0, 16), np.linspace(-2.0, 2.0, 16)) + points = np.column_stack((x.ravel(), y.ravel(), np.zeros(x.size))) + for index in range(samples): + stamp = time_from_nanoseconds(1_700_000_000_123_456_789 + index * 250_000_000) + header = Header(stamp=stamp, frame_id="world") + angle = index / 80.0 + pose = Pose( + position=Point(x=math.cos(angle), y=math.sin(angle)), + orientation=Quaternion(w=1.0), + ) + pixels = np.full((48, 64, 3), 40 + index % 180, dtype=np.uint8) + pixels[:, 16:32, 1] = 230 + image = image_from_array(pixels, encoding="rgb8", header=header) + ts = 1_700_000_000.1234567 + index / 4.0 + images.append(image, ts=ts, pose=pose) + poses.append(PoseStamped(header=header, pose=pose), ts=ts, pose=pose) + clouds.append(pointcloud_from_xyz(points, header=header), ts=ts, pose=pose) + return store diff --git a/dimos/memory/recording_policy.py b/dimos/memory/recording_policy.py new file mode 100644 index 00000000000..388a5b5e81f --- /dev/null +++ b/dimos/memory/recording_policy.py @@ -0,0 +1,24 @@ +# 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. + +"""File-existence policy shared by Python and native recorders.""" + +import enum + + +class OnExisting(str, enum.Enum): + OVERWRITE = "overwrite" + ERROR = "error" + BACKUP = "backup" + APPEND = "append" diff --git a/dimos/memory/replay_module.py b/dimos/memory/replay_module.py index b3259dc29e4..69a0358b9d9 100644 --- a/dimos/memory/replay_module.py +++ b/dimos/memory/replay_module.py @@ -72,21 +72,12 @@ def dataset_path(name: str, *, explicit: bool = True) -> str: return "" -def _wire_type(t: type) -> type: - """The class that owns *t*'s ``msg_name``: the type a port must declare so a consumer - of the base message shares its topic (unitree ``Odometry`` -> ``PoseStamped``).""" - for cls in t.__mro__: - if "msg_name" in vars(cls): - return cls - return t - - def stream_types_of(dataset: str) -> dict[str, type]: """Stream name -> port type for every stream in *dataset*.""" store = open_dataset(dataset) store.start() try: - return {n: _wire_type(t) for n, t in stream_payload_types(store).items()} + return stream_payload_types(store) finally: store.stop() @@ -233,10 +224,9 @@ def recorded_rerun_config(dataset: str) -> dict[str, Any]: def rerun_layout(stream_types: dict[str, type]) -> Any: """A Rerun blueprint showing every recorded stream: one 3D world view, one 2D view per image.""" + from dimos_generated.sensor_msgs.msg import Image import rerun.blueprint as rrb - from dimos.msgs.sensor_msgs.Image import Image - images = [ rrb.Spatial2DView(origin=f"world/{n}", name=n) for n, t in stream_types.items() diff --git a/dimos/memory/store/test_mcap.py b/dimos/memory/store/test_mcap.py index e6e3496e61b..788682ed21c 100644 --- a/dimos/memory/store/test_mcap.py +++ b/dimos/memory/store/test_mcap.py @@ -12,166 +12,131 @@ # See the License for the specific language governing permissions and # limitations under the License. -from __future__ import annotations - -from pathlib import Path - +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Vector3 +from dimos_generated.sensor_msgs.msg import CompressedImage, Imu +from dimos_generated.std_msgs.msg import Header +from mcap.writer import Writer import numpy as np import pytest -from dimos.memory.codecs.jpeg import JpegCodec -from dimos.memory.codecs.lcm import LcmCodec -from dimos.memory.codecs.lz4 import Lz4Codec +from dimos.memory.cli.dataset import open_store +from dimos.memory.codecs.cdr import CdrCodec from dimos.memory.store.mcap import McapStore -from dimos.memory.type.observation import Observation -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.sensor_msgs.Imu import Imu +from dimos.msgs.image import image_from_array, image_to_jpeg +from dimos.protocol.cdr_mcap import CdrMcapWriter -mcap_writer = pytest.importorskip("mcap.writer", reason="mcap not installed") - -def test_lcm_channel_decodes_with_explicit_codec(tmp_path: Path) -> None: +@pytest.mark.parametrize("explicit", [False, True]) +def test_cdr_channel_decodes_and_orders_by_declared_source_time(tmp_path, explicit): path = tmp_path / "recording.mcap" expected = Imu( - ts=12.5, - frame_id="imu_link", - angular_velocity=Vector3(1.0, 2.0, 3.0), + header=Header(stamp=Time(sec=12, nanosec=500000000), frame_id="imu_link"), + angular_velocity=Vector3(x=1, y=2, z=3), ) + earlier = Imu(header=Header(stamp=Time(sec=11, nanosec=500000000), frame_id="earlier")) with path.open("wb") as output: - writer = mcap_writer.Writer(output) - writer.start(profile="dimos", library="test") - channel_id = writer.register_channel( - topic="imu", - message_encoding="lcm", - schema_id=0, - metadata={ - "dimos.payload_type": "dimos.msgs.sensor_msgs.Imu.Imu", - "dimos.observation_time": "publish_time", - }, + writer = Writer(output) + writer.start(profile="ros2", library="test") + schema = writer.register_schema( + name=Imu.msg_name, encoding="ros2msg", data=Imu.schema.encode() ) - writer.add_message( - channel_id=channel_id, - log_time=13_000_000_000, - publish_time=12_500_000_000, - data=expected.lcm_encode(), - ) - writer.add_message( - channel_id=channel_id, - log_time=14_000_000_000, - publish_time=11_500_000_000, - data=Imu(ts=11.5, frame_id="earlier").lcm_encode(), + channel = writer.register_channel( + topic="imu", + message_encoding="cdr", + schema_id=schema, + metadata={"dimos.observation_time": "publish_time"}, ) + for index, (message, stamp) in enumerate([(expected, 12500000000), (earlier, 11500000000)]): + writer.add_message( + channel_id=channel, + log_time=13000000000 + index * 1000000000, + publish_time=stamp, + data=message.encode(), + ) writer.finish() - - with McapStore(path=str(path), codecs={"imu": LcmCodec(Imu)}) as store: + with McapStore(path=str(path), codecs={"imu": CdrCodec(Imu)} if explicit else None) as store: assert store.list_streams() == ["imu"] - observation: Observation[Imu] = store.stream("imu").order_by("ts").first() - assert observation.ts == 11.5 - assert observation.data.frame_id == "earlier" - - latest_observation: Observation[Imu] = store.stream("imu").order_by("ts", desc=True).first() - assert latest_observation.ts == 12.5 - assert latest_observation.data.lcm_encode() == expected.lcm_encode() - - -def test_wrapped_codec_decodes_with_explicit_codec(tmp_path: Path) -> None: - path = tmp_path / "recording.mcap" - expected = Imu( - ts=12.5, - frame_id="imu_link", - angular_velocity=Vector3(1.0, 2.0, 3.0), + observations = list(store.stream("imu").order_by("ts")) + assert [observation.ts for observation in observations] == [11.5, 12.5] + assert [observation.data for observation in observations] == [earlier, expected] + assert store.stream("imu").order_by("ts", desc=True).first().data == expected + + +def test_standard_compressed_image_uses_cdr_and_generic_dataset_dispatch(tmp_path): + path = tmp_path / "camera.mcap" + image = image_from_array( + np.full((8, 8, 3), [20, 80, 140], dtype=np.uint8), + encoding="rgb8", + header=Header(frame_id="camera", stamp=Time(sec=12, nanosec=123456789)), ) - codec = Lz4Codec(LcmCodec(Imu)) - with path.open("wb") as output: - writer = mcap_writer.Writer(output) - writer.start(profile="dimos", library="test") - channel_id = writer.register_channel( - topic="imu", - message_encoding="lz4+lcm", - schema_id=0, - metadata={ - "dimos.payload_type": "dimos.msgs.sensor_msgs.Imu.Imu", - "dimos.observation_time": "publish_time", - }, - ) - writer.add_message( - channel_id=channel_id, - log_time=13_000_000_000, - publish_time=12_500_000_000, - data=codec.encode(expected), - ) - writer.finish() - - with McapStore(path=str(path), codecs={"imu": codec}) as store: - observation: Observation[Imu] = store.stream("imu").first() - assert observation.ts == 12.5 - assert observation.data.lcm_encode() == expected.lcm_encode() - - -def test_self_describing_jpeg_channel_decodes_without_a_codec_registry(tmp_path: Path) -> None: - path = tmp_path / "recording.mcap" - expected = Image( - data=np.full((8, 8, 3), [20, 80, 140], dtype=np.uint8), - format=ImageFormat.RGB, - frame_id="camera", - ts=12.5, + expected = CompressedImage( + header=image.header, format="rgb8; jpeg compressed bgr8", data=image_to_jpeg(image) ) - codec = JpegCodec() - with path.open("wb") as output: - writer = mcap_writer.Writer(output) - writer.start(profile="dimos", library="test") - channel_id = writer.register_channel( - topic="color_image", - message_encoding="jpeg", - schema_id=0, - metadata={ - "dimos.payload_type": "dimos.msgs.sensor_msgs.Image.Image", - "dimos.observation_time": "publish_time", - }, + with CdrMcapWriter(path) as writer: + writer.write( + "color_image", + expected.encode(), + schema_name=expected.msg_name, + schema=expected.schema, + log_time_ns=13000000000, + publish_time_ns=12123456789, ) - writer.add_message( - channel_id=channel_id, - log_time=13_000_000_000, - publish_time=12_500_000_000, - data=codec.encode(expected), - ) - writer.finish() - - with McapStore(path=str(path)) as store: - observation: Observation[Image] = store.stream("color_image").first() - decoded = observation.data - assert observation.ts == 12.5 - assert decoded.frame_id == "camera" - assert decoded.format is ImageFormat.RGB - assert decoded.data.shape == expected.data.shape - assert np.mean(np.abs(decoded.data.astype(float) - expected.data.astype(float))) < 5 + with open_store(path) as store: + observation = store.stream("color_image").first() + assert observation.ts == 13.0 + assert observation.data == expected + assert observation.data.header.stamp.nanosec == 123456789 + assert observation.data.format == "rgb8; jpeg compressed bgr8" -def test_lcm_metadata_does_not_import_payload_module( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: +def test_unknown_schema_metadata_never_imports_a_payload_module(tmp_path, monkeypatch): path = tmp_path / "untrusted.mcap" with path.open("wb") as output: - writer = mcap_writer.Writer(output) - writer.start(profile="dimos", library="test") - channel_id = writer.register_channel( + writer = Writer(output) + writer.start(profile="ros2", library="test") + schema = writer.register_schema( + name="untrusted/msg/Payload", encoding="ros2msg", data=b"uint8 value\n" + ) + channel = writer.register_channel( topic="untrusted", - message_encoding="lcm", - schema_id=0, + message_encoding="cdr", + schema_id=schema, metadata={"dimos.payload_type": "untrusted_module.Payload"}, ) - writer.add_message( - channel_id=channel_id, - log_time=1, - publish_time=1, - data=b"raw payload", - ) + writer.add_message(channel_id=channel, log_time=1, publish_time=1, data=b"raw payload") writer.finish() - def fail_import(name: str) -> None: + def fail_import(name): raise AssertionError(f"artifact metadata imported {name!r}") - monkeypatch.setattr("dimos.memory.codecs.base.importlib.import_module", fail_import) + monkeypatch.setattr("dimos.memory.codecs.base.resolve_payload_type", fail_import) with McapStore(path=str(path)) as store: assert store.stream("untrusted").first().data == b"raw payload" + assert "untrusted/msg/Payload" in store.summary() + + +@pytest.mark.parametrize("second_type", [Imu, CompressedImage]) +def test_duplicate_topic_channels_count_together_or_reject_conflicting_schemas( + tmp_path, second_type +): + path = tmp_path / "channels.mcap" + with path.open("wb") as output: + writer = Writer(output) + writer.start(profile="ros2") + for cls in [Imu, second_type]: + schema = writer.register_schema( + name=cls.msg_name, encoding="ros2msg", data=cls.schema.encode() + ) + channel = writer.register_channel( + topic="sensor", message_encoding="cdr", schema_id=schema + ) + writer.add_message(channel_id=channel, log_time=1, publish_time=1, data=cls().encode()) + writer.finish() + if second_type is Imu: + with McapStore(path=str(path)) as store: + assert store.stream("sensor").count() == 2 + assert [observation.data for observation in store.stream("sensor")] == [Imu(), Imu()] + else: + with pytest.raises(ValueError, match="conflicting channel schemas"): + McapStore(path=str(path)) diff --git a/dimos/memory/stream.py b/dimos/memory/stream.py index da3aa9dc3db..6eb8f87072c 100644 --- a/dimos/memory/stream.py +++ b/dimos/memory/stream.py @@ -25,6 +25,9 @@ else: from typing_extensions import TypeVar +from dimos_generated.geometry_msgs.msg import PoseStamped, TransformStamped, Vector3 +from dimos_generated.nav_msgs.msg import Odometry + from dimos.core.resource import CompositeResource from dimos.memory.backend import Backend from dimos.memory.buffer import BackpressureBuffer, KeepLast @@ -41,7 +44,6 @@ TimeRangeFilter, ) from dimos.memory.type.observation import EmbeddedObservation, Observation -from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.utils.human import human_bytes from dimos.utils.logging_config import setup_logger @@ -242,11 +244,22 @@ def at_relative(self, t: float, tolerance: float = 1.0) -> Stream[T, O]: return self.at(t0 + t, tolerance=tolerance) def near(self, pose: Any, radius: float) -> Stream[T, O]: - # Accept Pose/PoseStamped (any object with `.position`), Vector3, - # numpy arrays, or (x, y, z) tuples — Vector3() handles the rest. + if isinstance(pose, PoseStamped): + pose = pose.pose + elif isinstance(pose, Odometry): + pose = pose.pose.pose + elif isinstance(pose, TransformStamped): + pose = pose.transform.translation if hasattr(pose, "position"): pose = pose.position - return self._with_filter(NearFilter(Vector3(pose), radius)) + if all(hasattr(pose, coordinate) for coordinate in ("x", "y", "z")): + xyz = (pose.x, pose.y, pose.z) + else: + xyz = tuple(pose) + if len(xyz) != 3: + raise ValueError("near requires exactly three position coordinates") + center = Vector3(x=float(xyz[0]), y=float(xyz[1]), z=float(xyz[2])) + return self._with_filter(NearFilter(center, radius)) def tags(self, **tags: Any) -> Stream[T, O]: return self._with_filter(TagsFilter(tags)) diff --git a/dimos/memory/test_replay_module.py b/dimos/memory/test_replay_module.py index bcfff4c96d1..356a335672a 100644 --- a/dimos/memory/test_replay_module.py +++ b/dimos/memory/test_replay_module.py @@ -16,12 +16,14 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.std_msgs.msg import Header import pytest from dimos.core.stream import Out from dimos.memory.replay_module import replay_module from dimos.memory.store.sqlite import SqliteStore -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.time import time_from_seconds, to_seconds @pytest.fixture @@ -32,8 +34,8 @@ def recording(tmp_path: Path) -> str: odom = store.stream("odom", PoseStamped) goal = store.stream("goal", PoseStamped) for ts in (1.0, 1.1, 1.2): - odom.append(PoseStamped(ts=ts), ts=ts) - goal.append(PoseStamped(ts=1.05), ts=1.05) + odom.append(PoseStamped(header=Header(stamp=time_from_seconds(ts))), ts=ts) + goal.append(PoseStamped(header=Header(stamp=time_from_seconds(1.05))), ts=1.05) store.stop() return str(path) @@ -52,8 +54,8 @@ def test_module_publishes_recorded_messages(recording: str) -> None: cls = replay_module(recording) module = cls(dataset=recording) got: dict[str, list[float]] = {"odom": [], "goal": []} - module.outputs["odom"].subscribe(lambda m: got["odom"].append(m.ts)) - module.outputs["goal"].subscribe(lambda m: got["goal"].append(m.ts)) + module.outputs["odom"].subscribe(lambda m: got["odom"].append(to_seconds(m.header.stamp))) + module.outputs["goal"].subscribe(lambda m: got["goal"].append(to_seconds(m.header.stamp))) module.start() deadline = time.time() + 5 while time.time() < deadline and len(got["odom"]) < 3: @@ -91,8 +93,12 @@ def test_stream_named_like_a_module_attribute_is_skipped( path = tmp_path / "memory.db" store = SqliteStore(path=str(path)) store.start() - store.stream("start", PoseStamped).append(PoseStamped(ts=1.0), ts=1.0) - store.stream("odom", PoseStamped).append(PoseStamped(ts=1.0), ts=1.0) + store.stream("start", PoseStamped).append( + PoseStamped(header=Header(stamp=time_from_seconds(1.0))), ts=1.0 + ) + store.stream("odom", PoseStamped).append( + PoseStamped(header=Header(stamp=time_from_seconds(1.0))), ts=1.0 + ) store.stop() assert list(replay_module(str(path)).__annotations__) == ["odom"] module = replay_module("")(dataset=str(path)) @@ -101,19 +107,6 @@ def test_stream_named_like_a_module_attribute_is_skipped( assert [c.args[1] for c in log.warning.call_args_list] == ["start", "start"] -def test_subclass_payload_gets_the_base_port_type(tmp_path: Path) -> None: - """A recorded unitree Odometry is a PoseStamped on the wire; the port must say so or a - PoseStamped consumer of the same name lands on a different topic.""" - from dimos.robot.unitree.type.odometry import Odometry - - path = tmp_path / "memory.db" - store = SqliteStore(path=str(path)) - store.start() - store.stream("odom", Odometry).append(Odometry(ts=1.0), ts=1.0) - store.stop() - assert replay_module(str(path)).__annotations__ == {"odom": Out[PoseStamped]} - - def test_single_timestamp_stream_is_republished_and_not_the_anchor(tmp_path: Path) -> None: """camera_info stamped at import sits 39 s before the run; the run must start at odom.""" path = tmp_path / "memory.db" @@ -121,8 +114,10 @@ def test_single_timestamp_stream_is_republished_and_not_the_anchor(tmp_path: Pat store.start() info = store.stream("camera_info", PoseStamped) for _ in range(3): - info.append(PoseStamped(ts=1.0), ts=1.0) - store.stream("odom", PoseStamped).append(PoseStamped(ts=40.0), ts=40.0) + info.append(PoseStamped(header=Header(stamp=time_from_seconds(1.0))), ts=1.0) + store.stream("odom", PoseStamped).append( + PoseStamped(header=Header(stamp=time_from_seconds(40.0))), ts=40.0 + ) store.stop() module = replay_module(str(path))(dataset=str(path)) got: dict[str, int] = {"camera_info": 0, "odom": 0} @@ -160,8 +155,8 @@ def slow(self: Any, obs: Any) -> Any: monkeypatch.setattr(replay_mod.ReplayStream, "_decode", slow) module = replay_module(recording)(dataset=recording) got: dict[str, list[float]] = {"odom": [], "goal": []} - module.outputs["odom"].subscribe(lambda m: got["odom"].append(m.ts)) - module.outputs["goal"].subscribe(lambda m: got["goal"].append(m.ts)) + module.outputs["odom"].subscribe(lambda m: got["odom"].append(to_seconds(m.header.stamp))) + module.outputs["goal"].subscribe(lambda m: got["goal"].append(to_seconds(m.header.stamp))) module.start() deadline = time.time() + 5 while time.time() < deadline and (len(got["odom"]) < 3 or not got["goal"]): diff --git a/dimos/memory/test_voxel_map.py b/dimos/memory/test_voxel_map.py index ae9adc3bfcf..f5c60a20876 100644 --- a/dimos/memory/test_voxel_map.py +++ b/dimos/memory/test_voxel_map.py @@ -14,15 +14,24 @@ from __future__ import annotations +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header import numpy as np from dimos.mapping.voxels.module import VoxelMapTransformer from dimos.memory.type.observation import Observation -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import time_from_seconds def _make_obs(obs_id: int, points: np.ndarray, ts: float = 0.0) -> Observation[PointCloud2]: - return Observation(id=obs_id, ts=ts, _data=PointCloud2.from_numpy(points)) + return Observation( + id=obs_id, + ts=ts, + _data=pointcloud_from_xyz( + points, header=Header(frame_id="world", stamp=time_from_seconds(ts)) + ), + ) def _unit_cube_points(n: int = 100) -> np.ndarray: @@ -36,18 +45,21 @@ def test_accumulate_two_frames() -> None: obs1 = _make_obs(0, pts, ts=1.0) obs2 = _make_obs(1, pts + 10.0, ts=2.0) # offset by 10m, no overlap - xf = VoxelMapTransformer(voxel_size=0.5, carve_columns=False) + xf = VoxelMapTransformer(voxel_size=0.5, carve_columns=False, device="CPU:0") results = list(xf(iter([obs1, obs2]))) assert len(results) == 2 # emit_every=1 default global_map = results[-1].data # last result has the full accumulated map - single_results = list(VoxelMapTransformer(voxel_size=0.5)(iter([obs1]))) - assert len(global_map) > len(single_results[0].data) + single_results = list(VoxelMapTransformer(voxel_size=0.5, device="CPU:0")(iter([obs1]))) + assert ( + global_map.width * global_map.height + > single_results[0].data.width * single_results[0].data.height + ) def test_empty_stream() -> None: - xf = VoxelMapTransformer(voxel_size=0.5) + xf = VoxelMapTransformer(voxel_size=0.5, device="CPU:0") assert list(xf(iter([]))) == [] diff --git a/dimos/memory/tf.py b/dimos/memory/tf.py index 53e720a9125..036a16b6c22 100644 --- a/dimos/memory/tf.py +++ b/dimos/memory/tf.py @@ -19,12 +19,14 @@ import math from typing import TYPE_CHECKING, Any, cast +from dimos_generated.tf2_msgs.msg import TFMessage + from dimos.memory.stream import Stream -from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.protocol.tf.tf import MultiTBuffer if TYPE_CHECKING: - from dimos.msgs.geometry_msgs.Transform import Transform + from dimos_generated.geometry_msgs.msg import TransformStamped + from dimos.protocol.tf.tf import TFLookup @@ -58,7 +60,7 @@ def from_store(cls, store: Any, stream: str = "tf") -> StreamTF | None: recorded = tf_stream(store, stream) return None if recorded is None else cls(recorded) - def publish(self, *args: Transform) -> None: + def publish(self, *args: TransformStamped) -> None: raise NotImplementedError("StreamTF is a read-only replay service.") def _load(self, lo: float, hi: float) -> None: @@ -86,7 +88,7 @@ def get( *, forward_tolerance: float = 0.0, warn: bool = True, - ) -> Transform | None: + ) -> TransformStamped | None: tp = time_point if tp is None: last = next(iter(self.stream.order_by("ts", desc=True).limit(1)), None) diff --git a/dimos/memory/tool_e2e.py b/dimos/memory/tool_e2e.py index a3d3ffa6767..42a8f5e80f8 100644 --- a/dimos/memory/tool_e2e.py +++ b/dimos/memory/tool_e2e.py @@ -12,46 +12,69 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""E2E test: import legacy pickle replays into memory SqliteStore.""" +"""E2E test: copy a deterministic CDR recording into a fresh SQLite store.""" from __future__ import annotations import bisect from typing import TYPE_CHECKING, Any +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from dimos.memory.embed import EmbedImages from dimos.memory.store.sqlite import SqliteStore +from dimos.memory.stream import Stream from dimos.memory.transform import QualityWindow from dimos.models.embedding.clip import CLIPModel -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.robot.unitree.type.odometry import Odometry -from dimos.utils.data import get_data_dir -from dimos.utils.testing.legacy_pickle import LegacyPickleStore +from dimos.msgs.image import image_brightness, image_from_array, image_sharpness +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import time_from_nanoseconds if TYPE_CHECKING: from collections.abc import Iterator -DB_PATH = get_data_dir() / "go2_short.db" + +@pytest.fixture(scope="module") +def source(tmp_path_factory: pytest.TempPathFactory) -> Iterator[SqliteStore]: + path = tmp_path_factory.mktemp("cdr-source") / "recording.db" + with SqliteStore(path=str(path)) as store: + video = store.stream("color_image", Image) + lidar = store.stream("lidar", PointCloud2) + odom = store.stream("odom", PoseStamped) + for index in range(1024): + ns = 1_700_000_000_123_456_789 + index * 25_000_000 + header = Header(stamp=time_from_nanoseconds(ns), frame_id="world") + pose = PoseStamped(header=header, pose=Pose(position=Point(x=index / 100))) + pixels = np.full((8, 8, 3), 64 + index % 128, dtype=np.uint8) + pixels[:, index % 8] = (255, 0, 127) + video.append(image_from_array(pixels, encoding="rgb8", header=header), ts=ns / 1e9) + lidar.append( + pointcloud_from_xyz(np.array([[1, 0, 0], [1, 1, 0], [2, 0, 0]]), header=header), + ts=ns / 1e9, + ) + odom.append(pose, ts=ns / 1e9) + yield store @pytest.fixture(scope="module") -def session() -> Iterator[SqliteStore]: - store = SqliteStore(path=str(DB_PATH)) - with store: +def session(tmp_path_factory: pytest.TempPathFactory) -> Iterator[SqliteStore]: + path = tmp_path_factory.mktemp("cdr-output") / "recording.db" + with SqliteStore(path=str(path)) as store: yield store - store.stop() class PoseIndex: """Preloaded odom data with O(log n) closest-timestamp lookup.""" - def __init__(self, replay: LegacyPickleStore[Any]) -> None: + def __init__(self, replay: Stream[Any]) -> None: self._timestamps: list[float] = [] self._data: list[Any] = [] - for ts, data in replay.iterate_ts(): + for obs in replay: + ts, data = obs.ts, obs.data self._timestamps.append(ts) self._data.append(data) @@ -73,22 +96,22 @@ def __iter__(self) -> Iterator[tuple[float, Any]]: @pytest.fixture(scope="module") -def video_replay() -> LegacyPickleStore[Image]: - return LegacyPickleStore("go2_short/video") +def video_replay(source: SqliteStore) -> Stream[Image]: + return source.stream("color_image", Image) @pytest.fixture(scope="module") -def odom_index() -> PoseIndex: - return PoseIndex(LegacyPickleStore("go2_short/odom")) +def odom_index(source: SqliteStore) -> PoseIndex: + return PoseIndex(source.stream("odom", PoseStamped)) @pytest.fixture(scope="module") -def lidar_replay() -> LegacyPickleStore[PointCloud2]: - return LegacyPickleStore("go2_short/lidar") +def lidar_replay(source: SqliteStore) -> Stream[PointCloud2]: + return source.stream("lidar", PointCloud2) class TestImportReplay: - """Import legacy pickle replay data into a memory SqliteStore. + """Copy generated CDR streams into a fresh memory SqliteStore. Lidar/odom are trimmed to start at video's first ts. The MemoryReplayAdapter scheduler anchors each stream to its own first_ts on subscribe, so aligning @@ -99,10 +122,10 @@ def test_import_odom( self, session: SqliteStore, odom_index: PoseIndex, - video_replay: LegacyPickleStore[Any], + video_replay: Stream[Any], ) -> None: - threshold = video_replay.first_timestamp() - with session.stream("odom", Odometry) as odom: + threshold = video_replay.first().ts + with session.stream("odom", PoseStamped) as odom: count = 0 skipped = 0 for ts, data in odom_index: @@ -119,12 +142,13 @@ def test_import_odom( def test_import_video( self, session: SqliteStore, - video_replay: LegacyPickleStore[Any], + video_replay: Stream[Any], odom_index: PoseIndex, ) -> None: with session.stream("color_image", Image) as video: count = 0 - for ts, frame in video_replay.iterate_ts(): + for obs in video_replay: + ts, frame = obs.ts, obs.data pose = odom_index.find_closest(ts) video.append(frame, ts=ts, pose=pose) count += 1 @@ -137,16 +161,17 @@ def test_import_video( def test_import_lidar( self, session: SqliteStore, - lidar_replay: LegacyPickleStore[Any], + lidar_replay: Stream[Any], odom_index: PoseIndex, - video_replay: LegacyPickleStore[Any], + video_replay: Stream[Any], ) -> None: - threshold = video_replay.first_timestamp() - lidar = session.stream("lidar", PointCloud2, codec="lz4+lcm") + threshold = video_replay.first().ts + lidar = session.stream("lidar", PointCloud2, codec="lz4+cdr") count = 0 skipped = 0 - for ts, frame in lidar_replay.iterate_ts(): + for obs in lidar_replay: + ts, frame = obs.ts, obs.data if ts < threshold: skipped += 1 continue @@ -173,9 +198,9 @@ def test_embed_and_save(self, session: SqliteStore, clip: CLIPModel) -> None: # Downsample to 2Hz, then embed pipeline = ( - video.filter(lambda obs: obs.data.brightness > 0.1) + video.filter(lambda obs: image_brightness(obs.data) > 0.1) .tap(print) - .transform(QualityWindow(lambda img: img.sharpness, window=0.5)) + .transform(QualityWindow(image_sharpness, window=0.5)) .transform(EmbedImages(clip)) .save(embedded) ) @@ -209,7 +234,7 @@ def test_query_imported_data(self, session: SqliteStore) -> None: class TestE2EQuery: - """Query operations against real robot replay data.""" + """Query operations against a deterministic generated CDR recording.""" def test_list_streams(self, session: SqliteStore) -> None: streams = session.list_streams() diff --git a/dimos/memory/tool_visualizer.py b/dimos/memory/tool_visualizer.py index 399abb71216..174955e153d 100644 --- a/dimos/memory/tool_visualizer.py +++ b/dimos/memory/tool_visualizer.py @@ -19,6 +19,7 @@ import pickle from typing import TYPE_CHECKING +from dimos_generated.sensor_msgs.msg import Image import pytest from dimos.memory.store.sqlite import SqliteStore @@ -26,8 +27,9 @@ from dimos.models.embedding.clip import CLIPModel from dimos.models.vl.florence import Florence2Model from dimos.models.vl.moondream import MoondreamVlModel -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.geometry import inverse_transform, transform_from_pose +from dimos.msgs.image import image_sharpness +from dimos.msgs.time import to_seconds from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC from dimos.robot.unitree.go2.connection import GO2Connection from dimos.utils.data import get_data, get_data_dir @@ -113,11 +115,10 @@ def test_detect_objects(self, store: SqliteStore, clip: CLIPModel) -> None: embedded = store.streams.color_image_embedded lidar = store.streams.lidar - for obs in embedded.search(clip.embed_text("bottle"), k=1).map( - lambda obs: obs.derive(data=vlm.query_detections(obs.data, "bottle")) - ): + for obs in embedded.search(clip.embed_text("bottle"), k=1): + detections = vlm.query_detections(obs.data, "bottle") print(f"ts={obs.ts:.2f} sim={obs.similarity:.3f} pose={obs.pose}") - for det in obs.data.detections: + for det in detections.detections: print(det) print( lidar.at(obs.ts).first().data @@ -140,7 +141,7 @@ def test_agent_visual_description_passive(self, store: SqliteStore) -> None: pipeline = ( store.streams.color_image.limit(200) .transform( - QualityWindow(lambda img: img.sharpness, window=5.0) + QualityWindow(image_sharpness, window=5.0) # we are batch processing images here, # so we can use the more efficient batch captioning API # (instead of using .map() and calling caption() for each image, @@ -167,20 +168,26 @@ def test_detect_objects_smart(self, store: SqliteStore, clip: CLIPModel) -> None # find a location in the world with highest semantic similarity to a bottle bottle_pos = embedded.search(clip.embed_text("bottle"), k=1).first().pose_stamped + if bottle_pos is None: + raise ValueError("The semantic hotspot has no pose for spatial detection") for obs in ( store.streams.color_image # find all frames within 60 seconds of the semantic hotspot - .at(bottle_pos.ts, tolerance=60.0) + .at(to_seconds(bottle_pos.header.stamp), tolerance=60.0) # filter the frames within 1m radius near the semantic hotspot .near(bottle_pos, radius=1.0) # select highest quality frames from these results (based on sharpness) - .transform(QualityWindow(lambda img: img.sharpness, window=1.0)) + .transform(QualityWindow(image_sharpness, window=1.0)) # run detection on these frames to find bottles .map(lambda obs: obs.derive(data=vlm.query_detections(obs.data, "bottle"))) ): print(f"ts={obs.ts:.2f} pose={obs.pose_stamped}") + image_pose = obs.pose_stamped + if image_pose is None: + continue + # find the lidar frame captured closest in time to an image lidar_frame = lidar.at(obs.ts).first().data @@ -192,10 +199,11 @@ def test_detect_objects_smart(self, store: SqliteStore, clip: CLIPModel) -> None det, lidar_frame, camera_info=GO2Connection.camera_info_static, - world_to_optical_transform=Transform( - ts=obs.ts, - translation=obs.pose_stamped.position, - rotation=obs.pose_stamped.orientation, - ).inverse(), + world_to_optical_transform=inverse_transform( + transform_from_pose( + image_pose, + child_frame_id=GO2Connection.camera_info_static.header.frame_id, + ) + ), ) print(det3d) diff --git a/dimos/memory/tool_voxel_map.py b/dimos/memory/tool_voxel_map.py index 90b992be103..c1bdc8edf2e 100644 --- a/dimos/memory/tool_voxel_map.py +++ b/dimos/memory/tool_voxel_map.py @@ -17,11 +17,11 @@ import time from typing import TYPE_CHECKING +from dimos_generated.sensor_msgs.msg import PointCloud2 import pytest from dimos.mapping.voxels.module import VoxelMapTransformer from dimos.memory.store.sqlite import SqliteStore -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.utils.data import get_data if TYPE_CHECKING: @@ -51,11 +51,11 @@ def test_build_global_map(store: SqliteStore) -> None: frame_count = result.tags["frame_count"] assert frame_count == n_frames - assert len(global_map) > 0 + assert global_map.width * global_map.height > 0 print( lidar.summary(), - f"\n{frame_count} frames -> {len(global_map)} voxels" + f"\n{frame_count} frames -> {global_map.width * global_map.height} voxels" f"\n transform: {t_transform:.2f}s ({t_transform / frame_count * 1000:.1f}ms/frame)" f"\n total wall: {t_total:.2f}s", ) diff --git a/dimos/memory/transform.py b/dimos/memory/transform.py index 2f23ec7859c..abcef5fe5b2 100644 --- a/dimos/memory/transform.py +++ b/dimos/memory/transform.py @@ -22,6 +22,7 @@ from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar, cast from dimos.memory.utils.formatting import FilterRepr +from dimos.msgs.geometry import point_distance, quaternion_angle if TYPE_CHECKING: from collections.abc import Callable, Iterator, Sequence @@ -183,7 +184,7 @@ def _speed(upstream: Iterator[Observation[Any]]) -> Iterator[Observation[float]] pp = prev.pose if prev is not None else None if prev is not None and p is not None and pp is not None: dt = obs.ts - prev.ts - v = (p.position - pp.position).length() / dt if dt > 0 else 0.0 + v = point_distance(p.position, pp.position) / dt if dt > 0 else 0.0 yield obs.derive(data=v) prev = obs @@ -231,10 +232,10 @@ def __call__(self, upstream: Iterator[Observation[T]]) -> Iterator[Observation[T if prev is not None and pp is not None: dt = obs.ts - prev.ts if dt > 0: - v = (p.position - pp.position).length() / dt + v = point_distance(p.position, pp.position) / dt ok = v <= max_mps if ok and max_rps is not None: - ok = (p.orientation.angle_to(pp.orientation) / dt) <= max_rps + ok = (quaternion_angle(p.orientation, pp.orientation) / dt) <= max_rps if ok: yield obs prev = obs diff --git a/dimos/memory/type/filter.py b/dimos/memory/type/filter.py index 37229571f5c..01ce60d2da6 100644 --- a/dimos/memory/type/filter.py +++ b/dimos/memory/type/filter.py @@ -19,7 +19,7 @@ from itertools import islice from typing import TYPE_CHECKING, Any -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos_generated.geometry_msgs.msg import Vector3 if TYPE_CHECKING: from collections.abc import Callable, Iterator @@ -84,8 +84,13 @@ def matches(self, obs: Observation[Any]) -> bool: p = obs.pose_tuple if p is None: return False - delta = self.position - Vector3(p[0], p[1], p[2]) - return delta.length_squared() <= self.radius * self.radius + distance_squared = sum( + (value - center) ** 2 + for value, center in zip( + p[:3], (self.position.x, self.position.y, self.position.z), strict=True + ) + ) + return distance_squared <= self.radius * self.radius @dataclass(frozen=True) diff --git a/dimos/memory/type/observation.py b/dimos/memory/type/observation.py index 3a42b4c80d0..ee3808f94ab 100644 --- a/dimos/memory/type/observation.py +++ b/dimos/memory/type/observation.py @@ -19,8 +19,10 @@ import threading from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, TransformStamped +from dimos_generated.std_msgs.msg import Header + +from dimos.msgs.time import time_from_seconds if sys.version_info >= (3, 11): from typing import Self @@ -65,9 +67,8 @@ def _to_tuple(p: Any) -> PoseTuple | None: Accepts: `None`, a 3- or 7-tuple, or any object with `translation`+`rotation` (Transform) or `position`+`orientation` - (Pose / PoseStamped) attributes. 3-tuples are padded with the identity - quaternion. Duck-typed to avoid importing `dimos.msgs.geometry_msgs` - at module load. + (Pose) attributes, or generated PoseStamped/TransformStamped values. + 3-tuples are padded with the identity quaternion. """ if p is None: return None @@ -78,6 +79,10 @@ def _to_tuple(p: Any) -> PoseTuple | None: x, y, z = p return (float(x), float(y), float(z), 0.0, 0.0, 0.0, 1.0) raise TypeError(f"Pose tuple must have length 3 or 7, got {len(p)}") + if isinstance(p, PoseStamped): + p = p.pose + elif isinstance(p, TransformStamped): + p = p.transform # Use ``is not None`` rather than ``or`` — a Vector3 at the origin is # falsy but valid, and falling through to ``position`` on a Transform # would crash. @@ -156,17 +161,21 @@ def pose(self) -> Pose | None: if self.pose_tuple is None: return None - return Pose(*self.pose_tuple) + x, y, z, qx, qy, qz, qw = self.pose_tuple + return Pose(position=Point(x=x, y=y, z=z), orientation=Quaternion(x=qx, y=qy, z=qz, w=qw)) @property def pose_stamped(self) -> PoseStamped | None: """Typed :class:`PoseStamped` (or None) carrying this observation's ts.""" if self.pose_tuple is None: return None - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - x, y, z, qx, qy, qz, qw = self.pose_tuple - return PoseStamped(ts=self.ts, position=(x, y, z), orientation=(qx, qy, qz, qw)) + return PoseStamped( + header=Header(stamp=time_from_seconds(self.ts)), + pose=Pose( + position=Point(x=x, y=y, z=z), orientation=Quaternion(x=qx, y=qy, z=qz, w=qw) + ), + ) @property def data(self) -> T: diff --git a/dimos/memory/type/test_generated_query.py b/dimos/memory/type/test_generated_query.py new file mode 100644 index 00000000000..50a57d8a547 --- /dev/null +++ b/dimos/memory/type/test_generated_query.py @@ -0,0 +1,54 @@ +# 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. + +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + PoseWithCovariance, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.nav_msgs.msg import Odometry +import numpy as np +import pytest + +from dimos.memory.store.memory import MemoryStore + + +@pytest.mark.parametrize( + "center", + [ + Point(x=1), + Vector3(x=1), + (1, 0, 0), + np.array([1, 0, 0]), + Pose(position=Point(x=1)), + PoseStamped(pose=Pose(position=Point(x=1))), + TransformStamped(transform=Transform(translation=Vector3(x=1))), + Odometry(pose=PoseWithCovariance(pose=Pose(position=Point(x=1)))), + ], +) +def test_near_accepts_generated_nested_positions_and_retains_radius_boundary(center): + with MemoryStore() as store: + stream = store.stream("values", str) + stream.append("center", ts=1, pose=(1, 0, 0)) + stream.append("boundary", ts=2, pose=(2, 0, 0)) + stream.append("outside", ts=3, pose=(2.001, 0, 0)) + stream.append("poseless", ts=4) + assert [obs.data for obs in stream.near(center, radius=1)] == ["center", "boundary"] diff --git a/dimos/memory/vis/space/elements.py b/dimos/memory/vis/space/elements.py index 2c4f052f639..7b7d84f8ecf 100644 --- a/dimos/memory/vis/space/elements.py +++ b/dimos/memory/vis/space/elements.py @@ -32,16 +32,16 @@ ColorLike = Union[str, Color, DeferredColor] if TYPE_CHECKING: + from dimos_generated.geometry_msgs.msg import ( + Point as GeoPoint, + Pose as GeoPose, + PoseStamped, + Vector3, + ) + from dimos_generated.nav_msgs.msg import OccupancyGrid, Path + from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 + from dimos.memory.type.observation import Observation - from dimos.msgs.geometry_msgs.Point import Point as GeoPoint - from dimos.msgs.geometry_msgs.Pose import Pose as GeoPose - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - from dimos.msgs.geometry_msgs.Vector3 import Vector3 - from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid - from dimos.msgs.nav_msgs.Path import Path - from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo - from dimos.msgs.sensor_msgs.Image import Image - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 @dataclass @@ -83,7 +83,7 @@ class Point: Rerun: rr.Points3D """ - msg: GeoPoint | GeoPose + msg: GeoPoint | GeoPose | PoseStamped color: ColorLike = "#e74c3c" radius: float = 0.05 label: str | None = None diff --git a/dimos/memory/vis/space/geometry.py b/dimos/memory/vis/space/geometry.py new file mode 100644 index 00000000000..9db4a82378c --- /dev/null +++ b/dimos/memory/vis/space/geometry.py @@ -0,0 +1,32 @@ +# 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. + +"""Geometry access shared by SVG and Rerun scene renderers.""" + +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped + +from dimos.msgs.geometry import yaw + + +def message_position(message: Point | Pose | PoseStamped) -> Point: + if isinstance(message, PoseStamped): + return message.pose.position + if isinstance(message, Pose): + return message.position + return message + + +def message_yaw(message: Pose | PoseStamped) -> float: + pose = message.pose if isinstance(message, PoseStamped) else message + return yaw(pose.orientation) diff --git a/dimos/memory/vis/space/rerun.py b/dimos/memory/vis/space/rerun.py index 0242774dc3d..643b6f363c7 100644 --- a/dimos/memory/vis/space/rerun.py +++ b/dimos/memory/vis/space/rerun.py @@ -19,14 +19,26 @@ import math from typing import TYPE_CHECKING, Any +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.nav_msgs.msg import OccupancyGrid +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header +import numpy as np + +from dimos.mapping.occupancy.visualizations import generate_rgba_texture from dimos.memory.type.observation import Observation from dimos.memory.vis.color import Color from dimos.memory.vis.space.elements import Arrow, Box3D, Camera, Point, Polyline, Pose, Text -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.memory.vis.space.geometry import message_position, message_yaw +from dimos.msgs.geometry import ( + compose_transforms, + pose_matrix, + transform_from_pose, + transform_matrix, +) +from dimos.msgs.occupancy import occupancy_extent, occupancy_view +from dimos.msgs.pointcloud import pointcloud_rgb, pointcloud_xyz +from dimos.visualization.rerun.message_helpers import camera_pinhole, image_archetype if TYPE_CHECKING: from dimos.memory.vis.space.space import Space @@ -40,16 +52,13 @@ def _rgba(el: Any) -> tuple[int, int, int, int]: # base_link → camera_optical extrinsics (applied at render time for image observations) -_BASE_TO_OPTICAL = Transform( - translation=Vector3(0.3, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="base_link", - child_frame_id="camera_link", -) + Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(-0.5, 0.5, -0.5, 0.5), - frame_id="camera_link", +_BASE_TO_OPTICAL = TransformStamped( + header=Header(frame_id="base_link"), child_frame_id="camera_optical", + transform=Transform( + translation=Vector3(x=0.3), + rotation=Quaternion(x=-0.5, y=0.5, z=-0.5, w=0.5), + ), ) @@ -127,17 +136,44 @@ def render(space: Space, app_id: str = "space", spawn: bool = True) -> None: # Log elements if grids: for i, el in enumerate(grids): - rr.log(f"scene/map/{i}", el.to_rerun(), static=True) + cells = occupancy_view(el) + if cells.size == 0: + continue + width, height = occupancy_extent(el) + matrix = pose_matrix(el.info.origin) + corners = np.array([[0, 0, 0], [width, 0, 0], [width, height, 0], [0, height, 0]]) + vertices = corners @ matrix[:3, :3].T + matrix[:3, 3] + rr.log( + f"scene/map/{i}", + rr.Mesh3D( + vertex_positions=vertices, + triangle_indices=[[0, 1, 2], [0, 2, 3]], + vertex_texcoords=[[0, 1], [1, 1], [1, 0], [0, 0]], + albedo_texture=np.ascontiguousarray(np.flipud(generate_rgba_texture(cells))), + ), + static=True, + ) if pointclouds: for i, el in enumerate(pointclouds): - rr.log(f"scene/pointcloud/{i}", el.to_rerun(), static=True) + rr.log( + f"scene/pointcloud/{i}", + rr.Points3D(pointcloud_xyz(el), colors=pointcloud_rgb(el)), + static=True, + ) if points: rr.log( "scene/points", rr.Points3D( - positions=[[p.msg.x, p.msg.y, p.msg.z] for p in points], + positions=[ + [ + message_position(p.msg).x, + message_position(p.msg).y, + message_position(p.msg).z, + ] + for p in points + ], colors=[_rgba(p) for p in points], radii=[max(p.radius, 0.05) for p in points], labels=[p.label or "" for p in points] if any(p.label for p in points) else None, @@ -149,7 +185,9 @@ def render(space: Space, app_id: str = "space", spawn: bool = True) -> None: rr.log( "scene/poses", rr.Points3D( - positions=[[p.msg.x, p.msg.y, 0] for p in poses], + positions=[ + [message_position(p.msg).x, message_position(p.msg).y, 0] for p in poses + ], colors=[_rgba(p) for p in poses], radii=[p.size * 0.3 for p in poses], labels=[p.label or "" for p in poses] if any(p.label for p in poses) else None, @@ -159,9 +197,14 @@ def render(space: Space, app_id: str = "space", spawn: bool = True) -> None: rr.log( "scene/poses/headings", rr.Arrows3D( - origins=[[p.msg.x, p.msg.y, 0] for p in poses], + origins=[[message_position(p.msg).x, message_position(p.msg).y, 0] for p in poses], vectors=[ - [math.cos(p.msg.yaw) * p.size, math.sin(p.msg.yaw) * p.size, 0] for p in poses + [ + math.cos(message_yaw(p.msg)) * p.size, + math.sin(message_yaw(p.msg)) * p.size, + 0, + ] + for p in poses ], colors=[_rgba(p) for p in poses], ), @@ -172,9 +215,13 @@ def render(space: Space, app_id: str = "space", spawn: bool = True) -> None: rr.log( "scene/arrows", rr.Arrows3D( - origins=[[a.msg.x, a.msg.y, 0] for a in arrows], + origins=[[message_position(a.msg).x, message_position(a.msg).y, 0] for a in arrows], vectors=[ - [math.cos(a.msg.yaw) * a.length, math.sin(a.msg.yaw) * a.length, 0] + [ + math.cos(message_yaw(a.msg)) * a.length, + math.sin(message_yaw(a.msg)) * a.length, + 0, + ] for a in arrows ], colors=[_rgba(a) for a in arrows], @@ -186,7 +233,7 @@ def render(space: Space, app_id: str = "space", spawn: bool = True) -> None: rr.log( "scene/boxes", rr.Boxes3D( - centers=[[b.center.x, b.center.y, 0] for b in boxes], + centers=[[b.center.position.x, b.center.position.y, 0] for b in boxes], half_sizes=[[b.size.x / 2, b.size.y / 2, b.size.z / 2] for b in boxes], colors=[_rgba(b) for b in boxes], labels=[b.label or "" for b in boxes] if any(b.label for b in boxes) else None, @@ -198,7 +245,7 @@ def render(space: Space, app_id: str = "space", spawn: bool = True) -> None: rr.log( f"scene/polylines/{i}", rr.LineStrips3D( - strips=[[[p.x, p.y, 0] for p in el.msg.poses]], + strips=[[[p.pose.position.x, p.pose.position.y, 0] for p in el.msg.poses]], colors=[_rgba(el)], radii=[el.width / 2], ), @@ -219,13 +266,14 @@ def render(space: Space, app_id: str = "space", spawn: bool = True) -> None: for i, el in enumerate(cameras): path = f"scene/cameras/{i}" - rr.log(path, el.pose.to_rerun(), static=True) + matrix = pose_matrix(el.pose.pose) + rr.log(path, rr.Transform3D(translation=matrix[:3, 3], mat3x3=matrix[:3, :3]), static=True) if el.camera_info: - pinhole = el.camera_info.to_rerun() + pinhole = camera_pinhole(el.camera_info) assert not isinstance(pinhole, list) rr.log(path, pinhole, static=True) elif el.image: - h, w = el.image.shape[:2] + h, w = el.image.height, el.image.width focal = max(w, h) rr.log( path, @@ -233,7 +281,7 @@ def render(space: Space, app_id: str = "space", spawn: bool = True) -> None: static=True, ) if el.image: - rr.log(f"{path}/image", el.image.to_rerun(), static=True) + rr.log(f"{path}/image", image_archetype(el.image), static=True) for i, obs in enumerate(observations): path = f"scene/observations/{i}" @@ -244,9 +292,14 @@ def render(space: Space, app_id: str = "space", spawn: bool = True) -> None: img = _as_image(data) if img is not None: # Apply base→optical extrinsics for camera frustum rendering - world_T_optical = Transform.from_pose("world", ps) + _BASE_TO_OPTICAL - rr.log(path, world_T_optical.to_pose().to_rerun(), static=True) - h, w = img.shape[:2] + world_T_optical = compose_transforms( + transform_from_pose(ps, child_frame_id="base_link"), _BASE_TO_OPTICAL + ) + matrix = transform_matrix(world_T_optical.transform) + rr.log( + path, rr.Transform3D(translation=matrix[:3, 3], mat3x3=matrix[:3, :3]), static=True + ) + h, w = img.height, img.width focal = max(w, h) rr.log( path, @@ -258,15 +311,22 @@ def render(space: Space, app_id: str = "space", spawn: bool = True) -> None: ), static=True, ) - rr.log(f"{path}/image", img.to_rerun(), static=True) + rr.log(f"{path}/image", image_archetype(img), static=True) elif isinstance(data, PointCloud2): - rr.log(path, ps.to_rerun(), static=True) - rr.log(f"{path}/pointcloud", data.to_rerun(), static=True) + matrix = pose_matrix(ps.pose) + rr.log( + path, rr.Transform3D(translation=matrix[:3, 3], mat3x3=matrix[:3, :3]), static=True + ) + rr.log( + f"{path}/pointcloud", + rr.Points3D(pointcloud_xyz(data), colors=pointcloud_rgb(data)), + static=True, + ) elif isinstance(data, (int, float)): rr.log( path, rr.Points3D( - positions=[[ps.x, ps.y, 0]], + positions=[[ps.pose.position.x, ps.pose.position.y, 0]], labels=[str(data)], radii=[0.025], ), @@ -286,7 +346,7 @@ def render(space: Space, app_id: str = "space", spawn: bool = True) -> None: if line: lines.append(line) label = "\n".join(lines) - x, y = ps.x, ps.y + x, y = ps.pose.position.x, ps.pose.position.y # Pin: line from ground up, label at the tip rr.log( f"{path}/pin", @@ -309,19 +369,18 @@ def render(space: Space, app_id: str = "space", spawn: bool = True) -> None: else: rr.log( path, - rr.Points3D(positions=[[ps.x, ps.y, 0]], radii=[0.05]), + rr.Points3D(positions=[[ps.pose.position.x, ps.pose.position.y, 0]], radii=[0.05]), static=True, ) for i, obs in enumerate(panels): img = _as_image(obs.data) if img is not None: - rr.log(f"scene/panels/{i}", img.to_rerun(), static=True) + rr.log(f"scene/panels/{i}", image_archetype(img), static=True) def _as_image(data: Any) -> Any | None: """Return an Image if data is an Image or ImageDetections, else None.""" - from dimos.msgs.sensor_msgs.Image import Image from dimos.perception.detection.type.imageDetections import ImageDetections if isinstance(data, Image): diff --git a/dimos/memory/vis/space/space.py b/dimos/memory/vis/space/space.py index 65afe7d0efc..864da7169a2 100644 --- a/dimos/memory/vis/space/space.py +++ b/dimos/memory/vis/space/space.py @@ -23,6 +23,11 @@ from typing import Any +from dimos_generated.geometry_msgs.msg import Point as GeoPoint, Pose as GeoPose, PoseStamped +from dimos_generated.nav_msgs.msg import OccupancyGrid, Path as NavPath +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.vision_msgs.msg import Detection3D + from dimos.memory.type.observation import EmbeddedObservation, Observation from dimos.memory.vis.color import ColorRange, resolve_deferred from dimos.memory.vis.space.elements import ( @@ -35,14 +40,7 @@ SpaceElement, Text, ) -from dimos.msgs.geometry_msgs.Point import Point as GeoPoint -from dimos.msgs.geometry_msgs.Pose import Pose as GeoPose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.nav_msgs.Path import Path as NavPath from dimos.msgs.protocol import DimosMsg -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.vision_msgs.Detection3D import Detection3D def _autocolor_value(item: Any) -> float | None: diff --git a/dimos/memory/vis/space/svg.py b/dimos/memory/vis/space/svg.py index aae639e5346..ee024a7f0ec 100644 --- a/dimos/memory/vis/space/svg.py +++ b/dimos/memory/vis/space/svg.py @@ -28,10 +28,14 @@ from pathlib import Path from typing import TYPE_CHECKING +from dimos_generated.nav_msgs.msg import OccupancyGrid +from dimos_generated.sensor_msgs.msg import PointCloud2 import numpy as np from PIL import Image as PILImage +from dimos.mapping.occupancy.inflation import simple_inflate from dimos.mapping.occupancy.visualizations import generate_rgba_texture +from dimos.mapping.pointclouds.occupancy import height_cost_occupancy from dimos.memory.type.observation import Observation from dimos.memory.vis.color import Color from dimos.memory.vis.space.elements import ( @@ -44,8 +48,9 @@ SpaceElement, Text, ) -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.memory.vis.space.geometry import message_position, message_yaw +from dimos.msgs.geometry import pose_matrix +from dimos.msgs.occupancy import occupancy_extent, occupancy_view if TYPE_CHECKING: from dimos.memory.vis.space.space import Space @@ -95,7 +100,8 @@ def _style(el: object) -> tuple[str, float]: def _render_point(el: Point, b: Bounds) -> str: - x, y = el.msg.x, _y(el.msg.y) + position = message_position(el.msg) + x, y = position.x, _y(position.y) r = el.radius b.include(x - r, y - r) b.include(x + r, y + r) @@ -110,8 +116,9 @@ def _render_point(el: Point, b: Bounds) -> str: def _render_arrow(el: Arrow, b: Bounds) -> str: - x, y = el.msg.x, _y(el.msg.y) - yaw = el.msg.yaw + position = message_position(el.msg) + x, y = position.x, _y(position.y) + yaw = message_yaw(el.msg) length = el.length half_base = length * 0.4 @@ -139,7 +146,8 @@ def _render_pose(el: Pose, b: Bounds) -> str: arrow = Arrow(msg=el.msg, length=el.size, color=el.color, opacity=el.opacity) parts = [_render_arrow(arrow, b)] if el.label: - x, y = el.msg.x, _y(el.msg.y) + position = message_position(el.msg) + x, y = position.x, _y(position.y) fill, alpha = _style(el) parts.append( f' str: def _render_polyline(el: Polyline, b: Bounds) -> str: pts = [] for p in el.msg.poses: - x, y = p.x, _y(p.y) + x, y = p.pose.position.x, _y(p.pose.position.y) b.include(x, y) pts.append(f"{x:.4f},{y:.4f}") stroke, alpha = _style(el) @@ -163,7 +171,7 @@ def _render_polyline(el: Polyline, b: Bounds) -> str: def _render_box3d(el: Box3D, b: Bounds) -> str: - cx, cy = el.center.x, el.center.y + cx, cy = el.center.position.x, el.center.position.y hw, hh = el.size.x / 2, el.size.y / 2 # Top-left in world → SVG x = cx - hw @@ -188,12 +196,13 @@ def _render_box3d(el: Box3D, b: Bounds) -> str: def _render_camera(el: Camera, b: Bounds) -> str: - x, y = el.pose.x, _y(el.pose.y) - yaw = el.pose.yaw + position = message_position(el.pose) + x, y = position.x, _y(position.y) + yaw = message_yaw(el.pose) stroke, alpha = _style(el) - if el.camera_info and el.camera_info.K[4] > 0: - fy = el.camera_info.K[4] + if el.camera_info and el.camera_info.k[4] > 0: + fy = el.camera_info.k[4] fov_y = 2 * math.atan(el.camera_info.height / (2 * fy)) fov_half = fov_y / 2 wedge_len = 0.8 @@ -239,35 +248,29 @@ def _render_text(el: Text, b: Bounds) -> str: def _render_occupancy_grid(el: OccupancyGrid, b: Bounds) -> str: - if el.grid.size == 0: + cells = occupancy_view(el) + if cells.size == 0: return "" - - rgba = np.flipud(generate_rgba_texture(el)) + rgba = np.flipud(generate_rgba_texture(cells)) img = PILImage.fromarray(rgba, "RGBA") buf = io.BytesIO() img.save(buf, format="PNG") b64 = base64.b64encode(buf.getvalue()).decode("ascii") - - ox, oy = el.origin.x, el.origin.y - world_w = el.width * el.resolution - world_h = el.height * el.resolution - - # SVG top-left: world top-left with Y-flip - sx = ox - sy = _y(oy + world_h) - - b.include(sx, sy) - b.include(sx + world_w, sy + world_h) - + matrix = pose_matrix(el.info.origin) + width, height = occupancy_extent(el) + # Image row zero is the top edge after flipud; project the ROS grid plane into SVG XY. + a, c = matrix[0, 0], -matrix[0, 1] + d, e = -matrix[1, 0], matrix[1, 1] + tx, ty = matrix[0, 3] + matrix[0, 1] * height, -matrix[1, 3] - matrix[1, 1] * height + for x, y in ((0, 0), (width, 0), (width, height), (0, height)): + b.include(a * x + c * y + tx, d * x + e * y + ty) return ( - f'' ) -# Dispatch + top-level render - - def _render_element(el: SpaceElement, b: Bounds) -> str: if isinstance(el, Point): return _render_point(el, b) @@ -286,9 +289,6 @@ def _render_element(el: SpaceElement, b: Bounds) -> str: elif isinstance(el, OccupancyGrid): return _render_occupancy_grid(el, b) elif isinstance(el, PointCloud2): - from dimos.mapping.occupancy.inflation import simple_inflate - from dimos.mapping.pointclouds.occupancy import height_cost_occupancy - return _render_occupancy_grid(simple_inflate(height_cost_occupancy(el), 0.05), b) elif isinstance(el, Observation): ps = el.pose_stamped diff --git a/dimos/memory/vis/space/test_generated_geometry.py b/dimos/memory/vis/space/test_generated_geometry.py new file mode 100644 index 00000000000..ac0a9f36027 --- /dev/null +++ b/dimos/memory/vis/space/test_generated_geometry.py @@ -0,0 +1,155 @@ +# 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. + +"""Generated geometry reaches both scene renderers without message convenience methods.""" + +import math +import xml.etree.ElementTree as ET + +from dimos_generated.geometry_msgs.msg import Point as GeoPoint, Pose as GeoPose, PoseStamped +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid, Path +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest +import rerun as rr + +from dimos.memory.vis.space.elements import Arrow, Camera, Point, Polyline, Pose +from dimos.memory.vis.space.rerun import render +from dimos.memory.vis.space.space import Space +from dimos.msgs.camera_info import camera_info_from_intrinsics +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.image import image_from_array +from dimos.msgs.pointcloud import pointcloud_from_xyz + + +def test_rerun_logs_generated_pose_point_and_path(mocker): + stamped = PoseStamped( + pose=GeoPose( + position=GeoPoint(x=1, y=2, z=3), orientation=quaternion_from_euler(0, 0, math.pi / 2) + ) + ) + path = Path(poses=[stamped, PoseStamped(pose=GeoPose(position=GeoPoint(x=4, y=5)))]) + space = Space() + space.add(Point(stamped)).add(Pose(stamped)).add(Arrow(stamped)) + space.add(Polyline(path)).add(Camera(stamped)) + mocker.patch("dimos.visualization.rerun.init.rerun_init") + log = mocker.spy(rr, "log") + + with rr.RecordingStream("generated-geometry-test") as recording: + memory = recording.memory_recording() + try: + render(space, spawn=False) + assert memory.num_msgs() > 0 + finally: + recording.disconnect() + + logged = {call.args[0]: call.args[1] for call in log.call_args_list} + assert logged["scene/points"].positions.as_arrow_array().to_pylist() == [[1, 2, 3]] + assert logged["scene/poses/headings"].vectors.as_arrow_array().to_pylist()[0] == pytest.approx( + [0, 0.3, 0], abs=1e-7 + ) + assert logged["scene/polylines/0"].strips.as_arrow_array().to_pylist() == [ + [[1, 2, 0], [4, 5, 0]] + ] + assert logged["scene/cameras/0"].translation.as_arrow_array().to_pylist() == [[1, 2, 3]] + + +def test_both_renderers_apply_rotated_occupancy_origin(mocker): + grid = OccupancyGrid( + info=MapMetaData( + width=2, + height=1, + resolution=1.0, + origin=GeoPose( + position=GeoPoint(x=2, y=3), orientation=quaternion_from_euler(0, 0, math.pi / 2) + ), + ), + data=[0, 100], + ) + space = Space().base_map(OccupancyGrid.decode(grid.encode())) + svg = ET.fromstring(space.to_svg()) + image = svg.find(".//{http://www.w3.org/2000/svg}image") + assert image is not None + transform = image.attrib["transform"].removeprefix("matrix(").removesuffix(")") + assert [float(value) for value in transform.split()] == pytest.approx([0, -1, 1, 0, 1, -3]) + mocker.patch("dimos.visualization.rerun.init.rerun_init") + log = mocker.spy(rr, "log") + with rr.RecordingStream("generated-grid-test") as recording: + memory = recording.memory_recording() + try: + render(space, spawn=False) + assert memory.num_msgs() > 0 + finally: + recording.disconnect() + meshes = [call.args[1] for call in log.call_args_list if call.args[0] == "scene/map/0"] + assert len(meshes) == 1 + np.testing.assert_allclose( + meshes[0].vertex_positions.as_arrow_array().to_pylist(), + [[2, 3, 0], [2, 5, 0], [1, 5, 0], [1, 3, 0]], + atol=1e-7, + ) + + +def test_generated_cloud_renders_in_svg_and_rerun(mocker): + cloud = pointcloud_from_xyz( + np.array([[0.0, 0.0, 0.0], [0.2, 0.3, 1.0], [1.0, 1.0, 0.0]]), + header=Header(frame_id="map"), + ) + space = Space().add(PointCloud2.decode(cloud.encode())) + svg = ET.fromstring(space.to_svg()) + assert svg.find(".//{http://www.w3.org/2000/svg}image") is not None + mocker.patch("dimos.visualization.rerun.init.rerun_init") + log = mocker.spy(rr, "log") + with rr.RecordingStream("generated-cloud-test") as recording: + memory = recording.memory_recording() + try: + render(space, spawn=False) + assert memory.num_msgs() > 0 + finally: + recording.disconnect() + logged = {call.args[0]: call.args[1] for call in log.call_args_list} + np.testing.assert_allclose( + logged["scene/pointcloud/0"].positions.as_arrow_array().to_pylist(), + [[0, 0, 0], [0.2, 0.3, 1], [1, 1, 0]], + ) + + +def test_rerun_camera_uses_generated_calibration_and_image_helpers(mocker): + calibration = camera_info_from_intrinsics(20, 30, 2, 2, 4, 4, header=Header(frame_id="optical")) + image = image_from_array( + np.zeros((4, 4, 3), dtype=np.uint8), encoding="rgb8", header=calibration.header + ) + pose = PoseStamped( + pose=GeoPose(position=GeoPoint(x=2), orientation=quaternion_from_euler(0, 0, 0)) + ) + space = Space().add(Camera(pose=pose, image=image, camera_info=calibration)) + original = (calibration.encode(), image.encode()) + mocker.patch("dimos.visualization.rerun.init.rerun_init") + log = mocker.spy(rr, "log") + with rr.RecordingStream("generated-camera-test") as recording: + recording.memory_recording() + try: + render(space, spawn=False) + finally: + recording.disconnect() + values = [call.args[1] for call in log.call_args_list if call.args[0] == "scene/cameras/0"] + assert any(isinstance(value, rr.Pinhole) for value in values) + pinhole = next(value for value in values if isinstance(value, rr.Pinhole)) + assert pinhole.parent_frame.as_arrow_array().to_pylist() == ["tf#/optical"] + assert any( + call.args[0] == "scene/cameras/0/image" and isinstance(call.args[1], rr.EncodedImage) + for call in log.call_args_list + ) + assert (calibration.encode(), image.encode()) == original diff --git a/dimos/memory/vis/space/test_space.py b/dimos/memory/vis/space/test_space.py index 8f09e0e21ea..cc289827a93 100644 --- a/dimos/memory/vis/space/test_space.py +++ b/dimos/memory/vis/space/test_space.py @@ -14,59 +14,71 @@ """Tests for Space builder and element types.""" +from dimos_generated.geometry_msgs.msg import ( + Point as GeoPoint, + Pose as GeoPose, + PoseStamped, + Quaternion, + Vector3, +) +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid, Path as Path +from dimos_generated.vision_msgs.msg import Detection3D import numpy as np import pytest from dimos.memory.type.observation import EmbeddedObservation, Observation +from dimos.memory.vis import color from dimos.memory.vis.space.elements import Arrow, Box3D, Camera, Point, Polyline, Pose, Text from dimos.memory.vis.space.space import Space -from dimos.msgs.geometry_msgs.Point import Point as GeoPoint -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.nav_msgs.Path import Path as Path -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.vision_msgs.Detection3D import Detection3D +from dimos.msgs.image import image_from_array class TestElementTypes: """Element types wrap msgs with rendering intent + style.""" def test_pose_wraps_posestamped(self): - ps = PoseStamped(3.2, 1.5, 0.0) + ps = PoseStamped( + pose=GeoPose(position=GeoPoint(x=3.2, y=1.5, z=0.0), orientation=Quaternion()) + ) p = Pose(ps, color="red", label="fridge") assert p.msg is ps assert p.color == "red" assert p.label == "fridge" def test_arrow_wraps_posestamped(self): - ps = PoseStamped(1, 2, 0, 0, 0, 0.1, 1) + ps = PoseStamped( + pose=GeoPose( + position=GeoPoint(x=1, y=2, z=0), orientation=Quaternion(x=0, y=0, z=0.1, w=1) + ) + ) a = Arrow(ps, color="orange", length=0.8) assert a.msg is ps assert a.length == 0.8 def test_point_wraps_geopoint(self): - gp = GeoPoint(7.1, 4.3, 0) + gp = GeoPoint(x=7.1, y=4.3, z=0) p = Point(gp, color="green", label="bottle") assert p.msg is gp assert p.msg.x == pytest.approx(7.1) def test_point_wraps_posestamped(self): - ps = PoseStamped(3, 1, 0) + ps = PoseStamped(pose=GeoPose(position=GeoPoint(x=3, y=1, z=0), orientation=Quaternion())) p = Point(ps, radius=0.5) - assert p.msg.x == pytest.approx(3.0) + assert p.msg.pose.position.x == pytest.approx(3.0) def test_box3d_from_center_size(self): - from dimos.msgs.geometry_msgs.Pose import Pose as GeoPose - from dimos.msgs.geometry_msgs.Vector3 import Vector3 - - b = Box3D(center=GeoPose(5, 3, 0), size=Vector3(2, 1, 0.5), label="table") - assert b.center.x == pytest.approx(5.0) + b = Box3D( + center=GeoPose(position=GeoPoint(x=5, y=3, z=0), orientation=Quaternion()), + size=Vector3(x=2, y=1, z=0.5), + label="table", + ) + assert b.center.position.x == pytest.approx(5.0) assert b.size.x == pytest.approx(2.0) assert b.label == "table" def test_camera_with_image(self): - ps = PoseStamped(1, 2, 0) - img = Image(np.zeros((480, 640, 3), dtype=np.uint8)) + ps = PoseStamped(pose=GeoPose(position=GeoPoint(x=1, y=2, z=0), orientation=Quaternion())) + img = image_from_array(np.zeros((480, 640, 3), dtype=np.uint8), encoding="rgb8") c = Camera(pose=ps, image=img, color="purple") assert c.pose is ps assert c.image is img @@ -83,7 +95,7 @@ class TestSpaceExplicitElements: def test_add_pose(self): s = Space() - ps = PoseStamped(3, 1, 0) + ps = PoseStamped(pose=GeoPose(position=GeoPoint(x=3, y=1, z=0), orientation=Quaternion())) pose = Pose(ps, color="red") s.add(pose) assert len(s) == 1 @@ -91,15 +103,15 @@ def test_add_pose(self): def test_add_multiple_types(self): s = Space() - ps = PoseStamped(3, 1, 0) + ps = PoseStamped(pose=GeoPose(position=GeoPoint(x=3, y=1, z=0), orientation=Quaternion())) s.add(Pose(ps, color="red")) s.add(Arrow(ps, color="orange")) - s.add(Point(GeoPoint(1, 2, 0), label="x")) + s.add(Point(GeoPoint(x=1, y=2, z=0), label="x")) s.add(Text((0, 0, 0), "hello")) assert len(s) == 4 def test_chaining(self): - ps = PoseStamped(1, 1, 0) + ps = PoseStamped(pose=GeoPose(position=GeoPoint(x=1, y=1, z=0), orientation=Quaternion())) s = Space().add(Pose(ps)).add(Arrow(ps)).add(Text((0, 0, 0), "hi")) assert len(s) == 3 @@ -109,7 +121,9 @@ class TestSpaceAutoWrap: def test_posestamped_becomes_pose(self): s = Space() - ps = PoseStamped(3.2, 1.5, 0) + ps = PoseStamped( + pose=GeoPose(position=GeoPoint(x=3.2, y=1.5, z=0), orientation=Quaternion()) + ) s.add(ps, color="blue", label="auto") assert len(s) == 1 el = s.elements[0] @@ -120,7 +134,7 @@ def test_posestamped_becomes_pose(self): def test_geopoint_becomes_point(self): s = Space() - gp = GeoPoint(7, 4, 0) + gp = GeoPoint(x=7, y=4, z=0) s.add(gp, color="yellow") el = s.elements[0] assert isinstance(el, Point) @@ -129,7 +143,14 @@ def test_geopoint_becomes_point(self): def test_path_becomes_polyline(self): s = Space() - p = Path(poses=[PoseStamped(i, 0, 0) for i in range(3)]) + p = Path( + poses=[ + PoseStamped( + pose=GeoPose(position=GeoPoint(x=i, y=0, z=0), orientation=Quaternion()) + ) + for i in range(3) + ] + ) s.add(p, color="blue", width=0.1) el = s.elements[0] assert isinstance(el, Polyline) @@ -169,7 +190,7 @@ class TestSpaceObservations: """Space.add() smart dispatch for Observation types.""" def test_image_observation_stored_as_observation(self): - img = Image(np.zeros((480, 640, 3), dtype=np.uint8)) + img = image_from_array(np.zeros((480, 640, 3), dtype=np.uint8), encoding="rgb8") obs = Observation(id=1, ts=1.0, pose=(3, 1, 0, 0, 0, 0, 1), _data=img) s = Space() @@ -188,15 +209,18 @@ def test_non_image_observation_stored_as_observation(self): assert el.data == "some_data" def test_posestamped_observation_stored_as_observation(self): - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped as PS - - obs = Observation(id=3, ts=3.0, pose=(1, 2, 0, 0, 0, 0, 1), _data=PS(5, 2, 0)) + obs = Observation( + id=3, + ts=3.0, + pose=(1, 2, 0, 0, 0, 0, 1), + _data=PoseStamped(pose=GeoPose(position=GeoPoint(x=5, y=2))), + ) s = Space() s.add(obs) el = s.elements[0] assert isinstance(el, Observation) - assert el.data.x == pytest.approx(5.0) + assert el.data.pose.position.x == pytest.approx(5.0) def test_embedded_observation_stored_as_arrow(self): obs = EmbeddedObservation( @@ -224,7 +248,10 @@ def test_base_map(self): assert isinstance(s.elements[0], OccupancyGrid) def test_add_list_of_msgs(self): - poses = [PoseStamped(i, 0, 0) for i in range(3)] + poses = [ + PoseStamped(pose=GeoPose(position=GeoPoint(x=i, y=0, z=0), orientation=Quaternion())) + for i in range(3) + ] s = Space() s.add(poses, color="red") assert len(s) == 3 @@ -239,7 +266,7 @@ def test_repr_empty(self): def test_repr_with_elements(self): s = Space() - ps = PoseStamped(0, 0, 0) + ps = PoseStamped(pose=GeoPose(position=GeoPoint(x=0, y=0, z=0), orientation=Quaternion())) s.add(Pose(ps)) s.add(Pose(ps)) s.add(Arrow(ps)) @@ -255,29 +282,40 @@ def test_empty_space(self): assert svg.endswith("") def test_point_renders_circle(self): - from dimos.memory.vis import color - s = Space() - s.add(Point(GeoPoint(3, 4, 0), color="red", label="hi")) + s.add(Point(GeoPoint(x=3, y=4, z=0), color="red", label="hi")) svg = s.to_svg() assert " None: self.vector = vector - if timestamp: + if timestamp is not None: self.timestamp = timestamp else: self.timestamp = time.time() diff --git a/dimos/models/embedding/clip.py b/dimos/models/embedding/clip.py index 87e7cefa5db..478cc5ee962 100644 --- a/dimos/models/embedding/clip.py +++ b/dimos/models/embedding/clip.py @@ -17,6 +17,7 @@ from functools import cached_property from typing import overload +from dimos_generated.sensor_msgs.msg import Image from PIL import Image as PILImage import torch import torch.nn.functional as functional @@ -24,7 +25,8 @@ from dimos.models.base import HuggingFaceModel from dimos.models.embedding.base import Embedding, EmbeddingModel, HuggingFaceEmbeddingModelConfig -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_to_rgb +from dimos.msgs.time import to_seconds class CLIPModelConfig(HuggingFaceEmbeddingModelConfig): @@ -57,7 +59,7 @@ def embed(self, *images: Image) -> Embedding | list[Embedding]: Returns embeddings as torch.Tensor on device for efficient GPU comparisons. """ # Convert to PIL images - pil_images = [PILImage.fromarray(img.to_rgb().data) for img in images] + pil_images = [PILImage.fromarray(image_to_rgb(img)) for img in images] # Process images with torch.inference_mode(): @@ -70,7 +72,7 @@ def embed(self, *images: Image) -> Embedding | list[Embedding]: # Create embeddings (keep as torch.Tensor on device) embeddings: list[Embedding] = [] for i, feat in enumerate(image_features): - timestamp = images[i].ts + timestamp = to_seconds(images[i].header.stamp) embeddings.append(Embedding(vector=feat, timestamp=timestamp)) return embeddings[0] if len(images) == 1 else embeddings diff --git a/dimos/models/embedding/dino.py b/dimos/models/embedding/dino.py index 54d500c693d..1cdb276f154 100644 --- a/dimos/models/embedding/dino.py +++ b/dimos/models/embedding/dino.py @@ -17,6 +17,7 @@ from functools import cached_property from typing import Any, overload +from dimos_generated.sensor_msgs.msg import Image from PIL import Image as PILImage import torch import torch.nn.functional as functional @@ -24,7 +25,8 @@ from dimos.models.base import HuggingFaceModel from dimos.models.embedding.base import Embedding, EmbeddingModel, HuggingFaceEmbeddingModelConfig -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_to_rgb +from dimos.msgs.time import to_seconds class DINOModelConfig(HuggingFaceEmbeddingModelConfig): @@ -51,7 +53,7 @@ def embed(self, image: Image, /) -> Embedding: ... @overload def embed(self, *images: Image) -> list[Embedding]: ... def embed(self, *images: Image) -> Embedding | list[Embedding]: - pil_images = [PILImage.fromarray(img.to_rgb().data) for img in images] + pil_images = [PILImage.fromarray(image_to_rgb(img)) for img in images] with torch.inference_mode(): inputs = self._processor(images=pil_images, return_tensors="pt").to(self.config.device) outputs = self._model(**inputs) @@ -68,7 +70,7 @@ def embed(self, *images: Image) -> Embedding | list[Embedding]: embeddings: list[Embedding] = [] for i, feat in enumerate(feats): - embeddings.append(Embedding(vector=feat, timestamp=images[i].ts)) + embeddings.append(Embedding(vector=feat, timestamp=to_seconds(images[i].header.stamp))) return embeddings[0] if len(images) == 1 else embeddings @overload diff --git a/dimos/models/embedding/siglip.py b/dimos/models/embedding/siglip.py index 225c98a04f4..71dbd0a1e13 100644 --- a/dimos/models/embedding/siglip.py +++ b/dimos/models/embedding/siglip.py @@ -17,6 +17,7 @@ from functools import cached_property from typing import overload +from dimos_generated.sensor_msgs.msg import Image from PIL import Image as PILImage import torch import torch.nn.functional as functional @@ -24,7 +25,8 @@ from dimos.models.base import HuggingFaceModel from dimos.models.embedding.base import Embedding, EmbeddingModel, HuggingFaceEmbeddingModelConfig -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_to_rgb +from dimos.msgs.time import to_seconds class SigLIPModelConfig(HuggingFaceEmbeddingModelConfig): @@ -59,7 +61,7 @@ def embed(self, image: Image, /) -> Embedding: ... def embed(self, *images: Image) -> list[Embedding]: ... def embed(self, *images: Image) -> Embedding | list[Embedding]: """Embed one or more images into the shared image-text space.""" - pil_images = [PILImage.fromarray(img.to_rgb().data) for img in images] + pil_images = [PILImage.fromarray(image_to_rgb(img)) for img in images] with torch.inference_mode(): inputs = self._processor(images=pil_images, return_tensors="pt").to(self.config.device) @@ -68,7 +70,8 @@ def embed(self, *images: Image) -> Embedding | list[Embedding]: image_features = functional.normalize(image_features, dim=-1) embeddings = [ - Embedding(vector=feat, timestamp=images[i].ts) for i, feat in enumerate(image_features) + Embedding(vector=feat, timestamp=to_seconds(images[i].header.stamp)) + for i, feat in enumerate(image_features) ] return embeddings[0] if len(images) == 1 else embeddings diff --git a/dimos/models/embedding/treid.py b/dimos/models/embedding/treid.py index 3c3e33dd27e..0221bdba538 100644 --- a/dimos/models/embedding/treid.py +++ b/dimos/models/embedding/treid.py @@ -19,13 +19,15 @@ from functools import cached_property +from dimos_generated.sensor_msgs.msg import Image import torch import torch.nn.functional as functional from torchreid import utils as torchreid_utils from dimos.models.base import LocalModel from dimos.models.embedding.base import Embedding, EmbeddingModel, EmbeddingModelConfig -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_to_bgr +from dimos.msgs.time import to_seconds from dimos.utils.data import get_data @@ -60,7 +62,7 @@ def embed(self, *images: Image) -> Embedding | list[Embedding]: Returns embeddings as torch.Tensor on device for efficient GPU comparisons. """ # Convert to numpy arrays - torchreid expects numpy arrays or file paths - np_images = [img.to_opencv() for img in images] + np_images = [image_to_bgr(img) for img in images] # Extract features with torch.inference_mode(): @@ -78,7 +80,7 @@ def embed(self, *images: Image) -> Embedding | list[Embedding]: # Create embeddings (keep as torch.Tensor on device) embeddings = [] for i, feat in enumerate(features_tensor): - timestamp = images[i].ts + timestamp = to_seconds(images[i].header.stamp) embeddings.append(Embedding(vector=feat, timestamp=timestamp)) return embeddings[0] if len(images) == 1 else embeddings diff --git a/dimos/models/qwen/video_query.py b/dimos/models/qwen/video_query.py index 35979addd43..fcf8c75b517 100644 --- a/dimos/models/qwen/video_query.py +++ b/dimos/models/qwen/video_query.py @@ -20,7 +20,7 @@ from dimos.models.qwen.bbox import BBox from dimos.models.vl.qwen import DEFAULT_QWEN_VL_MODEL, QwenVlModel -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_from_array def query_single_frame( @@ -51,7 +51,7 @@ def query_single_frame( model = QwenVlModel(model_name=model_name, api_key=api_key) # Wrap with the default BGR tag so Image.to_base64()'s to_bgr() is a no-op and the array # reaches Qwen's JPEG encoder unchanged, matching the prior cv2.imencode(".jpg", frame). - return model.query(Image.from_numpy(image), query) + return model.query(image_from_array(image, encoding="bgr8"), query) def get_bbox_from_qwen_frame(frame, object_name: str | None = None) -> BBox | None: # type: ignore[no-untyped-def] diff --git a/dimos/models/segmentation/edge_tam.py b/dimos/models/segmentation/edge_tam.py index 667b05b4047..4be04c120ef 100644 --- a/dimos/models/segmentation/edge_tam.py +++ b/dimos/models/segmentation/edge_tam.py @@ -20,6 +20,7 @@ import tempfile from typing import TYPE_CHECKING, Any, Protocol, TypedDict +from dimos_generated.sensor_msgs.msg import Image from hydra.utils import instantiate import numpy as np from numpy.typing import NDArray @@ -28,7 +29,7 @@ import torch from dimos.models.base import default_torch_device -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_to_bgr, image_to_jpeg from dimos.perception.detection.detectors.base import Detector from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D @@ -172,7 +173,7 @@ def propose_all( ) self._mask_generators[points_per_side] = generator - rgb = cv2.cvtColor(image.to_opencv(), cv2.COLOR_BGR2RGB) + rgb = cv2.cvtColor(image_to_bgr(image), cv2.COLOR_BGR2RGB) amp = ( torch.autocast("cuda", dtype=torch.bfloat16) if self._predictor.device.type == "cuda" @@ -203,7 +204,7 @@ def segment( return ImageDetections2D(detections.image, []) image = detections.image - rgb = cv2.cvtColor(image.to_opencv(), cv2.COLOR_BGR2RGB) + rgb = cv2.cvtColor(image_to_bgr(image), cv2.COLOR_BGR2RGB) # bfloat16 autocast is CUDA-only; MPS runs in float32 amp = ( @@ -254,7 +255,7 @@ def _prepare_frame(self, image: Image) -> torch.Tensor: """Prepare frame for SAM2 (resize, normalize, convert to tensor).""" import cv2 - cv_image = image.to_opencv() + cv_image = image_to_bgr(image) rgb_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB) pil_image = PILImage.fromarray(rgb_image) @@ -407,7 +408,7 @@ def stop(self) -> None: def _temp_dir_context(image: Image) -> Generator[str, None, None]: path = tempfile.mkdtemp() - image.save(f"{path}/00000.jpg") + Path(path, "00000.jpg").write_bytes(image_to_jpeg(image)) try: yield path diff --git a/dimos/models/segmentation/test_generated_edge_tam_images.py b/dimos/models/segmentation/test_generated_edge_tam_images.py new file mode 100644 index 00000000000..ac079411bb5 --- /dev/null +++ b/dimos/models/segmentation/test_generated_edge_tam_images.py @@ -0,0 +1,50 @@ +# 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. + +from pathlib import Path +from types import SimpleNamespace + +import cv2 +import numpy as np +import torch + +from dimos.models.segmentation import edge_tam +from dimos.msgs.image import image_from_array +from dimos.perception.detection.type.detection2d.seg import Detection2DSeg + + +def test_generated_tracker_frame_and_initial_jpeg(monkeypatch): + monkeypatch.setattr( + edge_tam, "_build_model", lambda: SimpleNamespace(image_size=8, device="cpu") + ) + processor = edge_tam.EdgeTAMProcessor() + pixels = np.zeros((12, 16, 3), dtype=np.uint8) + pixels[:, :, 0] = 255 + image = image_from_array(pixels, encoding="rgb8") + frame = processor._prepare_frame(image) + assert frame.shape == (3, 8, 8) + torch.testing.assert_close( + frame[:, 0, 0], torch.tensor([(1 - 0.485) / 0.229, -0.456 / 0.224, -0.406 / 0.225]) + ) + with edge_tam._temp_dir_context(image) as directory: + path = Path(directory) + decoded = cv2.imread(str(path / "00000.jpg")) + assert decoded.shape == (12, 16, 3) + assert decoded[0, 0, 2] > 240 + assert not path.exists() + mask = torch.zeros((12, 16)) + mask[2:8, 3:10] = 1 + detection = Detection2DSeg.from_sam2_result(mask, 7, image) + assert detection.bbox == (3, 2, 9, 7) + assert detection.to_ros_detection2d().header == image.header diff --git a/dimos/models/segmentation/test_yoloe.py b/dimos/models/segmentation/test_yoloe.py index 5848f1aeca5..f576719b4c8 100644 --- a/dimos/models/segmentation/test_yoloe.py +++ b/dimos/models/segmentation/test_yoloe.py @@ -16,22 +16,24 @@ from typing import Any from unittest.mock import MagicMock +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header import numpy as np from ultralytics.models.yolo.yoloe import YOLOEVPSegPredictor from dimos.models.segmentation.yoloe import YoloeBoxSegmenter -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_from_array, image_to_bgr +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.detection.type.detection2d.seg import Detection2DSeg def _image() -> Image: - return Image( - data=np.zeros((32, 48, 3), dtype=np.uint8), - format=ImageFormat.BGR, - frame_id="camera", - ts=12.5, + return image_from_array( + np.zeros((32, 48, 3), dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id="camera", stamp=time_from_seconds(12.5)), ) @@ -42,7 +44,7 @@ def _box(image: Image, *, bbox: tuple[float, float, float, float]) -> Detection2 class_id=7, confidence=0.73, name="white mug", - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, ) @@ -60,7 +62,7 @@ def _mask( class_id=prompt_id, confidence=0.11, name=f"prompt-{prompt_id}", - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, mask=np.full((32, 48), value, dtype=np.uint8), ) @@ -137,7 +139,7 @@ def test_visual_box_prompts_use_segmentation_predictor(mocker: Any) -> None: get_data.assert_called_once_with("models_yoloe") model_factory.assert_called_once_with(Path("/models/yoloe-11s-seg.pt")) kwargs = dict(model.predict.call_args.kwargs) - np.testing.assert_array_equal(kwargs.pop("source"), image.to_opencv()) + np.testing.assert_array_equal(kwargs.pop("source"), image_to_bgr(image)) visual_prompts = kwargs.pop("visual_prompts") np.testing.assert_array_equal( visual_prompts["bboxes"], np.asarray([detection.bbox], dtype=np.float64) diff --git a/dimos/models/segmentation/yoloe.py b/dimos/models/segmentation/yoloe.py index a13d0a9d009..b2be2676b24 100644 --- a/dimos/models/segmentation/yoloe.py +++ b/dimos/models/segmentation/yoloe.py @@ -19,11 +19,12 @@ from dataclasses import replace from typing import Protocol +from dimos_generated.sensor_msgs.msg import Image import numpy as np from numpy.typing import NDArray from ultralytics.models.yolo.yoloe import YOLOEVPSegPredictor -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_to_bgr from dimos.perception.detection.detectors.yoloe import Yoloe2DDetector, YoloePromptMode from dimos.perception.detection.type.detection2d.bbox import Bbox, Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D @@ -58,7 +59,7 @@ def predict_image(self, image: Image) -> ImageDetections2D: raise RuntimeError("YOLO-E segmentation requires visual box prompts") with self._lock: results = self.model.predict( - source=image.to_opencv(), + source=image_to_bgr(image), device=self.device, conf=self.conf, iou=0.6, diff --git a/dimos/models/vl/README.md b/dimos/models/vl/README.md index 12fdbfbc1ba..488cee1ef19 100644 --- a/dimos/models/vl/README.md +++ b/dimos/models/vl/README.md @@ -13,14 +13,14 @@ returns absolute pixel coordinates, which rules out the `qwen3-vl-*` models. ```python from dimos.models.vl.qwen import QwenVlModel -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_from_file, image_view # Initialize the model (requires ALIBABA_API_KEY environment variable) model = QwenVlModel() -image = Image.from_file("path/to/your/image.jpg") +image = image_from_file("path/to/your/image.jpg") -response = model.query(image.data, "What do you see in this image?") +response = model.query(image_view(image), "What do you see in this image?") print(response) ``` @@ -48,10 +48,10 @@ The model supports four modes of operation: ```python from dimos.models.vl.moondream_hosted import MoondreamHostedVlModel -from dimos.msgs.sensor_msgs import Image +from dimos.msgs.image import image_from_file, image_view model = MoondreamHostedVlModel() -image = Image.from_file("path/to/image.jpg") +image = image_from_file("path/to/image.jpg") # 1. Caption print(f"Caption: {model.caption(image)}") diff --git a/dimos/models/vl/base.py b/dimos/models/vl/base.py index 55b616176dc..3f1d4bad18d 100644 --- a/dimos/models/vl/base.py +++ b/dimos/models/vl/base.py @@ -19,8 +19,11 @@ from typing import Any import warnings +from dimos_generated.sensor_msgs.msg import Image + from dimos.core.resource import Resource -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_from_array, image_view +from dimos.msgs.time import to_seconds from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.detection.type.detection2d.point import Detection2DPoint @@ -116,7 +119,7 @@ def vlm_detection_to_detection2d( class_id=-1, confidence=1.0, name=name, - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, ) @@ -166,7 +169,7 @@ def vlm_point_to_detection2d_point( x=x, y=y, name=name, - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, track_id=track_id, ) @@ -196,9 +199,21 @@ def _prepare_image(self, image: Image) -> tuple[Image, float]: Returns: Tuple of (prepared_image, scale_factor). Scale factor is 1.0 if no resize. """ + import cv2 + if self.config.auto_resize is not None: max_w, max_h = self.config.auto_resize - return image.resize_to_fit(max_w, max_h) + if max_w <= 0 or max_h <= 0: + raise ValueError("auto_resize dimensions must be positive") + scale = min(1.0, max_w / image.width, max_h / image.height) + if scale == 1.0: + return image, scale + pixels = cv2.resize( + image_view(image), + (max(1, int(image.width * scale)), max(1, int(image.height * scale))), + interpolation=cv2.INTER_AREA, + ) + return image_from_array(pixels, encoding=image.encoding, header=image.header), scale return image, 1.0 @abstractmethod @@ -251,8 +266,13 @@ def caption(self, image: Image) -> str: def start(self) -> None: """Start the model by running a simple query (Resource interface).""" + import cv2 + try: - image = Image.from_file(get_data("cafe-smol.jpg")).to_rgb() + pixels = cv2.imread(str(get_data("cafe-smol.jpg"))) + if pixels is None: + raise ValueError("Could not read model warmup image") + image = image_from_array(pixels, encoding="bgr8") self.query(image, "What is this?") except Exception: pass diff --git a/dimos/models/vl/florence.py b/dimos/models/vl/florence.py index 8e964bb85f7..92ce9b2a353 100644 --- a/dimos/models/vl/florence.py +++ b/dimos/models/vl/florence.py @@ -15,13 +15,14 @@ from enum import Enum from functools import cached_property +from dimos_generated.sensor_msgs.msg import Image from PIL import Image as PILImage import torch from transformers import AutoModelForCausalLM, AutoProcessor from dimos.models.base import HuggingFaceModel from dimos.models.vl.base import Captioner -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_to_rgb class CaptionDetail(Enum): @@ -98,7 +99,7 @@ def caption(self, image: Image, detail: str | CaptionDetail | None = None) -> st task_prompt = CaptionDetail.from_str(detail).value # Convert to PIL - pil_image = PILImage.fromarray(image.to_rgb().data) + pil_image = PILImage.fromarray(image_to_rgb(image)) # Process inputs inputs = self._processor(text=task_prompt, images=pil_image, return_tensors="pt") @@ -137,7 +138,7 @@ def caption_batch(self, *images: Image) -> list[str]: task_prompt = self._task_prompt # Convert all to PIL - pil_images = [PILImage.fromarray(img.to_rgb().data) for img in images] + pil_images = [PILImage.fromarray(image_to_rgb(img)) for img in images] # Process batch inputs = self._processor( diff --git a/dimos/models/vl/moondream.py b/dimos/models/vl/moondream.py index e3cfe744ce8..0956d7844d1 100644 --- a/dimos/models/vl/moondream.py +++ b/dimos/models/vl/moondream.py @@ -16,6 +16,7 @@ from typing import Any import warnings +from dimos_generated.sensor_msgs.msg import Image import numpy as np from PIL import Image as PILImage import torch @@ -23,7 +24,8 @@ from dimos.models.base import HuggingFaceModel, HuggingFaceModelConfig from dimos.models.vl.base import VlModel, VlModelConfig -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_from_array, image_to_rgb +from dimos.msgs.time import to_seconds from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.detection.type.detection2d.point import Detection2DPoint @@ -63,11 +65,11 @@ def _to_pil(self, image: Image | np.ndarray[Any, Any]) -> PILImage.Image: DeprecationWarning, stacklevel=2, ) - image = Image.from_numpy(image) + image = image_from_array(image, encoding="bgr8") image, _ = self._prepare_image(image) - rgb_image = image.to_rgb() - return PILImage.fromarray(rgb_image.data) + rgb_image = image_to_rgb(image) + return PILImage.fromarray(rgb_image) def query(self, image: Image | np.ndarray, query: str, **kwargs) -> str: # type: ignore[no-untyped-def] pil_image = self._to_pil(image) @@ -182,7 +184,7 @@ def query_detections( class_id=-1, # Moondream doesn't provide class IDs confidence=1.0, # Moondream doesn't provide confidence scores name=query, # Use the query as the object name - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, ) @@ -222,7 +224,7 @@ def query_points( x=x, y=y, name=query, - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, track_id=track_id, ) diff --git a/dimos/models/vl/moondream_hosted.py b/dimos/models/vl/moondream_hosted.py index 76e55451a1e..c03d50832a1 100644 --- a/dimos/models/vl/moondream_hosted.py +++ b/dimos/models/vl/moondream_hosted.py @@ -17,12 +17,14 @@ from typing import Any import warnings +from dimos_generated.sensor_msgs.msg import Image import moondream as md # type: ignore[import-untyped] import numpy as np from PIL import Image as PILImage from dimos.models.vl.base import VlModel, VlModelConfig -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_from_array, image_to_rgb +from dimos.msgs.time import to_seconds from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.detection.type.detection2d.point import Detection2DPoint @@ -51,10 +53,10 @@ def _to_pil_image(self, image: Image | np.ndarray) -> PILImage.Image: DeprecationWarning, stacklevel=3, ) - image = Image.from_numpy(image) + image = image_from_array(image, encoding="bgr8") - rgb_image = image.to_rgb() - return PILImage.fromarray(rgb_image.data) + rgb_image = image_to_rgb(image) + return PILImage.fromarray(rgb_image) def query(self, image: Image | np.ndarray, query: str, **kwargs) -> str: # type: ignore[no-untyped-def] pil_image = self._to_pil_image(image) @@ -118,7 +120,7 @@ def query_detections( class_id=-1, confidence=1.0, name=query, - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, ) @@ -153,7 +155,7 @@ def query_points( x=x, y=y, name=query, - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, track_id=track_id, ) diff --git a/dimos/models/vl/openai.py b/dimos/models/vl/openai.py index 2303e20fb70..ae22e6a7a58 100644 --- a/dimos/models/vl/openai.py +++ b/dimos/models/vl/openai.py @@ -12,15 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +import base64 from functools import cached_property import os from typing import Any -import numpy as np +from dimos_generated.sensor_msgs.msg import Image from openai import OpenAI from dimos.models.vl.base import VlModel, VlModelConfig -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_to_jpeg from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -46,26 +47,15 @@ def _client(self) -> OpenAI: def query( self, - image: Image | np.ndarray, + image: Image, query: str, response_format: dict[str, Any] | None = None, **kwargs: Any, ) -> str: - if isinstance(image, np.ndarray): - import warnings - - warnings.warn( - "OpenAIVlModel.query should receive standard dimos Image type, not a numpy array", - DeprecationWarning, - stacklevel=2, - ) - - image = Image.from_numpy(image) - # Apply auto_resize if configured image, _ = self._prepare_image(image) - img_base64 = image.to_base64() + img_base64 = base64.b64encode(image_to_jpeg(image)).decode("ascii") api_kwargs: dict[str, Any] = { "model": self.config.model_name, @@ -105,7 +95,7 @@ def query_batch( { "type": "image_url", "image_url": { - "url": f"data:image/jpeg;base64,{self._prepare_image(img)[0].to_base64()}" + "url": f"data:image/jpeg;base64,{base64.b64encode(image_to_jpeg(self._prepare_image(img)[0])).decode('ascii')}" }, } for img in images diff --git a/dimos/models/vl/qwen.py b/dimos/models/vl/qwen.py index 84542b60c68..46a0ba2677f 100644 --- a/dimos/models/vl/qwen.py +++ b/dimos/models/vl/qwen.py @@ -12,15 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +import base64 from functools import cached_property import os from typing import Any -import numpy as np +from dimos_generated.sensor_msgs.msg import Image from openai import OpenAI from dimos.models.vl.base import VlModel, VlModelConfig -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_to_jpeg # Must be a model that returns ABSOLUTE pixel coordinates. The qwen3-vl-* models return # normalized 0-1000 coordinates and ignore prompts asking for pixels, which silently breaks @@ -51,22 +52,11 @@ def _client(self) -> OpenAI: api_key=api_key, ) - def query(self, image: Image | np.ndarray, query: str) -> str: # type: ignore[override] - if isinstance(image, np.ndarray): - import warnings - - warnings.warn( - "QwenVlModel.query should receive standard dimos Image type, not a numpy array", - DeprecationWarning, - stacklevel=2, - ) - - image = Image.from_numpy(image) - + def query(self, image: Image, query: str) -> str: # type: ignore[override] # Apply auto_resize if configured image, _ = self._prepare_image(image) - img_base64 = image.to_base64() + img_base64 = base64.b64encode(image_to_jpeg(image)).decode("ascii") response = self._client.chat.completions.create( model=self.config.model_name, @@ -76,7 +66,7 @@ def query(self, image: Image | np.ndarray, query: str) -> str: # type: ignore[o "content": [ { "type": "image_url", - "image_url": {"url": f"data:image/png;base64,{img_base64}"}, + "image_url": {"url": f"data:image/jpeg;base64,{img_base64}"}, }, {"type": "text", "text": query}, ], @@ -101,7 +91,7 @@ def query_batch( { "type": "image_url", "image_url": { - "url": f"data:image/png;base64,{self._prepare_image(img)[0].to_base64()}" + "url": f"data:image/jpeg;base64,{base64.b64encode(image_to_jpeg(self._prepare_image(img)[0])).decode('ascii')}" }, } for img in images diff --git a/dimos/models/vl/test_base.py b/dimos/models/vl/test_base.py index 92b8406f136..5d00e5395a0 100644 --- a/dimos/models/vl/test_base.py +++ b/dimos/models/vl/test_base.py @@ -17,7 +17,8 @@ import pytest from dimos.models.vl.qwen import QwenVlModel -from dimos.msgs.sensor_msgs.Image import Image +from dimos_generated.sensor_msgs.msg import Image +from dimos.msgs.image import image_from_file from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.utils.data import get_data @@ -44,7 +45,7 @@ def test_query_detections_mocked() -> None: """Test query_detections with mocked API response (no API key required).""" # Load test image - image = Image.from_file(get_data("cafe.jpg")) + image = image_from_file(get_data("cafe.jpg")) # Create model and mock the query method model = QwenVlModel() diff --git a/dimos/models/vl/test_generated_images.py b/dimos/models/vl/test_generated_images.py new file mode 100644 index 00000000000..01e560fb76a --- /dev/null +++ b/dimos/models/vl/test_generated_images.py @@ -0,0 +1,85 @@ +# 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. + +import base64 +from unittest.mock import MagicMock + +import cv2 +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header +from dimos_generated.vision_msgs.msg import Detection2DArray +import numpy as np +import pytest + +from dimos.models.vl.base import VlModel +from dimos.models.vl.openai import OpenAIVlModel +from dimos.models.vl.qwen import QwenVlModel +from dimos.msgs.image import image_from_array + + +class StaticModel(VlModel): + def stop(self) -> None: + pass + + def query(self, image: Image, query: str, **kwargs) -> str: + return '[["target", 2, 4, 10, 12]]' + + +def test_generated_vl_resize_and_detection(): + image = image_from_array( + np.zeros((40, 60, 3), dtype=np.uint8), + encoding="rgb8", + header=Header(frame_id="camera", stamp=Time(sec=1700000000, nanosec=123456789)), + ) + model = StaticModel(auto_resize=(30, 20)) + resized, scale = model._prepare_image(image) + assert scale == 0.5 + assert (resized.width, resized.height) == (30, 20) + assert resized.header == image.header + detections = model.query_detections(image, "target") + wire = Detection2DArray.decode(detections.to_ros_detection2d_array().encode()) + assert len(wire.detections) == 1 + assert wire.header == image.header + assert wire.detections[0].results[0].hypothesis.class_id == "-1" + + +@pytest.mark.parametrize("model_type", [QwenVlModel, OpenAIVlModel]) +@pytest.mark.parametrize("batch", [False, True]) +def test_provider_encodes_generated_images_as_jpeg(model_type, batch): + pixels = np.zeros((40, 60, 3), dtype=np.uint8) + pixels[:, :, 0] = 255 + image = image_from_array(pixels, encoding="rgb8") + model = model_type(auto_resize=(30, 20)) + client = MagicMock() + client.chat.completions.create.return_value.choices[0].message.content = "target" + model.__dict__["_client"] = client + if batch: + assert model.query_batch([image, image], "find target") == ["target", "target"] + else: + assert model.query(image, "find target") == "target" + request = client.chat.completions.create.call_args.kwargs + content = request["messages"][0]["content"] + images = [item for item in content if item["type"] == "image_url"] + assert len(images) == (2 if batch else 1) + for item in images: + prefix, encoded = item["image_url"]["url"].split(",", 1) + assert prefix == "data:image/jpeg;base64" + jpeg = base64.b64decode(encoded) + assert jpeg[:2] == b"\xff\xd8" + decoded = cv2.imdecode(np.frombuffer(jpeg, dtype=np.uint8), cv2.IMREAD_COLOR) + assert decoded.shape == (20, 30, 3) + assert decoded[10, 15, 2] > 240 + assert decoded[10, 15, 0] < 10 + model.stop() diff --git a/dimos/models/vl/tool_base.py b/dimos/models/vl/tool_base.py index 0a6b5b80c62..2cd64850e00 100644 --- a/dimos/models/vl/tool_base.py +++ b/dimos/models/vl/tool_base.py @@ -17,7 +17,8 @@ from dimos.core.transport import LCMTransport from dimos.models.vl.moondream import MoondreamVlModel from dimos.models.vl.qwen import QwenVlModel -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos_generated.sensor_msgs.msg import Image +from dimos.msgs.image import image_from_file from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.utils.data import get_data @@ -26,7 +27,7 @@ def test_query_detections_real() -> None: """Test query_detections with real API calls (requires API key).""" # Load test image - image = Image.from_file(get_data("cafe.jpg")) + image = image_from_file(get_data("cafe.jpg")) # Initialize the model (will use real API) model = QwenVlModel() @@ -55,7 +56,7 @@ def test_query_detections_real() -> None: def test_query_points() -> None: """Test query_points with real API calls (requires API key).""" # Load test image - image = Image.from_file(get_data("cafe.jpg"), format=ImageFormat.RGB).to_rgb() + image = image_from_file(get_data("cafe.jpg")) # Initialize the model (will use real API) model = MoondreamVlModel() diff --git a/dimos/models/vl/tool_vlm.py b/dimos/models/vl/tool_vlm.py index 7bc88c13677..a8b99dd445a 100644 --- a/dimos/models/vl/tool_vlm.py +++ b/dimos/models/vl/tool_vlm.py @@ -12,13 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os import time from typing import TYPE_CHECKING import pytest from dimos.models.vl.moondream import MoondreamVlModel -from dimos.msgs.sensor_msgs.Image import Image +from dimos_generated.sensor_msgs.msg import Image +from dimos.utils.testing.replay import TimedSensorReplay if TYPE_CHECKING: from dimos.models.vl.base import VlModel @@ -32,11 +34,14 @@ ) def test_vlm_query_batch(model_class: "type[VlModel]", model_name: str) -> None: """Test query_batch optimization - multiple images, same query.""" - from dimos.utils.testing.legacy_pickle import LegacyPickleStore - - # Load 5 frames at 1-second intervals using LegacyPickleStore - replay = LegacyPickleStore[Image]("unitree_go2_office_walk2/video") - images = [replay.find_closest_seek(i).to_rgb() for i in range(0, 10, 2)] + # Explicit CDR recording input: .db/. + recording = os.environ["DIMOS_VLM_RECORDING"] + replay = TimedSensorReplay[Image](recording) + images = [] + for offset in range(0, 10, 2): + image = replay.find_closest_seek(offset) + assert isinstance(image, Image), f"No generated image at offset {offset}" + images.append(image) print(f"\nTesting {model_name} query_batch with {len(images)} images") diff --git a/dimos/msgs/camera_info.py b/dimos/msgs/camera_info.py new file mode 100644 index 00000000000..7977a275b22 --- /dev/null +++ b/dimos/msgs/camera_info.py @@ -0,0 +1,169 @@ +# 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. + +"""Calibration construction and array access for generated ROS CameraInfo values.""" + +from copy import deepcopy +import math +from pathlib import Path +import re +from typing import Any, Literal + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.sensor_msgs.msg import CameraInfo +from dimos_generated.std_msgs.msg import Header +import numpy as np +from numpy.typing import NDArray +import yaml + + +def _dimensions(width: int, height: int) -> None: + if any( + isinstance(value, bool) or not isinstance(value, int) or value <= 0 + for value in (width, height) + ): + raise ValueError("Calibration width and height must be positive integers") + + +def camera_info_from_intrinsics( + fx: float, fy: float, cx: float, cy: float, width: int, height: int, *, header: Header +) -> CameraInfo: + """Construct an undistorted pinhole calibration with an explicit source header.""" + _dimensions(width, height) + if not all(math.isfinite(value) for value in (fx, fy, cx, cy)) or fx <= 0 or fy <= 0: + raise ValueError("Intrinsics must be finite with positive focal lengths") + return CameraInfo( + header=header, + width=width, + height=height, + distortion_model="plumb_bob", + d=[0.0] * 5, + k=[fx, 0, cx, 0, fy, cy, 0, 0, 1], + r=[1, 0, 0, 0, 1, 0, 0, 0, 1], + p=[fx, 0, cx, 0, 0, fy, cy, 0, 0, 0, 1, 0], + ) + + +def camera_info_from_fov( + fov_degrees: float, + width: int, + height: int, + *, + header: Header, + axis: Literal["horizontal", "vertical"] = "vertical", +) -> CameraInfo: + """Construct a centered pinhole model with square pixels from an axis's FOV.""" + _dimensions(width, height) + if not math.isfinite(fov_degrees) or not 0 < fov_degrees < 180: + raise ValueError("Field of view must be finite and between 0 and 180 degrees") + if axis not in ("horizontal", "vertical"): + raise ValueError("Field-of-view axis must be horizontal or vertical") + extent = width if axis == "horizontal" else height + focal = extent / (2 * math.tan(math.radians(fov_degrees) / 2)) + return camera_info_from_intrinsics( + focal, focal, width / 2, height / 2, width, height, header=header + ) + + +def _matrix(data: dict[str, Any], name: str, rows: int, columns: int | None) -> list[float]: + matrix = data[name] + if not isinstance(matrix, dict) or matrix.get("rows") != rows: + raise ValueError(f"{name} must declare {rows} rows") + actual_columns = matrix.get("cols") + if ( + not isinstance(actual_columns, int) + or actual_columns < 0 + or (columns is not None and actual_columns != columns) + ): + raise ValueError(f"{name} has invalid columns") + values = matrix.get("data") + if not isinstance(values, list) or len(values) != rows * actual_columns: + raise ValueError(f"{name} data length does not match its dimensions") + converted = [float(value) for value in values] + if not all(math.isfinite(value) for value in converted): + raise ValueError(f"{name} must contain finite values") + return converted + + +def camera_info_from_yaml(path: str | Path, *, header: Header) -> CameraInfo: + """Read a ROS calibration YAML file, validating declared matrix dimensions.""" + with Path(path).open() as source: + data = yaml.safe_load(source) + if not isinstance(data, dict): + raise ValueError("Calibration YAML must contain a mapping") + width, height = data["image_width"], data["image_height"] + _dimensions(width, height) + return CameraInfo( + header=header, + width=width, + height=height, + distortion_model=data["distortion_model"], + k=_matrix(data, "camera_matrix", 3, 3), + d=_matrix(data, "distortion_coefficients", 1, None), + r=_matrix(data, "rectification_matrix", 3, 3), + p=_matrix(data, "projection_matrix", 3, 4), + ) + + +def intrinsic_matrix(message: CameraInfo) -> NDArray[np.float64]: + """Copy the 3x3 intrinsic matrix into independently mutable NumPy storage.""" + return np.array(message.k, dtype=np.float64, copy=True).reshape(3, 3) + + +def camera_info_with_stamp(message: CameraInfo, stamp: Time) -> CameraInfo: + """Copy a calibration with an exact source stamp, without mutating prior publishes.""" + result = deepcopy(message) + result.header.stamp = deepcopy(stamp) + return result + + +class CalibrationProvider: + """Lazily load named YAML calibrations as plain generated CameraInfo values. + + Both SingleWebcam and single_webcam resolve single_webcam.yaml. This + helper lives outside generated messages; calibration files stay local. + """ + + def __init__(self, calibration_dir: str | Path, *, frame_id: str = "camera_optical") -> None: + self._calibration_dir = Path(calibration_dir) + self._frame_id = frame_id + self._cache: dict[str, CameraInfo] = {} + + @staticmethod + def _to_snake_case(name: str) -> str: + first = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name) + return re.sub("([a-z0-9])([A-Z])", r"\1_\2", first).lower() + + def __getattr__(self, name: str) -> CameraInfo: + if not re.fullmatch(r"[A-Za-z][A-Za-z0-9_]*", name): + raise AttributeError(f"No calibration file found for: {name}") + snake = self._to_snake_case(name) + for key in (name, snake): + if key in self._cache: + return self._cache[key] + for key in (name, snake): + path = self._calibration_dir / f"{key}.yaml" + if path.is_file(): + value = camera_info_from_yaml(path, header=Header(frame_id=self._frame_id)) + self._cache[name] = value + self._cache[snake] = value + return value + raise AttributeError(f"No calibration file found for: {name}") + + def __dir__(self) -> list[str]: + names: set[str] = set() + for path in self._calibration_dir.glob("*.yaml"): + names.add(path.stem) + names.add("".join(word.capitalize() for word in path.stem.split("_"))) + return sorted(names) diff --git a/dimos/msgs/foxglove_msgs/test_CompressedVideo.py b/dimos/msgs/foxglove_msgs/test_CompressedVideo.py index fcecd05a1cb..fa475e0c5b9 100644 --- a/dimos/msgs/foxglove_msgs/test_CompressedVideo.py +++ b/dimos/msgs/foxglove_msgs/test_CompressedVideo.py @@ -12,37 +12,44 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.foxglove_msgs.msg import CompressedVideo import pytest import rerun as rr -from dimos.msgs.foxglove_msgs.CompressedVideo import CompressedVideo from dimos.msgs.helpers import resolve_msg_type +from dimos.msgs.time import time_from_nanoseconds, to_nanoseconds +from dimos.visualization.rerun.message_helpers import video_archetype PACKET = b"\x00\x00\x00\x01\x65payload" -def test_lcm_encode_decode() -> None: - original = CompressedVideo(data=PACKET, format="h264", frame_id="front_camera", ts=1234567890.5) - decoded = CompressedVideo.lcm_decode(original.lcm_encode()) +def test_cdr_encode_decode() -> None: + original = CompressedVideo( + data=PACKET, + format="h264", + frame_id="front_camera", + timestamp=time_from_nanoseconds(1234567890123456789), + ) + decoded = CompressedVideo.decode(original.encode()) - assert decoded.data.tobytes() == PACKET + assert bytes(decoded.data) == PACKET assert decoded.format == "h264" assert decoded.frame_id == "front_camera" - assert abs(decoded.ts - original.ts) < 1e-6 + assert to_nanoseconds(decoded.timestamp) == 1234567890123456789 -def test_resolves_over_the_raw_lcm_type() -> None: - """The rerun bridge decodes by type name — it must land on the to_rerun() class.""" - resolved = resolve_msg_type("foxglove_msgs.CompressedVideo") +def test_resolves_generated_canonical_type() -> None: + """Typed transports discover generated schema-bearing values.""" + resolved = resolve_msg_type("foxglove_msgs/msg/CompressedVideo") assert resolved is CompressedVideo def test_to_rerun_video_stream() -> None: - stream = CompressedVideo(data=PACKET, format="h264").to_rerun() + stream = video_archetype(CompressedVideo(data=PACKET, format="h264")) assert isinstance(stream, rr.VideoStream) assert bytes(stream.sample.as_arrow_array().to_pylist()[0]) == PACKET # type: ignore[union-attr] def test_to_rerun_unknown_codec() -> None: with pytest.raises(ValueError, match="mjpeg"): - CompressedVideo(data=PACKET, format="mjpeg").to_rerun() + video_archetype(CompressedVideo(data=PACKET, format="mjpeg")) diff --git a/dimos/msgs/geometry.py b/dimos/msgs/geometry.py new file mode 100644 index 00000000000..a5f007eb114 --- /dev/null +++ b/dimos/msgs/geometry.py @@ -0,0 +1,374 @@ +# 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. + +"""Geometry operations on generated ROS value types.""" + +from collections.abc import Sequence +import math + +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, + Wrench, +) +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.std_msgs.msg import Header +import numpy as np +from numpy.typing import NDArray +from scipy.spatial.transform import Rotation + + +def yaw(rotation: Quaternion) -> float: + """Return rotation about Z in radians (ROS XYZW quaternion convention).""" + return math.atan2( + 2 * (rotation.w * rotation.z + rotation.x * rotation.y), + 1 - 2 * (rotation.y * rotation.y + rotation.z * rotation.z), + ) + + +def transform_from_odometry(message: Odometry) -> TransformStamped: + """Copy the pose and declared frames into TF, preserving the exact source stamp.""" + position = message.pose.pose.position + return TransformStamped( + header=message.header, + child_frame_id=message.child_frame_id, + transform=Transform( + translation=Vector3(x=position.x, y=position.y, z=position.z), + rotation=message.pose.pose.orientation, + ), + ) + + +def quaternion_from_euler(roll: float, pitch: float, yaw: float) -> Quaternion: + """Convert fixed-axis XYZ angles in radians to a normalized ROS quaternion.""" + if not all(math.isfinite(value) for value in (roll, pitch, yaw)): + raise ValueError("Euler angles must be finite") + x, y, z, w = Rotation.from_euler("xyz", [roll, pitch, yaw]).as_quat() + return Quaternion(x=x, y=y, z=z, w=w) + + +def _rotation(quaternion: Quaternion) -> Rotation: + values = [quaternion.x, quaternion.y, quaternion.z, quaternion.w] + if not all(math.isfinite(value) for value in values): + raise ValueError("Quaternion components must be finite") + # SciPy rejects zero norm and normalizes non-unit input. + return Rotation.from_quat(values) + + +def _translation(transform: Transform) -> NDArray[np.float64]: + translation = transform.translation + values = np.array([translation.x, translation.y, translation.z], dtype=np.float64) + if not np.isfinite(values).all(): + raise ValueError("Translation components must be finite") + return values + + +def _transform(translation: NDArray[np.float64], rotation: Rotation) -> Transform: + x, y, z, w = rotation.as_quat() + return Transform( + translation=Vector3(x=translation[0], y=translation[1], z=translation[2]), + rotation=Quaternion(x=x, y=y, z=z, w=w), + ) + + +def transform_matrix(transform: Transform) -> NDArray[np.float64]: + """Copy a generated transform into a homogeneous 4-by-4 matrix.""" + matrix = np.eye(4, dtype=np.float64) + matrix[:3, :3] = _rotation(transform.rotation).as_matrix() + matrix[:3, 3] = _translation(transform) + return matrix + + +def compose_transforms(first: TransformStamped, second: TransformStamped) -> TransformStamped: + """Compose A←B and B←C, preserving A's exact source stamp in the A←C result. + + Inputs are copied into a new value. The shared frame must match; callers + querying a time-varying chain are responsible for choosing each edge's time. + """ + if first.child_frame_id != second.header.frame_id: + raise ValueError( + f"Cannot compose frames {first.child_frame_id!r} and {second.header.frame_id!r}" + ) + rotation = _rotation(first.transform.rotation) + translation = _translation(first.transform) + rotation.apply(_translation(second.transform)) + return TransformStamped( + header=first.header, + child_frame_id=second.child_frame_id, + transform=_transform(translation, rotation * _rotation(second.transform.rotation)), + ) + + +def inverse_transform(message: TransformStamped) -> TransformStamped: + """Return B←A from A←B, swapping frame names and preserving the exact stamp.""" + rotation = _rotation(message.transform.rotation).inv() + return TransformStamped( + header=Header(stamp=message.header.stamp, frame_id=message.child_frame_id), + child_frame_id=message.header.frame_id, + transform=_transform(-rotation.apply(_translation(message.transform)), rotation), + ) + + +def pose_from_transform(message: TransformStamped) -> PoseStamped: + """Copy the child frame's pose in the parent frame into a generated PoseStamped.""" + translation = message.transform.translation + return PoseStamped( + header=message.header, + pose=Pose( + position=Point(x=translation.x, y=translation.y, z=translation.z), + orientation=message.transform.rotation, + ), + ) + + +def transform_from_pose(message: PoseStamped, *, child_frame_id: str) -> TransformStamped: + """Copy a stamped pose into TF with an explicit child and unchanged source header.""" + if not isinstance(message, PoseStamped): + raise TypeError("Expected a generated PoseStamped") + position = message.pose.position + return TransformStamped( + header=message.header, + child_frame_id=child_frame_id, + transform=Transform( + translation=Vector3(x=position.x, y=position.y, z=position.z), + rotation=message.pose.orientation, + ), + ) + + +def pose_matrix(pose: Pose) -> NDArray[np.float64]: + """Copy a generated pose into a homogeneous matrix, rejecting invalid geometry.""" + values = np.array([pose.position.x, pose.position.y, pose.position.z], dtype=np.float64) + if not np.isfinite(values).all(): + raise ValueError("Position components must be finite") + result = np.eye(4, dtype=np.float64) + result[:3, :3] = _rotation(pose.orientation).as_matrix() + result[:3, 3] = values + return result + + +def quaternion_from_matrix(matrix: NDArray[np.float64]) -> Quaternion: + """Convert a finite proper 3-by-3 rotation matrix into a generated quaternion.""" + values = np.asarray(matrix, dtype=np.float64) + if values.shape != (3, 3) or not np.isfinite(values).all(): + raise ValueError("Rotation matrix must be finite and 3-by-3") + if not np.allclose(values.T @ values, np.eye(3), atol=1e-6, rtol=0.0) or not np.isclose( + np.linalg.det(values), 1.0, atol=1e-6, rtol=0.0 + ): + raise ValueError("Rotation matrix must be orthonormal with determinant +1") + x, y, z, w = Rotation.from_matrix(values).as_quat() + return Quaternion(x=x, y=y, z=z, w=w) + + +def pose_from_matrix(matrix: NDArray[np.float64]) -> Pose: + """Copy a finite rigid 4-by-4 transform into a generated pose.""" + values = np.asarray(matrix, dtype=np.float64) + if values.shape != (4, 4) or not np.isfinite(values).all(): + raise ValueError("Pose matrix must be finite and 4-by-4") + if not np.allclose(values[3], [0.0, 0.0, 0.0, 1.0], atol=1e-9, rtol=0.0): + raise ValueError("Pose matrix must have homogeneous last row [0, 0, 0, 1]") + return Pose( + position=Point(x=values[0, 3], y=values[1, 3], z=values[2, 3]), + orientation=quaternion_from_matrix(values[:3, :3]), + ) + + +def transform_from_matrix(matrix: NDArray[np.float64]) -> Transform: + """Copy a finite rigid homogeneous matrix into a generated transform.""" + pose = pose_from_matrix(matrix) + return Transform( + translation=Vector3(x=pose.position.x, y=pose.position.y, z=pose.position.z), + rotation=pose.orientation, + ) + + +def quaternion_euler(rotation: Quaternion) -> tuple[float, float, float]: + """Return fixed-axis XYZ Euler angles in radians for a generated quaternion.""" + roll, pitch, yaw = _rotation(rotation).as_euler("xyz") + return float(roll), float(pitch), float(yaw) + + +def translate_pose_local(pose: Pose, offset: Vector3) -> Pose: + """Copy a pose displaced by an offset expressed in its local axes.""" + matrix = pose_matrix(pose) + values = np.array([offset.x, offset.y, offset.z], dtype=np.float64) + if not np.isfinite(values).all(): + raise ValueError("Offset components must be finite") + delta = matrix[:3, :3] @ values + return Pose( + position=Point( + x=pose.position.x + delta[0], y=pose.position.y + delta[1], z=pose.position.z + delta[2] + ), + orientation=pose.orientation, + ) + + +def point_distance(first: Point, second: Point) -> float: + """Euclidean distance between generated position values.""" + return math.dist((first.x, first.y, first.z), (second.x, second.y, second.z)) + + +def quaternion_angle(first: Quaternion, second: Quaternion) -> float: + """Return the shortest angular separation in radians, normalizing inputs.""" + return float((_rotation(first).inv() * _rotation(second)).magnitude()) + + +def wrench_from_array(values: Sequence[float] | NDArray[np.float64]) -> Wrench: + """Copy force in N and torque in Nm from [fx, fy, fz, tx, ty, tz].""" + array = np.asarray(values, dtype=np.float64) + if array.shape != (6,): + raise ValueError("Expected 6 elements [fx, fy, fz, tx, ty, tz]") + return Wrench( + force=Vector3(x=array[0], y=array[1], z=array[2]), + torque=Vector3(x=array[3], y=array[4], z=array[5]), + ) + + +def wrench_array(message: Wrench) -> NDArray[np.float64]: + """Return an independent numeric force/torque array for arithmetic.""" + return np.array( + [ + message.force.x, + message.force.y, + message.force.z, + message.torque.x, + message.torque.y, + message.torque.z, + ], + dtype=np.float64, + ) + + +def vector_array(value: Vector3 | Point) -> NDArray[np.float64]: + """Copy a generated spatial vector/point into an independent XYZ array.""" + return np.array([value.x, value.y, value.z], dtype=np.float64) + + +def vector_from_array(values: Sequence[float] | NDArray[np.float64]) -> Vector3: + """Copy exactly three coordinates into a generated vector.""" + array = np.asarray(values, dtype=np.float64) + if array.shape != (3,): + raise ValueError("Expected 3 vector components") + return Vector3(x=array[0], y=array[1], z=array[2]) + + +def normalized_vector(value: Vector3) -> Vector3: + """Return a unit vector; zero length remains zero.""" + array = vector_array(value) + length = float(np.linalg.norm(array)) + return vector_from_array(array / length) if length else Vector3() + + +def quaternion_array(value: Quaternion) -> NDArray[np.float64]: + """Copy XYZW components without normalizing their declared values.""" + return np.array([value.x, value.y, value.z, value.w], dtype=np.float64) + + +def quaternion_from_array(values: Sequence[float] | NDArray[np.float64]) -> Quaternion: + """Copy exactly four XYZW components into a generated quaternion.""" + array = np.asarray(values, dtype=np.float64) + if array.shape != (4,): + raise ValueError("Quaternion requires exactly 4 components") + return Quaternion(x=array[0], y=array[1], z=array[2], w=array[3]) + + +def quaternion_product(first: Quaternion, second: Quaternion) -> Quaternion: + """Hamilton product, retaining non-unit magnitudes and multiplication order.""" + if not isinstance(first, Quaternion) or not isinstance(second, Quaternion): + raise TypeError("Cannot multiply Quaternion with a non-quaternion value") + a, b = quaternion_array(first), quaternion_array(second) + xyz = a[3] * b[:3] + b[3] * a[:3] + np.cross(a[:3], b[:3]) + return quaternion_from_array(np.append(xyz, a[3] * b[3] - np.dot(a[:3], b[:3]))) + + +def quaternion_conjugate(value: Quaternion) -> Quaternion: + """Negate the imaginary components, preserving the scalar component.""" + return Quaternion(x=-value.x, y=-value.y, z=-value.z, w=value.w) + + +def quaternion_inverse(value: Quaternion) -> Quaternion: + """Algebraic inverse, including non-unit quaternions.""" + array = quaternion_array(value) + squared = float(np.dot(array, array)) + if squared == 0: + raise ZeroDivisionError("Cannot invert zero quaternion") + return quaternion_from_array(quaternion_array(quaternion_conjugate(value)) / squared) + + +def normalized_quaternion(value: Quaternion) -> Quaternion: + """Return unit components without changing sign; reject zero norm.""" + array = quaternion_array(value) + length = float(np.linalg.norm(array)) + if length == 0: + raise ZeroDivisionError("Cannot normalize zero quaternion") + return quaternion_from_array(array / length) + + +def rotate_vector(rotation: Quaternion, value: Vector3) -> Vector3: + """Rotate an XYZ vector using a normalized spatial rotation.""" + return vector_from_array(_rotation(rotation).apply(vector_array(value))) + + +def point_from_array(values: Sequence[float] | NDArray[np.float64]) -> Point: + """Copy exactly three Cartesian coordinates into a generated point.""" + vector = vector_from_array(values) + return Point(x=vector.x, y=vector.y, z=vector.z) + + +def compose_poses(first: Pose, second: Pose) -> Pose: + """Compose parent/local poses without mutating inputs or normalizing stored quaternions.""" + offset = _rotation(first.orientation).apply(vector_array(second.position)) + return Pose( + position=point_from_array(vector_array(first.position) + offset), + orientation=quaternion_product(first.orientation, second.orientation), + ) + + +def compose_transform_values(first: Transform, second: Transform) -> Transform: + """Compose unframed spatial transforms in parent/local order.""" + if not isinstance(first, Transform) or not isinstance(second, Transform): + raise TypeError("Expected generated Transform values") + return Transform( + translation=vector_from_array( + vector_array(first.translation) + + _rotation(first.rotation).apply(vector_array(second.translation)) + ), + rotation=quaternion_product(first.rotation, second.rotation), + ) + + +def transform_pose_local(pose: Pose, transform: Transform) -> Pose: + """Apply an offset expressed in a pose's local axes.""" + return compose_poses( + pose, + Pose( + position=point_from_array(vector_array(transform.translation)), + orientation=transform.rotation, + ), + ) + + +def relative_pose_transform(first: Pose, second: Pose) -> Transform: + """Find the local transform taking the first pose to the second.""" + rotation = quaternion_inverse(first.orientation) + return Transform( + translation=vector_from_array( + _rotation(rotation).apply(vector_array(second.position) - vector_array(first.position)) + ), + rotation=quaternion_product(rotation, second.orientation), + ) diff --git a/dimos/msgs/geometry_msgs/test_PointStamped.py b/dimos/msgs/geometry_msgs/test_PointStamped.py index 727d9190ed8..e4a27f14b88 100644 --- a/dimos/msgs/geometry_msgs/test_PointStamped.py +++ b/dimos/msgs/geometry_msgs/test_PointStamped.py @@ -12,53 +12,49 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for geometry_msgs.PointStamped — msg -> lcm bytes -> msg roundtrip.""" +"""Generated point/stamp roundtrips and explicit pose construction.""" -import time +from dimos_generated.geometry_msgs.msg import Point, PointStamped, Pose, PoseStamped, Quaternion +from dimos_generated.std_msgs.msg import Header +from rosbags.typesys import Stores, get_typestore -from dimos_lcm.geometry_msgs import Point as LCMPoint +from dimos.msgs.time import time_from_nanoseconds, time_from_seconds, to_nanoseconds -from dimos.msgs.geometry_msgs.Point import Point -from dimos.msgs.geometry_msgs.PointStamped import PointStamped +def test_point_has_standard_cdr_schema() -> None: + point = Point(x=1, y=2, z=3) + decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(point.encode(), Point.msg_name) + assert (decoded.x, decoded.y, decoded.z) == (1, 2, 3) -def test_point_inherits_lcm() -> None: - """Point wrapper inherits from LCMPoint.""" - assert isinstance(Point(1.0, 2.0, 3.0), LCMPoint) - -def test_lcm_encode_decode() -> None: - """Test encoding and decoding of PointStamped to/from binary LCM format.""" +def test_cdr_encode_decode() -> None: source = PointStamped( - x=1.5, - y=-2.5, - z=3.5, - ts=time.time(), - frame_id="/world/grid", + point=Point(x=1.5, y=-2.5, z=3.5), + header=Header(stamp=time_from_nanoseconds(1700000000123456789), frame_id="/world/grid"), ) - binary_msg = source.lcm_encode() - dest = PointStamped.lcm_decode(binary_msg) - + dest = PointStamped.decode(source.encode()) assert isinstance(dest, PointStamped) assert dest is not source - assert dest.x == source.x - assert dest.y == source.y - assert dest.z == source.z - assert abs(dest.ts - source.ts) < 1e-6 - assert dest.frame_id == source.frame_id - + assert dest.point.x == source.point.x + assert dest.point.y == source.point.y + assert dest.point.z == source.point.z + assert to_nanoseconds(dest.header.stamp) == to_nanoseconds(source.header.stamp) + assert dest.header.frame_id == source.header.frame_id -def test_to_pose_stamped() -> None: - """Test conversion to PoseStamped with identity orientation.""" - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - - pt = PointStamped(x=1.0, y=2.0, z=3.0, ts=500.0, frame_id="/map") - pose = pt.to_pose_stamped() +def test_explicit_pose_stamped_conversion() -> None: + point = PointStamped( + point=Point(x=1, y=2, z=3), header=Header(stamp=time_from_seconds(500), frame_id="/map") + ) + pose = PoseStamped( + header=point.header, pose=Pose(position=point.point, orientation=Quaternion(w=1)) + ) assert isinstance(pose, PoseStamped) - assert pose.x == 1.0 - assert pose.y == 2.0 - assert pose.z == 3.0 - assert pose.orientation.w == 1.0 - assert pose.ts == 500.0 - assert pose.frame_id == "/map" + assert pose.pose.position.x == 1 + assert pose.pose.position.y == 2 + assert pose.pose.position.z == 3 + assert pose.pose.orientation.w == 1 + assert to_nanoseconds(pose.header.stamp) == 500000000000 + assert pose.header.frame_id == "/map" + pose.pose.position.x = 9 + assert point.point.x == 1 diff --git a/dimos/msgs/geometry_msgs/test_Pose.py b/dimos/msgs/geometry_msgs/test_Pose.py index dac5ed6207d..9dd6bb12a88 100644 --- a/dimos/msgs/geometry_msgs/test_Pose.py +++ b/dimos/msgs/geometry_msgs/test_Pose.py @@ -14,112 +14,92 @@ import pickle -from dimos_lcm.geometry_msgs import Pose as LCMPose +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion, Vector3 import numpy as np import pytest - -from dimos.msgs.geometry_msgs.Pose import Pose, to_pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from rosbags.typesys import Stores, get_typestore + +from dimos.msgs.geometry import ( + compose_poses, + point_from_array, + quaternion_euler, + quaternion_from_array, + quaternion_from_euler, + vector_array, + vector_from_array, +) def test_pose_default_init() -> None: """Test that default initialization creates a pose at origin with identity orientation.""" pose = Pose() - - # Position should be at origin assert pose.position.x == 0.0 assert pose.position.y == 0.0 assert pose.position.z == 0.0 - - # Orientation should be identity quaternion assert pose.orientation.x == 0.0 assert pose.orientation.y == 0.0 assert pose.orientation.z == 0.0 assert pose.orientation.w == 1.0 - - # Test convenience properties - assert pose.x == 0.0 - assert pose.y == 0.0 - assert pose.z == 0.0 + assert pose.position.x == 0.0 + assert pose.position.y == 0.0 + assert pose.position.z == 0.0 def test_pose_pose_init() -> None: """Test initialization with position coordinates only (identity orientation).""" - pose_data = Pose(1.0, 2.0, 3.0) - - pose = to_pose(pose_data) - - # Position should be as specified + pose_data = Pose(position=Point(x=1.0, y=2.0, z=3.0)) + pose = pose_data assert pose.position.x == 1.0 assert pose.position.y == 2.0 assert pose.position.z == 3.0 - - # Orientation should be identity quaternion assert pose.orientation.x == 0.0 assert pose.orientation.y == 0.0 assert pose.orientation.z == 0.0 assert pose.orientation.w == 1.0 - - # Test convenience properties - assert pose.x == 1.0 - assert pose.y == 2.0 - assert pose.z == 3.0 + assert pose.position.x == 1.0 + assert pose.position.y == 2.0 + assert pose.position.z == 3.0 def test_pose_position_init() -> None: """Test initialization with position coordinates only (identity orientation).""" - pose = Pose(1.0, 2.0, 3.0) - - # Position should be as specified + pose = Pose(position=Point(x=1.0, y=2.0, z=3.0)) assert pose.position.x == 1.0 assert pose.position.y == 2.0 assert pose.position.z == 3.0 - - # Orientation should be identity quaternion assert pose.orientation.x == 0.0 assert pose.orientation.y == 0.0 assert pose.orientation.z == 0.0 assert pose.orientation.w == 1.0 - - # Test convenience properties - assert pose.x == 1.0 - assert pose.y == 2.0 - assert pose.z == 3.0 + assert pose.position.x == 1.0 + assert pose.position.y == 2.0 + assert pose.position.z == 3.0 def test_pose_full_init() -> None: """Test initialization with position and orientation coordinates.""" - pose = Pose(1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.9) - - # Position should be as specified + pose = Pose( + position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) + ) assert pose.position.x == 1.0 assert pose.position.y == 2.0 assert pose.position.z == 3.0 - - # Orientation should be as specified assert pose.orientation.x == 0.1 assert pose.orientation.y == 0.2 assert pose.orientation.z == 0.3 assert pose.orientation.w == 0.9 - - # Test convenience properties - assert pose.x == 1.0 - assert pose.y == 2.0 - assert pose.z == 3.0 + assert pose.position.x == 1.0 + assert pose.position.y == 2.0 + assert pose.position.z == 3.0 def test_pose_vector_position_init() -> None: """Test initialization with Vector3 position (identity orientation).""" - position = Vector3(4.0, 5.0, 6.0) - pose = Pose(position) - - # Position should match the vector + position = Vector3(x=4.0, y=5.0, z=6.0) + pose = Pose(position=point_from_array(vector_array(position))) assert pose.position.x == 4.0 assert pose.position.y == 5.0 assert pose.position.z == 6.0 - - # Orientation should be identity assert pose.orientation.x == 0.0 assert pose.orientation.y == 0.0 assert pose.orientation.z == 0.0 @@ -128,16 +108,12 @@ def test_pose_vector_position_init() -> None: def test_pose_vector_quaternion_init() -> None: """Test initialization with Vector3 position and Quaternion orientation.""" - position = Vector3(1.0, 2.0, 3.0) - orientation = Quaternion(0.1, 0.2, 0.3, 0.9) - pose = Pose(position, orientation) - - # Position should match the vector + position = Vector3(x=1.0, y=2.0, z=3.0) + orientation = Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) + pose = Pose(position=point_from_array(vector_array(position)), orientation=orientation) assert pose.position.x == 1.0 assert pose.position.y == 2.0 assert pose.position.z == 3.0 - - # Orientation should match the quaternion assert pose.orientation.x == 0.1 assert pose.orientation.y == 0.2 assert pose.orientation.z == 0.3 @@ -148,14 +124,13 @@ def test_pose_list_init() -> None: """Test initialization with lists for position and orientation.""" position_list = [1.0, 2.0, 3.0] orientation_list = [0.1, 0.2, 0.3, 0.9] - pose = Pose(position_list, orientation_list) - - # Position should match the list + pose = Pose( + position=point_from_array(position_list), + orientation=quaternion_from_array(orientation_list), + ) assert pose.position.x == 1.0 assert pose.position.y == 2.0 assert pose.position.z == 3.0 - - # Orientation should match the list assert pose.orientation.x == 0.1 assert pose.orientation.y == 0.2 assert pose.orientation.z == 0.3 @@ -167,14 +142,12 @@ def test_pose_tuple_init() -> None: position = [1.0, 2.0, 3.0] orientation = [0.1, 0.2, 0.3, 0.9] pose_tuple = (position, orientation) - pose = Pose(pose_tuple) - - # Position should match + pose = Pose( + position=point_from_array(pose_tuple[0]), orientation=quaternion_from_array(pose_tuple[1]) + ) assert pose.position.x == 1.0 assert pose.position.y == 2.0 assert pose.position.z == 3.0 - - # Orientation should match assert pose.orientation.x == 0.1 assert pose.orientation.y == 0.2 assert pose.orientation.z == 0.3 @@ -184,14 +157,13 @@ def test_pose_tuple_init() -> None: def test_pose_dict_init() -> None: """Test initialization from a dictionary with 'position' and 'orientation' keys.""" pose_dict = {"position": [1.0, 2.0, 3.0], "orientation": [0.1, 0.2, 0.3, 0.9]} - pose = Pose(pose_dict) - - # Position should match + pose = Pose( + position=point_from_array(pose_dict["position"]), + orientation=quaternion_from_array(pose_dict["orientation"]), + ) assert pose.position.x == 1.0 assert pose.position.y == 2.0 assert pose.position.z == 3.0 - - # Orientation should match assert pose.orientation.x == 0.1 assert pose.orientation.y == 0.2 assert pose.orientation.z == 0.3 @@ -200,45 +172,35 @@ def test_pose_dict_init() -> None: def test_pose_copy_init() -> None: """Test initialization from another Pose (copy constructor).""" - original = Pose(1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.9) - copy = Pose(original) - - # Position should match + original = Pose( + position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) + ) + copy = Pose.decode(original.encode()) assert copy.position.x == 1.0 assert copy.position.y == 2.0 assert copy.position.z == 3.0 - - # Orientation should match assert copy.orientation.x == 0.1 assert copy.orientation.y == 0.2 assert copy.orientation.z == 0.3 assert copy.orientation.w == 0.9 - - # Should be a copy, not the same object assert copy is not original assert copy == original -def test_pose_lcm_init() -> None: - """Test initialization from an LCM Pose.""" - # Create LCM pose - lcm_pose = LCMPose() - lcm_pose.position.x = 1.0 - lcm_pose.position.y = 2.0 - lcm_pose.position.z = 3.0 - lcm_pose.orientation.x = 0.1 - lcm_pose.orientation.y = 0.2 - lcm_pose.orientation.z = 0.3 - lcm_pose.orientation.w = 0.9 - - pose = Pose(lcm_pose) - - # Position should match +def test_pose_mutated_value_copy() -> None: + """Test initialization from a generated Pose.""" + source_pose = Pose() + source_pose.position.x = 1.0 + source_pose.position.y = 2.0 + source_pose.position.z = 3.0 + source_pose.orientation.x = 0.1 + source_pose.orientation.y = 0.2 + source_pose.orientation.z = 0.3 + source_pose.orientation.w = 0.9 + pose = Pose.decode(source_pose.encode()) assert pose.position.x == 1.0 assert pose.position.y == 2.0 assert pose.position.z == 3.0 - - # Orientation should match assert pose.orientation.x == 0.1 assert pose.orientation.y == 0.2 assert pose.orientation.z == 0.3 @@ -247,86 +209,63 @@ def test_pose_lcm_init() -> None: def test_pose_properties() -> None: """Test pose property access.""" - pose = Pose(1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.9) - - # Test position properties - assert pose.x == 1.0 - assert pose.y == 2.0 - assert pose.z == 3.0 - - # Test orientation properties (through quaternion's to_euler method) - euler = pose.orientation.to_euler() - assert pose.roll == euler.x - assert pose.pitch == euler.y - assert pose.yaw == euler.z + pose = Pose( + position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) + ) + assert pose.position.x == 1.0 + assert pose.position.y == 2.0 + assert pose.position.z == 3.0 + euler = vector_from_array(quaternion_euler(pose.orientation)) + assert quaternion_euler(pose.orientation)[0] == euler.x + assert quaternion_euler(pose.orientation)[1] == euler.y + assert quaternion_euler(pose.orientation)[2] == euler.z def test_pose_euler_properties_identity() -> None: """Test pose Euler angle properties with identity orientation.""" - pose = Pose(1.0, 2.0, 3.0) # Identity orientation - - # Identity quaternion should give zero Euler angles - assert np.isclose(pose.roll, 0.0, atol=1e-10) - assert np.isclose(pose.pitch, 0.0, atol=1e-10) - assert np.isclose(pose.yaw, 0.0, atol=1e-10) - - # Euler property should also be zeros - assert np.isclose(pose.orientation.euler.x, 0.0, atol=1e-10) - assert np.isclose(pose.orientation.euler.y, 0.0, atol=1e-10) - assert np.isclose(pose.orientation.euler.z, 0.0, atol=1e-10) - - -def test_pose_repr() -> None: - """Test pose string representation.""" - pose = Pose(1.234, 2.567, 3.891, 0.1, 0.2, 0.3, 0.9) - - repr_str = repr(pose) - - # Should contain position and orientation info - assert "Pose" in repr_str - assert "position" in repr_str - assert "orientation" in repr_str - - # Should contain the actual values (approximately) - assert "1.234" in repr_str or "1.23" in repr_str - assert "2.567" in repr_str or "2.57" in repr_str - - -def test_pose_str() -> None: - """Test pose string formatting.""" - pose = Pose(1.234, 2.567, 3.891, 0.1, 0.2, 0.3, 0.9) - - str_repr = str(pose) - - # Should contain position coordinates - assert "1.234" in str_repr - assert "2.567" in str_repr - assert "3.891" in str_repr - - # Should contain Euler angles - assert "euler" in str_repr - - # Should be formatted with specified precision - assert str_repr.count("Pose") == 1 + pose = Pose(position=Point(x=1.0, y=2.0, z=3.0)) + assert np.isclose(quaternion_euler(pose.orientation)[0], 0.0, atol=1e-10) + assert np.isclose(quaternion_euler(pose.orientation)[1], 0.0, atol=1e-10) + assert np.isclose(quaternion_euler(pose.orientation)[2], 0.0, atol=1e-10) + assert np.isclose(vector_from_array(quaternion_euler(pose.orientation)).x, 0.0, atol=1e-10) + assert np.isclose(vector_from_array(quaternion_euler(pose.orientation)).y, 0.0, atol=1e-10) + assert np.isclose(vector_from_array(quaternion_euler(pose.orientation)).z, 0.0, atol=1e-10) + + +def test_pose_independent_cdr_fields() -> None: + source = Pose( + position=Point(x=1.234, y=2.567, z=3.891), + orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9), + ) + decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(source.encode(), Pose.msg_name) + assert (decoded.position.x, decoded.position.y, decoded.position.z) == (1.234, 2.567, 3.891) + assert ( + decoded.orientation.x, + decoded.orientation.y, + decoded.orientation.z, + decoded.orientation.w, + ) == (0.1, 0.2, 0.3, 0.9) def test_pose_equality() -> None: """Test pose equality comparison.""" - pose1 = Pose(1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.9) - pose2 = Pose(1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.9) - pose3 = Pose(1.1, 2.0, 3.0, 0.1, 0.2, 0.3, 0.9) # Different position - pose4 = Pose(1.0, 2.0, 3.0, 0.11, 0.2, 0.3, 0.9) # Different orientation - - # Equal poses + pose1 = Pose( + position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) + ) + pose2 = Pose( + position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) + ) + pose3 = Pose( + position=Point(x=1.1, y=2.0, z=3.0), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) + ) + pose4 = Pose( + position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.11, y=0.2, z=0.3, w=0.9) + ) assert pose1 == pose2 assert pose2 == pose1 - - # Different poses assert pose1 != pose3 assert pose1 != pose4 assert pose3 != pose4 - - # Different types assert pose1 != "not a pose" assert pose1 != [1.0, 2.0, 3.0] assert pose1 is not None @@ -336,15 +275,13 @@ def test_pose_with_numpy_arrays() -> None: """Test pose initialization with numpy arrays.""" position_array = np.array([1.0, 2.0, 3.0]) orientation_array = np.array([0.1, 0.2, 0.3, 0.9]) - - pose = Pose(position_array, orientation_array) - - # Position should match + pose = Pose( + position=point_from_array(position_array), + orientation=quaternion_from_array(orientation_array), + ) assert pose.position.x == 1.0 assert pose.position.y == 2.0 assert pose.position.z == 3.0 - - # Orientation should match assert pose.orientation.x == 0.1 assert pose.orientation.y == 0.2 assert pose.orientation.z == 0.3 @@ -353,15 +290,13 @@ def test_pose_with_numpy_arrays() -> None: def test_pose_with_mixed_types() -> None: """Test pose initialization with mixed input types.""" - # Position as tuple, orientation as list - pose1 = Pose((1.0, 2.0, 3.0), [0.1, 0.2, 0.3, 0.9]) - - # Position as numpy array, orientation as Vector3/Quaternion + pose1 = Pose( + position=point_from_array((1.0, 2.0, 3.0)), + orientation=quaternion_from_array([0.1, 0.2, 0.3, 0.9]), + ) position = np.array([1.0, 2.0, 3.0]) - orientation = Quaternion(0.1, 0.2, 0.3, 0.9) - pose2 = Pose(position, orientation) - - # Both should result in the same pose + orientation = Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) + pose2 = Pose(position=point_from_array(position), orientation=orientation) assert pose1.position.x == pose2.position.x assert pose1.position.y == pose2.position.y assert pose1.position.z == pose2.position.z @@ -371,24 +306,20 @@ def test_pose_with_mixed_types() -> None: assert pose1.orientation.w == pose2.orientation.w -def test_to_pose_passthrough() -> None: - """Test to_pose function with Pose input (passthrough).""" - original = Pose(1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.9) - result = to_pose(original) - - # Should be the same object (passthrough) - assert result is original +def test_pose_nested_value_copy() -> None: + source = Pose(position=Point(x=1, y=2, z=3)) + copied = Pose(position=source.position, orientation=source.orientation) + copied.position.x = 10 + assert source.position == Point(x=1, y=2, z=3) + assert copied.position == Point(x=10, y=2, z=3) def test_to_pose_conversion() -> None: """Test to_pose function with convertible inputs.""" - # Note: The to_pose conversion function has type checking issues in the current implementation - # Test direct construction instead to verify the intended functionality - - # Test the intended functionality by creating poses directly pose_tuple = ([1.0, 2.0, 3.0], [0.1, 0.2, 0.3, 0.9]) - result1 = Pose(pose_tuple) - + result1 = Pose( + position=point_from_array(pose_tuple[0]), orientation=quaternion_from_array(pose_tuple[1]) + ) assert isinstance(result1, Pose) assert result1.position.x == 1.0 assert result1.position.y == 2.0 @@ -397,11 +328,11 @@ def test_to_pose_conversion() -> None: assert result1.orientation.y == 0.2 assert result1.orientation.z == 0.3 assert result1.orientation.w == 0.9 - - # Test with dictionary pose_dict = {"position": [1.0, 2.0, 3.0], "orientation": [0.1, 0.2, 0.3, 0.9]} - result2 = Pose(pose_dict) - + result2 = Pose( + position=point_from_array(pose_dict["position"]), + orientation=quaternion_from_array(pose_dict["orientation"]), + ) assert isinstance(result2, Pose) assert result2.position.x == 1.0 assert result2.position.y == 2.0 @@ -414,71 +345,56 @@ def test_to_pose_conversion() -> None: def test_pose_euler_roundtrip() -> None: """Test conversion from Euler angles to quaternion and back.""" - # Start with known Euler angles (small angles to avoid gimbal lock) roll = 0.1 pitch = 0.2 yaw = 0.3 - - # Create quaternion from Euler angles - euler_vector = Vector3(roll, pitch, yaw) - quaternion = euler_vector.to_quaternion() - - # Create pose with this quaternion - pose = Pose(Vector3(0, 0, 0), quaternion) - - # Convert back to Euler angles - result_euler = pose.orientation.euler - - # Should get back the original Euler angles (within tolerance) - assert np.isclose(result_euler.x, roll, atol=1e-6) - assert np.isclose(result_euler.y, pitch, atol=1e-6) - assert np.isclose(result_euler.z, yaw, atol=1e-6) + euler_vector = Vector3(x=roll, y=pitch, z=yaw) + quaternion = quaternion_from_euler(euler_vector.x, euler_vector.y, euler_vector.z) + pose = Pose( + position=point_from_array(vector_array(Vector3(x=0, y=0, z=0))), orientation=quaternion + ) + result_euler = vector_from_array(quaternion_euler(pose.orientation)) + assert np.isclose(result_euler.x, roll, atol=1e-06) + assert np.isclose(result_euler.y, pitch, atol=1e-06) + assert np.isclose(result_euler.z, yaw, atol=1e-06) def test_pose_zero_position() -> None: """Test pose with zero position vector.""" - # Use manual construction since Vector3.zeros has signature issues - pose = Pose(0.0, 0.0, 0.0) # Position at origin with identity orientation - - assert pose.x == 0.0 - assert pose.y == 0.0 - assert pose.z == 0.0 - assert np.isclose(pose.roll, 0.0, atol=1e-10) - assert np.isclose(pose.pitch, 0.0, atol=1e-10) - assert np.isclose(pose.yaw, 0.0, atol=1e-10) + pose = Pose(position=Point(x=0.0, y=0.0, z=0.0)) + assert pose.position.x == 0.0 + assert pose.position.y == 0.0 + assert pose.position.z == 0.0 + assert np.isclose(quaternion_euler(pose.orientation)[0], 0.0, atol=1e-10) + assert np.isclose(quaternion_euler(pose.orientation)[1], 0.0, atol=1e-10) + assert np.isclose(quaternion_euler(pose.orientation)[2], 0.0, atol=1e-10) def test_pose_unit_vectors() -> None: """Test pose with unit vector positions.""" - # Test unit x vector position - pose_x = Pose(Vector3.unit_x()) - assert pose_x.x == 1.0 - assert pose_x.y == 0.0 - assert pose_x.z == 0.0 - - # Test unit y vector position - pose_y = Pose(Vector3.unit_y()) - assert pose_y.x == 0.0 - assert pose_y.y == 1.0 - assert pose_y.z == 0.0 - - # Test unit z vector position - pose_z = Pose(Vector3.unit_z()) - assert pose_z.x == 0.0 - assert pose_z.y == 0.0 - assert pose_z.z == 1.0 + pose_x = Pose(position=point_from_array(vector_array(Vector3(x=1)))) + assert pose_x.position.x == 1.0 + assert pose_x.position.y == 0.0 + assert pose_x.position.z == 0.0 + pose_y = Pose(position=point_from_array(vector_array(Vector3(y=1)))) + assert pose_y.position.x == 0.0 + assert pose_y.position.y == 1.0 + assert pose_y.position.z == 0.0 + pose_z = Pose(position=point_from_array(vector_array(Vector3(z=1)))) + assert pose_z.position.x == 0.0 + assert pose_z.position.y == 0.0 + assert pose_z.position.z == 1.0 def test_pose_negative_coordinates() -> None: """Test pose with negative coordinates.""" - pose = Pose(-1.0, -2.0, -3.0, -0.1, -0.2, -0.3, 0.9) - - # Position should be negative - assert pose.x == -1.0 - assert pose.y == -2.0 - assert pose.z == -3.0 - - # Orientation should be as specified + pose = Pose( + position=Point(x=-1.0, y=-2.0, z=-3.0), + orientation=Quaternion(x=-0.1, y=-0.2, z=-0.3, w=0.9), + ) + assert pose.position.x == -1.0 + assert pose.position.y == -2.0 + assert pose.position.z == -3.0 assert pose.orientation.x == -0.1 assert pose.orientation.y == -0.2 assert pose.orientation.z == -0.3 @@ -488,13 +404,10 @@ def test_pose_negative_coordinates() -> None: def test_pose_large_coordinates() -> None: """Test pose with large coordinate values.""" large_value = 1000.0 - pose = Pose(large_value, large_value, large_value) - - assert pose.x == large_value - assert pose.y == large_value - assert pose.z == large_value - - # Orientation should still be identity + pose = Pose(position=Point(x=large_value, y=large_value, z=large_value)) + assert pose.position.x == large_value + assert pose.position.y == large_value + assert pose.position.z == large_value assert pose.orientation.x == 0.0 assert pose.orientation.y == 0.0 assert pose.orientation.z == 0.0 @@ -507,13 +420,10 @@ def test_pose_large_coordinates() -> None: ) def test_pose_parametrized_positions(x, y, z) -> None: """Parametrized test for various position values.""" - pose = Pose(x, y, z) - - assert pose.x == x - assert pose.y == y - assert pose.z == z - - # Should have identity orientation + pose = Pose(position=Point(x=x, y=y, z=z)) + assert pose.position.x == x + assert pose.position.y == y + assert pose.position.z == z assert pose.orientation.x == 0.0 assert pose.orientation.y == 0.0 assert pose.orientation.z == 0.0 @@ -523,53 +433,52 @@ def test_pose_parametrized_positions(x, y, z) -> None: @pytest.mark.parametrize( "qx,qy,qz,qw", [ - (0.0, 0.0, 0.0, 1.0), # Identity - (1.0, 0.0, 0.0, 0.0), # 180° around x - (0.0, 1.0, 0.0, 0.0), # 180° around y - (0.0, 0.0, 1.0, 0.0), # 180° around z - (0.5, 0.5, 0.5, 0.5), # Equal components + (0.0, 0.0, 0.0, 1.0), + (1.0, 0.0, 0.0, 0.0), + (0.0, 1.0, 0.0, 0.0), + (0.0, 0.0, 1.0, 0.0), + (0.5, 0.5, 0.5, 0.5), ], ) def test_pose_parametrized_orientations(qx, qy, qz, qw) -> None: """Parametrized test for various orientation values.""" - pose = Pose(0.0, 0.0, 0.0, qx, qy, qz, qw) - - # Position should be at origin - assert pose.x == 0.0 - assert pose.y == 0.0 - assert pose.z == 0.0 - - # Orientation should match + pose = Pose(position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=qx, y=qy, z=qz, w=qw)) + assert pose.position.x == 0.0 + assert pose.position.y == 0.0 + assert pose.position.z == 0.0 assert pose.orientation.x == qx assert pose.orientation.y == qy assert pose.orientation.z == qz assert pose.orientation.w == qw -def test_lcm_encode_decode() -> None: - """Test encoding and decoding of Pose to/from binary LCM format.""" +def test_cdr_encode_decode() -> None: + """Test encoding and decoding of Pose to/from binary CDR format.""" def encodepass() -> None: - pose_source = Pose(1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.9) - binary_msg = pose_source.lcm_encode() - pose_dest = Pose.lcm_decode(binary_msg) + pose_source = Pose( + position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) + ) + binary_msg = pose_source.encode() + pose_dest = Pose.decode(binary_msg) assert isinstance(pose_dest, Pose) assert pose_dest is not pose_source assert pose_dest == pose_source - # Verify we get our custom types back - assert isinstance(pose_dest.position, Vector3) + assert isinstance(pose_dest.position, Point) assert isinstance(pose_dest.orientation, Quaternion) import timeit - print(f"{timeit.timeit(encodepass, number=1000)} ms per cycle") + timeit.timeit(encodepass, number=1000) def test_pickle_encode_decode() -> None: - """Test encoding and decoding of Pose to/from binary LCM format.""" + """Test encoding and decoding of Pose to/from binary CDR format.""" def encodepass() -> None: - pose_source = Pose(1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.9) + pose_source = Pose( + position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) + ) binary_msg = pickle.dumps(pose_source) pose_dest = pickle.loads(binary_msg) assert isinstance(pose_dest, Pose) @@ -578,23 +487,17 @@ def encodepass() -> None: import timeit - print(f"{timeit.timeit(encodepass, number=1000)} ms per cycle") + timeit.timeit(encodepass, number=1000) def test_pose_addition_translation_only() -> None: """Test pose addition with translation only (identity rotations).""" - # Two poses with only translations - pose1 = Pose(1.0, 2.0, 3.0) # First translation - pose2 = Pose(4.0, 5.0, 6.0) # Second translation - - # Adding should combine translations - result = pose1 + pose2 - - assert result.position.x == 5.0 # 1 + 4 - assert result.position.y == 7.0 # 2 + 5 - assert result.position.z == 9.0 # 3 + 6 - - # Orientation should remain identity + pose1 = Pose(position=Point(x=1.0, y=2.0, z=3.0)) + pose2 = Pose(position=Point(x=4.0, y=5.0, z=6.0)) + result = compose_poses(pose1, pose2) + assert result.position.x == 5.0 + assert result.position.y == 7.0 + assert result.position.z == 9.0 assert result.orientation.x == 0.0 assert result.orientation.y == 0.0 assert result.orientation.z == 0.0 @@ -603,23 +506,16 @@ def test_pose_addition_translation_only() -> None: def test_pose_addition_with_rotation() -> None: """Test pose addition with rotation applied to translation.""" - # First pose: at origin, rotated 90 degrees around Z (yaw) - # 90 degree rotation quaternion around Z: (0, 0, sin(pi/4), cos(pi/4)) - angle = np.pi / 2 # 90 degrees - pose1 = Pose(0.0, 0.0, 0.0, 0.0, 0.0, np.sin(angle / 2), np.cos(angle / 2)) - - # Second pose: 1 unit forward (along X in its frame) - pose2 = Pose(1.0, 0.0, 0.0) - - # After rotation, the forward direction should be along Y - result = pose1 + pose2 - - # Position should be rotated + angle = np.pi / 2 + pose1 = Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=np.sin(angle / 2), w=np.cos(angle / 2)), + ) + pose2 = Pose(position=Point(x=1.0, y=0.0, z=0.0)) + result = compose_poses(pose1, pose2) assert np.isclose(result.position.x, 0.0, atol=1e-10) assert np.isclose(result.position.y, 1.0, atol=1e-10) assert np.isclose(result.position.z, 0.0, atol=1e-10) - - # Orientation should be same as pose1 (pose2 has identity rotation) assert np.isclose(result.orientation.x, 0.0, atol=1e-10) assert np.isclose(result.orientation.y, 0.0, atol=1e-10) assert np.isclose(result.orientation.z, np.sin(angle / 2), atol=1e-10) @@ -628,62 +524,45 @@ def test_pose_addition_with_rotation() -> None: def test_pose_addition_rotation_composition() -> None: """Test that rotations are properly composed.""" - # First pose: 45 degrees around Z - angle1 = np.pi / 4 # 45 degrees - pose1 = Pose(0.0, 0.0, 0.0, 0.0, 0.0, np.sin(angle1 / 2), np.cos(angle1 / 2)) - - # Second pose: another 45 degrees around Z - angle2 = np.pi / 4 # 45 degrees - pose2 = Pose(0.0, 0.0, 0.0, 0.0, 0.0, np.sin(angle2 / 2), np.cos(angle2 / 2)) - - # Result should be 90 degrees around Z - result = pose1 + pose2 - - # Check final angle is 90 degrees - expected_angle = angle1 + angle2 # 90 degrees + angle1 = np.pi / 4 + pose1 = Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=np.sin(angle1 / 2), w=np.cos(angle1 / 2)), + ) + angle2 = np.pi / 4 + pose2 = Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=np.sin(angle2 / 2), w=np.cos(angle2 / 2)), + ) + result = compose_poses(pose1, pose2) + expected_angle = angle1 + angle2 expected_qz = np.sin(expected_angle / 2) expected_qw = np.cos(expected_angle / 2) - assert np.isclose(result.orientation.z, expected_qz, atol=1e-10) assert np.isclose(result.orientation.w, expected_qw, atol=1e-10) def test_pose_addition_full_transform() -> None: """Test full pose composition with translation and rotation.""" - # Robot pose: at (2, 1, 0), facing 90 degrees left (positive yaw) - robot_yaw = np.pi / 2 # 90 degrees - robot_pose = Pose(2.0, 1.0, 0.0, 0.0, 0.0, np.sin(robot_yaw / 2), np.cos(robot_yaw / 2)) - - # Object in robot frame: 3 units forward, 1 unit right - object_in_robot = Pose(3.0, -1.0, 0.0) - - # Compose to get object in world frame - object_in_world = robot_pose + object_in_robot - - # Robot is facing left (90 degrees), so: - # - Robot's forward (X) is world's negative Y - # - Robot's right (negative Y) is world's X - # So object should be at: robot_pos + rotated_offset - # rotated_offset: (3, -1) rotated 90° CCW = (1, 3) - assert np.isclose(object_in_world.position.x, 3.0, atol=1e-10) # 2 + 1 - assert np.isclose(object_in_world.position.y, 4.0, atol=1e-10) # 1 + 3 + robot_yaw = np.pi / 2 + robot_pose = Pose( + position=Point(x=2.0, y=1.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=np.sin(robot_yaw / 2), w=np.cos(robot_yaw / 2)), + ) + object_in_robot = Pose(position=Point(x=3.0, y=-1.0, z=0.0)) + object_in_world = compose_poses(robot_pose, object_in_robot) + assert np.isclose(object_in_world.position.x, 3.0, atol=1e-10) + assert np.isclose(object_in_world.position.y, 4.0, atol=1e-10) assert np.isclose(object_in_world.position.z, 0.0, atol=1e-10) - - # Orientation should match robot's orientation (object has no rotation) - assert np.isclose(object_in_world.yaw, robot_yaw, atol=1e-10) + assert np.isclose(quaternion_euler(object_in_world.orientation)[2], robot_yaw, atol=1e-10) def test_pose_addition_chain() -> None: """Test chaining multiple pose additions.""" - # Create a chain of transformations - pose1 = Pose(1.0, 0.0, 0.0) # Move 1 unit in X - pose2 = Pose(0.0, 1.0, 0.0) # Move 1 unit in Y (relative to pose1) - pose3 = Pose(0.0, 0.0, 1.0) # Move 1 unit in Z (relative to pose1+pose2) - - # Chain them together - result = pose1 + pose2 + pose3 - - # Should accumulate all translations + pose1 = Pose(position=Point(x=1.0, y=0.0, z=0.0)) + pose2 = Pose(position=Point(x=0.0, y=1.0, z=0.0)) + pose3 = Pose(position=Point(x=0.0, y=0.0, z=1.0)) + result = compose_poses(compose_poses(pose1, pose2), pose3) assert result.position.x == 1.0 assert result.position.y == 1.0 assert result.position.z == 1.0 @@ -691,18 +570,26 @@ def test_pose_addition_chain() -> None: def test_pose_addition_with_convertible() -> None: """Test pose addition with convertible types.""" - pose1 = Pose(1.0, 2.0, 3.0) - - # Add with tuple + pose1 = Pose(position=Point(x=1.0, y=2.0, z=3.0)) pose_tuple = ([4.0, 5.0, 6.0], [0.0, 0.0, 0.0, 1.0]) - result1 = pose1 + pose_tuple + result1 = compose_poses( + pose1, + Pose( + position=point_from_array(pose_tuple[0]), + orientation=quaternion_from_array(pose_tuple[1]), + ), + ) assert result1.position.x == 5.0 assert result1.position.y == 7.0 assert result1.position.z == 9.0 - - # Add with dict pose_dict = {"position": [1.0, 0.0, 0.0], "orientation": [0.0, 0.0, 0.0, 1.0]} - result2 = pose1 + pose_dict + result2 = compose_poses( + pose1, + Pose( + position=point_from_array(pose_dict["position"]), + orientation=quaternion_from_array(pose_dict["orientation"]), + ), + ) assert result2.position.x == 2.0 assert result2.position.y == 2.0 assert result2.position.z == 3.0 @@ -710,12 +597,11 @@ def test_pose_addition_with_convertible() -> None: def test_pose_identity_addition() -> None: """Test that adding identity pose leaves pose unchanged.""" - pose = Pose(1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.9) - identity = Pose() # Identity pose at origin - - result = pose + identity - - # Should be unchanged + pose = Pose( + position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) + ) + identity = Pose() + result = compose_poses(pose, identity) assert result.position.x == pose.position.x assert result.position.y == pose.position.y assert result.position.z == pose.position.z @@ -727,23 +613,15 @@ def test_pose_identity_addition() -> None: def test_pose_addition_3d_rotation() -> None: """Test pose addition with 3D rotations.""" - # First pose: rotated around X axis (roll) - roll = np.pi / 4 # 45 degrees - pose1 = Pose(1.0, 0.0, 0.0, np.sin(roll / 2), 0.0, 0.0, np.cos(roll / 2)) - - # Second pose: movement along Y and Z in local frame - pose2 = Pose(0.0, 1.0, 1.0) - - # Compose transformations - result = pose1 + pose2 - - # The Y and Z movement should be rotated around X - # After 45° rotation around X: - # - Local Y -> world Y * cos(45°) - Z * sin(45°) - # - Local Z -> world Y * sin(45°) + Z * cos(45°) + roll = np.pi / 4 + pose1 = Pose( + position=Point(x=1.0, y=0.0, z=0.0), + orientation=Quaternion(x=np.sin(roll / 2), y=0.0, z=0.0, w=np.cos(roll / 2)), + ) + pose2 = Pose(position=Point(x=0.0, y=1.0, z=1.0)) + result = compose_poses(pose1, pose2) cos45 = np.cos(roll) sin45 = np.sin(roll) - - assert np.isclose(result.position.x, 1.0, atol=1e-10) # X unchanged + assert np.isclose(result.position.x, 1.0, atol=1e-10) assert np.isclose(result.position.y, cos45 - sin45, atol=1e-10) assert np.isclose(result.position.z, sin45 + cos45, atol=1e-10) diff --git a/dimos/msgs/geometry_msgs/test_PoseStamped.py b/dimos/msgs/geometry_msgs/test_PoseStamped.py index a486f333031..d43b2981eaa 100644 --- a/dimos/msgs/geometry_msgs/test_PoseStamped.py +++ b/dimos/msgs/geometry_msgs/test_PoseStamped.py @@ -13,43 +13,52 @@ # limitations under the License. import pickle -import time -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.std_msgs.msg import Header +from rosbags.typesys import Stores, get_typestore +from dimos.msgs.time import time_from_nanoseconds -def test_lcm_encode_decode() -> None: - """Test encoding and decoding of Pose to/from binary LCM format.""" - pose_source = PoseStamped( - ts=time.time(), - position=(1.0, 2.0, 3.0), - orientation=(0.1, 0.2, 0.3, 0.9), +def test_cdr_encode_decode() -> None: + source = PoseStamped( + header=Header(stamp=time_from_nanoseconds(1234567890123456789)), + pose=Pose( + position=Point(x=1, y=2, z=3), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) + ), ) - binary_msg = pose_source.lcm_encode() - pose_dest = PoseStamped.lcm_decode(binary_msg) - - assert isinstance(pose_dest, PoseStamped) - assert pose_dest is not pose_source - - print(pose_source.position) - print(pose_source.orientation) - - print(pose_dest.position) - print(pose_dest.orientation) - assert pose_dest == pose_source + binary = source.encode() + decoded = PoseStamped.decode(binary) + assert isinstance(decoded, PoseStamped) + assert decoded is not source + assert decoded.encode() == binary + independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(binary, source.msg_name) + assert (independent.header.stamp.sec, independent.header.stamp.nanosec) == ( + 1234567890, + 123456789, + ) + assert ( + independent.pose.position.x, + independent.pose.position.y, + independent.pose.position.z, + ) == (1, 2, 3) + assert ( + independent.pose.orientation.x, + independent.pose.orientation.y, + independent.pose.orientation.z, + independent.pose.orientation.w, + ) == (0.1, 0.2, 0.3, 0.9) def test_pickle_encode_decode() -> None: - """Test encoding and decoding of PoseStamped to/from binary LCM format.""" - - pose_source = PoseStamped( - ts=time.time(), - position=(1.0, 2.0, 3.0), - orientation=(0.1, 0.2, 0.3, 0.9), + source = PoseStamped( + header=Header(stamp=time_from_nanoseconds(1234567890123456789)), + pose=Pose( + position=Point(x=1, y=2, z=3), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) + ), ) - binary_msg = pickle.dumps(pose_source) - pose_dest = pickle.loads(binary_msg) - assert isinstance(pose_dest, PoseStamped) - assert pose_dest is not pose_source - assert pose_dest == pose_source + decoded = pickle.loads(pickle.dumps(source)) + assert isinstance(decoded, PoseStamped) + assert decoded is not source + assert decoded.encode() == source.encode() diff --git a/dimos/msgs/geometry_msgs/test_PoseWithCovariance.py b/dimos/msgs/geometry_msgs/test_PoseWithCovariance.py index f6936db9f7a..fa52d25c80a 100644 --- a/dimos/msgs/geometry_msgs/test_PoseWithCovariance.py +++ b/dimos/msgs/geometry_msgs/test_PoseWithCovariance.py @@ -12,311 +12,111 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos_lcm.geometry_msgs import PoseWithCovariance as LCMPoseWithCovariance -import numpy as np -import pytest - -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseWithCovariance import PoseWithCovariance - - -def test_pose_with_covariance_default_init() -> None: - """Test that default initialization creates a pose at origin with zero covariance.""" - pose_cov = PoseWithCovariance() - - # Pose should be at origin with identity orientation - assert pose_cov.pose.position.x == 0.0 - assert pose_cov.pose.position.y == 0.0 - assert pose_cov.pose.position.z == 0.0 - assert pose_cov.pose.orientation.x == 0.0 - assert pose_cov.pose.orientation.y == 0.0 - assert pose_cov.pose.orientation.z == 0.0 - assert pose_cov.pose.orientation.w == 1.0 - - # Covariance should be all zeros - assert np.all(pose_cov.covariance == 0.0) - assert pose_cov.covariance.shape == (36,) - - -def test_pose_with_covariance_pose_init() -> None: - """Test initialization with a Pose object.""" - pose = Pose(1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.9) - pose_cov = PoseWithCovariance(pose) - - # Pose should match - assert pose_cov.pose.position.x == 1.0 - assert pose_cov.pose.position.y == 2.0 - assert pose_cov.pose.position.z == 3.0 - assert pose_cov.pose.orientation.x == 0.1 - assert pose_cov.pose.orientation.y == 0.2 - assert pose_cov.pose.orientation.z == 0.3 - assert pose_cov.pose.orientation.w == 0.9 - - # Covariance should be zeros by default - assert np.all(pose_cov.covariance == 0.0) - - -def test_pose_with_covariance_pose_and_covariance_init() -> None: - """Test initialization with pose and covariance.""" - pose = Pose(1.0, 2.0, 3.0) - covariance = np.arange(36, dtype=float) - pose_cov = PoseWithCovariance(pose, covariance) - - # Pose should match - assert pose_cov.pose.position.x == 1.0 - assert pose_cov.pose.position.y == 2.0 - assert pose_cov.pose.position.z == 3.0 - - # Covariance should match - assert np.array_equal(pose_cov.covariance, covariance) - - -def test_pose_with_covariance_list_covariance() -> None: - """Test initialization with covariance as a list.""" - pose = Pose(1.0, 2.0, 3.0) - covariance_list = list(range(36)) - pose_cov = PoseWithCovariance(pose, covariance_list) - - # Covariance should be converted to numpy array - assert isinstance(pose_cov.covariance, np.ndarray) - assert np.array_equal(pose_cov.covariance, np.array(covariance_list)) - - -def test_pose_with_covariance_copy_init() -> None: - """Test copy constructor.""" - pose = Pose(1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.9) - covariance = np.arange(36, dtype=float) - original = PoseWithCovariance(pose, covariance) - copy = PoseWithCovariance(original) - - # Should be equal but not the same object - assert copy == original - assert copy is not original - assert copy.pose is not original.pose - assert copy.covariance is not original.covariance - - # Modify original to ensure they're independent - original.covariance[0] = 999.0 - assert copy.covariance[0] != 999.0 - - -def test_pose_with_covariance_lcm_init() -> None: - """Test initialization from LCM message.""" - lcm_msg = LCMPoseWithCovariance() - lcm_msg.pose.position.x = 1.0 - lcm_msg.pose.position.y = 2.0 - lcm_msg.pose.position.z = 3.0 - lcm_msg.pose.orientation.x = 0.1 - lcm_msg.pose.orientation.y = 0.2 - lcm_msg.pose.orientation.z = 0.3 - lcm_msg.pose.orientation.w = 0.9 - lcm_msg.covariance = list(range(36)) - - pose_cov = PoseWithCovariance(lcm_msg) - - # Pose should match - assert pose_cov.pose.position.x == 1.0 - assert pose_cov.pose.position.y == 2.0 - assert pose_cov.pose.position.z == 3.0 - assert pose_cov.pose.orientation.x == 0.1 - assert pose_cov.pose.orientation.y == 0.2 - assert pose_cov.pose.orientation.z == 0.3 - assert pose_cov.pose.orientation.w == 0.9 - - # Covariance should match - assert np.array_equal(pose_cov.covariance, np.arange(36)) - - -def test_pose_with_covariance_dict_init() -> None: - """Test initialization from dictionary.""" - pose_dict = {"pose": Pose(1.0, 2.0, 3.0), "covariance": list(range(36))} - pose_cov = PoseWithCovariance(pose_dict) - - assert pose_cov.pose.position.x == 1.0 - assert pose_cov.pose.position.y == 2.0 - assert pose_cov.pose.position.z == 3.0 - assert np.array_equal(pose_cov.covariance, np.arange(36)) - - -def test_pose_with_covariance_dict_init_no_covariance() -> None: - """Test initialization from dictionary without covariance.""" - pose_dict = {"pose": Pose(1.0, 2.0, 3.0)} - pose_cov = PoseWithCovariance(pose_dict) - - assert pose_cov.pose.position.x == 1.0 - assert np.all(pose_cov.covariance == 0.0) - - -def test_pose_with_covariance_tuple_init() -> None: - """Test initialization from tuple.""" - pose = Pose(1.0, 2.0, 3.0) - covariance = np.arange(36, dtype=float) - pose_tuple = (pose, covariance) - pose_cov = PoseWithCovariance(pose_tuple) - - assert pose_cov.pose.position.x == 1.0 - assert pose_cov.pose.position.y == 2.0 - assert pose_cov.pose.position.z == 3.0 - assert np.array_equal(pose_cov.covariance, covariance) - - -def test_pose_with_covariance_properties() -> None: - """Test convenience properties.""" - pose = Pose(1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.9) - pose_cov = PoseWithCovariance(pose) - - # Position properties - assert pose_cov.x == 1.0 - assert pose_cov.y == 2.0 - assert pose_cov.z == 3.0 - assert pose_cov.position.x == 1.0 - assert pose_cov.position.y == 2.0 - assert pose_cov.position.z == 3.0 - - # Orientation properties - assert pose_cov.orientation.x == 0.1 - assert pose_cov.orientation.y == 0.2 - assert pose_cov.orientation.z == 0.3 - assert pose_cov.orientation.w == 0.9 - - # Euler angle properties - assert pose_cov.roll == pose.roll - assert pose_cov.pitch == pose.pitch - assert pose_cov.yaw == pose.yaw - - -def test_pose_with_covariance_matrix_property() -> None: - """Test covariance matrix property.""" - pose = Pose() - covariance_array = np.arange(36, dtype=float) - pose_cov = PoseWithCovariance(pose, covariance_array) - - # Get as matrix - cov_matrix = pose_cov.covariance_matrix - assert cov_matrix.shape == (6, 6) - assert cov_matrix[0, 0] == 0.0 - assert cov_matrix[5, 5] == 35.0 - - # Set from matrix - new_matrix = np.eye(6) * 2.0 - pose_cov.covariance_matrix = new_matrix - assert np.array_equal(pose_cov.covariance[:6], [2.0, 0.0, 0.0, 0.0, 0.0, 0.0]) - - -def test_pose_with_covariance_repr() -> None: - """Test string representation.""" - pose = Pose(1.234, 2.567, 3.891) - pose_cov = PoseWithCovariance(pose) - - repr_str = repr(pose_cov) - assert "PoseWithCovariance" in repr_str - assert "pose=" in repr_str - assert "covariance=" in repr_str - assert "36 elements" in repr_str - - -def test_pose_with_covariance_str() -> None: - """Test string formatting.""" - pose = Pose(1.234, 2.567, 3.891) - covariance = np.eye(6).flatten() - pose_cov = PoseWithCovariance(pose, covariance) - - str_repr = str(pose_cov) - assert "PoseWithCovariance" in str_repr - assert "1.234" in str_repr - assert "2.567" in str_repr - assert "3.891" in str_repr - assert "cov_trace" in str_repr - assert "6.000" in str_repr # Trace of identity matrix is 6 - - -def test_pose_with_covariance_equality() -> None: - """Test equality comparison.""" - pose1 = Pose(1.0, 2.0, 3.0) - cov1 = np.arange(36, dtype=float) - pose_cov1 = PoseWithCovariance(pose1, cov1) - - pose2 = Pose(1.0, 2.0, 3.0) - cov2 = np.arange(36, dtype=float) - pose_cov2 = PoseWithCovariance(pose2, cov2) - - # Equal - assert pose_cov1 == pose_cov2 - - # Different pose - pose3 = Pose(1.1, 2.0, 3.0) - pose_cov3 = PoseWithCovariance(pose3, cov1) - assert pose_cov1 != pose_cov3 - - # Different covariance - cov3 = np.arange(36, dtype=float) + 1 - pose_cov4 = PoseWithCovariance(pose1, cov3) - assert pose_cov1 != pose_cov4 - - # Different type - assert pose_cov1 != "not a pose" - assert pose_cov1 is not None - - -def test_pose_with_covariance_lcm_encode_decode() -> None: - """Test LCM encoding and decoding.""" - pose = Pose(1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.9) - covariance = np.arange(36, dtype=float) - source = PoseWithCovariance(pose, covariance) - - # Encode and decode - binary_msg = source.lcm_encode() - decoded = PoseWithCovariance.lcm_decode(binary_msg) - - # Should be equal - assert decoded == source - assert isinstance(decoded, PoseWithCovariance) - assert isinstance(decoded.pose, Pose) - assert isinstance(decoded.covariance, np.ndarray) - - -def test_pose_with_covariance_zero_covariance() -> None: - """Test with zero covariance matrix.""" - pose = Pose(1.0, 2.0, 3.0) - pose_cov = PoseWithCovariance(pose) - - assert np.all(pose_cov.covariance == 0.0) - assert np.trace(pose_cov.covariance_matrix) == 0.0 - - -def test_pose_with_covariance_diagonal_covariance() -> None: - """Test with diagonal covariance matrix.""" - pose = Pose() - covariance = np.zeros(36) - # Set diagonal elements - for i in range(6): - covariance[i * 6 + i] = i + 1 - - pose_cov = PoseWithCovariance(pose, covariance) - - cov_matrix = pose_cov.covariance_matrix - assert np.trace(cov_matrix) == sum(range(1, 7)) # 1+2+3+4+5+6 = 21 - - # Check diagonal elements - for i in range(6): - assert cov_matrix[i, i] == i + 1 - - # Check off-diagonal elements are zero - for i in range(6): - for j in range(6): - if i != j: - assert cov_matrix[i, j] == 0.0 +"""Generated covariance values, with explicit NumPy matrix operations. +Legacy polymorphic positional constructors and presentation methods are retired. +""" +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseWithCovariance, Quaternion +import numpy as np +import pytest +from rosbags.typesys import Stores, get_typestore + + +def test_default_fields_and_covariance() -> None: + source = PoseWithCovariance() + value = source.pose + np.testing.assert_array_equal( + [ + value.position.x, + value.position.y, + value.position.z, + value.orientation.x, + value.orientation.y, + value.orientation.z, + value.orientation.w, + ], + [0, 0, 0, 0, 0, 0, 1], + ) + assert np.asarray(source.covariance).shape == (36,) + np.testing.assert_array_equal(source.covariance, np.zeros(36)) + + +@pytest.mark.parametrize("as_list", [False, True]) @pytest.mark.parametrize( - "x,y,z", - [(0.0, 0.0, 0.0), (1.0, 2.0, 3.0), (-1.0, -2.0, -3.0), (100.0, -100.0, 0.0)], + "covariance", + [np.zeros(36), np.arange(36, dtype=float), np.eye(6).ravel(), np.diag(np.arange(1, 7)).ravel()], ) -def test_pose_with_covariance_parametrized_positions(x, y, z) -> None: - """Parametrized test for various position values.""" - pose = Pose(x, y, z) - pose_cov = PoseWithCovariance(pose) - - assert pose_cov.x == x - assert pose_cov.y == y - assert pose_cov.z == z +def test_explicit_construction_and_independent_cdr(as_list: bool, covariance: np.ndarray) -> None: + source = PoseWithCovariance( + pose=Pose( + position=Point(x=1, y=2, z=3), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) + ), + covariance=covariance.tolist() if as_list else covariance, + ) + decoded = PoseWithCovariance.decode(source.encode()) + independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr( + source.encode(), PoseWithCovariance.msg_name + ) + for result in (source, decoded, independent): + value = result.pose + np.testing.assert_array_equal( + [ + value.position.x, + value.position.y, + value.position.z, + value.orientation.x, + value.orientation.y, + value.orientation.z, + value.orientation.w, + ], + [1, 2, 3, 0.1, 0.2, 0.3, 0.9], + ) + np.testing.assert_array_equal(result.covariance, covariance) + matrix = np.asarray(result.covariance).reshape(6, 6) + for row in range(6): + for col in range(6): + assert matrix[row, col] == covariance[row * 6 + col] + assert np.trace(matrix) == np.trace(covariance.reshape(6, 6)) + + +def test_copy_equality_and_independent_storage() -> None: + original = PoseWithCovariance( + pose=Pose( + position=Point(x=1, y=2, z=3), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) + ), + covariance=np.arange(36, dtype=float), + ) + copied = PoseWithCovariance.decode(original.encode()) + assert copied == original + assert copied is not original + assert copied.pose is not original.pose + assert copied.covariance is not original.covariance + original.covariance[0] = 999 + assert copied.covariance[0] == 0 + assert copied != original + assert copied != "not a message" + assert copied is not None + + +def test_matrix_assignment() -> None: + source = PoseWithCovariance(covariance=np.arange(36, dtype=float)) + matrix = np.asarray(source.covariance).reshape(6, 6) + assert matrix[0, 0] == 0 + assert matrix[5, 5] == 35 + source.covariance = (np.eye(6) * 2).ravel() + np.testing.assert_array_equal(list(source.covariance)[:6], [2, 0, 0, 0, 0, 0]) + assert np.trace(np.asarray(source.covariance).reshape(6, 6)) == 12 + + +@pytest.mark.parametrize("size", [0, 35, 37]) +def test_fixed_covariance_rejects_invalid_length(size: int) -> None: + with pytest.raises((ValueError, TypeError, RuntimeError)): + PoseWithCovariance(covariance=np.zeros(size)) + + +@pytest.mark.parametrize("xyz", [(0, 0, 0), (1, 2, 3), (-1, -2, -3), (100, -100, 0)]) +def test_parameterized_values(xyz: tuple[float, float, float]) -> None: + source = PoseWithCovariance(pose=Pose(position=Point(x=xyz[0], y=xyz[1], z=xyz[2]))) + position = source.pose.position + assert (position.x, position.y, position.z) == xyz diff --git a/dimos/msgs/geometry_msgs/test_PoseWithCovarianceStamped.py b/dimos/msgs/geometry_msgs/test_PoseWithCovarianceStamped.py index e72b96ad529..e2c1a33697d 100644 --- a/dimos/msgs/geometry_msgs/test_PoseWithCovarianceStamped.py +++ b/dimos/msgs/geometry_msgs/test_PoseWithCovarianceStamped.py @@ -12,211 +12,118 @@ # See the License for the specific language governing permissions and # limitations under the License. -import time - -import numpy as np - -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseWithCovariance import PoseWithCovariance -from dimos.msgs.geometry_msgs.PoseWithCovarianceStamped import PoseWithCovarianceStamped - - -def test_pose_with_covariance_stamped_default_init() -> None: - """Test default initialization.""" - pose_cov_stamped = PoseWithCovarianceStamped() - - # Should have current timestamp - assert pose_cov_stamped.ts > 0 - assert pose_cov_stamped.frame_id == "" +"""Nested generated Pose covariance with explicit stamps and array conversion. - # Pose should be at origin with identity orientation - assert pose_cov_stamped.pose.position.x == 0.0 - assert pose_cov_stamped.pose.position.y == 0.0 - assert pose_cov_stamped.pose.position.z == 0.0 - assert pose_cov_stamped.pose.orientation.w == 1.0 +Legacy inheritance, custom repr, and implicit wall-clock defaults are retired; +field values, covariance matrices, and time conversion remain checked explicitly. +""" - # Covariance should be all zeros - assert np.all(pose_cov_stamped.covariance == 0.0) - - -def test_pose_with_covariance_stamped_with_timestamp() -> None: - """Test initialization with specific timestamp.""" - ts = 1234567890.123456 - frame_id = "base_link" - pose_cov_stamped = PoseWithCovarianceStamped(ts=ts, frame_id=frame_id) - - assert pose_cov_stamped.ts == ts - assert pose_cov_stamped.frame_id == frame_id - - -def test_pose_with_covariance_stamped_with_pose() -> None: - """Test initialization with pose.""" - ts = 1234567890.123456 - frame_id = "map" - pose = Pose(1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.9) - covariance = np.arange(36, dtype=float) +import pickle +import time - pose_cov_stamped = PoseWithCovarianceStamped( - ts=ts, frame_id=frame_id, pose=pose, covariance=covariance +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseWithCovariance, + PoseWithCovarianceStamped, + Quaternion, +) +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest +from rosbags.typesys import Stores, get_typestore + +from dimos.msgs.time import header_now, time_from_nanoseconds, time_from_seconds, to_nanoseconds + + +def test_generated_defaults() -> None: + source = PoseWithCovarianceStamped() + assert source.header.frame_id == "" + assert to_nanoseconds(source.header.stamp) == 0 + assert isinstance(source.pose, PoseWithCovariance) + assert isinstance(source.pose.pose, Pose) + np.testing.assert_array_equal(source.pose.covariance, np.zeros(36)) + value = source.pose.pose + np.testing.assert_array_equal( + [ + value.position.x, + value.position.y, + value.position.z, + value.orientation.x, + value.orientation.y, + value.orientation.z, + value.orientation.w, + ], + [0, 0, 0, 0, 0, 0, 1], ) - assert pose_cov_stamped.ts == ts - assert pose_cov_stamped.frame_id == frame_id - assert pose_cov_stamped.pose.position.x == 1.0 - assert pose_cov_stamped.pose.position.y == 2.0 - assert pose_cov_stamped.pose.position.z == 3.0 - assert np.array_equal(pose_cov_stamped.covariance, covariance) - -def test_pose_with_covariance_stamped_properties() -> None: - """Test convenience properties.""" - pose = Pose(1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.9) - covariance = np.eye(6).flatten() - pose_cov_stamped = PoseWithCovarianceStamped( - ts=1234567890.0, frame_id="odom", pose=pose, covariance=covariance +@pytest.mark.parametrize("stamp", [0, 1234567890123456789]) +@pytest.mark.parametrize( + "covariance", + [np.zeros(36), np.eye(6).ravel(), np.arange(36, dtype=float), np.diag(np.arange(1, 7)).ravel()], +) +def test_fields_covariance_and_independent_cdr(stamp: int, covariance: np.ndarray) -> None: + source = PoseWithCovarianceStamped( + header=Header(stamp=time_from_nanoseconds(stamp), frame_id="camera_link"), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=1, y=2, z=3), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) + ), + covariance=covariance, + ), ) - - # Position properties - assert pose_cov_stamped.x == 1.0 - assert pose_cov_stamped.y == 2.0 - assert pose_cov_stamped.z == 3.0 - - # Orientation properties - assert pose_cov_stamped.orientation.x == 0.1 - assert pose_cov_stamped.orientation.y == 0.2 - assert pose_cov_stamped.orientation.z == 0.3 - assert pose_cov_stamped.orientation.w == 0.9 - - # Euler angles - assert pose_cov_stamped.roll == pose.roll - assert pose_cov_stamped.pitch == pose.pitch - assert pose_cov_stamped.yaw == pose.yaw - - # Covariance matrix - cov_matrix = pose_cov_stamped.covariance_matrix - assert cov_matrix.shape == (6, 6) - assert np.trace(cov_matrix) == 6.0 - - -def test_pose_with_covariance_stamped_str() -> None: - """Test string representation.""" - pose = Pose(1.234, 2.567, 3.891) - covariance = np.eye(6).flatten() * 2.0 - pose_cov_stamped = PoseWithCovarianceStamped( - ts=1234567890.0, frame_id="world", pose=pose, covariance=covariance + decoded = PoseWithCovarianceStamped.decode(source.encode()) + independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr( + source.encode(), PoseWithCovarianceStamped.msg_name ) - - str_repr = str(pose_cov_stamped) - assert "PoseWithCovarianceStamped" in str_repr - assert "1.234" in str_repr - assert "2.567" in str_repr - assert "3.891" in str_repr - assert "cov_trace" in str_repr - assert "12.000" in str_repr # Trace of 2*identity is 12 - - -def test_pose_with_covariance_stamped_lcm_encode_decode() -> None: - """Test LCM encoding and decoding.""" - ts = 1234567890.123456 - frame_id = "camera_link" - pose = Pose(1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.9) - covariance = np.arange(36, dtype=float) - - source = PoseWithCovarianceStamped(ts=ts, frame_id=frame_id, pose=pose, covariance=covariance) - - # Encode and decode - binary_msg = source.lcm_encode() - decoded = PoseWithCovarianceStamped.lcm_decode(binary_msg) - - # Check timestamp (may lose some precision) - assert abs(decoded.ts - ts) < 1e-6 - assert decoded.frame_id == frame_id - - # Check pose - assert decoded.pose.position.x == 1.0 - assert decoded.pose.position.y == 2.0 - assert decoded.pose.position.z == 3.0 - assert decoded.pose.orientation.x == 0.1 - assert decoded.pose.orientation.y == 0.2 - assert decoded.pose.orientation.z == 0.3 - assert decoded.pose.orientation.w == 0.9 - - # Check covariance - assert np.array_equal(decoded.covariance, covariance) - - -def test_pose_with_covariance_stamped_zero_timestamp() -> None: - """Test that an explicit zero timestamp is kept, and an omitted one stamps now.""" - assert PoseWithCovarianceStamped(ts=0.0).ts == 0.0 - - before = time.time() - assert before <= PoseWithCovarianceStamped().ts <= time.time() - - -def test_pose_with_covariance_stamped_inheritance() -> None: - """Test that it properly inherits from PoseWithCovariance and Timestamped.""" - pose = Pose(1.0, 2.0, 3.0) - covariance = np.eye(6).flatten() - pose_cov_stamped = PoseWithCovarianceStamped( - ts=1234567890.0, frame_id="test", pose=pose, covariance=covariance - ) - - # Should be instance of parent classes - assert isinstance(pose_cov_stamped, PoseWithCovariance) - - # Should have Timestamped attributes - assert hasattr(pose_cov_stamped, "ts") - assert hasattr(pose_cov_stamped, "frame_id") - - # Should have PoseWithCovariance attributes - assert hasattr(pose_cov_stamped, "pose") - assert hasattr(pose_cov_stamped, "covariance") - - -def test_pose_with_covariance_stamped_sec_nsec() -> None: - """Test the sec_nsec helper function.""" - from dimos.msgs.geometry_msgs.PoseWithCovarianceStamped import sec_nsec - - # Test integer seconds - s, ns = sec_nsec(1234567890.0) - assert s == 1234567890 - assert ns == 0 - - # Test fractional seconds - s, ns = sec_nsec(1234567890.123456789) - assert s == 1234567890 - assert abs(ns - 123456789) < 100 # Allow small rounding error - - # Test small fractional seconds - s, ns = sec_nsec(0.000000001) - assert s == 0 - assert ns == 1 - - # Test large timestamp - s, ns = sec_nsec(9999999999.999999999) - # Due to floating point precision, this might round to 10000000000 - assert s in [9999999999, 10000000000] - if s == 9999999999: - assert abs(ns - 999999999) < 10 - else: - assert ns == 0 - - -def test_pose_with_covariance_stamped_different_covariances() -> None: - """Test with different covariance patterns.""" - pose = Pose(1.0, 2.0, 3.0) - - # Zero covariance - zero_cov = np.zeros(36) - pose_cov1 = PoseWithCovarianceStamped(pose=pose, covariance=zero_cov) - assert np.all(pose_cov1.covariance == 0.0) - - # Identity covariance - identity_cov = np.eye(6).flatten() - pose_cov2 = PoseWithCovarianceStamped(pose=pose, covariance=identity_cov) - assert np.trace(pose_cov2.covariance_matrix) == 6.0 - - # Full covariance - full_cov = np.random.rand(36) - pose_cov3 = PoseWithCovarianceStamped(pose=pose, covariance=full_cov) - assert np.array_equal(pose_cov3.covariance, full_cov) + for result in (decoded, independent): + assert result.header.frame_id == "camera_link" + assert result.header.stamp.sec * 1000000000 + result.header.stamp.nanosec == stamp + value = result.pose.pose + np.testing.assert_array_equal( + [ + value.position.x, + value.position.y, + value.position.z, + value.orientation.x, + value.orientation.y, + value.orientation.z, + value.orientation.w, + ], + [1, 2, 3, 0.1, 0.2, 0.3, 0.9], + ) + np.testing.assert_array_equal(result.pose.covariance, covariance) + matrix = np.asarray(result.pose.covariance).reshape(6, 6) + assert matrix.shape == (6, 6) + assert np.trace(matrix) == np.trace(covariance.reshape(6, 6)) + copied = pickle.loads(pickle.dumps(source)) + assert copied.encode() == source.encode() + copied.pose.covariance[0] = 999 + assert source.pose.covariance[0] == covariance[0] + + +def test_explicit_current_header() -> None: + before = time.time_ns() + source = PoseWithCovarianceStamped(header=header_now("test")) + assert before <= to_nanoseconds(source.header.stamp) <= time.time_ns() + assert source.header.frame_id == "test" + + +@pytest.mark.parametrize( + "seconds,sec,nanosec", + [ + (1234567890.0, 1234567890, 0), + (1234567890.123456789, 1234567890, 123456789), + (0.000000001, 0, 1), + ], +) +def test_explicit_seconds_conversion(seconds: float, sec: int, nanosec: int) -> None: + stamp = time_from_seconds(seconds) + assert stamp.sec == sec + assert abs(stamp.nanosec - nanosec) < 100 + + +def test_seconds_outside_ros_int32_range_rejected() -> None: + with pytest.raises((TypeError, ValueError, OverflowError)): + time_from_seconds(9999999999.999999999) diff --git a/dimos/msgs/geometry_msgs/test_Quaternion.py b/dimos/msgs/geometry_msgs/test_Quaternion.py index 21c1e8caebf..06ba7e1ee2e 100644 --- a/dimos/msgs/geometry_msgs/test_Quaternion.py +++ b/dimos/msgs/geometry_msgs/test_Quaternion.py @@ -12,11 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos_lcm.geometry_msgs import Quaternion as LCMQuaternion +from dimos_generated.geometry_msgs.msg import Quaternion, Vector3 import numpy as np import pytest -from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry import ( + normalized_quaternion, + quaternion_array, + quaternion_conjugate, + quaternion_euler, + quaternion_from_array, + quaternion_inverse, + quaternion_product, + rotate_vector, + vector_from_array, +) def test_quaternion_default_init() -> None: @@ -26,33 +36,27 @@ def test_quaternion_default_init() -> None: assert q.y == 0.0 assert q.z == 0.0 assert q.w == 1.0 - assert q.to_tuple() == (0.0, 0.0, 0.0, 1.0) + assert tuple(quaternion_array(q)) == (0.0, 0.0, 0.0, 1.0) def test_quaternion_component_init() -> None: """Test initialization with four float components (x, y, z, w).""" - q = Quaternion(0.5, 0.5, 0.5, 0.5) + q = Quaternion(x=0.5, y=0.5, z=0.5, w=0.5) assert q.x == 0.5 assert q.y == 0.5 assert q.z == 0.5 assert q.w == 0.5 - - # Test with different values - q2 = Quaternion(1.0, 2.0, 3.0, 4.0) + q2 = Quaternion(x=1.0, y=2.0, z=3.0, w=4.0) assert q2.x == 1.0 assert q2.y == 2.0 assert q2.z == 3.0 assert q2.w == 4.0 - - # Test with negative values - q3 = Quaternion(-1.0, -2.0, -3.0, -4.0) + q3 = Quaternion(x=-1.0, y=-2.0, z=-3.0, w=-4.0) assert q3.x == -1.0 assert q3.y == -2.0 assert q3.z == -3.0 assert q3.w == -4.0 - - # Test with integers (should convert to float) - q4 = Quaternion(1, 2, 3, 4) + q4 = Quaternion(x=1, y=2, z=3, w=4) assert q4.x == 1.0 assert q4.y == 2.0 assert q4.z == 3.0 @@ -62,85 +66,67 @@ def test_quaternion_component_init() -> None: def test_quaternion_sequence_init() -> None: """Test initialization from sequence (list, tuple) of 4 numbers.""" - # From list - q1 = Quaternion([0.1, 0.2, 0.3, 0.4]) + q1 = quaternion_from_array([0.1, 0.2, 0.3, 0.4]) assert q1.x == 0.1 assert q1.y == 0.2 assert q1.z == 0.3 assert q1.w == 0.4 - - # From tuple - q2 = Quaternion((0.5, 0.6, 0.7, 0.8)) + q2 = quaternion_from_array((0.5, 0.6, 0.7, 0.8)) assert q2.x == 0.5 assert q2.y == 0.6 assert q2.z == 0.7 assert q2.w == 0.8 - - # Test with integers in sequence - q3 = Quaternion([1, 2, 3, 4]) + q3 = quaternion_from_array([1, 2, 3, 4]) assert q3.x == 1.0 assert q3.y == 2.0 assert q3.z == 3.0 assert q3.w == 4.0 - - # Test error with wrong length with pytest.raises(ValueError, match="Quaternion requires exactly 4 components"): - Quaternion([1, 2, 3]) # Only 3 components - + quaternion_from_array([1, 2, 3]) with pytest.raises(ValueError, match="Quaternion requires exactly 4 components"): - Quaternion([1, 2, 3, 4, 5]) # Too many components + quaternion_from_array([1, 2, 3, 4, 5]) def test_quaternion_numpy_init() -> None: """Test initialization from numpy array.""" - # From numpy array arr = np.array([0.1, 0.2, 0.3, 0.4]) - q1 = Quaternion(arr) + q1 = quaternion_from_array(arr) assert q1.x == 0.1 assert q1.y == 0.2 assert q1.z == 0.3 assert q1.w == 0.4 - - # Test with different dtypes arr_int = np.array([1, 2, 3, 4], dtype=int) - q2 = Quaternion(arr_int) + q2 = quaternion_from_array(arr_int) assert q2.x == 1.0 assert q2.y == 2.0 assert q2.z == 3.0 assert q2.w == 4.0 - - # Test error with wrong size with pytest.raises(ValueError, match="Quaternion requires exactly 4 components"): - Quaternion(np.array([1, 2, 3])) # Only 3 elements - + quaternion_from_array(np.array([1, 2, 3])) with pytest.raises(ValueError, match="Quaternion requires exactly 4 components"): - Quaternion(np.array([1, 2, 3, 4, 5])) # Too many elements + quaternion_from_array(np.array([1, 2, 3, 4, 5])) def test_quaternion_copy_init() -> None: """Test initialization from another Quaternion (copy constructor).""" - original = Quaternion(0.1, 0.2, 0.3, 0.4) - copy = Quaternion(original) - + original = Quaternion(x=0.1, y=0.2, z=0.3, w=0.4) + copy = Quaternion.decode(original.encode()) assert copy.x == 0.1 assert copy.y == 0.2 assert copy.z == 0.3 assert copy.w == 0.4 - - # Verify it's a copy, not the same object assert copy is not original assert copy == original -def test_quaternion_lcm_init() -> None: - """Test initialization from LCM Quaternion.""" - lcm_quat = LCMQuaternion() - lcm_quat.x = 0.1 - lcm_quat.y = 0.2 - lcm_quat.z = 0.3 - lcm_quat.w = 0.4 - - q = Quaternion(lcm_quat) +def test_quaternion_mutated_value_copy() -> None: + """Test initialization from generated Quaternion.""" + source_quaternion = Quaternion() + source_quaternion.x = 0.1 + source_quaternion.y = 0.2 + source_quaternion.z = 0.3 + source_quaternion.w = 0.4 + q = Quaternion.decode(source_quaternion.encode()) assert q.x == 0.1 assert q.y == 0.2 assert q.z == 0.3 @@ -149,62 +135,47 @@ def test_quaternion_lcm_init() -> None: def test_quaternion_properties() -> None: """Test quaternion component properties.""" - q = Quaternion(1.0, 2.0, 3.0, 4.0) - - # Test property access + q = Quaternion(x=1.0, y=2.0, z=3.0, w=4.0) assert q.x == 1.0 assert q.y == 2.0 assert q.z == 3.0 assert q.w == 4.0 - - # Test as_tuple property - assert q.to_tuple() == (1.0, 2.0, 3.0, 4.0) + assert tuple(quaternion_array(q)) == (1.0, 2.0, 3.0, 4.0) def test_quaternion_indexing() -> None: """Test quaternion indexing support.""" - q = Quaternion(1.0, 2.0, 3.0, 4.0) - - # Test indexing - assert q[0] == 1.0 - assert q[1] == 2.0 - assert q[2] == 3.0 - assert q[3] == 4.0 + q = Quaternion(x=1.0, y=2.0, z=3.0, w=4.0) + assert quaternion_array(q)[0] == 1.0 + assert quaternion_array(q)[1] == 2.0 + assert quaternion_array(q)[2] == 3.0 + assert quaternion_array(q)[3] == 4.0 def test_quaternion_euler() -> None: """Test quaternion to Euler angles conversion.""" - - # Test identity quaternion (should give zero angles) q_identity = Quaternion() - angles = q_identity.to_euler() - assert np.isclose(angles.x, 0.0, atol=1e-10) # roll - assert np.isclose(angles.y, 0.0, atol=1e-10) # pitch - assert np.isclose(angles.z, 0.0, atol=1e-10) # yaw - - # Test 90 degree rotation around Z-axis (yaw) - q_z90 = Quaternion(0, 0, np.sin(np.pi / 4), np.cos(np.pi / 4)) - angles_z90 = q_z90.to_euler() - assert np.isclose(angles_z90.roll, 0.0, atol=1e-10) # roll should be 0 - assert np.isclose(angles_z90.pitch, 0.0, atol=1e-10) # pitch should be 0 - assert np.isclose(angles_z90.yaw, np.pi / 2, atol=1e-10) # yaw should be π/2 (90 degrees) - - # Test 90 degree rotation around X-axis (roll) - q_x90 = Quaternion(np.sin(np.pi / 4), 0, 0, np.cos(np.pi / 4)) - angles_x90 = q_x90.to_euler() - assert np.isclose(angles_x90.x, np.pi / 2, atol=1e-10) # roll should be π/2 - assert np.isclose(angles_x90.y, 0.0, atol=1e-10) # pitch should be 0 - assert np.isclose(angles_x90.z, 0.0, atol=1e-10) # yaw should be 0 - - -def test_lcm_encode_decode() -> None: - """Test encoding and decoding of Quaternion to/from binary LCM format.""" - q_source = Quaternion(1.0, 2.0, 3.0, 4.0) - - binary_msg = q_source.lcm_encode() - - q_dest = Quaternion.lcm_decode(binary_msg) - + angles = vector_from_array(quaternion_euler(q_identity)) + assert np.isclose(angles.x, 0.0, atol=1e-10) + assert np.isclose(angles.y, 0.0, atol=1e-10) + assert np.isclose(angles.z, 0.0, atol=1e-10) + q_z90 = Quaternion(x=0, y=0, z=np.sin(np.pi / 4), w=np.cos(np.pi / 4)) + angles_z90 = vector_from_array(quaternion_euler(q_z90)) + assert np.isclose(angles_z90.x, 0.0, atol=1e-10) + assert np.isclose(angles_z90.y, 0.0, atol=1e-10) + assert np.isclose(angles_z90.z, np.pi / 2, atol=1e-10) + q_x90 = Quaternion(x=np.sin(np.pi / 4), y=0, z=0, w=np.cos(np.pi / 4)) + angles_x90 = vector_from_array(quaternion_euler(q_x90)) + assert np.isclose(angles_x90.x, np.pi / 2, atol=1e-10) + assert np.isclose(angles_x90.y, 0.0, atol=1e-10) + assert np.isclose(angles_x90.z, 0.0, atol=1e-10) + + +def test_cdr_encode_decode() -> None: + """Test encoding and decoding of Quaternion to/from binary CDR format.""" + q_source = Quaternion(x=1.0, y=2.0, z=3.0, w=4.0) + binary_msg = q_source.encode() + q_dest = Quaternion.decode(binary_msg) assert isinstance(q_dest, Quaternion) assert q_dest is not q_source assert q_dest == q_source @@ -212,32 +183,20 @@ def test_lcm_encode_decode() -> None: def test_quaternion_multiplication() -> None: """Test quaternion multiplication (Hamilton product).""" - # Test identity multiplication - q1 = Quaternion(0.5, 0.5, 0.5, 0.5) - identity = Quaternion(0, 0, 0, 1) - - result = q1 * identity + q1 = Quaternion(x=0.5, y=0.5, z=0.5, w=0.5) + identity = Quaternion(x=0, y=0, z=0, w=1) + result = quaternion_product(q1, identity) assert np.allclose([result.x, result.y, result.z, result.w], [q1.x, q1.y, q1.z, q1.w]) - - # Test multiplication order matters (non-commutative) - q2 = Quaternion(0.1, 0.2, 0.3, 0.4) - q3 = Quaternion(0.4, 0.3, 0.2, 0.1) - - result1 = q2 * q3 - result2 = q3 * q2 - - # Results should be different + q2 = Quaternion(x=0.1, y=0.2, z=0.3, w=0.4) + q3 = Quaternion(x=0.4, y=0.3, z=0.2, w=0.1) + result1 = quaternion_product(q2, q3) + result2 = quaternion_product(q3, q2) assert not np.allclose( [result1.x, result1.y, result1.z, result1.w], [result2.x, result2.y, result2.z, result2.w] ) - - # Test specific multiplication case - # 90 degree rotations around Z axis angle = np.pi / 2 - q_90z = Quaternion(0, 0, np.sin(angle / 2), np.cos(angle / 2)) - - # Two 90 degree rotations should give 180 degrees - result = q_90z * q_90z + q_90z = Quaternion(x=0, y=0, z=np.sin(angle / 2), w=np.cos(angle / 2)) + result = quaternion_product(q_90z, q_90z) expected_angle = np.pi assert np.isclose(result.x, 0, atol=1e-10) assert np.isclose(result.y, 0, atol=1e-10) @@ -247,50 +206,36 @@ def test_quaternion_multiplication() -> None: def test_quaternion_conjugate() -> None: """Test quaternion conjugate.""" - q = Quaternion(0.1, 0.2, 0.3, 0.4) - conj = q.conjugate() - - # Conjugate should negate x, y, z but keep w + q = Quaternion(x=0.1, y=0.2, z=0.3, w=0.4) + conj = quaternion_conjugate(q) assert conj.x == -q.x assert conj.y == -q.y assert conj.z == -q.z assert conj.w == q.w - - # Test that q * q^* gives a real quaternion (x=y=z=0) - result = q * conj + result = quaternion_product(q, conj) assert np.isclose(result.x, 0, atol=1e-10) assert np.isclose(result.y, 0, atol=1e-10) assert np.isclose(result.z, 0, atol=1e-10) - # w should be the squared norm expected_w = q.x**2 + q.y**2 + q.z**2 + q.w**2 assert np.isclose(result.w, expected_w, atol=1e-10) def test_quaternion_inverse() -> None: """Test quaternion inverse.""" - # Test with unit quaternion - q_unit = Quaternion(0, 0, 0, 1).normalize() # Already normalized but being explicit - inv = q_unit.inverse() - - # For unit quaternion, inverse equals conjugate - conj = q_unit.conjugate() + q_unit = normalized_quaternion(Quaternion(x=0, y=0, z=0, w=1)) + inv = quaternion_inverse(q_unit) + conj = quaternion_conjugate(q_unit) assert np.allclose([inv.x, inv.y, inv.z, inv.w], [conj.x, conj.y, conj.z, conj.w]) - - # Test that q * q^-1 = identity - q = Quaternion(0.5, 0.5, 0.5, 0.5) - inv = q.inverse() - result = q * inv - + q = Quaternion(x=0.5, y=0.5, z=0.5, w=0.5) + inv = quaternion_inverse(q) + result = quaternion_product(q, inv) assert np.isclose(result.x, 0, atol=1e-10) assert np.isclose(result.y, 0, atol=1e-10) assert np.isclose(result.z, 0, atol=1e-10) assert np.isclose(result.w, 1, atol=1e-10) - - # Test inverse of non-unit quaternion - q_non_unit = Quaternion(2, 0, 0, 0) # Non-unit quaternion - inv = q_non_unit.inverse() - result = q_non_unit * inv - + q_non_unit = Quaternion(x=2, y=0, z=0, w=0) + inv = quaternion_inverse(q_non_unit) + result = quaternion_product(q_non_unit, inv) assert np.isclose(result.x, 0, atol=1e-10) assert np.isclose(result.y, 0, atol=1e-10) assert np.isclose(result.z, 0, atol=1e-10) @@ -299,15 +244,10 @@ def test_quaternion_inverse() -> None: def test_quaternion_normalize() -> None: """Test quaternion normalization.""" - # Test non-unit quaternion - q = Quaternion(1, 2, 3, 4) - q_norm = q.normalize() - - # Check that magnitude is 1 + q = Quaternion(x=1, y=2, z=3, w=4) + q_norm = normalized_quaternion(q) magnitude = np.sqrt(q_norm.x**2 + q_norm.y**2 + q_norm.z**2 + q_norm.w**2) assert np.isclose(magnitude, 1.0, atol=1e-10) - - # Check that direction is preserved scale = np.sqrt(q.x**2 + q.y**2 + q.z**2 + q.w**2) assert np.isclose(q_norm.x, q.x / scale, atol=1e-10) assert np.isclose(q_norm.y, q.y / scale, atol=1e-10) @@ -317,44 +257,26 @@ def test_quaternion_normalize() -> None: def test_quaternion_rotate_vector() -> None: """Test rotating vectors with quaternions.""" - from dimos.msgs.geometry_msgs.Vector3 import Vector3 - - # Test rotation of unit vectors - # 90 degree rotation around Z axis angle = np.pi / 2 - q_rot = Quaternion(0, 0, np.sin(angle / 2), np.cos(angle / 2)) - - # Rotate X unit vector - v_x = Vector3(1, 0, 0) - v_rotated = q_rot.rotate_vector(v_x) - - # Should now point along Y axis + q_rot = Quaternion(x=0, y=0, z=np.sin(angle / 2), w=np.cos(angle / 2)) + v_x = Vector3(x=1, y=0, z=0) + v_rotated = rotate_vector(q_rot, v_x) assert np.isclose(v_rotated.x, 0, atol=1e-10) assert np.isclose(v_rotated.y, 1, atol=1e-10) assert np.isclose(v_rotated.z, 0, atol=1e-10) - - # Rotate Y unit vector - v_y = Vector3(0, 1, 0) - v_rotated = q_rot.rotate_vector(v_y) - - # Should now point along negative X axis + v_y = Vector3(x=0, y=1, z=0) + v_rotated = rotate_vector(q_rot, v_y) assert np.isclose(v_rotated.x, -1, atol=1e-10) assert np.isclose(v_rotated.y, 0, atol=1e-10) assert np.isclose(v_rotated.z, 0, atol=1e-10) - - # Test that Z vector is unchanged (rotation axis) - v_z = Vector3(0, 0, 1) - v_rotated = q_rot.rotate_vector(v_z) - + v_z = Vector3(x=0, y=0, z=1) + v_rotated = rotate_vector(q_rot, v_z) assert np.isclose(v_rotated.x, 0, atol=1e-10) assert np.isclose(v_rotated.y, 0, atol=1e-10) assert np.isclose(v_rotated.z, 1, atol=1e-10) - - # Test identity rotation - q_identity = Quaternion(0, 0, 0, 1) - v = Vector3(1, 2, 3) - v_rotated = q_identity.rotate_vector(v) - + q_identity = Quaternion(x=0, y=0, z=0, w=1) + v = Vector3(x=1, y=2, z=3) + v_rotated = rotate_vector(q_identity, v) assert np.isclose(v_rotated.x, v.x, atol=1e-10) assert np.isclose(v_rotated.y, v.y, atol=1e-10) assert np.isclose(v_rotated.z, v.z, atol=1e-10) @@ -362,26 +284,22 @@ def test_quaternion_rotate_vector() -> None: def test_quaternion_inverse_zero() -> None: """Test that inverting zero quaternion raises error.""" - q_zero = Quaternion(0, 0, 0, 0) - + q_zero = Quaternion(x=0, y=0, z=0, w=0) with pytest.raises(ZeroDivisionError, match="Cannot invert zero quaternion"): - q_zero.inverse() + quaternion_inverse(q_zero) def test_quaternion_normalize_zero() -> None: """Test that normalizing zero quaternion raises error.""" - q_zero = Quaternion(0, 0, 0, 0) - + q_zero = Quaternion(x=0, y=0, z=0, w=0) with pytest.raises(ZeroDivisionError, match="Cannot normalize zero quaternion"): - q_zero.normalize() + normalized_quaternion(q_zero) def test_quaternion_multiplication_type_error() -> None: """Test that multiplying quaternion with non-quaternion raises error.""" - q = Quaternion(1, 0, 0, 0) - + q = Quaternion(x=1, y=0, z=0, w=0) with pytest.raises(TypeError, match="Cannot multiply Quaternion with"): - q * 5.0 - + quaternion_product(q, 5.0) with pytest.raises(TypeError, match="Cannot multiply Quaternion with"): - q * [1, 2, 3, 4] + quaternion_product(q, [1, 2, 3, 4]) diff --git a/dimos/msgs/geometry_msgs/test_Transform.py b/dimos/msgs/geometry_msgs/test_Transform.py index 056238719a1..eb397dbe02c 100644 --- a/dimos/msgs/geometry_msgs/test_Transform.py +++ b/dimos/msgs/geometry_msgs/test_Transform.py @@ -13,20 +13,35 @@ # limitations under the License. import math -import time +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest +from rosbags.typesys import Stores, get_typestore -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.geometry import ( + compose_transform_values, + compose_transforms, + point_from_array, + quaternion_from_euler, + relative_pose_transform, + transform_from_pose, + transform_pose_local, + vector_array, +) def test_transform_initialization() -> None: - # Test default initialization (identity transform) tf = Transform() assert tf.translation.x == 0.0 assert tf.translation.y == 0.0 @@ -35,390 +50,253 @@ def test_transform_initialization() -> None: assert tf.rotation.y == 0.0 assert tf.rotation.z == 0.0 assert tf.rotation.w == 1.0 - - # Test initialization with Vector3 and Quaternion - trans = Vector3(1.0, 2.0, 3.0) - rot = Quaternion(0.0, 0.0, 0.707107, 0.707107) # 90 degrees around Z + trans = Vector3(x=1.0, y=2.0, z=3.0) + rot = Quaternion(x=0.0, y=0.0, z=0.707107, w=0.707107) tf2 = Transform(translation=trans, rotation=rot) assert tf2.translation == trans assert tf2.rotation == rot - - # Test initialization with only translation - tf5 = Transform(translation=Vector3(7.0, 8.0, 9.0)) + tf5 = Transform(translation=Vector3(x=7.0, y=8.0, z=9.0)) assert tf5.translation.x == 7.0 assert tf5.translation.y == 8.0 assert tf5.translation.z == 9.0 - assert tf5.rotation.w == 1.0 # Identity rotation - - # Test initialization with only rotation - tf6 = Transform(rotation=Quaternion(0.0, 0.0, 0.0, 1.0)) - assert tf6.translation.is_zero() # Zero translation + assert tf5.rotation.w == 1.0 + tf6 = Transform(rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)) + assert np.allclose(vector_array(tf6.translation), 0) assert tf6.rotation.w == 1.0 - - # Test keyword argument initialization - tf7 = Transform(translation=Vector3(1, 2, 3), rotation=Quaternion()) - assert tf7.translation == Vector3(1, 2, 3) + tf7 = Transform(translation=Vector3(x=1, y=2, z=3), rotation=Quaternion()) + assert tf7.translation == Vector3(x=1, y=2, z=3) assert tf7.rotation == Quaternion() - - # Test keyword with only translation - tf8 = Transform(translation=Vector3(4, 5, 6)) - assert tf8.translation == Vector3(4, 5, 6) + tf8 = Transform(translation=Vector3(x=4, y=5, z=6)) + assert tf8.translation == Vector3(x=4, y=5, z=6) assert tf8.rotation.w == 1.0 - - # Test keyword with only rotation - tf9 = Transform(rotation=Quaternion(0, 0, 1, 0)) - assert tf9.translation.is_zero() - assert tf9.rotation == Quaternion(0, 0, 1, 0) + tf9 = Transform(rotation=Quaternion(x=0, y=0, z=1, w=0)) + assert np.allclose(vector_array(tf9.translation), 0) + assert tf9.rotation == Quaternion(x=0, y=0, z=1, w=0) def test_transform_identity() -> None: - # Test identity class method - tf = Transform.identity() - assert tf.translation.is_zero() + tf = Transform() + assert np.allclose(vector_array(tf.translation), 0) assert tf.rotation.x == 0.0 assert tf.rotation.y == 0.0 assert tf.rotation.z == 0.0 assert tf.rotation.w == 1.0 - - # Identity should equal default constructor assert tf == Transform() def test_transform_equality() -> None: - tf1 = Transform(translation=Vector3(1, 2, 3), rotation=Quaternion(0, 0, 0, 1)) - tf2 = Transform(translation=Vector3(1, 2, 3), rotation=Quaternion(0, 0, 0, 1)) - tf3 = Transform(translation=Vector3(1, 2, 4), rotation=Quaternion(0, 0, 0, 1)) # Different z - tf4 = Transform( - translation=Vector3(1, 2, 3), rotation=Quaternion(0, 0, 1, 0) - ) # Different rotation - + tf1 = Transform(translation=Vector3(x=1, y=2, z=3), rotation=Quaternion(x=0, y=0, z=0, w=1)) + tf2 = Transform(translation=Vector3(x=1, y=2, z=3), rotation=Quaternion(x=0, y=0, z=0, w=1)) + tf3 = Transform(translation=Vector3(x=1, y=2, z=4), rotation=Quaternion(x=0, y=0, z=0, w=1)) + tf4 = Transform(translation=Vector3(x=1, y=2, z=3), rotation=Quaternion(x=0, y=0, z=1, w=0)) assert tf1 == tf2 assert tf1 != tf3 assert tf1 != tf4 assert tf1 != "not a transform" -def test_transform_string_representations() -> None: - tf = Transform( - translation=Vector3(1.5, -2.0, 3.14), rotation=Quaternion(0, 0, 0.707107, 0.707107) +def test_transform_independent_cdr_fields() -> None: + source = Transform( + translation=Vector3(x=1.5, y=-2.0, z=3.14), rotation=Quaternion(z=0.707107, w=0.707107) ) - - # Test repr - repr_str = repr(tf) - assert "Transform" in repr_str - assert "translation=" in repr_str - assert "rotation=" in repr_str - assert "1.5" in repr_str - - # Test str - str_str = str(tf) - assert "Translation:" in str_str - assert "Rotation:" in str_str + decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(source.encode(), Transform.msg_name) + assert [decoded.translation.x, decoded.translation.y, decoded.translation.z] == [ + 1.5, + -2.0, + 3.14, + ] + assert [decoded.rotation.x, decoded.rotation.y, decoded.rotation.z, decoded.rotation.w] == [ + 0, + 0, + 0.707107, + 0.707107, + ] def test_pose_add_transform() -> None: - initial_pose = Pose(1.0, 0.0, 0.0) - - # 90 degree rotation around Z axis + initial_pose = Pose(position=Point(x=1.0, y=0.0, z=0.0)) angle = np.pi / 2 transform = Transform( - translation=Vector3(2.0, 1.0, 0.0), - rotation=Quaternion(0.0, 0.0, np.sin(angle / 2), np.cos(angle / 2)), + translation=Vector3(x=2.0, y=1.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=np.sin(angle / 2), w=np.cos(angle / 2)), ) - - transformed_pose = initial_pose @ transform - - # - Translation (2, 1, 0) is added directly to position (1, 0, 0) - # - Result position: (3, 1, 0) + transformed_pose = transform_pose_local(initial_pose, transform) assert np.isclose(transformed_pose.position.x, 3.0, atol=1e-10) assert np.isclose(transformed_pose.position.y, 1.0, atol=1e-10) assert np.isclose(transformed_pose.position.z, 0.0, atol=1e-10) - - # Rotation should be 90 degrees around Z assert np.isclose(transformed_pose.orientation.x, 0.0, atol=1e-10) assert np.isclose(transformed_pose.orientation.y, 0.0, atol=1e-10) assert np.isclose(transformed_pose.orientation.z, np.sin(angle / 2), atol=1e-10) assert np.isclose(transformed_pose.orientation.w, np.cos(angle / 2), atol=1e-10) - - initial_pose_stamped = PoseStamped( - position=initial_pose.position, orientation=initial_pose.orientation - ) - transformed_pose_stamped = PoseStamped( - position=transformed_pose.position, orientation=transformed_pose.orientation - ) - - found_tf = initial_pose_stamped.find_transform(transformed_pose_stamped) - - assert found_tf.translation == transform.translation - assert found_tf.rotation == transform.rotation - assert found_tf.translation.x == transform.translation.x - assert found_tf.translation.y == transform.translation.y - assert found_tf.translation.z == transform.translation.z - - assert found_tf.rotation.x == transform.rotation.x - assert found_tf.rotation.y == transform.rotation.y - assert found_tf.rotation.z == transform.rotation.z - assert found_tf.rotation.w == transform.rotation.w - - print(found_tf.rotation, found_tf.translation) + found_tf = relative_pose_transform(initial_pose, transformed_pose) + assert found_tf == transform def test_pose_add_transform_with_rotation() -> None: - # Create a pose at (0, 0, 0) rotated 90 degrees around Z angle = np.pi / 2 - initial_pose = Pose(0.0, 0.0, 0.0, 0.0, 0.0, np.sin(angle / 2), np.cos(angle / 2)) - - # Add 45 degree rotation to transform1 - rotation_angle = np.pi / 4 # 45 degrees + initial_pose = Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=np.sin(angle / 2), w=np.cos(angle / 2)), + ) + rotation_angle = np.pi / 4 transform1 = Transform( - translation=Vector3(1.0, 0.0, 0.0), + translation=Vector3(x=1.0, y=0.0, z=0.0), rotation=Quaternion( - 0.0, 0.0, np.sin(rotation_angle / 2), np.cos(rotation_angle / 2) - ), # 45� around Z + x=0.0, y=0.0, z=np.sin(rotation_angle / 2), w=np.cos(rotation_angle / 2) + ), ) - transform2 = Transform( - translation=Vector3(0.0, 1.0, 1.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), # No rotation + translation=Vector3(x=0.0, y=1.0, z=1.0), rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) + transformed_pose1 = transform_pose_local(initial_pose, transform1) + transformed_pose2 = transform_pose_local( + transform_pose_local(initial_pose, transform1), transform2 ) - - transformed_pose1 = initial_pose @ transform1 - transformed_pose2 = initial_pose @ transform1 @ transform2 - - # Test transformed_pose1: initial_pose + transform1 - # Since the pose is rotated 90� (facing +Y), moving forward (local X) - # means moving in the +Y direction in world frame assert np.isclose(transformed_pose1.position.x, 0.0, atol=1e-10) assert np.isclose(transformed_pose1.position.y, 1.0, atol=1e-10) assert np.isclose(transformed_pose1.position.z, 0.0, atol=1e-10) - - # Orientation should be 90� + 45� = 135� around Z - total_angle1 = angle + rotation_angle # 135 degrees + total_angle1 = angle + rotation_angle assert np.isclose(transformed_pose1.orientation.x, 0.0, atol=1e-10) assert np.isclose(transformed_pose1.orientation.y, 0.0, atol=1e-10) assert np.isclose(transformed_pose1.orientation.z, np.sin(total_angle1 / 2), atol=1e-10) assert np.isclose(transformed_pose1.orientation.w, np.cos(total_angle1 / 2), atol=1e-10) - - # Test transformed_pose2: initial_pose + transform1 + transform2 - # Starting from (0, 0, 0) facing 90�: - # - # - Apply transform1: move 1 forward (along +Y) � (0, 1, 0), now facing 135� - # - # - Apply transform2: move 1 in local Y and 1 up - # At 135�, local Y points at 225� (135� + 90�) - # - # x += cos(225�) = -2/2, y += sin(225�) = -2/2 sqrt2_2 = np.sqrt(2) / 2 - expected_x = 0.0 - sqrt2_2 # 0 - 2/2 H -0.707 - expected_y = 1.0 - sqrt2_2 # 1 - 2/2 H 0.293 - expected_z = 1.0 # 0 + 1 - + expected_x = 0.0 - sqrt2_2 + expected_y = 1.0 - sqrt2_2 + expected_z = 1.0 assert np.isclose(transformed_pose2.position.x, expected_x, atol=1e-10) assert np.isclose(transformed_pose2.position.y, expected_y, atol=1e-10) assert np.isclose(transformed_pose2.position.z, expected_z, atol=1e-10) - - # Orientation should be 135� (only transform1 has rotation) - total_angle2 = total_angle1 # 135 degrees (transform2 has no rotation) + total_angle2 = total_angle1 assert np.isclose(transformed_pose2.orientation.x, 0.0, atol=1e-10) assert np.isclose(transformed_pose2.orientation.y, 0.0, atol=1e-10) assert np.isclose(transformed_pose2.orientation.z, np.sin(total_angle2 / 2), atol=1e-10) assert np.isclose(transformed_pose2.orientation.w, np.cos(total_angle2 / 2), atol=1e-10) -def test_lcm_encode_decode() -> None: +def test_encode_decode() -> None: angle = np.pi / 2 transform = Transform( - translation=Vector3(2.0, 1.0, 0.0), - rotation=Quaternion(0.0, 0.0, np.sin(angle / 2), np.cos(angle / 2)), + translation=Vector3(x=2.0, y=1.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=np.sin(angle / 2), w=np.cos(angle / 2)), ) - - data = transform.lcm_encode() - - decoded_transform = Transform.lcm_decode(data) - + data = transform.encode() + decoded_transform = Transform.decode(data) assert decoded_transform == transform def test_transform_addition() -> None: - # Test 1: Simple translation addition (no rotation) - t1 = Transform( - translation=Vector3(1, 0, 0), - rotation=Quaternion(0, 0, 0, 1), # identity rotation - ) - t2 = Transform( - translation=Vector3(2, 0, 0), - rotation=Quaternion(0, 0, 0, 1), # identity rotation - ) - t3 = t1 + t2 - assert t3.translation == Vector3(3, 0, 0) - assert t3.rotation == Quaternion(0, 0, 0, 1) - - # Test 2: 90-degree rotation composition - # First transform: move 1 unit in X - t1 = Transform( - translation=Vector3(1, 0, 0), - rotation=Quaternion(0, 0, 0, 1), # identity - ) - # Second transform: move 1 unit in X with 90-degree rotation around Z + t1 = Transform(translation=Vector3(x=1, y=0, z=0), rotation=Quaternion(x=0, y=0, z=0, w=1)) + t2 = Transform(translation=Vector3(x=2, y=0, z=0), rotation=Quaternion(x=0, y=0, z=0, w=1)) + t3 = compose_transform_values(t1, t2) + assert t3.translation == Vector3(x=3, y=0, z=0) + assert t3.rotation == Quaternion(x=0, y=0, z=0, w=1) + t1 = Transform(translation=Vector3(x=1, y=0, z=0), rotation=Quaternion(x=0, y=0, z=0, w=1)) angle = np.pi / 2 t2 = Transform( - translation=Vector3(1, 0, 0), - rotation=Quaternion(0, 0, np.sin(angle / 2), np.cos(angle / 2)), + translation=Vector3(x=1, y=0, z=0), + rotation=Quaternion(x=0, y=0, z=np.sin(angle / 2), w=np.cos(angle / 2)), ) - t3 = t1 + t2 - assert t3.translation == Vector3(2, 0, 0) - # Rotation should be 90 degrees around Z + t3 = compose_transform_values(t1, t2) + assert t3.translation == Vector3(x=2, y=0, z=0) assert np.isclose(t3.rotation.x, 0.0, atol=1e-10) assert np.isclose(t3.rotation.y, 0.0, atol=1e-10) assert np.isclose(t3.rotation.z, np.sin(angle / 2), atol=1e-10) assert np.isclose(t3.rotation.w, np.cos(angle / 2), atol=1e-10) - - # Test 3: Rotation affects translation - # First transform: 90-degree rotation around Z t1 = Transform( - translation=Vector3(0, 0, 0), - rotation=Quaternion(0, 0, np.sin(angle / 2), np.cos(angle / 2)), # 90° around Z - ) - # Second transform: move 1 unit in X - t2 = Transform( - translation=Vector3(1, 0, 0), - rotation=Quaternion(0, 0, 0, 1), # identity + translation=Vector3(x=0, y=0, z=0), + rotation=Quaternion(x=0, y=0, z=np.sin(angle / 2), w=np.cos(angle / 2)), ) - t3 = t1 + t2 - # X direction rotated 90° becomes Y direction + t2 = Transform(translation=Vector3(x=1, y=0, z=0), rotation=Quaternion(x=0, y=0, z=0, w=1)) + t3 = compose_transform_values(t1, t2) assert np.isclose(t3.translation.x, 0.0, atol=1e-10) assert np.isclose(t3.translation.y, 1.0, atol=1e-10) assert np.isclose(t3.translation.z, 0.0, atol=1e-10) - # Rotation remains 90° around Z assert np.isclose(t3.rotation.z, np.sin(angle / 2), atol=1e-10) assert np.isclose(t3.rotation.w, np.cos(angle / 2), atol=1e-10) - # Test 4: Frame tracking - t1 = Transform( - translation=Vector3(1, 0, 0), - rotation=Quaternion(0, 0, 0, 1), - frame_id="world", + +def test_transform_frame_tracking() -> None: + first = TransformStamped( + header=Header(frame_id="world"), child_frame_id="robot", + transform=Transform(translation=Vector3(x=1)), ) - t2 = Transform( - translation=Vector3(2, 0, 0), - rotation=Quaternion(0, 0, 0, 1), - frame_id="robot", + second = TransformStamped( + header=Header(frame_id="robot"), child_frame_id="sensor", + transform=Transform(translation=Vector3(x=2)), ) - t3 = t1 + t2 - assert t3.frame_id == "world" - assert t3.child_frame_id == "sensor" - - # Test 5: Type error + result = compose_transforms(first, second) + assert result.header.frame_id == "world" + assert result.child_frame_id == "sensor" + assert result.transform.translation == Vector3(x=3) with pytest.raises(TypeError): - t1 + "not a transform" + compose_transform_values(first.transform, "not a transform") def test_transform_from_pose() -> None: - """Test converting Pose to Transform""" - # Create a Pose with position and orientation - pose = Pose( - position=Vector3(1.0, 2.0, 3.0), - orientation=Quaternion(0.0, 0.0, 0.707, 0.707), # 90 degrees around Z + pose = Pose(position=Point(x=1, y=2, z=3), orientation=Quaternion(z=0.707, w=0.707)) + result = transform_from_pose( + PoseStamped(header=Header(frame_id="world"), pose=pose), child_frame_id="base_link" ) - - # Convert to Transform - transform = Transform.from_pose("base_link", pose) - - # Check that translation and rotation match - assert transform.translation == pose.position - assert transform.rotation == pose.orientation - assert transform.frame_id == "world" # default frame_id - assert transform.child_frame_id == "base_link" # passed as first argument + np.testing.assert_array_equal( + vector_array(result.transform.translation), vector_array(pose.position) + ) + assert result.transform.rotation == pose.orientation + assert result.header.frame_id == "world" + assert result.child_frame_id == "base_link" -# validating results from example @ -# https://foxglove.dev/blog/understanding-ros-transforms def test_transform_from_ros() -> None: - """Test converting PoseStamped to Transform""" - test_time = time.time() - pose_stamped = PoseStamped( - ts=test_time, - frame_id="base_link", - position=Vector3(1, -1, 0), - orientation=Quaternion.from_euler(Vector3(0, 0, math.pi / 6)), + stamp = Time(sec=123, nanosec=456789123) + first = transform_from_pose( + PoseStamped( + header=Header(stamp=stamp, frame_id="base_link"), + pose=Pose( + position=Point(x=1, y=-1), orientation=quaternion_from_euler(0, 0, math.pi / 6) + ), + ), + child_frame_id="arm_base_link", ) - transform_base_link_to_arm = Transform.from_pose("arm_base_link", pose_stamped) - - transform_arm_to_end = Transform.from_pose( - "end", + second = transform_from_pose( PoseStamped( - ts=test_time, - frame_id="arm_base_link", - position=Vector3(1, 1, 0), - orientation=Quaternion.from_euler(Vector3(0, 0, math.pi / 6)), + header=Header(stamp=stamp, frame_id="arm_base_link"), + pose=Pose( + position=Point(x=1, y=1), orientation=quaternion_from_euler(0, 0, math.pi / 6) + ), ), + child_frame_id="end", ) - - print(transform_base_link_to_arm) - print(transform_arm_to_end) - - end_effector_global_pose = transform_base_link_to_arm + transform_arm_to_end - - assert end_effector_global_pose.translation.x == pytest.approx(1.366, abs=1e-3) - assert end_effector_global_pose.translation.y == pytest.approx(0.366, abs=1e-3) + result = compose_transforms(first, second) + assert result.transform.translation.x == pytest.approx(1.366, abs=1e-3) + assert result.transform.translation.y == pytest.approx(0.366, abs=1e-3) + assert result.header.stamp == stamp def test_transform_from_pose_stamped() -> None: - """Test converting PoseStamped to Transform""" - # Create a PoseStamped with position, orientation, timestamp and frame - test_time = time.time() - pose_stamped = PoseStamped( - ts=test_time, - frame_id="map", - position=Vector3(4.0, 5.0, 6.0), - orientation=Quaternion(0.0, 0.707, 0.0, 0.707), # 90 degrees around Y + pose = PoseStamped( + header=Header(stamp=Time(sec=123, nanosec=456789123), frame_id="map"), + pose=Pose(position=Point(x=4, y=5, z=6), orientation=Quaternion(y=0.707, w=0.707)), ) - - # Convert to Transform - transform = Transform.from_pose("robot_base", pose_stamped) - - # Check that all fields match - assert transform.translation == pose_stamped.position - assert transform.rotation == pose_stamped.orientation - assert transform.frame_id == pose_stamped.frame_id - assert transform.ts == pose_stamped.ts - assert transform.child_frame_id == "robot_base" # passed as first argument - - -def test_transform_from_pose_variants() -> None: - """Test from_pose with different Pose initialization methods""" - # Test with Pose created from x,y,z - pose1 = Pose(1.0, 2.0, 3.0) - transform1 = Transform.from_pose("base_link", pose1) - assert transform1.translation.x == 1.0 - assert transform1.translation.y == 2.0 - assert transform1.translation.z == 3.0 - assert transform1.rotation.w == 1.0 # Identity quaternion - - # Test with Pose created from tuple - pose2 = Pose(([7.0, 8.0, 9.0], [0.0, 0.0, 0.0, 1.0])) - transform2 = Transform.from_pose("base_link", pose2) - assert transform2.translation.x == 7.0 - assert transform2.translation.y == 8.0 - assert transform2.translation.z == 9.0 - - # Test with Pose created from dict - pose3 = Pose({"position": [10.0, 11.0, 12.0], "orientation": [0.0, 0.0, 0.0, 1.0]}) - transform3 = Transform.from_pose("base_link", pose3) - assert transform3.translation.x == 10.0 - assert transform3.translation.y == 11.0 - assert transform3.translation.z == 12.0 + result = transform_from_pose(pose, child_frame_id="robot_base") + np.testing.assert_array_equal( + vector_array(result.transform.translation), vector_array(pose.pose.position) + ) + assert result.transform.rotation == pose.pose.orientation + assert result.header == pose.header + assert result.child_frame_id == "robot_base" -def test_transform_from_pose_invalid_type() -> None: - """Test that from_pose raises TypeError for invalid types""" - with pytest.raises(TypeError): - Transform.from_pose("not a pose") +@pytest.mark.parametrize("values", [[1.0, 2.0, 3.0], (7.0, 8.0, 9.0), [10.0, 11.0, 12.0]]) +def test_transform_from_coordinate_arrays(values) -> None: + pose = PoseStamped(pose=Pose(position=point_from_array(values))) + result = transform_from_pose(pose, child_frame_id="base_link") + np.testing.assert_array_equal(vector_array(result.transform.translation), values) + assert result.transform.rotation == Quaternion() - with pytest.raises(TypeError): - Transform.from_pose(42) +@pytest.mark.parametrize("value", ["not a pose", 42, None]) +def test_transform_from_pose_invalid_type(value) -> None: with pytest.raises(TypeError): - Transform.from_pose(None) + transform_from_pose(value, child_frame_id="base_link") diff --git a/dimos/msgs/geometry_msgs/test_Twist.py b/dimos/msgs/geometry_msgs/test_Twist.py index a4dc93f3cc0..bff834b4c57 100644 --- a/dimos/msgs/geometry_msgs/test_Twist.py +++ b/dimos/msgs/geometry_msgs/test_Twist.py @@ -12,16 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos_lcm.geometry_msgs import Twist as LCMTwist +from copy import deepcopy + +from dimos_generated.geometry_msgs.msg import Quaternion, Twist, Vector3 import numpy as np +from rosbags.typesys import Stores, get_typestore -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.geometry import quaternion_euler, vector_array, vector_from_array def test_twist_initialization() -> None: - # Test default initialization (zero twist) tw = Twist() assert tw.linear.x == 0.0 assert tw.linear.y == 0.0 @@ -29,173 +29,129 @@ def test_twist_initialization() -> None: assert tw.angular.x == 0.0 assert tw.angular.y == 0.0 assert tw.angular.z == 0.0 - - # Test initialization with Vector3 linear and angular - lin = Vector3(1.0, 2.0, 3.0) - ang = Vector3(0.1, 0.2, 0.3) - tw2 = Twist(lin, ang) + lin = Vector3(x=1.0, y=2.0, z=3.0) + ang = Vector3(x=0.1, y=0.2, z=0.3) + tw2 = Twist(linear=lin, angular=ang) assert tw2.linear == lin assert tw2.angular == ang - - # Test copy constructor - tw3 = Twist(tw2) + tw3 = deepcopy(tw2) assert tw3.linear == tw2.linear assert tw3.angular == tw2.angular assert tw3 == tw2 - # Ensure it's a deep copy tw3.linear.x = 10.0 assert tw2.linear.x == 1.0 - - # Test initialization from LCM Twist - lcm_tw = LCMTwist() - lcm_tw.linear = Vector3(4.0, 5.0, 6.0) - lcm_tw.angular = Vector3(0.4, 0.5, 0.6) - tw4 = Twist(lcm_tw) + source = Twist() + source.linear = Vector3(x=4.0, y=5.0, z=6.0) + source.angular = Vector3(x=0.4, y=0.5, z=0.6) + tw4 = Twist.decode(source.encode()) assert tw4.linear.x == 4.0 assert tw4.linear.y == 5.0 assert tw4.linear.z == 6.0 assert tw4.angular.x == 0.4 assert tw4.angular.y == 0.5 assert tw4.angular.z == 0.6 - - # Test initialization with linear and angular as quaternion - quat = Quaternion(0, 0, 0.707107, 0.707107) # 90 degrees around Z - tw5 = Twist(Vector3(1.0, 2.0, 3.0), quat) - assert tw5.linear == Vector3(1.0, 2.0, 3.0) - # Quaternion should be converted to euler angles - euler = quat.to_euler() + quat = Quaternion(x=0, y=0, z=0.707107, w=0.707107) + tw5 = Twist( + linear=Vector3(x=1.0, y=2.0, z=3.0), angular=vector_from_array(quaternion_euler(quat)) + ) + assert tw5.linear == Vector3(x=1.0, y=2.0, z=3.0) + euler = vector_from_array(quaternion_euler(quat)) assert np.allclose(tw5.angular.x, euler.x) assert np.allclose(tw5.angular.y, euler.y) assert np.allclose(tw5.angular.z, euler.z) - - # Test keyword argument initialization - tw7 = Twist(linear=Vector3(1, 2, 3), angular=Vector3(0.1, 0.2, 0.3)) - assert tw7.linear == Vector3(1, 2, 3) - assert tw7.angular == Vector3(0.1, 0.2, 0.3) - - # Test keyword with only linear - tw8 = Twist(linear=Vector3(4, 5, 6)) - assert tw8.linear == Vector3(4, 5, 6) - assert tw8.angular.is_zero() - - # Test keyword with only angular - tw9 = Twist(angular=Vector3(0.4, 0.5, 0.6)) - assert tw9.linear.is_zero() - assert tw9.angular == Vector3(0.4, 0.5, 0.6) - - # Test keyword with angular as quaternion - tw10 = Twist(angular=Quaternion(0, 0, 0.707107, 0.707107)) - assert tw10.linear.is_zero() - euler = Quaternion(0, 0, 0.707107, 0.707107).to_euler() + tw7 = Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.1, y=0.2, z=0.3)) + assert tw7.linear == Vector3(x=1, y=2, z=3) + assert tw7.angular == Vector3(x=0.1, y=0.2, z=0.3) + tw8 = Twist(linear=Vector3(x=4, y=5, z=6)) + assert tw8.linear == Vector3(x=4, y=5, z=6) + assert np.allclose(vector_array(tw8.angular), 0) + tw9 = Twist(angular=Vector3(x=0.4, y=0.5, z=0.6)) + assert np.allclose(vector_array(tw9.linear), 0) + assert tw9.angular == Vector3(x=0.4, y=0.5, z=0.6) + tw10 = Twist( + angular=vector_from_array(quaternion_euler(Quaternion(x=0, y=0, z=0.707107, w=0.707107))) + ) + assert np.allclose(vector_array(tw10.linear), 0) + euler = vector_from_array(quaternion_euler(Quaternion(x=0, y=0, z=0.707107, w=0.707107))) assert np.allclose(tw10.angular.x, euler.x) assert np.allclose(tw10.angular.y, euler.y) assert np.allclose(tw10.angular.z, euler.z) - - # Test keyword with linear and angular as quaternion - tw11 = Twist(linear=Vector3(1, 0, 0), angular=Quaternion(0, 0, 0, 1)) - assert tw11.linear == Vector3(1, 0, 0) - assert tw11.angular.is_zero() # Identity quaternion -> zero euler angles + tw11 = Twist( + linear=Vector3(x=1, y=0, z=0), + angular=vector_from_array(quaternion_euler(Quaternion(x=0, y=0, z=0, w=1))), + ) + assert tw11.linear == Vector3(x=1, y=0, z=0) + assert np.allclose(vector_array(tw11.angular), 0) def test_twist_zero() -> None: - # Test zero class method - tw = Twist.zero() - assert tw.linear.is_zero() - assert tw.angular.is_zero() - assert tw.is_zero() - - # Zero should equal default constructor + tw = Twist() + assert np.allclose(vector_array(tw.linear), 0) + assert np.allclose(vector_array(tw.angular), 0) + assert np.allclose(np.concatenate((vector_array(tw.linear), vector_array(tw.angular))), 0) assert tw == Twist() def test_twist_equality() -> None: - tw1 = Twist(Vector3(1, 2, 3), Vector3(0.1, 0.2, 0.3)) - tw2 = Twist(Vector3(1, 2, 3), Vector3(0.1, 0.2, 0.3)) - tw3 = Twist(Vector3(1, 2, 4), Vector3(0.1, 0.2, 0.3)) # Different linear z - tw4 = Twist(Vector3(1, 2, 3), Vector3(0.1, 0.2, 0.4)) # Different angular z - + tw1 = Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.1, y=0.2, z=0.3)) + tw2 = Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.1, y=0.2, z=0.3)) + tw3 = Twist(linear=Vector3(x=1, y=2, z=4), angular=Vector3(x=0.1, y=0.2, z=0.3)) + tw4 = Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.1, y=0.2, z=0.4)) assert tw1 == tw2 assert tw1 != tw3 assert tw1 != tw4 assert tw1 != "not a twist" -def test_twist_string_representations() -> None: - tw = Twist(Vector3(1.5, -2.0, 3.14), Vector3(0.1, -0.2, 0.3)) - - # Test repr - repr_str = repr(tw) - assert "Twist" in repr_str - assert "linear=" in repr_str - assert "angular=" in repr_str - assert "1.5" in repr_str - assert "0.1" in repr_str - - # Test str - str_str = str(tw) - assert "Twist:" in str_str - assert "Linear:" in str_str - assert "Angular:" in str_str +def test_twist_independent_cdr_fields() -> None: + tw = Twist(linear=Vector3(x=1.5, y=-2.0, z=3.14), angular=Vector3(x=0.1, y=-0.2, z=0.3)) + decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr( + tw.encode(), "geometry_msgs/msg/Twist" + ) + assert [decoded.linear.x, decoded.linear.y, decoded.linear.z] == [1.5, -2.0, 3.14] + assert [decoded.angular.x, decoded.angular.y, decoded.angular.z] == [0.1, -0.2, 0.3] def test_twist_is_zero() -> None: - # Test zero twist tw1 = Twist() - assert tw1.is_zero() + assert np.allclose(np.concatenate((vector_array(tw1.linear), vector_array(tw1.angular))), 0) + tw2 = Twist(linear=Vector3(x=0.1, y=0, z=0)) + assert not np.allclose(np.concatenate((vector_array(tw2.linear), vector_array(tw2.angular))), 0) + tw3 = Twist(angular=Vector3(x=0, y=0, z=0.1)) + assert not np.allclose(np.concatenate((vector_array(tw3.linear), vector_array(tw3.angular))), 0) + tw4 = Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.1, y=0.2, z=0.3)) + assert not np.allclose(np.concatenate((vector_array(tw4.linear), vector_array(tw4.angular))), 0) - # Test non-zero linear - tw2 = Twist(linear=Vector3(0.1, 0, 0)) - assert not tw2.is_zero() - # Test non-zero angular - tw3 = Twist(angular=Vector3(0, 0, 0.1)) - assert not tw3.is_zero() - - # Test both non-zero - tw4 = Twist(Vector3(1, 2, 3), Vector3(0.1, 0.2, 0.3)) - assert not tw4.is_zero() - - -def test_twist_bool() -> None: - # Test zero twist is False +def test_twist_explicit_nonzero_predicate() -> None: tw1 = Twist() - assert not tw1 - - # Test non-zero twist is True - tw2 = Twist(linear=Vector3(1, 0, 0)) - assert tw2 - - tw3 = Twist(angular=Vector3(0, 0, 0.1)) - assert tw3 - - tw4 = Twist(Vector3(1, 2, 3), Vector3(0.1, 0.2, 0.3)) - assert tw4 - - -def test_twist_lcm_encoding() -> None: - # Test encoding and decoding - tw = Twist(Vector3(1.5, 2.5, 3.5), Vector3(0.1, 0.2, 0.3)) - - # Encode - encoded = tw.lcm_encode() + assert not (np.any(vector_array(tw1.linear)) or np.any(vector_array(tw1.angular))) + tw2 = Twist(linear=Vector3(x=1, y=0, z=0)) + assert np.any(vector_array(tw2.linear)) or np.any(vector_array(tw2.angular)) + tw3 = Twist(angular=Vector3(x=0, y=0, z=0.1)) + assert np.any(vector_array(tw3.linear)) or np.any(vector_array(tw3.angular)) + tw4 = Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.1, y=0.2, z=0.3)) + assert np.any(vector_array(tw4.linear)) or np.any(vector_array(tw4.angular)) + + +def test_twist_cdr_encoding() -> None: + tw = Twist(linear=Vector3(x=1.5, y=2.5, z=3.5), angular=Vector3(x=0.1, y=0.2, z=0.3)) + encoded = tw.encode() assert isinstance(encoded, bytes) - - # Decode - decoded = Twist.lcm_decode(encoded) + decoded = Twist.decode(encoded) assert decoded.linear == tw.linear assert decoded.angular == tw.angular - assert isinstance(decoded.linear, Vector3) assert decoded == tw def test_twist_with_lists() -> None: - # Test initialization with lists instead of Vector3 - tw1 = Twist(linear=[1, 2, 3], angular=[0.1, 0.2, 0.3]) - assert tw1.linear == Vector3(1, 2, 3) - assert tw1.angular == Vector3(0.1, 0.2, 0.3) - - # Test with numpy arrays - tw2 = Twist(linear=np.array([4, 5, 6]), angular=np.array([0.4, 0.5, 0.6])) - assert tw2.linear == Vector3(4, 5, 6) - assert tw2.angular == Vector3(0.4, 0.5, 0.6) + tw1 = Twist(linear=vector_from_array([1, 2, 3]), angular=vector_from_array([0.1, 0.2, 0.3])) + assert tw1.linear == Vector3(x=1, y=2, z=3) + assert tw1.angular == Vector3(x=0.1, y=0.2, z=0.3) + tw2 = Twist( + linear=vector_from_array(np.array([4, 5, 6])), + angular=vector_from_array(np.array([0.4, 0.5, 0.6])), + ) + assert tw2.linear == Vector3(x=4, y=5, z=6) + assert tw2.angular == Vector3(x=0.4, y=0.5, z=0.6) diff --git a/dimos/msgs/geometry_msgs/test_TwistStamped.py b/dimos/msgs/geometry_msgs/test_TwistStamped.py index d37b4b27172..590f31291d8 100644 --- a/dimos/msgs/geometry_msgs/test_TwistStamped.py +++ b/dimos/msgs/geometry_msgs/test_TwistStamped.py @@ -13,42 +13,47 @@ # limitations under the License. import pickle -import time -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped +from dimos_generated.geometry_msgs.msg import Twist, TwistStamped, Vector3 +from dimos_generated.std_msgs.msg import Header +from rosbags.typesys import Stores, get_typestore +from dimos.msgs.time import time_from_nanoseconds -def test_lcm_encode_decode() -> None: - """Test encoding and decoding of TwistStamped to/from binary LCM format.""" - twist_source = TwistStamped( - ts=time.time(), - linear=(1.0, 2.0, 3.0), - angular=(0.1, 0.2, 0.3), - ) - binary_msg = twist_source.lcm_encode() - twist_dest = TwistStamped.lcm_decode(binary_msg) - - assert isinstance(twist_dest, TwistStamped) - assert twist_dest is not twist_source - print(twist_source.linear) - print(twist_source.angular) - - print(twist_dest.linear) - print(twist_dest.angular) - assert twist_dest == twist_source +def test_cdr_encode_decode() -> None: + source = TwistStamped( + header=Header(stamp=time_from_nanoseconds(1234567890123456789)), + twist=Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.1, y=0.2, z=0.3)), + ) + binary = source.encode() + decoded = TwistStamped.decode(binary) + assert isinstance(decoded, TwistStamped) + assert decoded is not source + assert decoded.encode() == binary + independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(binary, source.msg_name) + assert (independent.header.stamp.sec, independent.header.stamp.nanosec) == ( + 1234567890, + 123456789, + ) + assert (independent.twist.linear.x, independent.twist.linear.y, independent.twist.linear.z) == ( + 1, + 2, + 3, + ) + assert ( + independent.twist.angular.x, + independent.twist.angular.y, + independent.twist.angular.z, + ) == (0.1, 0.2, 0.3) def test_pickle_encode_decode() -> None: - """Test encoding and decoding of TwistStamped to/from binary pickle format.""" - - twist_source = TwistStamped( - ts=time.time(), - linear=(1.0, 2.0, 3.0), - angular=(0.1, 0.2, 0.3), + source = TwistStamped( + header=Header(stamp=time_from_nanoseconds(1234567890123456789)), + twist=Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.1, y=0.2, z=0.3)), ) - binary_msg = pickle.dumps(twist_source) - twist_dest = pickle.loads(binary_msg) - assert isinstance(twist_dest, TwistStamped) - assert twist_dest is not twist_source - assert twist_dest == twist_source + decoded = pickle.loads(pickle.dumps(source)) + assert isinstance(decoded, TwistStamped) + assert decoded is not source + assert decoded.encode() == source.encode() diff --git a/dimos/msgs/geometry_msgs/test_TwistWithCovariance.py b/dimos/msgs/geometry_msgs/test_TwistWithCovariance.py index 1d8ae820ad0..2814a6667c3 100644 --- a/dimos/msgs/geometry_msgs/test_TwistWithCovariance.py +++ b/dimos/msgs/geometry_msgs/test_TwistWithCovariance.py @@ -12,345 +12,120 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos_lcm.geometry_msgs import TwistWithCovariance as LCMTwistWithCovariance -import numpy as np -import pytest - -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.TwistWithCovariance import TwistWithCovariance -from dimos.msgs.geometry_msgs.Vector3 import Vector3 - - -def test_twist_with_covariance_default_init() -> None: - """Test that default initialization creates a zero twist with zero covariance.""" - twist_cov = TwistWithCovariance() - - # Twist should be zero - assert twist_cov.twist.linear.x == 0.0 - assert twist_cov.twist.linear.y == 0.0 - assert twist_cov.twist.linear.z == 0.0 - assert twist_cov.twist.angular.x == 0.0 - assert twist_cov.twist.angular.y == 0.0 - assert twist_cov.twist.angular.z == 0.0 - - # Covariance should be all zeros - assert np.all(twist_cov.covariance == 0.0) - assert twist_cov.covariance.shape == (36,) - - -def test_twist_with_covariance_twist_init() -> None: - """Test initialization with a Twist object.""" - linear = Vector3(1.0, 2.0, 3.0) - angular = Vector3(0.1, 0.2, 0.3) - twist = Twist(linear, angular) - twist_cov = TwistWithCovariance(twist) - - # Twist should match - assert twist_cov.twist.linear.x == 1.0 - assert twist_cov.twist.linear.y == 2.0 - assert twist_cov.twist.linear.z == 3.0 - assert twist_cov.twist.angular.x == 0.1 - assert twist_cov.twist.angular.y == 0.2 - assert twist_cov.twist.angular.z == 0.3 - - # Covariance should be zeros by default - assert np.all(twist_cov.covariance == 0.0) - - -def test_twist_with_covariance_twist_and_covariance_init() -> None: - """Test initialization with twist and covariance.""" - twist = Twist(Vector3(1.0, 2.0, 3.0), Vector3(0.1, 0.2, 0.3)) - covariance = np.arange(36, dtype=float) - twist_cov = TwistWithCovariance(twist, covariance) - - # Twist should match - assert twist_cov.twist.linear.x == 1.0 - assert twist_cov.twist.linear.y == 2.0 - assert twist_cov.twist.linear.z == 3.0 - - # Covariance should match - assert np.array_equal(twist_cov.covariance, covariance) - - -def test_twist_with_covariance_tuple_init() -> None: - """Test initialization with tuple of (linear, angular) velocities.""" - linear = [1.0, 2.0, 3.0] - angular = [0.1, 0.2, 0.3] - covariance = np.arange(36, dtype=float) - twist_cov = TwistWithCovariance((linear, angular), covariance) - - # Twist should match - assert twist_cov.twist.linear.x == 1.0 - assert twist_cov.twist.linear.y == 2.0 - assert twist_cov.twist.linear.z == 3.0 - assert twist_cov.twist.angular.x == 0.1 - assert twist_cov.twist.angular.y == 0.2 - assert twist_cov.twist.angular.z == 0.3 - - # Covariance should match - assert np.array_equal(twist_cov.covariance, covariance) - - -def test_twist_with_covariance_list_covariance() -> None: - """Test initialization with covariance as a list.""" - twist = Twist(Vector3(1.0, 2.0, 3.0), Vector3(0.1, 0.2, 0.3)) - covariance_list = list(range(36)) - twist_cov = TwistWithCovariance(twist, covariance_list) - - # Covariance should be converted to numpy array - assert isinstance(twist_cov.covariance, np.ndarray) - assert np.array_equal(twist_cov.covariance, np.array(covariance_list)) - - -def test_twist_with_covariance_copy_init() -> None: - """Test copy constructor.""" - twist = Twist(Vector3(1.0, 2.0, 3.0), Vector3(0.1, 0.2, 0.3)) - covariance = np.arange(36, dtype=float) - original = TwistWithCovariance(twist, covariance) - copy = TwistWithCovariance(original) - - # Should be equal but not the same object - assert copy == original - assert copy is not original - assert copy.twist is not original.twist - assert copy.covariance is not original.covariance - - # Modify original to ensure they're independent - original.covariance[0] = 999.0 - assert copy.covariance[0] != 999.0 - - -def test_twist_with_covariance_lcm_init() -> None: - """Test initialization from LCM message.""" - lcm_msg = LCMTwistWithCovariance() - lcm_msg.twist.linear.x = 1.0 - lcm_msg.twist.linear.y = 2.0 - lcm_msg.twist.linear.z = 3.0 - lcm_msg.twist.angular.x = 0.1 - lcm_msg.twist.angular.y = 0.2 - lcm_msg.twist.angular.z = 0.3 - lcm_msg.covariance = list(range(36)) - - twist_cov = TwistWithCovariance(lcm_msg) - - # Twist should match - assert twist_cov.twist.linear.x == 1.0 - assert twist_cov.twist.linear.y == 2.0 - assert twist_cov.twist.linear.z == 3.0 - assert twist_cov.twist.angular.x == 0.1 - assert twist_cov.twist.angular.y == 0.2 - assert twist_cov.twist.angular.z == 0.3 - - # Covariance should match - assert np.array_equal(twist_cov.covariance, np.arange(36)) - - -def test_twist_with_covariance_dict_init() -> None: - """Test initialization from dictionary.""" - twist_dict = { - "twist": Twist(Vector3(1.0, 2.0, 3.0), Vector3(0.1, 0.2, 0.3)), - "covariance": list(range(36)), - } - twist_cov = TwistWithCovariance(twist_dict) - - assert twist_cov.twist.linear.x == 1.0 - assert twist_cov.twist.linear.y == 2.0 - assert twist_cov.twist.linear.z == 3.0 - assert np.array_equal(twist_cov.covariance, np.arange(36)) - - -def test_twist_with_covariance_dict_init_no_covariance() -> None: - """Test initialization from dictionary without covariance.""" - twist_dict = {"twist": Twist(Vector3(1.0, 2.0, 3.0), Vector3(0.1, 0.2, 0.3))} - twist_cov = TwistWithCovariance(twist_dict) - - assert twist_cov.twist.linear.x == 1.0 - assert np.all(twist_cov.covariance == 0.0) - - -def test_twist_with_covariance_tuple_of_tuple_init() -> None: - """Test initialization from tuple of (twist_tuple, covariance).""" - twist_tuple = ([1.0, 2.0, 3.0], [0.1, 0.2, 0.3]) - covariance = np.arange(36, dtype=float) - twist_cov = TwistWithCovariance((twist_tuple, covariance)) - - assert twist_cov.twist.linear.x == 1.0 - assert twist_cov.twist.linear.y == 2.0 - assert twist_cov.twist.linear.z == 3.0 - assert twist_cov.twist.angular.x == 0.1 - assert twist_cov.twist.angular.y == 0.2 - assert twist_cov.twist.angular.z == 0.3 - assert np.array_equal(twist_cov.covariance, covariance) +"""Generated covariance values, with explicit NumPy matrix operations. +Legacy polymorphic positional constructors and presentation methods are retired. +""" -def test_twist_with_covariance_properties() -> None: - """Test convenience properties.""" - twist = Twist(Vector3(1.0, 2.0, 3.0), Vector3(0.1, 0.2, 0.3)) - twist_cov = TwistWithCovariance(twist) - - # Linear and angular properties - assert twist_cov.linear.x == 1.0 - assert twist_cov.linear.y == 2.0 - assert twist_cov.linear.z == 3.0 - assert twist_cov.angular.x == 0.1 - assert twist_cov.angular.y == 0.2 - assert twist_cov.angular.z == 0.3 - - -def test_twist_with_covariance_matrix_property() -> None: - """Test covariance matrix property.""" - twist = Twist() - covariance_array = np.arange(36, dtype=float) - twist_cov = TwistWithCovariance(twist, covariance_array) - - # Get as matrix - cov_matrix = twist_cov.covariance_matrix - assert cov_matrix.shape == (6, 6) - assert cov_matrix[0, 0] == 0.0 - assert cov_matrix[5, 5] == 35.0 - - # Set from matrix - new_matrix = np.eye(6) * 2.0 - twist_cov.covariance_matrix = new_matrix - assert np.array_equal(twist_cov.covariance[:6], [2.0, 0.0, 0.0, 0.0, 0.0, 0.0]) - - -def test_twist_with_covariance_repr() -> None: - """Test string representation.""" - twist = Twist(Vector3(1.234, 2.567, 3.891), Vector3(0.1, 0.2, 0.3)) - twist_cov = TwistWithCovariance(twist) - - repr_str = repr(twist_cov) - assert "TwistWithCovariance" in repr_str - assert "twist=" in repr_str - assert "covariance=" in repr_str - assert "36 elements" in repr_str - - -def test_twist_with_covariance_str() -> None: - """Test string formatting.""" - twist = Twist(Vector3(1.234, 2.567, 3.891), Vector3(0.1, 0.2, 0.3)) - covariance = np.eye(6).flatten() - twist_cov = TwistWithCovariance(twist, covariance) - - str_repr = str(twist_cov) - assert "TwistWithCovariance" in str_repr - assert "1.234" in str_repr - assert "2.567" in str_repr - assert "3.891" in str_repr - assert "cov_trace" in str_repr - assert "6.000" in str_repr # Trace of identity matrix is 6 - - -def test_twist_with_covariance_equality() -> None: - """Test equality comparison.""" - twist1 = Twist(Vector3(1.0, 2.0, 3.0), Vector3(0.1, 0.2, 0.3)) - cov1 = np.arange(36, dtype=float) - twist_cov1 = TwistWithCovariance(twist1, cov1) - - twist2 = Twist(Vector3(1.0, 2.0, 3.0), Vector3(0.1, 0.2, 0.3)) - cov2 = np.arange(36, dtype=float) - twist_cov2 = TwistWithCovariance(twist2, cov2) - - # Equal - assert twist_cov1 == twist_cov2 - - # Different twist - twist3 = Twist(Vector3(1.1, 2.0, 3.0), Vector3(0.1, 0.2, 0.3)) - twist_cov3 = TwistWithCovariance(twist3, cov1) - assert twist_cov1 != twist_cov3 - - # Different covariance - cov3 = np.arange(36, dtype=float) + 1 - twist_cov4 = TwistWithCovariance(twist1, cov3) - assert twist_cov1 != twist_cov4 - - # Different type - assert twist_cov1 != "not a twist" - assert twist_cov1 is not None - - -def test_twist_with_covariance_is_zero() -> None: - """Test is_zero method.""" - # Zero twist - twist_cov1 = TwistWithCovariance() - assert twist_cov1.is_zero() - assert not twist_cov1 # Boolean conversion - - # Non-zero twist - twist = Twist(Vector3(1.0, 0.0, 0.0), Vector3(0.0, 0.0, 0.0)) - twist_cov2 = TwistWithCovariance(twist) - assert not twist_cov2.is_zero() - assert twist_cov2 # Boolean conversion - - -def test_twist_with_covariance_lcm_encode_decode() -> None: - """Test LCM encoding and decoding.""" - twist = Twist(Vector3(1.0, 2.0, 3.0), Vector3(0.1, 0.2, 0.3)) - covariance = np.arange(36, dtype=float) - source = TwistWithCovariance(twist, covariance) - - # Encode and decode - binary_msg = source.lcm_encode() - decoded = TwistWithCovariance.lcm_decode(binary_msg) - - # Should be equal - assert decoded == source - assert isinstance(decoded, TwistWithCovariance) - assert isinstance(decoded.twist, Twist) - assert isinstance(decoded.covariance, np.ndarray) - - -def test_twist_with_covariance_zero_covariance() -> None: - """Test with zero covariance matrix.""" - twist = Twist(Vector3(1.0, 2.0, 3.0), Vector3(0.1, 0.2, 0.3)) - twist_cov = TwistWithCovariance(twist) - - assert np.all(twist_cov.covariance == 0.0) - assert np.trace(twist_cov.covariance_matrix) == 0.0 - - -def test_twist_with_covariance_diagonal_covariance() -> None: - """Test with diagonal covariance matrix.""" - twist = Twist() - covariance = np.zeros(36) - # Set diagonal elements - for i in range(6): - covariance[i * 6 + i] = i + 1 - - twist_cov = TwistWithCovariance(twist, covariance) - - cov_matrix = twist_cov.covariance_matrix - assert np.trace(cov_matrix) == sum(range(1, 7)) # 1+2+3+4+5+6 = 21 - - # Check diagonal elements - for i in range(6): - assert cov_matrix[i, i] == i + 1 - - # Check off-diagonal elements are zero - for i in range(6): - for j in range(6): - if i != j: - assert cov_matrix[i, j] == 0.0 +from dimos_generated.geometry_msgs.msg import Twist, TwistWithCovariance, Vector3 +import numpy as np +import pytest +from rosbags.typesys import Stores, get_typestore + + +def test_default_fields_and_covariance() -> None: + source = TwistWithCovariance() + value = source.twist + np.testing.assert_array_equal( + [ + value.linear.x, + value.linear.y, + value.linear.z, + value.angular.x, + value.angular.y, + value.angular.z, + ], + [0, 0, 0, 0, 0, 0], + ) + assert np.asarray(source.covariance).shape == (36,) + np.testing.assert_array_equal(source.covariance, np.zeros(36)) + + +@pytest.mark.parametrize("as_list", [False, True]) +@pytest.mark.parametrize( + "covariance", + [np.zeros(36), np.arange(36, dtype=float), np.eye(6).ravel(), np.diag(np.arange(1, 7)).ravel()], +) +def test_explicit_construction_and_independent_cdr(as_list: bool, covariance: np.ndarray) -> None: + source = TwistWithCovariance( + twist=Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.1, y=0.2, z=0.3)), + covariance=covariance.tolist() if as_list else covariance, + ) + decoded = TwistWithCovariance.decode(source.encode()) + independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr( + source.encode(), TwistWithCovariance.msg_name + ) + for result in (source, decoded, independent): + value = result.twist + np.testing.assert_array_equal( + [ + value.linear.x, + value.linear.y, + value.linear.z, + value.angular.x, + value.angular.y, + value.angular.z, + ], + [1, 2, 3, 0.1, 0.2, 0.3], + ) + np.testing.assert_array_equal(result.covariance, covariance) + matrix = np.asarray(result.covariance).reshape(6, 6) + for row in range(6): + for col in range(6): + assert matrix[row, col] == covariance[row * 6 + col] + assert np.trace(matrix) == np.trace(covariance.reshape(6, 6)) + + +def test_copy_equality_and_independent_storage() -> None: + original = TwistWithCovariance( + twist=Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.1, y=0.2, z=0.3)), + covariance=np.arange(36, dtype=float), + ) + copied = TwistWithCovariance.decode(original.encode()) + assert copied == original + assert copied is not original + assert copied.twist is not original.twist + assert copied.covariance is not original.covariance + original.covariance[0] = 999 + assert copied.covariance[0] == 0 + assert copied != original + assert copied != "not a message" + assert copied is not None + + +def test_matrix_assignment() -> None: + source = TwistWithCovariance(covariance=np.arange(36, dtype=float)) + matrix = np.asarray(source.covariance).reshape(6, 6) + assert matrix[0, 0] == 0 + assert matrix[5, 5] == 35 + source.covariance = (np.eye(6) * 2).ravel() + np.testing.assert_array_equal(list(source.covariance)[:6], [2, 0, 0, 0, 0, 0]) + assert np.trace(np.asarray(source.covariance).reshape(6, 6)) == 12 + + +@pytest.mark.parametrize("size", [0, 35, 37]) +def test_fixed_covariance_rejects_invalid_length(size: int) -> None: + with pytest.raises((ValueError, TypeError, RuntimeError)): + TwistWithCovariance(covariance=np.zeros(size)) @pytest.mark.parametrize( - "linear,angular", + "xyz,angular", [ - ([0.0, 0.0, 0.0], [0.0, 0.0, 0.0]), - ([1.0, 2.0, 3.0], [0.1, 0.2, 0.3]), - ([-1.0, -2.0, -3.0], [-0.1, -0.2, -0.3]), - ([100.0, -100.0, 0.0], [3.14, -3.14, 0.0]), + ((0, 0, 0), (0, 0, 0)), + ((1, 2, 3), (0.1, 0.2, 0.3)), + ((-1, -2, -3), (-0.1, -0.2, -0.3)), + ((100, -100, 0), (3.14, -3.14, 0)), ], ) -def test_twist_with_covariance_parametrized_velocities(linear, angular) -> None: - """Parametrized test for various velocity values.""" - twist = Twist(linear, angular) - twist_cov = TwistWithCovariance(twist) - - assert twist_cov.linear.x == linear[0] - assert twist_cov.linear.y == linear[1] - assert twist_cov.linear.z == linear[2] - assert twist_cov.angular.x == angular[0] - assert twist_cov.angular.y == angular[1] - assert twist_cov.angular.z == angular[2] +def test_parameterized_values( + xyz: tuple[float, float, float], angular: tuple[float, float, float] +) -> None: + source = TwistWithCovariance( + twist=Twist( + linear=Vector3(x=xyz[0], y=xyz[1], z=xyz[2]), + angular=Vector3(x=angular[0], y=angular[1], z=angular[2]), + ) + ) + assert (source.twist.linear.x, source.twist.linear.y, source.twist.linear.z) == xyz + assert (source.twist.angular.x, source.twist.angular.y, source.twist.angular.z) == angular diff --git a/dimos/msgs/geometry_msgs/test_TwistWithCovarianceStamped.py b/dimos/msgs/geometry_msgs/test_TwistWithCovarianceStamped.py index e5195fb24d8..da7c1382ba0 100644 --- a/dimos/msgs/geometry_msgs/test_TwistWithCovarianceStamped.py +++ b/dimos/msgs/geometry_msgs/test_TwistWithCovarianceStamped.py @@ -12,240 +12,113 @@ # See the License for the specific language governing permissions and # limitations under the License. -import time - -import numpy as np - -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.TwistWithCovariance import TwistWithCovariance -from dimos.msgs.geometry_msgs.TwistWithCovarianceStamped import TwistWithCovarianceStamped -from dimos.msgs.geometry_msgs.Vector3 import Vector3 - - -def test_twist_with_covariance_stamped_default_init() -> None: - """Test default initialization.""" - twist_cov_stamped = TwistWithCovarianceStamped() - - # Should have current timestamp - assert twist_cov_stamped.ts > 0 - assert twist_cov_stamped.frame_id == "" - - # Twist should be zero - assert twist_cov_stamped.twist.linear.x == 0.0 - assert twist_cov_stamped.twist.linear.y == 0.0 - assert twist_cov_stamped.twist.linear.z == 0.0 - assert twist_cov_stamped.twist.angular.x == 0.0 - assert twist_cov_stamped.twist.angular.y == 0.0 - assert twist_cov_stamped.twist.angular.z == 0.0 - - # Covariance should be all zeros - assert np.all(twist_cov_stamped.covariance == 0.0) - - -def test_twist_with_covariance_stamped_with_timestamp() -> None: - """Test initialization with specific timestamp.""" - ts = 1234567890.123456 - frame_id = "base_link" - twist_cov_stamped = TwistWithCovarianceStamped(ts=ts, frame_id=frame_id) - - assert twist_cov_stamped.ts == ts - assert twist_cov_stamped.frame_id == frame_id - - -def test_twist_with_covariance_stamped_with_twist() -> None: - """Test initialization with twist.""" - ts = 1234567890.123456 - frame_id = "odom" - twist = Twist(Vector3(1.0, 2.0, 3.0), Vector3(0.1, 0.2, 0.3)) - covariance = np.arange(36, dtype=float) +"""Nested generated Twist covariance with explicit stamps and array conversion. - twist_cov_stamped = TwistWithCovarianceStamped( - ts=ts, frame_id=frame_id, twist=twist, covariance=covariance - ) - - assert twist_cov_stamped.ts == ts - assert twist_cov_stamped.frame_id == frame_id - assert twist_cov_stamped.twist.linear.x == 1.0 - assert twist_cov_stamped.twist.linear.y == 2.0 - assert twist_cov_stamped.twist.linear.z == 3.0 - assert np.array_equal(twist_cov_stamped.covariance, covariance) - - -def test_twist_with_covariance_stamped_with_tuple() -> None: - """Test initialization with tuple of velocities.""" - ts = 1234567890.123456 - frame_id = "robot_base" - linear = [1.0, 2.0, 3.0] - angular = [0.1, 0.2, 0.3] - covariance = np.arange(36, dtype=float) - - twist_cov_stamped = TwistWithCovarianceStamped( - ts=ts, frame_id=frame_id, twist=(linear, angular), covariance=covariance - ) - - assert twist_cov_stamped.ts == ts - assert twist_cov_stamped.frame_id == frame_id - assert twist_cov_stamped.twist.linear.x == 1.0 - assert twist_cov_stamped.twist.angular.x == 0.1 - assert np.array_equal(twist_cov_stamped.covariance, covariance) - - -def test_twist_with_covariance_stamped_properties() -> None: - """Test convenience properties.""" - twist = Twist(Vector3(1.0, 2.0, 3.0), Vector3(0.1, 0.2, 0.3)) - covariance = np.eye(6).flatten() - twist_cov_stamped = TwistWithCovarianceStamped( - ts=1234567890.0, frame_id="cmd_vel", twist=twist, covariance=covariance - ) - - # Linear and angular properties - assert twist_cov_stamped.linear.x == 1.0 - assert twist_cov_stamped.linear.y == 2.0 - assert twist_cov_stamped.linear.z == 3.0 - assert twist_cov_stamped.angular.x == 0.1 - assert twist_cov_stamped.angular.y == 0.2 - assert twist_cov_stamped.angular.z == 0.3 - - # Covariance matrix - cov_matrix = twist_cov_stamped.covariance_matrix - assert cov_matrix.shape == (6, 6) - assert np.trace(cov_matrix) == 6.0 +Legacy inheritance, custom repr, and implicit wall-clock defaults are retired; +field values, covariance matrices, and time conversion remain checked explicitly. +""" +import pickle +import time -def test_twist_with_covariance_stamped_str() -> None: - """Test string representation.""" - twist = Twist(Vector3(1.234, 2.567, 3.891), Vector3(0.111, 0.222, 0.333)) - covariance = np.eye(6).flatten() * 2.0 - twist_cov_stamped = TwistWithCovarianceStamped( - ts=1234567890.0, frame_id="world", twist=twist, covariance=covariance +from dimos_generated.geometry_msgs.msg import ( + Twist, + TwistWithCovariance, + TwistWithCovarianceStamped, + Vector3, +) +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest +from rosbags.typesys import Stores, get_typestore + +from dimos.msgs.time import header_now, time_from_nanoseconds, time_from_seconds, to_nanoseconds + + +def test_generated_defaults() -> None: + source = TwistWithCovarianceStamped() + assert source.header.frame_id == "" + assert to_nanoseconds(source.header.stamp) == 0 + assert isinstance(source.twist, TwistWithCovariance) + assert isinstance(source.twist.twist, Twist) + np.testing.assert_array_equal(source.twist.covariance, np.zeros(36)) + value = source.twist.twist + np.testing.assert_array_equal( + [ + value.linear.x, + value.linear.y, + value.linear.z, + value.angular.x, + value.angular.y, + value.angular.z, + ], + np.zeros(6), ) - str_repr = str(twist_cov_stamped) - assert "TwistWithCovarianceStamped" in str_repr - assert "1.234" in str_repr - assert "2.567" in str_repr - assert "3.891" in str_repr - assert "cov_trace" in str_repr - assert "12.000" in str_repr # Trace of 2*identity is 12 - - -def test_twist_with_covariance_stamped_lcm_encode_decode() -> None: - """Test LCM encoding and decoding.""" - ts = 1234567890.123456 - frame_id = "camera_link" - twist = Twist(Vector3(1.0, 2.0, 3.0), Vector3(0.1, 0.2, 0.3)) - covariance = np.arange(36, dtype=float) +@pytest.mark.parametrize("stamp", [0, 1234567890123456789]) +@pytest.mark.parametrize( + "covariance", + [np.zeros(36), np.eye(6).ravel(), np.arange(36, dtype=float), np.diag(np.arange(1, 7)).ravel()], +) +def test_fields_covariance_and_independent_cdr(stamp: int, covariance: np.ndarray) -> None: source = TwistWithCovarianceStamped( - ts=ts, frame_id=frame_id, twist=twist, covariance=covariance + header=Header(stamp=time_from_nanoseconds(stamp), frame_id="camera_link"), + twist=TwistWithCovariance( + twist=Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.1, y=0.2, z=0.3)), + covariance=covariance, + ), ) - - # Encode and decode - binary_msg = source.lcm_encode() - decoded = TwistWithCovarianceStamped.lcm_decode(binary_msg) - - # Check timestamp (may lose some precision) - assert abs(decoded.ts - ts) < 1e-6 - assert decoded.frame_id == frame_id - - # Check twist - assert decoded.twist.linear.x == 1.0 - assert decoded.twist.linear.y == 2.0 - assert decoded.twist.linear.z == 3.0 - assert decoded.twist.angular.x == 0.1 - assert decoded.twist.angular.y == 0.2 - assert decoded.twist.angular.z == 0.3 - - # Check covariance - assert np.array_equal(decoded.covariance, covariance) - - -def test_twist_with_covariance_stamped_zero_timestamp() -> None: - """Test that an explicit zero timestamp is kept, and an omitted one stamps now.""" - assert TwistWithCovarianceStamped(ts=0.0).ts == 0.0 - - before = time.time() - assert before <= TwistWithCovarianceStamped().ts <= time.time() - - -def test_twist_with_covariance_stamped_inheritance() -> None: - """Test that it properly inherits from TwistWithCovariance and Timestamped.""" - twist = Twist(Vector3(1.0, 2.0, 3.0), Vector3(0.1, 0.2, 0.3)) - covariance = np.eye(6).flatten() - twist_cov_stamped = TwistWithCovarianceStamped( - ts=1234567890.0, frame_id="test", twist=twist, covariance=covariance + decoded = TwistWithCovarianceStamped.decode(source.encode()) + independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr( + source.encode(), TwistWithCovarianceStamped.msg_name ) - - # Should be instance of parent classes - assert isinstance(twist_cov_stamped, TwistWithCovariance) - - # Should have Timestamped attributes - assert hasattr(twist_cov_stamped, "ts") - assert hasattr(twist_cov_stamped, "frame_id") - - # Should have TwistWithCovariance attributes - assert hasattr(twist_cov_stamped, "twist") - assert hasattr(twist_cov_stamped, "covariance") - - -def test_twist_with_covariance_stamped_is_zero() -> None: - """Test is_zero method inheritance.""" - # Zero twist - twist_cov_stamped1 = TwistWithCovarianceStamped() - assert twist_cov_stamped1.is_zero() - assert not twist_cov_stamped1 # Boolean conversion - - # Non-zero twist - twist = Twist(Vector3(0.5, 0.0, 0.0), Vector3(0.0, 0.0, 0.0)) - twist_cov_stamped2 = TwistWithCovarianceStamped(twist=twist) - assert not twist_cov_stamped2.is_zero() - assert twist_cov_stamped2 # Boolean conversion - - -def test_twist_with_covariance_stamped_sec_nsec() -> None: - """Test the sec_nsec helper function.""" - from dimos.msgs.geometry_msgs.TwistWithCovarianceStamped import sec_nsec - - # Test integer seconds - s, ns = sec_nsec(1234567890.0) - assert s == 1234567890 - assert ns == 0 - - # Test fractional seconds - s, ns = sec_nsec(1234567890.123456789) - assert s == 1234567890 - assert abs(ns - 123456789) < 100 # Allow small rounding error - - # Test small fractional seconds - s, ns = sec_nsec(0.000000001) - assert s == 0 - assert ns == 1 - - # Test large timestamp - s, ns = sec_nsec(9999999999.999999999) - # Due to floating point precision, this might round to 10000000000 - assert s in [9999999999, 10000000000] - if s == 9999999999: - assert abs(ns - 999999999) < 10 - else: - assert ns == 0 - - -def test_twist_with_covariance_stamped_different_covariances() -> None: - """Test with different covariance patterns.""" - twist = Twist(Vector3(1.0, 0.0, 0.0), Vector3(0.0, 0.0, 0.5)) - - # Zero covariance - zero_cov = np.zeros(36) - twist_cov1 = TwistWithCovarianceStamped(twist=twist, covariance=zero_cov) - assert np.all(twist_cov1.covariance == 0.0) - - # Identity covariance - identity_cov = np.eye(6).flatten() - twist_cov2 = TwistWithCovarianceStamped(twist=twist, covariance=identity_cov) - assert np.trace(twist_cov2.covariance_matrix) == 6.0 - - # Full covariance - full_cov = np.random.rand(36) - twist_cov3 = TwistWithCovarianceStamped(twist=twist, covariance=full_cov) - assert np.array_equal(twist_cov3.covariance, full_cov) + for result in (decoded, independent): + assert result.header.frame_id == "camera_link" + assert result.header.stamp.sec * 1000000000 + result.header.stamp.nanosec == stamp + value = result.twist.twist + np.testing.assert_array_equal( + [ + value.linear.x, + value.linear.y, + value.linear.z, + value.angular.x, + value.angular.y, + value.angular.z, + ], + [1, 2, 3, 0.1, 0.2, 0.3], + ) + np.testing.assert_array_equal(result.twist.covariance, covariance) + matrix = np.asarray(result.twist.covariance).reshape(6, 6) + assert matrix.shape == (6, 6) + assert np.trace(matrix) == np.trace(covariance.reshape(6, 6)) + copied = pickle.loads(pickle.dumps(source)) + assert copied.encode() == source.encode() + copied.twist.covariance[0] = 999 + assert source.twist.covariance[0] == covariance[0] + + +def test_explicit_current_header() -> None: + before = time.time_ns() + source = TwistWithCovarianceStamped(header=header_now("test")) + assert before <= to_nanoseconds(source.header.stamp) <= time.time_ns() + assert source.header.frame_id == "test" + + +@pytest.mark.parametrize( + "seconds,sec,nanosec", + [ + (1234567890.0, 1234567890, 0), + (1234567890.123456789, 1234567890, 123456789), + (0.000000001, 0, 1), + ], +) +def test_explicit_seconds_conversion(seconds: float, sec: int, nanosec: int) -> None: + stamp = time_from_seconds(seconds) + assert stamp.sec == sec + assert abs(stamp.nanosec - nanosec) < 100 + + +def test_seconds_outside_ros_int32_range_rejected() -> None: + with pytest.raises((TypeError, ValueError, OverflowError)): + time_from_seconds(9999999999.999999999) diff --git a/dimos/msgs/geometry_msgs/test_Vector3.py b/dimos/msgs/geometry_msgs/test_Vector3.py index 099e35eb19e..05941ddea08 100644 --- a/dimos/msgs/geometry_msgs/test_Vector3.py +++ b/dimos/msgs/geometry_msgs/test_Vector3.py @@ -12,10 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.geometry_msgs.msg import Vector3 import numpy as np import pytest -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.geometry import ( + normalized_vector, + quaternion_euler, + quaternion_from_euler, + vector_array, + vector_from_array, +) def test_vector_default_init() -> None: @@ -24,55 +31,47 @@ def test_vector_default_init() -> None: assert v.x == 0.0 assert v.y == 0.0 assert v.z == 0.0 - assert len(v.data) == 3 - assert v.to_list() == [0.0, 0.0, 0.0] - assert v.is_zero() # Zero vector should be considered zero + assert len(vector_array(v)) == 3 + assert vector_array(v).tolist() == [0.0, 0.0, 0.0] + assert np.allclose(vector_array(v), 0) def test_vector_specific_init() -> None: """Test initialization with specific values and different input types.""" - - v1 = Vector3(1.0, 2.0) # 2D vector (now becomes 3D with z=0) + v1 = Vector3(x=1.0, y=2.0) assert v1.x == 1.0 assert v1.y == 2.0 assert v1.z == 0.0 - - v2 = Vector3(3.0, 4.0, 5.0) # 3D vector + v2 = Vector3(x=3.0, y=4.0, z=5.0) assert v2.x == 3.0 assert v2.y == 4.0 assert v2.z == 5.0 - - v3 = Vector3([6.0, 7.0, 8.0]) + v3 = vector_from_array([6.0, 7.0, 8.0]) assert v3.x == 6.0 assert v3.y == 7.0 assert v3.z == 8.0 - - v4 = Vector3((9.0, 10.0, 11.0)) + v4 = vector_from_array((9.0, 10.0, 11.0)) assert v4.x == 9.0 assert v4.y == 10.0 assert v4.z == 11.0 - - v5 = Vector3(np.array([12.0, 13.0, 14.0])) + v5 = vector_from_array(np.array([12.0, 13.0, 14.0])) assert v5.x == 12.0 assert v5.y == 13.0 assert v5.z == 14.0 - - original = Vector3([15.0, 16.0, 17.0]) - v6 = Vector3(original) + original = vector_from_array([15.0, 16.0, 17.0]) + v6 = Vector3.decode(original.encode()) assert v6.x == 15.0 assert v6.y == 16.0 assert v6.z == 17.0 - assert v6 is not original assert v6 == original def test_vector_addition() -> None: """Test vector addition.""" - v1 = Vector3(1.0, 2.0, 3.0) - v2 = Vector3(4.0, 5.0, 6.0) - - v_add = v1 + v2 + v1 = Vector3(x=1.0, y=2.0, z=3.0) + v2 = Vector3(x=4.0, y=5.0, z=6.0) + v_add = vector_from_array(vector_array(v1) + vector_array(v2)) assert v_add.x == 5.0 assert v_add.y == 7.0 assert v_add.z == 9.0 @@ -80,10 +79,9 @@ def test_vector_addition() -> None: def test_vector_subtraction() -> None: """Test vector subtraction.""" - v1 = Vector3(1.0, 2.0, 3.0) - v2 = Vector3(4.0, 5.0, 6.0) - - v_sub = v2 - v1 + v1 = Vector3(x=1.0, y=2.0, z=3.0) + v2 = Vector3(x=4.0, y=5.0, z=6.0) + v_sub = vector_from_array(vector_array(v2) - vector_array(v1)) assert v_sub.x == 3.0 assert v_sub.y == 3.0 assert v_sub.z == 3.0 @@ -91,15 +89,12 @@ def test_vector_subtraction() -> None: def test_vector_scalar_multiplication() -> None: """Test vector multiplication by a scalar.""" - v1 = Vector3(1.0, 2.0, 3.0) - - v_mul = v1 * 2.0 + v1 = Vector3(x=1.0, y=2.0, z=3.0) + v_mul = vector_from_array(vector_array(v1) * 2.0) assert v_mul.x == 2.0 assert v_mul.y == 4.0 assert v_mul.z == 6.0 - - # Test right multiplication - v_rmul = 2.0 * v1 + v_rmul = vector_from_array(2.0 * vector_array(v1)) assert v_rmul.x == 2.0 assert v_rmul.y == 4.0 assert v_rmul.z == 6.0 @@ -107,9 +102,8 @@ def test_vector_scalar_multiplication() -> None: def test_vector_scalar_division() -> None: """Test vector division by a scalar.""" - v2 = Vector3(4.0, 5.0, 6.0) - - v_div = v2 / 2.0 + v2 = Vector3(x=4.0, y=5.0, z=6.0) + v_div = vector_from_array(vector_array(v2) / 2.0) assert v_div.x == 2.0 assert v_div.y == 2.5 assert v_div.z == 3.0 @@ -117,67 +111,54 @@ def test_vector_scalar_division() -> None: def test_vector_dot_product() -> None: """Test vector dot product.""" - v1 = Vector3(1.0, 2.0, 3.0) - v2 = Vector3(4.0, 5.0, 6.0) - - dot = v1.dot(v2) + v1 = Vector3(x=1.0, y=2.0, z=3.0) + v2 = Vector3(x=4.0, y=5.0, z=6.0) + dot = np.dot(vector_array(v1), vector_array(v2)) assert dot == 32.0 def test_vector_length() -> None: """Test vector length calculation.""" - # 2D vector with length 5 (now 3D with z=0) - v1 = Vector3(3.0, 4.0) - assert v1.length() == 5.0 - - # 3D vector - v2 = Vector3(2.0, 3.0, 6.0) - assert v2.length() == pytest.approx(7.0, 0.001) - - # Test length_squared - assert v1.length_squared() == 25.0 - assert v2.length_squared() == 49.0 + v1 = Vector3(x=3.0, y=4.0) + assert np.linalg.norm(vector_array(v1)) == 5.0 + v2 = Vector3(x=2.0, y=3.0, z=6.0) + assert np.linalg.norm(vector_array(v2)) == pytest.approx(7.0, 0.001) + assert np.dot(vector_array(v1), vector_array(v1)) == 25.0 + assert np.dot(vector_array(v2), vector_array(v2)) == 49.0 def test_vector_normalize() -> None: """Test vector normalization.""" - v = Vector3(2.0, 3.0, 6.0) - assert not v.is_zero() - - v_norm = v.normalize() - length = v.length() + v = Vector3(x=2.0, y=3.0, z=6.0) + assert not np.allclose(vector_array(v), 0) + v_norm = normalized_vector(v) + length = np.linalg.norm(vector_array(v)) expected_x = 2.0 / length expected_y = 3.0 / length expected_z = 6.0 / length - assert np.isclose(v_norm.x, expected_x) assert np.isclose(v_norm.y, expected_y) assert np.isclose(v_norm.z, expected_z) - assert np.isclose(v_norm.length(), 1.0) - assert not v_norm.is_zero() - - # Test normalizing a zero vector - v_zero = Vector3(0.0, 0.0, 0.0) - assert v_zero.is_zero() - v_zero_norm = v_zero.normalize() + assert np.isclose(np.linalg.norm(vector_array(v_norm)), 1.0) + assert not np.allclose(vector_array(v_norm), 0) + v_zero = Vector3(x=0.0, y=0.0, z=0.0) + assert np.allclose(vector_array(v_zero), 0) + v_zero_norm = normalized_vector(v_zero) assert v_zero_norm.x == 0.0 assert v_zero_norm.y == 0.0 assert v_zero_norm.z == 0.0 - assert v_zero_norm.is_zero() + assert np.allclose(vector_array(v_zero_norm), 0) def test_vector_to_2d() -> None: """Test conversion to 2D vector.""" - v = Vector3(2.0, 3.0, 6.0) - - v_2d = v.to_2d() + v = Vector3(x=2.0, y=3.0, z=6.0) + v_2d = Vector3(x=v.x, y=v.y) assert v_2d.x == 2.0 assert v_2d.y == 3.0 - assert v_2d.z == 0.0 # z should be 0 for 2D conversion - - # Already 2D vector (z=0) - v2 = Vector3(4.0, 5.0) - v2_2d = v2.to_2d() + assert v_2d.z == 0.0 + v2 = Vector3(x=4.0, y=5.0) + v2_2d = Vector3(x=v2.x, y=v2.y) assert v2_2d.x == 4.0 assert v2_2d.y == 5.0 assert v2_2d.z == 0.0 @@ -185,46 +166,32 @@ def test_vector_to_2d() -> None: def test_vector_distance() -> None: """Test distance calculations between vectors.""" - v1 = Vector3(1.0, 2.0, 3.0) - v2 = Vector3(4.0, 6.0, 8.0) - - # Distance - dist = v1.distance(v2) - expected_dist = np.sqrt(9.0 + 16.0 + 25.0) # sqrt((4-1)² + (6-2)² + (8-3)²) + v1 = Vector3(x=1.0, y=2.0, z=3.0) + v2 = Vector3(x=4.0, y=6.0, z=8.0) + dist = np.linalg.norm(vector_array(v1) - vector_array(v2)) + expected_dist = np.sqrt(9.0 + 16.0 + 25.0) assert dist == pytest.approx(expected_dist) - - # Distance squared - dist_sq = v1.distance_squared(v2) - assert dist_sq == 50.0 # 9 + 16 + 25 + dist_sq = np.sum((vector_array(v1) - vector_array(v2)) ** 2) + assert dist_sq == 50.0 def test_vector_cross_product() -> None: """Test vector cross product.""" - v1 = Vector3(1.0, 0.0, 0.0) # Unit x vector - v2 = Vector3(0.0, 1.0, 0.0) # Unit y vector - - # v1 × v2 should be unit z vector - cross = v1.cross(v2) + v1 = Vector3(x=1.0, y=0.0, z=0.0) + v2 = Vector3(x=0.0, y=1.0, z=0.0) + cross = vector_from_array(np.cross(vector_array(v1), vector_array(v2))) assert cross.x == 0.0 assert cross.y == 0.0 assert cross.z == 1.0 - - # Test with more complex vectors - a = Vector3(2.0, 3.0, 4.0) - b = Vector3(5.0, 6.0, 7.0) - c = a.cross(b) - - # Cross product manually calculated: - # (3*7-4*6, 4*5-2*7, 2*6-3*5) + a = Vector3(x=2.0, y=3.0, z=4.0) + b = Vector3(x=5.0, y=6.0, z=7.0) + c = vector_from_array(np.cross(vector_array(a), vector_array(b))) assert c.x == -3.0 assert c.y == 6.0 assert c.z == -3.0 - - # Test with vectors that have z=0 (still works as they're 3D) - v_2d1 = Vector3(1.0, 2.0) # (1, 2, 0) - v_2d2 = Vector3(3.0, 4.0) # (3, 4, 0) - cross_2d = v_2d1.cross(v_2d2) - # (2*0-0*4, 0*3-1*0, 1*4-2*3) = (0, 0, -2) + v_2d1 = Vector3(x=1.0, y=2.0) + v_2d2 = Vector3(x=3.0, y=4.0) + cross_2d = vector_from_array(np.cross(vector_array(v_2d1), vector_array(v_2d2))) assert cross_2d.x == 0.0 assert cross_2d.y == 0.0 assert cross_2d.z == -2.0 @@ -232,18 +199,16 @@ def test_vector_cross_product() -> None: def test_vector_zeros() -> None: """Test Vector3.zeros class method.""" - # 3D zero vector - v_zeros = Vector3.zeros() + v_zeros = Vector3() assert v_zeros.x == 0.0 assert v_zeros.y == 0.0 assert v_zeros.z == 0.0 - assert v_zeros.is_zero() + assert np.allclose(vector_array(v_zeros), 0) def test_vector_ones() -> None: """Test Vector3.ones class method.""" - # 3D ones vector - v_ones = Vector3.ones() + v_ones = Vector3(x=1, y=1, z=1) assert v_ones.x == 1.0 assert v_ones.y == 1.0 assert v_ones.z == 1.0 @@ -251,198 +216,134 @@ def test_vector_ones() -> None: def test_vector_conversion_methods() -> None: """Test vector conversion methods (to_list, to_tuple, to_numpy).""" - v = Vector3(1.0, 2.0, 3.0) - - # to_list - assert v.to_list() == [1.0, 2.0, 3.0] - - # to_tuple - assert v.to_tuple() == (1.0, 2.0, 3.0) - - # to_numpy - np_array = v.to_numpy() + v = Vector3(x=1.0, y=2.0, z=3.0) + assert vector_array(v).tolist() == [1.0, 2.0, 3.0] + assert tuple(vector_array(v)) == (1.0, 2.0, 3.0) + np_array = vector_array(v) assert isinstance(np_array, np.ndarray) assert np.array_equal(np_array, np.array([1.0, 2.0, 3.0])) def test_vector_equality() -> None: """Test vector equality.""" - v1 = Vector3(1, 2, 3) - v2 = Vector3(1, 2, 3) - v3 = Vector3(4, 5, 6) - + v1 = Vector3(x=1, y=2, z=3) + v2 = Vector3(x=1, y=2, z=3) + v3 = Vector3(x=4, y=5, z=6) assert v1 == v2 assert v1 != v3 - assert v1 != Vector3(1, 2) # Now (1, 2, 0) vs (1, 2, 3) - assert v1 != Vector3(1.1, 2, 3) # Different values + assert v1 != Vector3(x=1, y=2) + assert v1 != Vector3(x=1.1, y=2, z=3) assert v1 != [1, 2, 3] def test_vector_is_zero() -> None: """Test is_zero method for vectors.""" - # Default zero vector v0 = Vector3() - assert v0.is_zero() - - # Explicit zero vector - v1 = Vector3(0.0, 0.0, 0.0) - assert v1.is_zero() - - # Zero vector with different initialization (now always 3D) - v2 = Vector3(0.0, 0.0) # Becomes (0, 0, 0) - assert v2.is_zero() - - # Non-zero vectors - v3 = Vector3(1.0, 0.0, 0.0) - assert not v3.is_zero() - - v4 = Vector3(0.0, 2.0, 0.0) - assert not v4.is_zero() - - v5 = Vector3(0.0, 0.0, 3.0) - assert not v5.is_zero() - - # Almost zero (within tolerance) - v6 = Vector3(1e-10, 1e-10, 1e-10) - assert v6.is_zero() - - # Almost zero (outside tolerance) - v7 = Vector3(1e-6, 1e-6, 1e-6) - assert not v7.is_zero() + assert np.allclose(vector_array(v0), 0) + v1 = Vector3(x=0.0, y=0.0, z=0.0) + assert np.allclose(vector_array(v1), 0) + v2 = Vector3(x=0.0, y=0.0) + assert np.allclose(vector_array(v2), 0) + v3 = Vector3(x=1.0, y=0.0, z=0.0) + assert not np.allclose(vector_array(v3), 0) + v4 = Vector3(x=0.0, y=2.0, z=0.0) + assert not np.allclose(vector_array(v4), 0) + v5 = Vector3(x=0.0, y=0.0, z=3.0) + assert not np.allclose(vector_array(v5), 0) + v6 = Vector3(x=1e-10, y=1e-10, z=1e-10) + assert np.allclose(vector_array(v6), 0) + v7 = Vector3(x=1e-06, y=1e-06, z=1e-06) + assert not np.allclose(vector_array(v7), 0) def test_vector_bool_conversion(): """Test boolean conversion of vectors.""" - # Zero vectors should be False v0 = Vector3() - assert not bool(v0) - - v1 = Vector3(0.0, 0.0, 0.0) - assert not bool(v1) - - # Almost zero vectors should be False - v2 = Vector3(1e-10, 1e-10, 1e-10) - assert not bool(v2) - - # Non-zero vectors should be True - v3 = Vector3(1.0, 0.0, 0.0) - assert bool(v3) - - v4 = Vector3(0.0, 2.0, 0.0) - assert bool(v4) - - v5 = Vector3(0.0, 0.0, 3.0) - assert bool(v5) - - # Direct use in if statements - if v0: + assert not not np.allclose(vector_array(v0), 0) + v1 = Vector3(x=0.0, y=0.0, z=0.0) + assert not not np.allclose(vector_array(v1), 0) + v2 = Vector3(x=1e-10, y=1e-10, z=1e-10) + assert not not np.allclose(vector_array(v2), 0) + v3 = Vector3(x=1.0, y=0.0, z=0.0) + assert not np.allclose(vector_array(v3), 0) + v4 = Vector3(x=0.0, y=2.0, z=0.0) + assert not np.allclose(vector_array(v4), 0) + v5 = Vector3(x=0.0, y=0.0, z=3.0) + assert not np.allclose(vector_array(v5), 0) + if not np.allclose(vector_array(v0), 0): raise AssertionError("Zero vector should be False in boolean context") else: - pass # Expected path - - if v3: - pass # Expected path + pass + if not np.allclose(vector_array(v3), 0): + pass else: raise AssertionError("Non-zero vector should be True in boolean context") def test_vector_add() -> None: """Test vector addition operator.""" - v1 = Vector3(1.0, 2.0, 3.0) - v2 = Vector3(4.0, 5.0, 6.0) - - # Using __add__ method - v_add = v1.__add__(v2) + v1 = Vector3(x=1.0, y=2.0, z=3.0) + v2 = Vector3(x=4.0, y=5.0, z=6.0) + v_add = vector_from_array(vector_array(v1) + vector_array(v2)) assert v_add.x == 5.0 assert v_add.y == 7.0 assert v_add.z == 9.0 - - # Using + operator - v_add_op = v1 + v2 + v_add_op = vector_from_array(vector_array(v1) + vector_array(v2)) assert v_add_op.x == 5.0 assert v_add_op.y == 7.0 assert v_add_op.z == 9.0 - - # Adding zero vector should return original vector - v_zero = Vector3.zeros() - assert (v1 + v_zero) == v1 + v_zero = Vector3() + assert vector_from_array(vector_array(v1) + vector_array(v_zero)) == v1 def test_vector_add_dim_mismatch() -> None: """Test vector addition with different input dimensions (now all vectors are 3D).""" - v1 = Vector3(1.0, 2.0) # Becomes (1, 2, 0) - v2 = Vector3(4.0, 5.0, 6.0) # (4, 5, 6) - - # Using + operator - should work fine now since both are 3D - v_add_op = v1 + v2 - assert v_add_op.x == 5.0 # 1 + 4 - assert v_add_op.y == 7.0 # 2 + 5 - assert v_add_op.z == 6.0 # 0 + 6 + v1 = Vector3(x=1.0, y=2.0) + v2 = Vector3(x=4.0, y=5.0, z=6.0) + v_add_op = vector_from_array(vector_array(v1) + vector_array(v2)) + assert v_add_op.x == 5.0 + assert v_add_op.y == 7.0 + assert v_add_op.z == 6.0 def test_yaw_pitch_roll_accessors() -> None: """Test yaw, pitch, and roll accessor properties.""" - # Test with a 3D vector - v = Vector3(1.0, 2.0, 3.0) - - # According to standard convention: - # roll = rotation around x-axis = x component - # pitch = rotation around y-axis = y component - # yaw = rotation around z-axis = z component - assert v.roll == 1.0 # Should return x component - assert v.pitch == 2.0 # Should return y component - assert v.yaw == 3.0 # Should return z component - - # Test with a 2D vector (z should be 0.0) - v_2d = Vector3(4.0, 5.0) - assert v_2d.roll == 4.0 # Should return x component - assert v_2d.pitch == 5.0 # Should return y component - assert v_2d.yaw == 0.0 # Should return z component (defaults to 0 for 2D) - - # Test with empty vector (all should be 0.0) + v = Vector3(x=1.0, y=2.0, z=3.0) + assert v.x == 1.0 + assert v.y == 2.0 + assert v.z == 3.0 + v_2d = Vector3(x=4.0, y=5.0) + assert v_2d.x == 4.0 + assert v_2d.y == 5.0 + assert v_2d.z == 0.0 v_empty = Vector3() - assert v_empty.roll == 0.0 - assert v_empty.pitch == 0.0 - assert v_empty.yaw == 0.0 - - # Test with negative values - v_neg = Vector3(-1.5, -2.5, -3.5) - assert v_neg.roll == -1.5 - assert v_neg.pitch == -2.5 - assert v_neg.yaw == -3.5 + assert v_empty.x == 0.0 + assert v_empty.y == 0.0 + assert v_empty.z == 0.0 + v_neg = Vector3(x=-1.5, y=-2.5, z=-3.5) + assert v_neg.x == -1.5 + assert v_neg.y == -2.5 + assert v_neg.z == -3.5 def test_vector_to_quaternion() -> None: """Test vector to quaternion conversion.""" - # Test with zero Euler angles (should produce identity quaternion) - v_zero = Vector3(0.0, 0.0, 0.0) - q_identity = v_zero.to_quaternion() - - # Identity quaternion should have w=1, x=y=z=0 + v_zero = Vector3(x=0.0, y=0.0, z=0.0) + q_identity = quaternion_from_euler(v_zero.x, v_zero.y, v_zero.z) assert np.isclose(q_identity.x, 0.0, atol=1e-10) assert np.isclose(q_identity.y, 0.0, atol=1e-10) assert np.isclose(q_identity.z, 0.0, atol=1e-10) assert np.isclose(q_identity.w, 1.0, atol=1e-10) - - # Test with small angles (to avoid gimbal lock issues) - v_small = Vector3(0.1, 0.2, 0.3) # Small roll, pitch, yaw - q_small = v_small.to_quaternion() - - # Quaternion should be normalized (magnitude = 1) + v_small = Vector3(x=0.1, y=0.2, z=0.3) + q_small = quaternion_from_euler(v_small.x, v_small.y, v_small.z) magnitude = np.sqrt(q_small.x**2 + q_small.y**2 + q_small.z**2 + q_small.w**2) assert np.isclose(magnitude, 1.0, atol=1e-10) - - # Test conversion back to Euler (should be close to original) - v_back = q_small.to_euler() - assert np.isclose(v_back.x, 0.1, atol=1e-6) - assert np.isclose(v_back.y, 0.2, atol=1e-6) - assert np.isclose(v_back.z, 0.3, atol=1e-6) - - # Test with π/2 rotation around x-axis - v_x_90 = Vector3(np.pi / 2, 0.0, 0.0) - q_x_90 = v_x_90.to_quaternion() - - # Should be approximately (sin(π/4), 0, 0, cos(π/4)) = (√2/2, 0, 0, √2/2) + v_back = vector_from_array(quaternion_euler(q_small)) + assert np.isclose(v_back.x, 0.1, atol=1e-06) + assert np.isclose(v_back.y, 0.2, atol=1e-06) + assert np.isclose(v_back.z, 0.3, atol=1e-06) + v_x_90 = Vector3(x=np.pi / 2, y=0.0, z=0.0) + q_x_90 = quaternion_from_euler(v_x_90.x, v_x_90.y, v_x_90.z) expected = np.sqrt(2) / 2 assert np.isclose(q_x_90.x, expected, atol=1e-10) assert np.isclose(q_x_90.y, 0.0, atol=1e-10) @@ -450,13 +351,10 @@ def test_vector_to_quaternion() -> None: assert np.isclose(q_x_90.w, expected, atol=1e-10) -def test_lcm_encode_decode() -> None: - v_source = Vector3(1.0, 2.0, 3.0) - - binary_msg = v_source.lcm_encode() - - v_dest = Vector3.lcm_decode(binary_msg) - +def test_cdr_encode_decode() -> None: + v_source = Vector3(x=1.0, y=2.0, z=3.0) + binary_msg = v_source.encode() + v_dest = Vector3.decode(binary_msg) assert isinstance(v_dest, Vector3) assert v_dest is not v_source assert v_dest == v_source diff --git a/dimos/msgs/geometry_msgs/test_WrenchStamped.py b/dimos/msgs/geometry_msgs/test_WrenchStamped.py index e493b1cd396..2c3b932f12f 100644 --- a/dimos/msgs/geometry_msgs/test_WrenchStamped.py +++ b/dimos/msgs/geometry_msgs/test_WrenchStamped.py @@ -12,129 +12,109 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Generated force/torque values and explicit external numeric operations.""" + import pickle import time +from dimos_generated.geometry_msgs.msg import Vector3, Wrench, WrenchStamped +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest +from rosbags.typesys import Stores, get_typestore -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.geometry_msgs.Wrench import Wrench -from dimos.msgs.geometry_msgs.WrenchStamped import WrenchStamped - - -def test_wrench_empty_is_zero(): - w = Wrench() - assert w.force.to_tuple() == (0.0, 0.0, 0.0) - assert w.torque.to_tuple() == (0.0, 0.0, 0.0) - assert w.is_zero() - assert not w - - -def test_wrench_from_two_sequences(): - w = Wrench([1, 2, 3], [4, 5, 6]) - assert w.force.to_tuple() == (1.0, 2.0, 3.0) - assert w.torque.to_tuple() == (4.0, 5.0, 6.0) - +from dimos.msgs.geometry import wrench_array, wrench_from_array +from dimos.msgs.time import header_now, time_from_seconds, to_nanoseconds -def test_wrench_from_two_vector3(): - w = Wrench(Vector3(1, 2, 3), Vector3(4, 5, 6)) - assert w.force.to_tuple() == (1.0, 2.0, 3.0) - assert w.torque.to_tuple() == (4.0, 5.0, 6.0) +def test_wrench_empty_is_zero() -> None: + np.testing.assert_array_equal(wrench_array(Wrench()), np.zeros(6)) -def test_wrench_keywords(): - assert Wrench(force=[1, 2, 3]).torque.to_tuple() == (0.0, 0.0, 0.0) - assert Wrench(torque=[4, 5, 6]).force.to_tuple() == (0.0, 0.0, 0.0) +def test_wrench_from_array_and_vectors() -> None: + source = wrench_from_array([1, 2, 3, 4, 5, 6]) + explicit = Wrench(force=Vector3(x=1, y=2, z=3), torque=Vector3(x=4, y=5, z=6)) + assert source == explicit + np.testing.assert_array_equal(wrench_array(source), [1, 2, 3, 4, 5, 6]) -def test_wrench_copy_constructor(): - src = Wrench([1, 2, 3], [4, 5, 6]) - copy = Wrench(src) - assert copy == src - assert copy.force is not src.force - -def test_wrench_unknown_keyword_raises(): +def test_wrench_keywords() -> None: + assert Wrench(force=Vector3(x=1, y=2, z=3)).torque == Vector3() + assert Wrench(torque=Vector3(x=4, y=5, z=6)).force == Vector3() with pytest.raises(TypeError): Wrench(bogus=1) -def test_wrench_array_roundtrip(): - ft = [1.0, 2.0, 3.0, 0.1, 0.2, 0.3] - assert np.allclose(Wrench.from_array(ft).to_array(), ft) - - -def test_wrench_from_array_wrong_length_raises(): - with pytest.raises(ValueError, match="6 elements"): - Wrench.from_array([1.0, 2.0, 3.0]) - - -def test_wrench_add_and_sub(): - a = Wrench([1, 2, 3], [4, 5, 6]) - b = Wrench([1, 1, 1], [1, 1, 1]) - assert (a + b).force.to_tuple() == (2.0, 3.0, 4.0) - assert (a - b).torque.to_tuple() == (3.0, 4.0, 5.0) +def test_wrench_copy_storage() -> None: + source = wrench_from_array([1, 2, 3, 4, 5, 6]) + copied = Wrench.decode(source.encode()) + assert copied == source + copied.force.x = 10 + assert source.force.x == 1 + array = wrench_array(source) + array[0] = 20 + assert source.force.x == 1 -def test_stamped_is_a_wrench(): - assert isinstance(WrenchStamped(), Wrench) +def test_wrench_array_roundtrip() -> None: + values = [1, 2, 3, 0.1, 0.2, 0.3] + np.testing.assert_allclose(wrench_array(wrench_from_array(values)), values) -def test_stamped_inherits_wrench_positional_form(): - ws = WrenchStamped([1, 2, 3], [4, 5, 6]) - assert ws.force.to_tuple() == (1.0, 2.0, 3.0) - assert ws.torque.to_tuple() == (4.0, 5.0, 6.0) - - -def test_stamped_keeps_explicit_ts_and_frame(): - ws = WrenchStamped(ts=5.0, frame_id="ft_sensor", force=[1, 2, 3], torque=[4, 5, 6]) - assert (ws.ts, ws.frame_id) == (5.0, "ft_sensor") - assert ws.force.to_tuple() == (1.0, 2.0, 3.0) - - -def test_stamped_positional_ts_and_frame(): - ws = WrenchStamped(5.0, "ft_sensor") - assert (ws.ts, ws.frame_id) == (5.0, "ft_sensor") +@pytest.mark.parametrize("values", [[1, 2, 3], [], [[1, 2, 3], [4, 5, 6]]]) +def test_wrench_from_array_wrong_shape_raises(values) -> None: + with pytest.raises(ValueError, match="6 elements"): + wrench_from_array(values) -def test_stamped_defaults_stamp_to_now(): - before = time.time() - assert before <= WrenchStamped().ts <= time.time() +def test_wrench_add_and_sub() -> None: + a = wrench_from_array([1, 2, 3, 4, 5, 6]) + b = wrench_from_array([1, 1, 1, 1, 1, 1]) + added = wrench_from_array(wrench_array(a) + wrench_array(b)) + subtracted = wrench_from_array(wrench_array(a) - wrench_array(b)) + assert (added.force.x, added.force.y, added.force.z) == (2, 3, 4) + assert (subtracted.torque.x, subtracted.torque.y, subtracted.torque.z) == (3, 4, 5) -def test_stamped_zero_ts_is_a_real_timestamp(): - assert WrenchStamped(ts=0.0).ts == 0.0 +def test_stamped_nested_wrench_and_explicit_header() -> None: + source = WrenchStamped( + header=Header(stamp=time_from_seconds(5), frame_id="tool0"), + wrench=wrench_from_array([1, 2, 3, 0.1, 0.2, 0.3]), + ) + assert isinstance(source.wrench, Wrench) + assert to_nanoseconds(source.header.stamp) == 5000000000 + assert source.header.frame_id == "tool0" + np.testing.assert_array_equal(wrench_array(source.wrench), [1, 2, 3, 0.1, 0.2, 0.3]) -def test_stamped_from_array(): - ws = WrenchStamped.from_array([1.0, 2.0, 3.0, 0.1, 0.2, 0.3], frame_id="tool0", ts=5.0) - assert (ws.ts, ws.frame_id) == (5.0, "tool0") - assert ws.force.to_tuple() == (1.0, 2.0, 3.0) - assert ws.torque.to_tuple() == (0.1, 0.2, 0.3) +def test_explicit_current_time_and_zero_default() -> None: + before = time.time_ns() + source = WrenchStamped(header=header_now()) + assert before <= to_nanoseconds(source.header.stamp) <= time.time_ns() + assert to_nanoseconds(WrenchStamped().header.stamp) == 0 -def test_lcm_encode_decode(): +@pytest.mark.parametrize("stamp", [0, 5.25]) +def test_cdr_and_independent_decoding(stamp: float) -> None: source = WrenchStamped( - ts=5.25, frame_id="ft_sensor", force=(1.0, 2.0, 3.0), torque=(0.1, 0.2, 0.3) + header=Header(stamp=time_from_seconds(stamp), frame_id="ft_sensor"), + wrench=wrench_from_array([1, 2, 3, 0.1, 0.2, 0.3]), ) - dest = WrenchStamped.lcm_decode(source.lcm_encode()) - - assert isinstance(dest, WrenchStamped) - assert dest is not source - assert dest == source - assert dest.frame_id == "ft_sensor" - assert dest.ts == pytest.approx(5.25) - - -def test_lcm_decode_of_an_unstamped_message_keeps_the_zero_stamp(): - decoded = WrenchStamped.lcm_decode(WrenchStamped(ts=0.0, frame_id="f").lcm_encode()) - assert decoded.ts == 0.0 + decoded = WrenchStamped.decode(source.encode()) + assert decoded is not source + assert decoded == source + independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr( + source.encode(), WrenchStamped.msg_name + ) + assert independent.header.frame_id == "ft_sensor" + assert independent.header.stamp.sec * 1000000000 + independent.header.stamp.nanosec == int( + stamp * 1000000000 + ) + np.testing.assert_array_equal(wrench_array(independent.wrench), [1, 2, 3, 0.1, 0.2, 0.3]) -def test_pickle_encode_decode(): - source = WrenchStamped(ts=time.time(), force=(1.0, 2.0, 3.0), torque=(0.1, 0.2, 0.3)) - dest = pickle.loads(pickle.dumps(source)) - assert isinstance(dest, WrenchStamped) - assert dest is not source - assert dest == source +def test_pickle_encode_decode() -> None: + source = WrenchStamped(header=header_now(), wrench=wrench_from_array([1, 2, 3, 0.1, 0.2, 0.3])) + decoded = pickle.loads(pickle.dumps(source)) + assert decoded is not source + assert decoded == source diff --git a/dimos/msgs/geometry_msgs/tool_publish.py b/dimos/msgs/geometry_msgs/tool_publish.py index a6276c19e56..a428513e86d 100644 --- a/dimos/msgs/geometry_msgs/tool_publish.py +++ b/dimos/msgs/geometry_msgs/tool_publish.py @@ -14,16 +14,15 @@ import time +from dimos_generated.geometry_msgs.msg import Vector3 import lcm -from dimos.msgs.geometry_msgs.Vector3 import Vector3 - def test_runpublish() -> None: for i in range(10): - msg = Vector3(-5 + i, -5 + i, i) + msg = Vector3(x=-5 + i, y=-5 + i, z=i) lc = lcm.LCM() - lc.publish("thing1_vector3#geometry_msgs.Vector3", msg.encode()) + lc.publish("thing1_vector3#geometry_msgs/msg/Vector3", msg.encode()) time.sleep(0.1) print(f"Published: {msg}") @@ -31,11 +30,10 @@ def test_runpublish() -> None: def test_receive() -> None: lc = lcm.LCM() - def receive(bla, msg) -> None: - # print("receive", bla, msg) + def receive(channel: str, msg: bytes) -> None: print(Vector3.decode(msg)) - lc.subscribe("thing1_vector3#geometry_msgs.Vector3", receive) + lc.subscribe("thing1_vector3#geometry_msgs/msg/Vector3", receive) def _loop() -> None: while True: diff --git a/dimos/msgs/helpers.py b/dimos/msgs/helpers.py index 59f9bbf0589..278092a47fa 100644 --- a/dimos/msgs/helpers.py +++ b/dimos/msgs/helpers.py @@ -15,50 +15,15 @@ from __future__ import annotations from functools import lru_cache -import importlib -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, cast + +from dimos.message_codegen.registry import message_types if TYPE_CHECKING: from dimos.msgs.protocol import DimosMsg -def lcm_msg_type(msg_name: str) -> type[Any]: - """The generated dimos_lcm class of a '.' name. ImportError - when the wheel has no such message.""" - package, name = msg_name.split(".") - lcm_type: type[Any] = getattr(importlib.import_module(f"dimos_lcm.{package}.{name}"), name) - return lcm_type - - @lru_cache(maxsize=256) def resolve_msg_type(type_name: str) -> type[DimosMsg] | None: - """Resolve a message type name to its class. - - Args: - type_name: Type name in format "module.ClassName" (e.g., "geometry_msgs.Vector3") - - Returns: - The message class or None if not found. - """ - try: - module_name, class_name = type_name.rsplit(".", 1) - except ValueError: - return None - - # Try different import paths - # First try the direct submodule path (e.g., dimos.msgs.geometry_msgs.Quaternion) - # then fall back to parent package (for dimos_lcm or other packages) - import_paths = [ - f"dimos.msgs.{module_name}.{class_name}", - f"dimos.msgs.{module_name}", - f"dimos_lcm.{module_name}", - ] - - for path in import_paths: - try: - module = importlib.import_module(path) - return getattr(module, class_name) # type: ignore[no-any-return] - except (ImportError, AttributeError): - continue - - return None + """Resolve a qualified package/msg/Type through installed schema providers.""" + return cast("type[DimosMsg] | None", message_types().get(type_name)) diff --git a/dimos/msgs/image.py b/dimos/msgs/image.py new file mode 100644 index 00000000000..8b3167ef29e --- /dev/null +++ b/dimos/msgs/image.py @@ -0,0 +1,317 @@ +# 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. + +"""Array views and compression for generated ROS image messages.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from dimos_generated.sensor_msgs.msg import CompressedImage, Image +from dimos_generated.std_msgs.msg import Header +import numpy as np +from numpy.typing import NDArray +from PIL import Image as PILImage + +_FORMATS = { + "rgb8": ("u1", 3), + "bgr8": ("u1", 3), + "rgba8": ("u1", 4), + "bgra8": ("u1", 4), + "mono8": ("u1", 1), + "mono16": ("u2", 1), + "16UC1": ("u2", 1), + "32FC1": ("f4", 1), + "64FC1": ("f8", 1), + "16SC1": ("i2", 1), + "32FC3": ("f4", 3), +} + + +def image_from_array(pixels: NDArray[Any], *, encoding: str, header: Header | None = None) -> Image: + """Copy an array into a generated image with an explicit pixel encoding.""" + if encoding not in _FORMATS: + raise ValueError(f"unsupported image encoding {encoding!r}") + kind, channels = _FORMATS[encoding] + expected = np.dtype(kind) + if pixels.dtype.kind != expected.kind or pixels.dtype.itemsize != expected.itemsize: + raise ValueError(f"array dtype does not match {encoding}") + if (channels == 1 and pixels.ndim != 2) or ( + channels != 1 and (pixels.ndim != 3 or pixels.shape[2] != channels) + ): + raise ValueError(f"array shape does not match {encoding}") + contiguous = np.ascontiguousarray(pixels) + return Image( + header=header if header is not None else Header(), + width=pixels.shape[1], + height=pixels.shape[0], + encoding=encoding, + is_bigendian=int(not pixels.dtype.isnative if np.little_endian else pixels.dtype.isnative), + step=pixels.shape[1] * channels * pixels.dtype.itemsize, + data=contiguous.view(np.uint8).ravel(), + ) + + +def image_view(msg: Image) -> NDArray[Any]: + """Borrow a read-only array, retaining the message and respecting row padding/endian.""" + if msg.encoding not in _FORMATS: + raise ValueError(f"unsupported image encoding {msg.encoding!r}") + kind, channels = _FORMATS[msg.encoding] + dtype = np.dtype((">" if msg.is_bigendian else "<") + kind) + row_bytes = msg.width * channels * dtype.itemsize + if msg.step < row_bytes or len(msg.data) != msg.height * msg.step: + raise ValueError("image data length/step does not match its dimensions") + shape = (msg.height, msg.width) if channels == 1 else (msg.height, msg.width, channels) + strides = ( + (msg.step, dtype.itemsize) + if channels == 1 + else (msg.step, channels * dtype.itemsize, dtype.itemsize) + ) + return np.ndarray(shape, dtype=dtype, buffer=msg.data.view(), strides=strides) + + +def image_to_jpeg(msg: Image, quality: int = 75) -> bytes: + """Encode an 8-bit gray or color image using OpenCV's JPEG encoder.""" + import cv2 + + if isinstance(quality, bool) or not 0 <= quality <= 100: + raise ValueError("JPEG quality must be in 0..100") + pixels = image_view(msg) + conversions = { + "rgb8": cv2.COLOR_RGB2BGR, + "rgba8": cv2.COLOR_RGBA2BGR, + "bgra8": cv2.COLOR_BGRA2BGR, + } + if msg.encoding in conversions: + pixels = cv2.cvtColor(pixels, conversions[msg.encoding]) + elif msg.encoding not in ("bgr8", "mono8"): + raise ValueError(f"JPEG requires 8-bit color or gray, got {msg.encoding!r}") + ok, encoded = cv2.imencode(".jpg", pixels, [cv2.IMWRITE_JPEG_QUALITY, quality]) + if not ok: + raise ValueError("JPEG encoding failed") + return bytes(encoded) + + +def compressed_image_from_image( + message: Image, + quality: int = 75, + *, + format: str = "jpeg", + max_width: int | None = None, + effort: int | None = None, +) -> CompressedImage: + """Compress generated pixels, retaining the exact header. + + PNG preserves uint16 depth. JXL preserves uint16/float32 depth losslessly; + it requires imagecodecs and is an explicit codec extension for consumers. + """ + if not isinstance(message, Image): + raise TypeError("compression expects a generated Image") + if max_width is not None: + message, _ = image_resize_to_fit(message, max_width, max_width) + if format == "jpeg": + output_encoding = "mono8" if message.encoding == "mono8" else "bgr8" + data = image_to_jpeg(message, quality=quality) + elif format == "png": + import cv2 + + pixels = image_view(message) + if pixels.dtype.kind != "u" or pixels.dtype.itemsize not in (1, 2): + raise ValueError("PNG cannot encode floating-point depth") + if message.encoding in ("rgb8", "rgba8"): + code = cv2.COLOR_RGB2BGR if message.encoding == "rgb8" else cv2.COLOR_RGBA2BGRA + pixels = cv2.cvtColor(pixels, code) + pixels = np.ascontiguousarray(pixels, dtype=pixels.dtype.newbyteorder("=")) + ok, encoded = cv2.imencode(".png", pixels) + if not ok: + raise ValueError("PNG encoding failed") + data = bytes(encoded) + output_encoding = {"rgb8": "bgr8", "rgba8": "bgra8"}.get(message.encoding, message.encoding) + elif format == "jxl": + import imagecodecs + + pixels = image_view(message) + if pixels.dtype.kind not in ("u", "f") or ( + pixels.dtype.kind, + pixels.dtype.itemsize, + ) not in (("u", 1), ("u", 2), ("f", 4)): + raise ValueError(f"JXL cannot encode dtype {pixels.dtype}") + if message.encoding == "bgr8": + pixels = pixels[..., ::-1] + elif message.encoding == "bgra8": + pixels = pixels[..., [2, 1, 0, 3]] + pixels = np.ascontiguousarray(pixels, dtype=pixels.dtype.newbyteorder("=")) + if pixels.dtype == np.uint8: + data = bytes(imagecodecs.jpegxl_encode(pixels, level=quality, effort=effort or 3)) + else: + data = bytes(imagecodecs.jpegxl_encode(pixels, lossless=True, effort=effort or 1)) + output_encoding = {"bgr8": "rgb8", "bgra8": "rgba8"}.get(message.encoding, message.encoding) + else: + raise ValueError(f"unsupported compression format {format!r}") + return CompressedImage( + header=message.header, + format=f"{message.encoding}; {format} compressed {output_encoding}", + data=data, + ) + + +def image_from_compressed(message: CompressedImage) -> Image: + """Decode JPEG/PNG or the JXL extension into explicitly encoded generated pixels.""" + import cv2 + + format_name = message.format.lower() + if "jxl" in format_name: + import imagecodecs + + jxl_pixels = imagecodecs.jpegxl_decode(bytes(message.data)) + if jxl_pixels.ndim == 2: + encoding = { + np.dtype("float32"): "32FC1", + np.dtype("uint16"): "mono16", + np.dtype("uint8"): "mono8", + }.get(jxl_pixels.dtype) + elif jxl_pixels.ndim == 3 and jxl_pixels.shape[2] in (3, 4): + encoding = "rgb8" if jxl_pixels.shape[2] == 3 else "rgba8" + else: + encoding = None + if encoding is None: + raise ValueError("unsupported decoded JXL layout") + if encoding == "mono16" and format_name.startswith("16uc1;"): + encoding = "16UC1" + return image_from_array(jxl_pixels, encoding=encoding, header=message.header) + if not any(name in format_name for name in ("jpeg", "jpg", "png")): + raise ValueError(f"unsupported compressed image format {message.format!r}") + pixels = cv2.imdecode(np.frombuffer(bytes(message.data), dtype=np.uint8), cv2.IMREAD_UNCHANGED) + if pixels is None: + raise ValueError("invalid compressed image data") + if pixels.ndim == 2: + encoding = "mono16" if pixels.dtype.itemsize == 2 else "mono8" + elif pixels.ndim == 3 and pixels.shape[2] == 3: + encoding = "bgr8" + elif pixels.ndim == 3 and pixels.shape[2] == 4: + encoding = "bgra8" + else: + raise ValueError("unsupported decoded image layout") + if encoding == "mono16" and format_name.startswith("16uc1;"): + encoding = "16UC1" + return image_from_array(pixels, encoding=encoding, header=message.header) + + +def image_brightness(message: Image) -> float: + """Sample mean pixel intensity in [0, 1] for 8/16-bit image encodings. + + Respect row padding and endian. Floating-point depth has no normalized + intensity scale and is rejected rather than treated as a color image. + """ + pixels = image_view(message) + if pixels.dtype.kind != "u": + raise ValueError("brightness requires unsigned integer pixels") + stride = max(1, max(message.height, message.width) // 256) + return float(pixels[::stride, ::stride].mean() / np.iinfo(pixels.dtype).max) + + +def image_sharpness(message: Image) -> float: + """Laplacian variance of an 8-bit visual image, downsampled to 160 pixels wide.""" + import cv2 + + pixels = image_view(message) + codes = { + "rgb8": cv2.COLOR_RGB2GRAY, + "bgr8": cv2.COLOR_BGR2GRAY, + "rgba8": cv2.COLOR_RGBA2GRAY, + "bgra8": cv2.COLOR_BGRA2GRAY, + } + if message.encoding == "mono8": + gray = pixels + elif message.encoding in codes: + gray = cv2.cvtColor(pixels, codes[message.encoding]) + else: + raise ValueError("Sharpness requires an 8-bit visual image encoding") + if gray.size == 0: + raise ValueError("Sharpness requires a nonempty image") + if message.width > 160: + gray = cv2.resize( + gray, + (160, max(1, round(message.height * 160 / message.width))), + interpolation=cv2.INTER_AREA, + ) + return float(cv2.Laplacian(gray, cv2.CV_64F).var()) + + +def image_to_bgr(message: Image) -> NDArray[np.uint8]: + """Return an independent BGR8 array for drawing or OpenCV color operations.""" + import cv2 + + pixels = image_view(message) + if message.encoding == "bgr8": + return pixels.copy() + codes = { + "rgb8": cv2.COLOR_RGB2BGR, + "rgba8": cv2.COLOR_RGBA2BGR, + "bgra8": cv2.COLOR_BGRA2BGR, + "mono8": cv2.COLOR_GRAY2BGR, + } + if message.encoding not in codes: + raise ValueError(f"Cannot convert {message.encoding!r} to BGR8") + return np.asarray(cv2.cvtColor(pixels, codes[message.encoding]), dtype=np.uint8) + + +def image_to_rgb(message: Image) -> NDArray[np.uint8]: + """Copy a display image as RGB8; 16-bit grayscale uses its high byte. + + Alpha is discarded. Floating-point depth needs an explicit visualization + scale and is rejected rather than assigning pixel values to meters. + """ + import cv2 + + pixels = image_view(message) + if message.encoding == "rgb8": + return pixels.copy() + if message.encoding in ("mono16", "16UC1"): + return np.asarray( + cv2.cvtColor((pixels / 256).astype(np.uint8), cv2.COLOR_GRAY2RGB), dtype=np.uint8 + ) + codes = { + "bgr8": cv2.COLOR_BGR2RGB, + "rgba8": cv2.COLOR_RGBA2RGB, + "bgra8": cv2.COLOR_BGRA2RGB, + "mono8": cv2.COLOR_GRAY2RGB, + } + if message.encoding not in codes: + raise ValueError(f"Cannot convert {message.encoding!r} to RGB8") + return np.asarray(cv2.cvtColor(pixels, codes[message.encoding]), dtype=np.uint8) + + +def image_resize_to_fit(message: Image, max_width: int, max_height: int) -> tuple[Image, float]: + """Downscale a generated image while preserving encoding, aspect and exact header.""" + import cv2 + + if min(max_width, max_height, message.width, message.height) <= 0: + raise ValueError("Image and target dimensions must be positive") + if message.width <= max_width and message.height <= max_height: + return message, 1.0 + scale = min(max_width / message.width, max_height / message.height) + pixels = cv2.resize( + image_view(message), + (max(1, int(message.width * scale)), max(1, int(message.height * scale))), + interpolation=cv2.INTER_LINEAR, + ) + return image_from_array(pixels, encoding=message.encoding, header=message.header), scale + + +def image_from_file(path: str | Path, *, header: Header | None = None) -> Image: + """Decode an image file as generated RGB8 with an explicit optional source header.""" + with PILImage.open(path) as image: + return image_from_array(np.asarray(image.convert("RGB")), encoding="rgb8", header=header) diff --git a/dimos/msgs/manipulation_msgs/test_grasp_candidate.py b/dimos/msgs/manipulation_msgs/test_grasp_candidate.py index 9538beda52e..eb676b1441e 100644 --- a/dimos/msgs/manipulation_msgs/test_grasp_candidate.py +++ b/dimos/msgs/manipulation_msgs/test_grasp_candidate.py @@ -12,16 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pytest +from dimos_generated.dimos_msgs.msg import GraspCandidate, GraspCandidateArray +from dimos_generated.geometry_msgs.msg import Point, Pose +from dimos_generated.std_msgs.msg import Header -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.manipulation_msgs.GraspCandidate import GraspCandidate -from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray -from dimos.msgs.std_msgs.Header import Header +from dimos.msgs.time import time_from_seconds def test_grasp_candidate_round_trip_preserves_pose_and_score() -> None: - candidate = GraspCandidate(Pose(0.4, -0.2, 0.3), 0.75) + candidate = GraspCandidate(pose=Pose(position=Point(x=0.4, y=-0.2, z=0.3)), score=0.75) decoded = GraspCandidate.decode(candidate.encode()) @@ -31,21 +30,18 @@ def test_grasp_candidate_round_trip_preserves_pose_and_score() -> None: assert decoded.score == 0.75 -def test_grasp_candidate_rejects_non_finite_score() -> None: - with pytest.raises(ValueError, match="score must be finite"): - GraspCandidate(score=float("nan")) - - def test_grasp_candidate_array_round_trip_preserves_header_and_order() -> None: candidates = [ - GraspCandidate(Pose(0.1, 0.0, 0.2), 0.9), - GraspCandidate(Pose(0.2, 0.0, 0.2), 0.7), + GraspCandidate(pose=Pose(position=Point(x=0.1, z=0.2)), score=0.9), + GraspCandidate(pose=Pose(position=Point(x=0.2, z=0.2)), score=0.7), ] - proposals = GraspCandidateArray(Header(123.0, "world"), candidates) + proposals = GraspCandidateArray( + header=Header(stamp=time_from_seconds(123.0), frame_id="world"), candidates=candidates + ) decoded = GraspCandidateArray.decode(proposals.encode()) - assert decoded.header.timestamp == 123.0 + assert decoded.header.stamp == time_from_seconds(123.0) assert decoded.header.frame_id == "world" - assert [candidate.score for candidate in decoded] == [0.9, 0.7] - assert [candidate.pose.position.x for candidate in decoded] == [0.1, 0.2] + assert [candidate.score for candidate in decoded.candidates] == [0.9, 0.7] + assert [candidate.pose.position.x for candidate in decoded.candidates] == [0.1, 0.2] diff --git a/dimos/msgs/nav_msgs/test_OccupancyGrid.py b/dimos/msgs/nav_msgs/test_OccupancyGrid.py index 98297f6f395..050c6420900 100644 --- a/dimos/msgs/nav_msgs/test_OccupancyGrid.py +++ b/dimos/msgs/nav_msgs/test_OccupancyGrid.py @@ -13,379 +13,276 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Test the OccupancyGrid convenience class.""" - -import pickle +"""Generated occupancy grids with external array and coordinate operations.""" +from dimos_generated.geometry_msgs.msg import Point, Pose +from dimos_generated.nav_msgs.msg import OccupancyGrid +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest -from dimos.mapping.occupancy.gradient import gradient from dimos.mapping.occupancy.inflation import simple_inflate from dimos.mapping.pointclouds.occupancy import general_occupancy -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid, block_max_reduce -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.utils.data import get_data +from dimos.msgs.geometry import quaternion_from_euler, yaw +from dimos.msgs.occupancy import ( + block_max_reduce, + grid_to_world, + occupancy_extent, + occupancy_from_array, + occupancy_view, + world_to_grid, +) +from dimos.msgs.pointcloud import pointcloud_from_xyz def test_empty_grid() -> None: """Test creating an empty grid.""" grid = OccupancyGrid() - assert grid.width == 0 - assert grid.height == 0 - assert grid.grid.shape == (0,) - assert grid.total_cells == 0 - assert grid.frame_id == "world" + assert grid.info.width == 0 + assert grid.info.height == 0 + assert occupancy_view(grid).shape == (0, 0) + assert len(grid.data) == 0 + assert grid.header.frame_id == "" def test_grid_with_dimensions() -> None: """Test creating a grid with specified dimensions.""" - grid = OccupancyGrid(width=10, height=10, resolution=0.1, frame_id="map") - assert grid.width == 10 - assert grid.height == 10 - assert grid.resolution == 0.1 - assert grid.frame_id == "map" - assert grid.grid.shape == (10, 10) - assert np.all(grid.grid == -1) # All unknown - assert grid.unknown_cells == 100 - assert grid.unknown_percent == 100.0 + grid = occupancy_from_array( + np.full((10, 10), -1, dtype=np.int8), resolution=0.1, header=Header(frame_id="map") + ) + assert grid.info.width == 10 + assert grid.info.height == 10 + assert grid.info.resolution == pytest.approx(0.1) + assert grid.header.frame_id == "map" + assert occupancy_view(grid).shape == (10, 10) + assert np.all(occupancy_view(grid) == -1) + assert np.count_nonzero(occupancy_view(grid) == -1) == 100 + assert (np.count_nonzero(occupancy_view(grid) == -1) / len(grid.data) * 100) == 100.0 def test_grid_from_numpy_array() -> None: """Test creating a grid from a numpy array.""" data = np.zeros((20, 30), dtype=np.int8) - data[5:10, 10:20] = 100 # Add some obstacles - data[15:18, 5:8] = -1 # Add unknown area - - origin = Pose(1.0, 2.0, 0.0) - grid = OccupancyGrid(grid=data, resolution=0.05, origin=origin, frame_id="odom") - - assert grid.width == 30 - assert grid.height == 20 - assert grid.resolution == 0.05 - assert grid.frame_id == "odom" - assert grid.origin.position.x == 1.0 - assert grid.origin.position.y == 2.0 - assert grid.grid.shape == (20, 30) - - # Check cell counts - assert grid.occupied_cells == 50 # 5x10 obstacle area - assert grid.free_cells == 541 # Total - occupied - unknown - assert grid.unknown_cells == 9 # 3x3 unknown area - - # Check percentages (approximately) - assert abs(grid.occupied_percent - 8.33) < 0.1 - assert abs(grid.free_percent - 90.17) < 0.1 - assert abs(grid.unknown_percent - 1.5) < 0.1 + data[5:10, 10:20] = 100 + data[15:18, 5:8] = -1 + origin = Pose(position=Point(x=1, y=2)) + grid = occupancy_from_array( + data, resolution=0.05, origin=origin, header=Header(frame_id="odom") + ) + assert grid.info.width == 30 + assert grid.info.height == 20 + assert grid.info.resolution == pytest.approx(0.05) + assert grid.header.frame_id == "odom" + assert grid.info.origin.position.x == 1.0 + assert grid.info.origin.position.y == 2.0 + assert occupancy_view(grid).shape == (20, 30) + assert np.count_nonzero(occupancy_view(grid) == 100) == 50 + assert np.count_nonzero(occupancy_view(grid) == 0) == 541 + assert np.count_nonzero(occupancy_view(grid) == -1) == 9 + assert abs((np.count_nonzero(occupancy_view(grid) == 100) / len(grid.data) * 100) - 8.33) < 0.1 + assert abs((np.count_nonzero(occupancy_view(grid) == 0) / len(grid.data) * 100) - 90.17) < 0.1 + assert abs((np.count_nonzero(occupancy_view(grid) == -1) / len(grid.data) * 100) - 1.5) < 0.1 def test_world_grid_coordinate_conversion() -> None: """Test converting between world and grid coordinates.""" data = np.zeros((20, 30), dtype=np.int8) - origin = Pose(1.0, 2.0, 0.0) - grid = OccupancyGrid(grid=data, resolution=0.05, origin=origin, frame_id="odom") - - # Test world to grid - grid_pos = grid.world_to_grid((2.5, 3.0)) - assert int(grid_pos.x) == 30 - assert int(grid_pos.y) == 20 - - # Test grid to world - world_pos = grid.grid_to_world((10, 5)) - assert world_pos.x == 1.5 - assert world_pos.y == 2.25 + origin = Pose(position=Point(x=1, y=2)) + grid = occupancy_from_array( + data, resolution=0.05, origin=origin, header=Header(frame_id="odom") + ) + grid_pos = world_to_grid(grid, Point(x=2.5, y=3.0)) + assert grid_pos[0] == pytest.approx(30) + assert grid_pos[1] == pytest.approx(20) + world_pos = grid_to_world(grid, (10, 5)) + assert world_pos.x == pytest.approx(1.5) + assert world_pos.y == pytest.approx(2.25) -def test_lcm_encode_decode() -> None: +def test_cdr_encode_decode() -> None: """Test LCM encoding and decoding.""" data = np.zeros((20, 30), dtype=np.int8) - data[5:10, 10:20] = 100 # Add some obstacles - data[15:18, 5:8] = -1 # Add unknown area - origin = Pose(1.0, 2.0, 0.0) - grid = OccupancyGrid(grid=data, resolution=0.05, origin=origin, frame_id="odom") - - # Set a specific value for testing - # Convert world coordinates to grid indices - grid_pos = grid.world_to_grid((1.5, 2.25)) - grid.grid[int(grid_pos.y), int(grid_pos.x)] = 50 - - # Encode - lcm_data = grid.lcm_encode() + data[5:10, 10:20] = 100 + data[15:18, 5:8] = -1 + origin = Pose(position=Point(x=1, y=2)) + grid = occupancy_from_array( + data, resolution=0.05, origin=origin, header=Header(frame_id="odom") + ) + grid_pos = world_to_grid(grid, Point(x=1.5, y=2.25)) + grid.data[round(grid_pos[1]) * grid.info.width + round(grid_pos[0])] = 50 + lcm_data = grid.encode() assert isinstance(lcm_data, bytes) assert len(lcm_data) > 0 - - # Decode - decoded = OccupancyGrid.lcm_decode(lcm_data) - - # Check that data matches exactly (grid arrays should be identical) - assert np.array_equal(grid.grid, decoded.grid) - assert grid.width == decoded.width - assert grid.height == decoded.height - assert abs(grid.resolution - decoded.resolution) < 1e-6 # Use approximate equality for floats - assert abs(grid.origin.position.x - decoded.origin.position.x) < 1e-6 - assert abs(grid.origin.position.y - decoded.origin.position.y) < 1e-6 - assert grid.frame_id == decoded.frame_id - - # Check that the actual grid data was preserved (don't rely on float conversions) - assert decoded.grid[5, 10] == 50 # Value we set should be preserved in grid - - -def test_lcm_decode_origin_is_dimos_pose() -> None: - """The decoded origin must be the dimos Pose (with .yaw), not the raw - generated pose from the wire (the relay costmap encoder reads origin.yaw).""" - quat = Quaternion.from_euler(Vector3(0.0, 0.0, 0.5)) - origin = Pose(1.0, 2.0, 0.0, quat.x, quat.y, quat.z, quat.w) - grid = OccupancyGrid(grid=np.zeros((2, 3), dtype=np.int8), resolution=0.05, origin=origin) - - decoded = OccupancyGrid.lcm_decode(grid.lcm_encode()) - - assert isinstance(decoded.origin, Pose) - assert decoded.origin.yaw == pytest.approx(0.5) + decoded = OccupancyGrid.decode(lcm_data) + assert np.array_equal(occupancy_view(grid), occupancy_view(decoded)) + assert grid.info.width == decoded.info.width + assert grid.info.height == decoded.info.height + assert abs(grid.info.resolution - decoded.info.resolution) < 1e-06 + assert abs(grid.info.origin.position.x - decoded.info.origin.position.x) < 1e-06 + assert abs(grid.info.origin.position.y - decoded.info.origin.position.y) < 1e-06 + assert grid.header.frame_id == decoded.header.frame_id + assert occupancy_view(decoded)[5, 10] == 50 + + +def test_cdr_decode_origin_is_dimos_pose() -> None: + """Generated origin orientation works with the external yaw helper.""" + quat = quaternion_from_euler(0, 0, 0.5) + origin = Pose(position=Point(x=1, y=2), orientation=quat) + grid = occupancy_from_array( + np.zeros((2, 3), dtype=np.int8), resolution=0.05, origin=origin, header=Header(frame_id="") + ) + decoded = OccupancyGrid.decode(grid.encode()) + assert isinstance(decoded.info.origin, Pose) + assert yaw(decoded.info.origin.orientation) == pytest.approx(0.5) -def test_lcm_decode_empty_grid() -> None: +def test_cdr_decode_empty_grid() -> None: """An empty grid (mapper warming up) must survive the wire; the decoder used to rebuild it as a 1-D array the constructor rejects.""" - decoded = OccupancyGrid.lcm_decode(OccupancyGrid().lcm_encode()) - assert decoded.grid.size == 0 - assert isinstance(decoded.origin, Pose) + decoded = OccupancyGrid.decode(OccupancyGrid().encode()) + assert occupancy_view(decoded).size == 0 + assert isinstance(decoded.info.origin, Pose) -def test_string_representation() -> None: - """Test string representations.""" - grid = OccupancyGrid(width=10, height=10, resolution=0.1, frame_id="map") - - # Test __str__ - str_repr = str(grid) - assert "OccupancyGrid[map]" in str_repr - assert "10x10" in str_repr - assert "1.0x1.0m" in str_repr - assert "10cm res" in str_repr - - # Test __repr__ - repr_str = repr(grid) - assert "OccupancyGrid(" in repr_str - assert "width=10" in repr_str - assert "height=10" in repr_str - assert "resolution=0.1" in repr_str +def test_physical_extent() -> None: + grid = occupancy_from_array( + np.full((10, 10), -1, dtype=np.int8), header=Header(frame_id="map"), resolution=0.1 + ) + assert occupancy_extent(grid) == pytest.approx((1, 1)) def test_grid_property_sync() -> None: """Test that the grid property works correctly.""" - grid = OccupancyGrid(width=5, height=5, resolution=0.1, frame_id="map") - - # Modify via numpy array - grid.grid[2, 3] = 100 - assert grid.grid[2, 3] == 100 - - # Check that we can access grid values - grid.grid[0, 0] = 50 - assert grid.grid[0, 0] == 50 + grid = occupancy_from_array( + np.full((5, 5), -1, dtype=np.int8), resolution=0.1, header=Header(frame_id="map") + ) + grid.data[2 * grid.info.width + 3] = 100 + assert occupancy_view(grid)[2, 3] == 100 + grid.data[0] = 50 + assert occupancy_view(grid)[0, 0] == 50 def test_invalid_grid_dimensions() -> None: """Test handling of invalid grid dimensions.""" - # Test with non-2D array - with pytest.raises(ValueError, match="Grid must be a 2D array"): - OccupancyGrid(grid=np.zeros(10), resolution=0.1) + with pytest.raises(ValueError, match="2D integer array"): + occupancy_from_array(np.zeros(10), resolution=0.1, header=Header(frame_id="")) def test_from_pointcloud() -> None: """Test creating OccupancyGrid from PointCloud2.""" - file_path = get_data("lcm_msgs") / "sensor_msgs/PointCloud2.pickle" - with open(file_path, "rb") as f: - lcm_msg = pickle.loads(f.read()) - - pointcloud = PointCloud2.lcm_decode(lcm_msg) - - # Convert pointcloud to occupancy grid + x, y = np.meshgrid(np.arange(10) * 0.05, np.arange(10) * 0.05) + points = np.column_stack((x.ravel(), y.ravel(), np.full(x.size, 0.5))) + pointcloud = pointcloud_from_xyz(points, header=Header(frame_id="world")) occupancygrid = general_occupancy(pointcloud, resolution=0.05, min_height=0.1, max_height=2.0) - # Apply inflation separately if needed occupancygrid = simple_inflate(occupancygrid, 0.1) - - # Check that grid was created with reasonable properties - assert occupancygrid.width > 0 - assert occupancygrid.height > 0 - assert occupancygrid.resolution == 0.05 - assert occupancygrid.frame_id == pointcloud.frame_id - assert occupancygrid.occupied_cells > 0 # Should have some occupied cells - - -def test_gradient() -> None: - """Test converting occupancy grid to gradient field.""" - # Create a small test grid with an obstacle in the middle - data = np.zeros((10, 10), dtype=np.int8) - data[4:6, 4:6] = 100 # 2x2 obstacle in center - - grid = OccupancyGrid(grid=data, resolution=0.1) # 0.1m per cell - - # Convert to gradient - gradient_grid = gradient(grid, obstacle_threshold=50, max_distance=1.0) - - # Check that we get an OccupancyGrid back - assert isinstance(gradient_grid, OccupancyGrid) - assert gradient_grid.grid.shape == (10, 10) - assert gradient_grid.resolution == grid.resolution - assert gradient_grid.frame_id == grid.frame_id - - # Obstacle cells should have value 100 - assert gradient_grid.grid[4, 4] == 100 - assert gradient_grid.grid[5, 5] == 100 - - # Adjacent cells should have high values (near obstacles) - assert gradient_grid.grid[3, 4] > 85 # Very close to obstacle - assert gradient_grid.grid[4, 3] > 85 # Very close to obstacle - - # Cells at moderate distance should have moderate values - assert 30 < gradient_grid.grid[0, 0] < 60 # Corner is ~0.57m away - - # Check that gradient decreases with distance - assert gradient_grid.grid[3, 4] > gradient_grid.grid[2, 4] # Closer is higher - assert gradient_grid.grid[2, 4] > gradient_grid.grid[0, 4] # Further is lower - - # Test with unknown cells - data_with_unknown = data.copy() - data_with_unknown[0:2, 0:2] = -1 # Add unknown area (close to obstacle) - data_with_unknown[8:10, 8:10] = -1 # Add unknown area (far from obstacle) - - grid_with_unknown = OccupancyGrid(data_with_unknown, resolution=0.1) - gradient_with_unknown = gradient(grid_with_unknown, max_distance=1.0) # 1m max distance - - # Unknown cells should remain unknown (new behavior - unknowns are preserved) - assert gradient_with_unknown.grid[0, 0] == -1 # Should remain unknown - assert gradient_with_unknown.grid[1, 1] == -1 # Should remain unknown - assert gradient_with_unknown.grid[8, 8] == -1 # Should remain unknown - assert gradient_with_unknown.grid[9, 9] == -1 # Should remain unknown - - # Unknown cells count should be preserved - assert gradient_with_unknown.unknown_cells == 8 # All unknowns preserved + assert occupancygrid.info.width > 0 + assert occupancygrid.info.height > 0 + assert occupancygrid.info.resolution == pytest.approx(0.05) + assert occupancygrid.header.frame_id == pointcloud.header.frame_id + assert np.count_nonzero(occupancy_view(occupancygrid) == 100) > 0 def test_filter_above() -> None: """Test filtering cells above threshold.""" - # Create test grid with various values data = np.array( [[-1, 0, 20, 50], [10, 30, 60, 80], [40, 70, 90, 100], [-1, 15, 25, -1]], dtype=np.int8 ) - - grid = OccupancyGrid(grid=data, resolution=0.1) - - # Filter to keep only values > 50 - filtered = grid.filter_above(50) - - # Check that values > 50 are preserved - assert filtered.grid[1, 2] == 60 - assert filtered.grid[1, 3] == 80 - assert filtered.grid[2, 1] == 70 - assert filtered.grid[2, 2] == 90 - assert filtered.grid[2, 3] == 100 - - # Check that values <= 50 are set to -1 (unknown) - assert filtered.grid[0, 1] == -1 # was 0 - assert filtered.grid[0, 2] == -1 # was 20 - assert filtered.grid[0, 3] == -1 # was 50 - assert filtered.grid[1, 0] == -1 # was 10 - assert filtered.grid[1, 1] == -1 # was 30 - assert filtered.grid[2, 0] == -1 # was 40 - - # Check that unknown cells are preserved - assert filtered.grid[0, 0] == -1 - assert filtered.grid[3, 0] == -1 - assert filtered.grid[3, 3] == -1 - - # Check dimensions and metadata preserved - assert filtered.width == grid.width - assert filtered.height == grid.height - assert filtered.resolution == grid.resolution - assert filtered.frame_id == grid.frame_id + grid = occupancy_from_array(data, resolution=0.1, header=Header(frame_id="")) + filtered = occupancy_from_array( + np.where(occupancy_view(grid) > 50, occupancy_view(grid), -1).astype(np.int8), + header=grid.header, + resolution=grid.info.resolution, + origin=grid.info.origin, + ) + assert occupancy_view(filtered)[1, 2] == 60 + assert occupancy_view(filtered)[1, 3] == 80 + assert occupancy_view(filtered)[2, 1] == 70 + assert occupancy_view(filtered)[2, 2] == 90 + assert occupancy_view(filtered)[2, 3] == 100 + assert occupancy_view(filtered)[0, 1] == -1 + assert occupancy_view(filtered)[0, 2] == -1 + assert occupancy_view(filtered)[0, 3] == -1 + assert occupancy_view(filtered)[1, 0] == -1 + assert occupancy_view(filtered)[1, 1] == -1 + assert occupancy_view(filtered)[2, 0] == -1 + assert occupancy_view(filtered)[0, 0] == -1 + assert occupancy_view(filtered)[3, 0] == -1 + assert occupancy_view(filtered)[3, 3] == -1 + assert filtered.info.width == grid.info.width + assert filtered.info.height == grid.info.height + assert filtered.info.resolution == grid.info.resolution + assert filtered.header.frame_id == grid.header.frame_id def test_filter_below() -> None: """Test filtering cells below threshold.""" - # Create test grid with various values data = np.array( [[-1, 0, 20, 50], [10, 30, 60, 80], [40, 70, 90, 100], [-1, 15, 25, -1]], dtype=np.int8 ) - - grid = OccupancyGrid(grid=data, resolution=0.1) - - # Filter to keep only values < 50 - filtered = grid.filter_below(50) - - # Check that values < 50 are preserved - assert filtered.grid[0, 1] == 0 - assert filtered.grid[0, 2] == 20 - assert filtered.grid[1, 0] == 10 - assert filtered.grid[1, 1] == 30 - assert filtered.grid[2, 0] == 40 - assert filtered.grid[3, 1] == 15 - assert filtered.grid[3, 2] == 25 - - # Check that values >= 50 are set to -1 (unknown) - assert filtered.grid[0, 3] == -1 # was 50 - assert filtered.grid[1, 2] == -1 # was 60 - assert filtered.grid[1, 3] == -1 # was 80 - assert filtered.grid[2, 1] == -1 # was 70 - assert filtered.grid[2, 2] == -1 # was 90 - assert filtered.grid[2, 3] == -1 # was 100 - - # Check that unknown cells are preserved - assert filtered.grid[0, 0] == -1 - assert filtered.grid[3, 0] == -1 - assert filtered.grid[3, 3] == -1 - - # Check dimensions and metadata preserved - assert filtered.width == grid.width - assert filtered.height == grid.height - assert filtered.resolution == grid.resolution - assert filtered.frame_id == grid.frame_id + grid = occupancy_from_array(data, resolution=0.1, header=Header(frame_id="")) + filtered = occupancy_from_array( + np.where(occupancy_view(grid) < 50, occupancy_view(grid), -1).astype(np.int8), + header=grid.header, + resolution=grid.info.resolution, + origin=grid.info.origin, + ) + assert occupancy_view(filtered)[0, 1] == 0 + assert occupancy_view(filtered)[0, 2] == 20 + assert occupancy_view(filtered)[1, 0] == 10 + assert occupancy_view(filtered)[1, 1] == 30 + assert occupancy_view(filtered)[2, 0] == 40 + assert occupancy_view(filtered)[3, 1] == 15 + assert occupancy_view(filtered)[3, 2] == 25 + assert occupancy_view(filtered)[0, 3] == -1 + assert occupancy_view(filtered)[1, 2] == -1 + assert occupancy_view(filtered)[1, 3] == -1 + assert occupancy_view(filtered)[2, 1] == -1 + assert occupancy_view(filtered)[2, 2] == -1 + assert occupancy_view(filtered)[2, 3] == -1 + assert occupancy_view(filtered)[0, 0] == -1 + assert occupancy_view(filtered)[3, 0] == -1 + assert occupancy_view(filtered)[3, 3] == -1 + assert filtered.info.width == grid.info.width + assert filtered.info.height == grid.info.height + assert filtered.info.resolution == grid.info.resolution + assert filtered.header.frame_id == grid.header.frame_id def test_max() -> None: """Test setting all non-unknown cells to maximum.""" - # Create test grid with various values data = np.array( [[-1, 0, 20, 50], [10, 30, 60, 80], [40, 70, 90, 100], [-1, 15, 25, -1]], dtype=np.int8 ) - - grid = OccupancyGrid(grid=data, resolution=0.1) - - # Apply max - maxed = grid.max() - - # Check that all non-unknown cells are set to 100 - assert maxed.grid[0, 1] == 100 # was 0 - assert maxed.grid[0, 2] == 100 # was 20 - assert maxed.grid[0, 3] == 100 # was 50 - assert maxed.grid[1, 0] == 100 # was 10 - assert maxed.grid[1, 1] == 100 # was 30 - assert maxed.grid[1, 2] == 100 # was 60 - assert maxed.grid[1, 3] == 100 # was 80 - assert maxed.grid[2, 0] == 100 # was 40 - assert maxed.grid[2, 1] == 100 # was 70 - assert maxed.grid[2, 2] == 100 # was 90 - assert maxed.grid[2, 3] == 100 # was 100 (already max) - assert maxed.grid[3, 1] == 100 # was 15 - assert maxed.grid[3, 2] == 100 # was 25 - - # Check that unknown cells are preserved - assert maxed.grid[0, 0] == -1 - assert maxed.grid[3, 0] == -1 - assert maxed.grid[3, 3] == -1 - - # Check dimensions and metadata preserved - assert maxed.width == grid.width - assert maxed.height == grid.height - assert maxed.resolution == grid.resolution - assert maxed.frame_id == grid.frame_id - - # Verify statistics - assert maxed.unknown_cells == 3 # Same as original - assert maxed.occupied_cells == 13 # All non-unknown cells - assert maxed.free_cells == 0 # No free cells + grid = occupancy_from_array(data, resolution=0.1, header=Header(frame_id="")) + maxed = occupancy_from_array( + np.where(occupancy_view(grid) >= 0, 100, -1).astype(np.int8), + header=grid.header, + resolution=grid.info.resolution, + origin=grid.info.origin, + ) + assert occupancy_view(maxed)[0, 1] == 100 + assert occupancy_view(maxed)[0, 2] == 100 + assert occupancy_view(maxed)[0, 3] == 100 + assert occupancy_view(maxed)[1, 0] == 100 + assert occupancy_view(maxed)[1, 1] == 100 + assert occupancy_view(maxed)[1, 2] == 100 + assert occupancy_view(maxed)[1, 3] == 100 + assert occupancy_view(maxed)[2, 0] == 100 + assert occupancy_view(maxed)[2, 1] == 100 + assert occupancy_view(maxed)[2, 2] == 100 + assert occupancy_view(maxed)[2, 3] == 100 + assert occupancy_view(maxed)[3, 1] == 100 + assert occupancy_view(maxed)[3, 2] == 100 + assert occupancy_view(maxed)[0, 0] == -1 + assert occupancy_view(maxed)[3, 0] == -1 + assert occupancy_view(maxed)[3, 3] == -1 + assert maxed.info.width == grid.info.width + assert maxed.info.height == grid.info.height + assert maxed.info.resolution == grid.info.resolution + assert maxed.header.frame_id == grid.header.frame_id + assert np.count_nonzero(occupancy_view(maxed) == -1) == 3 + assert np.count_nonzero(occupancy_view(maxed) == 100) == 13 + assert np.count_nonzero(occupancy_view(maxed) == 0) == 0 def test_block_max_reduce_preserves_lone_obstacle() -> None: @@ -408,10 +305,9 @@ def test_block_max_reduce_unknown_only_when_whole_block_unknown() -> None: def test_block_max_reduce_trims_remainder() -> None: cells = np.arange(35, dtype=np.int8).reshape(7, 5) reduced = block_max_reduce(cells, 2) - # 7x5 at factor 2 keeps rows 0-5 and cols 0-3 (origin corner side). assert reduced.shape == (3, 2) - assert reduced[0, 0] == 6 # max of rows 0-1, cols 0-1 - assert reduced[2, 1] == 28 # max of rows 4-5, cols 2-3 + assert reduced[0, 0] == 6 + assert reduced[2, 1] == 28 def test_block_max_reduce_thin_grid_passes_through() -> None: diff --git a/dimos/msgs/nav_msgs/test_Odometry.py b/dimos/msgs/nav_msgs/test_Odometry.py index 9a0e77c5587..73af9d94e05 100644 --- a/dimos/msgs/nav_msgs/test_Odometry.py +++ b/dimos/msgs/nav_msgs/test_Odometry.py @@ -14,50 +14,60 @@ import time +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseWithCovariance, + Quaternion, + Twist, + TwistWithCovariance, + Vector3, +) +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.std_msgs.msg import Header import numpy as np +from rosbags.typesys import Stores, get_typestore -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseWithCovariance import PoseWithCovariance -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.TwistWithCovariance import TwistWithCovariance -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.msgs.geometry import quaternion_euler +from dimos.msgs.time import header_now, time_from_seconds, to_seconds def test_odometry_default_init() -> None: odom = Odometry() - - assert odom.ts > 0 - assert odom.frame_id == "" + assert to_seconds(odom.header.stamp) == 0 + assert odom.header.frame_id == "" assert odom.child_frame_id == "" - - assert odom.pose.position.x == 0.0 - assert odom.pose.position.y == 0.0 - assert odom.pose.position.z == 0.0 - assert odom.pose.orientation.w == 1.0 - - assert odom.twist.linear.x == 0.0 - assert odom.twist.angular.x == 0.0 - - assert np.all(odom.pose.covariance == 0.0) - assert np.all(odom.twist.covariance == 0.0) + assert odom.pose.pose.position.x == 0.0 + assert odom.pose.pose.position.y == 0.0 + assert odom.pose.pose.position.z == 0.0 + assert odom.pose.pose.orientation.w == 1.0 + assert odom.twist.twist.linear.x == 0.0 + assert odom.twist.twist.angular.x == 0.0 + assert np.all(np.asarray(odom.pose.covariance) == 0.0) + assert np.all(np.asarray(odom.twist.covariance) == 0.0) def test_odometry_with_frames() -> None: ts = 1234567890.123456 - odom = Odometry(ts=ts, frame_id="odom", child_frame_id="base_link") - - assert odom.ts == ts - assert odom.frame_id == "odom" + odom = Odometry( + child_frame_id="base_link", header=Header(stamp=time_from_seconds(ts), frame_id="odom") + ) + assert to_seconds(odom.header.stamp) == ts + assert odom.header.frame_id == "odom" assert odom.child_frame_id == "base_link" def test_odometry_with_pose_and_twist() -> None: - pose = Pose(1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.9) - twist = Twist(Vector3(0.5, 0.0, 0.0), Vector3(0.0, 0.0, 0.1)) - - odom = Odometry(ts=1000.0, frame_id="odom", child_frame_id="base_link", pose=pose, twist=twist) - + pose = Pose( + position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) + ) + twist = Twist(linear=Vector3(x=0.5, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.1)) + odom = Odometry( + child_frame_id="base_link", + pose=PoseWithCovariance(pose=pose), + twist=TwistWithCovariance(twist=twist), + header=Header(stamp=time_from_seconds(1000.0), frame_id="odom"), + ) assert odom.pose.pose.position.x == 1.0 assert odom.pose.pose.position.y == 2.0 assert odom.pose.pose.position.z == 3.0 @@ -66,145 +76,149 @@ def test_odometry_with_pose_and_twist() -> None: def test_odometry_with_covariances() -> None: - pose = Pose(1.0, 2.0, 3.0) + pose = Pose(position=Point(x=1.0, y=2.0, z=3.0)) pose_cov = np.arange(36, dtype=float) - pose_with_cov = PoseWithCovariance(pose, pose_cov) - - twist = Twist(Vector3(0.5, 0.0, 0.0), Vector3(0.0, 0.0, 0.1)) + pose_with_cov = PoseWithCovariance(pose=pose, covariance=pose_cov) + twist = Twist(linear=Vector3(x=0.5, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.1)) twist_cov = np.arange(36, 72, dtype=float) - twist_with_cov = TwistWithCovariance(twist, twist_cov) - + twist_with_cov = TwistWithCovariance(twist=twist, covariance=twist_cov) odom = Odometry( - ts=1000.0, - frame_id="odom", child_frame_id="base_link", pose=pose_with_cov, twist=twist_with_cov, + header=Header(stamp=time_from_seconds(1000.0), frame_id="odom"), ) - - assert odom.pose.position.x == 1.0 + assert odom.pose.pose.position.x == 1.0 assert np.array_equal(odom.pose.covariance, pose_cov) - assert odom.twist.linear.x == 0.5 + assert odom.twist.twist.linear.x == 0.5 assert np.array_equal(odom.twist.covariance, twist_cov) def test_odometry_properties() -> None: - pose = Pose(1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.9) - twist = Twist(Vector3(0.5, 0.6, 0.7), Vector3(0.1, 0.2, 0.3)) - - odom = Odometry(ts=1000.0, frame_id="odom", child_frame_id="base_link", pose=pose, twist=twist) - - assert odom.x == 1.0 - assert odom.y == 2.0 - assert odom.z == 3.0 - assert odom.position.x == 1.0 - - assert odom.orientation.x == 0.1 - - assert odom.vx == 0.5 - assert odom.vy == 0.6 - assert odom.vz == 0.7 - assert odom.linear_velocity.x == 0.5 - - assert odom.wx == 0.1 - assert odom.wy == 0.2 - assert odom.wz == 0.3 - assert odom.angular_velocity.x == 0.1 - - assert odom.roll == pose.roll - assert odom.pitch == pose.pitch - assert odom.yaw == pose.yaw - - -def test_odometry_str_repr() -> None: + pose = Pose( + position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) + ) + twist = Twist(linear=Vector3(x=0.5, y=0.6, z=0.7), angular=Vector3(x=0.1, y=0.2, z=0.3)) odom = Odometry( - ts=1234567890.123456, - frame_id="odom", child_frame_id="base_link", - pose=Pose(1.234, 2.567, 3.891), - twist=Twist(Vector3(0.5, 0.0, 0.0), Vector3(0.0, 0.0, 0.1)), + pose=PoseWithCovariance(pose=pose), + twist=TwistWithCovariance(twist=twist), + header=Header(stamp=time_from_seconds(1000.0), frame_id="odom"), ) + assert odom.pose.pose.position.x == 1.0 + assert odom.pose.pose.position.y == 2.0 + assert odom.pose.pose.position.z == 3.0 + assert odom.pose.pose.position.x == 1.0 + assert odom.pose.pose.orientation.x == 0.1 + assert odom.twist.twist.linear.x == 0.5 + assert odom.twist.twist.linear.y == 0.6 + assert odom.twist.twist.linear.z == 0.7 + assert odom.twist.twist.linear.x == 0.5 + assert odom.twist.twist.angular.x == 0.1 + assert odom.twist.twist.angular.y == 0.2 + assert odom.twist.twist.angular.z == 0.3 + assert odom.twist.twist.angular.x == 0.1 + assert quaternion_euler(odom.pose.pose.orientation)[0] == quaternion_euler(pose.orientation)[0] + assert quaternion_euler(odom.pose.pose.orientation)[1] == quaternion_euler(pose.orientation)[1] + assert quaternion_euler(odom.pose.pose.orientation)[2] == quaternion_euler(pose.orientation)[2] - assert "Odometry" in repr(odom) - assert "1234567890.123456" in repr(odom) - s = str(odom) - assert "odom -> base_link" in s - assert "1.234" in s +def test_independent_ros_decoding() -> None: + source = Odometry( + header=Header(stamp=time_from_seconds(1000), frame_id="odom"), + child_frame_id="base_link", + pose=PoseWithCovariance(pose=Pose(position=Point(x=1.234, y=2.567, z=3.891))), + twist=TwistWithCovariance(twist=Twist(linear=Vector3(x=0.5), angular=Vector3(z=0.1))), + ) + decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(source.encode(), Odometry.msg_name) + assert decoded.header.frame_id == "odom" + assert decoded.child_frame_id == "base_link" + assert ( + decoded.pose.pose.position.x, + decoded.pose.pose.position.y, + decoded.pose.pose.position.z, + ) == (1.234, 2.567, 3.891) + assert decoded.twist.twist.linear.x == 0.5 + assert decoded.twist.twist.angular.z == 0.1 def test_odometry_equality() -> None: kwargs = dict( - ts=1000.0, - frame_id="odom", child_frame_id="base_link", - pose=Pose(1.0, 2.0, 3.0), - twist=Twist(Vector3(0.5, 0.0, 0.0), Vector3(0.0, 0.0, 0.1)), + pose=PoseWithCovariance(pose=Pose(position=Point(x=1.0, y=2.0, z=3.0))), + twist=TwistWithCovariance( + twist=Twist(linear=Vector3(x=0.5, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.1)) + ), + header=Header(stamp=time_from_seconds(1000.0), frame_id="odom"), ) - assert Odometry(**kwargs) == Odometry(**kwargs) - assert Odometry(**kwargs) != Odometry(**{**kwargs, "pose": Pose(1.1, 2.0, 3.0)}) + assert Odometry(**kwargs) != Odometry( + **{**kwargs, "pose": PoseWithCovariance(pose=Pose(position=Point(x=1.1, y=2.0, z=3.0)))} + ) assert Odometry(**kwargs) != "not an odometry" -def test_odometry_lcm_roundtrip() -> None: - pose = Pose(1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.9) +def test_odometry_cdr_roundtrip() -> None: + pose = Pose( + position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) + ) pose_cov = np.arange(36, dtype=float) - twist = Twist(Vector3(0.5, 0.6, 0.7), Vector3(0.1, 0.2, 0.3)) + twist = Twist(linear=Vector3(x=0.5, y=0.6, z=0.7), angular=Vector3(x=0.1, y=0.2, z=0.3)) twist_cov = np.arange(36, 72, dtype=float) - source = Odometry( - ts=1234567890.123456, - frame_id="odom", child_frame_id="base_link", - pose=PoseWithCovariance(pose, pose_cov), - twist=TwistWithCovariance(twist, twist_cov), + pose=PoseWithCovariance(pose=pose, covariance=pose_cov), + twist=TwistWithCovariance(twist=twist, covariance=twist_cov), + header=Header(stamp=time_from_seconds(1234567890.123456), frame_id="odom"), ) - - decoded = Odometry.lcm_decode(source.lcm_encode()) - - assert abs(decoded.ts - source.ts) < 1e-6 - assert decoded.frame_id == source.frame_id + decoded = Odometry.decode(source.encode()) + assert abs(to_seconds(decoded.header.stamp) - to_seconds(source.header.stamp)) < 1e-06 + assert decoded.header.frame_id == source.header.frame_id assert decoded.child_frame_id == source.child_frame_id assert decoded.pose == source.pose assert decoded.twist == source.twist def test_odometry_zero_timestamp() -> None: - """An explicit zero timestamp is kept; an omitted one stamps now.""" - assert Odometry(ts=0.0).ts == 0.0 - + assert to_seconds(Odometry().header.stamp) == 0 before = time.time() - assert before <= Odometry().ts <= time.time() + assert before <= to_seconds(Odometry(header=header_now()).header.stamp) <= time.time() def test_odometry_with_just_pose() -> None: - odom = Odometry(pose=Pose(1.0, 2.0, 3.0)) - - assert odom.pose.position.x == 1.0 - assert np.all(odom.pose.covariance == 0.0) - assert np.all(odom.twist.covariance == 0.0) + odom = Odometry(pose=PoseWithCovariance(pose=Pose(position=Point(x=1.0, y=2.0, z=3.0)))) + assert odom.pose.pose.position.x == 1.0 + assert np.all(np.asarray(odom.pose.covariance) == 0.0) + assert np.all(np.asarray(odom.twist.covariance) == 0.0) def test_odometry_with_just_twist() -> None: - odom = Odometry(twist=Twist(Vector3(0.5, 0.0, 0.0), Vector3(0.0, 0.0, 0.1))) - - assert odom.twist.linear.x == 0.5 - assert odom.twist.angular.z == 0.1 - assert np.all(odom.twist.covariance == 0.0) + odom = Odometry( + twist=TwistWithCovariance( + twist=Twist(linear=Vector3(x=0.5, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.1)) + ) + ) + assert odom.twist.twist.linear.x == 0.5 + assert odom.twist.twist.angular.z == 0.1 + assert np.all(np.asarray(odom.twist.covariance) == 0.0) def test_odometry_typical_robot_scenario() -> None: odom = Odometry( - ts=1000.0, - frame_id="odom", child_frame_id="base_footprint", - pose=Pose(10.0, 5.0, 0.0, 0.0, 0.0, np.sin(0.1), np.cos(0.1)), - twist=Twist(Vector3(0.5, 0.0, 0.0), Vector3(0.0, 0.0, 0.05)), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=10.0, y=5.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=np.sin(0.1), w=np.cos(0.1)), + ) + ), + twist=TwistWithCovariance( + twist=Twist(linear=Vector3(x=0.5, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.05)) + ), + header=Header(stamp=time_from_seconds(1000.0), frame_id="odom"), ) - - assert odom.x == 10.0 - assert odom.y == 5.0 - assert abs(odom.yaw - 0.2) < 0.01 - assert odom.vx == 0.5 - assert odom.wz == 0.05 + assert odom.pose.pose.position.x == 10.0 + assert odom.pose.pose.position.y == 5.0 + assert abs(quaternion_euler(odom.pose.pose.orientation)[2] - 0.2) < 0.01 + assert odom.twist.twist.linear.x == 0.5 + assert odom.twist.twist.angular.z == 0.05 diff --git a/dimos/msgs/nav_msgs/test_Path.py b/dimos/msgs/nav_msgs/test_Path.py index 9bd0cc92b6b..bd68698c5fc 100644 --- a/dimos/msgs/nav_msgs/test_Path.py +++ b/dimos/msgs/nav_msgs/test_Path.py @@ -13,35 +13,37 @@ # limitations under the License. -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.nav_msgs.Path import Path +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Header +from rosbags.typesys import Stores, get_typestore + +from dimos.msgs.time import time_from_nanoseconds def create_test_pose(x: float, y: float, z: float, frame_id: str = "map") -> PoseStamped: """Helper to create a test PoseStamped.""" return PoseStamped( - frame_id=frame_id, - position=[x, y, z], - orientation=Quaternion(0, 0, 0, 1), # Identity quaternion + header=Header(frame_id=frame_id), + pose=Pose(position=Point(x=x, y=y, z=z), orientation=Quaternion(w=1)), ) def test_init_empty() -> None: """Test creating an empty path.""" - path = Path(frame_id="map") - assert path.frame_id == "map" - assert len(path) == 0 - assert not path # Should be falsy when empty + path = Path(header=Header(frame_id="map")) + assert path.header.frame_id == "map" + assert len(path.poses) == 0 + assert not list(path.poses) # Should be falsy when empty assert path.poses == [] def test_init_with_poses() -> None: """Test creating a path with initial poses.""" poses = [create_test_pose(i, i, 0) for i in range(3)] - path = Path(frame_id="map", poses=poses) - assert len(path) == 3 - assert bool(path) # Should be truthy when has poses + path = Path(header=Header(frame_id="map"), poses=poses) + assert len(path.poses) == 3 + assert bool(list(path.poses)) # Should be truthy when has poses assert path.poses == poses @@ -49,74 +51,74 @@ def test_head() -> None: """Test getting the first pose.""" poses = [create_test_pose(i, i, 0) for i in range(3)] path = Path(poses=poses) - assert path.head() == poses[0] + assert (path.poses[0] if path.poses else None) == poses[0] # Test empty path empty_path = Path() - assert empty_path.head() is None + assert (empty_path.poses[0] if empty_path.poses else None) is None def test_last() -> None: """Test getting the last pose.""" poses = [create_test_pose(i, i, 0) for i in range(3)] path = Path(poses=poses) - assert path.last() == poses[-1] + assert (path.poses[-1] if path.poses else None) == poses[-1] # Test empty path empty_path = Path() - assert empty_path.last() is None + assert (empty_path.poses[-1] if empty_path.poses else None) is None def test_tail() -> None: """Test getting all poses except the first.""" poses = [create_test_pose(i, i, 0) for i in range(3)] path = Path(poses=poses) - tail = path.tail() - assert len(tail) == 2 + tail = Path(header=path.header, poses=list(path.poses)[1:]) + assert len(tail.poses) == 2 assert tail.poses == poses[1:] - assert tail.frame_id == path.frame_id + assert tail.header.frame_id == path.header.frame_id # Test single element path single_path = Path(poses=[poses[0]]) - assert len(single_path.tail()) == 0 + assert len(list(single_path.poses)[1:]) == 0 # Test empty path empty_path = Path() - assert len(empty_path.tail()) == 0 + assert len(list(empty_path.poses)[1:]) == 0 def test_push_immutable() -> None: """Test immutable push operation.""" - path = Path(frame_id="map") + path = Path(header=Header(frame_id="map")) pose1 = create_test_pose(1, 1, 0) pose2 = create_test_pose(2, 2, 0) # Push should return new path - path2 = path.push(pose1) - assert len(path) == 0 # Original unchanged - assert len(path2) == 1 + path2 = Path(header=path.header, poses=[*path.poses, pose1]) + assert len(path.poses) == 0 # Original unchanged + assert len(path2.poses) == 1 assert path2.poses[0] == pose1 # Chain pushes - path3 = path2.push(pose2) - assert len(path2) == 1 # Previous unchanged - assert len(path3) == 2 + path3 = Path(header=path2.header, poses=[*path2.poses, pose2]) + assert len(path2.poses) == 1 # Previous unchanged + assert len(path3.poses) == 2 assert path3.poses == [pose1, pose2] def test_push_mutable() -> None: """Test mutable push operation.""" - path = Path(frame_id="map") + path = Path(header=Header(frame_id="map")) pose1 = create_test_pose(1, 1, 0) pose2 = create_test_pose(2, 2, 0) # Push should modify in place - path.push_mut(pose1) - assert len(path) == 1 + path.poses.append(pose1) + assert len(path.poses) == 1 assert path.poses[0] == pose1 - path.push_mut(pose2) - assert len(path) == 2 + path.poses.append(pose2) + assert len(path.poses) == 2 assert path.poses == [pose1, pose2] @@ -126,13 +128,13 @@ def test_indexing() -> None: path = Path(poses=poses) # Single index - assert path[0] == poses[0] - assert path[-1] == poses[-1] + assert next(iter(path.poses)) == poses[0] + assert list(path.poses)[-1] == poses[-1] # Slicing - assert path[1:3] == poses[1:3] - assert path[:2] == poses[:2] - assert path[3:] == poses[3:] + assert list(path.poses)[1:3] == poses[1:3] + assert list(path.poses)[:2] == poses[:2] + assert list(path.poses)[3:] == poses[3:] def test_iteration() -> None: @@ -141,7 +143,7 @@ def test_iteration() -> None: path = Path(poses=poses) collected = [] - for pose in path: + for pose in path.poses: collected.append(pose) assert collected == poses @@ -149,15 +151,15 @@ def test_iteration() -> None: def test_slice_method() -> None: """Test slice method.""" poses = [create_test_pose(i, i, 0) for i in range(5)] - path = Path(frame_id="map", poses=poses) + path = Path(header=Header(frame_id="map"), poses=poses) - sliced = path.slice(1, 4) - assert len(sliced) == 3 + sliced = Path(header=path.header, poses=list(path.poses)[1:4]) + assert len(sliced.poses) == 3 assert sliced.poses == poses[1:4] - assert sliced.frame_id == "map" + assert sliced.header.frame_id == "map" # Test open-ended slice - sliced2 = path.slice(2) + sliced2 = Path(header=path.header, poses=list(path.poses)[2:]) assert sliced2.poses == poses[2:] @@ -166,14 +168,14 @@ def test_extend_immutable() -> None: poses1 = [create_test_pose(i, i, 0) for i in range(2)] poses2 = [create_test_pose(i + 2, i + 2, 0) for i in range(2)] - path1 = Path(frame_id="map", poses=poses1) - path2 = Path(frame_id="odom", poses=poses2) + path1 = Path(header=Header(frame_id="map"), poses=poses1) + path2 = Path(header=Header(frame_id="odom"), poses=poses2) - extended = path1.extend(path2) - assert len(path1) == 2 # Original unchanged - assert len(extended) == 4 + extended = Path(header=path1.header, poses=[*path1.poses, *path2.poses]) + assert len(path1.poses) == 2 # Original unchanged + assert len(extended.poses) == 4 assert extended.poses == poses1 + poses2 - assert extended.frame_id == "map" # Keeps first path's frame + assert extended.header.frame_id == "map" # Keeps first path's frame def test_extend_mutable() -> None: @@ -181,16 +183,18 @@ def test_extend_mutable() -> None: poses1 = [create_test_pose(i, i, 0) for i in range(2)] poses2 = [create_test_pose(i + 2, i + 2, 0) for i in range(2)] - path1 = Path(frame_id="map", poses=poses1.copy()) # Use copy to avoid modifying original - path2 = Path(frame_id="odom", poses=poses2) + path1 = Path( + header=Header(frame_id="map"), poses=poses1.copy() + ) # Use copy to avoid modifying original + path2 = Path(header=Header(frame_id="odom"), poses=poses2) - path1.extend_mut(path2) - assert len(path1) == 4 + path1.poses.extend(path2.poses) + assert len(path1.poses) == 4 # Check poses are the same as concatenation for _i, (p1, p2) in enumerate(zip(path1.poses, poses1 + poses2, strict=False)): - assert p1.x == p2.x - assert p1.y == p2.y - assert p1.z == p2.z + assert p1.pose.position.x == p2.pose.position.x + assert p1.pose.position.y == p2.pose.position.y + assert p1.pose.position.z == p2.pose.position.z def test_reverse() -> None: @@ -198,8 +202,8 @@ def test_reverse() -> None: poses = [create_test_pose(i, i, 0) for i in range(3)] path = Path(poses=poses) - reversed_path = path.reverse() - assert len(path) == 3 # Original unchanged + reversed_path = Path(header=path.header, poses=list(reversed(list(path.poses)))) + assert len(path.poses) == 3 # Original unchanged assert reversed_path.poses == list(reversed(poses)) @@ -208,80 +212,56 @@ def test_clear() -> None: poses = [create_test_pose(i, i, 0) for i in range(3)] path = Path(poses=poses) - path.clear() - assert len(path) == 0 + path.poses.clear() + assert len(path.poses) == 0 assert path.poses == [] -def test_lcm_encode_decode() -> None: - """Test encoding and decoding of Path to/from binary LCM format.""" - # Create path with poses - # Use timestamps that can be represented exactly in float64 - path_ts = 1234567890.5 +def test_cdr_roundtrip_preserves_per_pose_frames_and_stamps() -> None: poses = [ PoseStamped( - ts=1234567890.0 + i * 0.1, # Use simpler timestamps - frame_id=f"frame_{i}", - position=[i * 1.5, i * 2.5, i * 3.5], - orientation=(0.1 * i, 0.2 * i, 0.3 * i, 0.9), + header=Header( + stamp=time_from_nanoseconds(1234567890000000000 + i * 100000000), + frame_id=f"frame_{i}", + ), + pose=Pose( + position=Point(x=i * 1.5, y=i * 2.5, z=i * 3.5), + orientation=Quaternion(x=0.1 * i, y=0.2 * i, z=0.3 * i, w=0.9), + ), ) for i in range(3) ] - - path_source = Path(ts=path_ts, frame_id="world", poses=poses) - - # Encode to binary - binary_msg = path_source.lcm_encode() - - # Decode from binary - path_dest = Path.lcm_decode(binary_msg) - - assert isinstance(path_dest, Path) - assert path_dest is not path_source - - # Check header - assert path_dest.frame_id == path_source.frame_id - # Path timestamp should be preserved - assert abs(path_dest.ts - path_source.ts) < 1e-6 # Microsecond precision - - # Check poses - assert len(path_dest.poses) == len(path_source.poses) - - for orig, decoded in zip(path_source.poses, path_dest.poses, strict=False): - # Check pose timestamps - assert abs(decoded.ts - orig.ts) < 1e-6 - # All poses should have the path's frame_id - assert decoded.frame_id == path_dest.frame_id - - # Check position - assert decoded.x == orig.x - assert decoded.y == orig.y - assert decoded.z == orig.z - - # Check orientation - assert decoded.orientation.x == orig.orientation.x - assert decoded.orientation.y == orig.orientation.y - assert decoded.orientation.z == orig.orientation.z - assert decoded.orientation.w == orig.orientation.w - - -def test_lcm_encode_decode_empty() -> None: - """Test encoding and decoding of empty Path.""" - path_source = Path(frame_id="base_link") - - binary_msg = path_source.lcm_encode() - path_dest = Path.lcm_decode(binary_msg) - - assert isinstance(path_dest, Path) - assert path_dest.frame_id == path_source.frame_id - assert len(path_dest.poses) == 0 - - -def test_str_representation() -> None: - """Test string representation.""" - path = Path(frame_id="map") - assert str(path) == "Path(frame_id='map', poses=0)" - - path.push_mut(create_test_pose(1, 1, 0)) - path.push_mut(create_test_pose(2, 2, 0)) - assert str(path) == "Path(frame_id='map', poses=2)" + source = Path( + header=Header(stamp=time_from_nanoseconds(1234567890500000000), frame_id="world"), + poses=poses, + ) + decoded = Path.decode(source.encode()) + assert decoded is not source + assert decoded.header == source.header + assert list(decoded.poses) == poses + independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(source.encode(), Path.msg_name) + assert independent.header.frame_id == "world" + assert independent.header.stamp.sec == 1234567890 + assert independent.header.stamp.nanosec == 500000000 + for i, value in enumerate(independent.poses): + assert value.header.frame_id == f"frame_{i}" + assert value.header.stamp.sec == 1234567890 + assert value.header.stamp.nanosec == i * 100000000 + assert (value.pose.position.x, value.pose.position.y, value.pose.position.z) == ( + i * 1.5, + i * 2.5, + i * 3.5, + ) + assert ( + value.pose.orientation.x, + value.pose.orientation.y, + value.pose.orientation.z, + value.pose.orientation.w, + ) == (0.1 * i, 0.2 * i, 0.3 * i, 0.9) + + +def test_cdr_empty_path() -> None: + source = Path(header=Header(frame_id="base_link")) + decoded = Path.decode(source.encode()) + assert decoded.header.frame_id == "base_link" + assert len(decoded.poses) == 0 diff --git a/dimos/msgs/occupancy.py b/dimos/msgs/occupancy.py new file mode 100644 index 00000000000..773c34fe137 --- /dev/null +++ b/dimos/msgs/occupancy.py @@ -0,0 +1,132 @@ +# 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. + +"""Array operations on generated occupancy grids.""" + +import math +from pathlib import Path +from typing import Any + +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid +from dimos_generated.std_msgs.msg import Header +import numpy as np +from numpy.typing import NDArray + +from dimos.msgs.geometry import pose_matrix + + +def occupancy_view(msg: OccupancyGrid) -> NDArray[np.int8]: + """Borrow a read-only row-major occupancy grid, retaining its owner.""" + if len(msg.data) != msg.info.height * msg.info.width: + raise ValueError("occupancy data length does not match dimensions") + return msg.data.view().reshape(msg.info.height, msg.info.width) + + +def block_max_reduce(cells: NDArray[np.int8], factor: int) -> NDArray[np.int8]: + """Coarsen an occupancy grid by taking the max of factor x factor blocks. + + Block maximum (not mean) so coarsening never erases an obstacle; a block + is unknown (-1) only when every cell in it is unknown. Trailing remainder + rows/cols are trimmed, keeping row 0/col 0 - the origin corner - so the + grid origin stays valid. Grids thinner than `factor` pass through. + """ + h, w = cells.shape[:2] + new_h, new_w = h // factor, w // factor + if new_h == 0 or new_w == 0: + return cells + trimmed = cells[: new_h * factor, : new_w * factor] + blocks = trimmed.reshape(new_h, factor, new_w, factor) + # Sink unknown below every known value for the max, then map it back. + as_int = blocks.astype(np.int16) + known = np.where(as_int < 0, -1000, as_int) + reduced = known.max(axis=(1, 3)) + reduced[reduced == -1000] = -1 + result: NDArray[np.int8] = reduced.astype(np.int8) + return result + + +def occupancy_extent(message: OccupancyGrid) -> tuple[float, float]: + """Return physical width and height, rejecting an invalid grid resolution.""" + info = message.info + if not math.isfinite(info.resolution) or info.resolution <= 0: + raise ValueError("OccupancyGrid resolution must be finite and positive") + return info.width * info.resolution, info.height * info.resolution + + +def world_to_grid(message: OccupancyGrid, point: Point) -> tuple[float, float]: + """Convert a world point to continuous grid coordinates, including origin rotation.""" + occupancy_extent(message) + matrix = pose_matrix(message.info.origin) + local = matrix[:3, :3].T @ (np.array([point.x, point.y, point.z]) - matrix[:3, 3]) + if not np.isfinite(local).all(): + raise ValueError("grid coordinates must be finite") + return float(local[0] / message.info.resolution), float(local[1] / message.info.resolution) + + +def grid_to_world(message: OccupancyGrid, coordinates: tuple[float, float]) -> Point: + """Convert continuous grid coordinates to the grid plane in world coordinates.""" + occupancy_extent(message) + x, y = coordinates + if not math.isfinite(x) or not math.isfinite(y): + raise ValueError("grid coordinates must be finite") + result = pose_matrix(message.info.origin) @ np.array( + [x * message.info.resolution, y * message.info.resolution, 0.0, 1.0] + ) + return Point(x=float(result[0]), y=float(result[1]), z=float(result[2])) + + +def occupancy_from_file(path: Path, *, header: Header, resolution: float = 0.05) -> OccupancyGrid: + """Load a static .npy or grayscale PNG grid with explicit frame/time metadata. + + Cells must be ROS occupancy values (-1 unknown, 0..100 known); PNG 255 is + unknown. NumPy object/pickle loading is disabled. Row order is unchanged. + """ + if not math.isfinite(resolution) or resolution <= 0: + raise ValueError("OccupancyGrid resolution must be finite and positive") + cells: NDArray[Any] + if path.suffix == ".npy": + cells = np.load(path, allow_pickle=False) + elif path.suffix == ".png": + from PIL import Image + + with Image.open(path) as image: + cells = np.asarray(image.convert("L")).astype(np.int16) + cells[cells == 255] = -1 + else: + raise ValueError(f"Unsupported occupancy file format: {path.suffix}") + return occupancy_from_array(cells, header=header, resolution=resolution) + + +def occupancy_from_array( + cells: NDArray[Any], *, header: Header, resolution: float, origin: Pose | None = None +) -> OccupancyGrid: + """Copy a 2D ROS occupancy array with explicit map metadata.""" + if not math.isfinite(resolution) or resolution <= 0: + raise ValueError("OccupancyGrid resolution must be finite and positive") + if cells.ndim != 2 or cells.dtype.kind not in "iu": + raise ValueError("Occupancy cells must be a 2D integer array") + if np.any((cells < -1) | (cells > 100)): + raise ValueError("Occupancy cells must be -1 or 0..100") + return OccupancyGrid( + header=header, + info=MapMetaData( + map_load_time=header.stamp, + resolution=resolution, + height=cells.shape[0], + width=cells.shape[1], + origin=origin if origin is not None else Pose(orientation=Quaternion(w=1)), + ), + data=cells.astype(np.int8).reshape(-1), + ) diff --git a/dimos/msgs/pointcloud.py b/dimos/msgs/pointcloud.py new file mode 100644 index 00000000000..aeb81a58a56 --- /dev/null +++ b/dimos/msgs/pointcloud.py @@ -0,0 +1,348 @@ +# 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. + +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +"""NumPy conversions for generated PointCloud2 messages.""" + +import math +from typing import Any + +from dimos_generated.geometry_msgs.msg import TransformStamped +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2, PointField +from dimos_generated.std_msgs.msg import Header +import numpy as np +from numpy.typing import NDArray + +from dimos.msgs.camera_info import intrinsic_matrix +from dimos.msgs.geometry import transform_matrix +from dimos.msgs.image import image_to_rgb, image_view +from dimos.msgs.time import to_nanoseconds + +_FIELD_KINDS = {1: "i1", 2: "u1", 3: "i2", 4: "u2", 5: "i4", 6: "u4", 7: "f4", 8: "f8"} + + +def pointcloud_view(message: PointCloud2) -> NDArray[Any]: + """Borrow read-only structured (height, width) points, including row/field padding. + + Fields retain their declared counts and endianness. The returned array keeps + the message's storage alive. Use ``view.copy()`` for independent mutable data. + """ + if message.width and not message.point_step: + raise ValueError("nonempty point cloud requires a positive point_step") + if message.row_step < message.width * message.point_step: + raise ValueError("point cloud row_step is smaller than its row of points") + if len(message.data) != message.height * message.row_step: + raise ValueError("point cloud data length does not match height * row_step") + names: list[str] = [] + formats: list[Any] = [] + offsets: list[int] = [] + for field in message.fields: + if not field.name or field.name in names: + raise ValueError("point cloud field names must be nonempty and unique") + if field.datatype not in _FIELD_KINDS or not field.count: + raise ValueError(f"invalid point cloud datatype/count for {field.name!r}") + dtype = np.dtype((">" if message.is_bigendian else "<") + _FIELD_KINDS[field.datatype]) + if field.offset + dtype.itemsize * field.count > message.point_step: + raise ValueError(f"point cloud field {field.name!r} exceeds point_step") + names.append(field.name) + formats.append(dtype if field.count == 1 else (dtype, (field.count,))) + offsets.append(field.offset) + dtype = np.dtype( + {"names": names, "formats": formats, "offsets": offsets, "itemsize": message.point_step} + ) + return np.ndarray( + (message.height, message.width), + dtype=dtype, + buffer=message.data.view(), + strides=(message.row_step, message.point_step), + ) + + +def pointcloud_xyz(message: PointCloud2) -> NDArray[np.float64]: + """Copy scalar x/y/z fields into an N-by-3 float64 array, preserving point order.""" + points = pointcloud_view(message) + for name in ("x", "y", "z"): + if name not in (points.dtype.names or ()) or points[name].ndim != 2: + raise ValueError(f"point cloud requires a scalar {name!r} field") + return np.stack([points[name].ravel() for name in ("x", "y", "z")], axis=-1).astype(np.float64) + + +def pointcloud_from_xyz(points: NDArray[Any], *, header: Header) -> PointCloud2: + """Copy Nx3 or HxWx3 coordinates into a tightly packed little-endian XYZ cloud. + + XYZ fields are float32. NaNs/infinities remain missing-point markers and set + is_dense=False; finite coordinates outside float32 range are rejected. + """ + values = np.asarray(points) + if values.ndim not in (2, 3) or values.shape[-1] != 3: + raise ValueError("XYZ coordinates must have shape (N, 3) or (H, W, 3)") + if values.dtype.kind not in "fiu": + raise ValueError("XYZ coordinates must be numeric") + finite = np.isfinite(values) + if np.any(np.abs(values[finite].astype(np.float64)) > np.finfo(np.float32).max): + raise ValueError("XYZ coordinates exceed float32 range") + packed = np.ascontiguousarray(values, dtype=" PointCloud2: + """Copy selected point records in row order, retaining every field and point byte. + + The result is unorganized with no row padding. Point padding, endianness, + field metadata, source header, and the conservative density flag survive. + """ + pointcloud_view(message) # Validate the full declared layout before selecting bytes. + if keep.dtype != np.bool_ or keep.shape != (message.height * message.width,): + raise ValueError("point selection must be a flat boolean mask matching the point count") + records: NDArray[np.uint8] = np.ndarray( + (message.height, message.width, message.point_step), + dtype=np.uint8, + buffer=message.data.view(), + strides=(message.row_step, message.point_step, 1), + ) + selected = records.reshape(message.height * message.width, message.point_step)[keep] + width = len(selected) + return PointCloud2( + header=message.header, + height=1, + width=width, + fields=message.fields, + is_bigendian=message.is_bigendian, + point_step=message.point_step, + row_step=width * message.point_step, + data=selected.tobytes(), + is_dense=message.is_dense, + ) + + +def pointcloud_rgb(message: PointCloud2) -> NDArray[np.uint8] | None: + """Copy RGB bytes from the standard packed rgb/rgba field, if present. + + FLOAT32 fields carry color bits, not numeric floating-point color values. + Both representations preserve the declared byte order and row padding. + """ + points = pointcloud_view(message) + names = points.dtype.names or () + name = next((name for name in ("rgb", "rgba") if name in names), None) + if name is None: + return None + field = points[name] + if field.ndim != 2 or field.dtype.kind not in "fu" or field.dtype.itemsize != 4: + raise ValueError("packed point cloud color must be a scalar FLOAT32 or UINT32 field") + packed = field.view(np.dtype(">u4" if message.is_bigendian else "> shift) & 255 for shift in (16, 8, 0)], axis=-1).astype(np.uint8) + + +def concatenate_clouds(first: PointCloud2, second: PointCloud2) -> PointCloud2: + """Join compatible point records, preserving fields, endian, and point padding. + + Row padding is removed. The first frame and newest exact source stamp survive. + Empty inputs are identities; differing nonempty frames/layouts are rejected. + """ + pointcloud_view(first) + pointcloud_view(second) + if not first.width * first.height: + return PointCloud2.decode(second.encode()) + if not second.width * second.height: + return PointCloud2.decode(first.encode()) + if first.header.frame_id != second.header.frame_id: + raise ValueError("cannot concatenate point clouds in different frames") + if ( + first.fields != second.fields + or first.point_step != second.point_step + or first.is_bigendian != second.is_bigendian + ): + raise ValueError("cannot concatenate point clouds with different record layouts") + a = select_points(first, np.ones(first.width * first.height, dtype=np.bool_)) + b = select_points(second, np.ones(second.width * second.height, dtype=np.bool_)) + stamp = max((first.header.stamp, second.header.stamp), key=to_nanoseconds) + return PointCloud2( + header=Header(frame_id=first.header.frame_id, stamp=stamp), + height=1, + width=a.width + b.width, + fields=first.fields, + is_bigendian=first.is_bigendian, + point_step=first.point_step, + row_step=a.row_step + b.row_step, + data=bytes(a.data) + bytes(b.data), + is_dense=first.is_dense and second.is_dense, + ) + + +def transform_cloud(message: PointCloud2, transform: TransformStamped) -> PointCloud2: + """Apply parent←child to XYZ while retaining the complete organized record layout. + + Copy storage before changing coordinates. Keep all other fields, point/row + padding, endian and the cloud's exact source stamp. Frames must match; + integer XYZ is rejected because transformed coordinates would be truncated. + """ + if message.header.frame_id != transform.child_frame_id: + raise ValueError("point cloud frame does not match transform child frame") + view = pointcloud_view(message) + for name in ("x", "y", "z"): + if name not in (view.dtype.names or ()) or view.dtype[name].kind != "f": + raise ValueError("transform requires scalar floating-point XYZ fields") + matrix = transform_matrix(transform.transform) + xyz = pointcloud_xyz(message) + transformed = xyz @ matrix[:3, :3].T + matrix[:3, 3] + for index, name in enumerate(("x", "y", "z")): + values = transformed[:, index] + if np.any(np.abs(values[np.isfinite(values)]) > np.finfo(view.dtype[name]).max): + raise ValueError(f"transformed {name} coordinates exceed field range") + payload = bytearray(bytes(message.data)) + writable = np.ndarray( + (message.height, message.width), + dtype=view.dtype, + buffer=payload, + strides=(message.row_step, message.point_step), + ) + for index, name in enumerate(("x", "y", "z")): + writable[name] = transformed[:, index].reshape(message.height, message.width) + output = PointCloud2.decode(message.encode()) + output.data = bytes(payload) + output.header = Header(stamp=message.header.stamp, frame_id=transform.header.frame_id) + return output + + +def pointcloud_from_xyz_rgb( + points: NDArray[Any], colors: NDArray[np.uint8], *, header: Header +) -> PointCloud2: + """Copy XYZ and RGB8 into a declared XYZ FLOAT32 / packed RGB UINT32 layout.""" + cloud = pointcloud_from_xyz(points, header=header) + rgb = np.asarray(colors) + if rgb.dtype != np.uint8 or rgb.shape != np.asarray(points).shape: + raise ValueError("RGB must be uint8 with the same (..., 3) shape as XYZ") + records = np.empty((cloud.height, cloud.width), dtype=[("xyz", " PointCloud2: + """Project a rectified RGB/depth pair into the depth optical frame. + + Integer depth is multiplied by depth_scale to obtain meters; 32FC1 already + represents meters. Invalid/nonpositive/too-distant samples are removed. + Scale K when calibration resolution differs, retaining the depth header. + """ + if ( + not math.isfinite(depth_scale) + or depth_scale <= 0 + or not math.isfinite(depth_trunc) + or depth_trunc <= 0 + ): + raise ValueError("depth scale and truncation must be finite and positive") + if (color.width, color.height) != (depth.width, depth.height): + raise ValueError("color and depth dimensions do not match") + if depth.encoding not in ("16UC1", "mono16", "32FC1"): + raise ValueError("RGBD requires uint16 or float32 depth") + if calibration.width <= 0 or calibration.height <= 0: + raise ValueError("calibration dimensions must be positive") + matrix = intrinsic_matrix(calibration).copy() + matrix[0, :] *= depth.width / calibration.width + matrix[1, :] *= depth.height / calibration.height + distances = image_view(depth).astype(np.float64) + if depth.encoding != "32FC1": + distances *= depth_scale + valid = np.isfinite(distances) & (distances > 0) & (distances < depth_trunc) + rows, columns = np.indices(distances.shape) + z = distances[valid] + xyz = np.column_stack( + ( + (columns[valid] - matrix[0, 2]) * z / matrix[0, 0], + (rows[valid] - matrix[1, 2]) * z / matrix[1, 1], + z, + ) + ) + return pointcloud_from_xyz_rgb(xyz, image_to_rgb(color)[valid], header=depth.header) + + +def voxel_downsample_cloud(message: PointCloud2, voxel_size: float = 0.025) -> PointCloud2: + """Downsample XYZ/RGB clouds with the existing Open3D tensor backend. + + Non-positive sizes and clouds below twenty points retain the existing fast + path. Unsupported extra fields are rejected rather than silently dropped. + Source headers are preserved; generated messages own no numerical methods. + """ + if not math.isfinite(voxel_size): + raise ValueError("voxel_size must be finite") + if voxel_size <= 0 or message.width * message.height < 20: + return message + names = {field.name for field in message.fields} + if names not in ({"x", "y", "z"}, {"x", "y", "z", "rgb"}): + raise ValueError("voxel downsampling supports only XYZ and optional RGB fields") + points = pointcloud_xyz(message).astype(np.float32) + if not np.isfinite(points).all(): + raise ValueError("voxel downsampling requires finite XYZ positions") + colors = pointcloud_rgb(message) + import open3d as o3d # type: ignore[import-untyped] + + cloud = o3d.t.geometry.PointCloud(o3d.core.Tensor(points)) + if colors is not None: + cloud.point["colors"] = o3d.core.Tensor(colors.astype(np.float32) / 255.0) + reduced = cloud.voxel_down_sample(voxel_size) + xyz = reduced.point["positions"].numpy() + if colors is None: + return pointcloud_from_xyz(xyz, header=message.header) + rgb = np.clip(reduced.point["colors"].numpy() * 255.0, 0, 255).astype(np.uint8) + return pointcloud_from_xyz_rgb(xyz, rgb, header=message.header) + + +def pointcloud_to_open3d(message: PointCloud2) -> Any: + """Copy generated XYZ/RGB values into a native geometry algorithm input.""" + import open3d as o3d # type: ignore[import-untyped] + + result = o3d.geometry.PointCloud() + result.points = o3d.utility.Vector3dVector(pointcloud_xyz(message)) + colors = pointcloud_rgb(message) + if colors is not None: + result.colors = o3d.utility.Vector3dVector(colors.astype(np.float64) / 255.0) + return result + + +def cloud_bounds_intersect(first: PointCloud2, second: PointCloud2) -> bool: + """Test closed axis-aligned bounds; empty or entirely nonfinite clouds do not intersect.""" + a, b = pointcloud_xyz(first), pointcloud_xyz(second) + a, b = a[np.isfinite(a).all(axis=1)], b[np.isfinite(b).all(axis=1)] + if len(a) == 0 or len(b) == 0: + return False + return bool(np.all(a.min(axis=0) <= b.max(axis=0)) and np.all(a.max(axis=0) >= b.min(axis=0))) diff --git a/dimos/msgs/protocol.py b/dimos/msgs/protocol.py index 38d7ca57e29..11e218e9009 100644 --- a/dimos/msgs/protocol.py +++ b/dimos/msgs/protocol.py @@ -12,20 +12,23 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Protocol, runtime_checkable +from typing import ClassVar, Protocol, runtime_checkable + +from typing_extensions import Self @runtime_checkable class DimosMsg(Protocol): - """Protocol for dimos message types (LCM-based messages from dimos.msgs).""" + """Transport-neutral contract implemented by generated CDR messages.""" - msg_name: str + msg_name: ClassVar[str] + schema: ClassVar[str] @classmethod - def lcm_decode(cls, data: bytes) -> "DimosMsg": + def decode(cls, data: bytes) -> Self: """Decode bytes into a message instance.""" ... - def lcm_encode(self) -> bytes: + def encode(self, little_endian: bool = True) -> bytes: """Encode this message instance into bytes.""" ... diff --git a/dimos/msgs/sensor_msgs/test_CameraInfo.py b/dimos/msgs/sensor_msgs/test_CameraInfo.py index e8a1c3a789d..6967b9c0fd3 100644 --- a/dimos/msgs/sensor_msgs/test_CameraInfo.py +++ b/dimos/msgs/sensor_msgs/test_CameraInfo.py @@ -13,23 +13,27 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.sensor_msgs.msg import CameraInfo +from dimos_generated.std_msgs.msg import Header import numpy as np +from rosbags.typesys import Stores, get_typestore from dimos.constants import DIMOS_PROJECT_ROOT -from dimos.msgs.sensor_msgs.CameraInfo import CalibrationProvider, CameraInfo +from dimos.msgs.camera_info import CalibrationProvider, camera_info_from_yaml, intrinsic_matrix +from dimos.msgs.time import time_from_nanoseconds, to_nanoseconds -def test_lcm_encode_decode() -> None: - """Test LCM encode/decode preserves CameraInfo data.""" - print("Testing CameraInfo LCM encode/decode...") +def test_encode_decode() -> None: + """Test CDR encode/decode preserves CameraInfo data.""" + print("Testing CameraInfo CDR encode/decode...") # Create test camera info with sample calibration data original = CameraInfo( height=480, width=640, distortion_model="plumb_bob", - D=[-0.1, 0.05, 0.001, -0.002, 0.0], # 5 distortion coefficients - K=[ + d=[-0.1, 0.05, 0.001, -0.002, 0.0], # 5 distortion coefficients + k=[ 500.0, 0.0, 320.0, # fx, 0, cx @@ -40,8 +44,8 @@ def test_lcm_encode_decode() -> None: 0.0, 1.0, ], # 0, 0, 1 - R=[1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0], - P=[ + r=[1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0], + p=[ 500.0, 0.0, 320.0, @@ -57,20 +61,25 @@ def test_lcm_encode_decode() -> None: ], # 0, 0, 1, 0 binning_x=2, binning_y=2, - frame_id="camera_optical_frame", - ts=1234567890.123456, + header=Header( + frame_id="camera_optical_frame", stamp=time_from_nanoseconds(1234567890123456789) + ), ) # Set ROI - original.roi_x_offset = 100 - original.roi_y_offset = 50 - original.roi_height = 200 - original.roi_width = 300 - original.roi_do_rectify = True + original.roi.x_offset = 100 + original.roi.y_offset = 50 + original.roi.height = 200 + original.roi.width = 300 + original.roi.do_rectify = True # Encode and decode - binary_msg = original.lcm_encode() - decoded = CameraInfo.lcm_decode(binary_msg) + binary_msg = original.encode() + decoded = CameraInfo.decode(binary_msg) + independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(binary_msg, CameraInfo.msg_name) + assert independent.header.stamp.nanosec == 123456789 + assert independent.roi.x_offset == 100 + np.testing.assert_array_equal(independent.k, original.k) # Check basic properties assert original.height == decoded.height, ( @@ -85,27 +94,27 @@ def test_lcm_encode_decode() -> None: print(f"✓ Distortion model preserved: '{decoded.distortion_model}'") # Check distortion coefficients - assert len(original.D) == len(decoded.D), ( - f"D length mismatch: {len(original.D)} vs {len(decoded.D)}" + assert len(original.d) == len(decoded.d), ( + f"D length mismatch: {len(original.d)} vs {len(decoded.d)}" ) np.testing.assert_allclose( - original.D, decoded.D, rtol=1e-9, atol=1e-9, err_msg="Distortion coefficients don't match" + original.d, decoded.d, rtol=1e-9, atol=1e-9, err_msg="Distortion coefficients don't match" ) - print(f"✓ Distortion coefficients preserved: {len(decoded.D)} coefficients") + print(f"✓ Distortion coefficients preserved: {len(decoded.d)} coefficients") # Check camera matrices np.testing.assert_allclose( - original.K, decoded.K, rtol=1e-9, atol=1e-9, err_msg="K matrix doesn't match" + original.k, decoded.k, rtol=1e-9, atol=1e-9, err_msg="K matrix doesn't match" ) print("✓ Intrinsic matrix K preserved") np.testing.assert_allclose( - original.R, decoded.R, rtol=1e-9, atol=1e-9, err_msg="R matrix doesn't match" + original.r, decoded.r, rtol=1e-9, atol=1e-9, err_msg="R matrix doesn't match" ) print("✓ Rectification matrix R preserved") np.testing.assert_allclose( - original.P, decoded.P, rtol=1e-9, atol=1e-9, err_msg="P matrix doesn't match" + original.p, decoded.p, rtol=1e-9, atol=1e-9, err_msg="P matrix doesn't match" ) print("✓ Projection matrix P preserved") @@ -119,25 +128,25 @@ def test_lcm_encode_decode() -> None: print(f"✓ Binning preserved: {decoded.binning_x}x{decoded.binning_y}") # Check ROI - assert original.roi_x_offset == decoded.roi_x_offset, "ROI x_offset mismatch" - assert original.roi_y_offset == decoded.roi_y_offset, "ROI y_offset mismatch" - assert original.roi_height == decoded.roi_height, "ROI height mismatch" - assert original.roi_width == decoded.roi_width, "ROI width mismatch" - assert original.roi_do_rectify == decoded.roi_do_rectify, "ROI do_rectify mismatch" + assert original.roi.x_offset == decoded.roi.x_offset, "ROI x_offset mismatch" + assert original.roi.y_offset == decoded.roi.y_offset, "ROI y_offset mismatch" + assert original.roi.height == decoded.roi.height, "ROI height mismatch" + assert original.roi.width == decoded.roi.width, "ROI width mismatch" + assert original.roi.do_rectify == decoded.roi.do_rectify, "ROI do_rectify mismatch" print("✓ ROI preserved") # Check metadata - assert original.frame_id == decoded.frame_id, ( - f"Frame ID mismatch: '{original.frame_id}' vs '{decoded.frame_id}'" + assert original.header.frame_id == decoded.header.frame_id, ( + f"Frame ID mismatch: '{original.header.frame_id}' vs '{decoded.header.frame_id}'" ) - print(f"✓ Frame ID preserved: '{decoded.frame_id}'") + print(f"✓ Frame ID preserved: '{decoded.header.frame_id}'") - assert abs(original.ts - decoded.ts) < 1e-6, ( - f"Timestamp mismatch: {original.ts} vs {decoded.ts}" + assert to_nanoseconds(original.header.stamp) == to_nanoseconds(decoded.header.stamp), ( + f"Timestamp mismatch: {to_nanoseconds(original.header.stamp)} vs {to_nanoseconds(decoded.header.stamp)}" ) - print(f"✓ Timestamp preserved: {decoded.ts}") + print(f"✓ Timestamp preserved: {to_nanoseconds(decoded.header.stamp)}") - print("✓ LCM encode/decode test passed - all properties preserved!") + print("✓ CDR encode/decode test passed - all properties preserved!") def test_numpy_matrix_operations() -> None: @@ -148,29 +157,29 @@ def test_numpy_matrix_operations() -> None: # Test K matrix K = np.array([[525.0, 0.0, 319.5], [0.0, 525.0, 239.5], [0.0, 0.0, 1.0]]) - camera_info.set_K_matrix(K) - K_retrieved = camera_info.get_K_matrix() + camera_info.k = K.ravel() + K_retrieved = intrinsic_matrix(camera_info) np.testing.assert_allclose(K, K_retrieved, rtol=1e-9, atol=1e-9) print("✓ K matrix setter/getter works") # Test P matrix P = np.array([[525.0, 0.0, 319.5, 0.0], [0.0, 525.0, 239.5, 0.0], [0.0, 0.0, 1.0, 0.0]]) - camera_info.set_P_matrix(P) - P_retrieved = camera_info.get_P_matrix() + camera_info.p = P.ravel() + P_retrieved = np.asarray(camera_info.p).reshape(3, 4) np.testing.assert_allclose(P, P_retrieved, rtol=1e-9, atol=1e-9) print("✓ P matrix setter/getter works") # Test R matrix R = np.eye(3) - camera_info.set_R_matrix(R) - R_retrieved = camera_info.get_R_matrix() + camera_info.r = R.ravel() + R_retrieved = np.asarray(camera_info.r).reshape(3, 3) np.testing.assert_allclose(R, R_retrieved, rtol=1e-9, atol=1e-9) print("✓ R matrix setter/getter works") # Test D coefficients D = np.array([-0.2, 0.1, 0.001, -0.002, 0.05]) - camera_info.set_D_coeffs(D) - D_retrieved = camera_info.get_D_coeffs() + camera_info.d = D + D_retrieved = np.asarray(camera_info.d) np.testing.assert_allclose(D, D_retrieved, rtol=1e-9, atol=1e-9) print("✓ D coefficients setter/getter works") @@ -185,24 +194,24 @@ def test_equality() -> None: height=480, width=640, distortion_model="plumb_bob", - D=[-0.1, 0.05, 0.0, 0.0, 0.0], - frame_id="camera1", + d=[-0.1, 0.05, 0.0, 0.0, 0.0], + header=Header(frame_id="camera1"), ) info2 = CameraInfo( height=480, width=640, distortion_model="plumb_bob", - D=[-0.1, 0.05, 0.0, 0.0, 0.0], - frame_id="camera1", + d=[-0.1, 0.05, 0.0, 0.0, 0.0], + header=Header(frame_id="camera1"), ) info3 = CameraInfo( height=720, width=1280, # Different resolution distortion_model="plumb_bob", - D=[-0.1, 0.05, 0.0, 0.0, 0.0], - frame_id="camera1", + d=[-0.1, 0.05, 0.0, 0.0, 0.0], + header=Header(frame_id="camera1"), ) assert info1 == info2, "Identical CameraInfo objects should be equal" @@ -227,16 +236,16 @@ def test_camera_info_from_yaml() -> None: ) # Load CameraInfo from YAML - camera_info = CameraInfo.from_yaml(str(yaml_path)) + camera_info = camera_info_from_yaml(yaml_path, header=Header(frame_id="camera_optical")) # Verify loaded values assert camera_info.width == 640 assert camera_info.height == 376 assert camera_info.distortion_model == "plumb_bob" - assert camera_info.frame_id == "camera_optical" + assert camera_info.header.frame_id == "camera_optical" # Check camera matrix K - K = camera_info.get_K_matrix() + K = intrinsic_matrix(camera_info) assert K.shape == (3, 3) assert np.isclose(K[0, 0], 379.45267) # fx assert np.isclose(K[1, 1], 380.67871) # fy @@ -244,12 +253,12 @@ def test_camera_info_from_yaml() -> None: assert np.isclose(K[1, 2], 228.00954) # cy # Check distortion coefficients - D = camera_info.get_D_coeffs() + D = np.asarray(camera_info.d) assert len(D) == 5 assert np.isclose(D[0], -0.309435) # Check projection matrix P - P = camera_info.get_P_matrix() + P = np.asarray(camera_info.p).reshape(3, 4) assert P.shape == (3, 4) assert np.isclose(P[0, 0], 291.12888) diff --git a/dimos/msgs/sensor_msgs/test_CompressedImage.py b/dimos/msgs/sensor_msgs/test_CompressedImage.py index a7d54e3e735..ecae6d26b06 100644 --- a/dimos/msgs/sensor_msgs/test_CompressedImage.py +++ b/dimos/msgs/sensor_msgs/test_CompressedImage.py @@ -12,13 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.sensor_msgs.msg import CompressedImage +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest -from dimos.msgs.sensor_msgs.CompressedImage import CompressedImage -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat - -pytestmark = pytest.mark.skipif_no_turbojpeg +from dimos.msgs.image import ( + compressed_image_from_image, + image_from_array, + image_from_compressed, + image_to_bgr, + image_to_rgb, + image_view, +) +from dimos.msgs.time import time_from_seconds, to_nanoseconds +from dimos.visualization.rerun.message_helpers import image_archetype @pytest.fixture() @@ -27,40 +35,43 @@ def rgb_image(): gradient = np.linspace(0, 255, 640, dtype=np.uint8) data = np.broadcast_to(gradient, (480, 640)).copy() data = np.stack([data, data // 2, rng.randint(0, 50, (480, 640), dtype=np.uint8)], axis=-1) - return Image(data=data, format=ImageFormat.RGB, frame_id="cam", ts=1234.5678) + return image_from_array( + data, encoding="rgb8", header=Header(frame_id="cam", stamp=time_from_seconds(1234.5678)) + ) def test_jpeg_roundtrip_preserves_meta_and_pixels(rgb_image) -> None: - ci = CompressedImage.from_image(rgb_image, quality=90) - assert ci.format == "jpeg" - assert ci.frame_id == "cam" - assert ci.ts == rgb_image.ts - assert 0 < len(ci.data) < rgb_image.data.nbytes // 4 - - img = ci.decode() - assert img.format == ImageFormat.RGB - assert img.frame_id == "cam" - assert img.ts == rgb_image.ts - assert img.shape == rgb_image.shape - diff = np.abs(img.data.astype(int) - rgb_image.data.astype(int)).mean() + ci = compressed_image_from_image(rgb_image, quality=90) + assert "; jpeg compressed " in ci.format + assert ci.header.frame_id == "cam" + assert to_nanoseconds(ci.header.stamp) == to_nanoseconds(rgb_image.header.stamp) + assert 0 < len(ci.data) < image_view(rgb_image).nbytes // 4 + img = image_from_compressed(ci) + assert img.encoding in ("rgb8", "bgr8") + assert img.header.frame_id == "cam" + assert to_nanoseconds(img.header.stamp) == to_nanoseconds(rgb_image.header.stamp) + assert image_view(img).shape == image_view(rgb_image).shape + diff = np.abs(image_to_rgb(img).astype(int) - image_view(rgb_image).astype(int)).mean() assert diff < 5, f"JPEG q90 mean pixel error too high: {diff}" def test_png_roundtrip_is_lossless_bgr(rgb_image) -> None: - bgr = rgb_image.to_bgr() - ci = CompressedImage.from_image(bgr, format="png") - img = ci.decode() - assert img.format == ImageFormat.BGR - assert np.array_equal(img.data, bgr.data) - assert img.ts == bgr.ts + bgr = image_from_array(image_to_bgr(rgb_image), encoding="bgr8", header=rgb_image.header) + ci = compressed_image_from_image(bgr, format="png") + img = image_from_compressed(ci) + assert img.encoding == "bgr8" + assert np.array_equal(image_view(img), image_view(bgr)) + assert to_nanoseconds(img.header.stamp) == to_nanoseconds(bgr.header.stamp) def test_png_roundtrip_is_lossless_gray16() -> None: data = np.arange(100 * 80, dtype=np.uint16).reshape(100, 80) - src = Image(data=data, format=ImageFormat.GRAY16, frame_id="d", ts=1.0) - img = CompressedImage.from_image(src, format="png").decode() - assert img.format == ImageFormat.GRAY16 - assert np.array_equal(img.data, data) + src = image_from_array( + data, encoding="mono16", header=Header(frame_id="d", stamp=time_from_seconds(1.0)) + ) + img = image_from_compressed(compressed_image_from_image(src, format="png")) + assert img.encoding == "mono16" + assert np.array_equal(image_view(img), data) def test_png_decode_alpha_is_bgra() -> None: @@ -70,40 +81,44 @@ def test_png_decode_alpha_is_bgra() -> None: arr[..., 3] = 128 ok, buf = cv2.imencode(".png", arr) assert ok - img = CompressedImage(data=buf.tobytes(), format="png").decode() - assert img.format == ImageFormat.BGRA - assert np.array_equal(img.data, arr) + img = image_from_compressed(CompressedImage(data=buf.tobytes(), format="png")) + assert img.encoding == "bgra8" + assert np.array_equal(image_view(img), arr) def test_lcm_roundtrip(rgb_image) -> None: - ci = CompressedImage.from_image(rgb_image) - wire = ci.lcm_encode() - out = CompressedImage.lcm_decode(wire) + ci = compressed_image_from_image(rgb_image) + wire = ci.encode() + out = CompressedImage.decode(wire) assert out.data == ci.data - assert out.format == "jpeg" - assert out.frame_id == "cam" - assert abs(out.ts - ci.ts) < 1e-6 + assert "; jpeg compressed " in out.format + assert out.header.frame_id == "cam" + assert abs(to_nanoseconds(out.header.stamp) - to_nanoseconds(ci.header.stamp)) < 1e-06 def test_jpeg_rejects_depth_formats() -> None: - depth = Image(data=np.zeros((10, 10), dtype=np.uint16), format=ImageFormat.GRAY16) - with pytest.raises(ValueError, match="use format='png'"): - CompressedImage.from_image(depth) + depth = image_from_array( + np.zeros((10, 10), dtype=np.uint16), encoding="mono16", header=Header() + ) + with pytest.raises(ValueError, match="JPEG requires 8-bit"): + compressed_image_from_image(depth) def test_png_rejects_float_depth() -> None: - depth = Image(data=np.zeros((10, 10), dtype=np.float32), format=ImageFormat.DEPTH) + depth = image_from_array( + np.zeros((10, 10), dtype=np.float32), encoding="32FC1", header=Header() + ) with pytest.raises(ValueError, match="PNG cannot encode"): - CompressedImage.from_image(depth, format="png") + compressed_image_from_image(depth, format="png") def test_from_image_rejects_non_image() -> None: with pytest.raises(TypeError): - CompressedImage.from_image(b"not an image") # type: ignore[arg-type] + compressed_image_from_image(b"not an image") def test_max_width_resizes(rgb_image) -> None: - img = CompressedImage.from_image(rgb_image, max_width=320).decode() + img = image_from_compressed(compressed_image_from_image(rgb_image, max_width=320)) assert img.width <= 320 assert img.height <= 320 @@ -111,46 +126,73 @@ def test_max_width_resizes(rgb_image) -> None: def test_to_rerun_is_encoded_image(rgb_image) -> None: import rerun as rr - assert isinstance(CompressedImage.from_image(rgb_image).to_rerun(), rr.EncodedImage) + assert isinstance(image_archetype(compressed_image_from_image(rgb_image)), rr.EncodedImage) def test_jxl_roundtrip_is_lossless_gray16() -> None: data = np.arange(100 * 80, dtype=np.uint16).reshape(100, 80) - src = Image(data=data, format=ImageFormat.GRAY16, frame_id="d", ts=1.0) - ci = CompressedImage.from_image(src, format="jxl") - assert ci.format == "jxl" - img = ci.decode() - assert img.format == ImageFormat.GRAY16 - assert np.array_equal(img.data, data) + src = image_from_array( + data, encoding="mono16", header=Header(frame_id="d", stamp=time_from_seconds(1.0)) + ) + ci = compressed_image_from_image(src, format="jxl") + assert "; jxl compressed " in ci.format + img = image_from_compressed(ci) + assert img.encoding == "mono16" + assert np.array_equal(image_view(img), data) def test_jxl_roundtrip_is_lossless_float_depth() -> None: data = np.linspace(0.1, 10.0, 100 * 80, dtype=np.float32).reshape(100, 80) - src = Image(data=data, format=ImageFormat.DEPTH, frame_id="d", ts=2.0) - ci = CompressedImage.from_image(src, format="jxl") - img = ci.decode() - assert img.format == ImageFormat.DEPTH - assert np.array_equal(img.data, data) - assert img.ts == src.ts - assert img.frame_id == "d" + src = image_from_array( + data, encoding="32FC1", header=Header(frame_id="d", stamp=time_from_seconds(2.0)) + ) + ci = compressed_image_from_image(src, format="jxl") + img = image_from_compressed(ci) + assert img.encoding == "32FC1" + assert np.array_equal(image_view(img), data) + assert to_nanoseconds(img.header.stamp) == to_nanoseconds(src.header.stamp) + assert img.header.frame_id == "d" def test_jxl_lossy_rgb_honors_quality(rgb_image) -> None: - ci = CompressedImage.from_image(rgb_image, format="jxl", quality=90) - assert 0 < len(ci.data) < rgb_image.data.nbytes // 4 - img = ci.decode() - assert img.format == ImageFormat.RGB - diff = np.abs(img.data.astype(int) - rgb_image.data.astype(int)).mean() + ci = compressed_image_from_image(rgb_image, format="jxl", quality=90) + assert 0 < len(ci.data) < image_view(rgb_image).nbytes // 4 + img = image_from_compressed(ci) + assert img.encoding in ("rgb8", "bgr8") + diff = np.abs(image_to_rgb(img).astype(int) - image_view(rgb_image).astype(int)).mean() assert diff < 5, f"JXL q90 mean pixel error too high: {diff}" def test_jxl_effort_trades_cpu_for_size(rgb_image) -> None: - fast = CompressedImage.from_image(rgb_image, format="jxl", effort=1) - thorough = CompressedImage.from_image(rgb_image, format="jxl", effort=7) + fast = compressed_image_from_image(rgb_image, format="jxl", effort=1) + thorough = compressed_image_from_image(rgb_image, format="jxl", effort=7) assert len(thorough.data) <= len(fast.data) def test_jxl_rejects_float64() -> None: - depth = Image(data=np.zeros((10, 10), dtype=np.float64), format=ImageFormat.DEPTH) + depth = image_from_array(np.zeros((10, 10), dtype=np.float64), encoding="64FC1") with pytest.raises(ValueError, match="JXL cannot encode"): - CompressedImage.from_image(depth, format="jxl") + compressed_image_from_image(depth, format="jxl") + + +@pytest.mark.parametrize("codec", ["png", "jxl"]) +def test_compressed_uint16_depth_keeps_units_and_rerun_depth(codec: str) -> None: + import rerun as rr + + pixels = np.arange(80, dtype=np.uint16).reshape(10, 8) + source = image_from_array(pixels, encoding="16UC1", header=Header(frame_id="depth")) + compressed = compressed_image_from_image(source, format=codec) + decoded = image_from_compressed(compressed) + assert decoded.encoding == "16UC1" + assert decoded.header == source.header + np.testing.assert_array_equal(image_view(decoded), pixels) + if codec == "jxl": + assert isinstance(image_archetype(compressed), rr.DepthImage) + + +@pytest.mark.parametrize("codec", ["png", "jxl"]) +def test_compressed_big_endian_depth(codec: str) -> None: + pixels = np.arange(80, dtype=">u2").reshape(10, 8) * 257 + source = image_from_array(pixels.astype(">u2"), encoding="16UC1") + decoded = image_from_compressed(compressed_image_from_image(source, format=codec)) + np.testing.assert_array_equal(image_view(decoded), pixels) diff --git a/dimos/msgs/sensor_msgs/test_Joy.py b/dimos/msgs/sensor_msgs/test_Joy.py index 499c2bb860f..d49c81bee90 100644 --- a/dimos/msgs/sensor_msgs/test_Joy.py +++ b/dimos/msgs/sensor_msgs/test_Joy.py @@ -14,175 +14,133 @@ # limitations under the License. -from dimos.msgs.sensor_msgs.Joy import Joy +from dimos_generated.sensor_msgs.msg import Joy +from dimos_generated.std_msgs.msg import Header +import pytest +from rosbags.typesys import Stores, get_typestore +from dimos.msgs.time import time_from_nanoseconds, time_from_seconds, to_nanoseconds -def test_lcm_encode_decode() -> None: - """Test LCM encode/decode preserves Joy data.""" - print("Testing Joy LCM encode/decode...") - # Create test joy message with sample gamepad data +def test_cdr_encode_decode() -> None: + """Test CDR encode/decode preserves Joy data.""" + print("Testing Joy CDR encode/decode...") original = Joy( - ts=1234567890.123456789, - frame_id="gamepad", - axes=[0.5, -0.25, 1.0, -1.0, 0.0, 0.75], # 6 axes (e.g., left/right sticks + triggers) - buttons=[1, 0, 0, 1, 1, 0, 0, 0, 1, 0, 0, 0], # 12 buttons + axes=[0.5, -0.25, 1.0, -1.0, 0.0, 0.75], + buttons=[1, 0, 0, 1, 1, 0, 0, 0, 1, 0, 0, 0], + header=Header(stamp=time_from_seconds(1234567890.1234567), frame_id="gamepad"), ) - - # Encode to LCM bytes - encoded = original.lcm_encode() + encoded = original.encode() assert isinstance(encoded, bytes) assert len(encoded) > 0 - - # Decode back - decoded = Joy.lcm_decode(encoded) - - # Verify all fields match - assert abs(decoded.ts - original.ts) < 1e-9 - assert decoded.frame_id == original.frame_id + decoded = Joy.decode(encoded) + assert abs(to_nanoseconds(decoded.header.stamp) - to_nanoseconds(original.header.stamp)) < 1e-09 + assert decoded.header.frame_id == original.header.frame_id assert decoded.axes == original.axes assert decoded.buttons == original.buttons - - print("✓ Joy LCM encode/decode test passed") + print("✓ Joy CDR encode/decode test passed") def test_initialization_methods() -> None: """Test various initialization methods for Joy.""" print("Testing Joy initialization methods...") - - # Test default initialization joy1 = Joy() assert joy1.axes == [] assert joy1.buttons == [] - assert joy1.frame_id == "" - assert joy1.ts > 0 # Should have current time - - # Test full initialization - joy2 = Joy(ts=1234567890.0, frame_id="xbox_controller", axes=[0.1, 0.2, 0.3], buttons=[1, 0, 1]) - assert joy2.ts == 1234567890.0 - assert joy2.frame_id == "xbox_controller" - assert joy2.axes == [0.1, 0.2, 0.3] + assert joy1.header.frame_id == "" + assert to_nanoseconds(joy1.header.stamp) == 0 + joy2 = Joy( + axes=[0.1, 0.2, 0.3], + buttons=[1, 0, 1], + header=Header(stamp=time_from_seconds(1234567890.0), frame_id="xbox_controller"), + ) + assert to_nanoseconds(joy2.header.stamp) == 1234567890000000000 + assert joy2.header.frame_id == "xbox_controller" + assert list(joy2.axes) == pytest.approx([0.1, 0.2, 0.3], abs=1e-6) assert joy2.buttons == [1, 0, 1] - - # Test tuple initialization - joy3 = Joy(([0.5, -0.5], [1, 1, 0])) + joy3 = Joy(axes=[0.5, -0.5], buttons=[1, 1, 0]) assert joy3.axes == [0.5, -0.5] assert joy3.buttons == [1, 1, 0] - - # Test dict initialization - joy4 = Joy({"axes": [0.7, 0.8], "buttons": [0, 1], "frame_id": "ps4_controller"}) - assert joy4.axes == [0.7, 0.8] + joy4 = Joy(axes=[0.7, 0.8], buttons=[0, 1], header=Header(frame_id="ps4_controller")) + assert list(joy4.axes) == pytest.approx([0.7, 0.8], abs=1e-6) assert joy4.buttons == [0, 1] - assert joy4.frame_id == "ps4_controller" - - # Test copy constructor - joy5 = Joy(joy2) - assert joy5.ts == joy2.ts - assert joy5.frame_id == joy2.frame_id + assert joy4.header.frame_id == "ps4_controller" + joy5 = Joy.decode(joy2.encode()) + assert to_nanoseconds(joy5.header.stamp) == to_nanoseconds(joy2.header.stamp) + assert joy5.header.frame_id == joy2.header.frame_id assert joy5.axes == joy2.axes assert joy5.buttons == joy2.buttons - assert joy5 is not joy2 # Different objects - + assert joy5 is not joy2 print("✓ Joy initialization methods test passed") def test_equality() -> None: """Test Joy equality comparison.""" print("Testing Joy equality...") - - joy1 = Joy(ts=1000.0, frame_id="controller1", axes=[0.5, -0.5], buttons=[1, 0, 1]) - - joy2 = Joy(ts=1000.0, frame_id="controller1", axes=[0.5, -0.5], buttons=[1, 0, 1]) - + joy1 = Joy( + axes=[0.5, -0.5], + buttons=[1, 0, 1], + header=Header(stamp=time_from_seconds(1000.0), frame_id="controller1"), + ) + joy2 = Joy( + axes=[0.5, -0.5], + buttons=[1, 0, 1], + header=Header(stamp=time_from_seconds(1000.0), frame_id="controller1"), + ) joy3 = Joy( - ts=1000.0, - frame_id="controller2", # Different frame_id axes=[0.5, -0.5], buttons=[1, 0, 1], + header=Header(stamp=time_from_seconds(1000.0), frame_id="controller2"), ) - joy4 = Joy( - ts=1000.0, - frame_id="controller1", - axes=[0.6, -0.5], # Different axes + axes=[0.6, -0.5], buttons=[1, 0, 1], + header=Header(stamp=time_from_seconds(1000.0), frame_id="controller1"), ) - - # Same content should be equal assert joy1 == joy2 - - # Different frame_id should not be equal assert joy1 != joy3 - - # Different axes should not be equal assert joy1 != joy4 - - # Different type should not be equal assert joy1 != "not a joy" assert joy1 != 42 - print("✓ Joy equality test passed") -def test_string_representation() -> None: - """Test Joy string representations.""" - print("Testing Joy string representations...") - +def test_independent_schema_decoding() -> None: joy = Joy( - ts=1234567890.123, - frame_id="test_controller", + header=Header(stamp=time_from_nanoseconds(1234567890123456789), frame_id="test_controller"), axes=[0.1, -0.2, 0.3, 0.4], buttons=[1, 0, 1, 0, 0, 1], ) - - # Test __str__ - str_repr = str(joy) - assert "Joy" in str_repr - assert "axes=4 values" in str_repr - assert "buttons=6 values" in str_repr - assert "test_controller" in str_repr - - # Test __repr__ - repr_str = repr(joy) - assert "Joy" in repr_str - assert "1234567890.123" in repr_str - assert "test_controller" in repr_str - assert "[0.1, -0.2, 0.3, 0.4]" in repr_str - assert "[1, 0, 1, 0, 0, 1]" in repr_str - - print("✓ Joy string representation test passed") + decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(joy.encode(), Joy.msg_name) + assert decoded.header.frame_id == "test_controller" + assert decoded.header.stamp.sec == 1234567890 + assert decoded.header.stamp.nanosec == 123456789 + assert list(decoded.axes) == pytest.approx([0.1, -0.2, 0.3, 0.4], abs=1e-6) + assert list(decoded.buttons) == [1, 0, 1, 0, 0, 1] def test_edge_cases() -> None: """Test Joy with edge cases.""" print("Testing Joy edge cases...") - - # Empty axes and buttons joy1 = Joy(axes=[], buttons=[]) assert joy1.axes == [] assert joy1.buttons == [] - encoded = joy1.lcm_encode() - decoded = Joy.lcm_decode(encoded) + encoded = joy1.encode() + decoded = Joy.decode(encoded) assert decoded.axes == [] assert decoded.buttons == [] - - # Large number of axes and buttons many_axes = [float(i) / 100.0 for i in range(20)] many_buttons = [i % 2 for i in range(32)] joy2 = Joy(axes=many_axes, buttons=many_buttons) assert len(joy2.axes) == 20 assert len(joy2.buttons) == 32 - encoded = joy2.lcm_encode() - decoded = Joy.lcm_decode(encoded) - # Check axes with floating point tolerance + encoded = joy2.encode() + decoded = Joy.decode(encoded) assert len(decoded.axes) == len(many_axes) for i, (a, b) in enumerate(zip(decoded.axes, many_axes, strict=False)): - assert abs(a - b) < 1e-6, f"Axis {i}: {a} != {b}" + assert abs(a - b) < 1e-06, f"Axis {i}: {a} != {b}" assert decoded.buttons == many_buttons - - # Extreme axis values extreme_axes = [-1.0, 1.0, 0.0, -0.999999, 0.999999] joy3 = Joy(axes=extreme_axes) - assert joy3.axes == extreme_axes - + assert list(joy3.axes) == pytest.approx(extreme_axes, abs=1e-6) print("✓ Joy edge cases test passed") diff --git a/dimos/msgs/sensor_msgs/test_image.py b/dimos/msgs/sensor_msgs/test_image.py index 21590c5e658..be1cca8a22b 100644 --- a/dimos/msgs/sensor_msgs/test_image.py +++ b/dimos/msgs/sensor_msgs/test_image.py @@ -12,85 +12,90 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest.mock import MagicMock +"""File, color, CDR and reactive selection checks on generated images.""" +from pathlib import Path + +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header import numpy as np +from PIL import Image as PILImage import pytest from reactivex.testing import ReactiveTest, TestScheduler -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat, sharpness_barrier -from dimos.utils.data import get_data +from dimos.msgs.image import ( + image_from_array, + image_from_file, + image_sharpness, + image_to_bgr, + image_to_rgb, + image_view, +) +from dimos.msgs.time import time_from_nanoseconds, to_nanoseconds +from dimos.utils.reactive import quality_barrier @pytest.fixture -def img(): - image_file_path = get_data("cafe.jpg") - return Image.from_file(str(image_file_path)) +def img(tmp_path: Path) -> Image: + pixels = np.zeros((771, 1024, 3), dtype=np.uint8) + pixels[..., 0] = 201 + pixels[..., 1] = 37 + path = tmp_path / "camera.png" + PILImage.fromarray(pixels).save(path) + return image_from_file(path, header=Header(stamp=time_from_nanoseconds(1234567890123456789))) def test_file_load(img: Image) -> None: - assert isinstance(img.data, np.ndarray) - assert img.width == 1024 - assert img.height == 771 - assert img.channels == 3 - assert img.shape == (771, 1024, 3) - assert img.data.dtype == np.uint8 - assert img.format == ImageFormat.BGR - assert img.frame_id == "" - assert isinstance(img.ts, float) - assert img.ts > 0 - assert img.data.flags["C_CONTIGUOUS"] - + pixels = image_view(img) + assert isinstance(pixels, np.ndarray) + assert (img.width, img.height) == (1024, 771) + assert pixels.shape == (771, 1024, 3) + assert pixels.dtype == np.uint8 + assert img.encoding == "rgb8" + assert img.header.frame_id == "" + assert to_nanoseconds(img.header.stamp) == 1234567890123456789 + assert pixels.flags["C_CONTIGUOUS"] + np.testing.assert_array_equal(pixels[0, 0], [201, 37, 0]) -def test_lcm_encode_decode(img: Image) -> None: - binary_msg = img.lcm_encode() - decoded_img = Image.lcm_decode(binary_msg) - assert isinstance(decoded_img, Image) - assert decoded_img is not img - assert decoded_img == img +def test_cdr_encode_decode(img: Image) -> None: + decoded = Image.decode(img.encode()) + assert decoded is not img + assert decoded == img def test_rgb_bgr_conversion(img: Image) -> None: - rgb = img.to_rgb() - assert not rgb == img - assert rgb.to_bgr() == img + bgr = image_from_array(image_to_bgr(img), encoding="bgr8", header=img.header) + assert bgr != img + restored = image_from_array(image_to_rgb(bgr), encoding="rgb8", header=img.header) + assert restored == img def test_opencv_conversion(img: Image) -> None: - ocv = img.to_opencv() - decoded_img = Image.from_opencv(ocv) - - # artificially patch timestamp - decoded_img.ts = img.ts - assert decoded_img == img + pixels = image_to_bgr(img) + generated = image_from_array(pixels, encoding="bgr8", header=img.header) + np.testing.assert_array_equal(image_to_rgb(generated), image_view(img)) + assert generated.header == img.header def test_sharpness_barrier() -> None: - # Mock images with known sharpness values, avoiding real data from disk - sharpness_values = [0.3711, 0.3241, 0.3067, 0.2583, 0.3665] - mock_images = [] - for sharp in sharpness_values: - img = MagicMock() - img.sharpness = sharp - mock_images.append(img) - - # sharpness_barrier(20) -> 0.05s windows in virtual time. Subscription is at - # t=200, so windows close at 200.05, 200.10, ... Items 1-4 land in the first - # window, item 5 in the second. + # Real generated images: the first and fifth win the same two virtual windows. + pattern = (np.indices((8, 8)).sum(axis=0) % 2).astype(np.uint8) + images = [ + image_from_array(pattern * level, encoding="mono8") for level in (255, 200, 150, 50, 230) + ] + assert image_sharpness(images[0]) > max(image_sharpness(value) for value in images[1:4]) scheduler = TestScheduler() source = scheduler.create_hot_observable( - ReactiveTest.on_next(200.01, mock_images[0]), - ReactiveTest.on_next(200.02, mock_images[1]), - ReactiveTest.on_next(200.03, mock_images[2]), - ReactiveTest.on_next(200.04, mock_images[3]), - ReactiveTest.on_next(200.06, mock_images[4]), + ReactiveTest.on_next(200.01, images[0]), + ReactiveTest.on_next(200.02, images[1]), + ReactiveTest.on_next(200.03, images[2]), + ReactiveTest.on_next(200.04, images[3]), + ReactiveTest.on_next(200.06, images[4]), ReactiveTest.on_completed(200.08), ) - - results = scheduler.start(lambda: source.pipe(sharpness_barrier(20, scheduler=scheduler))) - - emitted = [m.value.value for m in results.messages if m.value.kind == "N"] - assert len(emitted) == 2, f"Expected one emission per window, got {len(emitted)}" - assert emitted[0].sharpness == 0.3711 # Sharpest of the 4 in the first window - assert emitted[1].sharpness == 0.3665 # Only item in the second window + results = scheduler.start(lambda: source.pipe(quality_barrier(image_sharpness, 20, scheduler))) + emitted = [message.value.value for message in results.messages if message.value.kind == "N"] + assert len(emitted) == 2 + assert emitted[0] is images[0] + assert emitted[1] is images[4] diff --git a/dimos/msgs/sensor_msgs/tool_image.py b/dimos/msgs/sensor_msgs/tool_image.py index af7f04073e4..1f81e8591bc 100644 --- a/dimos/msgs/sensor_msgs/tool_image.py +++ b/dimos/msgs/sensor_msgs/tool_image.py @@ -12,20 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos.msgs.sensor_msgs.Image import Image +from typing import Any + +from dimos.msgs.image import image_from_array, image_sharpness from dimos.utils.data import get_data from dimos.utils.testing.legacy_pickle import LegacyPickleStore def test_sharpness_stream() -> None: get_data("unitree_office_walk") # Preload data for testing - video_store = LegacyPickleStore( - "unitree_office_walk/video", autocast=lambda x: Image.from_numpy(x).to_rgb() - ) + video_store: LegacyPickleStore[Any] = LegacyPickleStore("unitree_office_walk/video") cnt = 0 - for image in video_store.iterate(): + for pixels in video_store.iterate(): + image = image_from_array(pixels, encoding="bgr8") cnt = cnt + 1 - print(image.sharpness) + print(image_sharpness(image)) if cnt > 30: return diff --git a/dimos/msgs/std_msgs/test_header.py b/dimos/msgs/std_msgs/test_header.py index 29f4ee2c0e9..a1849aae5fd 100644 --- a/dimos/msgs/std_msgs/test_header.py +++ b/dimos/msgs/std_msgs/test_header.py @@ -12,87 +12,51 @@ # See the License for the specific language governing permissions and # limitations under the License. -from datetime import datetime -import time - -from dimos.msgs.std_msgs.Header import Header - - -def test_header_initialization_methods() -> None: - """Test various ways to initialize a Header.""" - - # Method 1: With timestamp and frame_id - header1 = Header(123.456, "world") - assert header1.seq == 1 - assert header1.stamp.sec == 123 - assert header1.stamp.nsec == 456000000 - assert header1.frame_id == "world" - - # Method 2: With just frame_id (uses current time) - header2 = Header("base_link") - assert header2.seq == 1 - assert header2.frame_id == "base_link" - # Timestamp should be close to current time - assert abs(header2.timestamp - time.time()) < 0.1 - - # Method 3: Empty header (current time, empty frame_id) - header3 = Header() - assert header3.seq == 0 - assert header3.frame_id == "" +"""ROS2 headers use explicit time conversion; no legacy sequence counter.""" - # Method 4: With datetime object - dt = datetime(2025, 1, 18, 12, 30, 45, 500000) # 500ms - header4 = Header(dt, "sensor") - assert header4.seq == 1 - assert header4.frame_id == "sensor" - expected_timestamp = dt.timestamp() - assert abs(header4.timestamp - expected_timestamp) < 1e-6 - - # Method 5: With custom seq number - header5 = Header(999.123, "custom", seq=42) - assert header5.seq == 42 - assert header5.stamp.sec == 999 - assert header5.stamp.nsec == 123000000 - assert header5.frame_id == "custom" - - # Method 6: Using now() class method - header6 = Header.now("camera") - assert header6.seq == 1 - assert header6.frame_id == "camera" - assert abs(header6.timestamp - time.time()) < 0.1 - - # Method 7: now() with custom seq - header7 = Header.now("lidar", seq=99) - assert header7.seq == 99 - assert header7.frame_id == "lidar" +from datetime import datetime, timezone +import time +from dimos_generated.std_msgs.msg import Header +from rosbags.typesys import Stores, get_typestore -def test_header_properties() -> None: - """Test Header property accessors.""" - header = Header(1234567890.123456789, "test") +from dimos.msgs.time import header_now, time_from_nanoseconds, time_from_seconds, to_seconds - # Test timestamp property - assert abs(header.timestamp - 1234567890.123456789) < 1e-6 - # Test datetime property - dt = header.datetime - assert isinstance(dt, datetime) +def test_header_initialization_methods() -> None: + header = Header(stamp=time_from_seconds(123.456), frame_id="world") + assert header.stamp.sec == 123 + assert header.stamp.nanosec == 456000000 + assert header.frame_id == "world" + before = time.time_ns() + current = header_now("base_link") + after = time.time_ns() + assert before <= current.stamp.sec * 1000000000 + current.stamp.nanosec <= after + assert current.frame_id == "base_link" + empty = Header() + assert empty.stamp.sec == empty.stamp.nanosec == 0 + assert empty.frame_id == "" + dt = datetime(2025, 1, 18, 12, 30, 45, 500000, tzinfo=timezone.utc) + dated = Header(stamp=time_from_seconds(dt.timestamp()), frame_id="sensor") + assert dated.frame_id == "sensor" + assert abs(to_seconds(dated.stamp) - dt.timestamp()) < 1e-6 + custom = Header(stamp=time_from_seconds(999.123), frame_id="custom") + assert custom.stamp.sec == 999 + assert custom.stamp.nanosec == 123000000 + assert custom.frame_id == "custom" + + +def test_header_datetime_conversion() -> None: + header = Header(stamp=time_from_nanoseconds(1234567890123456789), frame_id="test") + dt = datetime.fromtimestamp(to_seconds(header.stamp), tz=timezone.utc) assert abs(dt.timestamp() - 1234567890.123456789) < 1e-6 + assert header.stamp.nanosec == 123456789 -def test_header_string_representation() -> None: - """Test Header string representations.""" - header = Header(100.5, "map", seq=10) - - # Test __str__ - str_repr = str(header) - assert "seq=10" in str_repr - assert "time=100.5" in str_repr - assert "frame_id='map'" in str_repr - - # Test __repr__ - repr_str = repr(header) - assert "Header(" in repr_str - assert "seq=10" in repr_str - assert "Time(sec=100, nsec=500000000)" in repr_str - assert "frame_id='map'" in repr_str +def test_header_independent_cdr_fields() -> None: + header = Header(stamp=time_from_seconds(100.5), frame_id="map") + decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(header.encode(), Header.msg_name) + assert decoded.stamp.sec == 100 + assert decoded.stamp.nanosec == 500000000 + assert decoded.frame_id == "map" + assert "seq" not in Header.schema diff --git a/dimos/msgs/test_camera_info.py b/dimos/msgs/test_camera_info.py new file mode 100644 index 00000000000..641f71e3220 --- /dev/null +++ b/dimos/msgs/test_camera_info.py @@ -0,0 +1,127 @@ +# 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. + +"""Calibration helpers preserve generated ROS fields and reject malformed input.""" + +import math +from pathlib import Path + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.sensor_msgs.msg import CameraInfo +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest +import yaml + +from dimos.msgs.camera_info import ( + camera_info_from_fov, + camera_info_from_intrinsics, + camera_info_from_yaml, + camera_info_with_stamp, + intrinsic_matrix, +) + +_CALIBRATION = Path(__file__).parents[1] / "robot/unitree/go2/front_camera_720.yaml" + + +def test_real_go2_calibration_survives_cdr_with_exact_header(): + header = Header(frame_id="camera_optical", stamp=Time(sec=1700000000, nanosec=123456789)) + source = yaml.safe_load(_CALIBRATION.read_text()) + message = camera_info_from_yaml(_CALIBRATION, header=header) + decoded = CameraInfo.decode(message.encode()) + assert decoded.header == header + assert (decoded.width, decoded.height) == (1280, 720) + assert decoded.distortion_model == "equidistant" + for field, key in [ + ("k", "camera_matrix"), + ("d", "distortion_coefficients"), + ("r", "rectification_matrix"), + ("p", "projection_matrix"), + ]: + np.testing.assert_array_equal(getattr(decoded, field), source[key]["data"]) + matrix = intrinsic_matrix(decoded) + assert matrix.shape == (3, 3) + matrix[0, 0] = 0 + assert decoded.k[0] == source["camera_matrix"]["data"][0] + decoded.header.frame_id = "changed" + assert header.frame_id == "camera_optical" + + +@pytest.mark.parametrize("axis,focal", [("vertical", 240), ("horizontal", 320)]) +def test_fov_converts_selected_axis_and_builds_projection(axis, focal): + message = camera_info_from_fov(90, 640, 480, header=Header(frame_id="optical"), axis=axis) + np.testing.assert_allclose(message.k, [focal, 0, 320, 0, focal, 240, 0, 0, 1]) + np.testing.assert_allclose(message.p, [focal, 0, 320, 0, 0, focal, 240, 0, 0, 0, 1, 0]) + np.testing.assert_array_equal(message.r, np.eye(3).reshape(-1)) + assert message.header.frame_id == "optical" + + +@pytest.mark.parametrize("fov", [0, -1, 180, 181, math.nan, math.inf]) +def test_fov_rejects_nonphysical_angles(fov): + with pytest.raises(ValueError, match="Field of view"): + camera_info_from_fov(fov, 640, 480, header=Header()) + + +@pytest.mark.parametrize("width,height", [(0, 480), (640, -1), (True, 480), (1.5, 480)]) +def test_intrinsics_reject_invalid_dimensions(width, height): + with pytest.raises(ValueError, match="positive integers"): + camera_info_from_intrinsics(300, 300, 320, 240, width, height, header=Header()) + + +@pytest.mark.parametrize("fx", [0, -1, math.nan, math.inf]) +def test_intrinsics_reject_invalid_focal_length(fx): + with pytest.raises(ValueError, match="positive focal"): + camera_info_from_intrinsics(fx, 300, 320, 240, 640, 480, header=Header()) + + +@pytest.mark.parametrize("mutation", ["rows", "length", "nonfinite"]) +def test_yaml_rejects_invalid_camera_matrix(tmp_path, mutation): + data = yaml.safe_load(_CALIBRATION.read_text()) + matrix = data["camera_matrix"] + if mutation == "rows": + matrix["rows"] = 1 + elif mutation == "length": + matrix["data"].pop() + else: + matrix["data"][0] = math.nan + path = tmp_path / "invalid.yaml" + path.write_text(yaml.safe_dump(data)) + with pytest.raises(ValueError, match="camera_matrix"): + camera_info_from_yaml(path, header=Header()) + + +def test_stamp_copy_preserves_calibration_and_does_not_mutate_prior_publish(): + original = camera_info_from_intrinsics( + 500, + 510, + 320, + 240, + 640, + 480, + header=Header(frame_id="optical", stamp=Time(sec=1, nanosec=987654321)), + ) + original.binning_x = 2 + original.roi.x_offset = 12 + stamp = Time(sec=2, nanosec=123456789) + copied = camera_info_with_stamp(original, stamp) + assert copied.header.stamp == stamp + assert copied.header.frame_id == original.header.frame_id + assert list(copied.k) == list(original.k) + assert copied.binning_x == 2 and copied.roi.x_offset == 12 + copied.k[0] = 100 + copied.roi.x_offset = 24 + stamp.nanosec = 0 + assert original.k[0] == 500 and original.roi.x_offset == 12 + assert original.header.stamp == Time(sec=1, nanosec=987654321) + assert copied.header.stamp.nanosec == 123456789 diff --git a/dimos/msgs/test_generated_projection.py b/dimos/msgs/test_generated_projection.py new file mode 100644 index 00000000000..3c200468093 --- /dev/null +++ b/dimos/msgs/test_generated_projection.py @@ -0,0 +1,98 @@ +# 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. + +"""Stored image projection uses generated transforms and padded depth pixels.""" + +from dimos_generated.geometry_msgs.msg import ( + Point, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.std_msgs.msg import Header +import numpy as np + +from dimos.memory.type.observation import Observation +from dimos.msgs.geometry import transform_matrix +from dimos.perception.detection.project import _world_to_optical, sees + + +def test_stored_pose_projection_inverts_world_camera_chain() -> None: + obs = Observation( + id=0, ts=0.0, data_type=Image, _data=Image(), pose_tuple=(2.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0) + ) + extrinsic = TransformStamped( + header=Header(frame_id="base_link"), + child_frame_id="camera_optical", + transform=Transform(translation=Vector3(x=1), rotation=Quaternion(w=1)), + ) + result = _world_to_optical(obs, "world", None, extrinsic, "camera_optical", 5.0) + assert result is not None + assert result.header.frame_id == "camera_optical" + assert result.child_frame_id == "world" + np.testing.assert_allclose(transform_matrix(result.transform) @ [3, 0, 2, 1], [0, 0, 2, 1]) + + +def test_depth_occlusion_reads_big_endian_rows_and_millimeters() -> None: + obs = Observation( + id=0, ts=0.0, data_type=Image, _data=Image(), pose_tuple=(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0) + ) + extrinsic = TransformStamped( + header=Header(frame_id="base_link"), + child_frame_id="camera_optical", + transform=Transform(rotation=Quaternion(w=1)), + ) + calibration = CameraInfo(width=2, height=2, k=[1, 0, 0, 0, 1, 0, 0, 0, 1]) + # Pixel (0, 1) = 1500mm; padding cannot be interpreted as pixels. + depth = Image( + width=2, + height=2, + step=6, + encoding="16UC1", + is_bigendian=1, + data=[0, 0, 0, 0, 99, 99, 5, 220, 0, 0, 99, 99], + ) + visible = sees( + Point(x=0, y=1, z=1), calibration, base_to_optical=extrinsic, depth=lambda _: depth + ) + hidden = sees( + Point(x=0, y=2, z=2), calibration, base_to_optical=extrinsic, depth=lambda _: depth + ) + assert visible(obs) + assert not hidden(obs) + + +def test_timestamp_adapter_matches_generated_cloud_without_mutating_header() -> None: + from dimos_generated.sensor_msgs.msg import PointCloud2 + from reactivex.subject import Subject + + from dimos.msgs.time import time_from_nanoseconds, to_seconds + from dimos.types.timestamped import TimestampedData, align_timestamped + + source = PointCloud2(header=Header(stamp=time_from_nanoseconds(1700000000123456789))) + primary = Subject() + secondary = Subject() + received = [] + disposable = align_timestamped(primary, secondary, match_tolerance=0.0).subscribe( + received.append + ) + timed = TimestampedData(source, to_seconds(source.header.stamp)) + secondary.on_next(timed) + primary.on_next(TimestampedData("detection", timed.ts)) + assert len(received) == 1 + assert received[0][1].value is source + assert source.header.stamp.nanosec == 123456789 + disposable.dispose() diff --git a/dimos/msgs/test_generated_views.py b/dimos/msgs/test_generated_views.py new file mode 100644 index 00000000000..54f600a24ee --- /dev/null +++ b/dimos/msgs/test_generated_views.py @@ -0,0 +1,232 @@ +# 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. + +"""Borrowed buffers preserve padding, byte order, lifetime, and ROS value semantics.""" + +import gc +import json +import math + +import cv2 +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid, Path +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header +import numpy as np +from PIL import Image as PILImage +import pytest + +from dimos.msgs.geometry import yaw +from dimos.msgs.image import ( + image_brightness, + image_from_array, + image_from_file, + image_sharpness, + image_to_jpeg, + image_to_rgb, + image_view, +) +from dimos.msgs.occupancy import block_max_reduce, occupancy_view +from dimos.msgs.time import time_from_nanoseconds +from dimos.web.relay_bridge.builtin_codecs import decode_point, encode_path, encode_pose + + +def test_padded_image_view_retains_owner_and_is_readonly() -> None: + msg = Image(width=2, height=2, step=4, encoding="mono8", data=[1, 2, 99, 99, 3, 4, 99, 99]) + pixels = image_view(msg) + np.testing.assert_array_equal(pixels, [[1, 2], [3, 4]]) + assert pixels.strides == (4, 1) + with pytest.raises(ValueError): + pixels[0, 0] = 9 + del msg + gc.collect() + np.testing.assert_array_equal(pixels, [[1, 2], [3, 4]]) + copied = pixels.copy() + copied[0, 0] = 9 + assert pixels[0, 0] == 1 + + +def test_big_endian_depth_view() -> None: + msg = Image( + width=2, height=1, step=4, encoding="16UC1", is_bigendian=1, data=[0x01, 0x02, 0x03, 0x04] + ) + np.testing.assert_array_equal(image_view(msg), [[258, 772]]) + + +def test_image_array_copy_preserves_endian_and_handles_strides() -> None: + pixels = np.arange(24, dtype=">u2").reshape(4, 6)[::2, ::2] + msg = image_from_array(pixels, encoding="16UC1") + assert (msg.width, msg.height, msg.step, msg.is_bigendian) == (3, 2, 6, 1) + expected = pixels.copy() + pixels[:] = 0 + np.testing.assert_array_equal(image_view(msg), expected) + + +@pytest.mark.parametrize( + "pixels,encoding", + [(np.zeros((2, 2), dtype=np.float32), "mono8"), (np.zeros((2, 2), dtype=np.uint8), "rgb8")], +) +def test_image_array_requires_matching_encoding(pixels: np.ndarray, encoding: str) -> None: + with pytest.raises(ValueError, match="does not match"): + image_from_array(pixels, encoding=encoding) + + +@pytest.mark.parametrize("step,data", [(1, [1]), (2, [1]), (2, [1, 2, 3])]) +def test_invalid_image_layout(step: int, data: list[int]) -> None: + with pytest.raises(ValueError, match="dimensions"): + image_view(Image(width=2, height=1, step=step, data=data, encoding="mono8")) + + +def test_jpeg_preserves_rgb_color() -> None: + msg = Image(width=16, height=16, step=48, encoding="rgb8", data=[255, 0, 0] * 256) + jpeg = image_to_jpeg(msg, quality=95) + bgr = cv2.imdecode(np.frombuffer(jpeg, np.uint8), cv2.IMREAD_COLOR) + assert bgr.shape == (16, 16, 3) + assert bgr[8, 8, 2] > 250 and bgr[8, 8, 0] < 5 + + +def test_occupancy_view_and_obstacle_preserving_reduction() -> None: + msg = OccupancyGrid(info=MapMetaData(width=4, height=2), data=[-1, -1, 0, 100, -1, -1, 10, 20]) + cells = occupancy_view(msg) + assert not cells.flags.writeable + np.testing.assert_array_equal(block_max_reduce(cells, 2), [[-1, 100]]) + msg.info.width = 3 + with pytest.raises(ValueError, match="dimensions"): + occupancy_view(msg) + + +def test_web_pose_path_and_goal_use_nested_ros_fields() -> None: + pose = PoseStamped( + header=Header(stamp=time_from_nanoseconds(1500000000)), + pose=Pose( + position=Point(x=2, y=3), orientation=Quaternion(z=math.sqrt(0.5), w=math.sqrt(0.5)) + ), + ) + assert yaw(pose.pose.orientation) == pytest.approx(math.pi / 2) + value = json.loads(encode_pose(pose)) + assert value == {"x": 2, "y": 3, "z": 0, "yaw": pytest.approx(math.pi / 2), "ts": 1.5} + assert json.loads(encode_path(Path(poses=[pose]))) == [[2, 3]] + goal = decode_point({"x": 4, "y": 5}) + assert goal.point == Point(x=4, y=5) and goal.header.frame_id == "world" + + +@pytest.mark.parametrize("encoding", ["rgb8", "bgr8", "rgba8", "bgra8", "mono8"]) +def test_sharpness_prefers_edges_to_blurred_pixels(encoding): + mono = ((np.indices((64, 64)).sum(axis=0) % 2) * 255).astype(np.uint8) + channels = 1 if encoding == "mono8" else (4 if "a" in encoding else 3) + pixels = mono if channels == 1 else np.repeat(mono[:, :, None], channels, axis=2) + blurred = cv2.GaussianBlur(pixels, (5, 5), 0) + sharp = image_from_array(pixels, encoding=encoding) + soft = image_from_array(blurred, encoding=encoding) + assert image_sharpness(sharp) > image_sharpness(soft) + + +def test_sharpness_rejects_depth_and_empty_images(): + depth = image_from_array(np.zeros((2, 2), dtype=np.float32), encoding="32FC1") + with pytest.raises(ValueError, match="8-bit visual"): + image_sharpness(depth) + with pytest.raises(ValueError, match="nonempty"): + image_sharpness(Image(encoding="mono8")) + + +@pytest.mark.parametrize( + "encoding,pixel", + [ + ("rgb8", [255, 0, 17]), + ("bgr8", [17, 0, 255]), + ("rgba8", [255, 0, 17, 23]), + ("bgra8", [17, 0, 255, 23]), + ], +) +def test_rgb_model_input_is_color_correct_and_independent(encoding: str, pixel: list[int]) -> None: + source = image_from_array(np.array([[pixel]], dtype=np.uint8), encoding=encoding) + restored = Image.decode(source.encode()) + rgb = image_to_rgb(restored) + np.testing.assert_array_equal(rgb, [[[255, 0, 17]]]) + rgb[0, 0] = 0 + np.testing.assert_array_equal(image_view(restored), [[pixel]]) + + +def test_rgb_model_input_respects_big_endian_padded_grayscale() -> None: + msg = Image( + width=2, + height=2, + step=6, + encoding="mono16", + is_bigendian=1, + data=[0x12, 0x34, 0xFF, 0xFF, 99, 99, 0x01, 0x00, 0, 0, 99, 99], + ) + np.testing.assert_array_equal( + image_to_rgb(msg), [[[18, 18, 18], [255, 255, 255]], [[1, 1, 1], [0, 0, 0]]] + ) + + +def test_rgb_model_input_requires_explicit_float_depth_scale() -> None: + msg = image_from_array(np.array([[2.5]], dtype=np.float32), encoding="32FC1") + with pytest.raises(ValueError, match="RGB8"): + image_to_rgb(msg) + + +def test_mosaic_uses_generated_pixels_and_preserves_rgb_colors() -> None: + from dimos.memory.vis.utils import mosaic + + red = image_from_array(np.array([[[255, 0, 0]]], dtype=np.uint8), encoding="rgb8") + blue = image_from_array(np.array([[[255, 0, 0]]], dtype=np.uint8), encoding="bgr8") + observation = mosaic([red, blue], cols=2, cell_height=1) + result = Image.decode(observation.data.encode()) + assert result.encoding == "bgr8" + np.testing.assert_array_equal(image_to_rgb(result), [[[255, 0, 0], [0, 0, 255]]]) + assert observation.tags == {"mosaic": True} + assert observation.pose is None + + +def test_brightness_ignores_padding_and_normalizes_big_endian_u16() -> None: + message = Image( + width=2, + height=1, + step=6, + encoding="mono16", + is_bigendian=1, + data=[0, 0, 255, 255, 255, 255], + ) + assert image_brightness(message) == pytest.approx(0.5) + depth = image_from_array(np.array([[1.0]], dtype=np.float32), encoding="32FC1") + with pytest.raises(ValueError, match="unsigned integer"): + image_brightness(depth) + + +@pytest.mark.parametrize("mode", ["L", "RGBA"]) +def test_image_file_decodes_rgb_and_preserves_source_header(tmp_path, mode): + path = tmp_path / "frame.png" + pixels = np.array([[7, 19]], dtype=np.uint8) + source = PILImage.fromarray(pixels).convert(mode) + source.save(path) + header = Header(stamp=time_from_nanoseconds(1234567890), frame_id="camera") + message = image_from_file(path, header=header) + assert message.encoding == "rgb8" + assert message.header.encode() == header.encode() + np.testing.assert_array_equal(image_view(message), [[[7, 7, 7], [19, 19, 19]]]) + + +@pytest.mark.parametrize( + "encoding,dtype,shape", + [("64FC1", ">f8", (3, 4)), ("16SC1", ">i2", (3, 4)), ("32FC3", ">f4", (3, 4, 3))], +) +def test_remaining_raw_encodings_keep_signed_float_and_endian_values(encoding, dtype, shape): + pixels = (np.arange(np.prod(shape)).reshape(shape) - 5).astype(dtype) + source = image_from_array(pixels, encoding=encoding) + decoded = Image.decode(source.encode()) + assert decoded.encoding == encoding + assert decoded.is_bigendian == 1 + np.testing.assert_array_equal(image_view(decoded), pixels) diff --git a/dimos/msgs/test_geometry.py b/dimos/msgs/test_geometry.py new file mode 100644 index 00000000000..38780853d5f --- /dev/null +++ b/dimos/msgs/test_geometry.py @@ -0,0 +1,244 @@ +# 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. + +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +import math + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest + +from dimos.msgs.geometry import ( + compose_transforms, + inverse_transform, + pose_from_matrix, + pose_from_transform, + pose_matrix, + quaternion_euler, + quaternion_from_euler, + transform_from_odometry, + transform_from_pose, + transform_matrix, + translate_pose_local, + yaw, +) + + +@pytest.mark.parametrize("seconds,nanoseconds", [(0, 0), (-1, 500000000), (1700000000, 123456789)]) +def test_odometry_transform_preserves_frames_stamp_and_pose(seconds, nanoseconds): + message = Odometry(child_frame_id="robot/lidar") + message.header.frame_id = "robot/odom" + message.header.stamp = Time(sec=seconds, nanosec=nanoseconds) + message.pose.pose.position = Point(x=1, y=2, z=3) + message.pose.pose.orientation = Quaternion(z=0.6, w=0.8) + + transformed = transform_from_odometry(message) + + assert transformed.header.frame_id == "robot/odom" + assert transformed.child_frame_id == "robot/lidar" + assert transformed.header.stamp.sec == seconds + assert transformed.header.stamp.nanosec == nanoseconds + assert ( + transformed.transform.translation.x, + transformed.transform.translation.y, + transformed.transform.translation.z, + ) == (1, 2, 3) + assert transformed.transform.rotation == message.pose.pose.orientation + transformed.header.frame_id = "changed" + transformed.transform.rotation.w = 0 + assert message.header.frame_id == "robot/odom" + assert message.pose.pose.orientation.w == 0.8 + + +def _edge(parent, child, xyz, rotation, stamp=Time(sec=1700000000, nanosec=123456789)): + return TransformStamped( + header=Header(frame_id=parent, stamp=stamp), + child_frame_id=child, + transform=Transform(translation=Vector3(x=xyz[0], y=xyz[1], z=xyz[2]), rotation=rotation), + ) + + +def test_transform_composition_rotates_translation_and_preserves_exact_source_stamp(): + first = _edge("world", "arm", (1, -1, 0), quaternion_from_euler(0, 0, math.pi / 2)) + second = _edge("arm", "tool", (2, 3, 4), Quaternion(w=1), Time(sec=7)) + + result = compose_transforms(first, second) + + assert result.header == first.header + assert result.child_frame_id == "tool" + translation = result.transform.translation + assert (translation.x, translation.y, translation.z) == pytest.approx((-2, 1, 4)) + assert yaw(result.transform.rotation) == pytest.approx(math.pi / 2) + result.header.stamp.nanosec = 0 + result.transform.translation.x = 99 + assert first.header.stamp.nanosec == 123456789 + assert first.transform.translation.x == 1 + assert second.transform.translation.x == 2 + + +def test_inverse_composes_to_identity_and_swaps_frames_without_losing_nanoseconds(): + original = _edge("map", "camera", (2, 3, -4), quaternion_from_euler(0.3, -0.2, 1.1)) + inverse = inverse_transform(original) + result = compose_transforms(original, inverse) + + assert inverse.header.frame_id == "camera" + assert inverse.child_frame_id == "map" + assert inverse.header.stamp == original.header.stamp + translation = result.transform.translation + assert (translation.x, translation.y, translation.z) == pytest.approx((0, 0, 0), abs=1e-12) + rotation = result.transform.rotation + assert (rotation.x, rotation.y, rotation.z, abs(rotation.w)) == pytest.approx((0, 0, 0, 1)) + + +def test_transform_composition_rejects_disconnected_frames(): + first = _edge("a", "b", (0, 0, 0), Quaternion(w=1)) + second = _edge("c", "d", (0, 0, 0), Quaternion(w=1)) + with pytest.raises(ValueError, match="Cannot compose frames"): + compose_transforms(first, second) + + +@pytest.mark.parametrize( + "rotation", [Quaternion(w=0), Quaternion(x=math.nan, w=1), Quaternion(w=math.inf)] +) +def test_transform_math_rejects_invalid_rotations(rotation): + edge = _edge("a", "b", (0, 0, 0), rotation) + with pytest.raises(ValueError): + inverse_transform(edge) + with pytest.raises(ValueError): + compose_transforms(edge, _edge("b", "c", (0, 0, 0), Quaternion(w=1))) + + +@pytest.mark.parametrize("value", [math.nan, math.inf, -math.inf]) +def test_transform_math_rejects_nonfinite_translations_and_angles(value): + edge = _edge("a", "b", (value, 0, 0), Quaternion(w=1)) + with pytest.raises(ValueError, match="Translation components"): + inverse_transform(edge) + with pytest.raises(ValueError, match="Euler angles"): + quaternion_from_euler(value, 0, 0) + + +def test_pose_conversion_copies_nested_fields_and_preserves_header(): + edge = _edge("map", "robot", (1, 2, 3), Quaternion(z=0.6, w=0.8)) + pose = pose_from_transform(edge) + assert pose.header == edge.header + assert pose.pose.position == Point(x=1, y=2, z=3) + assert pose.pose.orientation == edge.transform.rotation + pose.header.frame_id = "changed" + pose.pose.orientation.w = 0 + assert edge.header.frame_id == "map" + assert edge.transform.rotation.w == 0.8 + + +def test_pose_transform_preserves_header_and_does_not_alias_input(): + source = PoseStamped( + header=Header(frame_id="odom", stamp=Time(sec=-1, nanosec=987654321)), + pose=Pose(position=Point(x=1, y=2, z=3), orientation=Quaternion(z=0.6, w=0.8)), + ) + result = transform_from_pose(source, child_frame_id="robot/base") + decoded = TransformStamped.decode(result.encode()) + assert decoded.header == source.header + assert decoded.child_frame_id == "robot/base" + assert decoded.transform.translation == Vector3(x=1, y=2, z=3) + assert decoded.transform.rotation == source.pose.orientation + result.header.frame_id = "changed" + result.transform.rotation.w = 0 + assert source.header.frame_id == "odom" + assert source.pose.orientation.w == 0.8 + + +def test_transform_matrix_maps_rotation_and_translation_without_aliasing(): + value = Transform( + translation=Vector3(x=2.0, y=-1.0, z=3.0), + rotation=quaternion_from_euler(0.0, 0.0, math.pi / 2), + ) + matrix = transform_matrix(value) + np.testing.assert_allclose(matrix @ [1.0, 0.0, 0.0, 1.0], [2.0, 0.0, 3.0, 1.0], atol=1e-12) + value.translation.x = 10.0 + assert matrix[0, 3] == 2.0 + + +def test_pose_matrix_roundtrip_preserves_rigid_motion_without_aliasing(): + pose = Pose( + position=Point(x=2, y=-1, z=3), + orientation=quaternion_from_euler(0, 0, math.pi / 2), + ) + matrix = pose_matrix(pose) + np.testing.assert_allclose(matrix @ [1, 0, 0, 1], [2, 0, 3, 1], atol=1e-12) + decoded = Pose.decode(pose_from_matrix(matrix).encode()) + np.testing.assert_allclose(pose_matrix(decoded), matrix, atol=1e-12) + matrix[0, 3] = 99 + decoded.position.y = 42 + assert pose.position == Point(x=2, y=-1, z=3) + assert decoded.position.x == 2 + + +@pytest.mark.parametrize( + "matrix", + [ + np.eye(3), + np.full((4, 4), math.nan), + np.diag([1, 1, 1, 2]), + np.diag([2, 1, 1, 1]), + np.diag([-1, 1, 1, 1]), + ], +) +def test_pose_from_matrix_rejects_nonrigid_or_invalid_matrices(matrix): + with pytest.raises(ValueError): + pose_from_matrix(matrix) + + +@pytest.mark.parametrize( + "pose", + [ + Pose(position=Point(x=math.inf), orientation=Quaternion(w=1)), + Pose(orientation=Quaternion(w=0)), + Pose(orientation=Quaternion(w=math.nan)), + ], +) +def test_pose_matrix_rejects_invalid_geometry(pose): + with pytest.raises(ValueError): + pose_matrix(pose) + + +def test_local_offset_uses_pose_axes_and_copies_orientation(): + pose = Pose( + position=Point(x=1, y=2, z=3), + orientation=quaternion_from_euler(0, 0, math.pi / 2), + ) + translated = translate_pose_local(pose, Vector3(x=2, z=-1)) + position = translated.position + assert (position.x, position.y, position.z) == pytest.approx((1, 4, 2)) + assert quaternion_euler(translated.orientation) == pytest.approx((0, 0, math.pi / 2)) + translated.orientation.w = 0 + assert pose.orientation.w == pytest.approx(math.sqrt(0.5)) + assert pose.position == Point(x=1, y=2, z=3) + + +@pytest.mark.parametrize("value", [math.nan, math.inf, -math.inf]) +def test_local_offset_rejects_nonfinite_displacements(value): + with pytest.raises(ValueError, match="Offset components"): + translate_pose_local(Pose(orientation=Quaternion(w=1)), Vector3(x=value)) diff --git a/dimos/msgs/test_helpers.py b/dimos/msgs/test_helpers.py index 72958fc1525..b3a7c516b93 100644 --- a/dimos/msgs/test_helpers.py +++ b/dimos/msgs/test_helpers.py @@ -12,15 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos_lcm.geometry_msgs import Pose -import pytest +from dimos_generated.geometry_msgs.msg import Pose -from dimos.msgs.helpers import lcm_msg_type +from dimos.msgs.helpers import resolve_msg_type -def test_lcm_msg_type() -> None: - assert lcm_msg_type("geometry_msgs.Pose") is Pose - with pytest.raises(ImportError): - lcm_msg_type("geometry_msgs.Nope") - with pytest.raises(ValueError): - lcm_msg_type("Bare") +def test_installed_generated_type_resolves_by_schema_name() -> None: + assert resolve_msg_type("geometry_msgs/msg/Pose") is Pose + + +def test_unknown_names_do_not_import_a_legacy_message_package() -> None: + assert resolve_msg_type("geometry_msgs/msg/Nope") is None + assert resolve_msg_type("geometry_msgs.Pose") is None + assert resolve_msg_type("Bare") is None diff --git a/dimos/msgs/test_msg_construction.py b/dimos/msgs/test_msg_construction.py index c4f314753c5..3b49b4661fd 100644 --- a/dimos/msgs/test_msg_construction.py +++ b/dimos/msgs/test_msg_construction.py @@ -12,616 +12,199 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Construction semantics for the hot message types. +"""Generated field construction replaces polymorphic rich-message constructors. -Every accepted construction form for the types whose ``__init__`` used to be a -``plum`` multiple-dispatch stack, plus the error cases. These pin the contract -that the dispatch removal had to preserve. +Array/dictionary inputs are unpacked explicitly. Stamped types contain a header +and a payload; timestamps default to zero and never sample the wall clock. """ -import time +from copy import deepcopy -from dimos_lcm.geometry_msgs import ( - Point as LCMPoint, - Pose as LCMPose, - PoseWithCovariance as LCMPoseWithCovariance, - Quaternion as LCMQuaternion, - Twist as LCMTwist, - TwistWithCovariance as LCMTwistWithCovariance, - Vector3 as LCMVector3, +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + Point, + PointStamped, + Pose, + PoseStamped, + PoseWithCovariance, + PoseWithCovarianceStamped, + Quaternion, + TransformStamped, + Twist, + TwistStamped, + TwistWithCovariance, + TwistWithCovarianceStamped, + Vector3, ) -from dimos_lcm.sensor_msgs import JointState as LCMJointState, Joy as LCMJoy +from dimos_generated.nav_msgs.msg import Odometry, Path +from dimos_generated.sensor_msgs.msg import JointState, Joy +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest - -from dimos.msgs.geometry_msgs.PointStamped import PointStamped -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.PoseWithCovariance import PoseWithCovariance -from dimos.msgs.geometry_msgs.PoseWithCovarianceStamped import PoseWithCovarianceStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped -from dimos.msgs.geometry_msgs.TwistWithCovariance import TwistWithCovariance -from dimos.msgs.geometry_msgs.TwistWithCovarianceStamped import TwistWithCovarianceStamped -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.sensor_msgs.Joy import Joy - -IDENTITY = (0.0, 0.0, 0.0, 1.0) - - -# Quaternion - - -def test_quaternion_empty_is_identity() -> None: - assert Quaternion().to_tuple() == IDENTITY - - -def test_quaternion_four_positional_floats() -> None: - q = Quaternion(1.0, 2.0, 3.0, 4.0) - assert q.to_tuple() == (1.0, 2.0, 3.0, 4.0) - - -def test_quaternion_four_positional_ints_are_coerced_to_float() -> None: - q = Quaternion(1, 2, 3, 4) - assert q.to_tuple() == (1.0, 2.0, 3.0, 4.0) - assert all(type(c) is float for c in q.to_tuple()) - - -def test_quaternion_from_list() -> None: - assert Quaternion([1.0, 2.0, 3.0, 4.0]).to_tuple() == (1.0, 2.0, 3.0, 4.0) - - -def test_quaternion_from_tuple() -> None: - assert Quaternion((1.0, 2.0, 3.0, 4.0)).to_tuple() == (1.0, 2.0, 3.0, 4.0) - - -def test_quaternion_from_range() -> None: - assert Quaternion(range(4)).to_tuple() == (0.0, 1.0, 2.0, 3.0) - - -def test_quaternion_from_ndarray() -> None: - assert Quaternion(np.array([1.0, 2.0, 3.0, 4.0])).to_tuple() == (1.0, 2.0, 3.0, 4.0) - - -def test_quaternion_copy_constructor() -> None: - src = Quaternion(0.1, 0.2, 0.3, 0.4) - assert Quaternion(src).to_tuple() == src.to_tuple() - - -def test_quaternion_from_lcm_keeps_lcm_defaults() -> None: - """The LCM base defaults w to 0.0, unlike the dimos identity default.""" - assert Quaternion(LCMQuaternion()).to_tuple() == (0.0, 0.0, 0.0, 0.0) - - -def test_quaternion_from_lcm_with_values() -> None: - assert Quaternion(LCMQuaternion(1.0, 2.0, 3.0, 4.0)).to_tuple() == (1.0, 2.0, 3.0, 4.0) - - -def test_quaternion_sequence_wrong_length_raises() -> None: - for bad in ([1, 2, 3], [1, 2, 3, 4, 5], (1, 2)): - with pytest.raises(ValueError, match=r"exactly 4 components"): - Quaternion(bad) - - -def test_quaternion_ndarray_wrong_size_raises() -> None: - with pytest.raises(ValueError, match=r"exactly 4 components"): - Quaternion(np.array([1.0, 2.0, 3.0])) - - -def test_quaternion_lcm_roundtrip_uses_the_sequence_form() -> None: - """``_lcm_decode_one`` builds from a 4-tuple, so decoding exercises that path.""" - decoded = Quaternion.lcm_decode(Quaternion(1.0, 2.0, 3.0, 4.0).lcm_encode()) - assert decoded.to_tuple() == (1.0, 2.0, 3.0, 4.0) - - -# Pose - - -def test_pose_empty_is_origin_with_identity() -> None: - p = Pose() - assert p.position.to_tuple() == (0.0, 0.0, 0.0) - assert p.orientation.to_tuple() == IDENTITY - - -def test_pose_three_positional_floats() -> None: - p = Pose(1, 2, 3) - assert p.position.to_tuple() == (1.0, 2.0, 3.0) - assert p.orientation.to_tuple() == IDENTITY - - -def test_pose_seven_positional_floats() -> None: - p = Pose(1, 2, 3, 0, 0, 0, 1) - assert p.position.to_tuple() == (1.0, 2.0, 3.0) - assert p.orientation.to_tuple() == IDENTITY - - -def test_pose_from_position_sequence_only() -> None: - p = Pose([1, 2, 3]) - assert p.position.to_tuple() == (1.0, 2.0, 3.0) - assert p.orientation.to_tuple() == IDENTITY - - -def test_pose_from_position_and_orientation_sequences() -> None: - p = Pose([1, 2, 3], [0, 0, 0, 1]) - assert p.position.to_tuple() == (1.0, 2.0, 3.0) - assert p.orientation.to_tuple() == IDENTITY - - -def test_pose_from_vector3_and_quaternion() -> None: - p = Pose(Vector3(1, 2, 3), Quaternion(0, 0, 0, 1)) - assert p.position.to_tuple() == (1.0, 2.0, 3.0) - assert p.orientation.to_tuple() == IDENTITY - - -def test_pose_from_ndarrays() -> None: - p = Pose(np.array([1.0, 2.0, 3.0]), np.array([0.0, 0.0, 0.0, 1.0])) - assert p.position.to_tuple() == (1.0, 2.0, 3.0) - assert p.orientation.to_tuple() == IDENTITY - - -def test_pose_keyword_position_only() -> None: - p = Pose(position=[1, 2, 3]) - assert p.position.to_tuple() == (1.0, 2.0, 3.0) - assert p.orientation.to_tuple() == IDENTITY - - -def test_pose_keyword_orientation_only() -> None: - p = Pose(orientation=[0, 0, 0, 1]) - assert p.position.to_tuple() == (0.0, 0.0, 0.0) - assert p.orientation.to_tuple() == IDENTITY - - -def test_pose_keyword_both() -> None: - p = Pose(position=[1, 2, 3], orientation=[0, 0, 0, 1]) - assert p.position.to_tuple() == (1.0, 2.0, 3.0) - - -def test_pose_positional_position_with_keyword_orientation() -> None: - p = Pose([1, 2, 3], orientation=[0, 0, 0, 1]) - assert p.position.to_tuple() == (1.0, 2.0, 3.0) - assert p.orientation.to_tuple() == IDENTITY - - -def test_pose_explicit_none_falls_back_to_defaults() -> None: - p = Pose(None, None) - assert p.position.to_tuple() == (0.0, 0.0, 0.0) - assert p.orientation.to_tuple() == IDENTITY - - -def test_pose_from_pair_tuple() -> None: - p = Pose(([1, 2, 3], [0, 0, 0, 1])) - assert p.position.to_tuple() == (1.0, 2.0, 3.0) - assert p.orientation.to_tuple() == IDENTITY - - -def test_pose_from_three_tuple_is_a_position() -> None: - """A 3-tuple of numbers is a position, not a (position, orientation) pair.""" - p = Pose((1, 2, 3)) - assert p.position.to_tuple() == (1.0, 2.0, 3.0) - - -def test_pose_from_dict() -> None: - p = Pose({"position": [1, 2, 3], "orientation": [0, 0, 0, 1]}) - assert p.position.to_tuple() == (1.0, 2.0, 3.0) - assert p.orientation.to_tuple() == IDENTITY +from rosbags.typesys import Stores, get_typestore + +from dimos.msgs.geometry import ( + point_from_array, + quaternion_array, + quaternion_euler, + quaternion_from_array, + vector_array, + vector_from_array, +) +from dimos.msgs.time import to_nanoseconds -def test_pose_from_dict_missing_key_raises_keyerror() -> None: - with pytest.raises(KeyError): - Pose({"position": [1, 2, 3]}) - with pytest.raises(KeyError): - Pose({}) +@pytest.mark.parametrize( + "values", [[1.0, 2.0, 3.0, 4.0], (1, 2, 3, 4), range(4), np.array([1.0, 2.0, 3.0, 4.0])] +) +def test_quaternion_explicit_array_conversion(values) -> None: + q = quaternion_from_array(values) + np.testing.assert_array_equal(quaternion_array(q), values) + assert all(type(value) is float for value in (q.x, q.y, q.z, q.w)) + assert Quaternion.decode(q.encode()) == q -def test_pose_copy_constructor() -> None: - src = Pose(1, 2, 3, 0, 0, 0, 1) - copy = Pose(src) - assert copy.position.to_tuple() == (1.0, 2.0, 3.0) - assert copy.position is not src.position +def test_quaternion_defaults_and_partial_fields() -> None: + assert Quaternion() == Quaternion(x=0, y=0, z=0, w=1) + assert Quaternion(z=1) == Quaternion(x=0, y=0, z=1, w=1) + # Explicit zero w is retained, even though it is not a spatial rotation. + assert Quaternion.decode(Quaternion(w=0).encode()).w == 0 -def test_pose_from_posestamped_downcasts_to_pose() -> None: - src = PoseStamped(ts=5.0, frame_id="odom", position=[1, 2, 3]) - p = Pose(src) - assert type(p) is Pose - assert p.position.to_tuple() == (1.0, 2.0, 3.0) +@pytest.mark.parametrize("values", [[1, 2, 3], [1, 2, 3, 4, 5], (1, 2), np.array([1.0, 2.0, 3.0])]) +def test_quaternion_array_wrong_size(values) -> None: + with pytest.raises(ValueError, match="exactly 4 components"): + quaternion_from_array(values) -def test_pose_from_lcm() -> None: - # Build the sub-messages explicitly: the generated LCM classes share one - # mutable default instance across every construction, so LCMPose().position - # carries whatever the last writer left there. - p = Pose( - LCMPose(position=LCMPoint(1.0, 2.0, 3.0), orientation=LCMQuaternion(0.0, 0.0, 0.0, 1.0)) +@pytest.mark.parametrize("position", [[1, 2, 3], (1, 2, 3), np.array([1.0, 2.0, 3.0])]) +@pytest.mark.parametrize( + "orientation", [[0, 0, 0, 1], (0, 0, 0, 1), np.array([0.0, 0.0, 0.0, 1.0])] +) +def test_pose_explicit_numeric_fields(position, orientation) -> None: + pose = Pose(position=point_from_array(position), orientation=quaternion_from_array(orientation)) + np.testing.assert_array_equal(vector_array(pose.position), [1, 2, 3]) + np.testing.assert_array_equal(quaternion_array(pose.orientation), [0, 0, 0, 1]) + assert type(pose.position) is Point + assert type(pose.orientation) is Quaternion + + +def test_pose_explicit_dictionary_and_pair_unpacking() -> None: + data = {"position": [1, 2, 3], "orientation": [0, 0, 0, 1]} + pair = (data["position"], data["orientation"]) + expected = Pose(position=Point(x=1, y=2, z=3)) + assert ( + Pose( + position=point_from_array(data["position"]), + orientation=quaternion_from_array(data["orientation"]), + ) + == expected ) - assert p.position.to_tuple() == (1.0, 2.0, 3.0) - assert p.orientation.to_tuple() == IDENTITY - - -def test_pose_position_and_orientation_are_wrapper_types() -> None: - p = Pose(1, 2, 3) - assert isinstance(p.position, Vector3) - assert isinstance(p.orientation, Quaternion) - - -# Twist - - -def test_twist_empty_is_zero() -> None: - t = Twist() - assert t.linear.to_tuple() == (0.0, 0.0, 0.0) - assert t.angular.to_tuple() == (0.0, 0.0, 0.0) - - -def test_twist_from_two_sequences() -> None: - t = Twist([1, 2, 3], [4, 5, 6]) - assert t.linear.to_tuple() == (1.0, 2.0, 3.0) - assert t.angular.to_tuple() == (4.0, 5.0, 6.0) - - -def test_twist_from_two_vector3() -> None: - t = Twist(Vector3(1, 2, 3), Vector3(4, 5, 6)) - assert t.linear.to_tuple() == (1.0, 2.0, 3.0) - assert t.angular.to_tuple() == (4.0, 5.0, 6.0) - - -def test_twist_from_ndarrays() -> None: - t = Twist(np.array([1.0, 2.0, 3.0]), np.array([4.0, 5.0, 6.0])) - assert t.linear.to_tuple() == (1.0, 2.0, 3.0) - assert t.angular.to_tuple() == (4.0, 5.0, 6.0) - - -def test_twist_angular_quaternion_is_converted_to_euler() -> None: - t = Twist([1, 2, 3], Quaternion(0, 0, 0, 1)) - assert t.linear.to_tuple() == (1.0, 2.0, 3.0) - assert t.angular.to_tuple() == (0.0, 0.0, 0.0) - - -def test_twist_copy_constructor() -> None: - src = Twist([1, 2, 3], [4, 5, 6]) - copy = Twist(src) - assert copy.linear.to_tuple() == (1.0, 2.0, 3.0) - assert copy.linear is not src.linear - - -def test_twist_from_lcm() -> None: - t = Twist(LCMTwist(linear=LCMVector3(1.0, 2.0, 3.0), angular=LCMVector3(4.0, 5.0, 6.0))) - assert t.linear.to_tuple() == (1.0, 2.0, 3.0) - assert t.angular.to_tuple() == (4.0, 5.0, 6.0) - - -def test_twist_from_twiststamped_downcasts() -> None: - t = Twist(TwistStamped(linear=[1, 2, 3], angular=[4, 5, 6])) - assert type(t) is Twist - assert t.linear.to_tuple() == (1.0, 2.0, 3.0) - - -def test_twist_keyword_both() -> None: - t = Twist(linear=[1, 2, 3], angular=[4, 5, 6]) - assert t.linear.to_tuple() == (1.0, 2.0, 3.0) - assert t.angular.to_tuple() == (4.0, 5.0, 6.0) - - -def test_twist_keyword_linear_only() -> None: - t = Twist(linear=[1, 2, 3]) - assert t.linear.to_tuple() == (1.0, 2.0, 3.0) - assert t.angular.to_tuple() == (0.0, 0.0, 0.0) - - -def test_twist_keyword_angular_only() -> None: - t = Twist(angular=[4, 5, 6]) - assert t.linear.to_tuple() == (0.0, 0.0, 0.0) - assert t.angular.to_tuple() == (4.0, 5.0, 6.0) - - -def test_twist_keyword_angular_quaternion_is_converted_to_euler() -> None: - t = Twist(linear=[1, 2, 3], angular=Quaternion(0, 0, 0, 1)) - assert t.angular.to_tuple() == (0.0, 0.0, 0.0) - - -# PoseWithCovariance - - -def test_pwc_empty() -> None: - p = PoseWithCovariance() - assert p.pose.position.to_tuple() == (0.0, 0.0, 0.0) - assert np.array_equal(p.covariance, np.zeros(36)) - - -def test_pwc_from_pose_only() -> None: - p = PoseWithCovariance(Pose(1, 2, 3)) - assert p.pose.position.to_tuple() == (1.0, 2.0, 3.0) - assert np.array_equal(p.covariance, np.zeros(36)) - - -def test_pwc_from_pose_and_covariance() -> None: - p = PoseWithCovariance(Pose(1, 2, 3), np.arange(36.0)) - assert np.array_equal(p.covariance, np.arange(36.0)) - - -def test_pwc_explicit_none_covariance_is_zeros() -> None: - assert np.array_equal(PoseWithCovariance(Pose(1, 2, 3), None).covariance, np.zeros(36)) - - -def test_pwc_covariance_wrong_size_raises() -> None: - with pytest.raises(ValueError, match=r"reshape"): - PoseWithCovariance(Pose(1, 2, 3), [1, 2, 3]) - - -def test_pwc_copy_constructor_copies_covariance() -> None: - src = PoseWithCovariance(Pose(1, 2, 3), np.arange(36.0)) - copy = PoseWithCovariance(src) - assert np.array_equal(copy.covariance, src.covariance) - assert copy.covariance is not src.covariance - - -def test_pwc_from_lcm() -> None: - lcm_pose = LCMPose( - position=LCMPoint(1.0, 2.0, 3.0), orientation=LCMQuaternion(0.0, 0.0, 0.0, 1.0) + assert ( + Pose(position=point_from_array(pair[0]), orientation=quaternion_from_array(pair[1])) + == expected ) - p = PoseWithCovariance(LCMPoseWithCovariance(pose=lcm_pose, covariance=list(range(36)))) - assert p.pose.position.to_tuple() == (1.0, 2.0, 3.0) - assert np.array_equal(p.covariance, np.arange(36.0)) - - -def test_pwc_from_dict_with_pose_object() -> None: - p = PoseWithCovariance({"pose": Pose(1, 2, 3)}) - assert p.pose.position.to_tuple() == (1.0, 2.0, 3.0) - assert np.array_equal(p.covariance, np.zeros(36)) - - -def test_pwc_from_dict_with_covariance() -> None: - p = PoseWithCovariance({"pose": Pose(1, 2, 3), "covariance": list(range(36))}) - assert np.array_equal(p.covariance, np.arange(36.0)) - - -def test_pwc_covariance_matrix_view() -> None: - p = PoseWithCovariance(Pose(), np.arange(36.0)) - assert p.covariance_matrix.shape == (6, 6) - - -# TwistWithCovariance - - -def test_twc_empty() -> None: - t = TwistWithCovariance() - assert t.twist.linear.to_tuple() == (0.0, 0.0, 0.0) - assert np.array_equal(t.covariance, np.zeros(36)) - - -def test_twc_from_twist_only() -> None: - t = TwistWithCovariance(Twist([1, 2, 3], [4, 5, 6])) - assert t.twist.linear.to_tuple() == (1.0, 2.0, 3.0) - - -def test_twc_from_twist_and_covariance() -> None: - t = TwistWithCovariance(Twist([1, 2, 3], [4, 5, 6]), np.arange(36.0)) - assert np.array_equal(t.covariance, np.arange(36.0)) - - -def test_twc_from_linear_angular_pair() -> None: - t = TwistWithCovariance(([1, 2, 3], [4, 5, 6])) - assert t.twist.linear.to_tuple() == (1.0, 2.0, 3.0) - assert t.twist.angular.to_tuple() == (4.0, 5.0, 6.0) - + assert Pose(position=Point(x=5)).position == Point(x=5, y=0, z=0) -def test_twc_copy_constructor() -> None: - src = TwistWithCovariance(Twist([1, 2, 3], [4, 5, 6]), np.arange(36.0)) - copy = TwistWithCovariance(src) - assert np.array_equal(copy.covariance, src.covariance) - assert copy.covariance is not src.covariance +@pytest.mark.parametrize("values", [[1, 2, 3], (1, 2, 3), np.array([1.0, 2.0, 3.0])]) +def test_twist_explicit_arrays(values) -> None: + twist = Twist(linear=vector_from_array(values), angular=vector_from_array([4, 5, 6])) + np.testing.assert_array_equal(vector_array(twist.linear), [1, 2, 3]) + np.testing.assert_array_equal(vector_array(twist.angular), [4, 5, 6]) + assert Twist(linear=twist.linear).angular == Vector3() + assert Twist(angular=twist.angular).linear == Vector3() -def test_twc_from_lcm() -> None: - lcm_twist = LCMTwist(linear=LCMVector3(1.0, 2.0, 3.0), angular=LCMVector3(4.0, 5.0, 6.0)) - t = TwistWithCovariance(LCMTwistWithCovariance(twist=lcm_twist, covariance=list(range(36)))) - assert t.twist.linear.to_tuple() == (1.0, 2.0, 3.0) - assert np.array_equal(t.covariance, np.arange(36.0)) - -def test_twc_from_dict() -> None: - t = TwistWithCovariance({"twist": Twist([1, 2, 3], [4, 5, 6])}) - assert t.twist.linear.to_tuple() == (1.0, 2.0, 3.0) - assert np.array_equal(t.covariance, np.zeros(36)) - - -def test_twc_from_dict_with_covariance() -> None: - t = TwistWithCovariance({"twist": Twist([1, 2, 3], [4, 5, 6]), "covariance": list(range(36))}) - assert np.array_equal(t.covariance, np.arange(36.0)) - - -# JointState - - -def test_jointstate_empty_has_empty_lists() -> None: - j = JointState() - assert (j.name, j.position, j.velocity, j.effort) == ([], [], [], []) - assert j.frame_id == "" - - -def test_jointstate_keywords() -> None: - j = JointState(ts=5.0, frame_id="f", name=["a"], position=[1.0]) - assert (j.ts, j.frame_id, j.name, j.position) == (5.0, "f", ["a"], [1.0]) - assert (j.velocity, j.effort) == ([], []) - - -def test_jointstate_all_positional() -> None: - j = JointState(5.0, "f", ["a"], [1.0], [2.0], [3.0]) - assert (j.name, j.position, j.velocity, j.effort) == (["a"], [1.0], [2.0], [3.0]) - - -def test_jointstate_from_dict() -> None: - j = JointState({"name": ["a"], "position": [1.0]}) - assert (j.name, j.position) == (["a"], [1.0]) - - -def test_jointstate_from_empty_dict() -> None: - j = JointState({}) - assert (j.name, j.position, j.velocity, j.effort) == ([], [], [], []) - - -def test_jointstate_copy_constructor_copies_lists() -> None: - src = JointState(ts=5.0, frame_id="f", name=["a"], position=[1.0]) - copy = JointState(src) - assert copy.ts == 5.0 - assert copy.name == ["a"] - assert copy.name is not src.name - - -def test_jointstate_from_lcm() -> None: - j = JointState(LCMJointState()) - assert j.name == [] - - -# Joy - - -def test_joy_empty() -> None: - j = Joy() - assert (j.axes, j.buttons, j.frame_id) == ([], [], "") - - -def test_joy_keywords() -> None: - j = Joy(ts=5.0, frame_id="f", axes=[1.0], buttons=[1]) - assert (j.ts, j.frame_id, j.axes, j.buttons) == (5.0, "f", [1.0], [1]) - - -def test_joy_all_positional() -> None: - j = Joy(5.0, "f", [1.0], [1]) - assert (j.axes, j.buttons) == ([1.0], [1]) - - -def test_joy_from_axes_buttons_pair() -> None: - j = Joy(([1.0, 2.0], [1, 0])) - assert (j.axes, j.buttons) == ([1.0, 2.0], [1, 0]) - - -def test_joy_from_dict() -> None: - j = Joy({"axes": [1.0], "buttons": [1]}) - assert (j.axes, j.buttons) == ([1.0], [1]) - - -def test_joy_from_empty_dict() -> None: - j = Joy({}) - assert (j.axes, j.buttons) == ([], []) - - -def test_joy_copy_constructor_copies_lists() -> None: - src = Joy(ts=5.0, axes=[1.0], buttons=[1]) - copy = Joy(src) - assert copy.ts == 5.0 - assert copy.axes == [1.0] - assert copy.axes is not src.axes - - -def test_joy_from_lcm() -> None: - j = Joy(LCMJoy()) - assert j.axes == [] - - -# Stamped wrappers: ts / frame_id plumbing - - -def test_posestamped_defaults_stamp_to_now() -> None: - before = time.time() - ps = PoseStamped() - assert before <= ps.ts <= time.time() - assert ps.frame_id == "" - assert ps.position.to_tuple() == (0.0, 0.0, 0.0) - - -def test_posestamped_keeps_explicit_ts_and_frame() -> None: - ps = PoseStamped(ts=5.0, frame_id="odom", position=[1, 2, 3], orientation=[0, 0, 0, 1]) - assert (ps.ts, ps.frame_id) == (5.0, "odom") - assert ps.position.to_tuple() == (1.0, 2.0, 3.0) - - -def test_posestamped_positional_ts_and_frame() -> None: - ps = PoseStamped(5.0, "odom") - assert (ps.ts, ps.frame_id) == (5.0, "odom") - - -def test_posestamped_is_a_pose() -> None: - assert isinstance(PoseStamped(), Pose) - - -def test_posestamped_inherits_pose_positional_forms() -> None: - """PoseStamped(x, y, z) and the 7-float form come from Pose.""" - assert PoseStamped(3.2, 1.5, 0.0).position.to_tuple() == (3.2, 1.5, 0.0) - seven = PoseStamped(1, 2, 0, 0, 0, 0.1, 1) - assert seven.position.to_tuple() == (1.0, 2.0, 0.0) - assert seven.orientation.to_tuple() == (0.0, 0.0, 0.1, 1.0) - - -def test_posestamped_inherits_pose_sequence_forms() -> None: - ps = PoseStamped([1, 2, 3], [0, 0, 0, 1]) - assert ps.position.to_tuple() == (1.0, 2.0, 3.0) - assert ps.orientation.to_tuple() == IDENTITY - - -def test_twiststamped_inherits_twist_positional_form() -> None: - ts = TwistStamped([1, 2, 3], [4, 5, 6]) - assert ts.linear.to_tuple() == (1.0, 2.0, 3.0) - assert ts.angular.to_tuple() == (4.0, 5.0, 6.0) - - -def test_twiststamped_keeps_explicit_ts_and_frame() -> None: - ts = TwistStamped(ts=5.0, frame_id="base", linear=[1, 2, 3], angular=[4, 5, 6]) - assert (ts.ts, ts.frame_id) == (5.0, "base") - assert ts.linear.to_tuple() == (1.0, 2.0, 3.0) - - -def test_twiststamped_defaults_stamp_to_now() -> None: - before = time.time() - t = TwistStamped() - assert before <= t.ts <= time.time() - - -def test_pwcs_keeps_explicit_ts_and_frame() -> None: - p = PoseWithCovarianceStamped(ts=5.0, frame_id="map", pose=Pose(1, 2, 3)) - assert (p.ts, p.frame_id) == (5.0, "map") - assert p.pose.position.to_tuple() == (1.0, 2.0, 3.0) - - -def test_pwcs_with_covariance() -> None: - p = PoseWithCovarianceStamped(ts=5.0, pose=Pose(1, 2, 3), covariance=list(range(36))) - assert np.array_equal(p.covariance, np.arange(36.0)) - - -def test_pwcs_empty_defaults() -> None: - p = PoseWithCovarianceStamped() - assert p.pose.position.to_tuple() == (0.0, 0.0, 0.0) - assert np.array_equal(p.covariance, np.zeros(36)) - - -def test_twcs_keeps_explicit_ts_and_frame() -> None: - t = TwistWithCovarianceStamped(ts=5.0, frame_id="base", twist=Twist([1, 2, 3], [4, 5, 6])) - assert (t.ts, t.frame_id) == (5.0, "base") - assert t.twist.linear.to_tuple() == (1.0, 2.0, 3.0) - - -def test_twcs_with_covariance() -> None: - t = TwistWithCovarianceStamped( - ts=5.0, twist=Twist([1, 2, 3], [4, 5, 6]), covariance=list(range(36)) +def test_twist_euler_conversion_is_explicit() -> None: + twist = Twist( + linear=Vector3(x=1, y=2, z=3), angular=vector_from_array(quaternion_euler(Quaternion())) ) - assert np.array_equal(t.covariance, np.arange(36.0)) - - -def test_twcs_empty_defaults() -> None: - t = TwistWithCovarianceStamped() - assert t.twist.linear.to_tuple() == (0.0, 0.0, 0.0) - + assert twist.linear == Vector3(x=1, y=2, z=3) + assert twist.angular == Vector3() + + +@pytest.mark.parametrize( + "message", + [ + Quaternion(x=0.1, y=0.2, z=0.3, w=0.4), + Pose(position=Point(x=1, y=2, z=3)), + Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=4, y=5, z=6)), + PoseWithCovariance(pose=Pose(position=Point(x=1, y=2, z=3)), covariance=list(range(36))), + TwistWithCovariance(twist=Twist(linear=Vector3(x=1, y=2, z=3)), covariance=list(range(36))), + JointState(header=Header(stamp=Time(sec=5), frame_id="f"), name=["a"], position=[1.0]), + Joy(header=Header(stamp=Time(sec=5), frame_id="f"), axes=[1.0], buttons=[1]), + ], +) +def test_copy_and_cdr_roundtrip_preserve_fields(message) -> None: + copied = deepcopy(message) + assert copied == message + assert copied is not message + assert type(message).decode(message.encode()) == message + independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr( + message.encode(), message.msg_name + ) + assert independent.__msgtype__ == message.msg_name + + +def test_copy_mutation_is_independent() -> None: + pose = Pose(position=Point(x=1, y=2, z=3)) + pose_copy = deepcopy(pose) + pose_copy.position.x = 99 + assert pose.position.x == 1 + twist = Twist(linear=Vector3(x=1, y=2, z=3)) + twist_copy = deepcopy(twist) + twist_copy.linear.x = 99 + assert twist.linear.x == 1 + covariance = PoseWithCovariance(covariance=list(range(36))) + covariance_copy = deepcopy(covariance) + covariance_copy.covariance[0] = 99 + assert covariance.covariance[0] == 0 + joints = JointState(name=["a"], position=[1.0]) + joints_copy = deepcopy(joints) + joints_copy.name[0] = "b" + assert joints.name == ["a"] + joy = Joy(axes=[1.0], buttons=[1]) + joy_copy = deepcopy(joy) + joy_copy.axes[0] = 2 + assert joy.axes == [1.0] + + +@pytest.mark.parametrize("message_type", [PoseWithCovariance, TwistWithCovariance]) +def test_covariance_defaults_layout_and_length(message_type) -> None: + assert list(message_type().covariance) == [0.0] * 36 + covariance = np.arange(36.0).reshape(6, 6) + message = message_type(covariance=covariance.flatten().tolist()) + np.testing.assert_array_equal(np.asarray(message.covariance).reshape(6, 6), covariance) + for invalid in ([1, 2, 3], [0.0] * 35, [0.0] * 37): + with pytest.raises(RuntimeError, match="Unable to cast"): + message_type(covariance=invalid) + + +def test_sensor_explicit_fields_and_defaults() -> None: + joints = JointState( + header=Header(stamp=Time(sec=5), frame_id="f"), + name=["a"], + position=[1.0], + velocity=[2.0], + effort=[3.0], + ) + assert (joints.name, joints.position, joints.velocity, joints.effort) == ( + ["a"], + [1.0], + [2.0], + [3.0], + ) + assert joints.header == Header(stamp=Time(sec=5), frame_id="f") + assert JointState().name == [] + assert JointState().velocity == [] + assert JointState().effort == [] + pair = ([1.0, 2.0], [1, 0]) + joy = Joy(axes=pair[0], buttons=pair[1]) + assert (joy.axes, joy.buttons) == pair + assert (Joy().axes, Joy().buttons) == ([], []) + assert Joy().header.frame_id == "" -# The ``ts`` sentinel -# -# ``ts`` defaults to None and means "stamp me now". 0.0 is an ordinary -# timestamp and is stored as given — it used to be swallowed by a -# ``ts if ts != 0`` sentinel that made it unrepresentable. STAMPED_TYPES = [ PoseStamped, @@ -632,152 +215,77 @@ def test_twcs_empty_defaults() -> None: Joy, Path, PointStamped, - Transform, + TransformStamped, Odometry, ] -@pytest.mark.parametrize("msg_type", STAMPED_TYPES, ids=lambda t: t.__name__) -def test_zero_ts_is_a_real_timestamp(msg_type: type) -> None: - assert msg_type(ts=0.0).ts == 0.0 - - -@pytest.mark.parametrize("msg_type", STAMPED_TYPES, ids=lambda t: t.__name__) -def test_omitted_ts_stamps_now(msg_type: type) -> None: - before = time.time() - assert before <= msg_type().ts <= time.time() - - -@pytest.mark.parametrize("msg_type", STAMPED_TYPES, ids=lambda t: t.__name__) -def test_explicit_ts_is_honoured(msg_type: type) -> None: - assert msg_type(ts=5.0).ts == 5.0 - - -def test_lcm_decode_of_an_unstamped_message_keeps_the_zero_stamp() -> None: - """A producer that never set the header now decodes to ts=0.0, not receive-time.""" - decoded = PoseStamped.lcm_decode(PoseStamped(ts=0.0, frame_id="odom").lcm_encode()) - assert decoded.ts == 0.0 - - -# Forms that used to build a silently-broken object -# -# Under plum these calls matched no overload and fell through to the LCM base -# __init__, which stores its arguments verbatim — leaving fields holding raw -# scalars/lists instead of Vector3/Quaternion/Pose, to blow up later in -# __repr__ or on encode. They now either do the obvious thing or raise. - - -def test_pose_from_scalar_is_an_x_only_position() -> None: - """Vector3(x) means (x, 0, 0), so Pose(5) is a position, not a raw 5.""" - assert Pose(5).position.to_tuple() == (5.0, 0.0, 0.0) - - -def test_pose_from_string_raises() -> None: - with pytest.raises(TypeError): - Pose("x") - - -def test_pose_from_two_numeric_positionals_raises() -> None: - """(1, 2) is a position and an orientation, and 2 is not a quaternion.""" - with pytest.raises(TypeError): - Pose(1, 2) - - -def test_pose_from_three_tuple_of_sequences_raises() -> None: - with pytest.raises((TypeError, ValueError)): - Pose(([1, 2, 3], [0, 0, 0, 1], [9])) - - -def test_quaternion_partial_positionals_raise() -> None: - """A quaternion needs all four components; 3 no longer zero-fills w.""" - for bad in ((1, 2, 3), (1, 2), (5,)): - with pytest.raises((TypeError, ValueError)): - Quaternion(*bad) - - -def test_quaternion_accepts_field_keywords() -> None: - assert Quaternion(x=1, y=2, z=3, w=4).to_tuple() == (1.0, 2.0, 3.0, 4.0) - - -def test_quaternion_partial_field_keywords_default_to_identity() -> None: - assert Quaternion(z=1).to_tuple() == (0.0, 0.0, 1.0, 1.0) - - -def test_quaternion_unknown_keyword_raises() -> None: - with pytest.raises(TypeError): - Quaternion(bogus=1) - - -def test_twist_from_single_sequence_is_linear_only() -> None: - t = Twist([1, 2, 3]) - assert t.linear.to_tuple() == (1.0, 2.0, 3.0) - assert t.angular.to_tuple() == (0.0, 0.0, 0.0) - - -def test_twist_from_scalars_raises() -> None: - with pytest.raises(TypeError): - Twist("x", "y") - - -def test_twist_mixed_positional_and_keyword_converts_both() -> None: - t = Twist([1, 2, 3], angular=[4, 5, 6]) - assert isinstance(t.linear, Vector3) - assert t.linear.to_tuple() == (1.0, 2.0, 3.0) - assert t.angular.to_tuple() == (4.0, 5.0, 6.0) - - -def test_twist_unknown_keyword_raises() -> None: - with pytest.raises(TypeError): - Twist(bogus=1) +@pytest.mark.parametrize("message_type", STAMPED_TYPES) +@pytest.mark.parametrize( + "stamp", + [Time(), Time(sec=5), Time(sec=-1, nanosec=999999999), Time(sec=1700000000, nanosec=123456789)], +) +def test_exact_stamps_survive_cdr(message_type, stamp) -> None: + message = message_type(header=Header(stamp=stamp, frame_id="odom")) + decoded = message_type.decode(message.encode()) + assert decoded.header == message.header + assert to_nanoseconds(decoded.header.stamp) == stamp.sec * 1000000000 + stamp.nanosec + + +@pytest.mark.parametrize("message_type", STAMPED_TYPES) +def test_omitted_stamp_is_zero_without_wall_clock(message_type) -> None: + message = message_type() + assert message.header.stamp == Time() + assert message_type.decode(message.encode()).header.stamp == Time() + + +def test_stamped_payloads_are_nested_values() -> None: + pose = Pose(position=Point(x=1, y=2, z=3), orientation=Quaternion(z=0.1, w=1)) + stamped = PoseStamped(header=Header(stamp=Time(sec=5), frame_id="odom"), pose=pose) + assert not isinstance(stamped, Pose) + assert stamped.pose == pose + assert deepcopy(stamped.pose) == pose + twist = Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=4, y=5, z=6)) + assert TwistStamped(twist=twist).twist == twist + pwc = PoseWithCovarianceStamped(pose=PoseWithCovariance(pose=pose, covariance=list(range(36)))) + assert pwc.pose.pose == pose + assert list(pwc.pose.covariance) == list(range(36)) + twc = TwistWithCovarianceStamped( + twist=TwistWithCovariance(twist=twist, covariance=list(range(36))) + ) + assert twc.twist.twist == twist + assert list(twc.twist.covariance) == list(range(36)) -def test_pose_unknown_keyword_raises() -> None: +@pytest.mark.parametrize( + "message_type", + [Pose, Quaternion, Twist, PoseWithCovariance, TwistWithCovariance, JointState, Joy], +) +def test_unknown_fields_rejected(message_type) -> None: with pytest.raises(TypeError): - Pose(bogus=1) - - -def test_pwc_pair_tuple_is_a_pose_and_a_covariance() -> None: - """The documented (pose, covariance) tuple form now actually fires.""" - p = PoseWithCovariance((Pose(1, 2, 3), list(range(36)))) - assert p.pose.position.to_tuple() == (1.0, 2.0, 3.0) - assert np.array_equal(p.covariance, np.arange(36.0)) - - -def test_pwc_accepts_field_keywords() -> None: - p = PoseWithCovariance(pose=Pose(1, 2, 3), covariance=list(range(36))) - assert p.pose.position.to_tuple() == (1.0, 2.0, 3.0) - assert np.array_equal(p.covariance, np.arange(36.0)) - - -def test_twc_accepts_field_keywords() -> None: - t = TwistWithCovariance(twist=Twist([1, 2, 3], [4, 5, 6])) - assert t.twist.linear.to_tuple() == (1.0, 2.0, 3.0) - - -def test_twc_pair_tuple_is_a_twist_and_a_covariance() -> None: - t = TwistWithCovariance((Twist([1, 2, 3], [4, 5, 6]), list(range(36)))) - assert t.twist.linear.to_tuple() == (1.0, 2.0, 3.0) - assert np.array_equal(t.covariance, np.arange(36.0)) - - -def test_pwc_dict_pose_may_be_any_pose_convertable() -> None: - p = PoseWithCovariance({"pose": [1, 2, 3], "covariance": list(range(36))}) - assert p.pose.position.to_tuple() == (1.0, 2.0, 3.0) - assert np.array_equal(p.covariance, np.arange(36.0)) - - -def test_jointstate_dict_carries_ts_and_frame_id() -> None: - j = JointState({"ts": 5.0, "frame_id": "f"}) - assert (j.ts, j.frame_id) == (5.0, "f") - assert j.name == [] - - -def test_joy_dict_carries_ts_and_frame_id() -> None: - j = Joy({"ts": 5.0, "frame_id": "f"}) - assert (j.ts, j.frame_id) == (5.0, "f") - assert j.axes == [] + message_type(bogus=1) + + +@pytest.mark.parametrize( + "message_type,kwargs", + [ + (Pose, {"position": [1, 2, 3]}), + (Pose, {"position": "x"}), + (Pose, {"orientation": [0, 0, 0, 1]}), + (Pose, {"position": None}), + (Twist, {"linear": [1, 2, 3]}), + (Twist, {"angular": Quaternion()}), + (PoseStamped, {"ts": 5}), + (PoseStamped, {"frame_id": "odom"}), + (PoseStamped, {"position": Point(x=1)}), + ], +) +def test_retired_polymorphic_fields_rejected(message_type, kwargs) -> None: + with pytest.raises((TypeError, RuntimeError)): + message_type(**kwargs) -def test_posestamped_string_positional_is_a_timestamp_not_a_position() -> None: - """A str first positional is a (bogus) ts; it no longer corrupts the position.""" - assert PoseStamped("f").position.to_tuple() == (0.0, 0.0, 0.0) +@pytest.mark.parametrize("values", [[], [1, 2], [1, 2, 3, 4], [[1, 2, 3], [4, 5, 6]]]) +def test_position_array_shape_rejected(values) -> None: + with pytest.raises(ValueError, match="3 vector components"): + point_from_array(values) diff --git a/dimos/msgs/test_occupancy.py b/dimos/msgs/test_occupancy.py new file mode 100644 index 00000000000..8773aa1a0f9 --- /dev/null +++ b/dimos/msgs/test_occupancy.py @@ -0,0 +1,146 @@ +# 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. + +import math + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid +from dimos_generated.std_msgs.msg import Header +import numpy as np +from PIL import Image +import pytest + +from dimos.mapping.occupancy.gradient import gradient +from dimos.msgs.occupancy import occupancy_extent, occupancy_from_file, occupancy_view + + +def test_occupancy_cells_keep_ros_row_order_and_do_not_alias_message(): + message = OccupancyGrid( + info=MapMetaData(width=3, height=2, resolution=0.1), data=[-1, 0, 100, 25, 50, 75] + ) + decoded = OccupancyGrid.decode(message.encode()) + values = occupancy_view(decoded) + np.testing.assert_array_equal(values, [[-1, 0, 100], [25, 50, 75]]) + with pytest.raises(ValueError, match="read-only"): + values[0, 0] = 0 + copied = values.copy() + copied[0, 0] = 0 + assert decoded.data[0] == -1 + + +def test_occupancy_rejects_dimensions_that_do_not_match_payload(): + message = OccupancyGrid(info=MapMetaData(width=2, height=2, resolution=1.0), data=[0]) + with pytest.raises(ValueError, match="data length"): + occupancy_view(message) + + +@pytest.mark.parametrize("resolution", [0.0, -0.1, math.nan, math.inf]) +def test_occupancy_rejects_invalid_resolution_for_nonempty_grid(resolution): + message = OccupancyGrid(info=MapMetaData(width=1, height=1, resolution=resolution), data=[0]) + with pytest.raises(ValueError, match="resolution"): + occupancy_extent(message) + + +def test_empty_default_occupancy_grid_has_no_cells(): + assert occupancy_view(OccupancyGrid()).shape == (0, 0) + + +@pytest.mark.parametrize("suffix", [".npy", ".png"]) +def test_static_occupancy_loader_preserves_cells_header_and_identity(tmp_path, suffix): + cells = np.array([[-1, 0, 100], [25, 50, 75]], dtype=np.int8) + path = tmp_path / ("grid" + suffix) + if suffix == ".npy": + np.save(path, cells) + else: + Image.fromarray(cells.astype(np.uint8)).save(path) + header = Header(frame_id="map", stamp=Time(sec=123, nanosec=456)) + output = OccupancyGrid.decode(occupancy_from_file(path, header=header, resolution=0.2).encode()) + np.testing.assert_array_equal(occupancy_view(output), cells) + assert output.header == header + assert output.info.map_load_time == header.stamp + assert output.info.origin.orientation.w == 1 + assert output.info.resolution == float(np.float32(0.2)) + + +@pytest.mark.parametrize( + "cells", + [np.zeros((2, 2), dtype=float), np.zeros(3, dtype=int), np.array([[101]]), np.array([[-2]])], +) +def test_static_occupancy_loader_rejects_invalid_cells(tmp_path, cells): + path = tmp_path / "grid.npy" + np.save(path, cells) + with pytest.raises(ValueError): + occupancy_from_file(path, header=Header()) + + +def test_gradient() -> None: + """Test converting occupancy grid to gradient field.""" + # Create a small test grid with an obstacle in the middle + data = np.zeros((10, 10), dtype=np.int8) + data[4:6, 4:6] = 100 # 2x2 obstacle in center + + grid = OccupancyGrid( + header=Header(frame_id="world"), + info=MapMetaData(width=10, height=10, resolution=0.1), + data=data.ravel(), + ) # 0.1m per cell + + # Convert to gradient + gradient_grid = gradient(grid, obstacle_threshold=50, max_distance=1.0) + + # Check that we get an OccupancyGrid back + assert isinstance(gradient_grid, OccupancyGrid) + assert occupancy_view(gradient_grid).shape == (10, 10) + assert gradient_grid.info.resolution == grid.info.resolution + assert gradient_grid.header == grid.header + + # Obstacle cells should have value 100 + assert occupancy_view(gradient_grid)[4, 4] == 100 + assert occupancy_view(gradient_grid)[5, 5] == 100 + + # Adjacent cells should have high values (near obstacles) + assert occupancy_view(gradient_grid)[3, 4] > 85 # Very close to obstacle + assert occupancy_view(gradient_grid)[4, 3] > 85 # Very close to obstacle + + # Cells at moderate distance should have moderate values + assert 30 < occupancy_view(gradient_grid)[0, 0] < 60 # Corner is ~0.57m away + + # Check that gradient decreases with distance + assert ( + occupancy_view(gradient_grid)[3, 4] > occupancy_view(gradient_grid)[2, 4] + ) # Closer is higher + assert ( + occupancy_view(gradient_grid)[2, 4] > occupancy_view(gradient_grid)[0, 4] + ) # Further is lower + + # Test with unknown cells + data_with_unknown = data.copy() + data_with_unknown[0:2, 0:2] = -1 # Add unknown area (close to obstacle) + data_with_unknown[8:10, 8:10] = -1 # Add unknown area (far from obstacle) + + grid_with_unknown = OccupancyGrid( + header=grid.header, info=grid.info, data=data_with_unknown.ravel() + ) + gradient_with_unknown = gradient(grid_with_unknown, max_distance=1.0) # 1m max distance + + # Unknown cells should remain unknown (new behavior - unknowns are preserved) + assert occupancy_view(gradient_with_unknown)[0, 0] == -1 # Should remain unknown + assert occupancy_view(gradient_with_unknown)[1, 1] == -1 # Should remain unknown + assert occupancy_view(gradient_with_unknown)[8, 8] == -1 # Should remain unknown + assert occupancy_view(gradient_with_unknown)[9, 9] == -1 # Should remain unknown + + # Unknown cells count should be preserved + assert ( + np.count_nonzero(occupancy_view(gradient_with_unknown) == -1) == 8 + ) # All unknowns preserved diff --git a/dimos/msgs/test_pointcloud.py b/dimos/msgs/test_pointcloud.py new file mode 100644 index 00000000000..e38565a6a2a --- /dev/null +++ b/dimos/msgs/test_pointcloud.py @@ -0,0 +1,384 @@ +# 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. + +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +import gc +import struct + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import PointCloud2, PointField +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest + +from dimos.msgs.camera_info import camera_info_from_intrinsics +from dimos.msgs.image import image_from_array +from dimos.msgs.pointcloud import ( + concatenate_clouds, + pointcloud_from_rgbd, + pointcloud_from_xyz, + pointcloud_from_xyz_rgb, + pointcloud_rgb, + pointcloud_view, + pointcloud_xyz, + select_points, + transform_cloud, + voxel_downsample_cloud, +) + + +@pytest.fixture(params=[False, True], ids=["little-endian", "big-endian"]) +def padded_cloud(request): + # Two organized rows, two points per row; fields are deliberately reordered. + endian = ">" if request.param else "<" + payload = bytearray(128) + for row in range(2): + for column in range(2): + value = row * 2 + column + 1 + struct.pack_into( + endian + "df2Hf", + payload, + row * 64 + column * 24, + value * 3.0, + value * 1.0, + value, + value + 10, + value * 2.0, + ) + return PointCloud2( + height=2, + width=2, + point_step=24, + row_step=64, + is_bigendian=request.param, + data=bytes(payload), + fields=[ + PointField(name="z", offset=0, datatype=8, count=1), + PointField(name="x", offset=8, datatype=7, count=1), + PointField(name="tags", offset=12, datatype=4, count=2), + PointField(name="y", offset=16, datatype=7, count=1), + ], + ) + + +def test_field_counts_padding_and_endianness(padded_cloud): + view = pointcloud_view(padded_cloud) + assert view["tags"].tolist() == [[[1, 11], [2, 12]], [[3, 13], [4, 14]]] + assert pointcloud_xyz(padded_cloud).tolist() == [[1, 2, 3], [2, 4, 6], [3, 6, 9], [4, 8, 12]] + assert not view.flags.writeable + with pytest.raises(ValueError, match="WRITEABLE"): + view.setflags(write=True) + copied = view.copy() + copied["x"][0, 0] = 99 + assert view["x"][0, 0] == 1 + + +def test_view_retains_storage_after_message_is_deleted(): + message = PointCloud2( + height=1, + width=1, + point_step=4, + row_step=4, + fields=[PointField(name="x", datatype=7, count=1)], + data=struct.pack("u2"), encoding="16UC1", header=header) + # Add nonzero row padding to big-endian depth. + payload = bytes(depth.data) + depth.data = payload[:6] + b"xx" + payload[6:] + b"yy" + depth.step = 8 + calibration = camera_info_from_intrinsics(4, 8, 2, 2, 6, 4, header=header) + cloud = PointCloud2.decode( + pointcloud_from_rgbd(color, depth, calibration, depth_scale=0.001, depth_trunc=4).encode() + ) + np.testing.assert_allclose( + pointcloud_xyz(cloud), [[-0.5, -0.25, 1], [1, -0.5, 2], [-1.5, 0, 3]] + ) + np.testing.assert_array_equal(pointcloud_rgb(cloud), rgb.reshape(-1, 3)[[0, 2, 3]]) + assert cloud.header == header + + +@pytest.mark.parametrize( + "scale,trunc", [(0, 1), (-1, 1), (float("nan"), 1), (1, 0), (1, float("inf"))] +) +def test_rgbd_rejects_invalid_units(scale, trunc): + header = Header() + color = image_from_array(np.zeros((1, 1, 3), dtype=np.uint8), encoding="rgb8", header=header) + depth = image_from_array(np.ones((1, 1), dtype=np.uint16), encoding="16UC1", header=header) + calibration = camera_info_from_intrinsics(1, 1, 0, 0, 1, 1, header=header) + with pytest.raises(ValueError, match="finite and positive"): + pointcloud_from_rgbd(color, depth, calibration, depth_scale=scale, depth_trunc=trunc) + + +def test_transform_rejects_coordinate_overflow(): + cloud = pointcloud_from_xyz(np.array([[3e38, 0, 0]]), header=Header(frame_id="child")) + transform = TransformStamped( + header=Header(frame_id="parent"), + child_frame_id="child", + transform=Transform(translation=Vector3(x=3e38), rotation=Quaternion(w=1)), + ) + with pytest.raises(ValueError, match="exceed field range"): + transform_cloud(cloud, transform) + assert np.isfinite(pointcloud_xyz(cloud)).all() + + +def test_voxel_downsample_retains_existing_small_cloud_fast_path(): + message = pointcloud_from_xyz(np.zeros((19, 3)), header=Header(frame_id="lidar")) + assert voxel_downsample_cloud(message, 0.5) is message + assert voxel_downsample_cloud(message, 0.0) is message + with pytest.raises(ValueError, match="finite"): + voxel_downsample_cloud(message, float("nan")) + + +def test_voxel_downsample_native_centroids_colors_and_exact_header(): + pytest.importorskip("open3d", reason="native tensor downsampling requires Open3D") + points = np.repeat([[0.1, 0.1, 0.1], [0.3, 0.3, 0.3]], 20, axis=0) + colors = np.repeat([[0, 0, 0], [255, 255, 255]], 20, axis=0).astype(np.uint8) + header = Header(frame_id="optical", stamp=Time(sec=5, nanosec=123456789)) + source = pointcloud_from_xyz_rgb(points, colors, header=header) + result = voxel_downsample_cloud(source, 1.0) + assert result.header == header + assert result.width == 1 + np.testing.assert_allclose(pointcloud_xyz(result), [[0.2, 0.2, 0.2]], atol=1e-6) + np.testing.assert_array_equal(pointcloud_rgb(result), [[127, 127, 127]]) + assert source.width == 40 diff --git a/dimos/msgs/test_pointcloud_colors.py b/dimos/msgs/test_pointcloud_colors.py new file mode 100644 index 00000000000..fe2ea97cdd3 --- /dev/null +++ b/dimos/msgs/test_pointcloud_colors.py @@ -0,0 +1,65 @@ +# 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. + +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +from dimos_generated.sensor_msgs.msg import PointCloud2, PointField +import numpy as np +import pytest + +from dimos.msgs.pointcloud import pointcloud_rgb + + +@pytest.mark.parametrize("bigendian", [False, True]) +@pytest.mark.parametrize("name", ["rgb", "rgba"]) +@pytest.mark.parametrize("datatype", [PointField.FLOAT32, PointField.UINT32]) +def test_packed_colors_preserve_bits_and_row_padding(bigendian, name, datatype): + order = ">" if bigendian else "<" + data = bytearray(48) + values = np.ndarray((2, 2), dtype=order + "u4", buffer=data, offset=4, strides=(24, 8)) + values[:] = [[0xFF123456, 0x00ABCDEF], [0x00FF0000, 0xFF00FF00]] + cloud = PointCloud2( + height=2, + width=2, + point_step=8, + row_step=24, + is_bigendian=bigendian, + data=bytes(data), + fields=[PointField(name=name, offset=4, count=1, datatype=datatype)], + ) + decoded = PointCloud2.decode(cloud.encode()) + np.testing.assert_array_equal( + pointcloud_rgb(decoded), [[0x12, 0x34, 0x56], [0xAB, 0xCD, 0xEF], [255, 0, 0], [0, 255, 0]] + ) + + +def test_missing_colors(): + assert pointcloud_rgb(PointCloud2()) is None + + +@pytest.mark.parametrize( + "datatype,count", [(PointField.FLOAT64, 1), (PointField.INT32, 1), (PointField.UINT32, 2)] +) +def test_invalid_color_layout(datatype, count): + cloud = PointCloud2( + height=1, + width=1, + point_step=8, + row_step=8, + data=bytes(8), + fields=[PointField(name="rgb", offset=0, count=count, datatype=datatype)], + ) + with pytest.raises(ValueError, match="packed point cloud color"): + pointcloud_rgb(cloud) diff --git a/dimos/msgs/test_time.py b/dimos/msgs/test_time.py new file mode 100644 index 00000000000..209e3b5653b --- /dev/null +++ b/dimos/msgs/test_time.py @@ -0,0 +1,81 @@ +# 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. + +import pytest + +from dimos.msgs.time import ( + duration_from_nanoseconds, + duration_from_seconds, + header_now, + time_from_nanoseconds, + time_from_seconds, + to_nanoseconds, + to_seconds, +) + + +@pytest.mark.parametrize("factory", [time_from_nanoseconds, duration_from_nanoseconds]) +@pytest.mark.parametrize("value", [0, -1, -1_250_000_000, 1_700_000_000_123_456_789]) +def test_integer_time_round_trip_preserves_every_nanosecond(factory, value): + stamp = factory(value) + + assert 0 <= stamp.nanosec < 1_000_000_000 + assert to_nanoseconds(stamp) == value + assert to_nanoseconds(type(stamp).decode(stamp.encode())) == value + + +@pytest.mark.parametrize("factory", [time_from_seconds, duration_from_seconds]) +@pytest.mark.parametrize("value,expected", [(-1.25, (-2, 750_000_000)), (0.9999999996, (1, 0))]) +def test_float_conversion_normalizes_negative_values_and_fractional_carry(factory, value, expected): + stamp = factory(value) + + assert (stamp.sec, stamp.nanosec) == expected + assert to_seconds(stamp) == pytest.approx(value, abs=1e-9) + + +@pytest.mark.parametrize("value", [float("nan"), float("inf"), -float("inf")]) +def test_nonfinite_seconds_are_rejected(value): + with pytest.raises(ValueError, match="finite"): + time_from_seconds(value) + + +def test_invalid_ros_fraction_is_rejected(): + stamp = time_from_nanoseconds(0) + stamp.nanosec = 1_000_000_000 + + with pytest.raises(ValueError, match="nanosec"): + to_nanoseconds(stamp) + + +def test_now_reads_integer_clock_and_preserves_frame(monkeypatch): + monkeypatch.setattr("dimos.msgs.time.time.time_ns", lambda: 1_700_000_000_123_456_789) + + header = header_now("camera") + + assert to_nanoseconds(header.stamp) == 1_700_000_000_123_456_789 + assert header.frame_id == "camera" + + +@pytest.mark.parametrize("factory", [time_from_nanoseconds, duration_from_nanoseconds]) +@pytest.mark.parametrize("seconds", [-(2**31), 2**31 - 1]) +def test_ros_seconds_range_boundaries(factory, seconds): + value = seconds * 1_000_000_000 + 999_999_999 + assert to_nanoseconds(factory(value)) == value + + +@pytest.mark.parametrize("factory", [time_from_nanoseconds, duration_from_nanoseconds]) +@pytest.mark.parametrize("value", [-(2**31) * 1_000_000_000 - 1, 2**31 * 1_000_000_000]) +def test_ros_seconds_overflow_is_explicit(factory, value): + with pytest.raises(ValueError, match="signed 32-bit"): + factory(value) diff --git a/dimos/msgs/test_trajectory.py b/dimos/msgs/test_trajectory.py new file mode 100644 index 00000000000..922fa3a2e4b --- /dev/null +++ b/dimos/msgs/test_trajectory.py @@ -0,0 +1,75 @@ +# 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. + +from dimos_generated.builtin_interfaces.msg import Duration, Time +from dimos_generated.dimos_msgs.msg import TrajectoryStatus +from dimos_generated.std_msgs.msg import Header +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import pytest + +from dimos.msgs.time import duration_from_seconds, to_nanoseconds +from dimos.msgs.trajectory import TrajectoryState, sample_trajectory, trajectory_duration + + +@pytest.mark.parametrize( + "seconds,expected", [(-1.0, [1.0, 2.0]), (0.5, [2.0, 1.0]), (3.0, [3.0, 0.0])] +) +@pytest.mark.parametrize("with_velocities", [False, True]) +def test_cdr_trajectory_sampling_and_copy_isolation(seconds, expected, with_velocities): + trajectory = JointTrajectory( + header=Header(frame_id="robot", stamp=Time(sec=1700000000, nanosec=123456789)), + joint_names=["a", "b"], + points=[ + JointTrajectoryPoint( + positions=[1.0, 2.0], velocities=[0.0, 1.0] if with_velocities else [] + ), + JointTrajectoryPoint( + positions=[3.0, 0.0], + velocities=[2.0, 3.0] if with_velocities else [], + time_from_start=duration_from_seconds(1.0), + ), + ], + ) + decoded = JointTrajectory.decode(trajectory.encode()) + positions, velocities = sample_trajectory(decoded, seconds) + assert positions == expected + alpha = max(0.0, min(seconds, 1.0)) + assert velocities == ([2 * alpha, 1 + 2 * alpha] if with_velocities else [0.0, 0.0]) + positions[0] = 999.0 + assert list(decoded.points[0].positions) == [1.0, 2.0] + assert trajectory_duration(decoded) == 1.0 + assert to_nanoseconds(decoded.header.stamp) == 1700000000123456789 + + +def test_empty_trajectory_sampling_and_invalid_sample_time(): + assert trajectory_duration(JointTrajectory()) == 0.0 + assert sample_trajectory(JointTrajectory(), 0.0) == ([], []) + with pytest.raises(ValueError, match="finite"): + sample_trajectory(JointTrajectory(), float("nan")) + + +def test_generated_status_roundtrips_duration_fields_and_state_constants(): + status = TrajectoryStatus( + header=Header(stamp=Time(sec=0, nanosec=7)), + state=TrajectoryState.EXECUTING, + progress=0.25, + time_elapsed=Duration(sec=1, nanosec=123456789), + time_remaining=Duration(sec=3, nanosec=987654321), + ) + decoded = TrajectoryStatus.decode(status.encode()) + assert decoded == status + assert decoded.state == TrajectoryStatus.EXECUTING + assert TrajectoryState(decoded.state).name == "EXECUTING" + assert to_nanoseconds(decoded.time_elapsed) == 1123456789 + assert to_nanoseconds(decoded.time_remaining) == 3987654321 diff --git a/dimos/msgs/tf2_msgs/test_TFMessage.py b/dimos/msgs/tf2_msgs/test_TFMessage.py index c379481f1dd..327ade99ba7 100644 --- a/dimos/msgs/tf2_msgs/test_TFMessage.py +++ b/dimos/msgs/tf2_msgs/test_TFMessage.py @@ -12,131 +12,102 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos_lcm.tf2_msgs import TFMessage as LCMTFMessage +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage +from rosbags.typesys import Stores, get_typestore -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.time import time_from_seconds def test_tfmessage_initialization() -> None: - """Test TFMessage initialization with Transform objects.""" - # Create some transforms - tf1 = Transform( - translation=Vector3(1, 2, 3), rotation=Quaternion(0, 0, 0, 1), frame_id="world", ts=100.0 + first = TransformStamped( + header=Header(stamp=time_from_seconds(100), frame_id="world"), + transform=Transform(translation=Vector3(x=1, y=2, z=3), rotation=Quaternion(w=1)), ) - tf2 = Transform( - translation=Vector3(4, 5, 6), - rotation=Quaternion(0, 0, 0.707, 0.707), - frame_id="map", - ts=101.0, + second = TransformStamped( + header=Header(stamp=time_from_seconds(101), frame_id="map"), + transform=Transform( + translation=Vector3(x=4, y=5, z=6), rotation=Quaternion(z=0.707, w=0.707) + ), ) - - # Create TFMessage with transforms - msg = TFMessage(tf1, tf2) - - assert len(msg) == 2 - assert msg[0] == tf1 - assert msg[1] == tf2 - - # Test iteration - transforms = list(msg) - assert transforms == [tf1, tf2] + message = TFMessage(transforms=[first, second]) + assert len(message.transforms) == 2 + assert [item.encode() for item in message.transforms] == [first.encode(), second.encode()] def test_tfmessage_empty() -> None: - """Test empty TFMessage.""" - msg = TFMessage() - assert len(msg) == 0 - assert list(msg) == [] + message = TFMessage() + assert len(message.transforms) == 0 + assert list(message.transforms) == [] def test_tfmessage_add_transform() -> None: - """Test adding transforms to TFMessage.""" - msg = TFMessage() - - tf = Transform(translation=Vector3(1, 2, 3), frame_id="base", ts=200.0) - - msg.add_transform(tf) - assert len(msg) == 1 - assert msg[0] == tf + message = TFMessage() + transform = TransformStamped( + header=Header(stamp=time_from_seconds(200), frame_id="base"), + transform=Transform(translation=Vector3(x=1, y=2, z=3)), + ) + message.transforms.append(transform) + assert len(message.transforms) == 1 + assert message.transforms[0].encode() == transform.encode() def test_tfmessage_tree() -> None: - """Test adding transforms to TFMessage.""" - msg = TFMessage() - - msg.add_transform( - Transform( - translation=Vector3(1, 2, 0.5), frame_id="world", child_frame_id="robot", ts=200.0 + message = TFMessage() + edges = [ + ("world", "robot", (1, 2, 0.5)), + ("robot", "camera", (0.1, 0, 1)), + ("robot", "lidar", (0.2, 0, 0.2)), + ("lidar", "lidar_scanner", (0.05, 0, 0)), + ] + for parent, child, (x, y, z) in edges: + message.transforms.append( + TransformStamped( + header=Header(stamp=time_from_seconds(200), frame_id=parent), + child_frame_id=child, + transform=Transform(translation=Vector3(x=x, y=y, z=z)), + ) ) - ) - msg.add_transform( - Transform( - translation=Vector3(0.1, 0, 1.0), frame_id="robot", child_frame_id="camera", ts=200.0 + assert len(message.transforms) == 4 + assert [ + ( + t.header.frame_id, + t.child_frame_id, + (t.transform.translation.x, t.transform.translation.y, t.transform.translation.z), ) - ) - msg.add_transform( - Transform( - translation=Vector3(0.2, 0, 0.2), frame_id="robot", child_frame_id="lidar", ts=200.0 - ) - ) - msg.add_transform( - Transform( - translation=Vector3(0.05, 0, 0.0), - frame_id="lidar", - child_frame_id="lidar_scanner", - ts=200.0, - ) - ) - - assert len(msg) == 4 - print(msg) + for t in TFMessage.decode(message.encode()).transforms + ] == edges -def test_tfmessage_lcm_encode_decode() -> None: - """Test encoding TFMessage to LCM bytes.""" - # Create transforms - tf1 = Transform( - translation=Vector3(1.0, 2.0, 3.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), +def test_tfmessage_cdr_encode_decode() -> None: + first = TransformStamped( + header=Header(stamp=time_from_seconds(123.456), frame_id="world"), child_frame_id="robot", - frame_id="world", - ts=123.456, + transform=Transform(translation=Vector3(x=1, y=2, z=3), rotation=Quaternion(w=1)), ) - tf2 = Transform( - translation=Vector3(4.0, 5.0, 6.0), - rotation=Quaternion(0.0, 0.0, 0.707, 0.707), - frame_id="robot", + second = TransformStamped( + header=Header(stamp=time_from_seconds(124.567), frame_id="robot"), child_frame_id="target", - ts=124.567, + transform=Transform( + translation=Vector3(x=4, y=5, z=6), rotation=Quaternion(z=0.707, w=0.707) + ), + ) + message = TFMessage(transforms=[first, second]) + decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(message.encode(), TFMessage.msg_name) + assert len(decoded.transforms) == 2 + first_result, second_result = decoded.transforms + assert first_result.header.frame_id == "world" + assert first_result.child_frame_id == "robot" + assert (first_result.header.stamp.sec, first_result.header.stamp.nanosec) == (123, 456000000) + assert ( + first_result.transform.translation.x, + first_result.transform.translation.y, + first_result.transform.translation.z, + ) == (1, 2, 3) + assert second_result.header.frame_id == "robot" + assert second_result.child_frame_id == "target" + assert (second_result.transform.rotation.z, second_result.transform.rotation.w) == ( + 0.707, + 0.707, ) - - # Create TFMessage - msg = TFMessage(tf1, tf2) - - # Encode with custom child_frame_ids - encoded = msg.lcm_encode() - - # Decode using LCM to verify - lcm_msg = LCMTFMessage.lcm_decode(encoded) - - assert lcm_msg.transforms_length == 2 - - # Check first transform - ts1 = lcm_msg.transforms[0] - assert ts1.header.frame_id == "world" - assert ts1.child_frame_id == "robot" - assert ts1.header.stamp.sec == 123 - assert ts1.header.stamp.nsec == 456000000 - assert ts1.transform.translation.x == 1.0 - assert ts1.transform.translation.y == 2.0 - assert ts1.transform.translation.z == 3.0 - - # Check second transform - ts2 = lcm_msg.transforms[1] - assert ts2.header.frame_id == "robot" - assert ts2.child_frame_id == "target" - assert ts2.transform.rotation.z == 0.707 - assert ts2.transform.rotation.w == 0.707 diff --git a/dimos/msgs/tf2_msgs/test_TFMessage_lcmpub.py b/dimos/msgs/tf2_msgs/test_TFMessage_lcmpub.py index 6df5e499fc5..c20cc55b291 100644 --- a/dimos/msgs/tf2_msgs/test_TFMessage_lcmpub.py +++ b/dimos/msgs/tf2_msgs/test_TFMessage_lcmpub.py @@ -15,12 +15,12 @@ from collections.abc import Iterator import time +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage import pytest -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.time import time_from_seconds from dimos.protocol.pubsub.impl.lcmpubsub import LCM, Topic from dimos.utils.testing.collector import CallbackCollector @@ -38,7 +38,7 @@ def lcm(lcm_url: str) -> Iterator[LCM]: # Publishes a series of transforms representing a robot kinematic chain # to actual LCM messages, rerun running in parallel should render this def test_publish_transforms(lcm: LCM) -> None: - topic = Topic(topic="/tf", lcm_type=TFMessage) + topic = Topic(topic="/tf", msg_type=TFMessage) collector = CallbackCollector(2) lcm.subscribe(topic, collector) @@ -46,42 +46,45 @@ def test_publish_transforms(lcm: LCM) -> None: current_time = time.time() # 1. World to base_link transform (robot at position) - world_to_base = Transform( - translation=Vector3(4.0, 3.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.382683, 0.923880), # 45 degrees around Z - frame_id="world", + world_to_base = TransformStamped( + transform=Transform( + translation=Vector3(x=4.0, y=3.0, z=0.0), + rotation=Quaternion(z=0.382683, w=0.923880), # 45 degrees around Z + ), + header=Header(frame_id="world", stamp=time_from_seconds(current_time)), child_frame_id="base_link", - ts=current_time, ) # 2. Base to arm transform (arm lifted up) - base_to_arm = Transform( - translation=Vector3(0.2, 0.0, 1.5), - rotation=Quaternion(0.0, 0.258819, 0.0, 0.965926), # 30 degrees around Y - frame_id="base_link", + base_to_arm = TransformStamped( + transform=Transform( + translation=Vector3(x=0.2, y=0.0, z=1.5), + rotation=Quaternion(y=0.258819, w=0.965926), # 30 degrees around Y + ), + header=Header(frame_id="base_link", stamp=time_from_seconds(current_time)), child_frame_id="arm_link", - ts=current_time, ) # 3. Arm to gripper transform (gripper extended) - arm_to_gripper = Transform( - translation=Vector3(0.5, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), # No rotation - frame_id="arm_link", + arm_to_gripper = TransformStamped( + transform=Transform( + translation=Vector3(x=0.5, y=0.0, z=0.0), + rotation=Quaternion(w=1.0), # No rotation + ), + header=Header(frame_id="arm_link", stamp=time_from_seconds(current_time)), child_frame_id="gripper_link", - ts=current_time, ) - lcm.publish(topic, TFMessage(world_to_base, base_to_arm)) - lcm.publish(topic, TFMessage(world_to_base, arm_to_gripper)) + lcm.publish(topic, TFMessage(transforms=[world_to_base, base_to_arm])) + lcm.publish(topic, TFMessage(transforms=[world_to_base, arm_to_gripper])) collector.wait() assert len(collector.results) == 2 first, _ = collector.results[0] assert isinstance(first, TFMessage) - assert [t.child_frame_id for t in first] == ["base_link", "arm_link"] + assert [t.child_frame_id for t in first.transforms] == ["base_link", "arm_link"] second, _ = collector.results[1] assert isinstance(second, TFMessage) - assert [t.child_frame_id for t in second] == ["base_link", "gripper_link"] + assert [t.child_frame_id for t in second.transforms] == ["base_link", "gripper_link"] diff --git a/dimos/msgs/time.py b/dimos/msgs/time.py new file mode 100644 index 00000000000..020fcd2f2ab --- /dev/null +++ b/dimos/msgs/time.py @@ -0,0 +1,70 @@ +# 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. + +"""Explicit conversions for generated ROS time, duration, and header fields.""" + +import math +import time + +from dimos_generated.builtin_interfaces.msg import Duration, Time +from dimos_generated.std_msgs.msg import Header + +NANOSECONDS_PER_SECOND = 1_000_000_000 + + +def to_nanoseconds(stamp: Time | Duration) -> int: + """Convert a normalized ROS time or duration without losing integer precision.""" + if not 0 <= stamp.nanosec < NANOSECONDS_PER_SECOND: + raise ValueError("nanosec must be in [0, 1000000000)") + return stamp.sec * NANOSECONDS_PER_SECOND + stamp.nanosec + + +def to_seconds(stamp: Time | Duration) -> float: + """Convert to floating seconds for APIs that require them.""" + return to_nanoseconds(stamp) / NANOSECONDS_PER_SECOND + + +def time_from_nanoseconds(value: int) -> Time: + sec, nanosec = divmod(value, NANOSECONDS_PER_SECOND) + if not -(2**31) <= sec < 2**31: + raise ValueError("seconds must fit the ROS signed 32-bit field") + return Time(sec=sec, nanosec=nanosec) + + +def duration_from_nanoseconds(value: int) -> Duration: + sec, nanosec = divmod(value, NANOSECONDS_PER_SECOND) + if not -(2**31) <= sec < 2**31: + raise ValueError("seconds must fit the ROS signed 32-bit field") + return Duration(sec=sec, nanosec=nanosec) + + +def _seconds_to_nanoseconds(value: float) -> int: + if not math.isfinite(value): + raise ValueError("seconds must be finite") + sec = math.floor(value) + return sec * NANOSECONDS_PER_SECOND + round((value - sec) * NANOSECONDS_PER_SECOND) + + +def time_from_seconds(value: float) -> Time: + """Round to the nearest nanosecond, normalizing negative and carried fractions.""" + return time_from_nanoseconds(_seconds_to_nanoseconds(value)) + + +def duration_from_seconds(value: float) -> Duration: + return duration_from_nanoseconds(_seconds_to_nanoseconds(value)) + + +def header_now(frame_id: str = "") -> Header: + """Construct a wall-clock header; generated defaults themselves remain zero.""" + return Header(stamp=time_from_nanoseconds(time.time_ns()), frame_id=frame_id) diff --git a/dimos/msgs/trajectory.py b/dimos/msgs/trajectory.py new file mode 100644 index 00000000000..02bc27861aa --- /dev/null +++ b/dimos/msgs/trajectory.py @@ -0,0 +1,79 @@ +# 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. + +"""Execution state and interpolation for generated ROS2 trajectories.""" + +from enum import IntEnum +import math + +from dimos_generated.dimos_msgs.msg import TrajectoryStatus +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint + +from dimos.msgs.time import to_seconds + + +class TrajectoryState(IntEnum): + """Internal state-machine names matching generated wire constants.""" + + IDLE = TrajectoryStatus.IDLE + EXECUTING = TrajectoryStatus.EXECUTING + COMPLETED = TrajectoryStatus.COMPLETED + ABORTED = TrajectoryStatus.ABORTED + FAULT = TrajectoryStatus.FAULT + + +def trajectory_duration(trajectory: JointTrajectory) -> float: + """Return the last waypoint time in seconds; validated paths have increasing times.""" + return to_seconds(trajectory.points[-1].time_from_start) if trajectory.points else 0.0 + + +def _sample_point(point: JointTrajectoryPoint) -> tuple[list[float], list[float]]: + return list(point.positions), list(point.velocities) or [0.0] * len(point.positions) + + +def sample_trajectory( + trajectory: JointTrajectory, seconds: float +) -> tuple[list[float], list[float]]: + """Linearly sample a validated path, clamping outside its waypoint time range. + + ROS permits absent velocities; this position controller's feed-forward then + defaults to zero. Inputs remain unchanged and returned arrays are independent. + """ + if not math.isfinite(seconds): + raise ValueError("sample time must be finite") + if not trajectory.points: + return [], [] + first, last = trajectory.points[0], trajectory.points[-1] + if seconds <= to_seconds(first.time_from_start): + return _sample_point(first) + if seconds >= to_seconds(last.time_from_start): + return _sample_point(last) + for index in range(len(trajectory.points) - 1): + before, after = trajectory.points[index], trajectory.points[index + 1] + start, end = to_seconds(before.time_from_start), to_seconds(after.time_from_start) + if start <= seconds <= end: + alpha = (seconds - start) / (end - start) + positions0, velocities0 = _sample_point(before) + positions1, velocities1 = _sample_point(after) + return ( + [ + left + alpha * (right - left) + for left, right in zip(positions0, positions1, strict=True) + ], + [ + left + alpha * (right - left) + for left, right in zip(velocities0, velocities1, strict=True) + ], + ) + raise ValueError("trajectory times must be increasing") diff --git a/dimos/msgs/vision_msgs/test_detection3d_rerun.py b/dimos/msgs/vision_msgs/test_detection3d_rerun.py index b0b4247052d..df75cfe4936 100644 --- a/dimos/msgs/vision_msgs/test_detection3d_rerun.py +++ b/dimos/msgs/vision_msgs/test_detection3d_rerun.py @@ -12,16 +12,20 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos_lcm.vision_msgs import BoundingBox3D, ObjectHypothesis, ObjectHypothesisWithPose +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion, Vector3 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.vision_msgs.msg import ( + BoundingBox3D, + Detection3D, + Detection3DArray, + ObjectHypothesis, + ObjectHypothesisWithPose, +) import pytest import rerun as rr -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.std_msgs.Header import Header -from dimos.msgs.vision_msgs.Detection3D import Detection3D -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.msgs.time import time_from_seconds +from dimos.visualization.rerun.message_helpers import detection_boxes def _detection3d( @@ -32,7 +36,7 @@ def _detection3d( class_id: str = "DICT_APRILTAG_36h11:7", ) -> Detection3D: det = Detection3D() - det.header = Header(ts, frame_id) + det.header = Header(stamp=time_from_seconds(ts), frame_id=frame_id) det.id = marker_id det.results = [ ObjectHypothesisWithPose( @@ -42,13 +46,12 @@ def _detection3d( ) ) ] - det.results_length = len(det.results) det.bbox = BoundingBox3D( center=Pose( - position=Vector3(1.0, 2.0, 3.0), - orientation=Quaternion(0.0, 0.0, 0.70710678, 0.70710678), + position=Point(x=1.0, y=2.0, z=3.0), + orientation=Quaternion(z=0.70710678, w=0.70710678), ), - size=Vector3(0.2, 0.4, 0.0), + size=Vector3(x=0.2, y=0.4), ) return det @@ -56,19 +59,18 @@ def _detection3d( def test_detection3d_frame_id_comes_from_header() -> None: det = _detection3d(frame_id="map") - assert det.frame_id == "map" + assert det.header.frame_id == "map" def test_detection3darray_to_rerun_preserves_wire_pose_size_and_identity() -> None: msg = Detection3DArray( - header=Header(12.5, "world"), + header=Header(stamp=time_from_seconds(12.5), frame_id="world"), detections=[_detection3d()], - detections_length=1, ) - boxes = msg.to_rerun() + boxes = detection_boxes(msg) - assert msg.frame_id == "world" + assert msg.header.frame_id == "world" assert isinstance(boxes, rr.Boxes3D) assert boxes.centers.as_arrow_array().to_pylist() == [[1.0, 2.0, 3.0]] assert boxes.half_sizes.as_arrow_array().to_pylist()[0] == pytest.approx([0.1, 0.2, 0.0]) @@ -80,12 +82,11 @@ def test_detection3darray_to_rerun_preserves_wire_pose_size_and_identity() -> No def test_detection3darray_to_rerun_empty_array_is_safe() -> None: msg = Detection3DArray( - header=Header(12.5, "world"), + header=Header(stamp=time_from_seconds(12.5), frame_id="world"), detections=[], - detections_length=0, ) - boxes = msg.to_rerun() + boxes = detection_boxes(msg) assert isinstance(boxes, rr.Boxes3D) assert boxes.centers.as_arrow_array().to_pylist() == [] diff --git a/dimos/navigation/base.py b/dimos/navigation/base.py index dc9cbbe9c27..affc0ad1a53 100644 --- a/dimos/navigation/base.py +++ b/dimos/navigation/base.py @@ -16,7 +16,7 @@ from abc import ABC, abstractmethod from enum import Enum -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos_generated.geometry_msgs.msg import PoseStamped class NavigationState(Enum): diff --git a/dimos/navigation/experimental/dannav/holonomic_tc/test_dan_holonomic_tc.py b/dimos/navigation/experimental/dannav/holonomic_tc/test_dan_holonomic_tc.py index 2ab9e4562a2..bdf0485cf44 100644 --- a/dimos/navigation/experimental/dannav/holonomic_tc/test_dan_holonomic_tc.py +++ b/dimos/navigation/experimental/dannav/holonomic_tc/test_dan_holonomic_tc.py @@ -23,14 +23,13 @@ import time from typing import Any, Literal -from dimos_lcm.std_msgs import Bool +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Twist +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Bool, Header import pytest from dimos.core.stream import Stream, Transport -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.time import time_from_seconds from dimos.navigation.base import NavigationState from dimos.navigation.experimental.dannav.holonomic_tc.module import DanHolonomicTC @@ -45,7 +44,7 @@ def __init__(self) -> None: def broadcast(self, _selfstream: Any, value: Any) -> None: for callback in list(self._subscribers): - callback(value) + callback(type(value).decode(value.encode())) def subscribe( self, callback: Callable[[Any], Any], _selfstream: Stream[Any] | None = None @@ -71,15 +70,13 @@ class _Captured: def _yaw_quaternion(yaw_rad: float) -> Quaternion: - return Quaternion(0.0, 0.0, math.sin(yaw_rad / 2.0), math.cos(yaw_rad / 2.0)) + return Quaternion(z=math.sin(yaw_rad / 2.0), w=math.cos(yaw_rad / 2.0)) def _odom(x: float, y: float, yaw_rad: float, *, ts: float = 1.0) -> PoseStamped: return PoseStamped( - ts=ts, - frame_id="map", - position=[x, y, 0.0], - orientation=_yaw_quaternion(yaw_rad), + header=Header(stamp=time_from_seconds(ts), frame_id="map"), + pose=Pose(position=Point(x=x, y=y), orientation=_yaw_quaternion(yaw_rad)), ) @@ -92,15 +89,8 @@ def _path_from_points(points: list[tuple[float, float]]) -> Path: else: prev_point = points[index - 1] yaw = math.atan2(point[1] - prev_point[1], point[0] - prev_point[0]) - poses.append( - PoseStamped( - ts=1.0, - frame_id="map", - position=[point[0], point[1], 0.0], - orientation=_yaw_quaternion(yaw), - ) - ) - return Path(frame_id="map", poses=poses) + poses.append(_odom(point[0], point[1], yaw)) + return Path(header=Header(frame_id="map"), poses=poses) def _is_zero_twist(cmd: Twist) -> bool: @@ -135,10 +125,10 @@ def feed_path(self, path: Path) -> None: self.module.path.transport.broadcast(None, path) def feed_empty_path(self) -> None: - self.module.path.transport.broadcast(None, Path(frame_id="map", poses=[])) + self.module.path.transport.broadcast(None, Path(header=Header(frame_id="map"), poses=[])) def feed_stop(self, value: bool = True) -> None: - self.module.stop_movement.transport.broadcast(None, Bool(value)) + self.module.stop_movement.transport.broadcast(None, Bool(data=value)) def close(self) -> None: for unsub in self._unsubs: diff --git a/dimos/navigation/experimental/frontier_exploration/wavefront_frontier_goal_selector.py b/dimos/navigation/experimental/frontier_exploration/wavefront_frontier_goal_selector.py index 93af3b398ab..f42128b3c32 100644 --- a/dimos/navigation/experimental/frontier_exploration/wavefront_frontier_goal_selector.py +++ b/dimos/navigation/experimental/frontier_exploration/wavefront_frontier_goal_selector.py @@ -25,7 +25,9 @@ import threading from typing import Any -from dimos_lcm.std_msgs import Bool +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Vector3 +from dimos_generated.nav_msgs.msg import OccupancyGrid +from dimos_generated.std_msgs.msg import Bool import numpy as np from reactivex.disposable import Disposable @@ -36,11 +38,9 @@ from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out from dimos.mapping.occupancy.inflation import simple_inflate -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid +from dimos.msgs.geometry import point_distance +from dimos.msgs.occupancy import grid_to_world, occupancy_view, world_to_grid from dimos.utils.logging_config import setup_logger -from dimos.utils.transform_utils import get_distance logger = setup_logger() @@ -110,7 +110,7 @@ class WavefrontFrontierExplorer(Module): config: WavefrontConfig - # LCM inputs + # Typed inputs global_costmap: In[OccupancyGrid] odom: In[PoseStamped] goal_reached: In[Bool] @@ -135,7 +135,7 @@ def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) self._cache = FrontierCache() self.explored_goals = [] # type: ignore[var-annotated] # list of explored goals - self.exploration_direction = Vector3(0.0, 0.0, 0.0) # current exploration direction + self.exploration_direction = Vector3(x=0.0, y=0.0, z=0.0) # current exploration direction self.last_costmap = None # store last costmap for information comparison self.no_gain_counter = 0 # track consecutive no-gain attempts @@ -223,8 +223,8 @@ def _count_costmap_information(self, costmap: OccupancyGrid) -> int: Returns: Number of cells that are free space or obstacles (not unknown) """ - free_count = np.sum(costmap.grid == CostValues.FREE) - obstacle_count = np.sum(costmap.grid >= self.config.occupancy_threshold) + free_count = np.sum(occupancy_view(costmap) == 0) + obstacle_count = np.sum(occupancy_view(costmap) >= self.config.occupancy_threshold) return int(free_count + obstacle_count) def _get_neighbors(self, point: GridPoint, costmap: OccupancyGrid) -> list[GridPoint]: @@ -240,7 +240,7 @@ def _get_neighbors(self, point: GridPoint, costmap: OccupancyGrid) -> list[GridP nx, ny = point.x + dx, point.y + dy # Check bounds - if 0 <= nx < costmap.width and 0 <= ny < costmap.height: + if 0 <= nx < costmap.info.width and 0 <= ny < costmap.info.height: neighbors.append(self._cache.get_point(nx, ny)) return neighbors @@ -252,21 +252,21 @@ def _is_frontier_point(self, point: GridPoint, costmap: OccupancyGrid) -> bool: and not adjacent to any occupied cells. """ # Point must be unknown - cost = costmap.grid[point.y, point.x] - if cost != CostValues.UNKNOWN: + cost = occupancy_view(costmap)[point.y, point.x] + if cost != -1: return False has_free = False for neighbor in self._get_neighbors(point, costmap): - neighbor_cost = costmap.grid[neighbor.y, neighbor.x] + neighbor_cost = occupancy_view(costmap)[neighbor.y, neighbor.x] # If adjacent to occupied space, not a frontier if neighbor_cost > self.config.occupancy_threshold: return False # Check if adjacent to free space - if neighbor_cost == CostValues.FREE: + if neighbor_cost == 0: has_free = True return has_free @@ -288,7 +288,7 @@ def _find_free_space( visited.add((point.x, point.y)) # Check if this point is free space - if costmap.grid[point.y, point.x] == CostValues.FREE: + if occupancy_view(costmap)[point.y, point.x] == 0: return (point.x, point.y) # Add neighbors to search @@ -299,18 +299,18 @@ def _find_free_space( # If no free space found, return original position return (start_x, start_y) - def _compute_centroid(self, frontier_points: list[Vector3]) -> Vector3: + def _compute_centroid(self, frontier_points: list[Point]) -> Point: """Compute the centroid of a list of frontier points.""" if not frontier_points: - return Vector3(0.0, 0.0, 0.0) + return Point(x=0.0, y=0.0, z=0.0) # Vectorized approach using numpy - points_array = np.array([[point.x, point.y] for point in frontier_points]) + points_array = np.array([[point.x, point.y, point.z] for point in frontier_points]) centroid = np.mean(points_array, axis=0) - return Vector3(centroid[0], centroid[1], 0.0) + return Point(x=centroid[0], y=centroid[1], z=centroid[2]) - def detect_frontiers(self, robot_pose: Vector3, costmap: OccupancyGrid) -> list[Vector3]: + def detect_frontiers(self, robot_pose: Point, costmap: OccupancyGrid) -> list[Point]: """ Main frontier detection algorithm using wavefront exploration. @@ -324,8 +324,8 @@ def detect_frontiers(self, robot_pose: Vector3, costmap: OccupancyGrid) -> list[ self._cache.clear() # Convert robot pose to grid coordinates - grid_pos = costmap.world_to_grid(robot_pose) - grid_x, grid_y = int(grid_pos.x), int(grid_pos.y) + grid_pos = world_to_grid(costmap, robot_pose) + grid_x, grid_y = int(np.floor(round(grid_pos[0], 9))), int(np.floor(round(grid_pos[1], 9))) # Find nearest free space to start exploration free_x, free_y = self._find_free_space(grid_x, grid_y, costmap) @@ -386,13 +386,13 @@ def detect_frontiers(self, robot_pose: Vector3, costmap: OccupancyGrid) -> list[ # Check if we found a large enough frontier # Convert minimum perimeter to minimum number of cells based on resolution - min_cells = int(self.config.min_frontier_perimeter / costmap.resolution) + min_cells = max( + 1, int(self.config.min_frontier_perimeter / costmap.info.resolution) + ) if len(new_frontier) >= min_cells: world_points = [] for point in new_frontier: - world_pos = costmap.grid_to_world( - Vector3(float(point.x), float(point.y), 0.0) - ) + world_pos = grid_to_world(costmap, (float(point.x), float(point.y))) world_points.append(world_pos) # Compute centroid in world coordinates (already correctly scaled) @@ -407,10 +407,10 @@ def detect_frontiers(self, robot_pose: Vector3, costmap: OccupancyGrid) -> list[ & (PointClassification.MapOpen | PointClassification.MapClosed) ): # Check if neighbor is free space or unknown (explorable) - neighbor_cost = costmap.grid[neighbor.y, neighbor.x] + neighbor_cost = occupancy_view(costmap)[neighbor.y, neighbor.x] # Add free space and unknown space to exploration queue - if neighbor_cost == CostValues.FREE or neighbor_cost == CostValues.UNKNOWN: + if neighbor_cost == 0 or neighbor_cost == -1: neighbor.classification |= PointClassification.MapOpen map_queue.append(neighbor) @@ -428,25 +428,27 @@ def detect_frontiers(self, robot_pose: Vector3, costmap: OccupancyGrid) -> list[ return ranked_frontiers def _update_exploration_direction( - self, robot_pose: Vector3, goal_pose: Vector3 | None = None + self, robot_pose: Point, goal_pose: Point | None = None ) -> None: """Update the current exploration direction based on robot movement or selected goal.""" if goal_pose is not None: # Calculate direction from robot to goal - direction = Vector3(goal_pose.x - robot_pose.x, goal_pose.y - robot_pose.y, 0.0) + direction = Vector3(x=goal_pose.x - robot_pose.x, y=goal_pose.y - robot_pose.y, z=0.0) magnitude = np.sqrt(direction.x**2 + direction.y**2) if magnitude > 0.1: # Avoid division by zero for very close goals self.exploration_direction = Vector3( - direction.x / magnitude, direction.y / magnitude, 0.0 + x=direction.x / magnitude, y=direction.y / magnitude, z=0.0 ) - def _compute_direction_momentum_score(self, frontier: Vector3, robot_pose: Vector3) -> float: + def _compute_direction_momentum_score(self, frontier: Point, robot_pose: Point) -> float: """Compute direction momentum score for a frontier.""" if self.exploration_direction.x == 0 and self.exploration_direction.y == 0: return 0.0 # No momentum if no previous direction # Calculate direction from robot to frontier - frontier_direction = Vector3(frontier.x - robot_pose.x, frontier.y - robot_pose.y, 0.0) + frontier_direction = Vector3( + x=frontier.x - robot_pose.x, y=frontier.y - robot_pose.y, z=0.0 + ) magnitude = np.sqrt(frontier_direction.x**2 + frontier_direction.y**2) if magnitude < 0.1: @@ -454,7 +456,7 @@ def _compute_direction_momentum_score(self, frontier: Vector3, robot_pose: Vecto # Normalize frontier direction frontier_direction = Vector3( - frontier_direction.x / magnitude, frontier_direction.y / magnitude, 0.0 + x=frontier_direction.x / magnitude, y=frontier_direction.y / magnitude, z=0.0 ) # Calculate dot product for directional alignment @@ -466,7 +468,7 @@ def _compute_direction_momentum_score(self, frontier: Vector3, robot_pose: Vecto # Return momentum score (higher for same direction, lower for opposite) return max(0.0, dot_product) # Only positive momentum, no penalty for different directions - def _compute_distance_to_explored_goals(self, frontier: Vector3) -> float: + def _compute_distance_to_explored_goals(self, frontier: Point) -> float: """Compute distance from frontier to the nearest explored goal.""" if not self.explored_goals: return 5.0 # Default consistent value when no explored goals @@ -478,7 +480,7 @@ def _compute_distance_to_explored_goals(self, frontier: Vector3) -> float: return min_distance - def _compute_distance_to_obstacles(self, frontier: Vector3, costmap: OccupancyGrid) -> float: + def _compute_distance_to_obstacles(self, frontier: Point, costmap: OccupancyGrid) -> float: """ Compute the minimum distance from a frontier point to the nearest obstacle. @@ -490,16 +492,21 @@ def _compute_distance_to_obstacles(self, frontier: Vector3, costmap: OccupancyGr Minimum distance to nearest obstacle in meters """ # Convert frontier to grid coordinates - grid_pos = costmap.world_to_grid(frontier) - grid_x, grid_y = int(grid_pos.x), int(grid_pos.y) + grid_pos = world_to_grid(costmap, frontier) + grid_x, grid_y = int(np.floor(round(grid_pos[0], 9))), int(np.floor(round(grid_pos[1], 9))) # Check if frontier is within costmap bounds - if grid_x < 0 or grid_x >= costmap.width or grid_y < 0 or grid_y >= costmap.height: + if ( + grid_x < 0 + or grid_x >= costmap.info.width + or grid_y < 0 + or grid_y >= costmap.info.height + ): return 0.0 # Consider out-of-bounds as obstacle min_distance = float("inf") search_radius = ( - int(self.config.safe_distance / costmap.resolution) + 5 + int(self.config.safe_distance / costmap.info.resolution) + 5 ) # Search a bit beyond minimum # Search in a square around the frontier point @@ -511,16 +518,16 @@ def _compute_distance_to_obstacles(self, frontier: Vector3, costmap: OccupancyGr # Skip if out of bounds if ( check_x < 0 - or check_x >= costmap.width + or check_x >= costmap.info.width or check_y < 0 - or check_y >= costmap.height + or check_y >= costmap.info.height ): continue # Check if this cell is an obstacle - if costmap.grid[check_y, check_x] >= self.config.occupancy_threshold: + if occupancy_view(costmap)[check_y, check_x] >= self.config.occupancy_threshold: # Calculate distance in meters - distance = np.sqrt(dx**2 + dy**2) * costmap.resolution + distance = np.sqrt(dx**2 + dy**2) * costmap.info.resolution min_distance = min(min_distance, distance) # If no obstacles found within search radius, return the safe distance @@ -528,12 +535,12 @@ def _compute_distance_to_obstacles(self, frontier: Vector3, costmap: OccupancyGr return min_distance if min_distance != float("inf") else self.config.safe_distance def _compute_comprehensive_frontier_score( - self, frontier: Vector3, frontier_size: int, robot_pose: Vector3, costmap: OccupancyGrid + self, frontier: Point, frontier_size: int, robot_pose: Point, costmap: OccupancyGrid ) -> float: """Compute comprehensive score considering multiple criteria.""" # 1. Distance from robot (preference for moderate distances) - robot_distance = get_distance(frontier, robot_pose) + robot_distance = point_distance(frontier, robot_pose) # Distance score: prefer moderate distances (not too close, not too far) # Normalized to 0-1 range @@ -541,7 +548,9 @@ def _compute_comprehensive_frontier_score( # 2. Information gain (frontier size) # Normalize by a reasonable max frontier size - max_expected_frontier_size = self.config.min_frontier_perimeter / costmap.resolution * 10 + max_expected_frontier_size = ( + self.config.min_frontier_perimeter / costmap.info.resolution * 10 + ) info_gain_score = min(frontier_size / max_expected_frontier_size, 1.0) # 3. Distance to explored goals (bonus for being far from explored areas) @@ -577,11 +586,11 @@ def _compute_comprehensive_frontier_score( def _rank_frontiers( self, - frontier_centroids: list[Vector3], + frontier_centroids: list[Point], frontier_sizes: list[int], - robot_pose: Vector3, + robot_pose: Point, costmap: OccupancyGrid, - ) -> list[Vector3]: + ) -> list[Point]: """ Find the single best frontier using comprehensive scoring and filtering. @@ -619,7 +628,7 @@ def _rank_frontiers( # Extract just the frontiers (remove scores) and return as list return [frontier for frontier, _ in valid_frontiers] - def get_exploration_goal(self, robot_pose: Vector3, costmap: OccupancyGrid) -> Vector3 | None: + def get_exploration_goal(self, robot_pose: Point, costmap: OccupancyGrid) -> Point | None: """ Get the single best exploration goal using comprehensive frontier scoring. @@ -683,7 +692,7 @@ def get_exploration_goal(self, robot_pose: Vector3, costmap: OccupancyGrid) -> V self.last_costmap = costmap # type: ignore[assignment] return None - def mark_explored_goal(self, goal: Vector3) -> None: + def mark_explored_goal(self, goal: Point) -> None: """Mark a goal as explored.""" self.explored_goals.append(goal) @@ -695,7 +704,7 @@ def reset_exploration_session(self) -> None: needs to forget its previous exploration history. """ self.explored_goals.clear() # Clear all previously explored goals - self.exploration_direction = Vector3(0.0, 0.0, 0.0) # Reset exploration direction + self.exploration_direction = Vector3(x=0.0, y=0.0, z=0.0) # Reset exploration direction self.last_costmap = None # Clear last costmap comparison self.no_gain_counter = 0 # Reset no-gain attempt counter self._cache.clear() # Clear frontier point cache @@ -749,12 +758,7 @@ def stop_exploration(self) -> bool: # 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, - ) + goal = PoseStamped(header=self.latest_odometry.header, pose=self.latest_odometry.pose) self.goal_request.publish(goal) logger.info("Stopped autonomous frontier exploration") @@ -791,8 +795,10 @@ def _run_exploration_loop(self) -> None: continue # Get robot pose from odometry - robot_pose = Vector3( - self.latest_odometry.position.x, self.latest_odometry.position.y, 0.0 + robot_pose = Point( + x=self.latest_odometry.pose.position.x, + y=self.latest_odometry.pose.position.y, + z=0.0, ) # Get exploration goal @@ -801,13 +807,10 @@ def _run_exploration_loop(self) -> None: if goal: # Publish goal to navigator - goal_msg = PoseStamped() - goal_msg.position.x = goal.x - goal_msg.position.y = goal.y - goal_msg.position.z = 0.0 - goal_msg.orientation.w = 1.0 # No rotation - goal_msg.frame_id = "world" - goal_msg.ts = self.latest_costmap.ts + goal_msg = PoseStamped( + header=self.latest_costmap.header, + pose=Pose(position=goal, orientation=Quaternion(w=1)), + ) self.goal_request.publish(goal_msg) logger.info(f"Published frontier goal: ({goal.x:.2f}, {goal.y:.2f})") diff --git a/dimos/navigation/experimental/patrolling/routers/patrol_router.py b/dimos/navigation/experimental/patrolling/routers/patrol_router.py index 19aff4ca34a..6209dade381 100644 --- a/dimos/navigation/experimental/patrolling/routers/patrol_router.py +++ b/dimos/navigation/experimental/patrolling/routers/patrol_router.py @@ -14,8 +14,8 @@ from typing import Protocol -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.nav_msgs.msg import OccupancyGrid class PatrolRouter(Protocol): diff --git a/dimos/navigation/experimental/patrolling/routers/visitation_history.py b/dimos/navigation/experimental/patrolling/routers/visitation_history.py index 939da19ddbb..e3d9a65e015 100644 --- a/dimos/navigation/experimental/patrolling/routers/visitation_history.py +++ b/dimos/navigation/experimental/patrolling/routers/visitation_history.py @@ -12,10 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.geometry_msgs.msg import Point +from dimos_generated.nav_msgs.msg import OccupancyGrid import numpy as np from numpy.typing import NDArray -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid +from dimos.msgs.occupancy import occupancy_view, world_to_grid def _circular_disk(radius_cells: int) -> NDArray[np.bool_]: @@ -57,7 +59,7 @@ def clearance_disk(self) -> NDArray[np.bool_]: def update_grid(self, grid: OccupancyGrid) -> None: self._grid = grid - self._clearance_radius_cells = int(np.ceil(self._clearance_radius_m / grid.resolution)) + self._clearance_radius_cells = int(np.ceil(self._clearance_radius_m / grid.info.resolution)) self._clearance_disk = _circular_disk(self._clearance_radius_cells) self._rebuild() @@ -80,7 +82,7 @@ def get_saturation(self) -> float: visited = self._visited if grid is None or visited is None: return 0.0 - free_mask = grid.grid == 0 + free_mask = occupancy_view(grid) == 0 total = int(np.count_nonzero(free_mask)) if total == 0: return 0.0 @@ -96,7 +98,7 @@ def _rebuild(self) -> None: grid = self._grid if grid is None: return - self._visited = np.zeros((grid.height, grid.width), dtype=bool) + self._visited = np.zeros((grid.info.height, grid.info.width), dtype=bool) for x, y in self._points: self._stamp(x, y) @@ -106,14 +108,14 @@ def _stamp(self, x: float, y: float) -> None: if grid is None or visited is None: return r = self._clearance_radius_cells - grid_pos = grid.world_to_grid((x, y)) - col, row = int(grid_pos.x), int(grid_pos.y) - if row + r < 0 or row - r >= grid.height or col + r < 0 or col - r >= grid.width: + grid_pos = world_to_grid(grid, Point(x=x, y=y)) + col, row = int(np.floor(round(grid_pos[0], 9))), int(np.floor(round(grid_pos[1], 9))) + if row + r < 0 or row - r >= grid.info.height or col + r < 0 or col - r >= grid.info.width: return r_min = max(0, row - r) - r_max = min(grid.height, row + r + 1) + r_max = min(grid.info.height, row + r + 1) c_min = max(0, col - r) - c_max = min(grid.width, col + r + 1) + c_max = min(grid.info.width, col + r + 1) d_r_min = r_min - (row - r) d_r_max = d_r_min + (r_max - r_min) d_c_min = c_min - (col - r) diff --git a/dimos/navigation/experimental/patrolling/utilities.py b/dimos/navigation/experimental/patrolling/utilities.py index d7caffaa9c4..df9805bc85e 100644 --- a/dimos/navigation/experimental/patrolling/utilities.py +++ b/dimos/navigation/experimental/patrolling/utilities.py @@ -12,15 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.std_msgs.msg import Header -def point_to_pose_stamped(point: tuple[float, float]) -> PoseStamped: - pose = PoseStamped() - pose.position.x = point[0] - pose.position.y = point[1] - return pose +def point_to_pose_stamped(point: Point, header: Header) -> PoseStamped: + return PoseStamped(header=header, pose=Pose(position=point, orientation=Quaternion(w=1))) def pose_stamped_to_point(pose: PoseStamped) -> tuple[float, float]: - return (pose.position.x, pose.position.y) + return (pose.pose.position.x, pose.pose.position.y) diff --git a/dimos/navigation/global_planner/mls_planner/start_relay.py b/dimos/navigation/global_planner/mls_planner/start_relay.py index e55aeb5856a..04aedd9b93b 100644 --- a/dimos/navigation/global_planner/mls_planner/start_relay.py +++ b/dimos/navigation/global_planner/mls_planner/start_relay.py @@ -17,13 +17,13 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.tf2_msgs.msg import TFMessage from reactivex.disposable import Disposable from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.tf2_msgs.TFMessage import TFMessage class StartRelayConfig(ModuleConfig): diff --git a/dimos/navigation/go2/loop_closure/pgo_auto.py b/dimos/navigation/go2/loop_closure/pgo_auto.py index 1a5fe756d9f..9e172c1f6f2 100644 --- a/dimos/navigation/go2/loop_closure/pgo_auto.py +++ b/dimos/navigation/go2/loop_closure/pgo_auto.py @@ -18,14 +18,14 @@ estimate of where it is in the world slowly diverges from ground truth as small per-frame errors accumulate. PGO ("pose graph optimization") detects revisited places via loop closure and pulls the trajectory back into -self-consistency. The result is a stream of *corrections*: a Transform per +self-consistency. The result is a stream of *corrections*: a TransformStamped per keyframe that maps every drifted ("raw") pose to its corrected one. Pipeline: lidar: Stream[PointCloud2] # pose-stamped scans -> pgo_keyframes(...) -> Stream[Keyframe] # one per kf - -> keyframes_to_corrections(...) -> Stream[Transform] # world_corrected <- world_raw + -> keyframes_to_corrections(...) -> Stream[TransformStamped] # world_corrected <- world_raw -> apply_corrections(any_stream, corrections) -> Stream[T] # obs.pose shuffled `pgo_keyframes` runs ISAM2 + ICP loop closure across the input, emitting @@ -35,11 +35,11 @@ optimized poses + ICP fitness) — useful for visualizing the pose graph. `keyframes_to_corrections` is pure arithmetic: per keyframe, the drift -correction is `optimized @ local^-1` (Transform composition). The resulting +correction is `optimized @ local^-1` (TransformStamped composition). The resulting stream has one entry per keyframe — sparse in time. `make_interpolator` materializes the sparse correction stream into a -ts -> Transform function: SLERP for rotation, linear lerp for translation, +ts -> TransformStamped function: SLERP for rotation, linear lerp for translation, clipping out-of-range queries to the endpoints. `apply_corrections` shuffles `obs.pose` on any stream using the interpolated @@ -58,16 +58,30 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any, TypedDict, TypeVar +from dimos_generated.geometry_msgs.msg import Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header import numpy as np from scipy.spatial.transform import Rotation from dimos.memory.store.memory import MemoryStore from dimos.memory.stream import Stream from dimos.memory.type.observation import Observation -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.geometry import ( + compose_transforms, + inverse_transform, + quaternion_from_matrix, + transform_from_pose, + transform_matrix, +) +from dimos.msgs.pointcloud import ( + concatenate_clouds, + pointcloud_from_xyz, + pointcloud_to_open3d, + transform_cloud, + voxel_downsample_cloud, +) +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.protocol.service.spec import BaseConfig from dimos.utils.logging_config import setup_logger @@ -142,8 +156,8 @@ class Keyframe: """ ts: float - local: Transform - optimized: Transform + local: TransformStamped + optimized: TransformStamped @dataclass(frozen=True) @@ -156,8 +170,8 @@ class LoopClosure: tighter match. """ - source: Transform - target: Transform + source: TransformStamped + target: TransformStamped score: float @@ -199,27 +213,26 @@ def pgo_keyframes( return out -def keyframes_to_corrections(keyframes: Stream[Keyframe]) -> Stream[Transform]: - """Per-keyframe drift correction as Transform(world_corrected <- world_raw).""" +def keyframes_to_corrections(keyframes: Stream[Keyframe]) -> Stream[TransformStamped]: + """Per-keyframe drift correction as TransformStamped(world_corrected <- world_raw).""" mem = MemoryStore() - out: Stream[Transform] = mem.stream("corrections", Transform) + out: Stream[TransformStamped] = mem.stream("corrections", TransformStamped) for obs in keyframes: kf = obs.data # world_corrected <- body <- world_raw composes to world_corrected <- world_raw. - drift = kf.optimized + kf.local.inverse() + drift = compose_transforms(kf.optimized, inverse_transform(kf.local)) out.append(drift, ts=kf.ts) return out -def make_interpolator(corrections: Stream[Transform]) -> Callable[[float], Transform]: - """Materialize corrections once; return a fast ts -> Transform lookup.""" +def make_interpolator(corrections: Stream[TransformStamped]) -> Callable[[float], TransformStamped]: + """Materialize corrections once; return a fast ts -> TransformStamped lookup.""" ts_list: list[float] = [] R_list: list[np.ndarray] = [] t_list: list[np.ndarray] = [] for obs in corrections: R, t = _r_t_from_transform(obs.data) - # obs.ts is authoritative; obs.data.ts can be mutated by Transform's - # ts=0.0 -> time.time() fallback in its constructor. + # Observation time is the correction lookup key. ts_list.append(obs.ts) R_list.append(R) t_list.append(t) @@ -238,7 +251,7 @@ def make_interpolator(corrections: Stream[Transform]) -> Callable[[float], Trans R_stack = np.stack(R_list) t_stack = np.stack(t_list) - def interp(ts: float) -> Transform: + def interp(ts: float) -> TransformStamped: ts_clip = float(np.clip(ts, ts_arr[0], ts_arr[-1])) idx = int(np.searchsorted(ts_arr, ts_clip)) if idx == 0: @@ -261,7 +274,7 @@ def interp(ts: float) -> Transform: def apply_corrections( stream: Stream[T], - corrections: Stream[Transform], + corrections: Stream[TransformStamped], ) -> Stream[T]: """Shuffle obs.pose on `stream` by the interpolated correction at each obs.ts. @@ -275,10 +288,12 @@ def xf(upstream: Iterator[Observation[T]]) -> Iterator[Observation[T]]: if obs.pose is None: yield obs continue - raw_tf = Transform.from_pose(FRAME_BODY, obs.pose_stamped) - # Transform.__add__ composes: (T_corr + T_raw) applies T_corr after T_raw. - # Observation normalizes Transform back to 7-tuple via __post_init__. - corrected = interp(obs.ts) + raw_tf + pose = obs.pose_stamped + assert pose is not None + pose.header.frame_id = FRAME_WORLD_RAW + raw_tf = transform_from_pose(pose, child_frame_id=FRAME_BODY) + # Observation normalizes the generated transform to its storage pose tuple. + corrected = compose_transforms(interp(obs.ts), raw_tf) yield obs.derive(data=obs.data, pose=corrected) return stream.transform(xf) @@ -297,19 +312,19 @@ def __init__( self, keyframes: tuple[Keyframe, ...], loops: tuple[LoopClosure, ...], - correct: Callable[[float], Transform], + correct: Callable[[float], TransformStamped], ) -> None: self.keyframes = keyframes self.loops = loops self._correct = correct - def correction_at(self, ts: float) -> Transform: + def correction_at(self, ts: float) -> TransformStamped: """Drift correction (world_corrected <- world_raw) at `ts`.""" return self._correct(ts) - def correct(self, pose: Transform) -> Transform: + def correct(self, pose: TransformStamped) -> TransformStamped: """Lift a world_raw pose into world_corrected.""" - return self._correct(pose.ts) + pose + return compose_transforms(self._correct(to_seconds(pose.header.stamp)), pose) class PGO: @@ -333,27 +348,28 @@ def __call__( last_ts = keyframes[-1].ts else: # No valid keyframes (e.g. pose-less input): correction is identity. - correct = lambda ts: Transform( # noqa: E731 - frame_id=FRAME_WORLD_CORRECTED, child_frame_id=FRAME_WORLD_RAW, ts=ts - ) + correct = lambda ts: _transform_from_r_t(np.eye(3), np.zeros(3), ts=ts) # noqa: E731 last_ts = 0.0 yield Observation(ts=last_ts, _data=PoseGraph(keyframes, tuple(loops), correct)) -def _r_t_from_transform(tf: Transform) -> tuple[np.ndarray, np.ndarray]: - q = tf.rotation +def _r_t_from_transform(tf: TransformStamped) -> tuple[np.ndarray, np.ndarray]: + q = tf.transform.rotation R = Rotation.from_quat([q.x, q.y, q.z, q.w]).as_matrix() - t = np.array([tf.translation.x, tf.translation.y, tf.translation.z]) + t = np.array( + [tf.transform.translation.x, tf.transform.translation.y, tf.transform.translation.z] + ) return R, t -def _transform_from_r_t(R: np.ndarray, t: np.ndarray, *, ts: float) -> Transform: - return Transform( - translation=Vector3(float(t[0]), float(t[1]), float(t[2])), - rotation=Quaternion.from_rotation_matrix(R), - frame_id=FRAME_WORLD_CORRECTED, +def _transform_from_r_t(R: np.ndarray, t: np.ndarray, *, ts: float) -> TransformStamped: + return TransformStamped( + header=Header(frame_id=FRAME_WORLD_CORRECTED, stamp=time_from_seconds(ts)), child_frame_id=FRAME_WORLD_RAW, - ts=ts, + transform=Transform( + translation=Vector3(x=float(t[0]), y=float(t[1]), z=float(t[2])), + rotation=quaternion_from_matrix(R), + ), ) @@ -419,15 +435,24 @@ def process( ts: float, world_cloud: PointCloud2, ) -> None: - if len(world_cloud) == 0: + if world_cloud.width * world_cloud.height == 0: return if not self._is_keyframe(local_pose): return # Unregister: lift world-frame scan back into body frame using its # odom pose, so PGO can re-project it via the optimized pose later. - body_cloud = world_cloud.transform( - _pose3_to_transform(local_pose.inverse(), ts=ts) - ).voxel_downsample(self._cfg.submap_resolution) + body_cloud = voxel_downsample_cloud( + transform_cloud( + world_cloud, + _pose3_to_transform( + local_pose.inverse(), + ts=ts, + frame_id=FRAME_BODY, + child_frame_id=world_cloud.header.frame_id, + ), + ), + self._cfg.submap_resolution, + ) self._add_keyframe(local_pose, ts, body_cloud) self._search_for_loops() self._smooth_and_update() @@ -558,16 +583,25 @@ def _get_submap(self, idx: int, half_range: int) -> PointCloud2: lo = max(0, idx - half_range) hi = min(len(self._key_poses) - 1, idx + half_range) if lo > hi: - return PointCloud2() - cloud = self._key_poses[lo].body_cloud.transform( - _pose3_to_transform(self._key_poses[lo].optimized, ts=self._key_poses[lo].timestamp) - ) - for i in range(lo + 1, hi + 1): - kp = self._key_poses[i] - cloud = cloud + kp.body_cloud.transform( - _pose3_to_transform(kp.optimized, ts=kp.timestamp) + return pointcloud_from_xyz( + np.empty((0, 3)), header=Header(frame_id=FRAME_WORLD_CORRECTED) ) - return cloud.voxel_downsample(self._cfg.submap_resolution) + + def registered(kp: _KeyPose) -> PointCloud2: + return transform_cloud( + kp.body_cloud, + _pose3_to_transform( + kp.optimized, + ts=kp.timestamp, + frame_id=FRAME_WORLD_CORRECTED, + child_frame_id=FRAME_BODY, + ), + ) + + cloud = registered(self._key_poses[lo]) + for i in range(lo + 1, hi + 1): + cloud = concatenate_clouds(cloud, registered(self._key_poses[i])) + return voxel_downsample_cloud(cloud, self._cfg.submap_resolution) def _search_for_loops( self, @@ -668,7 +702,7 @@ def _collect(local_idxs: set[int], time_thresh: float) -> list[tuple[float, int] # already and extra loops add noise. Early-exit on a very tight # match to save redundant ICP runs in the common case. source = self._get_submap(cur_idx, source_half_range) - accepted: list[tuple[int, Transform, float]] = [] + accepted: list[tuple[int, TransformStamped, float]] = [] k = ( candidates_per_iter if candidates_per_iter is not None @@ -771,23 +805,24 @@ def _pose3_to_transform( ts: float, frame_id: str = "", child_frame_id: str = "", -) -> Transform: - """PGO-internal: build a Transform from a Pose3.""" +) -> TransformStamped: + """PGO-internal: build a TransformStamped from a Pose3.""" t = np.asarray(pose.translation()) - return Transform( - translation=Vector3(float(t[0]), float(t[1]), float(t[2])), - rotation=Quaternion.from_rotation_matrix(pose.rotation().matrix()), - frame_id=frame_id, + return TransformStamped( + header=Header(frame_id=frame_id, stamp=time_from_seconds(ts)), child_frame_id=child_frame_id, - ts=ts, + transform=Transform( + translation=Vector3(x=float(t[0]), y=float(t[1]), z=float(t[2])), + rotation=quaternion_from_matrix(pose.rotation().matrix()), + ), ) -def _transform_to_pose3(tf: Transform) -> gtsam.Pose3: - """PGO-internal: build a Pose3 from a Transform.""" +def _transform_to_pose3(tf: TransformStamped) -> gtsam.Pose3: + """PGO-internal: build a Pose3 from a TransformStamped.""" import gtsam # type: ignore[import-not-found,import-untyped] - return gtsam.Pose3(tf.to_matrix()) + return gtsam.Pose3(transform_matrix(tf.transform)) def _icp( @@ -797,8 +832,8 @@ def _icp( max_dist: float = 1.0, tol: float = 1e-6, min_inliers: int = 10, - init: Transform | None = None, -) -> tuple[Transform, float]: + init: TransformStamped | None = None, +) -> tuple[TransformStamped, float]: """Point-to-plane ICP using Open3D's tensor pipeline. Returns ``(tf, fitness)`` where ``fitness`` is mean squared inlier @@ -809,11 +844,13 @@ def _icp( import open3d as o3d # type: ignore[import-untyped] import open3d.core as o3c # type: ignore[import-untyped] - if len(source) < min_inliers or len(target) < min_inliers: - return Transform.identity(), float("inf") + if source.width * source.height < min_inliers or target.width * target.height < min_inliers: + return _transform_from_r_t( + np.eye(3), np.zeros(3), ts=to_seconds(source.header.stamp) + ), float("inf") - src_pcd = source.pointcloud_tensor - tgt_pcd = target.pointcloud_tensor + src_pcd = o3d.t.geometry.PointCloud.from_legacy(pointcloud_to_open3d(source)) + tgt_pcd = o3d.t.geometry.PointCloud.from_legacy(pointcloud_to_open3d(target)) # Normals on the target enable point-to-plane ICP — converges tighter # than point-to-point on indoor scenes (walls give unambiguous normals @@ -822,7 +859,7 @@ def _icp( device = src_pcd.device init_T = o3c.Tensor( - init.to_matrix() if init is not None else np.eye(4), + transform_matrix(init.transform) if init is not None else np.eye(4), dtype=o3c.float64, device=device, ) @@ -845,11 +882,17 @@ def _icp( ), ) except RuntimeError: - return Transform.identity(), float("inf") + return _transform_from_r_t( + np.eye(3), np.zeros(3), ts=to_seconds(source.header.stamp) + ), float("inf") if float(result.fitness) == 0.0: - return Transform.identity(), float("inf") + return _transform_from_r_t( + np.eye(3), np.zeros(3), ts=to_seconds(source.header.stamp) + ), float("inf") T_mat = result.transformation.numpy() rmse = float(result.inlier_rmse) - return _transform_from_r_t(T_mat[:3, :3], T_mat[:3, 3], ts=source.ts), rmse * rmse + return _transform_from_r_t( + T_mat[:3, :3], T_mat[:3, 3], ts=to_seconds(source.header.stamp) + ), rmse * rmse diff --git a/dimos/navigation/go2/loop_closure/utils/markers_rrd.py b/dimos/navigation/go2/loop_closure/utils/markers_rrd.py index 0a54e097373..36158eb2f09 100644 --- a/dimos/navigation/go2/loop_closure/utils/markers_rrd.py +++ b/dimos/navigation/go2/loop_closure/utils/markers_rrd.py @@ -30,17 +30,21 @@ from pathlib import Path +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 import typer from dimos.memory.store.sqlite import SqliteStore from dimos.memory.stream import Stream from dimos.memory.transform import QualityWindow, SpeedLimit from dimos.memory.vis.color import Color -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.image import image_sharpness from dimos.perception.fiducial.marker_transformer import DetectMarkers -from dimos.robot.unitree.go2.connection import _camera_info_static +from dimos.robot.unitree.go2.camera_calibration import front_camera_calibration from dimos.utils.data import resolve_named_path +from dimos.visualization.rerun.message_helpers import ( + camera_pinhole, + image_archetype, +) TIMELINE = "ts" @@ -65,14 +69,14 @@ def main( import rerun as rr db_path = resolve_named_path(dataset, ".db") - cam_info = _camera_info_static() + cam_info = front_camera_calibration() rr.init("dimos markers", recording_id=db_path.stem) rr.save(str(out)) # Static pinhole on the camera entity; per-frame Transform3D goes on the # same entity. Image is the child so it projects through the pinhole. - pinhole = cam_info.to_rerun() + pinhole = camera_pinhole(cam_info) assert not isinstance(pinhole, list) rr.log("world/camera", pinhole, static=True) @@ -106,7 +110,7 @@ def main( translation=[x, y, z], quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw]) ), ) - rr.log("world/camera/image", img_obs.data.to_rerun()) + rr.log("world/camera/image", image_archetype(img_obs.data)) # Clear any prior detection box at this ts; pass 3 will overwrite # this with the actual detection iff one fires for this frame. rr.log( @@ -119,7 +123,7 @@ def main( # Pass 3: marker detections, filtered the same way as `dimos map`. xf = DetectMarkers(camera_info=cam_info, marker_length_m=marker_size) pipeline: Stream[Image] = color_image.transform( - QualityWindow(lambda img: img.sharpness, window=quality_window) + QualityWindow(image_sharpness, window=quality_window) ) if marker_max_speed > 0: pipeline = pipeline.transform( @@ -185,7 +189,7 @@ def main( smoothing_window=smoothing_window, ) pipeline_tracked: Stream[Image] = color_image.transform( - QualityWindow(lambda img: img.sharpness, window=quality_window) + QualityWindow(image_sharpness, window=quality_window) ) if marker_max_speed > 0: pipeline_tracked = pipeline_tracked.transform( diff --git a/dimos/navigation/go2/replanning_a_star/controllers.py b/dimos/navigation/go2/replanning_a_star/controllers.py index c9902e278a4..1b7f269d096 100644 --- a/dimos/navigation/go2/replanning_a_star/controllers.py +++ b/dimos/navigation/go2/replanning_a_star/controllers.py @@ -15,13 +15,12 @@ import math from typing import Protocol +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist, Vector3 import numpy as np from numpy.typing import NDArray from dimos.core.global_config import GlobalConfig -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.geometry import quaternion_euler from dimos.utils.trigonometry import angle_diff @@ -52,7 +51,7 @@ def __init__(self, global_config: GlobalConfig, speed: float, control_frequency: self._control_frequency = control_frequency def advance(self, lookahead_point: NDArray[np.float64], current_odom: PoseStamped) -> Twist: - current_pos = np.array([current_odom.position.x, current_odom.position.y]) + current_pos = np.array([current_odom.pose.position.x, current_odom.pose.position.y]) direction = lookahead_point - current_pos distance = np.linalg.norm(direction) @@ -60,7 +59,7 @@ def advance(self, lookahead_point: NDArray[np.float64], current_odom: PoseStampe # Robot is coincidentally at the lookahead point; skip this cycle. return Twist() - robot_yaw = current_odom.orientation.euler[2] + robot_yaw = quaternion_euler(current_odom.pose.orientation)[2] desired_yaw = np.arctan2(direction[1], direction[0]) yaw_error = angle_diff(desired_yaw, robot_yaw) @@ -75,8 +74,8 @@ def advance(self, lookahead_point: NDArray[np.float64], current_odom: PoseStampe linear_velocity = self._apply_min_velocity(linear_velocity, self._min_linear_velocity) return Twist( - linear=Vector3(linear_velocity, 0.0, 0.0), - angular=Vector3(0.0, 0.0, angular_velocity), + linear=Vector3(x=linear_velocity), + angular=Vector3(z=angular_velocity), ) def rotate(self, yaw_error: float) -> Twist: @@ -109,8 +108,8 @@ def _angular_twist(self, angular_velocity: float) -> Twist: angular_velocity = 0.8 * np.sign(angular_velocity) return Twist( - linear=Vector3(0.0, 0.0, 0.0), - angular=Vector3(0.0, 0.0, angular_velocity), + linear=Vector3(), + angular=Vector3(z=angular_velocity), ) diff --git a/dimos/navigation/go2/replanning_a_star/global_planner.py b/dimos/navigation/go2/replanning_a_star/global_planner.py index 5d03dde01f7..b1c39ee55fc 100644 --- a/dimos/navigation/go2/replanning_a_star/global_planner.py +++ b/dimos/navigation/go2/replanning_a_star/global_planner.py @@ -16,7 +16,9 @@ from threading import Event, RLock, Thread, current_thread import time -from dimos_lcm.std_msgs import Bool +from dimos_generated.geometry_msgs.msg import Point, PoseStamped, Twist +from dimos_generated.nav_msgs.msg import OccupancyGrid, Path +from dimos_generated.std_msgs.msg import Bool import numpy as np from reactivex import Subject from reactivex.disposable import CompositeDisposable @@ -25,11 +27,8 @@ from dimos.core.global_config import GlobalConfig from dimos.core.resource import Resource from dimos.mapping.occupancy.path_resampling import smooth_resample_path -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import point_distance, quaternion_euler +from dimos.msgs.occupancy import occupancy_view, world_to_grid from dimos.navigation.base import NavigationState from dimos.navigation.go2.replanning_a_star.goal_validator import find_safe_goal from dimos.navigation.go2.replanning_a_star.local_planner import LocalPlanner, StopMessage @@ -167,7 +166,7 @@ def cancel_goal(self, *, but_will_try_again: bool = False, arrived: bool = False self._local_planner.stop_planning() if not but_will_try_again: - self.goal_reached.on_next(Bool(arrived)) + self.goal_reached.on_next(Bool(data=arrived)) def set_replanning_enabled(self, enabled: bool) -> None: with self._lock: @@ -221,9 +220,13 @@ def _thread_entrypoint(self) -> None: continue if ( - current_goal.position.distance(current_odom.position) < self._goal_tolerance + point_distance(current_goal.pose.position, current_odom.pose.position) + < self._goal_tolerance and abs( - angle_diff(current_goal.orientation.euler[2], current_odom.orientation.euler[2]) + angle_diff( + quaternion_euler(current_goal.pose.orientation)[2], + quaternion_euler(current_odom.pose.orientation)[2], + ) ) < self._rotation_tolerance ): @@ -293,7 +296,10 @@ def _replan_path(self) -> None: assert current_odom is not None assert current_goal is not None - if current_goal.position.distance(current_odom.position) < self._replan_goal_tolerance: + if ( + point_distance(current_goal.pose.position, current_odom.pose.position) + < self._replan_goal_tolerance + ): self.cancel_goal(arrived=True) return @@ -301,7 +307,7 @@ def _replan_path(self) -> None: self.cancel_goal() return - if not self._replan_limiter.can_retry(current_odom.position): + if not self._replan_limiter.can_retry(current_odom.pose.position): self.cancel_goal() return @@ -322,16 +328,18 @@ def _plan_path(self) -> None: logger.warning("Cannot handle goal request: missing odometry.") return - safe_goal = self._find_safe_goal(current_goal.position) + safe_goal = self._find_safe_goal(current_goal.pose.position) if not safe_goal: logger.warning( - "No safe goal found.", x=round(current_goal.x, 3), y=round(current_goal.y, 3) + "No safe goal found.", + x=round(current_goal.pose.position.x, 3), + y=round(current_goal.pose.position.y, 3), ) self.cancel_goal() return - path = self._find_wide_path(safe_goal, current_odom.position) + path = self._find_wide_path(safe_goal, current_odom.pose.position) if not path: logger.warning( @@ -340,18 +348,18 @@ def _plan_path(self) -> None: self.cancel_goal() return - resampled_path = smooth_resample_path(path, current_goal, 0.1) + resampled_path = smooth_resample_path(path, current_goal.pose, 0.1) self.path.on_next(resampled_path) self._local_planner.start_planning(resampled_path) - def _find_wide_path(self, goal: Vector3, robot_pos: Vector3) -> Path | None: + def _find_wide_path(self, goal: Point, robot_pos: Point) -> Path | None: # sizes_to_try: list[float] = [2.2, 1.7, 1.3, 1] sizes_to_try: list[float] = [1.1] for size in sizes_to_try: - distance = robot_pos.distance(goal) + distance = point_distance(robot_pos, goal) navigation_map = self._navigation_map if distance > 1.5 else self._navigation_map_near costmap = navigation_map.make_gradient_costmap(size) self._clear_robot_footprint(costmap, navigation_map.binary_costmap, robot_pos) @@ -363,7 +371,7 @@ def _find_wide_path(self, goal: Vector3, robot_pos: Vector3) -> Path | None: return None def _clear_robot_footprint( - self, costmap: OccupancyGrid, binary: OccupancyGrid, robot_pos: Vector3 + self, costmap: OccupancyGrid, binary: OccupancyGrid, robot_pos: Point ) -> None: """Make the cells under the robot passable for planning. @@ -376,39 +384,54 @@ def _clear_robot_footprint( the raw binary costmap (actually observed obstacle points) stay blocked. """ - if binary.grid.shape != costmap.grid.shape or binary.origin != costmap.origin: + if ( + occupancy_view(binary).shape != occupancy_view(costmap).shape + or binary.info.origin != costmap.info.origin + or binary.info.resolution != costmap.info.resolution + or binary.header.frame_id != costmap.header.frame_id + ): # A newer map update raced in between building the two grids; # skip clearing rather than misalign cells. return - center = costmap.world_to_grid(robot_pos) - center_x, center_y = int(center.x), int(center.y) - cells = int(self._global_config.robot_rotation_diameter / 2 / costmap.resolution) + 1 + center = world_to_grid(costmap, robot_pos) + center_x, center_y = int(np.floor(round(center[0], 9))), int(np.floor(round(center[1], 9))) + cells = int(self._global_config.robot_rotation_diameter / 2 / costmap.info.resolution) + 1 - height, width = costmap.grid.shape + height, width = occupancy_view(costmap).shape y0, y1 = max(0, center_y - cells), min(height, center_y + cells + 1) x0, x1 = max(0, center_x - cells), min(width, center_x + cells + 1) if y0 >= y1 or x0 >= x1: return - region = costmap.grid[y0:y1, x0:x1] - binary_region = binary.grid[y0:y1, x0:x1] + mutable_cells = occupancy_view(costmap).copy() + region = mutable_cells[y0:y1, x0:x1] + binary_region = occupancy_view(binary)[y0:y1, x0:x1] rows, columns = np.ogrid[y0:y1, x0:x1] disc = (rows - center_y) ** 2 + (columns - center_x) ** 2 <= cells**2 - clearable = disc & (region >= CostValues.OCCUPIED) & (binary_region < CostValues.OCCUPIED) - region[clearable] = CostValues.OCCUPIED - 1 - - def _find_safe_goal(self, goal: Vector3) -> Vector3 | None: + clearable = disc & (region >= 100) & (binary_region < 100) + region[clearable] = 99 + costmap.data = mutable_cells.ravel() + + @staticmethod + def _goal_is_unknown(costmap: OccupancyGrid, goal: Point) -> bool: + x, y = world_to_grid(costmap, goal) + gx, gy = math.floor(round(x, 9)), math.floor(round(y, 9)) + if not (0 <= gx < costmap.info.width and 0 <= gy < costmap.info.height): + return True + return bool(occupancy_view(costmap)[gy, gx] == -1) + + def _find_safe_goal(self, goal: Point) -> Point | None: costmap = self._navigation_map.binary_costmap - if costmap.cell_value(goal) == CostValues.UNKNOWN: + if self._goal_is_unknown(costmap, goal): return goal safe_goal = find_safe_goal( costmap, goal, algorithm="bfs_contiguous", - cost_threshold=CostValues.OCCUPIED, + cost_threshold=100, min_clearance=self._safe_goal_clearance, max_search_distance=self._safe_goal_tolerance, ) @@ -417,7 +440,7 @@ def _find_safe_goal(self, goal: Vector3) -> Vector3 | None: logger.warning("No safe goal found near requested target.") return None - goals_distance = safe_goal.distance(goal) + goals_distance = point_distance(safe_goal, goal) if goals_distance > 0.2: logger.warning(f"Travelling to goal {goals_distance}m away from requested goal.") diff --git a/dimos/navigation/go2/replanning_a_star/goal_validator.py b/dimos/navigation/go2/replanning_a_star/goal_validator.py index b717c762952..cb20eb7342a 100644 --- a/dimos/navigation/go2/replanning_a_star/goal_validator.py +++ b/dimos/navigation/go2/replanning_a_star/goal_validator.py @@ -14,21 +14,22 @@ from collections import deque +from dimos_generated.geometry_msgs.msg import Point +from dimos_generated.nav_msgs.msg import OccupancyGrid import numpy as np -from dimos.msgs.geometry_msgs.Vector3 import Vector3, VectorLike -from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid +from dimos.msgs.occupancy import grid_to_world, occupancy_view, world_to_grid def find_safe_goal( costmap: OccupancyGrid, - goal: VectorLike, + goal: Point, algorithm: str = "bfs", cost_threshold: int = 50, min_clearance: float = 0.3, max_search_distance: float = 5.0, connectivity_check_radius: int = 3, -) -> Vector3 | None: +) -> Point | None: """ Find a safe goal position when the original goal is in collision or too close to obstacles. @@ -69,12 +70,12 @@ def find_safe_goal( def _find_safe_goal_bfs( costmap: OccupancyGrid, - goal: VectorLike, + goal: Point, cost_threshold: int, min_clearance: float, max_search_distance: float, connectivity_check_radius: int, -) -> Vector3 | None: +) -> Point | None: """ BFS-based search for nearest safe goal position. This guarantees finding the closest valid position. @@ -90,12 +91,12 @@ def _find_safe_goal_bfs( """ # Convert goal to grid coordinates - goal_grid = costmap.world_to_grid(goal) - gx, gy = int(goal_grid.x), int(goal_grid.y) + goal_grid = world_to_grid(costmap, goal) + gx, gy = int(np.floor(round(goal_grid[0], 9))), int(np.floor(round(goal_grid[1], 9))) # Convert distances to grid cells - clearance_cells = int(np.ceil(min_clearance / costmap.resolution)) - max_search_cells = int(np.ceil(max_search_distance / costmap.resolution)) + clearance_cells = int(np.ceil(min_clearance / costmap.info.resolution)) + max_search_cells = int(np.ceil(max_search_distance / costmap.info.resolution)) # BFS queue and visited set queue = deque([(gx, gy, 0)]) @@ -116,14 +117,14 @@ def _find_safe_goal_bfs( costmap, x, y, cost_threshold, clearance_cells, connectivity_check_radius ): # Convert back to world coordinates - return costmap.grid_to_world((x, y)) + return grid_to_world(costmap, (x, y)) # Add neighbors to queue for dx, dy in neighbors: nx, ny = x + dx, y + dy # Check bounds - if 0 <= nx < costmap.width and 0 <= ny < costmap.height: + if 0 <= nx < costmap.info.width and 0 <= ny < costmap.info.height: if (nx, ny) not in visited: visited.add((nx, ny)) queue.append((nx, ny, dist + 1)) @@ -133,12 +134,12 @@ def _find_safe_goal_bfs( def _find_safe_goal_bfs_contiguous( costmap: OccupancyGrid, - goal: VectorLike, + goal: Point, cost_threshold: int, min_clearance: float, max_search_distance: float, connectivity_check_radius: int, -) -> Vector3 | None: +) -> Point | None: """ BFS-based search for nearest safe goal position, only following passable cells. Unlike regular BFS, this only expands through cells with occupancy < 100, @@ -154,13 +155,15 @@ def _find_safe_goal_bfs_contiguous( - Slightly slower than regular BFS due to additional checks """ + grid = occupancy_view(costmap) + # Convert goal to grid coordinates - goal_grid = costmap.world_to_grid(goal) - gx, gy = int(goal_grid.x), int(goal_grid.y) + goal_grid = world_to_grid(costmap, goal) + gx, gy = int(np.floor(round(goal_grid[0], 9))), int(np.floor(round(goal_grid[1], 9))) # Convert distances to grid cells - clearance_cells = int(np.ceil(min_clearance / costmap.resolution)) - max_search_cells = int(np.ceil(max_search_distance / costmap.resolution)) + clearance_cells = int(np.ceil(min_clearance / costmap.info.resolution)) + max_search_cells = int(np.ceil(max_search_distance / costmap.info.resolution)) # BFS queue and visited set queue = deque([(gx, gy, 0)]) @@ -181,17 +184,17 @@ def _find_safe_goal_bfs_contiguous( costmap, x, y, cost_threshold, clearance_cells, connectivity_check_radius ): # Convert back to world coordinates - return costmap.grid_to_world((x, y)) + return grid_to_world(costmap, (x, y)) # Add neighbors to queue for dx, dy in neighbors: nx, ny = x + dx, y + dy # Check bounds - if 0 <= nx < costmap.width and 0 <= ny < costmap.height: + if 0 <= nx < costmap.info.width and 0 <= ny < costmap.info.height: if (nx, ny) not in visited: # Only expand through passable cells (occupancy < 100) - if costmap.grid[ny, nx] < 100: + if grid[ny, nx] < 100: visited.add((nx, ny)) queue.append((nx, ny, dist + 1)) @@ -220,22 +223,24 @@ def _is_position_safe( True if position is safe, False otherwise """ + grid = occupancy_view(costmap) + # Check bounds first - if not (0 <= x < costmap.width and 0 <= y < costmap.height): + if not (0 <= x < costmap.info.width and 0 <= y < costmap.info.height): return False # Check if position itself is free - if costmap.grid[y, x] >= cost_threshold or costmap.grid[y, x] == CostValues.UNKNOWN: + if grid[y, x] >= cost_threshold or grid[y, x] == -1: return False # Check clearance around position for dy in range(-clearance_cells, clearance_cells + 1): for dx in range(-clearance_cells, clearance_cells + 1): nx, ny = x + dx, y + dy - if 0 <= nx < costmap.width and 0 <= ny < costmap.height: + if 0 <= nx < costmap.info.width and 0 <= ny < costmap.info.height: # Check if within circular clearance if dx * dx + dy * dy <= clearance_cells * clearance_cells: - if costmap.grid[ny, nx] >= cost_threshold: + if grid[ny, nx] >= cost_threshold: return False # Check connectivity (not surrounded by obstacles) @@ -249,12 +254,9 @@ def _is_position_safe( continue nx, ny = x + dx, y + dy - if 0 <= nx < costmap.width and 0 <= ny < costmap.height: + if 0 <= nx < costmap.info.width and 0 <= ny < costmap.info.height: total_count += 1 - if ( - costmap.grid[ny, nx] < cost_threshold - and costmap.grid[ny, nx] != CostValues.UNKNOWN - ): + if grid[ny, nx] < cost_threshold and grid[ny, nx] != -1: free_count += 1 # Require at least 50% of neighbors to be free (not surrounded) diff --git a/dimos/navigation/go2/replanning_a_star/local_planner.py b/dimos/navigation/go2/replanning_a_star/local_planner.py index bfa5033a11f..56438270de3 100644 --- a/dimos/navigation/go2/replanning_a_star/local_planner.py +++ b/dimos/navigation/go2/replanning_a_star/local_planner.py @@ -13,20 +13,20 @@ # limitations under the License. import os -from threading import Event, RLock, Thread +from threading import Event, RLock, Thread, current_thread import time import traceback from typing import Literal, TypeAlias +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist +from dimos_generated.nav_msgs.msg import OccupancyGrid, Path import numpy as np from reactivex import Subject +from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.global_config import GlobalConfig from dimos.core.resource import Resource -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import point_distance, quaternion_euler from dimos.navigation.base import NavigationState from dimos.navigation.go2.replanning_a_star.controllers import Controller, PController from dimos.navigation.go2.replanning_a_star.navigation_map import NavigationMap @@ -120,13 +120,17 @@ def start_planning(self, path: Path) -> None: self._thread.start() def stop_planning(self) -> None: - self.cmd_vel.on_next(Twist()) self._stop_planning_event.set() - + with self._lock: + thread = self._thread + if thread is not None and thread is not current_thread(): + thread.join(DEFAULT_THREAD_JOIN_TIMEOUT) + if thread.is_alive(): + raise RuntimeError("LocalPlanner thread did not stop in time") with self._lock: self._thread = None - self._reset_state() + self.cmd_vel.on_next(Twist()) def get_state(self) -> NavigationState: with self._lock: @@ -176,14 +180,15 @@ def _loop(self) -> None: # Determine initial state: skip initial_rotation if already aligned. new_state: PlannerState = "initial_rotation" if current_odom is not None and len(path.poses) > 0: - first_yaw = path.poses[0].orientation.euler[2] - robot_yaw = current_odom.orientation.euler[2] + first_yaw = quaternion_euler(path.poses[0].pose.orientation)[2] + robot_yaw = quaternion_euler(current_odom.pose.orientation)[2] initial_yaw_error = angle_diff(first_yaw, robot_yaw) self._controller.reset_yaw_error(initial_yaw_error) angle_in_tolerance = abs(initial_yaw_error) < self._orientation_tolerance if angle_in_tolerance: position_in_tolerance = ( - path.poses[0].position.distance(current_odom.position) < 0.01 + point_distance(path.poses[-1].pose.position, current_odom.pose.position) + < self._goal_tolerance ) if position_in_tolerance: new_state = "final_rotation" @@ -241,8 +246,8 @@ def _compute_initial_rotation(self) -> Twist: assert current_odom is not None first_pose = path.poses[0] - first_yaw = first_pose.orientation.euler[2] - robot_yaw = current_odom.orientation.euler[2] + first_yaw = quaternion_euler(first_pose.pose.orientation)[2] + robot_yaw = quaternion_euler(current_odom.pose.orientation)[2] yaw_error = angle_diff(first_yaw, robot_yaw) if abs(yaw_error) < self._orientation_tolerance: @@ -260,7 +265,7 @@ def get_distance_to_path(self) -> float | None: if path_distancer is None or current_odom is None: return None - current_pos = np.array([current_odom.position.x, current_odom.position.y]) + current_pos = np.array([current_odom.pose.position.x, current_odom.pose.position.y]) return path_distancer.get_distance_to_path(current_pos) @@ -272,7 +277,7 @@ def _compute_path_following(self) -> Twist: assert path_distancer is not None assert current_odom is not None - current_pos = np.array([current_odom.position.x, current_odom.position.y]) + current_pos = np.array([current_odom.pose.position.x, current_odom.pose.position.y]) if path_distancer.distance_to_goal(current_pos) < self._goal_tolerance: logger.info("Reached goal position, starting final rotation") @@ -297,8 +302,8 @@ def _compute_final_rotation(self) -> Twist: assert path is not None assert current_odom is not None - goal_yaw = path.poses[-1].orientation.euler[2] - robot_yaw = current_odom.orientation.euler[2] + goal_yaw = quaternion_euler(path.poses[-1].pose.orientation)[2] + robot_yaw = quaternion_euler(current_odom.pose.orientation)[2] yaw_error = angle_diff(goal_yaw, robot_yaw) if abs(yaw_error) < self._orientation_tolerance: diff --git a/dimos/navigation/go2/replanning_a_star/min_cost_astar.py b/dimos/navigation/go2/replanning_a_star/min_cost_astar.py index 03311db6897..54f1f744d39 100644 --- a/dimos/navigation/go2/replanning_a_star/min_cost_astar.py +++ b/dimos/navigation/go2/replanning_a_star/min_cost_astar.py @@ -13,12 +13,12 @@ # limitations under the License. import heapq +import math -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import VectorLike -from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid -from dimos.msgs.nav_msgs.Path import Path +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.nav_msgs.msg import OccupancyGrid, Path + +from dimos.msgs.occupancy import grid_to_world, occupancy_view, world_to_grid from dimos.utils.logging_config import setup_logger # Try to import C++ extension for faster pathfinding @@ -65,88 +65,55 @@ def _reconstruct_path( parents: dict[tuple[int, int], tuple[int, int]], current: tuple[int, int], costmap: OccupancyGrid, - start_tuple: tuple[int, int], - goal_tuple: tuple[int, int], ) -> Path: - frame_id = costmap.frame_id - waypoints: list[PoseStamped] = [] + coordinates = [current] while current in parents: - world_point = costmap.grid_to_world(current) - pose = PoseStamped( - frame_id=frame_id, - position=[world_point.x, world_point.y, 0.0], - orientation=Quaternion(0, 0, 0, 1), # Identity quaternion - ) - waypoints.append(pose) current = parents[current] - - start_world_point = costmap.grid_to_world(start_tuple) - start_pose = PoseStamped( - frame_id=frame_id, - position=[start_world_point.x, start_world_point.y, 0.0], - orientation=Quaternion(0, 0, 0, 1), - ) - waypoints.append(start_pose) - - waypoints.reverse() - - # Add the goal position if it's not already included - goal_point = costmap.grid_to_world(goal_tuple) - - if ( - not waypoints - or (waypoints[-1].x - goal_point.x) ** 2 + (waypoints[-1].y - goal_point.y) ** 2 > 1e-10 - ): - goal_pose = PoseStamped( - frame_id=frame_id, - position=[goal_point.x, goal_point.y, 0.0], - orientation=Quaternion(0, 0, 0, 1), - ) - waypoints.append(goal_pose) - - return Path(frame_id=frame_id, poses=waypoints) + coordinates.append(current) + coordinates.reverse() + return _reconstruct_path_from_coords(coordinates, costmap) def _reconstruct_path_from_coords( path_coords: list[tuple[int, int]], costmap: OccupancyGrid, ) -> Path: - frame_id = costmap.frame_id - waypoints: list[PoseStamped] = [] - - for gx, gy in path_coords: - world_point = costmap.grid_to_world((gx, gy)) - pose = PoseStamped( - frame_id=frame_id, - position=[world_point.x, world_point.y, 0.0], - orientation=Quaternion(0, 0, 0, 1), - ) - waypoints.append(pose) - - return Path(frame_id=frame_id, poses=waypoints) + return Path( + header=costmap.header, + poses=[ + PoseStamped( + header=costmap.header, + pose=Pose(position=grid_to_world(costmap, coordinate), orientation=Quaternion(w=1)), + ) + for coordinate in path_coords + ], + ) def min_cost_astar( costmap: OccupancyGrid, - goal: VectorLike, - start: VectorLike = (0.0, 0.0), + goal: Point, + start: Point | None = None, cost_threshold: int = 100, unknown_penalty: float = 0.8, use_cpp: bool = True, ) -> Path | None: - start_vector = costmap.world_to_grid(start) - goal_vector = costmap.world_to_grid(goal) - - start_tuple = (int(start_vector.x), int(start_vector.y)) - goal_tuple = (int(goal_vector.x), int(goal_vector.y)) - - if not (0 <= goal_tuple[0] < costmap.width and 0 <= goal_tuple[1] < costmap.height): + cells = occupancy_view(costmap) + start_vector = world_to_grid(costmap, start if start is not None else Point()) + goal_vector = world_to_grid(costmap, goal) + # Ignore sub-nanocell rotation roundoff before assigning a point to its cell. + start_tuple = (math.floor(round(start_vector[0], 9)), math.floor(round(start_vector[1], 9))) + goal_tuple = (math.floor(round(goal_vector[0], 9)), math.floor(round(goal_vector[1], 9))) + + if not (0 <= goal_tuple[0] < costmap.info.width and 0 <= goal_tuple[1] < costmap.info.height): + return None + if not (0 <= start_tuple[0] < costmap.info.width and 0 <= start_tuple[1] < costmap.info.height): return None if use_cpp: if _USE_CPP: path_coords = _astar_cpp( - costmap.grid, + cells, start_tuple[0], start_tuple[1], goal_tuple[0], @@ -185,7 +152,7 @@ def min_cost_astar( continue if current == goal_tuple: - return _reconstruct_path(parents, current, costmap, start_tuple, goal_tuple) + return _reconstruct_path(parents, current, costmap) closed_set.add(current) @@ -193,22 +160,22 @@ def min_cost_astar( neighbor_x, neighbor_y = current_x + dx, current_y + dy neighbor = (neighbor_x, neighbor_y) - if not (0 <= neighbor_x < costmap.width and 0 <= neighbor_y < costmap.height): + if not (0 <= neighbor_x < costmap.info.width and 0 <= neighbor_y < costmap.info.height): continue if neighbor in closed_set: continue - neighbor_val = costmap.grid[neighbor_y, neighbor_x] + neighbor_val = cells[neighbor_y, neighbor_x] if neighbor_val >= cost_threshold: continue - if neighbor_val == CostValues.UNKNOWN: + if neighbor_val == -1: cell_cost = cost_threshold * unknown_penalty if cell_cost >= cost_threshold: continue - elif neighbor_val == CostValues.FREE: + elif neighbor_val == 0: cell_cost = 0.0 else: cell_cost = neighbor_val diff --git a/dimos/navigation/go2/replanning_a_star/module.py b/dimos/navigation/go2/replanning_a_star/module.py index 36114874bb1..7423dc5d91a 100644 --- a/dimos/navigation/go2/replanning_a_star/module.py +++ b/dimos/navigation/go2/replanning_a_star/module.py @@ -15,18 +15,14 @@ import os from typing import Any -from dimos_lcm.std_msgs import Bool, String +from dimos_generated.geometry_msgs.msg import PointStamped, Pose, PoseStamped, Quaternion, Twist +from dimos_generated.nav_msgs.msg import OccupancyGrid, Odometry, Path +from dimos_generated.std_msgs.msg import Bool, String from reactivex.disposable import Disposable from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PointStamped import PointStamped -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.nav_msgs.Path import Path from dimos.navigation.base import NavigationInterface, NavigationState from dimos.navigation.go2.replanning_a_star.global_planner import GlobalPlanner from dimos.utils.logging_config import setup_logger @@ -81,7 +77,9 @@ def start(self) -> None: self.register_disposable( Disposable( self.odometry.subscribe( - lambda msg: self._planner.handle_odom(msg.to_pose_stamped()) + lambda msg: self._planner.handle_odom( + PoseStamped(header=msg.header, pose=msg.pose.pose) + ) ) ) ) @@ -98,7 +96,12 @@ def start(self) -> None: self.register_disposable( Disposable( self.clicked_point.subscribe( - lambda pt: self._planner.handle_goal_request(pt.to_pose_stamped()) + lambda pt: self._planner.handle_goal_request( + PoseStamped( + header=pt.header, + pose=Pose(position=pt.point, orientation=Quaternion(w=1)), + ) + ) ) ) ) diff --git a/dimos/navigation/go2/replanning_a_star/module_spec.py b/dimos/navigation/go2/replanning_a_star/module_spec.py index c9ec73dc473..a930142b719 100644 --- a/dimos/navigation/go2/replanning_a_star/module_spec.py +++ b/dimos/navigation/go2/replanning_a_star/module_spec.py @@ -14,7 +14,8 @@ from typing import Protocol -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos_generated.geometry_msgs.msg import PoseStamped + from dimos.navigation.base import NavigationState from dimos.spec.utils import Spec diff --git a/dimos/navigation/go2/replanning_a_star/navigation_map.py b/dimos/navigation/go2/replanning_a_star/navigation_map.py index fde75c0b0ee..25104c593c2 100644 --- a/dimos/navigation/go2/replanning_a_star/navigation_map.py +++ b/dimos/navigation/go2/replanning_a_star/navigation_map.py @@ -14,10 +14,11 @@ from threading import RLock +from dimos_generated.nav_msgs.msg import OccupancyGrid + from dimos.core.global_config import GlobalConfig from dimos.mapping.occupancy.gradient import GradientStrategy from dimos.mapping.occupancy.path_map import make_navigation_map -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid class NavigationMap: diff --git a/dimos/navigation/go2/replanning_a_star/path_clearance.py b/dimos/navigation/go2/replanning_a_star/path_clearance.py index 7dc08d49e0c..5774f85f156 100644 --- a/dimos/navigation/go2/replanning_a_star/path_clearance.py +++ b/dimos/navigation/go2/replanning_a_star/path_clearance.py @@ -12,23 +12,25 @@ # See the License for the specific language governing permissions and # limitations under the License. +from copy import deepcopy from threading import RLock +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid, Path import numpy as np from numpy.typing import NDArray from dimos.core.global_config import GlobalConfig from dimos.mapping.occupancy.path_mask import make_path_mask -from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import point_distance +from dimos.msgs.occupancy import occupancy_view class PathClearance: _costmap: OccupancyGrid | None = None - _last_costmap: OccupancyGrid | None = None _path_lookup_distance: float = 3.0 _max_distance_cache: float = 1.0 - _last_used_shape: tuple[int, ...] | None = None + _last_info: MapMetaData | None = None + _last_frame: str | None = None _last_mask: NDArray[np.bool_] | None = None _last_used_pose: int | None = None _global_config: GlobalConfig @@ -61,7 +63,8 @@ def mask(self) -> NDArray[np.bool_]: if ( self._last_mask is not None and self._last_used_pose is not None - and costmap.grid.shape == self._last_used_shape + and costmap.info == self._last_info + and costmap.header.frame_id == self._last_frame and self._pose_distance(self._last_used_pose, pose_index) < self._max_distance_cache ): return self._last_mask @@ -74,7 +77,8 @@ def mask(self) -> NDArray[np.bool_]: max_length=self._path_lookup_distance, ) - self._last_used_shape = costmap.grid.shape + self._last_info = deepcopy(costmap.info) + self._last_frame = costmap.header.frame_id self._last_used_pose = pose_index return self._last_mask @@ -86,9 +90,9 @@ def is_obstacle_ahead(self) -> bool: if costmap is None: return True - return bool(np.any(costmap.grid[self.mask] == CostValues.OCCUPIED)) + return bool(np.any(occupancy_view(costmap)[self.mask] == 100)) def _pose_distance(self, index1: int, index2: int) -> float: - p1 = self._path.poses[index1].position - p2 = self._path.poses[index2].position - return p1.distance(p2) + p1 = self._path.poses[index1].pose.position + p2 = self._path.poses[index2].pose.position + return point_distance(p1, p2) diff --git a/dimos/navigation/go2/replanning_a_star/path_distancer.py b/dimos/navigation/go2/replanning_a_star/path_distancer.py index 9a2c28dad78..f5a481e9b6d 100644 --- a/dimos/navigation/go2/replanning_a_star/path_distancer.py +++ b/dimos/navigation/go2/replanning_a_star/path_distancer.py @@ -14,11 +14,10 @@ from typing import cast +from dimos_generated.nav_msgs.msg import Path import numpy as np from numpy.typing import NDArray -from dimos.msgs.nav_msgs.Path import Path - class PathDistancer: _lookahead_dist: float @@ -27,7 +26,7 @@ class PathDistancer: def __init__(self, path: Path, lookahead_dist: float = 0.5) -> None: self._lookahead_dist = lookahead_dist - self._path = np.array([[p.position.x, p.position.y] for p in path.poses]) + self._path = np.array([[p.pose.position.x, p.pose.position.y] for p in path.poses]) self._cumulative_dists = _make_cumulative_distance_array(self._path) def find_lookahead_point(self, start_idx: int) -> NDArray[np.float64]: diff --git a/dimos/navigation/go2/replanning_a_star/position_tracker.py b/dimos/navigation/go2/replanning_a_star/position_tracker.py index 7d8249b562b..61a56d12039 100644 --- a/dimos/navigation/go2/replanning_a_star/position_tracker.py +++ b/dimos/navigation/go2/replanning_a_star/position_tracker.py @@ -16,11 +16,10 @@ import time from typing import cast +from dimos_generated.geometry_msgs.msg import PoseStamped import numpy as np from numpy.typing import NDArray -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - _max_points_per_second = 1000 @@ -51,7 +50,7 @@ def reset_data(self) -> None: def add_position(self, pose: PoseStamped) -> None: with self._lock: self._timestamps[self._index] = time.time() - self._positions[self._index] = (pose.position.x, pose.position.y) + self._positions[self._index] = (pose.pose.position.x, pose.pose.position.y) self._index = (self._index + 1) % self._max_points self._size = min(self._size + 1, self._max_points) diff --git a/dimos/navigation/go2/replanning_a_star/replan_limiter.py b/dimos/navigation/go2/replanning_a_star/replan_limiter.py index 8cc630f3dff..d8849b71823 100644 --- a/dimos/navigation/go2/replanning_a_star/replan_limiter.py +++ b/dimos/navigation/go2/replanning_a_star/replan_limiter.py @@ -14,7 +14,9 @@ from threading import RLock -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos_generated.geometry_msgs.msg import Point + +from dimos.msgs.geometry import point_distance from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -28,7 +30,7 @@ class ReplanLimiter: _max_attempts: int = 6 _reset_distance: float = 2.0 - _attempt_pos: Vector3 | None = None + _attempt_pos: Point | None = None _lock: RLock _attempt: int @@ -37,13 +39,13 @@ def __init__(self) -> None: self._lock = RLock() self._attempt = 0 - def can_retry(self, position: Vector3) -> bool: + def can_retry(self, position: Point) -> bool: with self._lock: if self._attempt == 0: self._attempt_pos = position if self._attempt >= 1 and self._attempt_pos: - distance = self._attempt_pos.distance(position) + distance = point_distance(self._attempt_pos, position) if distance >= self._reset_distance: logger.info( "Traveled enough to reset attempts", diff --git a/dimos/navigation/go2/replanning_a_star/spec.py b/dimos/navigation/go2/replanning_a_star/spec.py index d7aaab0ba7c..71b6697dc8b 100644 --- a/dimos/navigation/go2/replanning_a_star/spec.py +++ b/dimos/navigation/go2/replanning_a_star/spec.py @@ -14,7 +14,8 @@ from typing import Protocol -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos_generated.geometry_msgs.msg import PoseStamped + from dimos.navigation.base import NavigationState from dimos.spec.utils import Spec diff --git a/dimos/navigation/movement_manager/movement_manager.py b/dimos/navigation/movement_manager/movement_manager.py index 4f83428f9a7..da3cce4fed6 100644 --- a/dimos/navigation/movement_manager/movement_manager.py +++ b/dimos/navigation/movement_manager/movement_manager.py @@ -25,15 +25,14 @@ import time from typing import Any -from dimos_lcm.std_msgs import Bool +from dimos_generated.geometry_msgs.msg import Point, PointStamped, Twist, Vector3 +from dimos_generated.std_msgs.msg import Bool from reactivex.disposable import Disposable from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PointStamped import PointStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.time import header_now from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -45,7 +44,9 @@ class MovementManagerConfig(ModuleConfig): tele_cooldown_sec: float = 1.0 - tele_cmd_vel_scaling: Twist = Twist(Vector3(1, 1, 1), Vector3(1, 1, 1)) + tele_cmd_vel_scaling: Twist = Twist( + linear=Vector3(x=1, y=1, z=1), angular=Vector3(x=1, y=1, z=1) + ) class MovementManager(Module): @@ -82,18 +83,22 @@ def stop(self) -> None: super().stop() def _on_click(self, msg: PointStamped) -> None: - if not all(math.isfinite(v) for v in (msg.x, msg.y, msg.z)): - logger.warning("Ignored invalid click", x=msg.x, y=msg.y, z=msg.z) + if not all(math.isfinite(v) for v in (msg.point.x, msg.point.y, msg.point.z)): + logger.warning("Ignored invalid click", x=msg.point.x, y=msg.point.y, z=msg.point.z) return if ( - abs(msg.x) > MAX_CLICK_HORIZONTAL_M - or abs(msg.y) > MAX_CLICK_HORIZONTAL_M - or abs(msg.z) > MAX_CLICK_VERTICAL_M + abs(msg.point.x) > MAX_CLICK_HORIZONTAL_M + or abs(msg.point.y) > MAX_CLICK_HORIZONTAL_M + or abs(msg.point.z) > MAX_CLICK_VERTICAL_M ): - logger.warning("Ignored out-of-range click", x=msg.x, y=msg.y, z=msg.z) + logger.warning( + "Ignored out-of-range click", x=msg.point.x, y=msg.point.y, z=msg.point.z + ) return - logger.debug("Goal", x=round(msg.x, 1), y=round(msg.y, 1), z=round(msg.z, 1)) + logger.debug( + "Goal", x=round(msg.point.x, 1), y=round(msg.point.y, 1), z=round(msg.point.z, 1) + ) self.way_point.publish(msg) self.goal.publish(msg) @@ -103,7 +108,7 @@ def _cancel_goal(self) -> None: # It can be REALLY bad if a robot is supposed to stop moving but wont # we should probably think a more robust/strict requirement on planners cancel = PointStamped( - ts=time.time(), frame_id="map", x=float("nan"), y=float("nan"), z=float("nan") + header=header_now("map"), point=Point(x=float("nan"), y=float("nan"), z=float("nan")) ) self.way_point.publish(cancel) self.goal.publish(cancel) @@ -129,14 +134,14 @@ def _on_teleop(self, msg: Twist) -> None: scale = self.config.tele_cmd_vel_scaling scaled = Twist( linear=Vector3( - msg.linear.x * scale.linear.x, - msg.linear.y * scale.linear.y, - msg.linear.z * scale.linear.z, + x=msg.linear.x * scale.linear.x, + y=msg.linear.y * scale.linear.y, + z=msg.linear.z * scale.linear.z, ), angular=Vector3( - msg.angular.x * scale.angular.x, - msg.angular.y * scale.angular.y, - msg.angular.z * scale.angular.z, + x=msg.angular.x * scale.angular.x, + y=msg.angular.y * scale.angular.y, + z=msg.angular.z * scale.angular.z, ), ) self.cmd_vel.publish(scaled) diff --git a/dimos/navigation/movement_manager/test_movement_manager.py b/dimos/navigation/movement_manager/test_movement_manager.py index e266ad3e983..0c860e0d2ad 100644 --- a/dimos/navigation/movement_manager/test_movement_manager.py +++ b/dimos/navigation/movement_manager/test_movement_manager.py @@ -19,12 +19,11 @@ import math import time +from dimos_generated.geometry_msgs.msg import Point, PointStamped, Twist, Vector3 import pytest from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT -from dimos.msgs.geometry_msgs.PointStamped import PointStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.time import header_now from dimos.navigation.movement_manager.movement_manager import ( MovementManager, ) @@ -44,10 +43,16 @@ def _attach(module): """Subscribe to every Out port; return (captured, unsubscribers).""" captured = Captured() unsubs = [ - module.cmd_vel.subscribe(captured.cmd_vel.append), - module.stop_movement.subscribe(captured.stop_movement.append), - module.goal.subscribe(captured.goal.append), - module.way_point.subscribe(captured.way_point.append), + module.cmd_vel.subscribe( + lambda msg: captured.cmd_vel.append(type(msg).decode(msg.encode())) + ), + module.stop_movement.subscribe( + lambda msg: captured.stop_movement.append(type(msg).decode(msg.encode())) + ), + module.goal.subscribe(lambda msg: captured.goal.append(type(msg).decode(msg.encode()))), + module.way_point.subscribe( + lambda msg: captured.way_point.append(type(msg).decode(msg.encode())) + ), ] return captured, unsubs @@ -65,11 +70,11 @@ def manager_and_captured() -> Generator[tuple[MovementManager, Captured], None, def _twist(lx=0.0): - return Twist(linear=Vector3(lx, 0, 0), angular=Vector3(0, 0, 0)) + return Twist(linear=Vector3(x=lx, y=0, z=0), angular=Vector3(x=0, y=0, z=0)) def _click(x=1.0, y=2.0, z=0.0): - return PointStamped(ts=time.time(), frame_id="map", x=x, y=y, z=z) + return PointStamped(header=header_now("map"), point=Point(x=x, y=y, z=z)) def test_teleop_suppresses_nav_and_cancels_goal(manager_and_captured): @@ -88,7 +93,7 @@ def test_teleop_suppresses_nav_and_cancels_goal(manager_and_captured): # Goal cancelled with NaN assert len(captured.goal) == 1 - assert math.isnan(captured.goal[0].x) + assert math.isnan(captured.goal[0].point.x) def test_nav_resumes_after_cooldown(manager_and_captured): @@ -127,11 +132,11 @@ def test_invalid_clicks_rejected(manager_and_captured): def test_tele_cmd_vel_scaling(manager_and_captured): """tele_cmd_vel_scaling multiplies each teleop twist component independently.""" manager, captured = manager_and_captured - scaling = Twist(Vector3(0.5, 2.0, 0.0), Vector3(1.0, 1.0, 0.25)) + scaling = Twist(linear=Vector3(x=0.5, y=2.0, z=0.0), angular=Vector3(x=1.0, y=1.0, z=0.25)) manager.config.tele_cmd_vel_scaling = scaling manager.config.tele_cooldown_sec = 10.0 - manager._on_teleop(Twist(Vector3(1, 1, 1), Vector3(1, 1, 1))) + manager._on_teleop(Twist(linear=Vector3(x=1, y=1, z=1), angular=Vector3(x=1, y=1, z=1))) assert len(captured.cmd_vel) == 1 published = captured.cmd_vel[0] diff --git a/dimos/navigation/test_generated_bbox_navigation.py b/dimos/navigation/test_generated_bbox_navigation.py new file mode 100644 index 00000000000..43ab745a57d --- /dev/null +++ b/dimos/navigation/test_generated_bbox_navigation.py @@ -0,0 +1,73 @@ +# Copyright 2025-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. + +from threading import Event +from uuid import uuid4 + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import CameraInfo +from dimos_generated.std_msgs.msg import Header +from dimos_generated.vision_msgs.msg import ( + BoundingBox2D, + Detection2D, + Detection2DArray, + Point2D, + Pose2D, +) +import pytest + +from dimos.core.transport import LCMTransport +from dimos.navigation.bbox_navigation import BBoxNavigationModule + + +def test_generated_bbox_goal_preserves_header_and_projection(): + module = BBoxNavigationModule(goal_distance=2) + prefix = uuid4().hex[:8] + module.camera_info.transport = LCMTransport(f"/bc/{prefix}", CameraInfo) + module.detection2d.transport = LCMTransport(f"/bd/{prefix}", Detection2DArray) + goals = [] + ready = Event() + received = Event() + unsubscribe = module.goal_request.subscribe( + lambda value: (goals.append(PoseStamped.decode(value.encode())), received.set()) + ) + message = Detection2DArray( + header=Header(frame_id="base_link", stamp=Time(sec=1700000000, nanosec=123456789)), + detections=[Detection2D(bbox=BoundingBox2D(center=Pose2D(position=Point2D(x=420, y=290))))], + ) + try: + module.start() + module._on_detection(message) + assert not goals + camera = CameraInfo(k=[500, 0, 320, 0, 250, 240, 0, 0, 1]) + camera_unsubscribe = module.camera_info.subscribe(lambda _: ready.set()) + module.camera_info.transport.publish(CameraInfo.decode(camera.encode())) + assert ready.wait(2) + camera_unsubscribe() + module.detection2d.transport.publish(Detection2DArray.decode(message.encode())) + assert received.wait(2) + assert len(goals) == 1 + goal = goals[0] + assert goal.header == message.header + assert (goal.pose.position.x, goal.pose.position.y, goal.pose.position.z) == pytest.approx( + (2, -0.4, -0.4) + ) + module._on_detection(Detection2DArray()) + assert len(goals) == 1 + finally: + unsubscribe() + module.stop() + module.camera_info.transport.stop() + module.detection2d.transport.stop() diff --git a/dimos/perception/detection/conftest.py b/dimos/perception/detection/conftest.py index a34bcf914f8..ff065a0e1ea 100644 --- a/dimos/perception/detection/conftest.py +++ b/dimos/perception/detection/conftest.py @@ -14,18 +14,17 @@ from collections.abc import Callable, Generator import functools +import hashlib +import json from typing import TypedDict -from unittest import mock -from dimos_lcm.visualization_msgs.MarkerArray import MarkerArray +from dimos_generated.geometry_msgs.msg import PoseStamped, TransformStamped +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 +from dimos_generated.vision_msgs.msg import Detection2DArray +from dimos_generated.visualization_msgs.msg import MarkerArray import pytest from dimos.core.transport import LCMTransport -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray from dimos.perception.detection.module2D import Detection2DModule from dimos.perception.detection.module3D import Detection3DModule from dimos.perception.detection.type.detection2d.base import Detection2D @@ -35,17 +34,15 @@ from dimos.perception.experimental.moduleDB import ObjectDBModule from dimos.protocol.tf.tf import TF from dimos.robot.unitree.go2 import connection -from dimos.robot.unitree.type.odometry import Odometry from dimos.utils.data import get_data -from dimos.utils.testing.legacy_pickle import LegacyPickleStore class Moment(TypedDict, total=False): - odom_frame: Odometry + odom_frame: PoseStamped lidar_frame: PointCloud2 image_frame: Image camera_info: CameraInfo - transforms: list[Transform] + transforms: list[TransformStamped] tf: TF detections: ImageDetections3DPC | None markers: MarkerArray | None @@ -70,32 +67,23 @@ def tf(): def get_moment(tf): @functools.lru_cache(maxsize=1) def moment_provider(**kwargs) -> Moment: - print("MOMENT PROVIDER ARGS:", kwargs) seek = kwargs.get("seek", 10.0) - - data_dir = "unitree_go2_lidar_corrected" - get_data(data_dir) - - lidar_frame_result = LegacyPickleStore(f"{data_dir}/lidar").find_closest_seek(seek) - if lidar_frame_result is None: - raise ValueError("No lidar frame found") - lidar_frame: PointCloud2 = lidar_frame_result - - image_frame = LegacyPickleStore( - f"{data_dir}/video", - ).find_closest(lidar_frame.ts) - - if image_frame is None: - raise ValueError("No image frame found") - - image_frame.frame_id = "camera_optical" - - odom_frame = LegacyPickleStore(f"{data_dir}/odom", autocast=Odometry.from_msg).find_closest( - lidar_frame.ts - ) - - if odom_frame is None: - raise ValueError("No odom frame found") + data_dir = get_data("unitree_go2_detection_cdr") + manifest = json.loads((data_dir / "manifest.json").read_text()) + entry = next(item for item in manifest["moments"] if item["seek"] == seek) + + def payload(name): + stream = entry["streams"][name] + data = (data_dir / stream["file"]).read_bytes() + assert hashlib.sha256(data).hexdigest() == stream["cdr_sha256"] + return data + + lidar_frame = PointCloud2.decode(payload("lidar")) + image_frame = Image.decode(payload("video")) + image_frame.header.frame_id = "camera_optical" + odom_frame = PoseStamped.decode(payload("odom")) + # Match the Go2 connection's configured global parent for this recording. + odom_frame.header.frame_id = "world" transforms = connection.GO2Connection._odom_to_tf(odom_frame) @@ -105,7 +93,7 @@ def moment_provider(**kwargs) -> Moment: "odom_frame": odom_frame, "lidar_frame": lidar_frame, "image_frame": image_frame, - "camera_info": connection._camera_info_static(), + "camera_info": connection.front_camera_calibration(), "transforms": transforms, "tf": tf, } @@ -185,7 +173,7 @@ def detection3dpc(detections3dpc) -> Detection3DPC: def get_moment_2d(get_moment) -> Generator[Callable[[], Moment2D], None, None]: from dimos.perception.detection.detectors.yolo import Yolo2DDetector - c = mock.create_autospec(CameraInfo, spec_set=True, instance=True) + c = connection.front_camera_calibration() module = Detection2DModule(detector=lambda: Yolo2DDetector(device="cpu"), camera_info=c) @functools.lru_cache(maxsize=1) @@ -220,7 +208,7 @@ def moment_provider(**kwargs) -> Moment3D: if lidar_frame is None: raise ValueError("No lidar frame found") - camera_transform = moment["tf"].get("camera_optical", lidar_frame.frame_id) + camera_transform = moment["tf"].get("camera_optical", lidar_frame.header.frame_id) if camera_transform is None: raise ValueError("No camera_optical transform in tf") @@ -244,10 +232,10 @@ def object_db_module(get_moment): """Create and populate an ObjectDBModule with detections from multiple frames.""" from dimos.perception.detection.detectors.yolo import Yolo2DDetector - c = mock.create_autospec(CameraInfo, spec_set=True, instance=True) + c = connection.front_camera_calibration() module2d = Detection2DModule(detector=lambda: Yolo2DDetector(device="cpu"), camera_info=c) - module3d = Detection3DModule(camera_info=connection._camera_info_static()) - moduleDB = ObjectDBModule(camera_info=connection._camera_info_static()) + module3d = Detection3DModule(camera_info=connection.front_camera_calibration()) + moduleDB = ObjectDBModule(camera_info=connection.front_camera_calibration()) # Process 5 frames to build up object history for i in range(5): @@ -258,7 +246,7 @@ def object_db_module(get_moment): imageDetections2d = module2d.process_image_frame(moment["image_frame"]) # Get camera transform - camera_transform = moment["tf"].get("camera_optical", moment.get("lidar_frame").frame_id) + camera_transform = moment["tf"].get("camera_optical", moment["lidar_frame"].header.frame_id) # Process 3D detections imageDetections3d = module3d.process_frame( diff --git a/dimos/perception/detection/detectors/base.py b/dimos/perception/detection/detectors/base.py index 40aa82e5bdd..2de6d864b63 100644 --- a/dimos/perception/detection/detectors/base.py +++ b/dimos/perception/detection/detectors/base.py @@ -14,7 +14,8 @@ from abc import ABC, abstractmethod -from dimos.msgs.sensor_msgs.Image import Image +from dimos_generated.sensor_msgs.msg import Image + from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D diff --git a/dimos/perception/detection/detectors/conftest.py b/dimos/perception/detection/detectors/conftest.py index bb9a47e0eb9..d7e42317792 100644 --- a/dimos/perception/detection/detectors/conftest.py +++ b/dimos/perception/detection/detectors/conftest.py @@ -14,7 +14,7 @@ import pytest -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_from_file from dimos.perception.detection.detectors.person.yolo import YoloPersonDetector from dimos.perception.detection.detectors.yolo import Yolo2DDetector from dimos.perception.detection.detectors.yoloe import Yoloe2DDetector, YoloePromptMode @@ -24,7 +24,7 @@ @pytest.fixture(scope="session") def test_image(): """Load the test image used for detector tests.""" - return Image.from_file(get_data("cafe.jpg")) + return image_from_file(get_data("cafe.jpg")) @pytest.fixture(scope="session") diff --git a/dimos/perception/detection/detectors/omdet.py b/dimos/perception/detection/detectors/omdet.py index aeee0f7433a..6ea3d71e42c 100644 --- a/dimos/perception/detection/detectors/omdet.py +++ b/dimos/perception/detection/detectors/omdet.py @@ -18,12 +18,14 @@ from functools import cached_property +from dimos_generated.sensor_msgs.msg import Image import numpy as np from PIL import Image as PILImage import torch from dimos.models.base import HuggingFaceModel, HuggingFaceModelConfig -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_to_rgb +from dimos.msgs.time import to_seconds from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D @@ -79,7 +81,7 @@ def query_detections( ``confidence`` is the per-box score. ``class_id`` indexes into ``queries``. """ - pil = PILImage.fromarray(image.to_rgb().data) + pil = PILImage.fromarray(image_to_rgb(image)) with torch.inference_mode(): inputs = self._processor(images=pil, text=queries, return_tensors="pt").to( self.config.device @@ -105,7 +107,7 @@ def query_detections( class_id=int(label), confidence=float(score), name=queries[int(label)], - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, ) if det.is_valid(): @@ -135,7 +137,7 @@ def query_score_rows( from torchvision.ops.boxes import batched_nms # type: ignore[import-untyped] from transformers.image_transforms import center_to_corners_format - pil = PILImage.fromarray(image.to_rgb().data) + pil = PILImage.fromarray(image_to_rgb(image)) with torch.inference_mode(): inputs = self._processor( images=pil, text=queries, task="Detect all objects.", return_tensors="pt" diff --git a/dimos/perception/detection/detectors/person/yolo.py b/dimos/perception/detection/detectors/person/yolo.py index 635a3d119b7..649a1a489aa 100644 --- a/dimos/perception/detection/detectors/person/yolo.py +++ b/dimos/perception/detection/detectors/person/yolo.py @@ -12,7 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos.msgs.sensor_msgs.Image import Image +from dimos_generated.sensor_msgs.msg import Image + +from dimos.msgs.image import image_to_bgr from dimos.perception.detection.detectors.base import Detector from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.utils.data import get_data @@ -57,7 +59,7 @@ def process_image(self, image: Image) -> ImageDetections2D: ImageDetections2D containing Detection2DPerson objects with pose keypoints """ results = self.model.track( - source=image.to_opencv(), + source=image_to_bgr(image), verbose=False, conf=0.5, tracker=self.tracker, diff --git a/dimos/perception/detection/detectors/test_bbox_detectors.py b/dimos/perception/detection/detectors/test_bbox_detectors.py index 9ca3eca029a..d71abd7fe6a 100644 --- a/dimos/perception/detection/detectors/test_bbox_detectors.py +++ b/dimos/perception/detection/detectors/test_bbox_detectors.py @@ -12,10 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.sensor_msgs.msg import Image import pytest from dimos.core.transport import LCMTransport -from dimos.msgs.sensor_msgs.Image import Image from dimos.perception.detection.type.detection2d.base import Detection2D from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D @@ -89,9 +89,8 @@ def test_detection_cropped_image(detections, test_image) -> None: assert cropped is not None # Cropped image should be smaller than original (usually) - if test_image.shape: - assert cropped.shape[0] <= test_image.shape[0] - assert cropped.shape[1] <= test_image.shape[1] + assert cropped.height <= test_image.height + assert cropped.width <= test_image.width def test_detection_ros_conversion(detections) -> None: diff --git a/dimos/perception/detection/detectors/yolo.py b/dimos/perception/detection/detectors/yolo.py index 20dd756ed4d..f9ad70d883a 100644 --- a/dimos/perception/detection/detectors/yolo.py +++ b/dimos/perception/detection/detectors/yolo.py @@ -12,7 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos.msgs.sensor_msgs.Image import Image +from dimos_generated.sensor_msgs.msg import Image + +from dimos.msgs.image import image_to_bgr from dimos.perception.detection.detectors.base import Detector from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.utils.data import get_data @@ -59,7 +61,7 @@ def process_image(self, image: Image) -> ImageDetections2D: ImageDetections2D containing all detected objects """ results = self.model.track( - source=image.to_opencv(), + source=image_to_bgr(image), device=self.device, conf=0.5, iou=0.6, diff --git a/dimos/perception/detection/detectors/yoloe.py b/dimos/perception/detection/detectors/yoloe.py index 4311f6783f7..856c6ab65f0 100644 --- a/dimos/perception/detection/detectors/yoloe.py +++ b/dimos/perception/detection/detectors/yoloe.py @@ -16,11 +16,12 @@ import threading from typing import Any +from dimos_generated.sensor_msgs.msg import Image import numpy as np from numpy.typing import NDArray import torch -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_to_bgr from dimos.perception.detection.detectors.base import Detector from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.utils.data import get_data @@ -144,7 +145,7 @@ def predict_image(self, image: Image) -> "ImageDetections2D[Any]": def _run(self, image: Image, *, track: bool) -> "ImageDetections2D[Any]": """Run tracking or prediction and apply common filtering.""" kwargs = { - "source": image.to_opencv(), + "source": image_to_bgr(image), "device": self.device, "conf": self.conf, "iou": 0.6, diff --git a/dimos/perception/detection/module2D.py b/dimos/perception/detection/module2D.py index 34e3e675915..696b16ae2a3 100644 --- a/dimos/perception/detection/module2D.py +++ b/dimos/perception/detection/module2D.py @@ -14,6 +14,11 @@ from collections.abc import Callable, Sequence from typing import Annotated, Any +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage +from dimos_generated.vision_msgs.msg import Detection2DArray from pydantic.experimental.pipeline import validate_as from reactivex import operators as ops from reactivex.observable import Observable @@ -22,18 +27,14 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import IO, In, Out -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image, sharpness_barrier -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray +from dimos.msgs.image import image_sharpness +from dimos.msgs.time import to_seconds from dimos.perception.detection.detectors.base import Detector from dimos.perception.detection.detectors.yolo import Yolo2DDetector from dimos.perception.detection.type.detection2d.base import Filter2D from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.utils.decorators.decorators import simple_mcache -from dimos.utils.reactive import backpressure +from dimos.utils.reactive import backpressure, quality_barrier class Config(ModuleConfig): @@ -80,7 +81,7 @@ def process_image_frame(self, image: Image) -> ImageDetections2D: def sharp_image_stream(self) -> Observable[Image]: return backpressure( self.color_image.pure_observable().pipe( - sharpness_barrier(self.config.max_freq), + quality_barrier(image_sharpness, self.config.max_freq), ) ) @@ -89,7 +90,9 @@ def detection_stream_2d(self) -> Observable[ImageDetections2D]: return backpressure(self.sharp_image_stream().pipe(ops.map(self.process_image_frame))) def track(self, detections: ImageDetections2D) -> None: - sensor_frame = self.tfbuffer.get("sensor", "camera_optical", detections.image.ts, 5.0) + sensor_frame = self.tfbuffer.get( + "sensor", "camera_optical", to_seconds(detections.image.header.stamp), 5.0 + ) if not sensor_frame: return @@ -110,19 +113,19 @@ def track(self, detections: ImageDetections2D) -> None: detection = detections.detections[index] position_3d = self.pixel_to_3d( # type: ignore[attr-defined] detection.center_bbox, - self.config.camera_info, assumed_depth=1.0, ) else: # No detection at this index - publish zero transform - position_3d = Vector3(0.0, 0.0, 0.0) + position_3d = Vector3() transforms.append( - Transform( - frame_id=sensor_frame.child_frame_id, + TransformStamped( + header=Header( + frame_id=sensor_frame.child_frame_id, stamp=detections.image.header.stamp + ), child_frame_id=f"det_{index}", - ts=detections.image.ts, - translation=position_3d, + transform=Transform(translation=position_3d, rotation=Quaternion(w=1)), ) ) diff --git a/dimos/perception/detection/module3D.py b/dimos/perception/detection/module3D.py index 861dc14e873..f3a53e531ab 100644 --- a/dimos/perception/detection/module3D.py +++ b/dimos/perception/detection/module3D.py @@ -13,27 +13,36 @@ # limitations under the License. +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + TransformStamped, + Vector3, +) +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.vision_msgs.msg import Detection2DArray from reactivex import operators as ops from reactivex.observable import Observable from dimos.agents.annotation import skill from dimos.core.core import rpc from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray +from dimos.msgs.time import to_seconds from dimos.perception.detection.module2D import Detection2DModule from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.detection.type.detection3d.imageDetections3DPC import ImageDetections3DPC from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC -from dimos.types.timestamped import align_timestamped +from dimos.types.timestamped import TimestampedData, align_timestamped from dimos.utils.reactive import backpressure +def _timed_cloud(cloud: PointCloud2) -> TimestampedData[PointCloud2]: + return TimestampedData(cloud, to_seconds(cloud.header.stamp)) + + class Detection3DModule(Detection2DModule): color_image: In[Image] pointcloud: In[PointCloud2] @@ -56,7 +65,7 @@ def process_frame( self, detections: ImageDetections2D, pointcloud: PointCloud2, - transform: Transform | None, + transform: TransformStamped | None, ) -> ImageDetections3DPC: if not transform: return ImageDetections3DPC(detections.image, []) @@ -89,8 +98,8 @@ def pixel_to_3d( Vector3 position in camera optical frame coordinates """ # Extract camera intrinsics - fx, fy = self.config.camera_info.K[0], self.config.camera_info.K[4] - cx, cy = self.config.camera_info.K[2], self.config.camera_info.K[5] + fx, fy = self.config.camera_info.k[0], self.config.camera_info.k[4] + cx, cy = self.config.camera_info.k[2], self.config.camera_info.k[5] # Unproject pixel to normalized camera coordinates x_norm = (pixel[0] - cx) / fx @@ -98,7 +107,7 @@ def pixel_to_3d( # Create 3D point at assumed depth in camera optical frame # Camera optical frame: X right, Y down, Z forward - return Vector3(x_norm * assumed_depth, y_norm * assumed_depth, assumed_depth) + return Vector3(x=x_norm * assumed_depth, y=y_norm * assumed_depth, z=assumed_depth) @skill def ask_vlm(self, question: str) -> str: @@ -113,7 +122,7 @@ def ask_vlm(self, question: str) -> str: # @skill @rpc - def nav_vlm(self, question: str) -> str: + def nav_vlm(self, question: str) -> PoseStamped | None: """ query visual model about the view in front of the camera you can ask to mark objects like: @@ -131,17 +140,19 @@ def nav_vlm(self, question: str) -> str: print("VLM result:", result, "for", image, "and question", question) if isinstance(result, str) or not result or not len(result): - return None # type: ignore[return-value] + return None detections: ImageDetections2D = result print(detections) if not len(detections): print("No 2d detections") - return None # type: ignore[return-value] + return None pc = self.pointcloud.get_next() - transform = self.tfbuffer.get("camera_optical", pc.frame_id, detections.image.ts, 5.0) + transform = self.tfbuffer.get( + "camera_optical", pc.header.frame_id, to_seconds(detections.image.header.stamp), 5.0 + ) detections3d = self.process_frame(detections, pc, transform) @@ -150,11 +161,13 @@ def nav_vlm(self, question: str) -> str: print("No 3d detections, projecting 2d") center = detections[0].get_bbox_center() + position = self.pixel_to_3d(center, assumed_depth=1.5) return PoseStamped( - ts=detections.image.ts, - frame_id="world", - position=self.pixel_to_3d(center, assumed_depth=1.5), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + header=Header(stamp=detections.image.header.stamp, frame_id="world"), + pose=Pose( + position=Point(x=position.x, y=position.y, z=position.z), + orientation=Quaternion(w=1), + ), ) @rpc @@ -162,13 +175,16 @@ def start(self) -> None: super().start() def detection2d_to_3d(args): # type: ignore[no-untyped-def] - detections, pc = args - transform = self.tfbuffer.get("camera_optical", pc.frame_id, detections.image.ts, 5.0) + detections, timed_pc = args + pc = timed_pc.value + transform = self.tfbuffer.get( + "camera_optical", pc.header.frame_id, to_seconds(detections.image.header.stamp), 5.0 + ) return self.process_frame(detections, pc, transform) self.detection_stream_3d = align_timestamped( # type: ignore[type-var] backpressure(self.detection_stream_2d()), - self.pointcloud.observable(), + self.pointcloud.observable().pipe(ops.map(_timed_cloud)), match_tolerance=0.25, buffer_size=20.0, ).pipe(ops.map(detection2d_to_3d)) diff --git a/dimos/perception/detection/person_tracker.py b/dimos/perception/detection/person_tracker.py index 1fb28dd29e7..0fb287856c4 100644 --- a/dimos/perception/detection/person_tracker.py +++ b/dimos/perception/detection/person_tracker.py @@ -15,24 +15,42 @@ from typing import Any +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage +from dimos_generated.vision_msgs.msg import Detection2DArray from reactivex import operators as ops from reactivex.observable import Observable from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray +from dimos.msgs.geometry import compose_transforms, pose_from_transform, transform_from_pose +from dimos.msgs.time import to_seconds from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D -from dimos.types.timestamped import align_timestamped +from dimos.types.timestamped import TimestampedData, align_timestamped from dimos.utils.reactive import backpressure +def _timed_image(message: Image) -> TimestampedData[Image]: + return TimestampedData(message, to_seconds(message.header.stamp)) + + +def _timed_detections(message: Detection2DArray) -> TimestampedData[Detection2DArray]: + return TimestampedData(message, to_seconds(message.header.stamp)) + + +def _has_detections(message: Detection2DArray) -> bool: + return len(message.detections) > 0 + + +def _paired_detections( + pair: tuple[TimestampedData[Image], TimestampedData[Detection2DArray]], +) -> ImageDetections2D: + return ImageDetections2D.from_ros_detection2d_array(pair[0].value, pair[1].value) + + class PersonTracker(Module): detections: In[Detection2DArray] color_image: In[Image] @@ -61,8 +79,8 @@ def center_to_3d( Vector3 position in camera_link frame coordinates (Z up, X forward) """ # Extract camera intrinsics - fx, fy = camera_info.K[0], camera_info.K[4] - cx, cy = camera_info.K[2], camera_info.K[5] + fx, fy = camera_info.k[0], camera_info.k[4] + cx, cy = camera_info.k[2], camera_info.k[5] # Unproject pixel to normalized camera coordinates x_norm = (pixel[0] - cx) / fx @@ -78,22 +96,18 @@ def center_to_3d( # Optical: X right, Y down, Z forward # Link: X forward, Y left, Z up # Transformation: x_link = z_optical, y_link = -x_optical, z_link = -y_optical - return Vector3(z_optical, -x_optical, -y_optical) + return Vector3(x=z_optical, y=-x_optical, z=-y_optical) def detections_stream(self) -> Observable[ImageDetections2D]: return backpressure( align_timestamped( - self.color_image.pure_observable(), + self.color_image.pure_observable().pipe(ops.map(_timed_image)), self.detections.pure_observable().pipe( - ops.filter(lambda d: d.detections_length > 0) # type: ignore[attr-defined] + ops.filter(_has_detections), ops.map(_timed_detections) ), match_tolerance=0.0, buffer_size=2.0, - ).pipe( - ops.map( - lambda pair: ImageDetections2D.from_ros_detection2d_array(*pair) # type: ignore[arg-type, misc] - ) - ) + ).pipe(ops.map(_paired_detections)) ) @rpc @@ -112,17 +126,18 @@ def track(self, detections2D: ImageDetections2D) -> None: vector = self.center_to_3d(target.center_bbox, self.camera_info, 2.0) pose_in_camera = PoseStamped( - ts=detections2D.ts, - position=vector, - frame_id="camera_link", + header=Header(frame_id="camera_link", stamp=detections2D.image.header.stamp), + pose=Pose( + position=Point(x=vector.x, y=vector.y, z=vector.z), orientation=Quaternion(w=1) + ), ) tf_world_to_camera = self.tfbuffer.get("world", "camera_link", detections2D.ts, 5.0) if not tf_world_to_camera: return - tf_camera_to_target = Transform.from_pose("target", pose_in_camera) - tf_world_to_target = tf_world_to_camera + tf_camera_to_target - pose_in_world = tf_world_to_target.to_pose(ts=detections2D.ts) + tf_camera_to_target = transform_from_pose(pose_in_camera, child_frame_id="target") + tf_world_to_target = compose_transforms(tf_world_to_camera, tf_camera_to_target) + pose_in_world = pose_from_transform(tf_world_to_target) self.target.publish(pose_in_world) diff --git a/dimos/perception/detection/project.py b/dimos/perception/detection/project.py index 1f9a8561ab1..d54e71e5422 100644 --- a/dimos/perception/detection/project.py +++ b/dimos/perception/detection/project.py @@ -18,20 +18,25 @@ from typing import TYPE_CHECKING, Any +from dimos_generated.geometry_msgs.msg import TransformStamped import numpy as np from dimos.memory.transform import Transformer -from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.geometry import ( + compose_transforms, + inverse_transform, + transform_from_pose, + transform_matrix, +) +from dimos.msgs.image import image_view from dimos.perception.detection.type.detection3d.imageDetections3DPC import ImageDetections3DPC if TYPE_CHECKING: from collections.abc import Callable, Iterator - from dimos_lcm.sensor_msgs import CameraInfo + from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 from dimos.memory.type.observation import Observation - from dimos.msgs.sensor_msgs.Image import Image - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.detection.type.detection3d.pointcloud_filters import PointCloudFilter from dimos.protocol.tf.tf import TFLookup @@ -41,17 +46,20 @@ def _world_to_optical( obs: Observation[Any], world_frame: str, tf: TFLookup | None, - base_to_optical: Transform | None, + base_to_optical: TransformStamped | None, optical_frame: str, time_tolerance: float, -) -> Transform | None: +) -> TransformStamped | None: """Resolve the world→optical transform for an observation.""" if tf is not None: return tf.get(optical_frame, world_frame, obs.ts, time_tolerance) pose = obs.pose_stamped if base_to_optical is None or pose is None: return None - return -(Transform.from_pose("base_link", pose) + base_to_optical) + pose.header.frame_id = world_frame + return inverse_transform( + compose_transforms(transform_from_pose(pose, child_frame_id="base_link"), base_to_optical) + ) class ProjectDepthTo3D(Transformer["ImageDetections2D", ImageDetections3DPC]): @@ -67,7 +75,7 @@ def __init__( depth: Callable[[Observation[ImageDetections2D]], Image | None], camera_info: CameraInfo, tf: TFLookup | None = None, - base_to_optical: Transform | None = None, + base_to_optical: TransformStamped | None = None, world_frame: str = "world", optical_frame: str = "camera_optical", time_tolerance: float = 5.0, @@ -116,7 +124,7 @@ def sees( point: Any, camera_info: CameraInfo, tf: TFLookup | None = None, - base_to_optical: Transform | None = None, + base_to_optical: TransformStamped | None = None, world_frame: str = "world", optical_frame: str = "camera_optical", time_tolerance: float = 5.0, @@ -151,8 +159,8 @@ def sees( raise ValueError("sees needs either tf or base_to_optical") target = np.array([point.x, point.y, point.z]) if hasattr(point, "x") else np.asarray(point) - fx, fy = camera_info.K[0], camera_info.K[4] - cx, cy = camera_info.K[2], camera_info.K[5] + fx, fy = camera_info.k[0], camera_info.k[4] + cx, cy = camera_info.k[2], camera_info.k[5] width, height = float(camera_info.width), float(camera_info.height) if extent is not None: @@ -174,7 +182,7 @@ def _sees(obs: Observation[Any]) -> bool: ) if transform is None: return False - matrix = transform.to_matrix() + matrix = transform_matrix(transform.transform) pts = (matrix @ np.column_stack([samples, np.ones(len(samples))]).T).T[:, :3] in_front = pts[:, 2] > 0 if not in_front.any(): @@ -192,8 +200,9 @@ def _sees(obs: Observation[Any]) -> bool: if depth is not None and visible.any(): depth_frame = depth(obs) if depth_frame is not None: - depth_m = np.asarray(depth_frame.data, dtype=np.float32) - if depth_frame.data.dtype == np.uint16: + pixels = image_view(depth_frame) + depth_m = np.asarray(pixels, dtype=np.float32) + if pixels.dtype.kind == "u" and pixels.dtype.itemsize == 2: depth_m = depth_m * 0.001 rows = np.clip(v[visible].astype(int), 0, depth_m.shape[0] - 1) cols = np.clip(u[visible].astype(int), 0, depth_m.shape[1] - 1) @@ -225,7 +234,7 @@ def __init__( pointcloud: PointCloud2 | Callable[[Observation[ImageDetections2D]], PointCloud2 | None], camera_info: CameraInfo, tf: TFLookup | None = None, - base_to_optical: Transform | None = None, + base_to_optical: TransformStamped | None = None, optical_frame: str = "camera_optical", time_tolerance: float = 5.0, filters: list[PointCloudFilter] | None = None, @@ -249,7 +258,7 @@ def __call__( continue transform = _world_to_optical( obs, - pointcloud.frame_id, + pointcloud.header.frame_id, self.tf, self.base_to_optical, self.optical_frame, diff --git a/dimos/perception/detection/reid/module.py b/dimos/perception/detection/reid/module.py index c5420a42e23..e409e87f33f 100644 --- a/dimos/perception/detection/reid/module.py +++ b/dimos/perception/detection/reid/module.py @@ -12,21 +12,40 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.vision_msgs.msg import Detection2DArray from reactivex import operators as ops from reactivex.observable import Observable from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray +from dimos.msgs.time import to_seconds from dimos.perception.detection.reid.embedding_id_system import EmbeddingIDSystem from dimos.perception.detection.reid.type import IDSystem from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D -from dimos.types.timestamped import align_timestamped +from dimos.types.timestamped import TimestampedData, align_timestamped from dimos.utils.reactive import backpressure +def _timed_image(image: Image) -> TimestampedData[Image]: + return TimestampedData(image, to_seconds(image.header.stamp)) + + +def _timed_detections(detections: Detection2DArray) -> TimestampedData[Detection2DArray]: + return TimestampedData(detections, to_seconds(detections.header.stamp)) + + +def _has_detections(detections: Detection2DArray) -> bool: + return bool(detections.detections) + + +def _image_detections( + pair: tuple[TimestampedData[Image], TimestampedData[Detection2DArray]], +) -> ImageDetections2D: + return ImageDetections2D.from_ros_detection2d_array(pair[0].value, pair[1].value) + + class Config(ModuleConfig): idsystem: IDSystem @@ -37,7 +56,6 @@ class ReidModule(Module): image: In[Image] def __init__(self, idsystem: IDSystem | None = None, **kwargs) -> None: # type: ignore[no-untyped-def] - super().__init__(**kwargs) if idsystem is None: try: from dimos.models.embedding.treid import TorchReIDModel @@ -48,18 +66,19 @@ def __init__(self, idsystem: IDSystem | None = None, **kwargs) -> None: # type: "TorchReIDModel not available. Please install with: pip install dimos[torchreid]" ) from e + super().__init__(idsystem=idsystem, **kwargs) self.idsystem = idsystem def detections_stream(self) -> Observable[ImageDetections2D]: return backpressure( align_timestamped( - self.image.pure_observable(), + self.image.pure_observable().pipe(ops.map(_timed_image)), self.detections.pure_observable().pipe( - ops.filter(lambda d: d.detections_length > 0) # type: ignore[attr-defined] + ops.filter(_has_detections), ops.map(_timed_detections) ), match_tolerance=0.0, buffer_size=2.0, - ).pipe(ops.map(lambda pair: ImageDetections2D.from_ros_detection2d_array(*pair))) # type: ignore[arg-type, misc] + ).pipe(ops.map(_image_detections)) ) @rpc diff --git a/dimos/perception/detection/reid/test_generated_module.py b/dimos/perception/detection/reid/test_generated_module.py new file mode 100644 index 00000000000..1443f2d0c9d --- /dev/null +++ b/dimos/perception/detection/reid/test_generated_module.py @@ -0,0 +1,91 @@ +# 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. + +from threading import Event +from unittest.mock import MagicMock + +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header +from dimos_generated.vision_msgs.msg import ( + BoundingBox2D, + Detection2D, + Detection2DArray, + Point2D, + Pose2D, +) +import numpy as np +import pytest +from reactivex.scheduler import ThreadPoolScheduler +from reactivex.subject import Subject + +from dimos.msgs.image import image_from_array +from dimos.msgs.time import time_from_seconds +from dimos.perception.detection.reid.module import ReidModule +from dimos.perception.detection.reid.type import PassthroughIDSystem + + +def test_generated_detection_alignment_uses_headers_and_filters_empty_arrays( + monkeypatch: pytest.MonkeyPatch, +) -> None: + scheduler = ThreadPoolScheduler(max_workers=2) + monkeypatch.setattr("dimos.utils.reactive.get_scheduler", lambda: scheduler) + module = ReidModule(idsystem=PassthroughIDSystem()) + images: Subject[Image] = Subject() + detections: Subject[Detection2DArray] = Subject() + module.image = MagicMock() + module.image.pure_observable.return_value = images + module.detections = MagicMock() + module.detections.pure_observable.return_value = detections + received = [] + errors = [] + ready = Event() + + def receive(value: object) -> None: + received.append(value) + ready.set() + + subscription = module.detections_stream().subscribe(receive, errors.append) + header = Header(stamp=time_from_seconds(1.25), frame_id="camera") + image = image_from_array(np.zeros((32, 32, 3), dtype=np.uint8), encoding="rgb8", header=header) + try: + detections.on_next(Detection2DArray(header=header)) + detections.on_next( + Detection2DArray( + header=header, + detections=[ + Detection2D( + header=header, + id="7", + bbox=BoundingBox2D( + center=Pose2D(position=Point2D(x=16, y=16)), size_x=8, size_y=8 + ), + ) + ], + ) + ) + images.on_next(image) + assert ready.wait(timeout=2.0) + assert not errors + (result,) = received + assert result.image is image + (detection,) = result.detections + assert detection.track_id == 7 + assert detection.bbox == (12, 12, 20, 20) + assert result.ts == 1.25 + finally: + subscription.dispose() + images.dispose() + detections.dispose() + module.stop() + scheduler.executor.shutdown(wait=True) diff --git a/dimos/perception/detection/type/detection2d/base.py b/dimos/perception/detection/type/detection2d/base.py index 7ee1d5bf1ad..0abc4a7329b 100644 --- a/dimos/perception/detection/type/detection2d/base.py +++ b/dimos/perception/detection/type/detection2d/base.py @@ -15,9 +15,9 @@ from abc import abstractmethod from collections.abc import Callable -from dimos_lcm.vision_msgs import Detection2D as ROSDetection2D +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.vision_msgs.msg import Detection2D as ROSDetection2D -from dimos.msgs.sensor_msgs.Image import Image from dimos.types.timestamped import Timestamped diff --git a/dimos/perception/detection/type/detection2d/bbox.py b/dimos/perception/detection/type/detection2d/bbox.py index 8facd3ac776..301fc1b4bf8 100644 --- a/dimos/perception/detection/type/detection2d/bbox.py +++ b/dimos/perception/detection/type/detection2d/bbox.py @@ -19,12 +19,11 @@ from typing import TYPE_CHECKING, Any if TYPE_CHECKING: + from dimos_generated.sensor_msgs.msg import Image from typing_extensions import Self from ultralytics.engine.results import Results - from dimos.msgs.sensor_msgs.Image import Image - -from dimos_lcm.vision_msgs import ( +from dimos_generated.vision_msgs.msg import ( BoundingBox2D, Detection2D as ROSDetection2D, ObjectHypothesis, @@ -32,12 +31,13 @@ Point2D, Pose2D, ) +import numpy as np from rich.console import Console from rich.text import Text -from dimos.msgs.std_msgs.Header import Header +from dimos.msgs.image import image_from_array, image_to_bgr, image_view +from dimos.msgs.time import to_seconds from dimos.perception.detection.type.detection2d.base import Detection2D -from dimos.types.timestamped import to_timestamp from dimos.utils.decorators.decorators import simple_mcache Bbox = tuple[float, float, float, float] @@ -94,7 +94,6 @@ def to_repr_dict(self) -> dict[str, Any]: def draw_on(self, img: Any, scale: float = 1.0) -> None: """Draw this detection's bbox and label onto a BGR numpy array (in-place).""" import cv2 - import numpy as np x1, y1, x2, y2 = map(int, self.bbox) @@ -127,11 +126,9 @@ def draw_on(self, img: Any, scale: float = 1.0) -> None: def annotated_image(self, scale: float = 1.0) -> Image: """Return the full image with this detection's bbox and label drawn on it.""" - img = self.image.to_opencv().copy() + img = image_to_bgr(self.image) self.draw_on(img, scale=scale) - from dimos.msgs.sensor_msgs.Image import Image - - return Image.from_opencv(img, ts=self.ts) + return image_from_array(img, encoding="bgr8", header=self.image.header) # return focused image, only on the bbox def cropped_image(self, padding: int = 20) -> Image: @@ -144,9 +141,12 @@ def cropped_image(self, padding: int = 20) -> Image: Cropped Image containing only the detection area plus padding """ x1, y1, x2, y2 = map(int, self.bbox) - return self.image.crop( - x1 - padding, y1 - padding, x2 - x1 + 2 * padding, y2 - y1 + 2 * padding - ) + pixels = image_view(self.image) + cropped = pixels[ + max(0, y1 - padding) : min(self.image.height, y2 + padding), + max(0, x1 - padding) : min(self.image.width, x2 + padding), + ] + return image_from_array(cropped, encoding=self.image.encoding, header=self.image.header) def __str__(self) -> str: console = Console(force_terminal=True, legacy_windows=False) @@ -202,8 +202,8 @@ def is_valid(self) -> bool: return False # Check if within image bounds (if image has shape) - if self.image.shape: - h, w = self.image.shape[:2] + if self.image.height and self.image.width: + h, w = self.image.height, self.image.width if not (0 <= x1 <= w and 0 <= y1 <= h and 0 <= x2 <= w and 0 <= y2 <= h): return False @@ -255,7 +255,7 @@ def from_ultralytics_result(cls, result: Results, idx: int, image: Image) -> Det class_id=class_id, confidence=confidence, name=name, - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, ) @@ -299,14 +299,14 @@ def from_ros_detection2d(cls, ros_det: ROSDetection2D, **kwargs) -> Self: # typ confidence = 0.0 if ros_det.results: hypothesis = ros_det.results[0].hypothesis - class_id = hypothesis.class_id + class_id = int(hypothesis.class_id) confidence = hypothesis.score # Extract track_id track_id = int(ros_det.id) if ros_det.id.isdigit() else 0 # Extract timestamp - ts = to_timestamp(ros_det.header.stamp) + ts = to_seconds(ros_det.header.stamp) name = kwargs.pop("name", f"class_{class_id}") @@ -322,12 +322,12 @@ def from_ros_detection2d(cls, ros_det: ROSDetection2D, **kwargs) -> Self: # typ def to_ros_detection2d(self) -> ROSDetection2D: return ROSDetection2D( - header=Header(self.ts, "camera_link"), + header=self.image.header, bbox=self.to_ros_bbox(), results=[ ObjectHypothesisWithPose( - ObjectHypothesis( - class_id=self.class_id, + hypothesis=ObjectHypothesis( + class_id=str(self.class_id), score=self.confidence, ) ) diff --git a/dimos/perception/detection/type/detection2d/imageDetections2D.py b/dimos/perception/detection/type/detection2d/imageDetections2D.py index 507125c333f..7b775cff544 100644 --- a/dimos/perception/detection/type/detection2d/imageDetections2D.py +++ b/dimos/perception/detection/type/detection2d/imageDetections2D.py @@ -25,11 +25,10 @@ from dimos.perception.detection.type.imageDetections import ImageDetections if TYPE_CHECKING: + from dimos_generated.sensor_msgs.msg import Image + from dimos_generated.vision_msgs.msg import Detection2DArray from ultralytics.engine.results import Results - from dimos.msgs.sensor_msgs.Image import Image - from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray - T2D = TypeVar("T2D", bound=Detection2D, default=Detection2DBBox) diff --git a/dimos/perception/detection/type/detection2d/person.py b/dimos/perception/detection/type/detection2d/person.py index 30100756e1e..94efb9e97b0 100644 --- a/dimos/perception/detection/type/detection2d/person.py +++ b/dimos/perception/detection/type/detection2d/person.py @@ -17,9 +17,10 @@ # Import for type checking only to avoid circular imports from typing import TYPE_CHECKING +from dimos_generated.sensor_msgs.msg import Image import numpy as np -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.time import to_seconds from dimos.perception.detection.type.detection2d.bbox import Bbox, Detection2DBBox from dimos.utils.decorators.decorators import simple_mcache @@ -152,7 +153,7 @@ def from_ultralytics_result( class_id=class_id, confidence=confidence, name=name, - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, # Person specific fields keypoints=keypoints, diff --git a/dimos/perception/detection/type/detection2d/point.py b/dimos/perception/detection/type/detection2d/point.py index 4e6d14ed27e..4042fd76ea0 100644 --- a/dimos/perception/detection/type/detection2d/point.py +++ b/dimos/perception/detection/type/detection2d/point.py @@ -17,7 +17,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING -from dimos_lcm.vision_msgs import ( +from dimos_generated.vision_msgs.msg import ( BoundingBox2D, Detection2D as ROSDetection2D, ObjectHypothesis, @@ -26,11 +26,11 @@ Pose2D, ) -from dimos.msgs.std_msgs.Header import Header +from dimos.msgs.image import image_from_array, image_view from dimos.perception.detection.type.detection2d.base import Detection2D if TYPE_CHECKING: - from dimos.msgs.sensor_msgs.Image import Image + from dimos_generated.sensor_msgs.msg import Image @dataclass @@ -65,17 +65,17 @@ def cropped_image(self, padding: int = 20) -> Image: Cropped Image containing the area around the point """ x, y = int(self.x), int(self.y) - return self.image.crop( - x - padding, - y - padding, - 2 * padding, - 2 * padding, - ) + pixels = image_view(self.image) + cropped = pixels[ + max(0, y - padding) : min(self.image.height, y + padding), + max(0, x - padding) : min(self.image.width, x + padding), + ] + return image_from_array(cropped, encoding=self.image.encoding, header=self.image.header) def to_ros_detection2d(self) -> ROSDetection2D: """Convert point to ROS Detection2D message (as zero-size bbox at point).""" return ROSDetection2D( - header=Header(self.ts, "camera_link"), + header=self.image.header, bbox=BoundingBox2D( center=Pose2D( position=Point2D(x=self.x, y=self.y), @@ -86,8 +86,8 @@ def to_ros_detection2d(self) -> ROSDetection2D: ), results=[ ObjectHypothesisWithPose( - ObjectHypothesis( - class_id=self.class_id, + hypothesis=ObjectHypothesis( + class_id=str(self.class_id), score=self.confidence, ) ) @@ -97,7 +97,7 @@ def to_ros_detection2d(self) -> ROSDetection2D: def is_valid(self) -> bool: """Check if the point is within image bounds.""" - if self.image.shape: - h, w = self.image.shape[:2] + if self.image.height and self.image.width: + h, w = self.image.height, self.image.width return bool(0 <= self.x <= w and 0 <= self.y <= h) return True diff --git a/dimos/perception/detection/type/detection2d/seg.py b/dimos/perception/detection/type/detection2d/seg.py index d96f2a8e7a1..852e0af12da 100644 --- a/dimos/perception/detection/type/detection2d/seg.py +++ b/dimos/perception/detection/type/detection2d/seg.py @@ -19,14 +19,14 @@ import numpy as np +from dimos.msgs.time import to_seconds from dimos.perception.detection.type.detection2d.bbox import Bbox, Detection2DBBox if TYPE_CHECKING: + from dimos_generated.sensor_msgs.msg import Image import torch from ultralytics.engine.results import Results - from dimos.msgs.sensor_msgs.Image import Image - @dataclass class Detection2DSeg(Detection2DBBox): @@ -90,7 +90,7 @@ def from_sam2_result( class_id=class_id, confidence=confidence, name=name, - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, mask=mask.astype(np.uint8), ) @@ -164,7 +164,7 @@ def from_ultralytics_result(cls, result: Results, idx: int, image: Image) -> Det class_id=class_id, confidence=confidence, name=name, - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, mask=mask, ) diff --git a/dimos/perception/detection/type/detection2d/test_bbox.py b/dimos/perception/detection/type/detection2d/test_bbox.py index cca4c8e59c2..312841ff60a 100644 --- a/dimos/perception/detection/type/detection2d/test_bbox.py +++ b/dimos/perception/detection/type/detection2d/test_bbox.py @@ -13,6 +13,8 @@ # limitations under the License. import pytest +from dimos.msgs.image import image_view + @pytest.mark.self_hosted def test_detection2d(detection2d) -> None: @@ -69,7 +71,7 @@ def test_detection2d(detection2d) -> None: # The actual cropped image is (260, 192, 3) assert cropped.width == 192 assert cropped.height == 260 - assert cropped.shape == (260, 192, 3) + assert image_view(cropped).shape == (260, 192, 3) # def test_to_ros_bbox(detection2d): """Test ROS bounding box conversion.""" diff --git a/dimos/perception/detection/type/detection2d/test_generated_bbox.py b/dimos/perception/detection/type/detection2d/test_generated_bbox.py new file mode 100644 index 00000000000..56648f90c71 --- /dev/null +++ b/dimos/perception/detection/type/detection2d/test_generated_bbox.py @@ -0,0 +1,101 @@ +# Copyright 2025-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. + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.std_msgs.msg import Header +from dimos_generated.vision_msgs.msg import Detection2D, Detection2DArray +import numpy as np + +from dimos.msgs.image import image_from_array, image_view +from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox +from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D +from dimos.perception.detection.type.detection2d.point import Detection2DPoint + + +def test_generated_bbox_image_and_wire_round_trip(): + header = Header(frame_id="camera_optical", stamp=Time(sec=1700000000, nanosec=123456789)) + pixels = np.zeros((40, 60, 3), dtype=np.uint8) + pixels[:, :, 0] = 250 + image = image_from_array(pixels, encoding="rgb8", header=header) + det = Detection2DBBox( + bbox=(10, 10, 30, 30), + track_id=7, + class_id=2, + confidence=0.9, + name="target", + ts=1.0, + image=image, + ) + assert det.is_valid() + crop = det.cropped_image(padding=2) + assert (crop.width, crop.height) == (24, 24) + assert crop.header == header + np.testing.assert_array_equal(image_view(crop), pixels[8:32, 8:32]) + annotated = det.annotated_image() + assert annotated.encoding == "bgr8" + assert annotated.header == header + np.testing.assert_array_equal(image_view(annotated)[39, 59], [0, 0, 250]) + np.testing.assert_array_equal(image_view(image), pixels) + wire = Detection2D.decode(det.to_ros_detection2d().encode()) + assert wire.header == header + assert wire.results[0].hypothesis.class_id == "2" + restored = Detection2DBBox.from_ros_detection2d(wire, image=image) + assert restored.bbox == det.bbox + assert restored.track_id == 7 + assert restored.class_id == 2 + + +def test_generated_detection_collection_round_trip(): + pixels = np.zeros((40, 60, 3), dtype=np.uint8) + image = image_from_array( + pixels, encoding="rgb8", header=Header(frame_id="camera", stamp=Time(nanosec=123)) + ) + det = Detection2DBBox( + bbox=(10, 10, 30, 30), + track_id=7, + class_id=2, + confidence=0.9, + name="target", + ts=0, + image=image, + ) + collection = ImageDetections2D(image=image, detections=[det]) + wire = Detection2DArray.decode(collection.to_ros_detection2d_array().encode()) + assert wire.header == image.header + assert len(wire.detections) == 1 + restored = ImageDetections2D.from_ros_detection2d_array(image, wire) + assert len(restored) == 1 + assert restored[0].bbox == det.bbox + annotated = collection.annotated_image() + assert annotated.header == image.header + assert annotated.encoding == "bgr8" + assert np.any(image_view(annotated)) + assert not np.any(image_view(image)) + + +def test_generated_point_detection_crop_and_cdr(): + pixels = np.arange(100, dtype=np.uint8).reshape(10, 10) + image = image_from_array( + pixels, encoding="mono8", header=Header(frame_id="camera", stamp=Time(nanosec=789)) + ) + point = Detection2DPoint(x=1, y=1, name="point", ts=0, image=image) + assert point.is_valid() + crop = point.cropped_image(padding=2) + np.testing.assert_array_equal(image_view(crop), pixels[:3, :3]) + assert crop.header == image.header + wire = Detection2D.decode(point.to_ros_detection2d().encode()) + assert wire.header == image.header + assert wire.bbox.center.position.x == wire.bbox.center.position.y == 1 + assert wire.bbox.size_x == wire.bbox.size_y == 0 + assert wire.results[0].hypothesis.class_id == "-1" diff --git a/dimos/perception/detection/type/detection2d/test_person.py b/dimos/perception/detection/type/detection2d/test_person.py index 3223380f4dc..59fc51e9a0d 100644 --- a/dimos/perception/detection/type/detection2d/test_person.py +++ b/dimos/perception/detection/type/detection2d/test_person.py @@ -18,14 +18,14 @@ def test_person_ros_confidence() -> None: """Test that Detection2DPerson preserves confidence when converting to ROS format.""" - from dimos.msgs.sensor_msgs.Image import Image + from dimos.msgs.image import image_from_file from dimos.perception.detection.detectors.person.yolo import YoloPersonDetector from dimos.perception.detection.type.detection2d.person import Detection2DPerson from dimos.utils.data import get_data # Load test image image_path = get_data("cafe.jpg") - image = Image.from_file(image_path) + image = image_from_file(image_path) # Run pose detection detector = YoloPersonDetector(device="cpu") diff --git a/dimos/perception/detection/type/detection3d/base.py b/dimos/perception/detection/type/detection3d/base.py index afe37aac6e2..f3b95df4d50 100644 --- a/dimos/perception/detection/type/detection3d/base.py +++ b/dimos/perception/detection/type/detection3d/base.py @@ -18,11 +18,12 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING -from dimos.msgs.geometry_msgs.Transform import Transform +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped + from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox if TYPE_CHECKING: - from dimos_lcm.sensor_msgs import CameraInfo + from dimos_generated.sensor_msgs.msg import CameraInfo @dataclass @@ -30,7 +31,9 @@ class Detection3D(Detection2DBBox): """Abstract base class for 3D detections.""" frame_id: str = "" - transform: Transform = field(default_factory=Transform.identity) + transform: TransformStamped = field( + default_factory=lambda: TransformStamped(transform=Transform(rotation=Quaternion(w=1))) + ) @classmethod @abstractmethod @@ -39,7 +42,7 @@ def from_2d( det: Detection2DBBox, distance: float, camera_info: CameraInfo, - world_to_optical_transform: Transform, + world_to_optical_transform: TransformStamped, ) -> Detection3D | None: """Create a 3D detection from a 2D detection.""" ... diff --git a/dimos/perception/detection/type/detection3d/bbox.py b/dimos/perception/detection/type/detection3d/bbox.py index be3b9b2f96d..e12078effe7 100644 --- a/dimos/perception/detection/type/detection3d/bbox.py +++ b/dimos/perception/detection/type/detection3d/bbox.py @@ -18,15 +18,23 @@ import functools from typing import Any -from dimos_lcm.vision_msgs import BoundingBox3D, ObjectHypothesis, ObjectHypothesisWithPose - -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.std_msgs.Header import Header -from dimos.msgs.vision_msgs.Detection3D import Detection3D +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + TransformStamped, + Vector3, +) +from dimos_generated.std_msgs.msg import Header +from dimos_generated.vision_msgs.msg import ( + BoundingBox3D, + Detection3D, + ObjectHypothesis, + ObjectHypothesisWithPose, +) + +from dimos.msgs.geometry import point_distance from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox @@ -39,9 +47,9 @@ class Detection3DBBox(Detection2DBBox): center: Vector3 # Center point in world frame size: Vector3 # Width, height, depth - transform: Transform | None = None # Camera to world transform + transform: TransformStamped | None = None # Camera to world transform frame_id: str = "" # Frame ID (e.g., "world", "map") - orientation: Quaternion = field(default_factory=lambda: Quaternion(0.0, 0.0, 0.0, 1.0)) + orientation: Quaternion = field(default_factory=lambda: Quaternion(w=1)) @functools.cached_property def pose(self) -> PoseStamped: @@ -50,16 +58,17 @@ def pose(self) -> PoseStamped: Returns pose in world frame with the detection's orientation. """ return PoseStamped( - ts=self.ts, - frame_id=self.frame_id, - position=self.center, - orientation=self.orientation, + header=Header(stamp=self.image.header.stamp, frame_id=self.frame_id), + pose=Pose( + position=Point(x=self.center.x, y=self.center.y, z=self.center.z), + orientation=self.orientation, + ), ) def to_detection3d_msg(self) -> Detection3D: """Convert to ROS Detection3D message.""" msg = Detection3D() - msg.header = Header(self.ts, self.frame_id) + msg.header = Header(stamp=self.image.header.stamp, frame_id=self.frame_id) # Results msg.results = [ @@ -70,13 +79,10 @@ def to_detection3d_msg(self) -> Detection3D: ) ) ] - msg.results_length = len(msg.results) - # The generated LCM constructor reuses a shared default bbox object. - # Assign a new BoundingBox3D so this message does not share bbox state with others. msg.bbox = BoundingBox3D( center=Pose( - position=self.center, + position=Point(x=self.center.x, y=self.center.y, z=self.center.z), orientation=self.orientation, ), size=self.size, @@ -88,8 +94,11 @@ def to_repr_dict(self) -> dict[str, Any]: # Calculate distance from camera if self.transform is None: return super().to_repr_dict() - camera_pos = self.transform.translation - distance = (self.center - camera_pos).magnitude() + camera_pos = self.transform.transform.translation + distance = point_distance( + Point(x=self.center.x, y=self.center.y, z=self.center.z), + Point(x=camera_pos.x, y=camera_pos.y, z=camera_pos.z), + ) parent_dict = super().to_repr_dict() # Remove bbox key if present diff --git a/dimos/perception/detection/type/detection3d/imageDetections3D.py b/dimos/perception/detection/type/detection3d/imageDetections3D.py index 93b33d15660..c6e55afd9ed 100644 --- a/dimos/perception/detection/type/detection3d/imageDetections3D.py +++ b/dimos/perception/detection/type/detection3d/imageDetections3D.py @@ -16,10 +16,10 @@ from typing import Generic +from dimos_generated.std_msgs.msg import Header +from dimos_generated.vision_msgs.msg import Detection3DArray from typing_extensions import TypeVar -from dimos.msgs.std_msgs.Header import Header -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray from dimos.perception.detection.type.detection3d.bbox import Detection3DBBox from dimos.perception.detection.type.imageDetections import ImageDetections @@ -36,7 +36,6 @@ def to_ros_detection3d_array(self, frame_id: str | None = None) -> Detection3DAr detections = [det.to_detection3d_msg() for det in self.detections] return Detection3DArray( - detections_length=len(detections), - header=Header(self.image.ts, resolved_frame_id), + header=Header(stamp=self.image.header.stamp, frame_id=resolved_frame_id), detections=detections, ) diff --git a/dimos/perception/detection/type/detection3d/imageDetections3DPC.py b/dimos/perception/detection/type/detection3d/imageDetections3DPC.py index 310b2651da0..e45ba357aed 100644 --- a/dimos/perception/detection/type/detection3d/imageDetections3DPC.py +++ b/dimos/perception/detection/type/detection3d/imageDetections3DPC.py @@ -20,11 +20,9 @@ from dimos.perception.detection.type.imageDetections import ImageDetections if TYPE_CHECKING: - from dimos_lcm.sensor_msgs import CameraInfo + from dimos_generated.geometry_msgs.msg import TransformStamped + from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 - from dimos.msgs.geometry_msgs.Transform import Transform - from dimos.msgs.sensor_msgs.Image import Image - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.detection.type.detection3d.pointcloud_filters import PointCloudFilter @@ -38,7 +36,7 @@ def from_2d( detections_2d: ImageDetections2D, world_pointcloud: PointCloud2, camera_info: CameraInfo, - world_to_optical_transform: Transform, + world_to_optical_transform: TransformStamped, filters: list[PointCloudFilter] | None = None, ) -> ImageDetections3DPC: """Project every 2D detection into 3D, dropping any that yield no valid points. @@ -77,7 +75,7 @@ def from_depth( detections_2d: ImageDetections2D, depth: Image, camera_info: CameraInfo, - world_to_optical_transform: Transform, + world_to_optical_transform: TransformStamped, filters: list[PointCloudFilter] | None = None, ) -> ImageDetections3DPC: """Unproject every detection's depth pixels into 3D, dropping empty results.""" diff --git a/dimos/perception/detection/type/detection3d/marker.py b/dimos/perception/detection/type/detection3d/marker.py index 9fd58cc6d85..f28104ddfa1 100644 --- a/dimos/perception/detection/type/detection3d/marker.py +++ b/dimos/perception/detection/type/detection3d/marker.py @@ -17,9 +17,9 @@ from dataclasses import dataclass, field from typing import Any +from dimos_generated.vision_msgs.msg import Detection3D import numpy as np -from dimos.msgs.vision_msgs.Detection3D import Detection3D from dimos.perception.detection.type.detection3d.bbox import Detection3DBBox @@ -35,7 +35,9 @@ class Detection3DMarker(Detection3DBBox): """ marker_id: int = -1 - corners_px: np.ndarray = field(default_factory=lambda: np.zeros((4, 2), dtype=np.float32)) + corners_px: np.ndarray[Any, np.dtype[np.float32]] = field( + default_factory=lambda: np.zeros((4, 2), dtype=np.float32) + ) dictionary: str = "" reprojection_error: float = 0.0 @@ -52,8 +54,9 @@ def to_detection3d_msg(self) -> Detection3D: msg = super().to_detection3d_msg() msg.id = str(self.marker_id) if msg.results: - msg.results[0].hypothesis.class_id = self.marker_label - msg.results_length = len(msg.results) + result = msg.results[0] + result.hypothesis.class_id = self.marker_label + msg.results[0] = result return msg def to_repr_dict(self) -> dict[str, Any]: diff --git a/dimos/perception/detection/type/detection3d/pointcloud_filters.py b/dimos/perception/detection/type/detection3d/pointcloud_filters.py index f55d0d09f36..1c83be49d03 100644 --- a/dimos/perception/detection/type/detection3d/pointcloud_filters.py +++ b/dimos/perception/detection/type/detection3d/pointcloud_filters.py @@ -15,34 +15,51 @@ from __future__ import annotations from collections.abc import Callable +from typing import TYPE_CHECKING -from dimos_lcm.sensor_msgs import CameraInfo +from dimos_generated.geometry_msgs.msg import TransformStamped +from dimos_generated.sensor_msgs.msg import CameraInfo, PointCloud2 +import numpy as np -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox +from dimos.msgs.geometry import inverse_transform +from dimos.msgs.pointcloud import pointcloud_xyz, select_points + +if TYPE_CHECKING: + import open3d as o3d + + from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox -# Filters take Detection2DBBox, PointCloud2, CameraInfo, Transform and return filtered PointCloud2 or None PointCloudFilter = Callable[ - [Detection2DBBox, PointCloud2, CameraInfo, Transform], PointCloud2 | None + ["Detection2DBBox", PointCloud2, CameraInfo, TransformStamped], PointCloud2 | None ] +def _open3d_cloud(pc: PointCloud2) -> o3d.geometry.PointCloud: + import open3d as o3d + + return o3d.geometry.PointCloud(o3d.utility.Vector3dVector(pointcloud_xyz(pc))) + + +def _selected(pc: PointCloud2, indices: list[int]) -> PointCloud2: + keep = np.zeros(pc.width * pc.height, dtype=bool) + keep[indices] = True + return select_points(pc, keep) + + def height_filter(height: float = 0.1) -> PointCloudFilter: - return lambda det, pc, ci, tf: pc.filter_by_height(height) + return lambda det, pc, ci, tf: select_points(pc, pointcloud_xyz(pc)[:, 2] >= height) def statistical(nb_neighbors: int = 40, std_ratio: float = 0.5) -> PointCloudFilter: def filter_func( - det: Detection2DBBox, pc: PointCloud2, ci: CameraInfo, tf: Transform + det: Detection2DBBox, pc: PointCloud2, ci: CameraInfo, tf: TransformStamped ) -> PointCloud2 | None: try: - statistical, _removed = pc.pointcloud.remove_statistical_outlier( + _, indices = _open3d_cloud(pc).remove_statistical_outlier( nb_neighbors=nb_neighbors, std_ratio=std_ratio ) - return PointCloud2(statistical, pc.frame_id, pc.ts) - except Exception: - # print("statistical filter failed:", e) + return _selected(pc, indices) + except RuntimeError: return None return filter_func @@ -82,33 +99,27 @@ def filter_func( def raycast() -> PointCloudFilter: def filter_func( - det: Detection2DBBox, pc: PointCloud2, ci: CameraInfo, tf: Transform + det: Detection2DBBox, pc: PointCloud2, ci: CameraInfo, tf: TransformStamped ) -> PointCloud2 | None: try: - camera_pos = tf.inverse().translation - camera_pos_np = camera_pos.to_numpy() - _, visible_indices = pc.pointcloud.hidden_point_removal(camera_pos_np, radius=100.0) - visible_pcd = pc.pointcloud.select_by_index(visible_indices) - return PointCloud2(visible_pcd, pc.frame_id, pc.ts) - except Exception: - # print("raycast filter failed:", e) + camera_pos = inverse_transform(tf).transform.translation + _, indices = _open3d_cloud(pc).hidden_point_removal( + np.array([camera_pos.x, camera_pos.y, camera_pos.z]), radius=100.0 + ) + return _selected(pc, indices) + except RuntimeError: return None return filter_func def radius_outlier(min_neighbors: int = 20, radius: float = 0.3) -> PointCloudFilter: - """ - Remove isolated points: keep only points that have at least `min_neighbors` - neighbors within `radius` meters (same units as your point cloud). - """ + """Keep points with at least ``min_neighbors`` within ``radius`` meters.""" def filter_func( - det: Detection2DBBox, pc: PointCloud2, ci: CameraInfo, tf: Transform + det: Detection2DBBox, pc: PointCloud2, ci: CameraInfo, tf: TransformStamped ) -> PointCloud2 | None: - filtered_pcd, _removed = pc.pointcloud.remove_radius_outlier( - nb_points=min_neighbors, radius=radius - ) - return PointCloud2(filtered_pcd, pc.frame_id, pc.ts) + _, indices = _open3d_cloud(pc).remove_radius_outlier(nb_points=min_neighbors, radius=radius) + return _selected(pc, indices) return filter_func diff --git a/dimos/perception/detection/type/detection3d/test_generated_filters.py b/dimos/perception/detection/type/detection3d/test_generated_filters.py new file mode 100644 index 00000000000..7ae5a158ec2 --- /dev/null +++ b/dimos/perception/detection/type/detection3d/test_generated_filters.py @@ -0,0 +1,68 @@ +# Copyright 2025-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. + +from types import SimpleNamespace + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped +from dimos_generated.sensor_msgs.msg import CameraInfo, PointCloud2, PointField +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest + +from dimos.msgs.pointcloud import pointcloud_view +from dimos.perception.detection.type.detection3d.pointcloud_filters import ( + height_filter, + radius_outlier, + raycast, + statistical, +) + + +@pytest.mark.parametrize("make_filter", [height_filter, statistical, radius_outlier, raycast]) +def test_filters_preserve_custom_point_records(make_filter): + rng = np.random.default_rng(42) + values = np.zeros(101, dtype=[("x", " Image: - return Image( - data=np.zeros((80, 100, 3), dtype=np.uint8), - format=ImageFormat.BGR, - frame_id="camera_optical", - ts=ts, + return image_from_array( + np.zeros((80, 100, 3), dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id="camera_optical", stamp=time_from_seconds(ts)), ) @@ -52,12 +53,12 @@ def _marker( class_id=marker_id, confidence=confidence, name="", - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, - center=Vector3(float(marker_id), 2.0, 3.0), - size=Vector3(0.16, 0.16, 0.0), + center=Vector3(x=float(marker_id), y=2.0, z=3.0), + size=Vector3(x=0.16, y=0.16, z=0.0), frame_id=frame_id, - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), marker_id=marker_id, corners_px=np.array( [[x1, y1], [x2, y1], [x2, y2], [x1, y2]], @@ -76,22 +77,22 @@ def test_to_ros_detection3d_array_serializes_plain_bbox_results() -> None: class_id=9, confidence=0.75, name="box", - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, - center=Vector3(1.0, 2.0, 3.0), - size=Vector3(0.4, 0.5, 0.6), + center=Vector3(x=1.0, y=2.0, z=3.0), + size=Vector3(x=0.4, y=0.5, z=0.6), frame_id="world", - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), ) msg = ImageDetections3D(image, [det]).to_ros_detection3d_array() - assert msg.detections_length == 1 - assert msg.detections[0].results_length == 1 + assert len(msg.detections) == 1 + assert len(msg.detections[0].results) == 1 - decoded = Detection3DArray.lcm_decode(msg.lcm_encode()) + decoded = Detection3DArray.decode(msg.encode()) decoded_det = decoded.detections[0] - assert decoded_det.results_length == 1 + assert len(decoded_det.results) == 1 assert decoded_det.results[0].hypothesis.class_id == "9" assert decoded_det.results[0].hypothesis.score == pytest.approx(0.75) assert decoded_det.bbox.center.position.x == pytest.approx(1.0) @@ -111,23 +112,23 @@ def test_to_ros_detection3d_array_preserves_marker_wire_identity() -> None: msg = detections.to_ros_detection3d_array() assert msg.header.frame_id == "world" - assert msg.ts == pytest.approx(image.ts) - assert msg.detections_length == 2 + assert to_seconds(msg.header.stamp) == pytest.approx(to_seconds(image.header.stamp)) + assert len(msg.detections) == 2 assert len(msg.detections) == 2 first = msg.detections[0] assert first.header.frame_id == "world" assert first.id == "7" - assert first.results_length == 1 + assert len(first.results) == 1 assert first.results[0].hypothesis.class_id == "DICT_APRILTAG_36h11:7" assert first.results[0].hypothesis.score == pytest.approx(1.0) assert first.bbox.center.position.x == pytest.approx(7.0) assert first.bbox.size.x == pytest.approx(0.16) assert first.bbox.size.z == pytest.approx(0.0) - decoded = Detection3DArray.lcm_decode(msg.lcm_encode()) + decoded = Detection3DArray.decode(msg.encode()) assert decoded.header.frame_id == "world" - assert decoded.detections_length == 2 + assert len(decoded.detections) == 2 assert decoded.detections[1].id == "42" assert decoded.detections[1].results[0].hypothesis.class_id == "DICT_APRILTAG_36h11:42" @@ -140,15 +141,15 @@ def test_to_ros_detection3d_array_uses_override_and_handles_empty_frames() -> No ) assert msg.header.frame_id == "map" - assert msg.ts == pytest.approx(image.ts) - assert msg.detections_length == 1 + assert to_seconds(msg.header.stamp) == pytest.approx(to_seconds(image.header.stamp)) + assert len(msg.detections) == 1 empty = ImageDetections3D(image, []).to_ros_detection3d_array(frame_id="world") assert empty.header.frame_id == "world" - assert empty.ts == pytest.approx(image.ts) - assert empty.detections_length == 0 - assert empty.detections == [] + assert to_seconds(empty.header.stamp) == pytest.approx(to_seconds(image.header.stamp)) + assert len(empty.detections) == 0 + assert len(empty.detections) == 0 def test_filter_and_annotated_image_work_for_3d_marker_detections() -> None: @@ -162,10 +163,28 @@ def test_filter_and_annotated_image_work_for_3d_marker_detections() -> None: assert isinstance(filtered, ImageDetections3D) assert filtered.detections == [keep] ros_msg = filtered.to_ros_detection3d_array() - assert ros_msg.detections_length == 1 + assert len(ros_msg.detections) == 1 assert ros_msg.detections[0].id == "5" annotated = filtered.annotated_image() - assert annotated.ts == pytest.approx(image.ts) - assert np.count_nonzero(annotated.data) > 0 - assert np.count_nonzero(image.data) == 0 + assert to_seconds(annotated.header.stamp) == pytest.approx(to_seconds(image.header.stamp)) + assert np.count_nonzero(image_view(annotated)) > 0 + assert np.count_nonzero(image_view(image)) == 0 + + +def test_marker_pose_and_array_keep_exact_source_stamp() -> None: + image = _image() + image.header.stamp.sec = 1700000000 + image.header.stamp.nanosec = 123456789 + marker = _marker(image, marker_id=42) + pose = marker.pose + message = marker.to_detection3d_msg() + array = ImageDetections3D(image, [marker]).to_ros_detection3d_array() + decoded = Detection3DArray.decode(array.encode()) + assert pose.header.stamp == image.header.stamp + assert message.header.stamp == image.header.stamp + assert decoded.header.stamp == image.header.stamp + assert decoded.detections[0].header.stamp == image.header.stamp + marker.center.x = 999 + assert pose.pose.position.x == 42 + assert decoded.detections[0].bbox.center.position.x == 42 diff --git a/dimos/perception/detection/type/detection3d/test_marker.py b/dimos/perception/detection/type/detection3d/test_marker.py index 79a71f0debb..6d9cb9e4e43 100644 --- a/dimos/perception/detection/type/detection3d/test_marker.py +++ b/dimos/perception/detection/type/detection3d/test_marker.py @@ -13,13 +13,14 @@ # limitations under the License. import cv2 +from dimos_generated.geometry_msgs.msg import Quaternion, Vector3 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.vision_msgs.msg import Detection3D import numpy as np import pytest -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.vision_msgs.Detection3D import Detection3D +from dimos.msgs.image import image_from_array +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.perception.detection.type.detection3d.marker import Detection3DMarker from dimos.perception.fiducial.marker_pose import marker_reprojection_error @@ -41,9 +42,9 @@ def _project_synthetic_marker_corners( def test_marker_detection3d_msg_preserves_marker_identity_on_wire() -> None: - center = Vector3(1.0, 2.0, 3.0) - orientation = Quaternion(0.1, 0.2, 0.3, 0.9) - size = Vector3(0.16, 0.16, 0.0) + center = Vector3(x=1.0, y=2.0, z=3.0) + orientation = Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) + size = Vector3(x=0.16, y=0.16, z=0.0) marker_length_m = 0.16 camera_matrix = np.array( [[420.0, 0.0, 60.0], [0.0, 420.0, 50.0], [0.0, 0.0, 1.0]], @@ -68,10 +69,10 @@ def test_marker_detection3d_msg_preserves_marker_identity_on_wire() -> None: rvec, tvec, ) - image = Image( - data=np.zeros((100, 120, 3), dtype=np.uint8), - frame_id="camera_optical", - ts=123.456, + image = image_from_array( + np.zeros((100, 120, 3), dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id="camera_optical", stamp=time_from_seconds(123.456)), ) det = Detection3DMarker( @@ -80,7 +81,7 @@ def test_marker_detection3d_msg_preserves_marker_identity_on_wire() -> None: class_id=999, confidence=1.0, name="marker_42", - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, center=center, size=size, @@ -99,7 +100,7 @@ def test_marker_detection3d_msg_preserves_marker_identity_on_wire() -> None: assert msg.header.frame_id == "world" assert msg.id == "42" - assert msg.results_length == 1 + assert len(msg.results) == 1 assert msg.results[0].hypothesis.class_id == "DICT_APRILTAG_36h11:42" assert msg.results[0].hypothesis.score == pytest.approx(1.0) assert msg.bbox.center.position.x == pytest.approx(center.x) @@ -113,7 +114,7 @@ def test_marker_detection3d_msg_preserves_marker_identity_on_wire() -> None: assert msg.bbox.size.y == pytest.approx(size.y) assert msg.bbox.size.z == pytest.approx(size.z) - decoded = Detection3D.lcm_decode(msg.lcm_encode()) + decoded = Detection3D.decode(msg.encode()) assert decoded.id == "42" - assert decoded.results_length == 1 + assert len(decoded.results) == 1 assert decoded.results[0].hypothesis.class_id == "DICT_APRILTAG_36h11:42" diff --git a/dimos/perception/detection/type/imageDetections.py b/dimos/perception/detection/type/imageDetections.py index 98fd0e5388f..199e17690b3 100644 --- a/dimos/perception/detection/type/imageDetections.py +++ b/dimos/perception/detection/type/imageDetections.py @@ -22,15 +22,17 @@ else: from typing_extensions import Self -from dimos_lcm.vision_msgs import Detection2DArray +from dimos_generated.vision_msgs.msg import Detection2DArray -from dimos.msgs.std_msgs.Header import Header +from dimos.msgs.image import image_from_array, image_to_bgr +from dimos.msgs.time import to_seconds from dimos.perception.detection.type.utils import TableStr if TYPE_CHECKING: from collections.abc import Callable, Iterator - from dimos.msgs.sensor_msgs.Image import Image + from dimos_generated.sensor_msgs.msg import Image + from dimos.perception.detection.type.detection2d.base import Detection2D T = TypeVar("T", bound=Detection2D) @@ -46,14 +48,14 @@ class ImageDetections(Generic[T], TableStr): @property def ts(self) -> float: - return self.image.ts + return to_seconds(self.image.header.stamp) def __init__(self, image: Image, detections: list[T] | None = None) -> None: self.image = image self.detections = detections or [] for det in self.detections: if not det.ts: - det.ts = image.ts + det.ts = to_seconds(image.header.stamp) def __len__(self) -> int: return len(self.detections) @@ -80,18 +82,15 @@ def filter(self, *predicates: Callable[[T], bool]) -> Self: def to_ros_detection2d_array(self) -> Detection2DArray: return Detection2DArray( - detections_length=len(self.detections), - header=Header(self.image.ts, "camera_optical"), + header=self.image.header, detections=[det.to_ros_detection2d() for det in self.detections], ) def annotated_image(self, scale: float = 1.0) -> Image: """Return the image with all detection bboxes and labels drawn on it.""" - img = self.image.to_opencv().copy() + img = image_to_bgr(self.image) for det in self.detections: if hasattr(det, "draw_on"): det.draw_on(img, scale=scale) - from dimos.msgs.sensor_msgs.Image import Image as ImageMsg - - return ImageMsg.from_opencv(img, ts=self.image.ts) + return image_from_array(img, encoding="bgr8", header=self.image.header) diff --git a/dimos/perception/experimental/moduleDB.py b/dimos/perception/experimental/moduleDB.py index 305bca063a0..4f5d6ddc266 100644 --- a/dimos/perception/experimental/moduleDB.py +++ b/dimos/perception/experimental/moduleDB.py @@ -16,17 +16,16 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Vector3 +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.vision_msgs.msg import Detection2DArray +import numpy as np from reactivex.observable import Observable from dimos.core.core import rpc from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray +from dimos.msgs.pointcloud import concatenate_clouds from dimos.perception.detection.module3D import Detection3DModule from dimos.perception.detection.type.detection3d.imageDetections3DPC import ImageDetections3DPC from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC @@ -45,7 +44,11 @@ def to_repr_dict(self) -> dict[str, Any]: center_str = "None" else: center_str = ( - "[" + ", ".join(list(map(lambda n: f"{n:1f}", self.center.to_list()))) + "]" + "[" + + ", ".join( + list(map(lambda n: f"{n:1f}", [self.center.x, self.center.y, self.center.z])) + ) + + "]" ) return { "object_id": self.track_id, @@ -82,9 +85,16 @@ def __add__(self, detection: Detection3DPC) -> "Object3D": new_object.class_id = self.class_id new_object.name = self.name new_object.transform = self.transform - new_object.pointcloud = self.pointcloud + detection.pointcloud + new_object.pointcloud = concatenate_clouds(self.pointcloud, detection.pointcloud) new_object.frame_id = self.frame_id - new_object.center = (self.center + detection.center) / 2 + if self.center is None: + raise ValueError("Cannot merge an object without a center") + center = detection.center + new_object.center = Vector3( + x=(self.center.x + center.x) / 2, + y=(self.center.y + center.y) / 2, + z=(self.center.z + center.z) / 2, + ) new_object.detections = self.detections + 1 if detection.bbox_2d_volume() > self.bbox_2d_volume(): @@ -113,26 +123,14 @@ def to_pose(self) -> PoseStamped: if self.best_detection is None or self.center is None: raise ValueError("Cannot compute pose without best_detection and center") - optical_inverse = Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(-0.5, 0.5, -0.5, 0.5), - frame_id="camera_link", - child_frame_id="camera_optical", - ).inverse() - - print("transform is", self.best_detection.transform) - - global_transform = optical_inverse + self.best_detection.transform - - print("inverse optical is", global_transform) - - print("obj center is", self.center) - global_pose = global_transform.to_pose() - print("Global pose:", global_pose) - global_pose.frame_id = self.best_detection.frame_id - print("remap to", self.best_detection.frame_id) return PoseStamped( - position=self.center, orientation=Quaternion(), frame_id=self.best_detection.frame_id + header=Header( + frame_id=self.best_detection.frame_id, stamp=self.pointcloud.header.stamp + ), + pose=Pose( + position=Point(x=self.center.x, y=self.center.y, z=self.center.z), + orientation=Quaternion(w=1), + ), ) @@ -177,13 +175,26 @@ def __init__(self, *args, **kwargs) -> None: # type: ignore[no-untyped-def] def closest_object(self, detection: Detection3DPC) -> Object3D | None: # Filter objects to only those with matching names - matching_objects = [obj for obj in self.objects.values() if obj.name == detection.name] + matching_objects = [ + obj + for obj in self.objects.values() + if obj.name == detection.name and obj.center is not None + ] if not matching_objects: return None # Sort by distance - distances = sorted(matching_objects, key=lambda obj: detection.center.distance(obj.center)) + def distance(obj: Object3D) -> float: + assert obj.center is not None + center = detection.center + return float( + np.linalg.norm( + [center.x - obj.center.x, center.y - obj.center.y, center.z - obj.center.z] + ) + ) + + distances = sorted(matching_objects, key=distance) return distances[0] diff --git a/dimos/perception/experimental/object.py b/dimos/perception/experimental/object.py index 76f4cc425c1..2c568201481 100644 --- a/dimos/perception/experimental/object.py +++ b/dimos/perception/experimental/object.py @@ -19,25 +19,41 @@ from typing import TYPE_CHECKING, Any import uuid -from dimos_lcm.geometry_msgs import Pose -from dimos_lcm.vision_msgs import BoundingBox3D, ObjectHypothesis, ObjectHypothesisWithPose +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + TransformStamped, + Vector3, +) +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.vision_msgs.msg import ( + BoundingBox3D, + Detection3D as ROSDetection3D, + Detection3DArray, + ObjectHypothesis, + ObjectHypothesisWithPose, +) import numpy as np -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.std_msgs.Header import Header -from dimos.msgs.vision_msgs.Detection3D import Detection3D as ROSDetection3D -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.msgs.geometry import quaternion_from_matrix, transform_matrix +from dimos.msgs.image import image_to_rgb, image_view +from dimos.msgs.pointcloud import ( + concatenate_clouds, + pointcloud_from_xyz, + pointcloud_from_xyz_rgb, + pointcloud_rgb, + pointcloud_to_open3d, + pointcloud_xyz, + voxel_downsample_cloud, +) +from dimos.msgs.time import time_from_seconds from dimos.perception.detection.type.detection2d.seg import Detection2DSeg from dimos.perception.detection.type.detection3d.base import Detection3D if TYPE_CHECKING: - from dimos_lcm.sensor_msgs import CameraInfo - from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D @@ -54,7 +70,7 @@ class Object(Detection3D): size: Vector3 pose: PoseStamped pointcloud: PointCloud2 - camera_transform: Transform | None = None + camera_transform: TransformStamped | None = None mask: np.ndarray[Any, np.dtype[np.uint8]] | None = None detections_count: int = 1 visual_embedding: np.ndarray[Any, np.dtype[np.float32]] | None = None @@ -68,8 +84,8 @@ def __post_init__(self) -> None: def set_center(self, center: Vector3) -> None: """Update the canonical center and its pose representation together.""" - self.center = Vector3(center) - self.pose.position = self.center + self.center = Vector3(x=center.x, y=center.y, z=center.z) + self.pose.pose.position = Point(x=center.x, y=center.y, z=center.z) def update_object(self, other: Object, *, accumulate_pointcloud: bool = True) -> None: """Update this object with data from another detection. @@ -89,7 +105,7 @@ def update_object(self, other: Object, *, accumulate_pointcloud: bool = True) -> and other.camera_transform is not None and self.camera_transform is not None ): - self.pointcloud = self.pointcloud + other.pointcloud + self.pointcloud = concatenate_clouds(self.pointcloud, other.pointcloud) else: self.pointcloud = other.pointcloud @@ -115,20 +131,20 @@ def update_object(self, other: Object, *, accumulate_pointcloud: bool = True) -> def get_oriented_bounding_box(self) -> Any: """Get oriented bounding box of the pointcloud.""" - return self.pointcloud.oriented_bounding_box + return pointcloud_to_open3d(self.pointcloud).get_oriented_bounding_box() def _detection3d_bbox_components(self) -> tuple[Vector3, Quaternion, Vector3]: """Return canonical geometry without refitting the point cloud.""" center = self.center size = self.size - orientation = self.pose.orientation + orientation = self.pose.pose.orientation return ( - Vector3(center.x, center.y, center.z), - Quaternion(orientation.x, orientation.y, orientation.z, orientation.w), + Vector3(x=center.x, y=center.y, z=center.z), + Quaternion(x=orientation.x, y=orientation.y, z=orientation.z, w=orientation.w), Vector3( - max(float(size.x), 1e-3), - max(float(size.y), 1e-3), - max(float(size.z), 1e-3), + x=max(float(size.x), 1e-3), + y=max(float(size.y), 1e-3), + z=max(float(size.z), 1e-3), ), ) @@ -143,7 +159,7 @@ def to_detection3d_msg(self) -> ROSDetection3D: center, orientation, size = self._detection3d_bbox_components() msg = ROSDetection3D() - msg.header = Header(self.ts, self.frame_id) + msg.header = self.pose.header msg.id = self.object_id msg.results = [ ObjectHypothesisWithPose( @@ -153,10 +169,9 @@ def to_detection3d_msg(self) -> ROSDetection3D: ) ) ] - msg.results_length = len(msg.results) msg.bbox = BoundingBox3D( center=Pose( - position=center, + position=Point(x=center.x, y=center.y, z=center.z), orientation=orientation, ), size=size, @@ -179,6 +194,7 @@ def agent_encode(self) -> dict[str, Any]: def to_dict(self) -> dict[str, Any]: """Convert object to dictionary with all relevant data.""" + colors = pointcloud_rgb(self.pointcloud) return { "object_id": self.object_id, "track_id": self.track_id, @@ -188,8 +204,11 @@ def to_dict(self) -> dict[str, Any]: "identity_basis": self.identity_basis, "last_seen_ts": self.last_seen_ts, "mask": self.mask, - "pointcloud": self.pointcloud.as_numpy(), - "image": self.image.as_numpy() if self.image else None, + "pointcloud": ( + pointcloud_xyz(self.pointcloud), + None if colors is None else colors.astype(np.float64) / 255.0, + ), + "image": image_view(self.image) if self.image else None, } @classmethod @@ -199,7 +218,7 @@ def from_2d_to_list( color_image: Image, depth_image: Image, camera_info: CameraInfo, - camera_transform: Transform | None = None, + camera_transform: TransformStamped | None = None, depth_scale: float = 1.0, depth_trunc: float = 10.0, statistical_nb_neighbors: int = 10, @@ -241,16 +260,13 @@ def from_2d_to_list( import cv2 import open3d as o3d # type: ignore[import-untyped] - color_cv = color_image.to_opencv() - if color_cv.ndim == 3 and color_cv.shape[2] == 3: - color_cv = cv2.cvtColor(color_cv, cv2.COLOR_BGR2RGB) - - depth_cv = depth_image.to_opencv() + color_cv = image_to_rgb(color_image) + depth_cv = image_view(depth_image) h, w = depth_cv.shape[:2] # Build Open3D camera intrinsics - fx, fy = camera_info.K[0], camera_info.K[4] - cx, cy = camera_info.K[2], camera_info.K[5] + fx, fy = camera_info.k[0], camera_info.k[4] + cx, cy = camera_info.k[2], camera_info.k[5] intrinsic_o3d = o3d.camera.PinholeCameraIntrinsic(w, h, fx, fy, cx, cy) objects: list[Object] = [] @@ -289,56 +305,56 @@ def from_2d_to_list( ) pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd, intrinsic_o3d) - pc0 = PointCloud2( - pcd, - frame_id=depth_image.frame_id, - ts=depth_image.ts, - ).voxel_downsample(voxel_downsample) - - pcd_filtered, _ = pc0.pointcloud.remove_statistical_outlier( + initial = pointcloud_from_xyz_rgb( + np.asarray(pcd.points), + (np.asarray(pcd.colors) * 255).astype(np.uint8), + header=depth_image.header, + ) + pc0 = voxel_downsample_cloud(initial, voxel_downsample) + pcd_filtered, _ = pointcloud_to_open3d(pc0).remove_statistical_outlier( nb_neighbors=statistical_nb_neighbors, std_ratio=statistical_std_ratio, ) - if len(pcd_filtered.points) < 10: continue - - pc = PointCloud2( - pcd_filtered, - frame_id=depth_image.frame_id, - ts=depth_image.ts, - ) - - # Transform pointcloud to world frame if camera_transform is provided + header = depth_image.header if camera_transform is not None: - pc = pc.transform(camera_transform) - frame_id = camera_transform.frame_id - else: - frame_id = depth_image.frame_id + if camera_transform.child_frame_id != header.frame_id: + raise ValueError("Camera transform source frame does not match depth image") + pcd_filtered.transform(transform_matrix(camera_transform.transform)) + header = Header( + stamp=depth_image.header.stamp, frame_id=camera_transform.header.frame_id + ) + pc = pointcloud_from_xyz_rgb( + np.asarray(pcd_filtered.points), + (np.asarray(pcd_filtered.colors) * 255).astype(np.uint8), + header=header, + ) + frame_id = header.frame_id # Compute bounding box: AABB for stable upright obstacles, OBB for tighter fit if use_aabb: - aabb = pc.pointcloud.get_axis_aligned_bounding_box() + aabb = pcd_filtered.get_axis_aligned_bounding_box() aabb_center = (aabb.min_bound + aabb.max_bound) / 2.0 aabb_extent = aabb.max_bound - aabb.min_bound - center = Vector3(aabb_center[0], aabb_center[1], aabb_center[2]) + center = Vector3(x=aabb_center[0], y=aabb_center[1], z=aabb_center[2]) sx, sy, sz = float(aabb_extent[0]), float(aabb_extent[1]), float(aabb_extent[2]) - orientation = Quaternion(0.0, 0.0, 0.0, 1.0) + orientation = Quaternion(w=1.0) else: - obb = pc.pointcloud.get_oriented_bounding_box() - center = Vector3(obb.center[0], obb.center[1], obb.center[2]) + obb = pcd_filtered.get_oriented_bounding_box() + center = Vector3(x=obb.center[0], y=obb.center[1], z=obb.center[2]) sx, sy, sz = float(obb.extent[0]), float(obb.extent[1]), float(obb.extent[2]) - orientation = Quaternion.from_rotation_matrix(obb.R) + orientation = quaternion_from_matrix(np.asarray(obb.R)) if max_obstacle_width > 0: sx = min(sx, max_obstacle_width) sy = min(sy, max_obstacle_width) - size = Vector3(sx, sy, sz) + size = Vector3(x=sx, y=sy, z=sz) pose = PoseStamped( - ts=det.ts, - frame_id=frame_id, - position=center, - orientation=orientation, + header=Header(stamp=time_from_seconds(det.ts), frame_id=frame_id), + pose=Pose( + position=Point(x=center.x, y=center.y, z=center.z), orientation=orientation + ), ) # Skip objects too far from origin (background detections) @@ -380,52 +396,38 @@ def aggregate_pointclouds(objects: list[Object]) -> PointCloud2: Returns: Combined PointCloud2 with all points colored by object (empty if no points). """ - import open3d as o3d # type: ignore[import-untyped] - - if not objects: - return PointCloud2(pointcloud=o3d.geometry.PointCloud(), frame_id="", ts=0.0) - + header = ( + Header() + if not objects + else Header(frame_id=objects[0].frame_id, stamp=objects[0].pointcloud.header.stamp) + ) all_points = [] all_colors = [] - for obj in objects: - points, colors = obj.pointcloud.as_numpy() + if obj.frame_id != header.frame_id: + raise ValueError("Cannot aggregate object clouds in different frames") + points = pointcloud_xyz(obj.pointcloud) if len(points) == 0: continue - + colors = pointcloud_rgb(obj.pointcloud) try: seed = int(obj.object_id, 16) except (ValueError, TypeError): seed = abs(hash(obj.object_id)) rng = np.random.default_rng(abs(seed)) track_color = rng.integers(50, 255, 3) / 255.0 - - if colors is not None: - blended = np.clip(0.6 * colors + 0.4 * track_color, 0.0, 1.0) - else: - blended = np.tile(track_color, (len(points), 1)) - + blended = ( + np.clip(0.6 * colors.astype(np.float64) / 255.0 + 0.4 * track_color, 0.0, 1.0) + if colors is not None + else np.tile(track_color, (len(points), 1)) + ) all_points.append(points) all_colors.append(blended) - if not all_points: - return PointCloud2( - pointcloud=o3d.geometry.PointCloud(), frame_id=objects[0].frame_id, ts=objects[0].ts - ) - - combined_points = np.vstack(all_points) - combined_colors = np.vstack(all_colors) - - pc = PointCloud2.from_numpy( - combined_points, - frame_id=objects[0].frame_id, - timestamp=objects[0].ts, + return pointcloud_from_xyz(np.empty((0, 3), dtype=np.float32), header=header) + return pointcloud_from_xyz_rgb( + np.vstack(all_points), (np.vstack(all_colors) * 255).astype(np.uint8), header=header ) - pcd = pc.pointcloud - pcd.colors = o3d.utility.Vector3dVector(combined_colors) - pc.pointcloud = pcd - - return pc def to_detection3d_array( @@ -452,7 +454,6 @@ def to_detection3d_array( if resolved_ts is None: resolved_ts = objects[0].ts if objects else 0.0 return Detection3DArray( - detections_length=len(detections), - header=Header(resolved_ts, resolved_frame_id), + header=Header(stamp=time_from_seconds(resolved_ts), frame_id=resolved_frame_id), detections=detections, ) diff --git a/dimos/perception/experimental/objectDB.py b/dimos/perception/experimental/objectDB.py index d2d8e7e1bea..2dda55d4924 100644 --- a/dimos/perception/experimental/objectDB.py +++ b/dimos/perception/experimental/objectDB.py @@ -14,15 +14,19 @@ from __future__ import annotations +import math import threading import time from typing import TYPE_CHECKING, Any -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +import numpy as np + +from dimos.msgs.pointcloud import pointcloud_from_xyz from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: - from dimos.msgs.geometry_msgs.Vector3 import Vector3 + from dimos_generated.geometry_msgs.msg import Vector3 + from dimos.perception.experimental.object import Object logger = setup_logger() @@ -176,20 +180,19 @@ def find_nearest( if not candidates: return None - return min(candidates, key=lambda obj: position.distance(obj.center)) + return min( + candidates, + key=lambda obj: math.dist( + (position.x, position.y, position.z), (obj.center.x, obj.center.y, obj.center.z) + ), + ) def clear(self) -> None: """Clear all objects from the database.""" - import open3d as o3d # type: ignore[import-untyped] - with self._lock: - # Drop Open3D pointcloud references before clearing to reduce shutdown warnings. + # Release large generated point buffers, retaining metadata on held objects. for obj in list(self._pending_objects.values()) + list(self._objects.values()): - obj.pointcloud = PointCloud2( - pointcloud=o3d.geometry.PointCloud(), - frame_id=obj.pointcloud.frame_id, - ts=obj.pointcloud.ts, - ) + obj.pointcloud = pointcloud_from_xyz(np.empty((0, 3)), header=obj.pointcloud.header) self._pending_objects.clear() self._objects.clear() self._track_id_map.clear() @@ -283,13 +286,22 @@ def _match_by_distance(self, obj: Object) -> Object | None: candidates = [ o for o in all_objects - if o.center is not None and obj.center.distance(o.center) < self._distance_threshold + if o.center is not None + and math.dist( + (obj.center.x, obj.center.y, obj.center.z), (o.center.x, o.center.y, o.center.z) + ) + < self._distance_threshold ] if not candidates: return None - return min(candidates, key=lambda o: obj.center.distance(o.center)) + return min( + candidates, + key=lambda o: math.dist( + (obj.center.x, obj.center.y, obj.center.z), (o.center.x, o.center.y, o.center.z) + ), + ) def _prune_stale_pending(self, now: float) -> None: if self._pending_ttl_s <= 0: diff --git a/dimos/perception/experimental/object_scene_registration.py b/dimos/perception/experimental/object_scene_registration.py index 5e43cb490d0..82c6c3a7337 100644 --- a/dimos/perception/experimental/object_scene_registration.py +++ b/dimos/perception/experimental/object_scene_registration.py @@ -19,8 +19,13 @@ import time from typing import Any, Literal +from dimos_generated.geometry_msgs.msg import TransformStamped +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 +from dimos_generated.tf2_msgs.msg import TFMessage +from dimos_generated.vision_msgs.msg import Detection2DArray, Detection3DArray import numpy as np from numpy.typing import NDArray +from reactivex import operators as ops from dimos.agents.annotation import skill from dimos.core.core import rpc @@ -28,14 +33,9 @@ from dimos.core.stream import In, Out from dimos.models.segmentation.edge_tam import BoxPromptImageSegmenter from dimos.models.segmentation.yoloe import YoloeBoxSegmenter -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.std_msgs.Header import Header -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.msgs.image import image_from_array, image_view +from dimos.msgs.pointcloud import pointcloud_from_xyz, transform_cloud +from dimos.msgs.time import to_seconds from dimos.perception.detection.detectors.yoloe import Yoloe2DDetector, YoloePromptMode from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.experimental.object import ( @@ -45,13 +45,23 @@ to_detection3d_array, ) from dimos.perception.experimental.objectDB import ObjectDB -from dimos.types.timestamped import align_timestamped +from dimos.types.timestamped import TimestampedData, align_timestamped from dimos.utils.logging_config import setup_logger from dimos.utils.reactive import backpressure logger = setup_logger() +def _timed_image(image: Image) -> TimestampedData[Image]: + return TimestampedData(image, to_seconds(image.header.stamp)) + + +def _aligned_images( + pair: tuple[TimestampedData[Image], TimestampedData[Image]], +) -> tuple[Image, Image]: + return pair[0].value, pair[1].value + + class ObjectSceneRegistrationConfig(ModuleConfig): target_frame: str = "base_link" prompt_mode: YoloePromptMode = YoloePromptMode.LRPC @@ -89,7 +99,7 @@ class ObjectSceneRegistrationModule(Module): _processing_lock: threading.RLock # A tuple assignment/read is atomic, so depth and its transform cannot be # observed from different frames by get_full_scene_pointcloud(). - _latest_scene_snapshot: tuple[Image, Transform | None] | None = None + _latest_scene_snapshot: tuple[Image, TransformStamped | None] | None = None config: ObjectSceneRegistrationConfig @@ -155,11 +165,11 @@ def start(self) -> None: self.camera_info.subscribe(lambda msg: setattr(self, "_camera_info", msg)) aligned_frames = align_timestamped( - self.color_image.observable(), - self.depth_image.observable(), + self.color_image.observable().pipe(ops.map(_timed_image)), + self.depth_image.observable().pipe(ops.map(_timed_image)), buffer_size=2.0, match_tolerance=0.1, - ) + ).pipe(ops.map(_aligned_images)) backpressure(aligned_frames).subscribe(self._on_aligned_frames) def _create_segmenter(self) -> BoxPromptImageSegmenter | None: @@ -232,7 +242,7 @@ def scan_scene(self, text: list[str] | None = None) -> Detection3DArray: return to_detection3d_array( objects, frame_id=self._target_frame, - ts=frames[0].ts, + ts=to_seconds(frames[0].header.stamp), ) def _scan_scene_objects(self, frames: tuple[Image, Image] | None = None) -> list[DetObject]: @@ -275,7 +285,7 @@ def get_object_pointcloud_by_object_id(self, object_id: str) -> PointCloud2 | No logger.warning(f"No object found with object_id='{object_id}'") return None pc = obj.pointcloud - num_points = len(pc.pointcloud.points) if pc else 0 + num_points = pc.width * pc.height logger.info(f"Found object '{object_id}' ({obj.name}) with {num_points} points") return pc @@ -313,7 +323,7 @@ def get_full_scene_pointcloud( return None depth_image, camera_transform = scene_snapshot - depth_cv = depth_image.to_opencv() + depth_cv = image_view(depth_image) h, w = depth_cv.shape[:2] # Zero out excluded object's depth @@ -324,8 +334,8 @@ def get_full_scene_pointcloud( depth_cv[exclude_mask > 0] = 0 # Build pointcloud from depth - fx, fy = self._camera_info.K[0], self._camera_info.K[4] - cx, cy = self._camera_info.K[2], self._camera_info.K[5] + fx, fy = self._camera_info.k[0], self._camera_info.k[4] + cx, cy = self._camera_info.k[2], self._camera_info.k[5] intrinsic = o3d.camera.PinholeCameraIntrinsic(w, h, fx, fy, cx, cy) depth_o3d = o3d.geometry.Image(depth_cv.astype(np.float32)) @@ -338,14 +348,9 @@ def get_full_scene_pointcloud( pcd = pcd.voxel_down_sample(voxel_size) - pc = PointCloud2( - pcd, - frame_id=depth_image.frame_id, - ts=depth_image.ts, - ) - + pc = pointcloud_from_xyz(np.asarray(pcd.points), header=depth_image.header) if camera_transform is not None: - pc = pc.transform(camera_transform) + pc = transform_cloud(pc, camera_transform) return pc @@ -412,21 +417,19 @@ def _process_images(self, color_msg: Image, depth_msg: Image) -> list[DetObject] color_image = color_msg # Convert depth to meters (float32) - depth_cv = depth_msg.to_opencv() - if depth_msg.format == ImageFormat.DEPTH16: + depth_cv = image_view(depth_msg) + if depth_msg.encoding == "16UC1": depth_cv = depth_cv.astype(np.float32) / 1000.0 elif depth_cv.dtype != np.float32: depth_cv = depth_cv.astype(np.float32) - depth_image = Image( - data=depth_cv, format=ImageFormat.DEPTH, frame_id=depth_msg.frame_id, ts=depth_msg.ts - ) + depth_image = image_from_array(depth_cv, encoding="32FC1", header=depth_msg.header) camera_transform = None - if self._target_frame != color_image.frame_id: + if self._target_frame != color_image.header.frame_id: camera_transform = self.tfbuffer.get( self._target_frame, - color_image.frame_id, - color_image.ts, + color_image.header.frame_id, + to_seconds(color_image.header.stamp), 0.1, forward_tolerance=0.2, ) @@ -457,8 +460,7 @@ def _process_images(self, color_msg: Image, depth_msg: Image) -> list[DetObject] detections_2d = self._segmenter.segment(detections_2d) detections_2d_msg = Detection2DArray( - detections_length=len(detections_2d.detections), - header=Header(color_image.ts, color_image.frame_id or ""), + header=color_image.header, detections=[det.to_ros_detection2d() for det in detections_2d.detections], ) self.detections_2d.publish(detections_2d_msg) @@ -473,13 +475,13 @@ def _process_3d_detections( detections_2d: ImageDetections2D[Any], color_image: Image, depth_image: Image, - camera_transform: Transform | None, + camera_transform: TransformStamped | None, ) -> list[DetObject]: """Convert 2D detections to 3D and publish.""" if self._camera_info is None: return [] - if self._target_frame != color_image.frame_id and camera_transform is None: + if self._target_frame != color_image.header.frame_id and camera_transform is None: logger.warning("Failed to lookup transform from camera frame to target frame") return [] diff --git a/dimos/perception/experimental/object_scene_registration_spec.py b/dimos/perception/experimental/object_scene_registration_spec.py index 5f0cbbf8e1b..f8869c87364 100644 --- a/dimos/perception/experimental/object_scene_registration_spec.py +++ b/dimos/perception/experimental/object_scene_registration_spec.py @@ -14,8 +14,9 @@ from typing import Protocol -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.vision_msgs.msg import Detection3DArray + from dimos.spec.utils import Spec diff --git a/dimos/perception/experimental/object_tracker.py b/dimos/perception/experimental/object_tracker.py index 2331dbfd8ca..045619a24dd 100644 --- a/dimos/perception/experimental/object_tracker.py +++ b/dimos/perception/experimental/object_tracker.py @@ -16,31 +16,38 @@ import time from typing import Any -# Import LCM messages -from dimos_lcm.vision_msgs import ( +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage + +# Generated detection values +from dimos_generated.vision_msgs.msg import ( Detection2D, + Detection2DArray, Detection3D, + Detection3DArray, ObjectHypothesisWithPose, ) import numpy as np from numpy.typing import NDArray +from reactivex import operators as ops from reactivex.disposable import Disposable from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.std_msgs.Header import Header -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray -from dimos.types.timestamped import align_timestamped +from dimos.msgs.image import image_from_array, image_to_rgb, image_view +from dimos.msgs.time import time_from_seconds, to_seconds +from dimos.types.timestamped import TimestampedData, align_timestamped from dimos.utils.logging_config import setup_logger from dimos.utils.transform_utils import ( euler_to_quaternion, @@ -48,6 +55,11 @@ yaw_towards_point, ) + +def _timed_image(image: Image) -> TimestampedData[Image]: + return TimestampedData(image, to_seconds(image.header.stamp)) + + logger = setup_logger() @@ -128,20 +140,20 @@ def start(self) -> None: # Subscribe to aligned rgb and depth streams def on_aligned_frames(frames_tuple) -> None: # type: ignore[no-untyped-def] - rgb_msg, depth_msg = frames_tuple + rgb_msg, depth_msg = (item.value for item in frames_tuple) with self._frame_lock: - self._latest_rgb_frame = rgb_msg.data + self._latest_rgb_frame = image_to_rgb(rgb_msg) - depth_data = depth_msg.data + depth_data = image_view(depth_msg) # Convert from millimeters to meters if depth is DEPTH16 format - if depth_msg.format == ImageFormat.DEPTH16: + if depth_msg.encoding in ("16UC1", "mono16"): depth_data = depth_data.astype(np.float32) / 1000.0 self._latest_depth_frame = depth_data # Create aligned observable for RGB and depth aligned_frames = align_timestamped( - self.color_image.observable(), - self.depth.observable(), + self.color_image.pure_observable().pipe(ops.map(_timed_image)), + self.depth.pure_observable().pipe(ops.map(_timed_image)), buffer_size=2.0, # 2 second buffer match_tolerance=0.5, # 500ms tolerance ) @@ -154,10 +166,10 @@ def on_camera_info(camera_info_msg: CameraInfo) -> None: # Extract intrinsics from camera info K matrix # K is a 3x3 matrix in row-major order: [fx, 0, cx, 0, fy, cy, 0, 0, 1] self.camera_intrinsics = [ # type: ignore[assignment] - camera_info_msg.K[0], - camera_info_msg.K[4], - camera_info_msg.K[2], - camera_info_msg.K[5], + camera_info_msg.k[0], + camera_info_msg.k[4], + camera_info_msg.k[2], + camera_info_msg.k[5], ] unsub = self.camera_info.subscribe(on_camera_info) # type: ignore[assignment] @@ -309,8 +321,8 @@ def _reset_tracking_state(self) -> None: self.tracking_frame_count = 0 # Reset frame counter # Publish empty detections to clear any visualizations - empty_2d = Detection2DArray(detections_length=0, header=Header(), detections=[]) - empty_3d = Detection3DArray(detections_length=0, header=Header(), detections=[]) + empty_2d = Detection2DArray(header=Header(), detections=[]) + empty_3d = Detection3DArray(header=Header(), detections=[]) self._latest_detection2d = empty_2d self._latest_detection3d = empty_3d self._detection_event.clear() @@ -409,9 +421,9 @@ def _process_tracking(self) -> None: return # Create detections if tracking succeeded - header = Header(self.frame_id) - detection2darray = Detection2DArray(detections_length=0, header=header, detections=[]) - detection3darray = Detection3DArray(detections_length=0, header=header, detections=[]) + header = Header(stamp=time_from_seconds(time.time()), frame_id=self.config.frame_id) + detection2darray = Detection2DArray(header=header, detections=[]) + detection3darray = Detection3DArray(header=header, detections=[]) if final_success and current_bbox_x1y1x2y2 is not None: x1, y1, x2, y2 = current_bbox_x1y1x2y2 @@ -423,7 +435,6 @@ def _process_tracking(self) -> None: # Create Detection2D detection_2d = Detection2D() detection_2d.id = "0" - detection_2d.results_length = 1 detection_2d.header = header # Create hypothesis @@ -440,7 +451,6 @@ def _process_tracking(self) -> None: detection_2d.bbox.size_y = height detection2darray = Detection2DArray() - detection2darray.detections_length = 1 detection2darray.header = header detection2darray.detections = [detection_2d] @@ -462,15 +472,15 @@ def _process_tracking(self) -> None: # Create pose in optical frame optical_pose = Pose() - optical_pose.position = Vector3(x_optical, y_optical, z_optical) - optical_pose.orientation = Quaternion(0.0, 0.0, 0.0, 1.0) # Identity for now + optical_pose.position = Point(x=x_optical, y=y_optical, z=z_optical) + optical_pose.orientation = Quaternion(w=1.0) # Identity for now # Convert to robot frame robot_pose = optical_to_robot_frame(optical_pose) # Calculate orientation: object facing towards camera (origin) yaw = yaw_towards_point(robot_pose.position) - euler = Vector3(0.0, 0.0, yaw) # Only yaw, no roll/pitch + euler = Vector3(z=yaw) # Only yaw, no roll/pitch robot_pose.orientation = euler_to_quaternion(euler) # Estimate object size in meters @@ -481,7 +491,6 @@ def _process_tracking(self) -> None: # Create Detection3D detection_3d = Detection3D() detection_3d.id = "0" - detection_3d.results_length = 1 detection_3d.header = header # Reuse hypothesis from 2D @@ -491,30 +500,34 @@ def _process_tracking(self) -> None: detection_3d.bbox.center = Pose() detection_3d.bbox.center.position = robot_pose.position detection_3d.bbox.center.orientation = robot_pose.orientation - detection_3d.bbox.size = Vector3(size_x, size_y, size_z) + detection_3d.bbox.size = Vector3(x=size_x, y=size_y, z=size_z) detection3darray = Detection3DArray() - detection3darray.detections_length = 1 detection3darray.header = header detection3darray.detections = [detection_3d] # Publish transform for tracked object # The optical pose is in camera optical frame, so publish it relative to the camera frame - tracked_object_tf = Transform( - translation=robot_pose.position, - rotation=robot_pose.orientation, - frame_id=self.frame_id, # Use configured camera frame + tracked_object_tf = TransformStamped( + header=header, child_frame_id="tracked_object", - ts=header.ts, + transform=Transform( + translation=Vector3( + x=robot_pose.position.x, + y=robot_pose.position.y, + z=robot_pose.position.z, + ), + rotation=robot_pose.orientation, + ), ) - self.tf.publish(TFMessage(tracked_object_tf)) + self.tf.publish(TFMessage(transforms=[tracked_object_tf])) # Store latest detections for RPC access self._latest_detection2d = detection2darray self._latest_detection3d = detection3darray # Signal that new detections are available - if detection2darray.detections_length > 0 or detection3darray.detections_length > 0: + if bool(detection2darray.detections) or bool(detection3darray.detections): self._detection_event.set() # Publish detections @@ -524,12 +537,8 @@ def _process_tracking(self) -> None: # Create and publish visualization if tracking is active if self.tracking_initialized: # Convert single detection to list for visualization - detections_3d = ( - detection3darray.detections if detection3darray.detections_length > 0 else [] - ) - detections_2d = ( - detection2darray.detections if detection2darray.detections_length > 0 else [] - ) + detections_3d = detection3darray.detections if bool(detection3darray.detections) else [] + detections_2d = detection2darray.detections if bool(detection2darray.detections) else [] if detections_3d and detections_2d: # Extract 2D bbox for visualization @@ -555,7 +564,9 @@ def _process_tracking(self) -> None: viz_image = self._draw_reid_matches(viz_image) # Convert to Image message and publish - viz_msg = Image.from_numpy(viz_image) + viz_msg = image_from_array( + viz_image, encoding="rgb8", header=self._latest_detection2d.header + ) self.tracked_overlay.publish(viz_msg) def _draw_reid_matches(self, image: NDArray[np.uint8]) -> NDArray[np.uint8]: diff --git a/dimos/perception/experimental/object_tracker_2d.py b/dimos/perception/experimental/object_tracker_2d.py index 385d0ff98d8..abb4b413729 100644 --- a/dimos/perception/experimental/object_tracker_2d.py +++ b/dimos/perception/experimental/object_tracker_2d.py @@ -17,10 +17,14 @@ import time from typing import Any -# Import LCM messages -from dimos_lcm.vision_msgs import ( +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header + +# Generated detection values +from dimos_generated.vision_msgs.msg import ( BoundingBox2D, Detection2D, + Detection2DArray, ObjectHypothesis, ObjectHypothesisWithPose, Point2D, @@ -34,9 +38,8 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.std_msgs.Header import Header -from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray +from dimos.msgs.image import image_from_array, image_to_rgb +from dimos.msgs.time import time_from_seconds from dimos.utils.logging_config import setup_logger logger = setup_logger(level=logging.INFO) @@ -91,7 +94,7 @@ def start(self) -> None: def on_frame(frame_msg: Image) -> None: arrival_time = time.perf_counter() with self._frame_lock: - self._latest_rgb_frame = frame_msg.data + self._latest_rgb_frame = image_to_rgb(frame_msg) self._frame_arrival_time = arrival_time unsub = self.color_image.subscribe(on_frame) @@ -176,7 +179,8 @@ def _reset_tracking_state(self) -> None: # Publish empty detection empty_2d = Detection2DArray( - detections_length=0, header=Header(time.time(), self.frame_id), detections=[] + header=Header(stamp=time_from_seconds(time.time()), frame_id=self.config.frame_id), + detections=[], ) self._latest_detection2d = empty_2d self.detection2darray.publish(empty_2d) @@ -260,12 +264,11 @@ def _process_tracking(self) -> None: height = float(y2 - y1) # Create 2D detection header - header = Header(time.time(), self.frame_id) + header = Header(stamp=time_from_seconds(time.time()), frame_id=self.config.frame_id) # Create Detection2D with all fields in constructors detection_2d = Detection2D( id="0", - results_length=1, header=header, bbox=BoundingBox2D( center=Pose2D(position=Point2D(x=center_x, y=center_y), theta=0.0), @@ -279,9 +282,7 @@ def _process_tracking(self) -> None: ], ) - detection2darray = Detection2DArray( - detections_length=1, header=header, detections=[detection_2d] - ) + detection2darray = Detection2DArray(header=header, detections=[detection_2d]) # Store and publish self._latest_detection2d = detection2darray @@ -290,7 +291,7 @@ def _process_tracking(self) -> None: # Create visualization viz_image = self._draw_visualization(frame, current_bbox_x1y1x2y2) viz_copy = viz_image.copy() # Force copy needed to prevent frame reuse - viz_msg = Image.from_numpy(viz_copy, format=ImageFormat.RGB) + viz_msg = image_from_array(viz_copy, encoding="rgb8", header=header) self.tracked_overlay.publish(viz_msg) def _draw_visualization(self, image: NDArray[np.uint8], bbox: list[int]) -> NDArray[np.uint8]: diff --git a/dimos/perception/experimental/object_tracker_3d.py b/dimos/perception/experimental/object_tracker_3d.py index f2f5a116b10..5377f259874 100644 --- a/dimos/perception/experimental/object_tracker_3d.py +++ b/dimos/perception/experimental/object_tracker_3d.py @@ -13,27 +13,33 @@ # limitations under the License. -# Import LCM messages -from dimos_lcm.sensor_msgs import CameraInfo -from dimos_lcm.vision_msgs import ( +# Generated detection values +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage +from dimos_generated.vision_msgs.msg import ( + Detection2DArray, Detection3D, + Detection3DArray, ObjectHypothesisWithPose, ) import numpy as np +from reactivex import operators as ops from dimos.core.core import rpc from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.std_msgs.Header import Header -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.msgs.image import image_from_array, image_to_rgb, image_view +from dimos.msgs.time import to_seconds from dimos.perception.experimental.object_tracker_2d import ObjectTracker2D -from dimos.types.timestamped import align_timestamped +from dimos.types.timestamped import TimestampedData, align_timestamped from dimos.utils.logging_config import setup_logger from dimos.utils.transform_utils import ( euler_to_quaternion, @@ -41,6 +47,11 @@ yaw_towards_point, ) + +def _timed_image(image: Image) -> TimestampedData[Image]: + return TimestampedData(image, to_seconds(image.header.stamp)) + + logger = setup_logger() @@ -78,20 +89,20 @@ def start(self) -> None: # Subscribe to aligned RGB and depth streams def on_aligned_frames(frames_tuple) -> None: # type: ignore[no-untyped-def] - rgb_msg, depth_msg = frames_tuple + rgb_msg, depth_msg = (item.value for item in frames_tuple) with self._frame_lock: - self._latest_rgb_frame = rgb_msg.data + self._latest_rgb_frame = image_to_rgb(rgb_msg) - depth_data = depth_msg.data + depth_data = image_view(depth_msg) # Convert from millimeters to meters if depth is DEPTH16 format - if depth_msg.format == ImageFormat.DEPTH16: + if depth_msg.encoding in ("16UC1", "mono16"): depth_data = depth_data.astype(np.float32) / 1000.0 self._latest_depth_frame = depth_data # Create aligned observable for RGB and depth aligned_frames = align_timestamped( - self.color_image.observable(), - self.depth.observable(), + self.color_image.pure_observable().pipe(ops.map(_timed_image)), + self.depth.pure_observable().pipe(ops.map(_timed_image)), buffer_size=2.0, # 2 second buffer match_tolerance=0.5, # 500ms tolerance ) @@ -103,10 +114,10 @@ def on_camera_info(camera_info_msg: CameraInfo) -> None: self._latest_camera_info = camera_info_msg # Extract intrinsics: K is [fx, 0, cx, 0, fy, cy, 0, 0, 1] self.camera_intrinsics = [ # type: ignore[assignment] - camera_info_msg.K[0], - camera_info_msg.K[4], - camera_info_msg.K[2], - camera_info_msg.K[5], + camera_info_msg.k[0], + camera_info_msg.k[4], + camera_info_msg.k[2], + camera_info_msg.k[5], ] self.camera_info.subscribe(on_camera_info) @@ -127,7 +138,7 @@ def _process_tracking(self) -> None: # Enhance with 3D if we have depth and a valid 2D detection if ( self._latest_detection2d - and self._latest_detection2d.detections_length > 0 + and bool(self._latest_detection2d.detections) and self._latest_depth_frame is not None and self.camera_intrinsics is not None ): @@ -161,12 +172,14 @@ def _process_tracking(self) -> None: if self.last_good_matches and self.last_roi_kps and self.last_roi_bbox: viz_image = self._draw_reid_overlay(viz_image) - viz_msg = Image.from_numpy(viz_image) + viz_msg = image_from_array( + viz_image, encoding="rgb8", header=self._latest_detection2d.header + ) self.tracked_overlay.publish(viz_msg) def _create_detection3d_from_2d(self, detection2d: Detection2DArray) -> Detection3DArray | None: """Create 3D detection from 2D detection using depth.""" - if detection2d.detections_length == 0: + if not detection2d.detections: return None det_2d = detection2d.detections[0] @@ -198,15 +211,15 @@ def _create_detection3d_from_2d(self, detection2d: Detection2DArray) -> Detectio # Create pose in optical frame optical_pose = Pose() - optical_pose.position = Vector3(x_optical, y_optical, z_optical) - optical_pose.orientation = Quaternion(0.0, 0.0, 0.0, 1.0) + optical_pose.position = Point(x=x_optical, y=y_optical, z=z_optical) + optical_pose.orientation = Quaternion(w=1.0) # Convert to robot frame robot_pose = optical_to_robot_frame(optical_pose) # Calculate orientation: object facing towards camera yaw = yaw_towards_point(robot_pose.position) - euler = Vector3(0.0, 0.0, yaw) + euler = Vector3(z=yaw) robot_pose.orientation = euler_to_quaternion(euler) # Estimate object size in meters @@ -215,10 +228,9 @@ def _create_detection3d_from_2d(self, detection2d: Detection2DArray) -> Detectio size_z = 0.1 # Default depth size # Create Detection3D - header = Header(self.frame_id) + header = Header(stamp=detection2d.header.stamp, frame_id=self.config.frame_id) detection_3d = Detection3D() detection_3d.id = "0" - detection_3d.results_length = 1 detection_3d.header = header # Create hypothesis @@ -231,22 +243,24 @@ def _create_detection3d_from_2d(self, detection2d: Detection2DArray) -> Detectio detection_3d.bbox.center = Pose() detection_3d.bbox.center.position = robot_pose.position detection_3d.bbox.center.orientation = robot_pose.orientation - detection_3d.bbox.size = Vector3(size_x, size_y, size_z) + detection_3d.bbox.size = Vector3(x=size_x, y=size_y, z=size_z) detection3darray = Detection3DArray() - detection3darray.detections_length = 1 detection3darray.header = header detection3darray.detections = [detection_3d] # Publish TF for tracked object - tracked_object_tf = Transform( - translation=robot_pose.position, - rotation=robot_pose.orientation, - frame_id=self.frame_id, + tracked_object_tf = TransformStamped( + header=header, child_frame_id="tracked_object", - ts=header.ts, + transform=Transform( + translation=Vector3( + x=robot_pose.position.x, y=robot_pose.position.y, z=robot_pose.position.z + ), + rotation=robot_pose.orientation, + ), ) - self.tf.publish(TFMessage(tracked_object_tf)) + self.tf.publish(TFMessage(transforms=[tracked_object_tf])) return detection3darray diff --git a/dimos/perception/experimental/spatial_perception.py b/dimos/perception/experimental/spatial_perception.py index 54f15cbd7b4..9e8110395f1 100644 --- a/dimos/perception/experimental/spatial_perception.py +++ b/dimos/perception/experimental/spatial_perception.py @@ -19,9 +19,12 @@ from datetime import datetime import os import time -from typing import TYPE_CHECKING, Any +from typing import Any import uuid +from dimos_generated.geometry_msgs.msg import Point +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.tf2_msgs.msg import TFMessage import numpy as np from reactivex import Observable, interval, operators as ops from reactivex.disposable import Disposable @@ -30,17 +33,14 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.geometry import pose_from_transform, quaternion_euler +from dimos.msgs.image import image_to_bgr from dimos.perception.experimental.image_embedding import ImageEmbeddingProvider from dimos.perception.experimental.spatial_vector_db import SpatialVectorDB from dimos.perception.experimental.visual_memory import VisualMemory from dimos.types.robot_location import RobotLocation from dimos.utils.logging_config import setup_logger -if TYPE_CHECKING: - from dimos.msgs.geometry_msgs.Vector3 import Vector3 - _OUTPUT_DIR = DIMOS_PROJECT_ROOT / "assets" / "output" _MEMORY_DIR = _OUTPUT_DIR / "memory" _SPATIAL_MEMORY_DIR = _MEMORY_DIR / "spatial_memory" @@ -154,7 +154,7 @@ def __init__(self, **kwargs: Any) -> None: embedding_provider=self.embedding_provider, ) - self.last_position: Vector3 | None = None + self.last_position: Point | None = None self.last_record_time: float | None = None self.frame_count: int = 0 @@ -169,19 +169,11 @@ def __init__(self, **kwargs: Any) -> None: @rpc def start(self) -> None: - import cv2 - super().start() # Subscribe to LCM streams def set_video(image_msg: Image) -> None: - # Convert Image message to numpy array - if hasattr(image_msg, "data"): - frame = image_msg.data - frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) - self._latest_video_frame = frame - else: - logger.warning("Received image message without data attribute") + self._latest_video_frame = image_to_bgr(image_msg) self.register_disposable(Disposable(self.color_image.subscribe(set_video))) @@ -211,7 +203,7 @@ def _process_frame(self) -> None: return # Create Pose object with position and orientation - current_pose = tf.to_pose() + current_pose = pose_from_transform(tf).pose # Process the frame directly try: @@ -242,16 +234,16 @@ def _process_frame(self) -> None: frame_id = f"frame_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}" # Get euler angles from quaternion orientation for metadata - euler = tf.rotation.to_euler() + roll, pitch, yaw = quaternion_euler(tf.transform.rotation) # Create metadata dictionary with primitive types only metadata = { "pos_x": float(current_pose.position.x), "pos_y": float(current_pose.position.y), "pos_z": float(current_pose.position.z), - "rot_x": float(euler.x), - "rot_y": float(euler.y), - "rot_z": float(euler.z), + "rot_x": roll, + "rot_y": pitch, + "rot_z": yaw, "timestamp": current_time, "frame_id": frame_id, } @@ -271,7 +263,7 @@ def _process_frame(self) -> None: logger.info( f"Stored frame at position ({current_pose.position.x:.2f}, {current_pose.position.y:.2f}, {current_pose.position.z:.2f}), " - f"rotation ({euler.x:.2f}, {euler.y:.2f}, {euler.z:.2f}) " + f"rotation ({roll:.2f}, {pitch:.2f}, {yaw:.2f}) " f"stored {self.stored_frame_count}/{self.frame_count} frames" ) @@ -374,7 +366,7 @@ def process_combined_data(data): # type: ignore[no-untyped-def] vector_id=frame_id, image=frame, embedding=frame_embedding, metadata=metadata ) - self.last_position = position_v3 + self.last_position = Point(x=position_v3.x, y=position_v3.y, z=position_v3.z) self.last_record_time = current_time self.stored_frame_count += 1 @@ -477,11 +469,12 @@ def add_named_location( return False # Create RobotLocation object - location = RobotLocation( # type: ignore[call-arg] + translation = tf.transform.translation + location = RobotLocation( name=name, - position=tf.translation, - rotation=tf.rotation.to_euler(), - description=description or f"Location: {name}", + position=(translation.x, translation.y, translation.z), + rotation=quaternion_euler(tf.transform.rotation), + metadata={"description": description or f"Location: {name}"}, timestamp=time.time(), ) diff --git a/dimos/perception/experimental/temporal_memory/clip_filter.py b/dimos/perception/experimental/temporal_memory/clip_filter.py index cb89c31c55d..19b2a142113 100644 --- a/dimos/perception/experimental/temporal_memory/clip_filter.py +++ b/dimos/perception/experimental/temporal_memory/clip_filter.py @@ -16,9 +16,10 @@ from typing import Any +from dimos_generated.sensor_msgs.msg import Image import numpy as np -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_view from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -36,9 +37,7 @@ def _get_image_data(image: Image) -> np.ndarray[Any, Any]: """Extract numpy array from Image.""" - if not hasattr(image, "data"): - raise AttributeError(f"Image missing .data attribute: {type(image)}") - return image.data + return image_view(image) if CLIP_AVAILABLE: diff --git a/dimos/perception/experimental/temporal_memory/entity_graph_db.py b/dimos/perception/experimental/temporal_memory/entity_graph_db.py index 54f188034a4..3c7d929f75e 100644 --- a/dimos/perception/experimental/temporal_memory/entity_graph_db.py +++ b/dimos/perception/experimental/temporal_memory/entity_graph_db.py @@ -37,8 +37,9 @@ from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: + from dimos_generated.sensor_msgs.msg import Image + from dimos.models.vl.base import VlModel - from dimos.msgs.sensor_msgs.Image import Image logger = setup_logger() diff --git a/dimos/perception/experimental/temporal_memory/frame_window_accumulator.py b/dimos/perception/experimental/temporal_memory/frame_window_accumulator.py index 4c910a1b88a..674c23caaee 100644 --- a/dimos/perception/experimental/temporal_memory/frame_window_accumulator.py +++ b/dimos/perception/experimental/temporal_memory/frame_window_accumulator.py @@ -24,8 +24,10 @@ import threading from typing import TYPE_CHECKING +from dimos.msgs.time import to_seconds + if TYPE_CHECKING: - from dimos.msgs.sensor_msgs.Image import Image + from dimos_generated.sensor_msgs.msg import Image @dataclass @@ -87,10 +89,7 @@ def add_frame(self, image: Image, wall_time: float) -> None: with self._lock: if self._video_start_wall_time is None: return - if image.ts is not None: - timestamp_s = image.ts - self._video_start_wall_time - else: - timestamp_s = wall_time - self._video_start_wall_time + timestamp_s = to_seconds(image.header.stamp) - self._video_start_wall_time frame = Frame( frame_index=self._frame_count, timestamp_s=timestamp_s, diff --git a/dimos/perception/experimental/temporal_memory/temporal_utils/helpers.py b/dimos/perception/experimental/temporal_memory/temporal_utils/helpers.py index 683c88376de..010845850bd 100644 --- a/dimos/perception/experimental/temporal_memory/temporal_utils/helpers.py +++ b/dimos/perception/experimental/temporal_memory/temporal_utils/helpers.py @@ -18,6 +18,8 @@ import numpy as np +from dimos.msgs.image import image_to_rgb + if TYPE_CHECKING: from dimos.perception.experimental.temporal_memory.frame_window_accumulator import Frame @@ -68,7 +70,9 @@ def is_scene_stale(frames: list["Frame"], stale_threshold: float = 5.0) -> bool: last_img = frames[-1].image if first_img is None or last_img is None: return False - if not hasattr(first_img, "data") or not hasattr(last_img, "data"): + first = image_to_rgb(first_img) + last = image_to_rgb(last_img) + if first.shape != last.shape: return False - diff = np.abs(first_img.data.astype(float) - last_img.data.astype(float)) + diff = np.abs(first.astype(float) - last.astype(float)) return bool(diff.mean() < stale_threshold) diff --git a/dimos/perception/experimental/test_generated_tracker.py b/dimos/perception/experimental/test_generated_tracker.py new file mode 100644 index 00000000000..3d8895afe80 --- /dev/null +++ b/dimos/perception/experimental/test_generated_tracker.py @@ -0,0 +1,105 @@ +# 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. + +from unittest.mock import MagicMock + +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage +from dimos_generated.vision_msgs.msg import ( + BoundingBox2D, + Detection2D, + Detection2DArray, + Detection3DArray, + Point2D, + Pose2D, +) +import numpy as np +import pytest + +from dimos.msgs.image import image_view +from dimos.msgs.time import time_from_seconds +from dimos.perception.experimental.object_tracker_2d import ObjectTracker2D +from dimos.perception.experimental.object_tracker_3d import ObjectTracker3D + + +@pytest.mark.parametrize("succeeded", [True, False]) +def test_tracker_publishes_generated_detection_and_clears_failed_track(succeeded: bool) -> None: + module = ObjectTracker2D(frame_id="camera") + module.tracker = MagicMock() + module.tracker.update.return_value = (succeeded, (10, 12, 20, 16)) + module.tracking_initialized = True + module._latest_rgb_frame = np.zeros((64, 64, 3), dtype=np.uint8) + module.detection2darray = MagicMock() + module.tracked_overlay = MagicMock() + try: + module._process_tracking() + result = Detection2DArray.decode(module.detection2darray.publish.call_args.args[0].encode()) + assert result.header.frame_id == "camera" + if not succeeded: + assert not result.detections + assert not module.tracking_initialized + module.tracked_overlay.publish.assert_not_called() + return + (detection,) = result.detections + assert (detection.bbox.center.position.x, detection.bbox.center.position.y) == (20, 20) + assert (detection.bbox.size_x, detection.bbox.size_y) == (20, 16) + assert detection.results[0].hypothesis.class_id == "tracked_object" + overlay = Image.decode(module.tracked_overlay.publish.call_args.args[0].encode()) + assert overlay.header == result.header + assert image_view(overlay).shape == (64, 64, 3) + assert overlay.encoding == "rgb8" + finally: + module.stop() + + +@pytest.mark.parametrize("depth", [2.0, 0.0, float("nan")]) +def test_tracker3d_generated_projection_preserves_depth_filter_and_source_stamp( + depth: float, +) -> None: + module = ObjectTracker3D(frame_id="camera") + module.camera_intrinsics = [100.0, 100.0, 32.0, 32.0] + module._latest_depth_frame = np.full((64, 64), depth, dtype=np.float32) + module.tf = MagicMock() + module.detection2darray = MagicMock() + header = Header(stamp=time_from_seconds(12.25), frame_id="camera") + detection = Detection2DArray( + header=header, + detections=[ + Detection2D( + header=header, + bbox=BoundingBox2D( + center=Pose2D(position=Point2D(x=32, y=32)), size_x=20, size_y=16 + ), + ) + ], + ) + try: + output = module._create_detection3d_from_2d(detection) + if not np.isfinite(depth) or depth <= 0: + assert output is None + module.tf.publish.assert_not_called() + return + result = Detection3DArray.decode(output.encode()) + assert result.header == header + box = result.detections[0].bbox + assert (box.center.position.x, box.center.position.y, box.center.position.z) == (2, 0, 0) + assert np.allclose([box.size.x, box.size.y, box.size.z], [0.4, 0.32, 0.1]) + transform = TFMessage.decode(module.tf.publish.call_args.args[0].encode()).transforms[0] + assert transform.header == header + assert transform.child_frame_id == "tracked_object" + assert transform.transform.translation.x == 2.0 + assert transform.transform.rotation.w == 1.0 + finally: + module.stop() diff --git a/dimos/perception/experimental/test_object_generated.py b/dimos/perception/experimental/test_object_generated.py new file mode 100644 index 00000000000..be288947c6d --- /dev/null +++ b/dimos/perception/experimental/test_object_generated.py @@ -0,0 +1,96 @@ +# 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. + +"""Object geometry and deduplication use generated ROS values.""" + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import PoseStamped, Vector3 +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.vision_msgs.msg import Detection3D +import numpy as np + +from dimos.msgs.image import image_from_array +from dimos.msgs.pointcloud import pointcloud_from_xyz_rgb, pointcloud_rgb, pointcloud_xyz +from dimos.perception.experimental.object import Object, aggregate_pointclouds +from dimos.perception.experimental.objectDB import ObjectDB + + +def _object(object_id, x, timestamp): + return Object( + object_id=object_id, + name="cup", + center=Vector3(x=x), + size=Vector3(x=0.1, y=0.2, z=0.3), + pose=PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=1700000000, nanosec=123456789)) + ), + pointcloud=PointCloud2(), + bbox=(0.0, 0.0, 1.0, 1.0), + track_id=-1, + class_id=0, + confidence=0.9, + ts=timestamp, + frame_id="world", + image=Image(), + ) + + +def test_detection_cdr_preserves_source_pose_header_and_current_center(): + obj = _object("cup-1", 0.2, 1.0) + obj.set_center(Vector3(x=0.4, y=0.5, z=0.6)) + decoded = Detection3D.decode(obj.to_detection3d_msg().encode()) + assert decoded.header == obj.pose.header + assert decoded.bbox.center == obj.pose.pose + assert decoded.bbox.size == obj.size + assert decoded.results[0].hypothesis.class_id == "cup" + assert decoded.results[0].hypothesis.score == 0.9 + decoded.bbox.center.position.x = 99 + assert obj.center.x == 0.4 + assert obj.pose.pose.position.x == 0.4 + + +def test_spatial_deduplication_and_nearest_lookup_use_generated_centers(): + database = ObjectDB(distance_threshold=0.2, min_detections_for_permanent=1) + first = _object("first", 0.0, 1.0) + nearby = _object("nearby", 0.1, 2.0) + distant = _object("distant", 2.0, 3.0) + assert database.add_objects([first]) == [first] + assert database.add_objects([nearby]) == [first] + assert first.center.x == 0.1 + assert database.add_objects([distant]) == [distant] + assert database.find_nearest(Vector3(x=1.8)) is distant + assert database.find_nearest(Vector3(x=0.2)) is first + + +def test_generated_object_cloud_aggregation_and_dictionary_preserve_values(): + obj = _object("abc", 0.2, 1.0) + stamp = Time(sec=1700000000, nanosec=987654321) + obj.pointcloud = pointcloud_from_xyz_rgb( + np.array([[1.0, 2.0, 3.0]], dtype=np.float32), + np.array([[255, 128, 0]], dtype=np.uint8), + header=Header(frame_id="world", stamp=stamp), + ) + pixels = np.zeros((2, 3, 3), dtype=np.uint8) + obj.image = image_from_array(pixels, encoding="rgb8", header=obj.pointcloud.header) + exported = obj.to_dict() + np.testing.assert_array_equal(exported["pointcloud"][0], [[1.0, 2.0, 3.0]]) + np.testing.assert_allclose(exported["pointcloud"][1], [[1.0, 128 / 255.0, 0.0]]) + np.testing.assert_array_equal(exported["image"], pixels) + merged = aggregate_pointclouds([obj]) + decoded = PointCloud2.decode(merged.encode()) + assert decoded.header.stamp == stamp and decoded.header.frame_id == "world" + np.testing.assert_array_equal(pointcloud_xyz(decoded), [[1.0, 2.0, 3.0]]) + assert pointcloud_rgb(decoded) is not None + assert aggregate_pointclouds([]).width == 0 diff --git a/dimos/perception/experimental/test_object_scene_registration_temporal.py b/dimos/perception/experimental/test_object_scene_registration_temporal.py index 82f14417f9c..4aa1bade0e8 100644 --- a/dimos/perception/experimental/test_object_scene_registration_temporal.py +++ b/dimos/perception/experimental/test_object_scene_registration_temporal.py @@ -20,11 +20,18 @@ from typing import Any from unittest.mock import ANY, MagicMock, call +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header +from dimos_generated.vision_msgs.msg import Detection3DArray import numpy as np import pytest -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.msgs.camera_info import camera_info_from_intrinsics +from dimos.msgs.image import image_from_array +from dimos.msgs.pointcloud import pointcloud_xyz +from dimos.msgs.time import time_from_seconds +from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.experimental.object_scene_registration import ObjectSceneRegistrationModule from dimos.perception.experimental.objectDB import ObjectDB @@ -44,18 +51,19 @@ def dispose(self) -> None: def _image(timestamp: float) -> Image: - return Image( - data=np.ones((2, 2), dtype=np.float32), - format=ImageFormat.DEPTH, - frame_id="camera", - ts=timestamp, + return image_from_array( + np.ones((2, 2), dtype=np.float32), + encoding="32FC1", + header=Header(frame_id="camera", stamp=time_from_seconds(timestamp)), ) @pytest.fixture def module() -> Iterator[ObjectSceneRegistrationModule]: module = ObjectSceneRegistrationModule(target_frame="map") - module._camera_info = MagicMock(K=[1.0, 0.0, 0.0, 0.0, 1.0, 0.0]) + module._camera_info = camera_info_from_intrinsics( + 1, 1, 0, 0, 2, 2, header=Header(frame_id="camera") + ) module._latest_scene_snapshot = None yield module module.stop() @@ -78,11 +86,10 @@ def get(self, *args: Any, **kwargs: Any) -> Any: module._detector = MagicMock() module._text_prompts = ["cup"] module.detections_2d = MagicMock() - color = Image( - data=np.zeros((2, 2, 3), dtype=np.uint8), - format=ImageFormat.BGR, - frame_id="camera", - ts=12.5, + color = image_from_array( + np.zeros((2, 2, 3), dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id="camera", stamp=time_from_seconds(12.5)), ) detections = ImageDetections2D(color, []) @@ -138,35 +145,30 @@ def test_full_scene_pointcloud_uses_one_coherent_scene_snapshot( monkeypatch: Any, module: ObjectSceneRegistrationModule ) -> None: depth = _image(3.0) - transform = MagicMock(name="transform") + transform = TransformStamped( + header=Header(frame_id="map"), + child_frame_id="camera", + transform=Transform(translation=Vector3(x=1), rotation=Quaternion(w=1)), + ) module._tf = _FakeTF(transform) # type: ignore[assignment] module._latest_scene_snapshot = (depth, transform) class _PointCloud: - points = list(range(100)) + points = np.zeros((100, 3), dtype=np.float64) def voxel_down_sample(self, voxel_size: float) -> _PointCloud: return self - pointcloud = _PointCloud() fake_o3d = MagicMock() - fake_o3d.camera.PinholeCameraIntrinsic.return_value = MagicMock() - fake_o3d.geometry.Image.return_value = MagicMock() - fake_o3d.geometry.PointCloud.create_from_depth_image.return_value = pointcloud - # open3d is imported inside the method under test, so swap the module itself + fake_o3d.geometry.PointCloud.create_from_depth_image.return_value = _PointCloud() + # Only the native depth projection is mocked; generated output and transformation are real. monkeypatch.setitem(sys.modules, "open3d", fake_o3d) - - result = MagicMock() - result.transform.side_effect = lambda used_transform: ( - result if used_transform is transform else pytest.fail("mixed scene snapshot") - ) - monkeypatch.setattr( - "dimos.perception.experimental.object_scene_registration.PointCloud2", - lambda *_args, **_kwargs: result, - ) - - module.get_full_scene_pointcloud() - result.transform.assert_called_once_with(transform) + result = module.get_full_scene_pointcloud() + assert result is not None + assert result.header.frame_id == "map" + assert result.header.stamp == depth.header.stamp + np.testing.assert_array_equal(pointcloud_xyz(result), np.tile([1, 0, 0], (100, 1))) + assert module._tf.calls == [] # The cached coherent snapshot supplied the transform. def test_owlv2_queries_configured_prompts(monkeypatch: Any) -> None: @@ -177,11 +179,10 @@ def test_owlv2_queries_configured_prompts(monkeypatch: Any) -> None: module._detector = MagicMock() module._text_prompts = ["mug"] module.detections_2d = MagicMock() - color = Image( - data=np.zeros((2, 2, 3), dtype=np.uint8), - format=ImageFormat.BGR, - frame_id="camera", - ts=4.0, + color = image_from_array( + np.zeros((2, 2, 3), dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id="camera", stamp=time_from_seconds(4.0)), ) detections = ImageDetections2D(color, []) module._detector.query_detections.return_value = detections @@ -219,27 +220,30 @@ def test_owlv2_yolo_constructs_box_prompt_segmenter(mocker: Any) -> None: def test_moondream_queries_each_configured_prompt(monkeypatch: Any) -> None: module = ObjectSceneRegistrationModule(target_frame="camera", detector_backend="moondream") module._camera_info = MagicMock() - module._detector = MagicMock() + detector = MagicMock() + module._detector = detector module._text_prompts = ["cup", "bottle"] module.detections_2d = MagicMock() - color = Image( - data=np.zeros((2, 2, 3), dtype=np.uint8), - format=ImageFormat.BGR, - frame_id="camera", - ts=4.0, + color = image_from_array( + np.zeros((2, 2, 3), dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id="camera", stamp=time_from_seconds(4.0)), ) - cup = MagicMock(track_id=0, class_id=-1) - bottle = MagicMock(track_id=0, class_id=-1) - module._detector.query_detections.side_effect = [ + cup = Detection2DBBox((0, 0, 1, 1), 0, -1, 0.9, "cup", 4.0, color) + bottle = Detection2DBBox((1, 1, 2, 2), 0, -1, 0.8, "bottle", 4.0, color) + detector.query_detections.side_effect = [ ImageDetections2D(color, [cup]), ImageDetections2D(color, [bottle]), ] process_3d = MagicMock() monkeypatch.setattr(module, "_process_3d_detections", process_3d) - module._process_images(color, _image(4.0)) + try: + module._process_images(color, _image(4.0)) + finally: + module.stop() - assert module._detector.query_detections.call_args_list == [ + assert detector.query_detections.call_args_list == [ call(color, "cup"), call(color, "bottle"), ] @@ -247,17 +251,15 @@ def test_moondream_queries_each_configured_prompt(monkeypatch: Any) -> None: assert combined.detections == [cup, bottle] assert (cup.track_id, cup.class_id) == (-1, 0) assert (bottle.track_id, bottle.class_id) == (-1, 1) - module.stop() def test_edgetam_refines_detector_output(monkeypatch: Any) -> None: module = ObjectSceneRegistrationModule(target_frame="camera", segmentation_backend="edgetam") module._camera_info = MagicMock() - color = Image( - data=np.zeros((2, 2, 3), dtype=np.uint8), - format=ImageFormat.BGR, - frame_id="camera", - ts=4.0, + color = image_from_array( + np.zeros((2, 2, 3), dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id="camera", stamp=time_from_seconds(4.0)), ) raw_detections = ImageDetections2D(color, []) segmented_detections = ImageDetections2D(color, []) @@ -281,11 +283,10 @@ def test_request_driven_scan_processes_latest_aligned_frame(monkeypatch: Any) -> target_frame="camera", detector_backend="owlv2", detect_on_request=True ) module._object_db = MagicMock() - color = Image( - data=np.zeros((2, 2, 3), dtype=np.uint8), - format=ImageFormat.BGR, - frame_id="camera", - ts=4.0, + color = image_from_array( + np.zeros((2, 2, 3), dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id="camera", stamp=time_from_seconds(4.0)), ) depth = _image(4.0) module._on_aligned_frames((color, depth)) @@ -378,11 +379,10 @@ def test_concurrent_request_scans_do_not_overlap(monkeypatch: Any) -> None: target_frame="camera", detector_backend="owlv2", detect_on_request=True ) module._object_db = MagicMock() - color = Image( - data=np.zeros((2, 2, 3), dtype=np.uint8), - format=ImageFormat.BGR, - frame_id="camera", - ts=4.0, + color = image_from_array( + np.zeros((2, 2, 3), dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id="camera", stamp=time_from_seconds(4.0)), ) module._on_aligned_frames((color, _image(4.0))) first_started = Event() @@ -454,10 +454,9 @@ def test_object_db_counts_each_source_frame_once(monkeypatch: Any) -> None: last_seen_ts=None, detections_count=1, ) - first.center = MagicMock() + first.center = Vector3() duplicate = MagicMock(object_id="duplicate-id", track_id=-1, ts=4.0) - duplicate.center = MagicMock() - duplicate.center.distance.return_value = 0.0 + duplicate.center = Vector3() observed = object_db.add_objects([first, duplicate]) @@ -466,8 +465,7 @@ def test_object_db_counts_each_source_frame_once(monkeypatch: Any) -> None: assert first.last_seen_ts == 1000.0 newer = MagicMock(object_id="newer-id", track_id=-1, ts=5.0) - newer.center = MagicMock() - newer.center.distance.return_value = 0.0 + newer.center = Vector3() first.update_object.side_effect = lambda _: setattr(first, "detections_count", 2) now[0] = 1001.0 diff --git a/dimos/perception/fiducial/blueprints/desk_marker_tf.py b/dimos/perception/fiducial/blueprints/desk_marker_tf.py index 0a63f90ca57..2367589457c 100644 --- a/dimos/perception/fiducial/blueprints/desk_marker_tf.py +++ b/dimos/perception/fiducial/blueprints/desk_marker_tf.py @@ -20,6 +20,12 @@ import threading import time +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage +from dimos_generated.vision_msgs.msg import Detection3DArray + from dimos.core.coordination.blueprints import autoconnect from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig @@ -27,12 +33,9 @@ from dimos.core.transport import LCMTransport from dimos.hardware.sensors.camera.module import CameraModule from dimos.hardware.sensors.camera.webcam import Webcam -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.msgs.camera_info import camera_info_from_yaml +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.time import time_from_nanoseconds from dimos.perception.fiducial.marker_detection_stream_module import MarkerDetectionStreamModule from dimos.perception.fiducial.marker_tf_module import MarkerTfModule @@ -61,8 +64,9 @@ def create_desk_webcam( def create_desk_camera_info( camera_info_yaml: str | Path = DEFAULT_DESK_CAMERA_INFO_YAML, ) -> CameraInfo: - camera_info = CameraInfo.from_yaml(str(camera_info_yaml)) - camera_info.frame_id = DESK_CAMERA_FRAME_ID + camera_info = camera_info_from_yaml( + camera_info_yaml, header=Header(frame_id=DESK_CAMERA_FRAME_ID) + ) return camera_info @@ -124,28 +128,27 @@ def stop(self) -> None: super().stop() def publish_static_chain(self) -> None: - ts = time.time() - self._last_publish_ts = ts + stamp = time_from_nanoseconds(time.time_ns()) + self._last_publish_ts = stamp.sec + stamp.nanosec / 1e9 roll, pitch, yaw = self.config.camera_rotation_rpy_rad x, y, z = self.config.camera_translation_m - self.tf.publish( TFMessage( - Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id=self.config.world_frame, - child_frame_id=self.config.base_frame, - ts=ts, - ), - Transform( - # Default desk camera pose: about 25 cm forward and 15 cm above base_link. - translation=Vector3(x, y, z), - rotation=Quaternion.from_euler(Vector3(roll, pitch, yaw)), - frame_id=self.config.base_frame, - child_frame_id=self.config.camera_optical_frame, - ts=ts, - ), + transforms=[ + TransformStamped( + header=Header(stamp=stamp, frame_id=self.config.world_frame), + child_frame_id=self.config.base_frame, + transform=Transform(rotation=Quaternion(w=1)), + ), + TransformStamped( + header=Header(stamp=stamp, frame_id=self.config.base_frame), + child_frame_id=self.config.camera_optical_frame, + transform=Transform( + translation=Vector3(x=x, y=y, z=z), + rotation=quaternion_from_euler(roll, pitch, yaw), + ), + ), + ] ) ) diff --git a/dimos/perception/fiducial/blueprints/test_desk_marker_tf.py b/dimos/perception/fiducial/blueprints/test_desk_marker_tf.py index fc7e4ecf639..2d820d94fa4 100644 --- a/dimos/perception/fiducial/blueprints/test_desk_marker_tf.py +++ b/dimos/perception/fiducial/blueprints/test_desk_marker_tf.py @@ -41,7 +41,9 @@ def test_desk_marker_tf_blueprint_declares_static_tf_module() -> None: assert desk_marker_tf.blueprints[2].module is MarkerDetectionStreamModule assert desk_marker_tf.blueprints[2].kwargs["marker_length_m"] == DESK_MARKER_LENGTH_M assert desk_marker_tf.blueprints[2].kwargs["aruco_dictionary"] == DESK_MARKER_ARUCO_DICTIONARY - assert desk_marker_tf.blueprints[2].kwargs["camera_info"].frame_id == DESK_CAMERA_FRAME_ID + assert ( + desk_marker_tf.blueprints[2].kwargs["camera_info"].header.frame_id == DESK_CAMERA_FRAME_ID + ) assert desk_marker_tf.blueprints[3].module is MarkerTfModule assert ( desk_marker_tf.blueprints[3].kwargs["marker_namespace_prefix"] @@ -87,10 +89,10 @@ def test_create_desk_webcam_loads_camera_info_yaml(tmp_path: Path) -> None: assert camera.config.width == 1920 assert camera.config.height == 1080 assert camera.config.fps == 7.5 - assert camera.config.camera_info.frame_id == DESK_CAMERA_FRAME_ID + assert camera.config.camera_info.header.frame_id == DESK_CAMERA_FRAME_ID camera_info = create_desk_camera_info(camera_info_yaml) - assert camera_info.frame_id == DESK_CAMERA_FRAME_ID + assert camera_info.header.frame_id == DESK_CAMERA_FRAME_ID assert camera_info.width == 1920 @@ -106,10 +108,11 @@ def test_desk_static_tf_module_publishes_world_to_camera_optical_chain() -> None world_camera = view.get("world", "camera_optical", mod._last_publish_ts, 1.0) assert world_camera is not None - assert world_camera.frame_id == "world" + assert world_camera.header.frame_id == "world" assert world_camera.child_frame_id == "camera_optical" - assert world_camera.translation.x == 0.3 - assert world_camera.translation.y == 0.0 - assert world_camera.translation.z == 0.2 + assert world_camera.transform.translation.x == 0.3 + assert world_camera.transform.translation.y == 0.0 + assert world_camera.transform.translation.z == 0.2 finally: + view.dispose() mod.stop() diff --git a/dimos/perception/fiducial/marker_detect.py b/dimos/perception/fiducial/marker_detect.py index 9abb07ff334..c002a5ba4fe 100644 --- a/dimos/perception/fiducial/marker_detect.py +++ b/dimos/perception/fiducial/marker_detect.py @@ -18,12 +18,14 @@ from typing import Any +from dimos_generated.geometry_msgs.msg import TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.std_msgs.msg import Header import numpy as np -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.geometry import compose_transforms +from dimos.msgs.image import image_to_bgr +from dimos.msgs.time import to_seconds from dimos.perception.detection.type.detection3d.marker import Detection3DMarker from dimos.perception.fiducial.marker_pose import ( camera_info_to_cv_matrices, @@ -40,16 +42,18 @@ def detect_markers_in_image( image: Image, *, camera_info: CameraInfo, - world_T_optical: Transform, + world_T_optical: TransformStamped, marker_length_m: float, aruco_dictionary: str, world_frame: str = "world", detect_inverted: bool = False, detector: Any | None = None, - camera_matrix: np.ndarray | None = None, - dist_coeffs: np.ndarray | None = None, + camera_matrix: np.ndarray[Any, np.dtype[Any]] | None = None, + dist_coeffs: np.ndarray[Any, np.dtype[Any]] | None = None, ) -> list[Detection3DMarker]: """Detect markers in one image and return rich world-frame 3D detections.""" + import cv2 + if marker_length_m <= 0: raise ValueError(f"marker_length_m must be > 0, got {marker_length_m}") if ( @@ -66,20 +70,18 @@ def detect_markers_in_image( if camera_matrix is None or dist_coeffs is None: camera_matrix, dist_coeffs = camera_info_to_cv_matrices(camera_info) - gray = image.to_grayscale().as_numpy() + gray = cv2.cvtColor(image_to_bgr(image), cv2.COLOR_BGR2GRAY) corners, ids, _ = detector.detectMarkers(gray) if ids is None or len(ids) == 0: return [] optical_frame = camera_optical_frame_id(image, camera_info) - t_world_optical = Transform( - translation=world_T_optical.translation, - rotation=world_T_optical.rotation, - frame_id=world_frame, + t_world_optical = TransformStamped( + header=Header(stamp=image.header.stamp, frame_id=world_frame), child_frame_id=optical_frame, - ts=image.ts, + transform=world_T_optical.transform, ) - marker_size = Vector3(marker_length_m, marker_length_m, 0.0) + marker_size = Vector3(x=marker_length_m, y=marker_length_m) detections: list[Detection3DMarker] = [] for corner_set, mid_arr in zip(corners, ids, strict=True): @@ -98,11 +100,10 @@ def detect_markers_in_image( t_optical_marker = rvec_tvec_to_transform( rvec, tvec, - frame_id=optical_frame, + header=Header(stamp=image.header.stamp, frame_id=optical_frame), child_frame_id=f"marker_{mid}", - ts=image.ts, ) - t_world_marker = t_world_optical + t_optical_marker + t_world_marker = compose_transforms(t_world_optical, t_optical_marker) corners_2d = corner_set.reshape(4, 2).astype(np.float32) bbox = marker_corners_to_bbox(corners_2d) @@ -123,13 +124,13 @@ def detect_markers_in_image( class_id=mid, confidence=1.0, name="", - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, - center=t_world_marker.translation, + center=t_world_marker.transform.translation, size=marker_size, transform=t_world_optical, frame_id=world_frame, - orientation=t_world_marker.rotation, + orientation=t_world_marker.transform.rotation, marker_id=mid, corners_px=corners_2d, dictionary=aruco_dictionary, diff --git a/dimos/perception/fiducial/marker_detection_stream_module.py b/dimos/perception/fiducial/marker_detection_stream_module.py index e200326c75a..2158daa5c11 100644 --- a/dimos/perception/fiducial/marker_detection_stream_module.py +++ b/dimos/perception/fiducial/marker_detection_stream_module.py @@ -21,9 +21,11 @@ from __future__ import annotations -import time from typing import Any, cast +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.tf2_msgs.msg import TFMessage +from dimos_generated.vision_msgs.msg import Detection3DArray from pydantic import Field from reactivex.disposable import Disposable @@ -34,10 +36,8 @@ from dimos.memory.store.null import NullStore from dimos.memory.stream import Stream from dimos.memory.transform import QualityWindow, SpeedLimit -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.msgs.image import image_sharpness +from dimos.msgs.time import to_seconds from dimos.perception.detection.type.detection3d.marker import Detection3DMarker from dimos.perception.fiducial.marker_pose import camera_optical_frame_id, is_fisheye_model from dimos.perception.fiducial.marker_transformer import DetectMarkers, MarkersPerFrame @@ -78,7 +78,7 @@ def __init__(self, **kwargs: Any) -> None: def pipeline(self, stream: Stream[Image]) -> Stream[Detection3DArray]: result: Stream[Any] = stream.transform( - QualityWindow(lambda img: img.sharpness, window=self.config.quality_window_s) + QualityWindow(image_sharpness, window=self.config.quality_window_s) ) if self.config.speed_limit_enabled: result = result.transform( @@ -121,7 +121,7 @@ def _append_image_with_pose(self, stream: Stream[Image], image: Image) -> None: logger.debug("MarkerDetectionStreamModule: no CameraInfo yet; skipping frame") return - ts = getattr(image, "ts", None) or time.time() + ts = to_seconds(image.header.stamp) optical = camera_optical_frame_id(image, info) t_world_optical = self.tfbuffer.get( self.config.world_frame, @@ -142,13 +142,13 @@ def _append_image_with_pose(self, stream: Stream[Image], image: Image) -> None: image, ts=ts, pose=( - t_world_optical.translation.x, - t_world_optical.translation.y, - t_world_optical.translation.z, - t_world_optical.rotation.x, - t_world_optical.rotation.y, - t_world_optical.rotation.z, - t_world_optical.rotation.w, + t_world_optical.transform.translation.x, + t_world_optical.transform.translation.y, + t_world_optical.transform.translation.z, + t_world_optical.transform.rotation.x, + t_world_optical.transform.rotation.y, + t_world_optical.transform.rotation.z, + t_world_optical.transform.rotation.w, ), ) diff --git a/dimos/perception/fiducial/marker_pose.py b/dimos/perception/fiducial/marker_pose.py index bfe7f115e18..6e7c9b838ed 100644 --- a/dimos/perception/fiducial/marker_pose.py +++ b/dimos/perception/fiducial/marker_pose.py @@ -16,15 +16,14 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any +from dimos_generated.geometry_msgs.msg import Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.std_msgs.msg import Header import numpy as np -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.geometry import quaternion_from_matrix if TYPE_CHECKING: import cv2 @@ -38,26 +37,28 @@ def is_fisheye_model(distortion_model: str | None) -> bool: return (distortion_model or "").strip().lower() in _FISHEYE_MODELS -def camera_info_to_cv_matrices(camera_info: CameraInfo) -> tuple[np.ndarray, np.ndarray]: +def camera_info_to_cv_matrices( + camera_info: CameraInfo, +) -> tuple[np.ndarray[Any, np.dtype[Any]], np.ndarray[Any, np.dtype[Any]]]: """Build OpenCV ``cameraMatrix`` and ``distCoeffs`` from ``CameraInfo``.""" - k = np.array(camera_info.K, dtype=np.float64).reshape(3, 3) - d = np.array(camera_info.D if camera_info.D else [], dtype=np.float64).reshape(-1, 1) + k = np.array(camera_info.k, dtype=np.float64).reshape(3, 3) + d = np.array(camera_info.d if camera_info.d else [], dtype=np.float64).reshape(-1, 1) return k, d def camera_optical_frame_id(image: Image, camera_info: CameraInfo) -> str: """Frame in which image pixels and intrinsics apply (optical convention in ROS). - Prefer ``Image.frame_id`` so TF lookups match the stream that produced the - pixels. Fall back to ``CameraInfo.frame_id``, then a conventional default. + Prefer ``Image.header.frame_id`` so TF lookups match the stream that produced the + pixels. Fall back to ``CameraInfo.header.frame_id``, then a conventional default. """ - for fid in (image.frame_id, camera_info.frame_id): + for fid in (image.header.frame_id, camera_info.header.frame_id): if fid and fid.strip(): return fid.strip() return "camera_optical" -def _aruco_marker_object_points(marker_length_m: float) -> np.ndarray: +def _aruco_marker_object_points(marker_length_m: float) -> np.ndarray[Any, np.dtype[Any]]: """Corner order matches OpenCV ArUco / solvePnP convention (planar square, Z=0).""" h = marker_length_m / 2.0 return np.array( @@ -72,13 +73,13 @@ def _aruco_marker_object_points(marker_length_m: float) -> np.ndarray: def estimate_marker_pose( - corners_px: np.ndarray, + corners_px: np.ndarray[Any, np.dtype[Any]], marker_length_m: float, - camera_matrix: np.ndarray, - dist_coeffs: np.ndarray, + camera_matrix: np.ndarray[Any, np.dtype[Any]], + dist_coeffs: np.ndarray[Any, np.dtype[Any]], *, distortion_model: str | None = None, -) -> tuple[np.ndarray, np.ndarray] | None: +) -> tuple[np.ndarray[Any, np.dtype[Any]], np.ndarray[Any, np.dtype[Any]]] | None: """Return ``(rvec, tvec)`` for camera optical <- marker from undistorted solvePnP. For fisheye/equidistant intrinsics, corners are first undistorted into the @@ -88,17 +89,17 @@ def estimate_marker_pose( import cv2 obj = _aruco_marker_object_points(marker_length_m) - img: np.ndarray = corners_px.reshape(4, 1, 2).astype(np.float32) + img: np.ndarray[Any, np.dtype[Any]] = corners_px.reshape(4, 1, 2).astype(np.float32) if is_fisheye_model(distortion_model): d_flat = np.asarray(dist_coeffs, dtype=np.float64).reshape(-1) if d_flat.size < 4: raise ValueError( f"Fisheye/equidistant distortion model requires at least 4 coefficients; " - f"got {d_flat.size}. Check CameraInfo.D." + f"got {d_flat.size}. Check CameraInfo.d." ) d_fisheye = d_flat[:4].reshape(4, 1) img = cv2.fisheye.undistortPoints(img, camera_matrix, d_fisheye, P=camera_matrix) - solve_dist: np.ndarray = np.zeros((0, 1), dtype=np.float64) + solve_dist: np.ndarray[Any, np.dtype[Any]] = np.zeros((0, 1), dtype=np.float64) else: solve_dist = dist_coeffs ok, rvec, tvec = cv2.solvePnP( @@ -114,25 +115,24 @@ def estimate_marker_pose( def rvec_tvec_to_transform( - rvec: np.ndarray, - tvec: np.ndarray, + rvec: np.ndarray[Any, np.dtype[Any]], + tvec: np.ndarray[Any, np.dtype[Any]], *, - frame_id: str, + header: Header, child_frame_id: str, - ts: float, -) -> Transform: +) -> TransformStamped: """Build ``Transform`` for ``frame_id`` <- ``child_frame_id`` (camera <- marker).""" import cv2 rot_mat, _ = cv2.Rodrigues(rvec) - quat = Quaternion.from_rotation_matrix(rot_mat) + quat = quaternion_from_matrix(np.asarray(rot_mat, dtype=np.float64)) tx, ty, tz = tvec.reshape(3) - return Transform( - translation=Vector3(float(tx), float(ty), float(tz)), - rotation=quat, - frame_id=frame_id, + return TransformStamped( + header=header, child_frame_id=child_frame_id, - ts=ts, + transform=Transform( + translation=Vector3(x=float(tx), y=float(ty), z=float(tz)), rotation=quat + ), ) @@ -150,7 +150,9 @@ def create_aruco_detector( return cv2.aruco.ArucoDetector(dictionary, parameters) -def marker_corners_to_bbox(corners_px: np.ndarray) -> tuple[float, float, float, float]: +def marker_corners_to_bbox( + corners_px: np.ndarray[Any, np.dtype[Any]], +) -> tuple[float, float, float, float]: """Return the axis-aligned image bbox around a marker's four pixel corners.""" corners_2d = np.asarray(corners_px, dtype=np.float32).reshape(4, 2) xy_min = corners_2d.min(axis=0) @@ -159,12 +161,12 @@ def marker_corners_to_bbox(corners_px: np.ndarray) -> tuple[float, float, float, def marker_reprojection_error( - corners_px: np.ndarray, + corners_px: np.ndarray[Any, np.dtype[Any]], marker_length_m: float, - camera_matrix: np.ndarray, - dist_coeffs: np.ndarray, - rvec: np.ndarray, - tvec: np.ndarray, + camera_matrix: np.ndarray[Any, np.dtype[Any]], + dist_coeffs: np.ndarray[Any, np.dtype[Any]], + rvec: np.ndarray[Any, np.dtype[Any]], + tvec: np.ndarray[Any, np.dtype[Any]], *, distortion_model: str | None = None, ) -> float: @@ -175,7 +177,9 @@ def marker_reprojection_error( """ import cv2 - observed: np.ndarray = np.asarray(corners_px, dtype=np.float32).reshape(4, 1, 2) + observed: np.ndarray[Any, np.dtype[Any]] = np.asarray(corners_px, dtype=np.float32).reshape( + 4, 1, 2 + ) project_dist = dist_coeffs if is_fisheye_model(distortion_model): @@ -183,7 +187,7 @@ def marker_reprojection_error( if d_flat.size < 4: raise ValueError( f"Fisheye/equidistant distortion model requires at least 4 coefficients; " - f"got {d_flat.size}. Check CameraInfo.D." + f"got {d_flat.size}. Check CameraInfo.d." ) d_fisheye = d_flat[:4].reshape(4, 1) observed = cv2.fisheye.undistortPoints( diff --git a/dimos/perception/fiducial/marker_tf_module.py b/dimos/perception/fiducial/marker_tf_module.py index f5184ab9430..573245a21bc 100644 --- a/dimos/perception/fiducial/marker_tf_module.py +++ b/dimos/perception/fiducial/marker_tf_module.py @@ -40,17 +40,15 @@ from __future__ import annotations +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage +from dimos_generated.vision_msgs.msg import Detection3D, Detection3DArray from reactivex.disposable import Disposable from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.msgs.vision_msgs.Detection3D import Detection3D -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -83,22 +81,16 @@ def _marker_child_frame(self, marker_id: str) -> str: return f"{p}/{name}" if p else name def _process_detections(self, detections: Detection3DArray) -> None: - if detections.detections_length == 0: + if not detections.detections: return - - marker_detections = detections.detections[: detections.detections_length] - if not marker_detections: - return - + marker_detections = detections.detections markers_parent = self._markers_parent_frame() - ts = detections.ts - out: list[Transform] = [ - Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id=self.config.world_frame, + stamp = detections.header.stamp + out = [ + TransformStamped( + header=Header(stamp=stamp, frame_id=self.config.world_frame), child_frame_id=markers_parent, - ts=ts, + transform=Transform(rotation=Quaternion(w=1)), ) ] @@ -110,12 +102,15 @@ def _process_detections(self, detections: Detection3DArray) -> None: pose = detection.bbox.center out.append( - Transform( - translation=pose.position, - rotation=pose.orientation, - frame_id=markers_parent, + TransformStamped( + header=Header(stamp=stamp, frame_id=markers_parent), child_frame_id=self._marker_child_frame(marker_id), - ts=ts, + transform=Transform( + translation=Vector3( + x=pose.position.x, y=pose.position.y, z=pose.position.z + ), + rotation=pose.orientation, + ), ) ) @@ -123,7 +118,7 @@ def _process_detections(self, detections: Detection3DArray) -> None: # In that case, skip TF entirely rather than publishing only the # namespace anchor without any marker child frames. if len(out) > 1: - self.tf.publish(TFMessage(*out)) + self.tf.publish(TFMessage(transforms=out)) @staticmethod def _marker_id_from_detection(detection: Detection3D) -> str | None: @@ -131,7 +126,7 @@ def _marker_id_from_detection(detection: Detection3D) -> str | None: if marker_id: return marker_id - for result in detection.results[: detection.results_length]: + for result in detection.results: class_id = str(result.hypothesis.class_id).strip() if ":" not in class_id: continue diff --git a/dimos/perception/fiducial/marker_transformer.py b/dimos/perception/fiducial/marker_transformer.py index 0a0f7bfd036..a2b8a447644 100644 --- a/dimos/perception/fiducial/marker_transformer.py +++ b/dimos/perception/fiducial/marker_transformer.py @@ -32,17 +32,15 @@ from collections.abc import Callable import dataclasses import math -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, TypeAlias, cast +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.std_msgs.msg import Header +from dimos_generated.vision_msgs.msg import Detection3DArray import numpy as np from dimos.memory.transform import Transformer -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray from dimos.perception.detection.type.detection3d.imageDetections3D import ImageDetections3D from dimos.perception.detection.type.detection3d.marker import Detection3DMarker from dimos.perception.fiducial.marker_detect import ( @@ -64,7 +62,7 @@ logger = setup_logger() -CameraInfoSource = CameraInfo | Callable[[], CameraInfo | None] | None +CameraInfoSource: TypeAlias = CameraInfo | Callable[[], CameraInfo | None] | None def _camera_info_key(info: CameraInfo) -> tuple[Any, ...]: @@ -73,25 +71,24 @@ def _camera_info_key(info: CameraInfo) -> tuple[Any, ...]: info.width, info.height, info.distortion_model, - tuple(info.K), - tuple(info.D), + tuple(info.k), + tuple(info.d), ) def _pose_tuple_to_transform( pose: tuple[float, float, float, float, float, float, float], *, - frame_id: str, + header: Header, child_frame_id: str, - ts: float, -) -> Transform: +) -> TransformStamped: x, y, z, qx, qy, qz, qw = pose - return Transform( - translation=Vector3(x, y, z), - rotation=Quaternion(qx, qy, qz, qw), - frame_id=frame_id, + return TransformStamped( + header=header, child_frame_id=child_frame_id, - ts=ts, + transform=Transform( + translation=Vector3(x=x, y=y, z=z), rotation=Quaternion(x=qx, y=qy, z=qz, w=qw) + ), ) @@ -124,10 +121,10 @@ def _average_marker_pose( norm = math.sqrt(qsx * qsx + qsy * qsy + qsz * qsz + qsw * qsw) if norm < 1e-12: # Signs cancelled exactly — pathological, fall back to the hemisphere ref. - return (Vector3(cx, cy, cz), ref) + return (Vector3(x=cx, y=cy, z=cz), ref) return ( - Vector3(cx, cy, cz), - Quaternion(qsx / norm, qsy / norm, qsz / norm, qsw / norm), + Vector3(x=cx, y=cy, z=cz), + Quaternion(x=qsx / norm, y=qsy / norm, z=qsz / norm, w=qsw / norm), ) @@ -158,8 +155,8 @@ def __init__( self._detector = create_aruco_detector(aruco_dictionary, detect_inverted=detect_inverted) self._camera_info_key: tuple[Any, ...] | None = None self._resolved_camera_info: CameraInfo | None = None - self._cam_mtx: np.ndarray | None = None - self._dist: np.ndarray | None = None + self._cam_mtx: np.ndarray[Any, np.dtype[Any]] | None = None + self._dist: np.ndarray[Any, np.dtype[Any]] | None = None # Per marker_id sliding-window buffer of raw detections, used to emit # smoothed pose updates when ``smoothing_window > 0``. self._buffers: dict[int, TimestampedBufferCollection[Detection3DMarker]] = {} @@ -220,9 +217,8 @@ def __call__( optical_frame = camera_optical_frame_id(image, info) t_world_optical = _pose_tuple_to_transform( pose_tuple, - frame_id=self.world_frame, + header=Header(stamp=image.header.stamp, frame_id=self.world_frame), child_frame_id=optical_frame, - ts=obs.ts, ) detections = _detect_markers_in_image( @@ -264,12 +260,10 @@ def __call__( track_id = -1 det = dataclasses.replace(det, track_id=track_id) - yielded_pose = Transform( - translation=det.center, - rotation=det.orientation, - frame_id=self.world_frame, + yielded_pose = TransformStamped( + header=Header(stamp=image.header.stamp, frame_id=self.world_frame), child_frame_id=f"marker_{mid}", - ts=obs.ts, + transform=Transform(translation=det.center, rotation=det.orientation), ) yielded_det = det @@ -289,12 +283,10 @@ def __call__( yielded_det = dataclasses.replace( det, center=avg_center, orientation=avg_orient, transform=None ) - yielded_pose = Transform( - translation=avg_center, - rotation=avg_orient, - frame_id=self.world_frame, + yielded_pose = TransformStamped( + header=Header(stamp=image.header.stamp, frame_id=self.world_frame), child_frame_id=f"marker_{mid}", - ts=obs.ts, + transform=Transform(translation=avg_center, rotation=avg_orient), ) yield obs.derive(data=yielded_det, pose=yielded_pose).tag( @@ -401,12 +393,12 @@ def _source_pose( if detections and detections[0].transform is not None: transform = detections[0].transform return ( - transform.translation.x, - transform.translation.y, - transform.translation.z, - transform.rotation.x, - transform.rotation.y, - transform.rotation.z, - transform.rotation.w, + transform.transform.translation.x, + transform.transform.translation.y, + transform.transform.translation.z, + transform.transform.rotation.x, + transform.transform.rotation.y, + transform.transform.rotation.z, + transform.transform.rotation.w, ) return obs.pose diff --git a/dimos/perception/fiducial/test_fixture_verification.py b/dimos/perception/fiducial/test_fixture_verification.py index 295fd1ef375..35047537e7b 100644 --- a/dimos/perception/fiducial/test_fixture_verification.py +++ b/dimos/perception/fiducial/test_fixture_verification.py @@ -15,14 +15,15 @@ from __future__ import annotations import cv2 +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.camera_info import camera_info_from_intrinsics +from dimos.msgs.geometry import transform_matrix +from dimos.msgs.image import image_from_array +from dimos.msgs.time import time_from_seconds from dimos.perception.detection.type.detection3d.imageDetections3D import ImageDetections3D from dimos.perception.fiducial.fixture_verification import ( BoardLayout, @@ -159,24 +160,15 @@ def test_marker_tf_replay_synthetic_packed_board_publishes_twelve_markers( width, height = 1920, 1080 bgr = _synthetic_packed_apriltag_board_bgr(layout, width=width, height=height) ts = 10_000.0 - cam_info = CameraInfo.from_intrinsics( - 1400.0, - 1400.0, - width / 2.0, - height / 2.0, - width, - height, - frame_id="camera_optical", + header = Header(stamp=time_from_seconds(ts), frame_id="camera_optical") + cam_info = camera_info_from_intrinsics( + fx=1400, fy=1400, cx=width / 2, cy=height / 2, width=width, height=height, header=header ) - cam_info.ts = ts - - image = Image.from_opencv(bgr, frame_id="camera_optical", ts=ts) - world_T_optical = Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="world", + image = image_from_array(bgr, encoding="bgr8", header=header) + world_T_optical = TransformStamped( + header=Header(stamp=header.stamp, frame_id="world"), child_frame_id="camera_optical", - ts=ts, + transform=Transform(rotation=Quaternion(w=1)), ) detections = detect_markers_in_image( image, @@ -199,8 +191,9 @@ def test_marker_tf_replay_synthetic_packed_board_publishes_twelve_markers( for tag_id in range(12): tr = view.get(marker_parent, f"fixture/marker_{tag_id}", ts, 0.1) assert tr is not None, f"missing marker {tag_id}" - assert np.all(np.isfinite(tr.to_matrix())) + assert np.all(np.isfinite(transform_matrix(tr.transform))) finally: + view.dispose() mod.stop() diff --git a/dimos/perception/fiducial/test_generated_marker_detect.py b/dimos/perception/fiducial/test_generated_marker_detect.py new file mode 100644 index 00000000000..178185eed28 --- /dev/null +++ b/dimos/perception/fiducial/test_generated_marker_detect.py @@ -0,0 +1,55 @@ +# Copyright 2025-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. + +import cv2 +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.std_msgs.msg import Header +from dimos_generated.vision_msgs.msg import Detection3D +import numpy as np +import pytest + +from dimos.msgs.image import image_from_array +from dimos.perception.fiducial.marker_detect import detect_markers_in_image + + +def test_generated_image_detects_real_aruco_and_composes_pose(): + dictionary = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_50) + pixels = np.full((400, 400), 255, dtype=np.uint8) + pixels[100:300, 100:300] = cv2.aruco.generateImageMarker(dictionary, 7, 200) + header = Header(frame_id="camera", stamp=Time(sec=1700000000, nanosec=123456789)) + image = image_from_array(pixels, encoding="mono8", header=header) + image = Image.decode(image.encode()) + camera = CameraInfo(width=400, height=400, k=[400, 0, 200, 0, 400, 200, 0, 0, 1]) + tf = TransformStamped( + header=Header(frame_id="world"), + child_frame_id="camera", + transform=Transform(translation=Vector3(x=2), rotation=Quaternion(w=1)), + ) + found = detect_markers_in_image( + image, + camera_info=camera, + world_T_optical=tf, + marker_length_m=0.2, + aruco_dictionary="DICT_4X4_50", + ) + assert len(found) == 1 + wire = Detection3D.decode(found[0].to_detection3d_msg().encode()) + assert wire.id == "7" + assert wire.results[0].hypothesis.class_id == "DICT_4X4_50:7" + assert wire.header.stamp == header.stamp + assert wire.header.frame_id == "world" + assert wire.bbox.center.position.x == pytest.approx(2, abs=0.01) + assert wire.bbox.center.position.z == pytest.approx(0.4, abs=0.01) diff --git a/dimos/perception/fiducial/test_helpers.py b/dimos/perception/fiducial/test_helpers.py index 12b4f890450..bf2449aa042 100644 --- a/dimos/perception/fiducial/test_helpers.py +++ b/dimos/perception/fiducial/test_helpers.py @@ -17,35 +17,31 @@ from __future__ import annotations import cv2 +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.std_msgs.msg import Header import numpy as np -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_from_array +from dimos.msgs.time import time_from_seconds def camera_info(ts: float = 10.0) -> CameraInfo: - info = CameraInfo.from_intrinsics( - fx=600.0, - fy=600.0, - cx=320.0, - cy=240.0, + info = CameraInfo( width=640, height=480, - frame_id="camera_optical", + k=[600, 0, 320, 0, 600, 240, 0, 0, 1], + header=Header(frame_id="camera_optical"), ) - info.ts = ts + info.header.stamp = time_from_seconds(ts) return info def blank_image(ts: float = 10.0) -> Image: - return Image( - data=np.full((480, 640, 3), 255, dtype=np.uint8), - format=ImageFormat.BGR, - frame_id="camera_optical", - ts=ts, + return image_from_array( + np.full((480, 640, 3), 255, dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id="camera_optical", stamp=time_from_seconds(ts)), ) @@ -70,19 +66,19 @@ def synthetic_marker_image(marker_id: int = 7, ts: float = 10.0, inverted: bool y0 = (canvas.shape[0] - side_px) // 2 x0 = (canvas.shape[1] - side_px) // 2 canvas[y0 : y0 + side_px, x0 : x0 + side_px] = tile - return Image( - data=cv2.cvtColor(canvas, cv2.COLOR_GRAY2BGR), - format=ImageFormat.BGR, - frame_id="camera_optical", - ts=ts, + return image_from_array( + cv2.cvtColor(canvas, cv2.COLOR_GRAY2BGR), + encoding="bgr8", + header=Header(frame_id="camera_optical", stamp=time_from_seconds(ts)), ) -def world_T_optical(ts: float = 10.0) -> Transform: - return Transform( - translation=Vector3(1.0, 2.0, 3.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="world", +def world_T_optical(ts: float = 10.0) -> TransformStamped: + return TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(ts)), child_frame_id="camera_optical", - ts=ts, + transform=Transform( + translation=Vector3(x=1.0, y=2.0, z=3.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) diff --git a/dimos/perception/fiducial/test_marker_detection_stream_module.py b/dimos/perception/fiducial/test_marker_detection_stream_module.py index 59779c5fc24..402ba2a789c 100644 --- a/dimos/perception/fiducial/test_marker_detection_stream_module.py +++ b/dimos/perception/fiducial/test_marker_detection_stream_module.py @@ -16,16 +16,17 @@ from typing import Any +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header +from dimos_generated.vision_msgs.msg import Detection3DArray import numpy as np import pytest from dimos.memory.store.memory import MemoryStore +from dimos.memory.transform import SpeedLimit from dimos.memory.type.observation import Observation -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.perception.detection.type.detection3d.marker import Detection3DMarker from dimos.perception.fiducial.marker_detection_stream_module import MarkerDetectionStreamModule from dimos.perception.fiducial.marker_transformer import MarkersPerFrame @@ -43,19 +44,22 @@ def _marker(image: Image, marker_id: int) -> Detection3DMarker: class_id=marker_id, confidence=1.0, name="", - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, - center=Vector3(float(marker_id), 2.0, 3.0), - size=Vector3(0.18, 0.18, 0.0), - transform=Transform( - translation=Vector3(1.0, 2.0, 3.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="world", + center=Vector3(x=float(marker_id), y=2.0, z=3.0), + size=Vector3(x=0.18, y=0.18, z=0.0), + transform=TransformStamped( + header=Header( + frame_id="world", stamp=time_from_seconds(to_seconds(image.header.stamp)) + ), child_frame_id="camera_optical", - ts=image.ts, + transform=Transform( + translation=Vector3(x=1.0, y=2.0, z=3.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), frame_id="world", - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), marker_id=marker_id, corners_px=np.array( [[10.0, 20.0], [40.0, 20.0], [40.0, 50.0], [10.0, 50.0]], @@ -82,7 +86,7 @@ def _marker_obs( tags["marker_frame_index"] = marker_index return Observation( id=obs_id, - ts=image.ts, + ts=to_seconds(image.header.stamp), data_type=Detection3DMarker if marker is not None else type(None), pose=(1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 1.0), tags=tags, @@ -120,15 +124,15 @@ def test_markers_per_frame_groups_markers_and_preserves_empty_frames() -> None: assert len(outputs) == 2 first = outputs[0].data assert first.header.frame_id == "world" - assert first.ts == pytest.approx(image.ts) - assert first.detections_length == 2 + assert to_seconds(first.header.stamp) == pytest.approx(to_seconds(image.header.stamp)) + assert len(first.detections) == 2 assert [det.id for det in first.detections] == ["7", "42"] assert outputs[0].pose_tuple == pytest.approx((1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 1.0)) empty = outputs[1].data assert empty.header.frame_id == "world" - assert empty.ts == pytest.approx(empty_image.ts) - assert empty.detections_length == 0 + assert to_seconds(empty.header.stamp) == pytest.approx(to_seconds(empty_image.header.stamp)) + assert len(empty.detections) == 0 assert empty.detections == [] @@ -140,7 +144,7 @@ def test_marker_detection_stream_pipeline_outputs_arrays_for_marker_and_empty_fr module = MarkerDetectionStreamModule( marker_length_m=marker_length_m, - camera_info=camera_info(marker_image.ts), + camera_info=camera_info(to_seconds(marker_image.header.stamp)), quality_window_s=0.01, ) try: @@ -148,12 +152,12 @@ def test_marker_detection_stream_pipeline_outputs_arrays_for_marker_and_empty_fr stream = store.stream("color_image", Image) stream.append( marker_image, - ts=marker_image.ts, + ts=to_seconds(marker_image.header.stamp), pose=(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0), ) stream.append( empty_image, - ts=empty_image.ts, + ts=to_seconds(empty_image.header.stamp), pose=(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0), ) @@ -162,15 +166,17 @@ def test_marker_detection_stream_pipeline_outputs_arrays_for_marker_and_empty_fr module.stop() assert len(outputs) == 2 - assert outputs[0].detections_length == 1 + assert len(outputs[0].detections) == 1 assert outputs[0].detections[0].id == str(marker_id) assert outputs[0].detections[0].results[0].hypothesis.class_id == ( f"DICT_APRILTAG_36h11:{marker_id}" ) assert outputs[0].detections[0].bbox.size.x == pytest.approx(marker_length_m) - assert outputs[1].ts == pytest.approx(empty_image.ts) - assert outputs[1].detections_length == 0 + assert to_seconds(outputs[1].header.stamp) == pytest.approx( + to_seconds(empty_image.header.stamp) + ) + assert len(outputs[1].detections) == 0 assert outputs[1].detections == [] @@ -199,7 +205,7 @@ def run_pipeline(*, speed_limit_enabled: bool) -> list[Detection3DArray]: with MemoryStore() as store: stream = store.stream("color_image", Image) for image, pose in zip(images, poses, strict=True): - stream.append(image, ts=image.ts, pose=pose) + stream.append(image, ts=to_seconds(image.header.stamp), pose=pose) return [obs.data for obs in module.pipeline(stream).to_list()] finally: module.stop() @@ -207,21 +213,23 @@ def run_pipeline(*, speed_limit_enabled: bool) -> list[Detection3DArray]: disabled = run_pipeline(speed_limit_enabled=False) enabled = run_pipeline(speed_limit_enabled=True) - assert [msg.ts for msg in disabled] == pytest.approx([10.0, 11.0, 12.0]) - assert all(msg.detections_length == 0 for msg in disabled) - assert [msg.ts for msg in enabled] == pytest.approx([12.0]) - assert enabled[0].detections_length == 0 + assert [to_seconds(msg.header.stamp) for msg in disabled] == pytest.approx([10.0, 11.0, 12.0]) + assert all(len(msg.detections) == 0 for msg in disabled) + assert [to_seconds(msg.header.stamp) for msg in enabled] == pytest.approx([12.0]) + assert len(enabled[0].detections) == 0 -def test_append_image_with_pose_uses_camera_optical_tf_without_recomputing_pose() -> None: - image = blank_image(ts=12.0) - info = camera_info(image.ts) - t_world_optical = Transform( - translation=Vector3(4.0, 5.0, 6.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="world", +@pytest.mark.parametrize("timestamp", [0.0, 12.0]) +def test_append_image_with_pose_uses_camera_optical_tf_without_recomputing_pose(timestamp) -> None: + image = blank_image(ts=timestamp) + info = camera_info(to_seconds(image.header.stamp)) + t_world_optical = TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(to_seconds(image.header.stamp))), child_frame_id="camera_optical", - ts=image.ts, + transform=Transform( + translation=Vector3(x=4.0, y=5.0, z=6.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) class FakeTf: @@ -234,7 +242,7 @@ def get( child_frame: str, time_point: float | None = None, time_tolerance: float | None = None, - ) -> Transform: + ) -> TransformStamped: self.calls.append((parent_frame, child_frame, time_point, time_tolerance)) return t_world_optical @@ -258,10 +266,10 @@ def dispose(self) -> None: finally: module.stop() - assert fake_tf.calls == [("world", "camera_optical", image.ts, 0.25)] + assert fake_tf.calls == [("world", "camera_optical", to_seconds(image.header.stamp), 0.25)] assert len(observations) == 1 assert observations[0].data is image - assert observations[0].ts == pytest.approx(image.ts) + assert observations[0].ts == pytest.approx(to_seconds(image.header.stamp)) assert observations[0].pose_tuple == pytest.approx((4.0, 5.0, 6.0, 0.0, 0.0, 0.0, 1.0)) @@ -289,7 +297,9 @@ def dispose(self) -> None: pass missing_tf = MissingTf() - module = MarkerDetectionStreamModule(marker_length_m=0.18, camera_info=camera_info(image.ts)) + module = MarkerDetectionStreamModule( + marker_length_m=0.18, camera_info=camera_info(to_seconds(image.header.stamp)) + ) module._tf = missing_tf try: with MemoryStore() as store: @@ -300,3 +310,14 @@ def dispose(self) -> None: module.stop() assert missing_tf.calls == 1 + + +def test_generated_speed_limit_handles_quaternion_sign_and_rotation(): + image = blank_image() + poses = [(0, 0, 0, 0, 0, 0, 1), (0, 0, 0, 0, 0, 0, -1), (0, 0, 0, 0, 0, 1, 0)] + observations = [ + Observation(id=i, ts=float(i), data_type=Image, _data=image, pose=pose) + for i, pose in enumerate(poses) + ] + accepted = list(SpeedLimit(max_mps=1, max_dps=15)(iter(observations))) + assert [obs.id for obs in accepted] == [1] diff --git a/dimos/perception/fiducial/test_marker_pose.py b/dimos/perception/fiducial/test_marker_pose.py index 78c22751805..31f64f6121e 100644 --- a/dimos/perception/fiducial/test_marker_pose.py +++ b/dimos/perception/fiducial/test_marker_pose.py @@ -13,10 +13,11 @@ # limitations under the License. import cv2 +from dimos_generated.sensor_msgs.msg import CameraInfo +from dimos_generated.std_msgs.msg import Header import numpy as np -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_from_array from dimos.perception.fiducial.marker_pose import ( camera_optical_frame_id, estimate_marker_pose, @@ -26,25 +27,14 @@ def test_camera_optical_frame_id_resolution() -> None: - ts = 1.0 - fx, fy, cx, cy = 600.0, 600.0, 320.0, 240.0 - info_named = CameraInfo.from_intrinsics(fx, fy, cx, cy, 640, 480, frame_id="cam_info_optical") - info_named.ts = ts - info_empty = CameraInfo.from_intrinsics(fx, fy, cx, cy, 640, 480) - info_empty.ts = ts - img_custom = Image( - data=np.zeros((480, 640, 3), dtype=np.uint8), - format=ImageFormat.BGR, - ts=ts, - frame_id="custom_optical", + info_named = CameraInfo(header=Header(frame_id="cam_info_optical")) + info_empty = CameraInfo() + pixels = np.zeros((480, 640, 3), dtype=np.uint8) + img_custom = image_from_array(pixels, encoding="bgr8", header=Header(frame_id="custom_optical")) + img_whitespace = image_from_array( + pixels, encoding="bgr8", header=Header(frame_id=" custom_optical ") ) - img_whitespace = Image( - data=np.zeros((480, 640, 3), dtype=np.uint8), - format=ImageFormat.BGR, - ts=ts, - frame_id=" custom_optical ", - ) - img_empty = Image(data=np.zeros((480, 640, 3), dtype=np.uint8), format=ImageFormat.BGR, ts=ts) + img_empty = image_from_array(pixels, encoding="bgr8") assert camera_optical_frame_id(img_custom, info_named) == "custom_optical" assert camera_optical_frame_id(img_whitespace, info_named) == "custom_optical" diff --git a/dimos/perception/fiducial/test_marker_tf_integration.py b/dimos/perception/fiducial/test_marker_tf_integration.py index 458edd22e67..43fa7fbe453 100644 --- a/dimos/perception/fiducial/test_marker_tf_integration.py +++ b/dimos/perception/fiducial/test_marker_tf_integration.py @@ -19,25 +19,28 @@ import time import uuid -from dimos_lcm.vision_msgs import BoundingBox3D, ObjectHypothesis, ObjectHypothesisWithPose +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion, Vector3 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage +from dimos_generated.vision_msgs.msg import ( + BoundingBox3D, + Detection3D, + Detection3DArray, + ObjectHypothesis, + ObjectHypothesisWithPose, +) import pytest from dimos.core.transport import LCMTransport from dimos.core.transport_factory import make_transport -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.std_msgs.Header import Header -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.msgs.vision_msgs.Detection3D import Detection3D -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.msgs.time import time_from_seconds from dimos.perception.fiducial.marker_tf_module import MarkerTfModule from dimos.protocol.tf.tf import TF def _marker_detection_array(ts: float) -> Detection3DArray: det = Detection3D() - det.header = Header(ts, "world") + det.header = Header(stamp=time_from_seconds(ts), frame_id="world") det.id = "11" det.results = [ ObjectHypothesisWithPose( @@ -47,18 +50,16 @@ def _marker_detection_array(ts: float) -> Detection3DArray: ) ) ] - det.results_length = len(det.results) det.bbox = BoundingBox3D( center=Pose( - position=Vector3(0.5, -0.2, 1.25), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + position=Point(x=0.5, y=-0.2, z=1.25), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), ), - size=Vector3(0.18, 0.18, 0.0), + size=Vector3(x=0.18, y=0.18, z=0.0), ) return Detection3DArray( - header=Header(ts, "world"), + header=Header(stamp=time_from_seconds(ts), frame_id="world"), detections=[det], - detections_length=1, ) @@ -89,13 +90,13 @@ def test_marker_tf_consumes_detection_array_over_lcm() -> None: assert w_markers is not None, "Timed out waiting for world -> marker_tf/markers" assert w_marker is not None, "Timed out waiting for world -> marker_tf/marker_11" - assert w_markers.frame_id == "world" + assert w_markers.header.frame_id == "world" assert w_markers.child_frame_id == "marker_tf/markers" - assert w_marker.frame_id == "world" + assert w_marker.header.frame_id == "world" assert w_marker.child_frame_id == "marker_tf/marker_11" - assert w_marker.translation.x == pytest.approx(0.5) - assert w_marker.translation.y == pytest.approx(-0.2) - assert w_marker.translation.z == pytest.approx(1.25) + assert w_marker.transform.translation.x == pytest.approx(0.5) + assert w_marker.transform.translation.y == pytest.approx(-0.2) + assert w_marker.transform.translation.z == pytest.approx(1.25) finally: host_tf.dispose() host_transport.stop() diff --git a/dimos/perception/fiducial/test_marker_tf_module.py b/dimos/perception/fiducial/test_marker_tf_module.py index 90d2c10ff36..2c4aeb3c261 100644 --- a/dimos/perception/fiducial/test_marker_tf_module.py +++ b/dimos/perception/fiducial/test_marker_tf_module.py @@ -14,15 +14,20 @@ from unittest.mock import patch -from dimos_lcm.vision_msgs import BoundingBox3D, ObjectHypothesis, ObjectHypothesisWithPose +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion, Vector3 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage +from dimos_generated.vision_msgs.msg import ( + BoundingBox3D, + Detection3D, + Detection3DArray, + ObjectHypothesis, + ObjectHypothesisWithPose, +) import pytest -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.std_msgs.Header import Header -from dimos.msgs.vision_msgs.Detection3D import Detection3D -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.time import time_from_seconds from dimos.perception.fiducial.marker_tf_module import MarkerTfModule from dimos.protocol.tf.tf import TF @@ -32,16 +37,16 @@ def _detection_array( ts: float, marker_id: str = "0", class_id: str = "DICT_APRILTAG_36h11:0", - center: Vector3 | None = None, + center: Point | None = None, orientation: Quaternion | None = None, ) -> Detection3DArray: if center is None: - center = Vector3(1.2, -0.3, 0.8) + center = Point(x=1.2, y=-0.3, z=0.8) if orientation is None: - orientation = Quaternion(0.0, 0.0, 0.0, 1.0) + orientation = Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) det = Detection3D() - det.header = Header(ts, "world") + det.header = Header(stamp=time_from_seconds(ts), frame_id="world") det.id = marker_id det.results = [ ObjectHypothesisWithPose( @@ -51,25 +56,23 @@ def _detection_array( ) ) ] - det.results_length = len(det.results) det.bbox = BoundingBox3D( center=Pose( position=center, orientation=orientation, ), - size=Vector3(0.18, 0.18, 0.0), + size=Vector3(x=0.18, y=0.18, z=0.0), ) return Detection3DArray( - header=Header(ts, "world"), + header=Header(stamp=time_from_seconds(ts), frame_id="world"), detections=[det], - detections_length=1, ) def test_marker_tf_module_publishes_world_markers_chain() -> None: ts = 1_000_000.0 - center = Vector3(1.3, -0.2, 0.4) - orientation = Quaternion(0.1, 0.2, 0.3, 0.9) + center = Point(x=1.3, y=-0.2, z=0.4) + orientation = quaternion_from_euler(0.1, 0.2, 0.3) mod = MarkerTfModule() view = TF(mod.tf) @@ -86,21 +89,22 @@ def test_marker_tf_module_publishes_world_markers_chain() -> None: wm = view.get("world", "markers", ts, 1.0) assert wm is not None - assert abs(wm.translation.x) < 1e-6 - assert abs(wm.translation.y) < 1e-6 - assert abs(wm.translation.z) < 1e-6 + assert abs(wm.transform.translation.x) < 1e-6 + assert abs(wm.transform.translation.y) < 1e-6 + assert abs(wm.transform.translation.z) < 1e-6 w_m7 = view.get("world", "marker_7", ts, 1.0) assert w_m7 is not None - assert w_m7.translation.x == pytest.approx(center.x) - assert w_m7.translation.y == pytest.approx(center.y) - assert w_m7.translation.z == pytest.approx(center.z) - assert w_m7.rotation.x == pytest.approx(orientation.x) - assert w_m7.rotation.y == pytest.approx(orientation.y) - assert w_m7.rotation.z == pytest.approx(orientation.z) - assert w_m7.rotation.w == pytest.approx(orientation.w) + assert w_m7.transform.translation.x == pytest.approx(center.x) + assert w_m7.transform.translation.y == pytest.approx(center.y) + assert w_m7.transform.translation.z == pytest.approx(center.z) + assert w_m7.transform.rotation.x == pytest.approx(orientation.x) + assert w_m7.transform.rotation.y == pytest.approx(orientation.y) + assert w_m7.transform.rotation.z == pytest.approx(orientation.z) + assert w_m7.transform.rotation.w == pytest.approx(orientation.w) assert view.get("world", "marker_99", ts, 0.1) is None finally: + view.dispose() mod.stop() @@ -116,6 +120,7 @@ def test_marker_tf_parses_class_id_when_detection_id_empty() -> None: assert view.get("world", "marker_42", ts, 1.0) is not None finally: + view.dispose() mod.stop() @@ -126,14 +131,14 @@ def test_marker_tf_empty_array_skips_publication() -> None: try: mod._process_detections( Detection3DArray( - header=Header(ts, "world"), + header=Header(stamp=time_from_seconds(ts), frame_id="world"), detections=[], - detections_length=0, ) ) assert view.get("world", "markers", ts, 0.1) is None finally: + view.dispose() mod.stop() @@ -146,6 +151,7 @@ def test_marker_tf_non_empty_array_without_marker_id_skips_publication() -> None assert view.get("world", "markers", ts, 0.1) is None finally: + view.dispose() mod.stop() @@ -160,6 +166,7 @@ def test_marker_tf_does_not_recompute_marker_pose() -> None: mock_estimate.assert_not_called() assert view.get("world", "marker_4", ts, 1.0) is not None finally: + view.dispose() mod.stop() @@ -174,4 +181,22 @@ def test_marker_namespace_prefix_child_frames() -> None: assert view.get("world", "r1/markers", ts, 1.0) is not None assert view.get("world", "r1/marker_0", ts, 1.0) is not None finally: + view.dispose() mod.stop() + + +def test_marker_tf_wire_preserves_exact_timestamp(): + module = MarkerTfModule() + outputs = [] + unsubscribe = module.tf.subscribe(lambda msg: outputs.append(TFMessage.decode(msg.encode()))) + detections = _detection_array(ts=0) + detections.header.stamp.sec = 1700000000 + detections.header.stamp.nanosec = 123456789 + try: + module._process_detections(detections) + assert len(outputs) == 1 + assert len(outputs[0].transforms) == 2 + assert all(t.header.stamp == detections.header.stamp for t in outputs[0].transforms) + finally: + unsubscribe() + module.stop() diff --git a/dimos/perception/fiducial/test_marker_transformer.py b/dimos/perception/fiducial/test_marker_transformer.py index 16830e35561..cef20524a59 100644 --- a/dimos/perception/fiducial/test_marker_transformer.py +++ b/dimos/perception/fiducial/test_marker_transformer.py @@ -12,12 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from dimos.memory.type.observation import Observation -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_from_array +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.perception.fiducial.marker_detect import detect_markers_in_image from dimos.perception.fiducial.marker_transformer import DetectMarkers from dimos.perception.fiducial.test_helpers import ( @@ -31,12 +33,12 @@ def test_detect_markers_in_image_builds_rich_marker_detection() -> None: marker_id = 7 marker_length_m = 0.18 image = synthetic_marker_image(marker_id) - info = camera_info(image.ts) + info = camera_info(to_seconds(image.header.stamp)) detections = detect_markers_in_image( image, camera_info=info, - world_T_optical=world_T_optical(image.ts), + world_T_optical=world_T_optical(to_seconds(image.header.stamp)), marker_length_m=marker_length_m, aruco_dictionary="DICT_APRILTAG_36h11", ) @@ -65,11 +67,10 @@ def test_detect_markers_in_image_builds_rich_marker_detection() -> None: def test_detect_markers_in_image_returns_empty_for_no_marker_frame() -> None: ts = 11.0 - image = Image( - data=np.full((480, 640, 3), 255, dtype=np.uint8), - format=ImageFormat.BGR, - frame_id="camera_optical", - ts=ts, + image = image_from_array( + np.full((480, 640, 3), 255, dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id="camera_optical", stamp=time_from_seconds(ts)), ) detections = detect_markers_in_image( @@ -86,8 +87,8 @@ def test_detect_markers_in_image_returns_empty_for_no_marker_frame() -> None: def test_detect_markers_in_image_needs_detect_inverted_for_negative_tags() -> None: image = synthetic_marker_image(7, inverted=True) kwargs = { - "camera_info": camera_info(image.ts), - "world_T_optical": world_T_optical(image.ts), + "camera_info": camera_info(to_seconds(image.header.stamp)), + "world_T_optical": world_T_optical(to_seconds(image.header.stamp)), "marker_length_m": 0.18, "aruco_dictionary": "DICT_APRILTAG_36h11", } @@ -104,13 +105,13 @@ def test_detect_markers_transformer_preserves_observation_context_and_tags() -> image = synthetic_marker_image(marker_id, ts=12.0) obs = Observation[Image]( id=42, - ts=image.ts, + ts=to_seconds(image.header.stamp), data_type=Image, pose=(1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 1.0), _data=image, ) transformer = DetectMarkers( - camera_info=camera_info(image.ts), + camera_info=camera_info(to_seconds(image.header.stamp)), marker_length_m=0.18, aruco_dictionary="DICT_APRILTAG_36h11", ) @@ -130,21 +131,20 @@ def test_detect_markers_transformer_preserves_observation_context_and_tags() -> def test_detect_markers_transformer_can_emit_empty_frame_sentinel() -> None: - image = Image( - data=np.full((480, 640, 3), 255, dtype=np.uint8), - format=ImageFormat.BGR, - frame_id="camera_optical", - ts=13.0, + image = image_from_array( + np.full((480, 640, 3), 255, dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id="camera_optical", stamp=time_from_seconds(13.0)), ) obs = Observation[Image]( id=43, - ts=image.ts, + ts=to_seconds(image.header.stamp), data_type=Image, pose=(1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 1.0), _data=image, ) transformer = DetectMarkers( - camera_info=camera_info(image.ts), + camera_info=camera_info(to_seconds(image.header.stamp)), marker_length_m=0.18, aruco_dictionary="DICT_APRILTAG_36h11", emit_empty_frames=True, @@ -165,7 +165,7 @@ def test_detect_markers_transformer_uses_callablecamera_info_source() -> None: image = synthetic_marker_image(marker_id=7, ts=14.0) obs = Observation[Image]( id=44, - ts=image.ts, + ts=to_seconds(image.header.stamp), data_type=Image, pose=(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0), _data=image, @@ -180,7 +180,7 @@ def test_detect_markers_transformer_uses_callablecamera_info_source() -> None: assert list(transformer(iter([obs]))) == [] - latest_info = camera_info(image.ts) + latest_info = camera_info(to_seconds(image.header.stamp)) results = list(transformer(iter([obs]))) assert len(results) == 1 @@ -193,24 +193,23 @@ def test_detect_markers_rebuilds_intrinsics_without_resetting_smoothing_track() marker_length_m = 0.18 image_a = synthetic_marker_image(marker_id=marker_id, ts=15.0) image_b = synthetic_marker_image(marker_id=marker_id, ts=15.2) - info_a = camera_info(image_a.ts) - info_b = camera_info(image_b.ts) - info_b.K = info_b.K.copy() - info_b.P = info_b.P.copy() - info_b.K[0] = info_b.K[4] = 900.0 - info_b.P[0] = info_b.P[5] = 900.0 + info_a = camera_info(to_seconds(image_a.header.stamp)) + info_b = camera_info(to_seconds(image_b.header.stamp)) + + info_b.k[0] = info_b.k[4] = 900.0 + info_b.p[0] = info_b.p[5] = 900.0 latest_info: CameraInfo | None = info_a obs_a = Observation[Image]( id=45, - ts=image_a.ts, + ts=to_seconds(image_a.header.stamp), data_type=Image, pose=(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0), _data=image_a, ) obs_b = Observation[Image]( id=46, - ts=image_b.ts, + ts=to_seconds(image_b.header.stamp), data_type=Image, pose=(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0), _data=image_b, diff --git a/dimos/perception/memory/test_generated_inventory.py b/dimos/perception/memory/test_generated_inventory.py new file mode 100644 index 00000000000..8c3223fc4b8 --- /dev/null +++ b/dimos/perception/memory/test_generated_inventory.py @@ -0,0 +1,73 @@ +# 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. + +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.std_msgs.msg import Header +import numpy as np + +from dimos.msgs.camera_info import camera_info_from_intrinsics +from dimos.msgs.pointcloud import pointcloud_from_xyz_rgb, pointcloud_rgb, pointcloud_xyz +from dimos.msgs.time import time_from_nanoseconds, to_nanoseconds +from dimos.perception.memory.inventory import _absorb_into, _pixel_bbox, _split_oversized +from dimos.perception.memory.types import InventoryPolicy, SupportObservation + + +def _observation(offset: float, timestamp_ns: int) -> SupportObservation: + points = np.array([[offset, 0, 1], [offset + 0.1, 0, 1]]) + cloud = pointcloud_from_xyz_rgb( + points, + np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint8), + header=Header(frame_id="world", stamp=time_from_nanoseconds(timestamp_ns)), + ) + return SupportObservation( + ts=timestamp_ns / 1e9, + cloud=cloud, + centroid=points.mean(axis=0), + aabb_min=points.min(axis=0), + aabb_max=points.max(axis=0), + n_points=2, + mask_area_px=2, + camera_position=np.zeros(3), + ) + + +def test_generated_support_merge_preserves_colors_and_source_time(): + first = _observation(0, 1700000000123456789) + second = _observation(0.2, 1700000000123456790) + _absorb_into(first, second) + assert first.n_points == 4 + assert to_nanoseconds(first.cloud.header.stamp) == 1700000000123456790 + np.testing.assert_array_equal( + pointcloud_rgb(first.cloud), [[1, 2, 3], [4, 5, 6], [1, 2, 3], [4, 5, 6]] + ) + np.testing.assert_allclose(first.centroid, pointcloud_xyz(first.cloud).mean(axis=0)) + assert second.n_points == 2 and second.cloud.width == 2 + + +def test_generated_support_projection_reads_nested_transform_and_intrinsics(): + calibration = camera_info_from_intrinsics(10, 20, 5, 6, 100, 100, header=Header()) + transform = TransformStamped( + header=Header(frame_id="optical"), + child_frame_id="world", + transform=Transform(translation=Vector3(x=-1), rotation=Quaternion(w=1)), + ) + assert _pixel_bbox(np.array([[1, 0, 2], [2, 1, 2]]), calibration, transform) == (5, 6, 10, 16) + + +def test_unsplit_support_returns_source_point_mask(): + points = np.array([[0, 0, 0], [0.01, 0.01, 0.01]]) + masks = _split_oversized(points, None, InventoryPolicy()) + assert len(masks) == 1 + assert masks[0].dtype == np.bool_ + np.testing.assert_array_equal(points[masks[0]], points) diff --git a/dimos/protocol/pubsub/benchmark/test_replay_bench.py b/dimos/protocol/pubsub/benchmark/test_replay_bench.py index 4e1d772ed29..6d663567950 100644 --- a/dimos/protocol/pubsub/benchmark/test_replay_bench.py +++ b/dimos/protocol/pubsub/benchmark/test_replay_bench.py @@ -20,17 +20,17 @@ import time from types import SimpleNamespace +from dimos_generated.sensor_msgs.msg import CompressedImage, Image +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from dimos.core.transport import LCMTransport, ZenohTransport -from dimos.msgs.sensor_msgs.CompressedImage import CompressedImage -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_from_array, image_to_jpeg, image_view +from dimos.msgs.time import time_from_seconds from dimos.protocol.pubsub.benchmark.tool_replay_bench import CompressedCodec from dimos.protocol.service.zenohservice import ZenohSessionPool -pytestmark = pytest.mark.skipif_no_turbojpeg - def make_image(width: int = 1280, height: int = 720) -> Image: """720p-ish frame: gradient + noise, roughly what a real camera compresses like.""" @@ -38,7 +38,9 @@ def make_image(width: int = 1280, height: int = 720) -> Image: gradient = np.broadcast_to(np.linspace(0, 255, width, dtype=np.uint8), (height, width)) noise = rng.randint(0, 60, (height, width), dtype=np.uint8) data = np.stack([gradient, np.minimum(gradient, 128) + noise // 2, noise], axis=-1) - return Image(data=data, format=ImageFormat.RGB, frame_id="cam", ts=42.125) + return image_from_array( + data, encoding="rgb8", header=Header(frame_id="cam", stamp=time_from_seconds(42.125)) + ) @pytest.fixture() @@ -60,30 +62,38 @@ def session_pool(): pool.close_all() -def test_roundtrip_over_lcm(retry_until, collector) -> None: - t = CompressedCodec(LCMTransport("dimos/test/codec_lcm", CompressedImage)) +@pytest.fixture +def codec_transports(): + transports = [] + yield transports + for transport in transports: + transport.stop() + + +def test_roundtrip_over_lcm(retry_until, collector, codec_transports, lcm_url) -> None: + t = CompressedCodec(LCMTransport("dimos/test/codec_lcm", CompressedImage, url=lcm_url)) + codec_transports.append(t) t.subscribe(collector.callback) src = make_image(320, 240) retry_until(collector.event, lambda: t.broadcast(None, src)) img = collector.received[0] assert isinstance(img, Image) - assert img.frame_id == "cam" - assert abs(img.ts - src.ts) < 1e-6 - assert img.shape == src.shape - t.stop() + assert img.header.frame_id == "cam" + assert img.header.stamp == src.header.stamp + assert image_view(img).shape == image_view(src).shape -def test_roundtrip_over_zenoh(retry_until, collector, session_pool) -> None: +def test_roundtrip_over_zenoh(retry_until, collector, session_pool, codec_transports) -> None: t = CompressedCodec( ZenohTransport("dimos/test/codec_zenoh", CompressedImage, session_pool=session_pool) ) + codec_transports.append(t) t.subscribe(collector.callback) src = make_image(320, 240) retry_until(collector.event, lambda: t.broadcast(None, src)) img = collector.received[0] assert isinstance(img, Image) - assert abs(img.ts - src.ts) < 1e-6 - t.stop() + assert img.header.stamp == src.header.stamp def test_pickle_roundtrip() -> None: @@ -184,8 +194,10 @@ def inner(topic: str, typ: type): return LCMTransport(topic, typ, url=f"{lcm_url}&recv_buf_size={2 * 1024 * 1024}") return ZenohTransport(topic, typ, session_pool=session_pool) - raw_wire = len(frame.lcm_encode()) - jpeg_wire = len(CompressedImage.from_image(frame).lcm_encode()) + raw_wire = len(frame.encode()) + jpeg_wire = len( + CompressedImage(header=frame.header, format="jpeg", data=image_to_jpeg(frame)).encode() + ) raw = _bench(inner(f"dimos/bench/{proto}_raw", Image), frame, n, raw_wire) codec = _bench( diff --git a/dimos/protocol/pubsub/benchmark/tool_replay_bench.py b/dimos/protocol/pubsub/benchmark/tool_replay_bench.py index ee80e5854e2..da4796d0eee 100644 --- a/dimos/protocol/pubsub/benchmark/tool_replay_bench.py +++ b/dimos/protocol/pubsub/benchmark/tool_replay_bench.py @@ -35,21 +35,23 @@ from pathlib import Path import threading import time -from typing import Any, TypeVar +from typing import Any + +from dimos_generated.sensor_msgs.msg import CompressedImage, Image +import numpy as np from dimos.core.core import rpc from dimos.core.module import Module -from dimos.core.stream import In, Out, Stream +from dimos.core.stream import In, Stream from dimos.core.transport import PubSubTransport -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_from_array, image_to_jpeg, image_view +from dimos.msgs.time import to_seconds from dimos.utils.logging_config import setup_logger logger = setup_logger() -T = TypeVar("T") - -class CompressedCodec(PubSubTransport[T]): +class CompressedCodec(PubSubTransport[Image]): """Image→CompressedImage jpeg codec over any inner transport. Bench-only utility: rejected as public API in #2831 (transports shouldn't @@ -58,7 +60,7 @@ class CompressedCodec(PubSubTransport[T]): Subscribers always receive a decoded Image. """ - def __init__(self, inner: PubSubTransport[Any], quality: int = 75) -> None: + def __init__(self, inner: PubSubTransport[CompressedImage], quality: int = 75) -> None: super().__init__(inner.topic) self.inner = inner self.quality = quality @@ -66,21 +68,30 @@ def __init__(self, inner: PubSubTransport[Any], quality: int = 75) -> None: def __reduce__(self): # type: ignore[no-untyped-def] return (CompressedCodec, (self.inner, self.quality)) - def broadcast(self, stream: Out[T] | None, msg: T) -> None: - from dimos.msgs.sensor_msgs.CompressedImage import ( - CompressedImage, - ) # deferred to avoid pulling in cv2/rerun - - if not isinstance(msg, CompressedImage): - msg = CompressedImage.from_image(msg, quality=self.quality) # type: ignore[assignment, arg-type] - self.inner.broadcast(stream, msg) + def broadcast(self, stream: Stream[Image] | None, msg: Image | CompressedImage) -> None: + compressed = ( + msg + if isinstance(msg, CompressedImage) + else CompressedImage( + header=msg.header, format="jpeg", data=image_to_jpeg(msg, quality=self.quality) + ) + ) + self.inner.broadcast(None, compressed) def subscribe( - self, callback: Callable[[T], Any], selfstream: Stream[T] | None = None + self, callback: Callable[[Image], Any], selfstream: Stream[Image] | None = None ) -> Callable[[], None]: - return self.inner.subscribe( # type: ignore[no-any-return] - lambda m: callback(m.decode()), selfstream - ) + def decoded(message: CompressedImage) -> Any: + import cv2 + + pixels = cv2.imdecode( + np.frombuffer(message.data.view(), dtype=np.uint8), cv2.IMREAD_COLOR + ) + if pixels is None: + raise ValueError("JPEG benchmark received an invalid compressed image") + return callback(image_from_array(pixels, encoding="bgr8", header=message.header)) + + return self.inner.subscribe(decoded) def start(self) -> None: self.inner.start() @@ -116,13 +127,14 @@ def stop(self) -> None: def _on_image(self, msg: Image) -> None: now = time.time() self._count += 1 - rec = {"t": now, "ts": msg.ts, "age": now - msg.ts, "i": self._count} + ts = to_seconds(msg.header.stamp) + rec = {"t": now, "ts": ts, "age": now - ts, "i": self._count} if self._work_ms > 0: import cv2 deadline = time.perf_counter() + self._work_ms / 1000 while time.perf_counter() < deadline: - cv2.GaussianBlur(msg.data, (31, 31), 5) + cv2.GaussianBlur(image_view(msg), (31, 31), 5) rec["done_t"] = time.time() with self._lock: self._buf.append(json.dumps(rec)) @@ -202,7 +214,6 @@ def main() -> None: from dimos.core.coordination.blueprints import autoconnect from dimos.core.coordination.module_coordinator import ModuleCoordinator from dimos.core.transport import LCMTransport - from dimos.msgs.sensor_msgs.CompressedImage import CompressedImage from dimos.robot.get_all_blueprints import get_blueprint_by_name bp = get_blueprint_by_name(args.blueprint) @@ -220,7 +231,7 @@ def main() -> None: } ) - coordinator = ModuleCoordinator.build(bp, {"g": {"replay": True, "viewer": "none"}}) + coordinator = ModuleCoordinator.build(bp.global_config(replay=True, viewer="none")) logger.info("benchmark run started", mode=args.mode, blueprint=args.blueprint) stop = threading.Event() diff --git a/dimos/protocol/pubsub/encoders.py b/dimos/protocol/pubsub/encoders.py index ff5c78a2f52..3ef7845d39e 100644 --- a/dimos/protocol/pubsub/encoders.py +++ b/dimos/protocol/pubsub/encoders.py @@ -94,33 +94,29 @@ def decode(self, msg: bytes, _: TopicT) -> MsgT: return cast("MsgT", pickle.loads(msg)) -class LCMTopicProto(Protocol): - """Protocol for topics usable with LCM encoders.""" +class TypedTopicProto(Protocol): + """A topic's explicit generated message type, independent of its transport.""" topic: str # At decode time, always concrete string - lcm_type: type[DimosMsg] | None + msg_type: type[DimosMsg] | None -class LCMEncoderMixin(PubSubEncoderMixin[LCMTopicProto, DimosMsg, bytes]): - """Encoder mixin for DimosMsg using LCM binary encoding.""" +class CDREncoderMixin(PubSubEncoderMixin[TypedTopicProto, DimosMsg, bytes]): + """Encode typed streams as standard encapsulated CDR.""" - def subscribe( - self, topic: LCMTopicProto, callback: Callable[[DimosMsg, LCMTopicProto], None] - ) -> Callable[[], None]: - # Import the type's heavy decode dependencies now, on the subscriber's - # thread. Left to the first decode, the import would run on the LCM - # handler thread and block it for seconds, dropping messages meanwhile. - warmup = getattr(topic.lcm_type, "lcm_warmup", None) - if warmup is not None: - warmup() - return super().subscribe(topic, callback) - - def encode(self, msg: DimosMsg | bytes, _: LCMTopicProto) -> bytes: + def encode(self, msg: DimosMsg | bytes, topic: TypedTopicProto) -> bytes: if isinstance(msg, bytes): return msg - return msg.lcm_encode() - - def decode(self, msg: bytes, topic: LCMTopicProto) -> DimosMsg: - if topic.lcm_type is None: - raise DecodingError(f"Cannot decode: topic {topic.topic!r} has no lcm_type") - return topic.lcm_type.lcm_decode(msg) + if topic.msg_type is not None and msg.msg_name != topic.msg_type.msg_name: + raise ValueError( + f"Message {msg.msg_name} does not match topic type {topic.msg_type.msg_name}" + ) + return msg.encode() + + def decode(self, msg: bytes, topic: TypedTopicProto) -> DimosMsg: + if topic.msg_type is None: + raise DecodingError(f"Cannot decode: topic {topic.topic!r} has no msg_type") + try: + return topic.msg_type.decode(msg) + except ValueError as exc: + raise DecodingError(f"Invalid CDR on topic {topic.topic!r}") from exc diff --git a/dimos/protocol/pubsub/impl/jpeg_lcm.py b/dimos/protocol/pubsub/impl/jpeg_lcm.py index b87905584ee..bafb414a22f 100644 --- a/dimos/protocol/pubsub/impl/jpeg_lcm.py +++ b/dimos/protocol/pubsub/impl/jpeg_lcm.py @@ -21,25 +21,29 @@ from __future__ import annotations -from typing import cast +from dimos_generated.sensor_msgs.msg import CompressedImage, Image -from dimos.msgs.sensor_msgs.Image import Image -from dimos.protocol.pubsub.encoders import DecodingError, LCMTopicProto, PubSubEncoderMixin +from dimos.msgs.image import compressed_image_from_image, image_from_compressed +from dimos.protocol.pubsub.encoders import DecodingError, PubSubEncoderMixin, TypedTopicProto from dimos.protocol.pubsub.impl.lcmpubsub import LCMPubSubBase -class JpegEncoderMixin(PubSubEncoderMixin[LCMTopicProto, Image, bytes]): - """Encoder mixin for DimosMsg using JPEG encoding (for images).""" +class JpegEncoderMixin(PubSubEncoderMixin[TypedTopicProto, Image, bytes]): + """Transcode generated Image values using CompressedImage CDR on the wire.""" - def encode(self, msg: Image, _: LCMTopicProto) -> bytes: - return msg.lcm_jpeg_encode() + def encode(self, msg: Image, topic: TypedTopicProto) -> bytes: + if topic.msg_type is not CompressedImage: + raise ValueError("JPEG wire topics must declare generated CompressedImage") + return compressed_image_from_image(msg).encode() - def decode(self, msg: bytes, topic: LCMTopicProto) -> Image: + def decode(self, msg: bytes, topic: TypedTopicProto) -> Image: if topic.topic == "LCM_SELF_TEST": raise DecodingError("Ignoring LCM_SELF_TEST topic") - if topic.lcm_type is None: - raise DecodingError(f"Cannot decode: topic {topic.topic!r} has no lcm_type") - return cast("type[Image]", topic.lcm_type).lcm_jpeg_decode(msg) + if topic.msg_type is None: + raise DecodingError(f"Cannot decode: topic {topic.topic!r} has no msg_type") + if topic.msg_type is not CompressedImage: + raise DecodingError("JPEG wire topics must declare generated CompressedImage") + return image_from_compressed(CompressedImage.decode(msg)) class JpegLCM( # type: ignore[misc] diff --git a/dimos/protocol/pubsub/impl/lcmpubsub.py b/dimos/protocol/pubsub/impl/lcmpubsub.py index f47280f35cb..34f28d4ca03 100644 --- a/dimos/protocol/pubsub/impl/lcmpubsub.py +++ b/dimos/protocol/pubsub/impl/lcmpubsub.py @@ -20,9 +20,10 @@ import threading from typing import Any +from dimos.msgs.helpers import resolve_msg_type from dimos.msgs.protocol import DimosMsg from dimos.protocol.pubsub.encoders import ( - LCMEncoderMixin, + CDREncoderMixin, PickleEncoderMixin, ) from dimos.protocol.pubsub.patterns import Glob @@ -36,7 +37,7 @@ @dataclass class Topic: topic: str | re.Pattern[str] | Glob - lcm_type: type[DimosMsg] | None = None + msg_type: type[DimosMsg] | None = None @property def is_pattern(self) -> bool: @@ -51,25 +52,23 @@ def pattern(self) -> str: return self.topic def __str__(self) -> str: - if self.lcm_type is None: + if self.msg_type is None: return self.pattern - return f"{self.pattern}#{self.lcm_type.msg_name}" + return f"{self.pattern}#{self.msg_type.msg_name}" @staticmethod - def from_channel_str(channel: str, default_lcm_type: type[DimosMsg] | None = None) -> Topic: + def from_channel_str(channel: str, default_msg_type: type[DimosMsg] | None = None) -> Topic: """Create Topic from channel string. - Channel format: /topic#module.ClassName - Falls back to default_lcm_type if type cannot be parsed. + Channel format: /topic#package/msg/Type. + A supplied type applies only to channels without a type suffix. """ - from dimos.msgs.helpers import resolve_msg_type - if "#" not in channel: - return Topic(topic=channel, lcm_type=default_lcm_type) + return Topic(topic=channel, msg_type=default_msg_type) topic_str, type_name = channel.rsplit("#", 1) - lcm_type = resolve_msg_type(type_name) - return Topic(topic=topic_str, lcm_type=lcm_type or default_lcm_type) + msg_type = resolve_msg_type(type_name) + return Topic(topic=topic_str, msg_type=msg_type) class LCMPubSubBase(LCMService, AllPubSub[Topic, Any]): @@ -89,6 +88,10 @@ def publish(self, topic: Topic | str, message: bytes) -> None: return topic_str = str(topic) if isinstance(topic, Topic) else topic + if "\0" in topic_str or len(topic_str.encode("utf-8")) > 63: + raise ValueError( + "LCM channel (including type suffix) must fit in 63 bytes and contain no NUL" + ) self.l.publish(topic_str, message) def subscribe_all(self, callback: Callable[[bytes, Topic], Any]) -> Callable[[], None]: @@ -115,11 +118,11 @@ def noop() -> None: def handler(channel: str, msg: bytes) -> None: if not alive or channel == "LCM_SELF_TEST": return - callback(msg, Topic.from_channel_str(channel, topic.lcm_type)) + callback(msg, Topic.from_channel_str(channel, topic.msg_type)) pattern_str = str(topic) - if not pattern_str.endswith("*"): - pattern_str = f"{pattern_str}(#.*)?" + if topic.msg_type is None: + pattern_str = f"(?:{pattern_str})(#.*)?" lcm_subscription = self.l.subscribe(pattern_str, handler) else: @@ -149,7 +152,7 @@ def unsubscribe() -> None: class LCM( # type: ignore[misc] - LCMEncoderMixin, + CDREncoderMixin, LCMPubSubBase, ): ... diff --git a/dimos/protocol/pubsub/impl/rospubsub_conversion.py b/dimos/protocol/pubsub/impl/rospubsub_conversion.py index 6853c57ca91..16438801eda 100644 --- a/dimos/protocol/pubsub/impl/rospubsub_conversion.py +++ b/dimos/protocol/pubsub/impl/rospubsub_conversion.py @@ -12,353 +12,48 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Conversion functions between dimos messages and ROS messages. - -This module provides conversion functions between dimos message types and ROS messages. -It handles three categories of types: - -1. Complex types (different internal representation) - use LCM roundtrip -2. Simple types (field structures match) - use direct field copy -3. No dimos.msgs equivalent - return dimos_lcm type -""" +"""Convert generated DimOS and ROS2 messages through their shared CDR format.""" from __future__ import annotations import importlib -import re -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, TypeVar, cast -from dimos.msgs.helpers import lcm_msg_type +try: + import rclpy.serialization as ros_serialization +except ImportError: + ros_serialization = None # type: ignore[assignment] if TYPE_CHECKING: from dimos.msgs.protocol import DimosMsg from dimos.protocol.pubsub.impl.rospubsub import ROSMessage -# Complex types that need LCM roundtrip (explicit list) -# These types have different internal representations in dimos vs ROS/LCM -COMPLEX_TYPES: set[str] = { - "sensor_msgs.PointCloud2", - "sensor_msgs.Image", - "sensor_msgs.CompressedImage", - "sensor_msgs.CameraInfo", - "geometry_msgs.PoseStamped", -} - -# Cache for dynamic imports of dimos types -_dimos_type_cache: dict[str, type[DimosMsg] | None] = {} - -# Cache for LCM type derivation -_lcm_type_cache: dict[str, type[Any]] = {} - -# Field name mappings between ROS and LCM (ROS name -> LCM name) -# This is some mixup in dimos_lcm having ROS1 and ROS2 message definitions? -# Would be good to clarify later, but this works for now -_ROS_TO_LCM_FIELD_MAP: dict[str, str] = { - "nanosec": "nsec", # ROS2 Time.nanosec -> LCM Time.nsec -} - -# Reverse mapping (LCM name -> ROS name) -_LCM_TO_ROS_FIELD_MAP: dict[str, str] = {v: k for k, v in _ROS_TO_LCM_FIELD_MAP.items()} - - -def get_dimos_type(msg_name: str) -> type[DimosMsg] | None: - """Try to import dimos.msgs type, return None if not found. Cached. - - Args: - msg_name: Message name in format "package.MessageName" (e.g., "geometry_msgs.Vector3") - - Returns: - The dimos message type, or None if not found - """ - if msg_name in _dimos_type_cache: - return _dimos_type_cache[msg_name] - - try: - package, name = msg_name.split(".") - module = importlib.import_module(f"dimos.msgs.{package}.{name}") - dimos_type = cast("type[DimosMsg]", getattr(module, name)) - _dimos_type_cache[msg_name] = dimos_type - return dimos_type - except (ImportError, AttributeError, ValueError): - _dimos_type_cache[msg_name] = None - return None - - -def derive_lcm_type(dimos_type: type[DimosMsg]) -> type[Any]: - """Derive the LCM message type from a dimos message type. - - Args: - dimos_type: A dimos message type (e.g., dimos.msgs.sensor_msgs.PointCloud2) - - Returns: - The corresponding LCM message type (e.g., dimos_lcm.sensor_msgs.PointCloud2) - """ - msg_name = dimos_type.msg_name # e.g., "sensor_msgs.PointCloud2" - - if msg_name in _lcm_type_cache: - return _lcm_type_cache[msg_name] - - parts = msg_name.split(".") - if len(parts) != 2: - raise ValueError(f"Invalid msg_name format: {msg_name}, expected 'package.MessageName'") - - lcm_type = lcm_msg_type(msg_name) - _lcm_type_cache[msg_name] = lcm_type - return lcm_type +MessageT = TypeVar("MessageT", bound="DimosMsg") def derive_ros_type(dimos_type: type[DimosMsg]) -> type[ROSMessage]: - """Derive the ROS message type from a dimos message type. - - Args: - dimos_type: A dimos message type (e.g., dimos.msgs.geometry_msgs.Vector3) - - Returns: - The corresponding ROS message type (e.g., geometry_msgs.msg.Vector3) - - Example: - msg_name = "geometry_msgs.Vector3" -> geometry_msgs.msg.Vector3 - """ - msg_name = dimos_type.msg_name # e.g., "geometry_msgs.Vector3" - parts = msg_name.split(".") - if len(parts) != 2: - raise ValueError(f"Invalid msg_name format: {msg_name}, expected 'package.MessageName'") - - package, message_name = parts - ros_module = importlib.import_module(f"{package}.msg") - return cast("type[ROSMessage]", getattr(ros_module, message_name)) - - -def _copy_ros_to_lcm_recursive(ros_msg: Any, lcm_msg: Any) -> None: - """Recursively copy fields from ROS message to LCM message. - - Handles nested messages, arrays, and primitive types. - - Args: - ros_msg: Source ROS message - lcm_msg: Target LCM message (modified in place) - """ - if not hasattr(ros_msg, "get_fields_and_field_types"): - raise TypeError(f"Expected ROS message, got {type(ros_msg).__name__}") - - field_types = ros_msg.get_fields_and_field_types() - for ros_field_name in field_types: - # Map ROS field name to LCM field name - lcm_field_name = _ROS_TO_LCM_FIELD_MAP.get(ros_field_name, ros_field_name) - - if not hasattr(lcm_msg, lcm_field_name): - continue - - ros_value = getattr(ros_msg, ros_field_name) - lcm_value = getattr(lcm_msg, lcm_field_name) - - # Handle nested messages - if hasattr(ros_value, "get_fields_and_field_types"): - _copy_ros_to_lcm_recursive(ros_value, lcm_value) - # Handle arrays of messages - elif isinstance(ros_value, (list, tuple)) and len(ros_value) > 0: - if hasattr(ros_value[0], "get_fields_and_field_types"): - # Array of nested messages - create LCM instances - lcm_array = [] - for ros_item in ros_value: - # Get the LCM element type from the first lcm_value element if available - # Otherwise try to derive from ros item - if isinstance(lcm_value, list) and len(lcm_value) > 0: - lcm_item = type(lcm_value[0])() - else: - # Try to create matching LCM type - lcm_item = _create_lcm_instance_for_ros_msg(ros_item) - _copy_ros_to_lcm_recursive(ros_item, lcm_item) - lcm_array.append(lcm_item) - setattr(lcm_msg, lcm_field_name, lcm_array) - else: - # Array of primitives - direct copy - setattr(lcm_msg, lcm_field_name, list(ros_value)) - # Handle bytes/data fields - elif isinstance(ros_value, (bytes, bytearray)): - setattr(lcm_msg, lcm_field_name, bytes(ros_value)) - # Handle array.array (ROS uses this for data fields) - elif hasattr(ros_value, "tobytes"): - setattr(lcm_msg, lcm_field_name, ros_value.tobytes()) - else: - # Primitive type - direct copy - setattr(lcm_msg, lcm_field_name, ros_value) - - # Update length fields if present (LCM convention: field_name_length) - length_field = f"{lcm_field_name}_length" - if hasattr(lcm_msg, length_field): - value = getattr(lcm_msg, lcm_field_name) - if isinstance(value, (list, tuple, bytes, bytearray)): - setattr(lcm_msg, length_field, len(value)) - - -def _copy_lcm_to_ros_recursive(lcm_msg: Any, ros_msg: Any) -> None: - """Recursively copy fields from LCM message to ROS message. - - Handles nested messages, arrays, and primitive types. - - Args: - lcm_msg: Source LCM message - ros_msg: Target ROS message (modified in place) - """ - if not hasattr(ros_msg, "get_fields_and_field_types"): - raise TypeError(f"Expected ROS message, got {type(ros_msg).__name__}") - - field_types = ros_msg.get_fields_and_field_types() - for ros_field_name in field_types: - # Map ROS field name to LCM field name - lcm_field_name = _ROS_TO_LCM_FIELD_MAP.get(ros_field_name, ros_field_name) - - if not hasattr(lcm_msg, lcm_field_name): - continue - - lcm_value = getattr(lcm_msg, lcm_field_name) - ros_value = getattr(ros_msg, ros_field_name) - - # Handle nested messages - if hasattr(ros_value, "get_fields_and_field_types"): - _copy_lcm_to_ros_recursive(lcm_value, ros_value) - # Handle arrays of messages - elif isinstance(lcm_value, (list, tuple)) and len(lcm_value) > 0: - if hasattr(lcm_value[0], "lcm_encode"): - # Array of nested LCM messages - ros_array = [] - for lcm_item in lcm_value: - ros_item = _create_ros_instance_for_lcm_msg( - lcm_item, field_types[ros_field_name] - ) - _copy_lcm_to_ros_recursive(lcm_item, ros_item) - ros_array.append(ros_item) - setattr(ros_msg, ros_field_name, ros_array) - else: - # Array of primitives - direct copy - setattr(ros_msg, ros_field_name, list(lcm_value)) - # Handle bytes/data fields - elif isinstance(lcm_value, (bytes, bytearray)): - # ROS data fields might expect array.array - if hasattr(ros_value, "frombytes"): - import array - - arr = array.array("B") - arr.frombytes(lcm_value) - setattr(ros_msg, ros_field_name, arr) - else: - setattr(ros_msg, ros_field_name, bytes(lcm_value)) - else: - # Primitive type - direct copy - setattr(ros_msg, ros_field_name, lcm_value) - - -def _create_lcm_instance_for_ros_msg(ros_msg: Any) -> Any: - """Create an LCM message instance that matches the ROS message type. - - Args: - ros_msg: ROS message to match - - Returns: - New LCM message instance - """ - # Get the ROS type name (e.g., "std_msgs.msg.Header" -> "std_msgs.Header") - ros_type = type(ros_msg) - module_name = ros_type.__module__ # e.g., "std_msgs.msg" - class_name = ros_type.__name__ # e.g., "Header" - - # std_msgs.msg.Header -> dimos_lcm.std_msgs.Header - package = module_name.split(".")[0] # e.g., "std_msgs" - return lcm_msg_type(f"{package}.{class_name}")() - - -def _create_ros_instance_for_lcm_msg(lcm_msg: Any, ros_type_hint: str) -> Any: - """Create a ROS message instance that matches the LCM message type. - - Args: - lcm_msg: LCM message to match - ros_type_hint: ROS type hint string (e.g., "sequence") - - Returns: - New ROS message instance - """ - # Parse the type hint to get the message type - # e.g., "sequence" -> "sensor_msgs", "PointField" - # e.g., "sensor_msgs/PointField" -> "sensor_msgs", "PointField" - - match = re.search(r"(\w+)/(\w+)", ros_type_hint) - if match: - package, class_name = match.groups() - ros_module = importlib.import_module(f"{package}.msg") - ros_type = getattr(ros_module, class_name) - return ros_type() - - # Fallback: try to derive from LCM type - lcm_type = type(lcm_msg) - module_name = lcm_type.__module__ # e.g., "dimos_lcm.std_msgs.Header" - class_name = lcm_type.__name__ - parts = module_name.split(".") - if len(parts) >= 2: - package = parts[1] # e.g., "std_msgs" - ros_module = importlib.import_module(f"{package}.msg") - ros_type = getattr(ros_module, class_name) - return ros_type() - - raise ValueError(f"Cannot determine ROS type for LCM message: {lcm_type}") + """Resolve the matching installed ROS type by canonical package/msg/Type name.""" + parts = dimos_type.msg_name.split("/") + if len(parts) != 3 or parts[1] != "msg" or not all(part.isidentifier() for part in parts): + raise ValueError(f"Invalid message name {dimos_type.msg_name!r}; expected package/msg/Type") + package, _, name = parts + return cast("type[ROSMessage]", getattr(importlib.import_module(f"{package}.msg"), name)) def dimos_to_ros(msg: DimosMsg, ros_type: type[ROSMessage]) -> ROSMessage: - """Convert a dimos message to a ROS message. - - For complex types (PointCloud2, Image, CameraInfo), uses LCM roundtrip - to properly convert internal representations. For simple types, uses - direct field copy. - - Args: - msg: Dimos message instance - ros_type: Target ROS message type - - Returns: - ROS message instance - """ - msg_name = type(msg).msg_name - - if msg_name in COMPLEX_TYPES: - # Complex: dimos → encode → decode LCM → copy to ROS - lcm_type = derive_lcm_type(type(msg)) - lcm_bytes = msg.lcm_encode() - lcm_msg = lcm_type.lcm_decode(lcm_bytes) - ros_msg = ros_type() - _copy_lcm_to_ros_recursive(lcm_msg, ros_msg) - return ros_msg - - # Simple: recursive field copy (handles nested messages) - ros_msg = ros_type() - _copy_lcm_to_ros_recursive(msg, ros_msg) - return ros_msg - - -def ros_to_dimos(msg: Any, dimos_type: type[DimosMsg]) -> DimosMsg: - """Convert a ROS message to a dimos message. - - For complex types (PointCloud2, Image, CameraInfo), uses LCM roundtrip - to properly build the dimos internal representation. For simple types, - uses direct field copy. - - Args: - msg: ROS message instance - dimos_type: Target dimos message type - - Returns: - Dimos message instance - """ - msg_name = dimos_type.msg_name - - if msg_name in COMPLEX_TYPES: - # Complex: ROS → LCM → encode → decode → dimos - lcm_type = derive_lcm_type(dimos_type) - lcm_msg = lcm_type() - _copy_ros_to_lcm_recursive(msg, lcm_msg) - return dimos_type.lcm_decode(lcm_msg.lcm_encode()) - - # Simple type: recursive field copy (handles nested messages) - dimos_msg = dimos_type() - _copy_ros_to_lcm_recursive(msg, dimos_msg) - return dimos_msg + """Deserialize generated CDR with ROS's own type support, preserving every field.""" + if ros_serialization is None: + raise ImportError("ROS message conversion requires rclpy; install and source ROS 2.") + if derive_ros_type(type(msg)) is not ros_type: + raise TypeError(f"ROS target does not match {msg.msg_name}") + return cast("ROSMessage", ros_serialization.deserialize_message(msg.encode(), ros_type)) + + +def ros_to_dimos(msg: Any, dimos_type: type[MessageT]) -> MessageT: + """Decode ROS CDR with the generated codec; no LCM or field-copy adapter is needed.""" + if ros_serialization is None: + raise ImportError("ROS message conversion requires rclpy; install and source ROS 2.") + if type(msg) is not derive_ros_type(dimos_type): + raise TypeError(f"ROS source does not match {dimos_type.msg_name}") + return dimos_type.decode(ros_serialization.serialize_message(msg)) diff --git a/dimos/protocol/pubsub/impl/test_jpeg_cdr.py b/dimos/protocol/pubsub/impl/test_jpeg_cdr.py new file mode 100644 index 00000000000..b1827a2abf9 --- /dev/null +++ b/dimos/protocol/pubsub/impl/test_jpeg_cdr.py @@ -0,0 +1,76 @@ +# 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. + +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +import pickle +from types import SimpleNamespace + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.sensor_msgs.msg import CompressedImage, Image +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest + +from dimos.core.transport import JpegLcmTransport +from dimos.msgs.image import image_from_array, image_from_compressed, image_view +from dimos.protocol.pubsub.encoders import DecodingError +from dimos.protocol.pubsub.impl.jpeg_lcm import JpegEncoderMixin +from dimos.protocol.pubsub.impl.jpeg_shm import JpegSharedMemoryEncoderMixin + + +@pytest.mark.parametrize("codec", [JpegEncoderMixin(), JpegSharedMemoryEncoderMixin(quality=90)]) +def test_jpeg_wire_is_standard_cdr_and_preserves_exact_image_header(codec): + header = Header(frame_id="optical", stamp=Time(sec=-1, nanosec=987654321)) + image = image_from_array( + np.full((16, 16, 3), [20, 80, 140], dtype=np.uint8), encoding="rgb8", header=header + ) + topic = SimpleNamespace(topic="/image", msg_type=CompressedImage) + before = image.encode() + encoded = codec.encode(image, topic) + wire = CompressedImage.decode(encoded) + assert encoded[:4] == b"\x00\x01\x00\x00" + assert wire.header == header + assert wire.format == "rgb8; jpeg compressed bgr8" + decoded = codec.decode(encoded, topic) + assert decoded.encoding == "bgr8" and decoded.header == header + np.testing.assert_allclose(image_view(decoded)[0, 0], [140, 80, 20], atol=3) + assert image.encode() == before + + +def test_jpeg_lcm_requires_truthful_compressed_wire_type(): + codec = JpegEncoderMixin() + with pytest.raises(DecodingError, match="CompressedImage"): + codec.decode(b"", SimpleNamespace(topic="/image", msg_type=Image)) + with pytest.raises(DecodingError, match="SELF_TEST"): + codec.decode(b"", SimpleNamespace(topic="LCM_SELF_TEST", msg_type=CompressedImage)) + + +def test_jpeg_transport_wire_metadata_and_pickle_keep_separate_application_type(monkeypatch): + monkeypatch.setattr( + "dimos.protocol.pubsub.impl.jpeg_lcm.JpegLCM", lambda **kwargs: SimpleNamespace() + ) + transport = JpegLcmTransport("/image", Image) + assert transport.topic.msg_type is CompressedImage + restored = pickle.loads(pickle.dumps(transport)) + assert restored.image_type is Image + assert restored.topic.msg_type is CompressedImage + + +def test_compressed_image_rejects_invalid_data_and_unknown_format(): + with pytest.raises(ValueError, match="invalid compressed"): + image_from_compressed(CompressedImage(format="jpeg", data=b"not jpeg")) + with pytest.raises(ValueError, match="unsupported compressed"): + image_from_compressed(CompressedImage(format="h264", data=b"not jpeg")) diff --git a/dimos/protocol/pubsub/impl/test_lcmpubsub.py b/dimos/protocol/pubsub/impl/test_lcmpubsub.py index 4ab7030b107..92a827d0289 100644 --- a/dimos/protocol/pubsub/impl/test_lcmpubsub.py +++ b/dimos/protocol/pubsub/impl/test_lcmpubsub.py @@ -13,14 +13,12 @@ # limitations under the License. from collections.abc import Iterator -import threading from typing import Any +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion, Vector3 +from dimos_generated.std_msgs.msg import String import pytest -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.protocol.pubsub.impl.lcmpubsub import ( LCM, LCMPubSubBase, @@ -54,34 +52,15 @@ def lcm(lcm_url: str) -> Iterator[LCM]: lcm.stop() -class MockLCMMessage: - """Mock LCM message for testing""" - - msg_name = "geometry_msgs.Mock" - - def __init__(self, data: Any) -> None: - self.data = data - - def lcm_encode(self) -> bytes: - return str(self.data).encode("utf-8") - - @classmethod - def lcm_decode(cls, data: bytes) -> "MockLCMMessage": - return cls(data.decode("utf-8")) - - def __eq__(self, other: object) -> bool: - return isinstance(other, MockLCMMessage) and self.data == other.data - - def test_LCMPubSubBase_pubsub(lcm_pub_sub_base: LCMPubSubBase) -> None: lcm = lcm_pub_sub_base collector = CallbackCollector(1) - topic = Topic(topic="/test_topic", lcm_type=MockLCMMessage) - test_message = MockLCMMessage("test_data") + topic = Topic(topic="/test_topic", msg_type=String) + test_message = String(data="test_data") lcm.subscribe(topic, collector) - lcm.publish(topic, test_message.lcm_encode()) + lcm.publish(topic, test_message.encode()) collector.wait() assert len(collector.results) == 1 @@ -90,31 +69,17 @@ def test_LCMPubSubBase_pubsub(lcm_pub_sub_base: LCMPubSubBase) -> None: received_topic = collector.results[0][1] assert isinstance(received_data, bytes) - assert received_data.decode() == "test_data" + assert String.decode(received_data).data == "test_data" assert isinstance(received_topic, Topic) assert received_topic == topic -def test_subscribe_calls_lcm_warmup(lcm: LCM) -> None: - warmup_threads = [] - - class WarmupMessage(MockLCMMessage): - @classmethod - def lcm_warmup(cls) -> None: - warmup_threads.append(threading.current_thread()) - - lcm.subscribe(Topic(topic="/warmup", lcm_type=WarmupMessage), lambda msg, topic: None) - - # Warmed synchronously on the subscriber's thread, not the LCM handler thread. - assert warmup_threads == [threading.current_thread()] - - def test_lcm_autodecoder_pubsub(lcm: LCM) -> None: collector = CallbackCollector(1) - topic = Topic(topic="/test_topic", lcm_type=MockLCMMessage) - test_message = MockLCMMessage("test_data") + topic = Topic(topic="/test_topic", msg_type=String) + test_message = String(data="test_data") lcm.subscribe(topic, collector) lcm.publish(topic, test_message) @@ -125,17 +90,47 @@ def test_lcm_autodecoder_pubsub(lcm: LCM) -> None: received_data = collector.results[0][0] received_topic = collector.results[0][1] - assert isinstance(received_data, MockLCMMessage) - assert received_data == test_message + assert isinstance(received_data, String) + assert received_data.encode() == test_message.encode() assert isinstance(received_topic, Topic) assert received_topic == topic +def test_invalid_cdr_does_not_stop_the_subscriber(lcm: LCM) -> None: + topic = Topic("/test_corrupt", String) + collector = CallbackCollector(1) + lcm.subscribe(topic, collector) + + lcm.publish(topic, b"invalid CDR") + lcm.publish(topic, String(data="after malformed payload")) + collector.wait() + + assert [message.data for message, _ in collector.results] == ["after malformed payload"] + + +def test_different_message_types_cannot_be_published_on_a_typed_topic(lcm: LCM) -> None: + with pytest.raises(ValueError, match="does not match"): + lcm.publish(Topic("/point", Vector3), Quaternion(w=1)) + + +def test_explicit_unknown_type_does_not_use_default_decoder() -> None: + topic = Topic.from_channel_str("/point#unknown_msgs/msg/Point", Vector3) + + assert topic.topic == "/point" + assert topic.msg_type is None + + +@pytest.mark.parametrize("channel", ["x" * 64, "invalid\0channel", "é" * 32]) +def test_invalid_lcm_channel_fails_before_sending(lcm: LCM, channel: str) -> None: + with pytest.raises(ValueError, match="63 bytes"): + lcm.publish(Topic(channel), b"payload") + + test_msgs = [ - (Vector3(1, 2, 3)), - (Quaternion(1, 2, 3, 4)), - (Pose(Vector3(1, 2, 3), Quaternion(0, 0, 0, 1))), + (Vector3(x=1, y=2, z=3)), + (Quaternion(x=1, y=2, z=3, w=4)), + (Pose(position=Point(x=1, y=2, z=3), orientation=Quaternion(w=1))), ] @@ -144,7 +139,7 @@ def test_lcm_autodecoder_pubsub(lcm: LCM) -> None: def test_lcm_geometry_msgs_pubsub(test_message: Any, lcm: LCM) -> None: collector = CallbackCollector(1) - topic = Topic(topic="/test_topic", lcm_type=test_message.__class__) + topic = Topic(topic="/test_topic", msg_type=test_message.__class__) lcm.subscribe(topic, collector) lcm.publish(topic, test_message) @@ -156,7 +151,7 @@ def test_lcm_geometry_msgs_pubsub(test_message: Any, lcm: LCM) -> None: received_topic = collector.results[0][1] assert isinstance(received_data, test_message.__class__) - assert received_data == test_message + assert received_data.encode() == test_message.encode() assert isinstance(received_topic, Topic) assert received_topic == topic @@ -180,7 +175,7 @@ def test_lcm_geometry_msgs_autopickle_pubsub(test_message: Any, pickle_lcm: Pick received_topic = collector.results[0][1] assert isinstance(received_data, test_message.__class__) - assert received_data == test_message + assert received_data.encode() == test_message.encode() assert isinstance(received_topic, Topic) assert received_topic == topic diff --git a/dimos/protocol/pubsub/impl/test_rospubsub.py b/dimos/protocol/pubsub/impl/test_rospubsub.py index 89be14cd296..48878767833 100644 --- a/dimos/protocol/pubsub/impl/test_rospubsub.py +++ b/dimos/protocol/pubsub/impl/test_rospubsub.py @@ -12,237 +12,73 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Generated CDR streams through the optional ROS2 pub/sub bridge.""" + from collections.abc import Generator +import threading +import uuid -from dimos_lcm.geometry_msgs import PointStamped -import numpy as np +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, PointStamped, PoseStamped, Twist, Vector3 +from dimos_generated.sensor_msgs.msg import Image, PointCloud2, PointField +from dimos_generated.std_msgs.msg import Header import pytest -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.protocol.pubsub.impl.rospubsub import DimosROS, ROSTopic +from dimos.protocol.pubsub.impl.rospubsub import ROS_AVAILABLE, DimosROS, ROSTopic -# Add msg_name to LCM PointStamped for testing nested message conversion -PointStamped.msg_name = "geometry_msgs.PointStamped" -from dimos.utils.testing.collector import CallbackCollector -from dimos.utils.testing.replay import TimedSensorReplay +pytestmark = pytest.mark.skipif(not ROS_AVAILABLE, reason="requires ROS 2") -def ros_node(): - ros = DimosROS() - ros.start() +@pytest.fixture() +def nodes() -> Generator[tuple[DimosROS, DimosROS], None, None]: + publisher, subscriber = DimosROS(), DimosROS() try: - yield ros + publisher.start() + subscriber.start() + yield publisher, subscriber finally: - ros.stop() - - -@pytest.fixture() -def publisher() -> Generator[DimosROS, None, None]: - yield from ros_node() - - -@pytest.fixture() -def subscriber() -> Generator[DimosROS, None, None]: - yield from ros_node() - - -@pytest.mark.skipif_no_ros -def test_basic_conversion(publisher, subscriber): - """Test Vector3 publish/subscribe through ROS. - - Simple flat dimos.msgs type with no nesting (just x/y/z floats). - """ - topic = ROSTopic("/test_ros_topic", Vector3) - collector = CallbackCollector(1) - - subscriber.subscribe(topic, collector) - publisher.publish(topic, Vector3(1.0, 2.0, 3.0)) - - collector.wait() - assert len(collector.results) == 1 - msg = collector.results[0][0] - assert msg.x == 1.0 - assert msg.y == 2.0 - assert msg.z == 3.0 - - -@pytest.mark.skipif_no_ros -@pytest.mark.self_hosted -def test_pointcloud2_pubsub(publisher, subscriber): - """Test PointCloud2 publish/subscribe through ROS. - - COMPLEX_TYPE - has non-standard attributes (numpy arrays, custom accessors) - that can't be treated like a standard message with direct field copy. - Uses LCM encode/decode roundtrip to properly convert internal representation. - """ - # Load real lidar data from replay (5 seconds in) - replay = TimedSensorReplay("go2_bigoffice/lidar") - original = replay.find_closest_seek(5.0) - - assert original is not None, "Failed to load lidar data from replay" - assert len(original) > 0, "Loaded empty pointcloud" - - topic = ROSTopic("/test_pointcloud2", PointCloud2) - collector = CallbackCollector(1, timeout=5.0) - - subscriber.subscribe(topic, collector) - publisher.publish(topic, original) - - collector.wait() - assert len(collector.results) == 1 - - converted = collector.results[0][0] - - # Verify point cloud data is preserved - original_points, _ = original.as_numpy() - converted_points, _ = converted.as_numpy() - - assert len(original_points) == len(converted_points), ( - f"Point count mismatch: {len(original_points)} vs {len(converted_points)}" - ) - - np.testing.assert_allclose( - original_points, - converted_points, - rtol=1e-5, - atol=1e-5, - err_msg="Points don't match after ROS pubsub roundtrip", - ) - - # Verify frame_id is preserved - assert converted.frame_id == original.frame_id - - # Verify timestamp is preserved (within 1ms tolerance) - assert abs(original.ts - converted.ts) < 0.001 - - -@pytest.mark.skipif_no_ros -def test_pointcloud2_empty_pubsub(publisher, subscriber): - """Test empty PointCloud2 publish/subscribe. - - Edge case for COMPLEX_TYPE with zero points. - """ - original = PointCloud2.from_numpy( - np.array([]).reshape(0, 3), - frame_id="empty_frame", - timestamp=1234567890.0, - ) - - topic = ROSTopic("/test_empty_pointcloud", PointCloud2) - collector = CallbackCollector(1) - - subscriber.subscribe(topic, collector) - publisher.publish(topic, original) - - collector.wait() - assert len(collector.results) == 1 - assert len(collector.results[0][0]) == 0 - - -@pytest.mark.skipif_no_ros -def test_posestamped_pubsub(publisher, subscriber): - """Test PoseStamped publish/subscribe through ROS. - - COMPLEX_TYPE with custom dimos.msgs implementation and nested messages - (Header, Pose containing Point and Quaternion). Uses LCM roundtrip. - """ - original = PoseStamped( - ts=1234567890.123456, - frame_id="base_link", - position=[1.0, 2.0, 3.0], - orientation=[0.0, 0.0, 0.7071068, 0.7071068], # 90 degree yaw - ) - - topic = ROSTopic("/test_posestamped", PoseStamped) - collector = CallbackCollector(1) - - subscriber.subscribe(topic, collector) - publisher.publish(topic, original) - - collector.wait() - assert len(collector.results) == 1 - - converted = collector.results[0][0] - - # Verify all fields preserved - assert converted.frame_id == original.frame_id - assert abs(converted.ts - original.ts) < 0.001 # 1ms tolerance - assert converted.x == original.x - assert converted.y == original.y - assert converted.z == original.z - np.testing.assert_allclose(converted.orientation.z, original.orientation.z, rtol=1e-5) - np.testing.assert_allclose(converted.orientation.w, original.orientation.w, rtol=1e-5) - - -@pytest.mark.skipif_no_ros -def test_pointstamped_pubsub(publisher, subscriber): - """Test PointStamped publish/subscribe through ROS. - - Raw LCM type with nested messages (Header, Point) but NO custom dimos.msgs - implementation. Tests recursive field copy for non-COMPLEX_TYPES. - """ - original = PointStamped() - original.header.stamp.sec = 1234567890 - original.header.stamp.nsec = 123456000 - original.header.frame_id = "map" - original.point.x = 1.5 - original.point.y = 2.5 - original.point.z = 3.5 - - topic = ROSTopic("/test_pointstamped", PointStamped) - collector = CallbackCollector(1) - - subscriber.subscribe(topic, collector) - publisher.publish(topic, original) - - collector.wait() - assert len(collector.results) == 1 - - converted = collector.results[0][0] - - # Verify nested header fields are preserved - assert converted.header.frame_id == original.header.frame_id - assert converted.header.stamp.sec == original.header.stamp.sec - assert converted.header.stamp.nsec == original.header.stamp.nsec - - # Verify point coordinates are preserved - assert converted.point.x == original.point.x - assert converted.point.y == original.point.y - assert converted.point.z == original.point.z - - -@pytest.mark.skipif_no_ros -def test_twist_pubsub(publisher, subscriber): - """Test Twist publish/subscribe through ROS. - - dimos.msgs type with nested Vector3 messages (linear, angular). - Tests recursive field copy with custom dimos.msgs nested types. - """ - original = Twist( - linear=[1.0, 2.0, 3.0], - angular=[0.1, 0.2, 0.3], - ) - - topic = ROSTopic("/test_twist", Twist) - collector = CallbackCollector(1) - - subscriber.subscribe(topic, collector) - publisher.publish(topic, original) - - collector.wait() - assert len(collector.results) == 1 - - converted = collector.results[0][0] - - # Verify linear velocity preserved - assert converted.linear.x == original.linear.x - assert converted.linear.y == original.linear.y - assert converted.linear.z == original.linear.z - - # Verify angular velocity preserved - assert converted.angular.x == original.angular.x - assert converted.angular.y == original.angular.y - assert converted.angular.z == original.angular.z + subscriber.stop() + publisher.stop() + + +@pytest.mark.parametrize( + "original", + [ + Vector3(x=1.0, y=2.0, z=3.0), + Twist(linear=Vector3(x=1.0), angular=Vector3(z=-0.5)), + PoseStamped(header=Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id="map")), + PointStamped(header=Header(stamp=Time(sec=-1, nanosec=999999999)), point=Point(x=1.5)), + PointCloud2(), + PointCloud2( + height=1, + width=1, + fields=[PointField(name="x", datatype=7, count=1)], + point_step=4, + row_step=4, + data=bytes([0, 0, 128, 63]), + ), + Image(height=1, width=2, encoding="rgb8", step=6, data=bytes([0, 1, 2, 253, 254, 255])), + ], + ids=lambda msg: msg.msg_name, +) +def test_generated_pubsub(nodes, original): + publisher, subscriber = nodes + topic = ROSTopic(f"/dimos_cdr_test_{uuid.uuid4().hex}", type(original)) + received = threading.Event() + values = [] + + def collect(message, _topic): + values.append(message) + received.set() + + unsubscribe = subscriber.subscribe(topic, collect) + node = publisher._raw._node + assert node is not None + # A live periodic source tolerates asynchronous DDS discovery without a fixed sleep. + timer = node.create_timer(0.05, lambda: publisher.publish(topic, original)) + try: + assert received.wait(10), "ROS subscription did not receive the generated message" + assert values[0] == original + finally: + node.destroy_timer(timer) + unsubscribe() diff --git a/dimos/protocol/pubsub/impl/test_rospubsub_conversion.py b/dimos/protocol/pubsub/impl/test_rospubsub_conversion.py new file mode 100644 index 00000000000..c9d025e947f --- /dev/null +++ b/dimos/protocol/pubsub/impl/test_rospubsub_conversion.py @@ -0,0 +1,133 @@ +# Copyright 2025-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. + +"""Conformance with installed ROS type support; skipped when ROS is absent.""" + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + Point, + PointStamped, + Pose, + PoseStamped, + Quaternion, + TransformStamped, + Vector3, +) +from dimos_generated.sensor_msgs.msg import ( + CameraInfo, + CompressedImage, + Image, + PointCloud2, + PointField, +) +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage +import pytest + +from dimos.protocol.pubsub.impl import rospubsub_conversion as conversion + +requires_ros = pytest.mark.skipif(conversion.ros_serialization is None, reason="requires ROS 2") + + +def messages(): + header = Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id="map/相機") + return [ + Point(x=1.25, y=-2.5, z=3.75), + PointStamped(header=header, point=Point(x=2.0)), + PoseStamped(header=header, pose=Pose(orientation=Quaternion(z=0.6, w=0.8))), + Image( + header=header, + height=1, + width=2, + encoding="rgb8", + step=6, + data=bytes([0, 1, 2, 253, 254, 255]), + ), + CompressedImage(header=header, format="jpeg", data=bytes([0, 255, 127])), + CameraInfo( + header=header, + width=640, + height=480, + distortion_model="plumb_bob", + d=[0.1, -0.2], + k=[1.0, 0.0, 3.0, 0.0, 2.0, 4.0, 0.0, 0.0, 1.0], + ), + PointCloud2( + header=header, + height=1, + width=1, + fields=[PointField(name="x", offset=4, datatype=7, count=1)], + is_bigendian=True, + point_step=8, + row_step=12, + data=bytes([9, 8, 7, 6, 0x3F, 0x80, 0, 0, 5, 4, 3, 2]), + ), + PointCloud2(header=header), + TFMessage(), + TFMessage(transforms=[TransformStamped(header=header, child_frame_id="camera")]), + ] + + +@requires_ros +@pytest.mark.parametrize("original", messages(), ids=lambda msg: msg.msg_name) +def test_generated_fields_survive_ros_serialization(original): + ros_type = conversion.derive_ros_type(type(original)) + ros_message = conversion.dimos_to_ros(original, ros_type) + converted = conversion.ros_to_dimos(ros_message, type(original)) + assert converted == original + + +@requires_ros +@pytest.mark.parametrize("sec,nanosec", [(0, 0), (-1, 999999999), (1700000000, 123456789)]) +def test_ros_origin_timestamp_and_nested_fields_are_exact(sec, nanosec): + ros_type = conversion.derive_ros_type(PointStamped) + source = ros_type() + source.header.stamp.sec = sec + source.header.stamp.nanosec = nanosec + source.header.frame_id = "odom" + source.point.x, source.point.y, source.point.z = 1.5, -2.5, 3.5 + + converted = conversion.ros_to_dimos(source, PointStamped) + assert converted == PointStamped( + header=Header(stamp=Time(sec=sec, nanosec=nanosec), frame_id="odom"), + point=Point(x=1.5, y=-2.5, z=3.5), + ) + assert conversion.dimos_to_ros(converted, ros_type) == source + + +@requires_ros +def test_matching_layout_does_not_allow_wrong_type(): + from_type = conversion.derive_ros_type(Point) + to_type = conversion.derive_ros_type(Vector3) + with pytest.raises(TypeError, match="target does not match"): + conversion.dimos_to_ros(Point(), to_type) + with pytest.raises(TypeError, match="source does not match"): + conversion.ros_to_dimos(from_type(), Vector3) + + +def test_missing_ros_is_explicit_and_does_not_affect_generated_messages(monkeypatch): + monkeypatch.setattr(conversion, "ros_serialization", None) + message = Point(x=1.0) + assert Point.decode(message.encode()) == message + with pytest.raises(ImportError, match="requires rclpy"): + conversion.dimos_to_ros(message, Point) + with pytest.raises(ImportError, match="requires rclpy"): + conversion.ros_to_dimos(message, Point) + + +@pytest.mark.parametrize("name", ["geometry_msgs.Point", "geometry_msgs/srv/Point", "../msg/Point"]) +def test_invalid_type_name_fails_before_import(name): + invalid = type("Invalid", (), {"msg_name": name}) + with pytest.raises(ValueError, match="expected package/msg/Type"): + conversion.derive_ros_type(invalid) diff --git a/dimos/protocol/pubsub/impl/test_zenohpubsub.py b/dimos/protocol/pubsub/impl/test_zenohpubsub.py index ecf407c7373..4dd88ad1354 100644 --- a/dimos/protocol/pubsub/impl/test_zenohpubsub.py +++ b/dimos/protocol/pubsub/impl/test_zenohpubsub.py @@ -25,10 +25,10 @@ pytest.importorskip("zenoh") +from dimos_generated.geometry_msgs.msg import Twist +from dimos_generated.sensor_msgs.msg import Image import zenoh -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.sensor_msgs.Image import Image from dimos.protocol.pubsub.impl.lcmpubsub import Topic as LCMTopic from dimos.protocol.pubsub.impl.zenohpubsub import ( UNCONFLATED_QUEUE_LEN, @@ -232,7 +232,7 @@ def race() -> None: def test_concrete_subscription_passes_topic_through(self, pubsub, retry_until) -> None: received: list[Topic] = [] event = threading.Event() - topic = Topic("dimos/test/passthrough", lcm_type=Twist) + topic = Topic("dimos/test/passthrough", msg_type=Twist) def callback(msg: bytes, t: Topic) -> None: received.append(t) @@ -247,7 +247,7 @@ def test_untyped_topic_with_dotted_segment_round_trips(self, pubsub, retry_until # must not re-parse it into base+type on receive. received: list[Topic] = [] event = threading.Event() - topic = Topic("dimos/test/geometry_msgs.Twist") + topic = Topic("dimos/test/geometry_msgs/msg/Twist") def callback(msg: bytes, t: Topic) -> None: received.append(t) @@ -450,9 +450,9 @@ class TestTopicKeyExprConversion: """Tests for _topic_to_key_expr and _key_expr_to_topic round-trip.""" def test_typed_topic_to_key_expr(self) -> None: - topic = Topic("dimos/cmd_vel", lcm_type=Twist) + topic = Topic("dimos/cmd_vel", msg_type=Twist) key = _topic_to_key_expr(topic) - assert key == "dimos/cmd_vel/geometry_msgs.Twist" + assert key == "dimos/cmd_vel/geometry_msgs/msg/Twist" def test_untyped_topic_to_key_expr(self) -> None: topic = Topic("dimos/data") @@ -460,36 +460,36 @@ def test_untyped_topic_to_key_expr(self) -> None: assert key == "dimos/data" def test_key_expr_to_topic_with_known_type(self) -> None: - topic = _key_expr_to_topic("dimos/cmd_vel/geometry_msgs.Twist") + topic = _key_expr_to_topic("dimos/cmd_vel/geometry_msgs/msg/Twist") assert topic.topic == "dimos/cmd_vel" - assert topic.lcm_type is Twist + assert topic.msg_type is Twist def test_key_expr_to_topic_with_unknown_type(self) -> None: - topic = _key_expr_to_topic("dimos/data/unknown.FooBar") - # Last segment doesn't resolve — entire string becomes the topic - assert topic.topic == "dimos/data/unknown.FooBar" - assert topic.lcm_type is None + topic = _key_expr_to_topic("dimos/data/unknown/msg/FooBar", default_msg_type=Twist) + # Unknown declared types cannot fall back to a subscriber codec + assert topic.topic == "dimos/data" + assert topic.msg_type is None def test_key_expr_to_topic_with_no_slash(self) -> None: topic = _key_expr_to_topic("simple_topic") assert topic.topic == "simple_topic" - assert topic.lcm_type is None + assert topic.msg_type is None def test_key_expr_to_topic_uses_default_type(self) -> None: - topic = _key_expr_to_topic("dimos/data", default_lcm_type=Twist) + topic = _key_expr_to_topic("dimos/data", default_msg_type=Twist) assert topic.topic == "dimos/data" - assert topic.lcm_type is Twist + assert topic.msg_type is Twist def test_round_trip_typed(self) -> None: - original = Topic("dimos/color_image", lcm_type=Image) + original = Topic("dimos/color_image", msg_type=Image) key = _topic_to_key_expr(original) reconstructed = _key_expr_to_topic(key) assert reconstructed.topic == original.topic - assert reconstructed.lcm_type is original.lcm_type + assert reconstructed.msg_type is original.msg_type def test_round_trip_untyped(self) -> None: original = Topic("dimos/gps_location") key = _topic_to_key_expr(original) reconstructed = _key_expr_to_topic(key) assert reconstructed.topic == original.topic - assert reconstructed.lcm_type is None + assert reconstructed.msg_type is None diff --git a/dimos/protocol/pubsub/impl/webrtc/providers/broker.py b/dimos/protocol/pubsub/impl/webrtc/providers/broker.py index b8637442e0f..17e2219633a 100644 --- a/dimos/protocol/pubsub/impl/webrtc/providers/broker.py +++ b/dimos/protocol/pubsub/impl/webrtc/providers/broker.py @@ -104,7 +104,7 @@ class BrokerProvider(AsyncProviderBase): Inbound (operator → robot): ``cmd_unreliable`` + ``state_reliable``; subscribers get the bytes of the channel matching their topic, and typed - demux by LCM fingerprint happens at the transport layer. Outbound + demux by explicit frame type happens at the transport layer. Outbound (robot → operator): ``publish()`` on ``state_reliable_back`` / ``map_unreliable``; while no operator is connected the channel doesn't exist and messages drop, which is normal pubsub behaviour. Media rides the @@ -637,7 +637,7 @@ def set_video_frame(self, img: Any) -> None: def subscribe(self, topic: str, callback: Callable[[bytes, str], None]) -> Callable[[], None]: """Subscribers receive the bytes of the inbound channel matching - their topic; the transport layer filters by LCM fingerprint.""" + their topic; the transport layer filters by explicit frame channel and type.""" if not self.is_connected: self.start() with self._lock: diff --git a/dimos/protocol/pubsub/impl/webrtc/providers/test_video_track_generated.py b/dimos/protocol/pubsub/impl/webrtc/providers/test_video_track_generated.py new file mode 100644 index 00000000000..375a635b1cf --- /dev/null +++ b/dimos/protocol/pubsub/impl/webrtc/providers/test_video_track_generated.py @@ -0,0 +1,44 @@ +# 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. + +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio + +import numpy as np +import pytest + +from dimos.msgs.image import image_from_array +from dimos.protocol.pubsub.impl.webrtc.providers.video_track import CameraVideoTrack + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "encoding,channels", [("rgb8", 3), ("bgr8", 3), ("rgba8", 4), ("bgra8", 4), ("mono8", 1)] +) +async def test_video_track_uses_generated_encoding_without_network(encoding, channels): + pixels = np.full((4, 4) if channels == 1 else (4, 4, channels), 127, dtype=np.uint8) + image = image_from_array(pixels, encoding=encoding) + original = image.encode() + track = CameraVideoTrack(asyncio.get_running_loop()) + try: + track.arm() + track.set_latest(image) + frame = await asyncio.wait_for(track.recv(), timeout=1) + assert frame.width == 4 and frame.height == 4 + assert frame.pts == 0 + assert image.encode() == original + finally: + track.stop() diff --git a/dimos/protocol/pubsub/impl/webrtc/test_transport.py b/dimos/protocol/pubsub/impl/webrtc/test_transport.py index d2a81dd5103..c3bc29b1302 100644 --- a/dimos/protocol/pubsub/impl/webrtc/test_transport.py +++ b/dimos/protocol/pubsub/impl/webrtc/test_transport.py @@ -18,8 +18,8 @@ from collections.abc import Callable import pickle -import struct +from dimos_generated.geometry_msgs.msg import Point, TwistStamped, Vector3 from pydantic import ValidationError import pytest @@ -28,7 +28,6 @@ from dimos.core.coordination.blueprints import Blueprint from dimos.core.coordination.module_coordinator import _materialize_transports from dimos.core.transport import WebRTCTransport -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped from dimos.protocol.pubsub.impl.webrtc.providers import spec from dimos.protocol.pubsub.impl.webrtc.providers.broker import BrokerConfig from dimos.protocol.pubsub.impl.webrtc.providers.spec import ( @@ -36,6 +35,7 @@ ProviderConfig, ) from dimos.protocol.pubsub.impl.webrtc.webrtcpubsub import WebRTCPubSub +from dimos.web.relay_bridge.protocol import FrameHeader, decode_data_frame, encode_data_frame # ─── Mock provider ─────────────────────────────────────────────────── @@ -85,51 +85,6 @@ class MockTransport(WebRTCTransport): _config_cls = MockConfig -# ─── Fake LCM messages ─────────────────────────────────────────────── - - -class FakeLCMMsg: - msg_name = "test.FakeLCMMsg" - _FINGERPRINT = b"\x01\x02\x03\x04\x05\x06\x07\x08" - - def __init__(self, value: float = 0.0): - self.value = value - - @classmethod - def _get_packed_fingerprint(cls) -> bytes: - return cls._FINGERPRINT - - def lcm_encode(self) -> bytes: - return self._FINGERPRINT + struct.pack(" FakeLCMMsg: - if data[:8] != cls._FINGERPRINT: # like real LCM generated code - raise ValueError("Decode error") - return cls(struct.unpack(" bytes: - return cls._FINGERPRINT - - def lcm_encode(self) -> bytes: - return self._FINGERPRINT + self.text.encode() - - @classmethod - def lcm_decode(cls, data: bytes) -> OtherLCMMsg: - if data[:8] != cls._FINGERPRINT: # like real LCM generated code - raise ValueError("Decode error") - return cls(data[8:].decode()) - - # ─── Transport modes ───────────────────────────────────────────────── @@ -142,46 +97,71 @@ def test_raw_bytes_mode() -> None: def test_typed_encode_decode() -> None: - transport = MockTransport("cmd_unreliable", FakeLCMMsg, name="typed") - received: list[FakeLCMMsg] = [] + transport = MockTransport("cmd_unreliable", Point, name="typed") + received: list[Point] = [] transport.subscribe(lambda msg: received.append(msg)) - transport.broadcast(None, FakeLCMMsg(3.14)) + transport.broadcast(None, Point(x=3.14)) assert len(received) == 1 - assert abs(received[0].value - 3.14) < 1e-9 + assert abs(received[0].x - 3.14) < 1e-9 def test_multiple_types_multiplexed() -> None: """Typed transports sharing one channel each receive only their own type.""" - t1 = MockTransport("cmd_unreliable", FakeLCMMsg, name="mux") - t2 = MockTransport("cmd_unreliable", OtherLCMMsg, name="mux") - r1: list[FakeLCMMsg] = [] - r2: list[OtherLCMMsg] = [] + t1 = MockTransport("cmd_unreliable", Point, name="mux") + t2 = MockTransport("cmd_unreliable", Vector3, name="mux") + r1: list[Point] = [] + r2: list[Vector3] = [] t1.subscribe(lambda msg: r1.append(msg)) t2.subscribe(lambda msg: r2.append(msg)) - t1.broadcast(None, FakeLCMMsg(1.0)) - t2.broadcast(None, OtherLCMMsg("world")) - assert [m.value for m in r1] == [1.0] - assert [m.text for m in r2] == ["world"] - + t1.broadcast(None, Point(x=1.0)) + t2.broadcast(None, Vector3(x=2.0)) + assert [m.x for m in r1] == [1.0] + assert [m.x for m in r2] == [2.0] -def test_wire_fingerprint_matches_encoding() -> None: - """Demux must follow the wire format, not _get_packed_fingerprint(). - TwistStamped inherits Twist's fingerprint but encodes as LCM TwistStamped — - any filter keyed on the class fingerprint would drop every real message. - The try-decode demux delegates the check to lcm_decode, which gets it right. - """ - transport = MockTransport("cmd_unreliable", TwistStamped, name="wire") - received: list[TwistStamped] = [] - transport.subscribe(lambda msg: received.append(msg)) - - wire = TwistStamped(linear=[0.5, 0, 0], angular=[0, 0, 0.1], frame_id="keyboard").lcm_encode() - assert wire[:8] != TwistStamped._get_packed_fingerprint() +def test_generated_frames_declare_type_and_sequence() -> None: + publisher = MockTransport("data", Point, name="framed") + raw = MockTransport("data", name="framed") + received: list[bytes] = [] + raw.subscribe(received.append) + publisher.broadcast(None, Point(x=1)) + publisher.broadcast(None, Point(x=2)) + frames = [decode_data_frame(data) for data in received] + assert [frame.header.seq for frame in frames] == [0, 1] + assert all(frame.header.ch == "data" for frame in frames) + assert all(frame.header.meta == {"type": Point.msg_name, "encoding": "cdr"} for frame in frames) + assert [Point.decode(frame.payload).x for frame in frames] == [1, 2] + + +@pytest.mark.parametrize("bad", ["type", "topic", "encoding", "truncated", "trailing", "raw"]) +def test_invalid_frame_does_not_poison_subscription(bad: str) -> None: + transport = MockTransport("cmd_unreliable", Point, name=f"invalid-{bad}") + received: list[Point] = [] + transport.subscribe(received.append) + header = FrameHeader( + ch="cmd_unreliable", + seq=1, + ts=0, + delivery="latest", + meta={"type": Point.msg_name, "encoding": "cdr"}, + ) + if bad == "topic": + header.ch = "other" + elif bad in ("type", "encoding"): + assert header.meta is not None + header.meta[bad] = "unknown" + wire = encode_data_frame(header, Point(x=1).encode()) + if bad == "truncated": + wire = wire[:-1] + elif bad == "trailing": + wire += b"extra" + elif bad == "raw": + wire = Point(x=1).encode() transport._config.provider().publish("cmd_unreliable", wire) - assert len(received) == 1 - assert abs(received[0].linear.x - 0.5) < 1e-9 - assert received[0].frame_id == "keyboard" + assert received == [] + transport.broadcast(None, Point(x=2)) + assert received == [Point(x=2)] # ─── Pickling + provider sharing ───────────────────────────────────── @@ -190,20 +170,20 @@ def test_wire_fingerprint_matches_encoding() -> None: def test_pickle_roundtrip_preserves_everything() -> None: """Transports are pickled into module worker processes; topic, type, and provider config must all survive.""" - t1 = MockTransport("cmd_unreliable", FakeLCMMsg, name="pickled") + t1 = MockTransport("cmd_unreliable", Point, name="pickled") t2 = pickle.loads(pickle.dumps(t1)) assert type(t2) is MockTransport assert t2.topic == t1.topic - assert t2._msg_type is FakeLCMMsg + assert t2._msg_type is Point assert t2._config == t1._config # Same config → same per-process provider, so the two halves interoperate. - received: list[FakeLCMMsg] = [] + received: list[Point] = [] t2.subscribe(lambda msg: received.append(msg)) - t1.broadcast(None, FakeLCMMsg(42.0)) + t1.broadcast(None, Point(x=42.0)) assert len(received) == 1 - assert abs(received[0].value - 42.0) < 1e-9 + assert abs(received[0].x - 42.0) < 1e-9 def test_provider_singleton_per_config() -> None: @@ -261,7 +241,7 @@ def test_provider_config_is_frozen_and_hashable() -> None: def test_blueprint_config_exposes_transport_fields() -> None: """Each unique `_config_cls` becomes a `transports.` config section.""" - bp = Blueprint(blueprints=()).transports({("topic", FakeLCMMsg): MockTransport.spec("topic")}) + bp = Blueprint(blueprints=()).transports({("topic", Point): MockTransport.spec("topic")}) parser = BlueprintConfigParser(bp) parsed = parser.parse(environ={}, overrides={"transports": {"mock": {"name": "override"}}}) assert parsed.transport_overrides() == {"mock": {"name": "override"}} @@ -273,8 +253,8 @@ def test_blueprint_config_exposes_transport_fields() -> None: # Multiple transports sharing one `_config_cls` collapse to one section. bp_shared = Blueprint(blueprints=()).transports( { - ("a", FakeLCMMsg): MockTransport.spec("a"), - ("b", FakeLCMMsg): MockTransport.spec("b"), + ("a", Point): MockTransport.spec("a"), + ("b", Point): MockTransport.spec("b"), } ) shared = BlueprintConfigParser(bp_shared).parse( @@ -285,32 +265,32 @@ def test_blueprint_config_exposes_transport_fields() -> None: def test_transport_overrides_apply_and_survive_pickle() -> None: """Materialization builds each transport with its resolved config; pickle-safe.""" - bp = Blueprint(blueprints=()).transports({("topic", FakeLCMMsg): MockTransport.spec("topic")}) + bp = Blueprint(blueprints=()).transports({("topic", Point): MockTransport.spec("topic")}) - transport = _materialize_transports(bp, {"mock": {"name": "overridden"}})[("topic", FakeLCMMsg)] + transport = _materialize_transports(bp, {"mock": {"name": "overridden"}})[("topic", Point)] assert transport._config.name == "overridden" assert pickle.loads(pickle.dumps(transport))._config.name == "overridden" # No override → config built from the spec's defaults. - default = _materialize_transports(bp, {})[("topic", FakeLCMMsg)] + default = _materialize_transports(bp, {})[("topic", Point)] assert default._config.name == "default" def test_materialize_uses_resolved_config() -> None: """The config reaching the transport is built straight from the overrides, not a default instance that is later mutated.""" - bp = Blueprint(blueprints=()).transports({("topic", FakeLCMMsg): MockTransport.spec("topic")}) + bp = Blueprint(blueprints=()).transports({("topic", Point): MockTransport.spec("topic")}) - transport = _materialize_transports(bp, {"mock": {"name": "resolved"}})[("topic", FakeLCMMsg)] + transport = _materialize_transports(bp, {"mock": {"name": "resolved"}})[("topic", Point)] assert transport._config == MockConfig(name="resolved") def test_transport_overrides_coerce_string_values() -> None: """CLI/env overrides arrive as raw strings; non-str fields must coerce, not pass through.""" - bp = Blueprint(blueprints=()).transports({("topic", FakeLCMMsg): MockTransport.spec("topic")}) - transport = _materialize_transports(bp, {"mock": {"count": "5"}})[("topic", FakeLCMMsg)] + bp = Blueprint(blueprints=()).transports({("topic", Point): MockTransport.spec("topic")}) + transport = _materialize_transports(bp, {"mock": {"count": "5"}})[("topic", Point)] assert transport._config.count == 5 @@ -323,11 +303,11 @@ def test_raw_transport_pins_still_work() -> None: # Pinned on the active backend: a plain LCM/Zenoh pin that does not match it # is deliberately rebuilt by the backend switch, which is a different path. - raw = make_transport("/raw_topic", FakeLCMMsg) + raw = make_transport("/raw_topic", Point) bp = Blueprint(blueprints=()).transports( { - ("raw", FakeLCMMsg): raw, - ("speced", FakeLCMMsg): MockTransport.spec("topic"), + ("raw", Point): raw, + ("speced", Point): MockTransport.spec("topic"), } ) # Raw pins contribute no transports.* config sections; the spec still does. @@ -337,8 +317,8 @@ def test_raw_transport_pins_still_work() -> None: assert parsed.transport_overrides() == {"mock": {"name": "x"}} materialized = _materialize_transports(bp, {}) - assert materialized[("raw", FakeLCMMsg)] is raw - assert isinstance(materialized[("speced", FakeLCMMsg)], MockTransport) + assert materialized[("raw", Point)] is raw + assert isinstance(materialized[("speced", Point)], MockTransport) # ─── Broker credential validation ──────────────────────────────────── diff --git a/dimos/protocol/pubsub/impl/webrtc/tool_blueprint_e2e.py b/dimos/protocol/pubsub/impl/webrtc/tool_blueprint_e2e.py index 286e2a3b264..c00d5d53df3 100644 --- a/dimos/protocol/pubsub/impl/webrtc/tool_blueprint_e2e.py +++ b/dimos/protocol/pubsub/impl/webrtc/tool_blueprint_e2e.py @@ -27,6 +27,7 @@ import os import time +from dimos_generated.geometry_msgs.msg import Twist, TwistStamped, Vector3 import pytest from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser @@ -35,7 +36,6 @@ from dimos.core.module import Module from dimos.core.stream import In, Out from dimos.core.transport import WebRTCTransport -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped from dimos.protocol.pubsub.impl.webrtc.providers.cloudflare import CloudflareConfig from dimos.protocol.pubsub.impl.webrtc.providers.spec import WEBRTC_AVAILABLE @@ -54,7 +54,7 @@ class TwistSource(Module): @rpc def send(self, x: float) -> None: - self.cmd_webrtc.publish(TwistStamped(linear=[x, 0, 0], angular=[0, 0, 0])) + self.cmd_webrtc.publish(TwistStamped(twist=Twist(linear=Vector3(x=x)))) class TwistSink(Module): @@ -64,7 +64,7 @@ class TwistSink(Module): def start(self) -> None: super().start() self._received: list[float] = [] - self.cmd_webrtc.subscribe(lambda msg: self._received.append(msg.linear.x)) + self.cmd_webrtc.subscribe(lambda msg: self._received.append(msg.twist.linear.x)) @rpc def received(self) -> list[float]: diff --git a/dimos/protocol/pubsub/impl/webrtc/tool_broker_e2e.py b/dimos/protocol/pubsub/impl/webrtc/tool_broker_e2e.py index 20c79310488..d82d7ae03d2 100644 --- a/dimos/protocol/pubsub/impl/webrtc/tool_broker_e2e.py +++ b/dimos/protocol/pubsub/impl/webrtc/tool_broker_e2e.py @@ -83,12 +83,13 @@ def test_operator_to_transport_e2e() -> None: RTCPeerConnection, RTCSessionDescription, ) + from dimos_generated.geometry_msgs.msg import Twist, TwistStamped, Vector3 import numpy as np from dimos.core.transport import CloudflareTransport, CloudflareVideoTransport - from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped - from dimos.msgs.sensor_msgs.Image import Image, ImageFormat + from dimos.msgs.image import image_from_array from dimos.protocol.pubsub.impl.webrtc.providers.spec import wait_connected, wait_open + from dimos.web.relay_bridge.protocol import FrameHeader, decode_data_frame, encode_data_frame # The TELEOP_* env fallback inside BrokerProvider was removed when # transport config moved to the blueprint flow — pass the key explicitly. @@ -128,7 +129,7 @@ async def operator() -> int: # from flowing RTP, so (as with a real camera stream) frames must be # arriving before the operator bridges, or the pull offer comes back # without a usable video m-line. - frame = Image(data=np.full((120, 160, 3), 128, dtype=np.uint8), format=ImageFormat.BGR) + frame = image_from_array(np.full((120, 160, 3), 128, dtype=np.uint8), encoding="bgr8") feeding = True async def _feed() -> None: @@ -230,16 +231,27 @@ def _on_back(payload: object) -> None: # Robot → operator: telemetry through the broker-bridged back channel. for i in range(10): - back_transport.broadcast(None, TwistStamped(linear=[0.0, 0.0, 1.0 + i])) + back_transport.broadcast(None, TwistStamped(twist=Twist(linear=Vector3(z=1.0 + i)))) await asyncio.sleep(0.05) await _wait_for(lambda: bool(back_bytes), 5.0, "robot->operator telemetry") - back_msg = TwistStamped.lcm_decode(back_bytes[-1]) - assert back_msg.linear.z >= 1.0, back_msg.linear + back_msg = TwistStamped.decode(decode_data_frame(back_bytes[-1]).payload) + assert back_msg.twist.linear.z >= 1.0, back_msg.twist.linear sent = 0 for i in range(40): - msg = TwistStamped(linear=[0.5, 0.0, 0.0], angular=[0.0, 0.0, i * 0.01]) - ch.send(msg.lcm_encode()) + msg = TwistStamped(twist=Twist(linear=Vector3(x=0.5), angular=Vector3(z=i * 0.01))) + ch.send( + encode_data_frame( + FrameHeader( + ch="cmd_unreliable", + seq=i, + ts=time.time(), + delivery="latest", + meta={"type": TwistStamped.msg_name, "encoding": "cdr"}, + ), + msg.encode(), + ) + ) sent += 1 await asyncio.sleep(0.05) # Unreliable channel: wait for the pass condition itself (early exit), @@ -261,7 +273,7 @@ def _on_back(payload: object) -> None: assert len(received) >= sent * 0.8, f"sent={sent} received={len(received)}" sample = received[-1] assert isinstance(sample, TwistStamped) - assert abs(sample.linear.x - 0.5) < 1e-9 + assert abs(sample.twist.linear.x - 0.5) < 1e-9 finally: transport.stop() back_transport.stop() diff --git a/dimos/protocol/pubsub/impl/webrtc/webrtcpubsub.py b/dimos/protocol/pubsub/impl/webrtc/webrtcpubsub.py index fa05e47ddf0..6ec24323fc4 100644 --- a/dimos/protocol/pubsub/impl/webrtc/webrtcpubsub.py +++ b/dimos/protocol/pubsub/impl/webrtc/webrtcpubsub.py @@ -38,7 +38,7 @@ class WebRTCPubSub(AllPubSub[str, bytes]): """AllPubSub[str, bytes] over a WebRTC DataChannel provider. WebRTC DataChannels are inherently "receive all" — messages arrive on a - shared multiplexed channel and are demuxed by topic/fingerprint. This + shared multiplexed channel and are demuxed by frame channel/type. This matches LCM multicast semantics, hence AllPubSub. """ diff --git a/dimos/protocol/pubsub/impl/zenohpubsub.py b/dimos/protocol/pubsub/impl/zenohpubsub.py index 6c28ad66830..9b6c677a9f0 100644 --- a/dimos/protocol/pubsub/impl/zenohpubsub.py +++ b/dimos/protocol/pubsub/impl/zenohpubsub.py @@ -25,7 +25,7 @@ import zenoh from dimos.msgs.helpers import resolve_msg_type -from dimos.protocol.pubsub.encoders import LCMEncoderMixin, PickleEncoderMixin +from dimos.protocol.pubsub.encoders import CDREncoderMixin, PickleEncoderMixin from dimos.protocol.pubsub.impl.lcmpubsub import Topic as LCMTopic from dimos.protocol.pubsub.spec import AllPubSub from dimos.protocol.service.zenohservice import ZenohService @@ -87,14 +87,14 @@ class Topic(LCMTopic): @property def key_expr(self) -> str: - """The Zenoh key expression for this topic, embedding lcm_type after a '/'. + """The Zenoh key expression for this topic, embedding msg_type after a '/'. Examples: - Topic("dimos/cmd_vel", Twist) -> "dimos/cmd_vel/geometry_msgs.Twist" + Topic("dimos/cmd_vel", Twist) -> "dimos/cmd_vel/geometry_msgs/msg/Twist" Topic("dimos/data") -> "dimos/data" """ - if self.lcm_type is not None: - return f"{self.pattern}/{self.lcm_type.msg_name}" + if self.msg_type is not None: + return f"{self.pattern}/{self.msg_type.msg_name}" return self.pattern @@ -102,32 +102,22 @@ def _topic_to_key_expr(topic: LCMTopic) -> str: """Convert any LCM-compatible Topic to a Zenoh key expression.""" if isinstance(topic, Topic): return topic.key_expr - return Topic(topic=topic.topic, lcm_type=topic.lcm_type).key_expr + return Topic(topic=topic.topic, msg_type=topic.msg_type).key_expr @lru_cache(maxsize=1024) -def _key_expr_to_topic(key_expr: str, default_lcm_type: type | None = None) -> Topic: - """Reconstruct a Topic from a Zenoh key expression. +def _key_expr_to_topic(key_expr: str, default_msg_type: type | None = None) -> Topic: + """Resolve a qualified package/msg/Type suffix; never guess from payload bytes. - Parses the last '/' segment and attempts to resolve it as a DimosMsg - type via resolve_msg_type(). If resolution succeeds, the segment is - treated as the type suffix and the remainder as the base topic. - - Results are cached; callers must treat the returned Topic as immutable. - - Examples: - "dimos/cmd_vel/geometry_msgs.Twist" -> Topic("dimos/cmd_vel", Twist) - "dimos/data" -> Topic("dimos/data", default_lcm_type) - "dimos/data/unknown.Foo" -> Topic("dimos/data/unknown.Foo", default_lcm_type) + Untyped keys retain their complete path. An explicit but unknown type suffix + cannot fall back to the subscriber's type. """ - # Try to resolve the last segment as a message type - parts = key_expr.rsplit("/", 1) - if len(parts) == 2: - base, maybe_type = parts - lcm_type = resolve_msg_type(maybe_type) - if lcm_type is not None: - return Topic(topic=base, lcm_type=lcm_type) - return Topic(topic=key_expr, lcm_type=default_lcm_type) + parts = key_expr.rsplit("/", 3) + if len(parts) == 4 and parts[2] == "msg": + base, package, _, name = parts + msg_type = resolve_msg_type(f"{package}/msg/{name}") + return Topic(topic=base, msg_type=msg_type) + return Topic(topic=key_expr, msg_type=default_msg_type) # Messages a keyed channel may hold between drains before the oldest is dropped. @@ -223,7 +213,7 @@ def on_sample(sample: zenoh.Sample) -> None: if sample_key == key_expr: recv_topic = topic else: - recv_topic = _key_expr_to_topic(sample_key, topic.lcm_type) + recv_topic = _key_expr_to_topic(sample_key, topic.msg_type) callback(data, recv_topic) sub = self.session.declare_subscriber(key_expr, on_sample) @@ -346,10 +336,10 @@ def stop(self) -> None: class Zenoh( # type: ignore[misc] - LCMEncoderMixin, + CDREncoderMixin, ZenohPubSubBase, ): - """Zenoh pub/sub with LCM encoding for typed DimosMsg.""" + """Zenoh pub/sub with CDR encoding for typed DimosMsg.""" ... diff --git a/dimos/protocol/pubsub/test_pattern_sub.py b/dimos/protocol/pubsub/test_pattern_sub.py index b1d4827c37a..10bc1d5290a 100644 --- a/dimos/protocol/pubsub/test_pattern_sub.py +++ b/dimos/protocol/pubsub/test_pattern_sub.py @@ -22,11 +22,9 @@ import time from typing import Any, Generic, TypeVar +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion, Vector3 import pytest -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.protocol.pubsub.impl.lcmpubsub import LCM, LCMPubSubBase, Topic from dimos.protocol.pubsub.patterns import Glob from dimos.protocol.pubsub.spec import AllPubSub, PubSub @@ -83,9 +81,12 @@ def lcm_bytes_context(url: str) -> Iterator[tuple[LCMPubSubBase, LCMPubSubBase]] name="lcm_typed", pubsub_context=lcm_typed_context, topic_values=[ - (Topic("/sensor/position", Vector3), Vector3(1, 2, 3)), - (Topic("/sensor/orientation", Quaternion), Quaternion(0, 0, 0, 1)), - (Topic("/robot/arm", Pose), Pose(Vector3(4, 5, 6), Quaternion(0, 0, 0, 1))), + (Topic("/sensor/position", Vector3), Vector3(x=1, y=2, z=3)), + (Topic("/sensor/orientation", Quaternion), Quaternion(w=1)), + ( + Topic("/robot/arm", Pose), + Pose(position=Point(x=4, y=5, z=6), orientation=Quaternion(w=1)), + ), ], tags={"all", "glob", "regex"}, glob_patterns=[ diff --git a/dimos/protocol/pubsub/test_spec.py b/dimos/protocol/pubsub/test_spec.py index e3e21437ae1..78a576a4778 100644 --- a/dimos/protocol/pubsub/test_spec.py +++ b/dimos/protocol/pubsub/test_spec.py @@ -21,9 +21,9 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import Vector3 import pytest -from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.protocol.pubsub.impl.lcmpubsub import LCM, Topic from dimos.protocol.pubsub.impl.memory import Memory from dimos.protocol.pubsub.impl.zenohpubsub import PickleZenoh, Zenoh @@ -122,8 +122,12 @@ def lcm_context() -> Generator[LCM, None, None]: testdata.append( ( lcm_context, - Topic(topic="/test_topic", lcm_type=Vector3), - [Vector3(1, 2, 3), Vector3(4, 5, 6), Vector3(7, 8, 9)], # Using Vector3 as mock data, + Topic(topic="/test_topic", msg_type=Vector3), + [ + Vector3(x=1, y=2, z=3), + Vector3(x=4, y=5, z=6), + Vector3(x=7, y=8, z=9), + ], # Using Vector3 as mock data, ) ) @@ -163,8 +167,8 @@ def zenoh_lcm_context() -> Generator[Zenoh, None, None]: testdata.append( ( zenoh_lcm_context, - Topic(topic="dimos/test/spec", lcm_type=Vector3), - [Vector3(1, 2, 3), Vector3(4, 5, 6), Vector3(7, 8, 9)], + Topic(topic="dimos/test/spec", msg_type=Vector3), + [Vector3(x=1, y=2, z=3), Vector3(x=4, y=5, z=6), Vector3(x=7, y=8, z=9)], ) ) diff --git a/dimos/protocol/service/test_lcmservice.py b/dimos/protocol/service/test_lcmservice.py index d04ad0ec679..78f1145b0f8 100644 --- a/dimos/protocol/service/test_lcmservice.py +++ b/dimos/protocol/service/test_lcmservice.py @@ -341,7 +341,7 @@ def test_str_without_lcm_type(self) -> None: def test_str_with_lcm_type(self) -> None: mock_type = MagicMock() mock_type.msg_name = "TestMessage" - topic = Topic(topic="my_topic", lcm_type=mock_type) + topic = Topic(topic="my_topic", msg_type=mock_type) assert str(topic) == "my_topic#TestMessage" diff --git a/dimos/protocol/tf/static_tf_publisher.py b/dimos/protocol/tf/static_tf_publisher.py index 0f5b6d12625..587fae4349b 100644 --- a/dimos/protocol/tf/static_tf_publisher.py +++ b/dimos/protocol/tf/static_tf_publisher.py @@ -27,15 +27,16 @@ import asyncio import time +from dimos_generated.geometry_msgs.msg import Transform, TransformStamped, Vector3 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage from pydantic import Field from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import Out -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.time import time_from_nanoseconds from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -44,23 +45,25 @@ FrameSpec = tuple[str, str | None, tuple[float, float, float], tuple[float, float, float]] -def frames_to_edge_transforms(frames: list[FrameSpec]) -> list[Transform]: - """Build a ``parent -> child`` Transform for each non-root edge of a frame tree. +def frames_to_edge_transforms(frames: list[FrameSpec]) -> list[TransformStamped]: + """Build a ``parent -> child`` TransformStamped for each non-root edge of a frame tree. This is the static mount tree (the rigid sensor offsets); a tf buffer composes these edges to answer any ``world <- frame`` query once odometry supplies the moving ``world <- root`` edge. """ - transforms: list[Transform] = [] + transforms: list[TransformStamped] = [] for name, parent, translation, rpy in frames: if parent is None: continue transforms.append( - Transform( - translation=Vector3(*translation), - rotation=Quaternion.from_euler(Vector3(*rpy)), - frame_id=parent, + TransformStamped( + header=Header(frame_id=parent), child_frame_id=name, + transform=Transform( + translation=Vector3(x=translation[0], y=translation[1], z=translation[2]), + rotation=quaternion_from_euler(*rpy), + ), ) ) return transforms @@ -77,9 +80,9 @@ class StaticTfPublisher(Module): tf: Out[TFMessage] _running: bool = False - _transforms: list[Transform] = [] + _transforms: list[TransformStamped] = [] - def transforms(self) -> list[Transform]: + def transforms(self) -> list[TransformStamped]: """The static transforms to publish. Override in a rig-specific subclass.""" raise NotImplementedError( f"{type(self).__name__} must override transforms() with its mount frames" @@ -104,10 +107,10 @@ def start(self) -> None: async def _publish_loop(self) -> None: period = 1.0 / self.config.publish_hz while self._running: - now = time.time() + now = time_from_nanoseconds(time.time_ns()) for transform in self._transforms: - transform.ts = now - self.tf.publish(TFMessage(*self._transforms)) + transform.header.stamp = now + self.tf.publish(TFMessage(transforms=self._transforms)) await asyncio.sleep(period) @rpc diff --git a/dimos/protocol/tf/test_static_tf_publisher.py b/dimos/protocol/tf/test_static_tf_publisher.py new file mode 100644 index 00000000000..d274f6bcc54 --- /dev/null +++ b/dimos/protocol/tf/test_static_tf_publisher.py @@ -0,0 +1,47 @@ +# 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. + +"""Static frame definitions become generated ROS TF edges without convenience types.""" + +import math + +from dimos_generated.tf2_msgs.msg import TFMessage +import pytest + +from dimos.msgs.geometry import compose_transforms +from dimos.protocol.tf.static_tf_publisher import frames_to_edge_transforms + + +def test_frame_tree_skips_root_and_composes_rotated_mounts(): + edges = frames_to_edge_transforms( + [ + ("world", None, (0, 0, 0), (0, 0, 0)), + ("body", "world", (1, 0, 0), (0, 0, math.pi / 2)), + ("camera", "body", (2, 0, 3), (0, 0, 0)), + ] + ) + assert [(edge.header.frame_id, edge.child_frame_id) for edge in edges] == [ + ("world", "body"), + ("body", "camera"), + ] + decoded = TFMessage.decode(TFMessage(transforms=edges).encode()) + combined = compose_transforms(*decoded.transforms) + translation = combined.transform.translation + assert (translation.x, translation.y, translation.z) == pytest.approx((1, 2, 3)) + assert combined.header.stamp.sec == 0 + assert combined.header.stamp.nanosec == 0 + + +def test_empty_frame_tree_produces_no_edges(): + assert frames_to_edge_transforms([]) == [] diff --git a/dimos/protocol/tf/test_tf.py b/dimos/protocol/tf/test_tf.py index 47e791ec94c..cfc395d639d 100644 --- a/dimos/protocol/tf/test_tf.py +++ b/dimos/protocol/tf/test_tf.py @@ -18,17 +18,18 @@ import threading import time +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage import pytest from dimos.core.transport_factory import make_transport from dimos.memory.store.memory import MemoryStore from dimos.memory.store.sqlite import SqliteStore from dimos.memory.tf import StreamTF -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.protocol.tf.tf import TF, MultiTBuffer, TBuffer, TFLookup @@ -36,20 +37,22 @@ def test_tf_ros_example() -> None: tf = TF() - base_link_to_arm = Transform( - translation=Vector3(1.0, -1.0, 0.0), - rotation=Quaternion.from_euler(Vector3(0, 0, math.pi / 6)), - frame_id="base_link", + base_link_to_arm = TransformStamped( + header=Header(frame_id="base_link", stamp=time_from_seconds(time.time())), child_frame_id="arm", - ts=time.time(), + transform=Transform( + translation=Vector3(x=1.0, y=-1.0, z=0.0), + rotation=quaternion_from_euler(0, 0, math.pi / 6), + ), ) - arm_to_end = Transform( - translation=Vector3(1.0, 1.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), # Identity rotation - frame_id="arm", + arm_to_end = TransformStamped( + header=Header(frame_id="arm", stamp=time_from_seconds(time.time())), child_frame_id="end_effector", - ts=time.time(), + transform=Transform( + translation=Vector3(x=1.0, y=1.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) tf.publish(base_link_to_arm, arm_to_end) @@ -58,8 +61,8 @@ def test_tf_ros_example() -> None: end_effector_global_pose = tf.get("base_link", "end_effector") assert end_effector_global_pose is not None - assert end_effector_global_pose.translation.x == pytest.approx(1.366, abs=1e-3) - assert end_effector_global_pose.translation.y == pytest.approx(0.366, abs=1e-3) + assert end_effector_global_pose.transform.translation.x == pytest.approx(1.366, abs=1e-3) + assert end_effector_global_pose.transform.translation.y == pytest.approx(0.366, abs=1e-3) tf.dispose() @@ -79,20 +82,21 @@ def test_tf_main() -> None: # Create a transform from world to robot current_time = time.time() - world_to_charger = Transform( - translation=Vector3(2.0, -2.0, 0.0), - rotation=Quaternion.from_euler(Vector3(0, 0, 2)), - frame_id="world", + world_to_charger = TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(current_time)), child_frame_id="charger", - ts=current_time, + transform=Transform( + translation=Vector3(x=2.0, y=-2.0, z=0.0), rotation=quaternion_from_euler(0, 0, 2) + ), ) - world_to_robot = Transform( - translation=Vector3(1.0, 2.0, 3.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), # Identity rotation - frame_id="world", + world_to_robot = TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(current_time)), child_frame_id="robot", - ts=current_time, + transform=Transform( + translation=Vector3(x=1.0, y=2.0, z=3.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) # Broadcast the transform @@ -107,12 +111,13 @@ def test_tf_main() -> None: assert "robot" in frames # Add another transform in the chain - robot_to_sensor = Transform( - translation=Vector3(0.5, 0.0, 0.2), - rotation=Quaternion(0.0, 0.0, 0.707107, 0.707107), # 90 degrees around Z - frame_id="robot", + robot_to_sensor = TransformStamped( + header=Header(frame_id="robot", stamp=time_from_seconds(current_time)), child_frame_id="sensor", - ts=current_time, + transform=Transform( + translation=Vector3(x=0.5, y=0.0, z=0.2), + rotation=Quaternion(x=0.0, y=0.0, z=0.707107, w=0.707107), + ), ) broadcaster.publish(robot_to_sensor) @@ -128,27 +133,17 @@ def test_tf_main() -> None: # The chain should compose: world->robot (1,2,3) + robot->sensor (0.5,0,0.2) # Expected translation: (1.5, 2.0, 3.2) assert chain_transform is not None - assert abs(chain_transform.translation.x - 1.5) < 0.001 - assert abs(chain_transform.translation.y - 2.0) < 0.001 - assert abs(chain_transform.translation.z - 3.2) < 0.001 - - # we see something on camera - random_object_in_view = PoseStamped( - frame_id="random_object", - position=Vector3(1, 0, 0), + assert abs(chain_transform.transform.translation.x - 1.5) < 0.001 + assert abs(chain_transform.transform.translation.y - 2.0) < 0.001 + assert abs(chain_transform.transform.translation.z - 3.2) < 0.001 + + # A perceived object is expressed as a stamped sensor -> object transform. + random_t = TransformStamped( + header=Header(frame_id="sensor", stamp=time_from_seconds(time.time())), + child_frame_id="random_object", + transform=Transform(translation=Vector3(x=1, y=0, z=0)), ) - print("Random obj", random_object_in_view) - - # random_object is perceived by the sensor - # we create a transform pointing from sensor to object - random_t = random_object_in_view.new_transform_from("sensor") - - # we could have also done - assert random_t == random_object_in_view.new_transform_to("sensor").inverse() - - print("randm t", random_t) - # we broadcast our object location broadcaster.publish(random_t) @@ -176,21 +171,21 @@ def test_tf_main() -> None: print(broadcaster.graph()) assert world_object is not None - assert abs(world_object.translation.x - 1.5) < 0.001 - assert abs(world_object.translation.y - 3.0) < 0.001 - assert abs(world_object.translation.z - 3.2) < 0.001 + assert abs(world_object.transform.translation.x - 1.5) < 0.001 + assert abs(world_object.transform.translation.y - 3.0) < 0.001 + assert abs(world_object.transform.translation.z - 3.2) < 0.001 # this doesn't work atm robot_to_charger = broadcaster.get("robot", "charger") assert robot_to_charger is not None # Expected: robot->world->charger - print(f"robot_to_charger translation: {robot_to_charger.translation}") - print(f"robot_to_charger rotation: {robot_to_charger.rotation}") + print(f"robot_to_charger translation: {robot_to_charger.transform.translation}") + print(f"robot_to_charger rotation: {robot_to_charger.transform.rotation}") - assert abs(robot_to_charger.translation.x - 1.0) < 0.001 - assert abs(robot_to_charger.translation.y - (-4.0)) < 0.001 - assert abs(robot_to_charger.translation.z - (-3.0)) < 0.001 + assert abs(robot_to_charger.transform.translation.x - 1.0) < 0.001 + assert abs(robot_to_charger.transform.translation.y - (-4.0)) < 0.001 + assert abs(robot_to_charger.transform.translation.z - (-3.0)) < 0.001 # Stop services (they were autostarted but don't know how to autostop) broadcaster.dispose() @@ -202,12 +197,13 @@ def test_tf_main() -> None: class TestTBuffer: def test_add_transform(self) -> None: buffer = TBuffer(buffer_size=10.0) - transform = Transform( - translation=Vector3(1.0, 2.0, 3.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="world", + transform = TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(time.time())), child_frame_id="robot", - ts=time.time(), + transform=Transform( + translation=Vector3(x=1.0, y=2.0, z=3.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) buffer.add(transform) @@ -220,23 +216,22 @@ def test_get(self) -> None: # Add transforms at different times for i in range(3): - transform = Transform( - translation=Vector3(float(i), 0.0, 0.0), - frame_id="world", + transform = TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(base_time + i * 0.5)), child_frame_id="robot", - ts=base_time + i * 0.5, + transform=Transform(translation=Vector3(x=float(i), y=0.0, z=0.0)), ) buffer.add(transform) # Test getting latest transform latest = buffer.get() assert latest is not None - assert latest.translation.x == 2.0 + assert latest.transform.translation.x == 2.0 # Test getting transform at specific time middle = buffer.get(time_point=base_time + 0.75) assert middle is not None - assert middle.translation.x == 2.0 # Closest to i=1 + assert middle.transform.translation.x == 2.0 # Closest to i=1 # Test time tolerance result = buffer.get(time_point=base_time + 10.0, time_tolerance=0.1) @@ -247,20 +242,18 @@ def test_buffer_pruning(self) -> None: # Add old transform old_time = time.time() - 2.0 - old_transform = Transform( - translation=Vector3(1.0, 0.0, 0.0), - frame_id="world", + old_transform = TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(old_time)), child_frame_id="robot", - ts=old_time, + transform=Transform(translation=Vector3(x=1.0, y=0.0, z=0.0)), ) buffer.add(old_transform) # Add recent transform - recent_transform = Transform( - translation=Vector3(2.0, 0.0, 0.0), - frame_id="world", + recent_transform = TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(time.time())), child_frame_id="robot", - ts=time.time(), + transform=Transform(translation=Vector3(x=2.0, y=0.0, z=0.0)), ) buffer.add(recent_transform) @@ -268,7 +261,7 @@ def test_buffer_pruning(self) -> None: assert len(buffer) == 1 first = buffer.first() assert first is not None - assert first.translation.x == 2.0 + assert first.transform.translation.x == 2.0 class TestMultiTBuffer: @@ -276,18 +269,16 @@ def test_multiple_frame_pairs(self) -> None: ttbuffer = MultiTBuffer(buffer_size=10.0) # Add transforms for different frame pairs - transform1 = Transform( - translation=Vector3(1.0, 0.0, 0.0), - frame_id="world", + transform1 = TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(time.time())), child_frame_id="robot1", - ts=time.time(), + transform=Transform(translation=Vector3(x=1.0, y=0.0, z=0.0)), ) - transform2 = Transform( - translation=Vector3(2.0, 0.0, 0.0), - frame_id="world", + transform2 = TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(time.time())), child_frame_id="robot2", - ts=time.time(), + transform=Transform(translation=Vector3(x=2.0, y=0.0, z=0.0)), ) ttbuffer.receive_transform(transform1, transform2) @@ -301,18 +292,16 @@ def test_graph(self) -> None: ttbuffer = MultiTBuffer(buffer_size=10.0) # Add transforms for different frame pairs - transform1 = Transform( - translation=Vector3(1.0, 0.0, 0.0), - frame_id="world", + transform1 = TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(time.time())), child_frame_id="robot1", - ts=time.time(), + transform=Transform(translation=Vector3(x=1.0, y=0.0, z=0.0)), ) - transform2 = Transform( - translation=Vector3(2.0, 0.0, 0.0), - frame_id="world", + transform2 = TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(time.time())), child_frame_id="robot2", - ts=time.time(), + transform=Transform(translation=Vector3(x=2.0, y=0.0, z=0.0)), ) ttbuffer.receive_transform(transform1, transform2) @@ -326,11 +315,10 @@ def test_forward_tolerance_returns_when_buffer_fills(self) -> None: def publish_after_delay() -> None: time.sleep(0.05) ttbuffer.receive_transform( - Transform( - translation=Vector3(1.0, 0.0, 0.0), - frame_id="world", + TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(base_time)), child_frame_id="robot", - ts=base_time, + transform=Transform(translation=Vector3(x=1.0, y=0.0, z=0.0)), ) ) @@ -345,7 +333,7 @@ def publish_after_delay() -> None: publisher.join() assert result is not None - assert result.translation.x == 1.0 + assert result.transform.translation.x == 1.0 assert elapsed < 0.5 def test_forward_tolerance_times_out(self) -> None: @@ -360,11 +348,10 @@ def test_forward_tolerance_fast_path_when_already_available(self) -> None: ttbuffer = MultiTBuffer() base_time = time.time() ttbuffer.receive_transform( - Transform( - translation=Vector3(2.0, 0.0, 0.0), - frame_id="world", + TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(base_time)), child_frame_id="robot", - ts=base_time, + transform=Transform(translation=Vector3(x=2.0, y=0.0, z=0.0)), ) ) @@ -373,7 +360,7 @@ def test_forward_tolerance_fast_path_when_already_available(self) -> None: elapsed = time.monotonic() - t0 assert result is not None - assert result.translation.x == 2.0 + assert result.transform.translation.x == 2.0 assert elapsed < 0.05 def test_forward_tolerance_wakes_on_chain_completion(self) -> None: @@ -381,22 +368,20 @@ def test_forward_tolerance_wakes_on_chain_completion(self) -> None: base_time = time.time() ttbuffer.receive_transform( - Transform( - translation=Vector3(1.0, 0.0, 0.0), - frame_id="world", + TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(base_time)), child_frame_id="robot", - ts=base_time, + transform=Transform(translation=Vector3(x=1.0, y=0.0, z=0.0)), ) ) def publish_after_delay() -> None: time.sleep(0.05) ttbuffer.receive_transform( - Transform( - translation=Vector3(0.0, 2.0, 0.0), - frame_id="robot", + TransformStamped( + header=Header(frame_id="robot", stamp=time_from_seconds(base_time)), child_frame_id="sensor", - ts=base_time, + transform=Transform(translation=Vector3(x=0.0, y=2.0, z=0.0)), ) ) @@ -409,19 +394,18 @@ def publish_after_delay() -> None: publisher.join() assert result is not None - assert result.translation.x == 1.0 - assert result.translation.y == 2.0 + assert result.transform.translation.x == 1.0 + assert result.transform.translation.y == 2.0 def test_get_transform_search_direct(self) -> None: ttbuffer = MultiTBuffer() base_time = time.time() # Add direct transform - transform = Transform( - translation=Vector3(1.0, 0.0, 0.0), - frame_id="world", + transform = TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(base_time)), child_frame_id="robot", - ts=base_time, + transform=Transform(translation=Vector3(x=1.0, y=0.0, z=0.0)), ) ttbuffer.receive_transform(transform) @@ -429,24 +413,22 @@ def test_get_transform_search_direct(self) -> None: result = ttbuffer.get_transform_search("world", "robot") assert result is not None assert len(result) == 1 - assert result[0].translation.x == 1.0 + assert result[0].transform.translation.x == 1.0 def test_get_transform_search_chain(self) -> None: ttbuffer = MultiTBuffer() base_time = time.time() # Create transform chain: world -> robot -> sensor - transform1 = Transform( - translation=Vector3(1.0, 0.0, 0.0), - frame_id="world", + transform1 = TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(base_time)), child_frame_id="robot", - ts=base_time, + transform=Transform(translation=Vector3(x=1.0, y=0.0, z=0.0)), ) - transform2 = Transform( - translation=Vector3(0.0, 2.0, 0.0), - frame_id="robot", + transform2 = TransformStamped( + header=Header(frame_id="robot", stamp=time_from_seconds(base_time)), child_frame_id="sensor", - ts=base_time, + transform=Transform(translation=Vector3(x=0.0, y=2.0, z=0.0)), ) ttbuffer.receive_transform(transform1, transform2) @@ -454,8 +436,8 @@ def test_get_transform_search_chain(self) -> None: result = ttbuffer.get_transform_search("world", "sensor") assert result is not None assert len(result) == 2 - assert result[0].translation.x == 1.0 # world -> robot - assert result[1].translation.y == 2.0 # robot -> sensor + assert result[0].transform.translation.x == 1.0 # world -> robot + assert result[1].transform.translation.y == 2.0 # robot -> sensor def test_get_transform_search_complex_chain(self) -> None: ttbuffer = MultiTBuffer() @@ -465,35 +447,30 @@ def test_get_transform_search_complex_chain(self) -> None: # world -> base -> arm -> hand # \-> robot -> sensor transforms = [ - Transform( - frame_id="world", + TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(base_time)), child_frame_id="base", - translation=Vector3(1.0, 0.0, 0.0), - ts=base_time, + transform=Transform(translation=Vector3(x=1.0, y=0.0, z=0.0)), ), - Transform( - frame_id="base", + TransformStamped( + header=Header(frame_id="base", stamp=time_from_seconds(base_time)), child_frame_id="arm", - translation=Vector3(0.0, 1.0, 0.0), - ts=base_time, + transform=Transform(translation=Vector3(x=0.0, y=1.0, z=0.0)), ), - Transform( - frame_id="arm", + TransformStamped( + header=Header(frame_id="arm", stamp=time_from_seconds(base_time)), child_frame_id="hand", - translation=Vector3(0.0, 0.0, 1.0), - ts=base_time, + transform=Transform(translation=Vector3(x=0.0, y=0.0, z=1.0)), ), - Transform( - frame_id="world", + TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(base_time)), child_frame_id="robot", - translation=Vector3(2.0, 0.0, 0.0), - ts=base_time, + transform=Transform(translation=Vector3(x=2.0, y=0.0, z=0.0)), ), - Transform( - frame_id="robot", + TransformStamped( + header=Header(frame_id="robot", stamp=time_from_seconds(base_time)), child_frame_id="sensor", - translation=Vector3(0.0, 2.0, 0.0), - ts=base_time, + transform=Transform(translation=Vector3(x=0.0, y=2.0, z=0.0)), ), ] @@ -513,8 +490,16 @@ def test_get_transform_search_no_path(self) -> None: base_time = time.time() # Create disconnected transforms - transform1 = Transform(frame_id="world", child_frame_id="robot", ts=base_time) - transform2 = Transform(frame_id="base", child_frame_id="sensor", ts=base_time) + transform1 = TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(base_time)), + child_frame_id="robot", + transform=Transform(), + ) + transform2 = TransformStamped( + header=Header(frame_id="base", stamp=time_from_seconds(base_time)), + child_frame_id="sensor", + transform=Transform(), + ) ttbuffer.receive_transform(transform1, transform2) # No path exists @@ -526,24 +511,22 @@ def test_get_transform_search_with_time(self) -> None: base_time = time.time() # Add transforms at different times - old_transform = Transform( - frame_id="world", + old_transform = TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(base_time - 10.0)), child_frame_id="robot", - translation=Vector3(1.0, 0.0, 0.0), - ts=base_time - 10.0, + transform=Transform(translation=Vector3(x=1.0, y=0.0, z=0.0)), ) - new_transform = Transform( - frame_id="world", + new_transform = TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(base_time)), child_frame_id="robot", - translation=Vector3(2.0, 0.0, 0.0), - ts=base_time, + transform=Transform(translation=Vector3(x=2.0, y=0.0, z=0.0)), ) ttbuffer.receive_transform(old_transform, new_transform) # Search at specific time result = ttbuffer.get_transform_search("world", "robot", time_point=base_time) assert result is not None - assert result[0].translation.x == 2.0 + assert result[0].transform.translation.x == 2.0 # Search with time tolerance result = ttbuffer.get_transform_search( @@ -559,10 +542,26 @@ def test_get_transform_search_shortest_path(self) -> None: # world -> A -> B -> target (3 hops) # world -> target (direct, 1 hop) transforms = [ - Transform(frame_id="world", child_frame_id="A", ts=base_time), - Transform(frame_id="A", child_frame_id="B", ts=base_time), - Transform(frame_id="B", child_frame_id="target", ts=base_time), - Transform(frame_id="world", child_frame_id="target", ts=base_time), + TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(base_time)), + child_frame_id="A", + transform=Transform(), + ), + TransformStamped( + header=Header(frame_id="A", stamp=time_from_seconds(base_time)), + child_frame_id="B", + transform=Transform(), + ), + TransformStamped( + header=Header(frame_id="B", stamp=time_from_seconds(base_time)), + child_frame_id="target", + transform=Transform(), + ), + TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(base_time)), + child_frame_id="target", + transform=Transform(), + ), ] for t in transforms: @@ -586,11 +585,10 @@ def test_string_representations(self) -> None: buffer = TBuffer() base_time = time.time() for i in range(3): - transform = Transform( - translation=Vector3(float(i), 0.0, 0.0), - frame_id="world", + transform = TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(base_time + i * 0.1)), child_frame_id="robot", - ts=base_time + i * 0.1, + transform=Transform(translation=Vector3(x=float(i), y=0.0, z=0.0)), ) buffer.add(transform) @@ -602,9 +600,21 @@ def test_string_representations(self) -> None: # Test MultiTBuffer with multiple frame pairs ttbuffer = MultiTBuffer() transforms = [ - Transform(frame_id="world", child_frame_id="robot1", ts=base_time), - Transform(frame_id="world", child_frame_id="robot2", ts=base_time + 0.5), - Transform(frame_id="robot1", child_frame_id="sensor", ts=base_time + 1.0), + TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(base_time)), + child_frame_id="robot1", + transform=Transform(), + ), + TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(base_time + 0.5)), + child_frame_id="robot2", + transform=Transform(), + ), + TransformStamped( + header=Header(frame_id="robot1", stamp=time_from_seconds(base_time + 1.0)), + child_frame_id="sensor", + transform=Transform(), + ), ] for t in transforms: @@ -626,12 +636,11 @@ def test_string_representations(self) -> None: T0 = 1_700_000_000.0 -def _t(parent: str, child: str, x: float, ts: float) -> Transform: - return Transform( - frame_id=parent, +def _t(parent: str, child: str, x: float, ts: float) -> TransformStamped: + return TransformStamped( + header=Header(frame_id=parent, stamp=time_from_seconds(ts)), child_frame_id=child, - translation=Vector3(x, 0.0, 0.0), - ts=ts, + transform=Transform(translation=Vector3(x=x, y=0.0, z=0.0)), ) @@ -649,7 +658,7 @@ def make_tf(request, tmp_path): # type: ignore[no-untyped-def] """Builder fixture: feed it transforms, get back a TFLookup over them.""" stores = [] - def build(*transforms: Transform) -> TFLookup: + def build(*transforms: TransformStamped) -> TFLookup: if request.param == "live": buf = MultiTBuffer() buf.receive_transform(*transforms) @@ -662,7 +671,7 @@ def build(*transforms: Transform) -> TFLookup: stores.append(store) stream = store.stream("tf", TFMessage) for t in transforms: - stream.append(TFMessage(t), ts=t.ts, pose=None) + stream.append(TFMessage(transforms=[t]), ts=to_seconds(t.header.stamp), pose=None) return StreamTF(store.stream("tf", TFMessage)) yield build @@ -677,20 +686,20 @@ def test_latest(self, make_tf) -> None: # type: ignore[no-untyped-def] tf = make_tf(*(_t("world", "robot", float(i), T0 + i * 0.1) for i in range(3))) got = tf.get("world", "robot") assert got is not None - assert got.translation.x == 2.0 + assert got.transform.translation.x == 2.0 def test_nearest_in_time(self, make_tf) -> None: # type: ignore[no-untyped-def] tf = make_tf(*(_t("world", "robot", float(i), T0 + i * 0.5) for i in range(5))) got = tf.get("world", "robot", time_point=T0 + 1.25) assert got is not None # Equidistant between i=2 (t=1.0) and i=3 (t=1.5) — the later one wins. - assert got.translation.x == 3.0 + assert got.transform.translation.x == 3.0 def test_inverse(self, make_tf) -> None: # type: ignore[no-untyped-def] tf = make_tf(_t("world", "robot", 5.0, T0)) got = tf.get("robot", "world", time_point=T0) assert got is not None - assert got.translation.x == pytest.approx(-5.0) + assert got.transform.translation.x == pytest.approx(-5.0) def test_time_tolerance(self, make_tf) -> None: # type: ignore[no-untyped-def] tf = make_tf(_t("world", "robot", 1.0, T0)) @@ -701,10 +710,10 @@ def test_same_frame_identity(self, make_tf) -> None: # type: ignore[no-untyped- tf = make_tf(_t("world", "robot", 1.0, T0)) got = tf.get("world", "world", time_point=T0) assert got is not None - assert got.frame_id == "world" + assert got.header.frame_id == "world" assert got.child_frame_id == "world" - assert got.translation.x == 0.0 - assert got.rotation.w == 1.0 + assert got.transform.translation.x == 0.0 + assert got.transform.rotation.w == 1.0 def test_unknown_frame(self, make_tf) -> None: # type: ignore[no-untyped-def] tf = make_tf(_t("world", "robot", 1.0, T0)) @@ -714,29 +723,33 @@ def test_chain_composition(self, make_tf) -> None: # type: ignore[no-untyped-de # world -> robot: translate (1, 0, 0); robot -> sensor: translate # (0, 2, 0) and rotate 90° around Z. tf = make_tf( - Transform( - translation=Vector3(1.0, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="world", + TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(T0)), child_frame_id="robot", - ts=T0, + transform=Transform( + translation=Vector3(x=1.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), - Transform( - translation=Vector3(0.0, 2.0, 0.0), - rotation=Quaternion(0.0, 0.0, math.sin(math.pi / 4), math.cos(math.pi / 4)), - frame_id="robot", + TransformStamped( + header=Header(frame_id="robot", stamp=time_from_seconds(T0)), child_frame_id="sensor", - ts=T0, + transform=Transform( + translation=Vector3(x=0.0, y=2.0, z=0.0), + rotation=Quaternion( + x=0.0, y=0.0, z=math.sin(math.pi / 4), w=math.cos(math.pi / 4) + ), + ), ), ) result = tf.get("world", "sensor", time_point=T0) assert result is not None - assert result.translation.x == pytest.approx(1.0) - assert result.translation.y == pytest.approx(2.0) - assert result.translation.z == pytest.approx(0.0) - assert result.rotation.z == pytest.approx(math.sin(math.pi / 4)) - assert result.rotation.w == pytest.approx(math.cos(math.pi / 4)) - assert result.frame_id == "world" + assert result.transform.translation.x == pytest.approx(1.0) + assert result.transform.translation.y == pytest.approx(2.0) + assert result.transform.translation.z == pytest.approx(0.0) + assert result.transform.rotation.z == pytest.approx(math.sin(math.pi / 4)) + assert result.transform.rotation.w == pytest.approx(math.cos(math.pi / 4)) + assert result.header.frame_id == "world" assert result.child_frame_id == "sensor" def test_chain_with_sparse_static_edge(self, make_tf) -> None: # type: ignore[no-untyped-def] @@ -748,7 +761,7 @@ def test_chain_with_sparse_static_edge(self, make_tf) -> None: # type: ignore[n ) got = tf.get("world", "base", time_point=T0 + 9.0) assert got is not None - assert got.translation.x == pytest.approx(109.0) + assert got.transform.translation.x == pytest.approx(109.0) def test_conforms_to_tf_lookup(self, make_tf) -> None: # type: ignore[no-untyped-def] assert isinstance(make_tf(_t("world", "robot", 1.0, T0)), TFLookup) @@ -763,10 +776,10 @@ def store(self): # type: ignore[no-untyped-def] stream = store.stream("tf", TFMessage) # Startup burst: a static stamped exactly at the stream start (regression: # strict `ts >` range queries used to drop it), then dynamic map→base. - stream.append(TFMessage(_t("world", "map", 100.0, T0)), ts=T0, pose=None) + stream.append(TFMessage(transforms=[_t("world", "map", 100.0, T0)]), ts=T0, pose=None) for i in range(100): ts = T0 + i / 10 - stream.append(TFMessage(_t("map", "base", i / 10, ts)), ts=ts, pose=None) + stream.append(TFMessage(transforms=[_t("map", "base", i / 10, ts)]), ts=ts, pose=None) yield store store.stop() @@ -803,7 +816,7 @@ def test_get_pose(self, store) -> None: # type: ignore[no-untyped-def] tf = StreamTF(store.stream("tf", TFMessage)) pose = tf.get_pose("map", "base", time_point=T0 + 5.0) assert pose is not None - assert pose.position.x == pytest.approx(5.0) + assert pose.pose.position.x == pytest.approx(5.0) def test_cache_prefetch_and_eviction(self) -> None: # 40 s of data with a small cache_span: a miss caches the query window @@ -813,25 +826,69 @@ def test_cache_prefetch_and_eviction(self) -> None: stream = store.stream("tf", TFMessage) for i in range(400): ts = T0 + i / 10 - stream.append(TFMessage(_t("map", "base", i / 10, ts)), ts=ts, pose=None) + stream.append(TFMessage(transforms=[_t("map", "base", i / 10, ts)]), ts=ts, pose=None) try: tf = StreamTF(store.stream("tf", TFMessage), cache_span=2.0) got = tf.get("map", "base", time_point=T0 + 35.0, time_tolerance=0.5) assert got is not None - assert got.translation.x == pytest.approx(35.0) + assert got.transform.translation.x == pytest.approx(35.0) covered = tf._covered assert covered == pytest.approx((T0 + 34.5, T0 + 37.5)) # Inside the prefetched span: served from cache, no re-query. got = tf.get("map", "base", time_point=T0 + 37.0, time_tolerance=0.5) assert got is not None - assert got.translation.x == pytest.approx(37.0) + assert got.transform.translation.x == pytest.approx(37.0) assert tf._covered == covered # Past the span: evict and re-cache around the new query. got = tf.get("map", "base", time_point=T0 + 39.0, time_tolerance=0.5) assert got is not None - assert got.translation.x == pytest.approx(39.0) + assert got.transform.translation.x == pytest.approx(39.0) assert tf._covered == pytest.approx((T0 + 38.5, T0 + 41.5)) # Bounded: the buffer holds the re-cached span, not the stream. assert len(tf.buffers[("map", "base")]) < 100 finally: store.stop() + + +def test_buffer_retains_distinct_nanoseconds_and_owns_its_snapshots(): + buffer = TBuffer() + first = TransformStamped( + header=Header(frame_id="a", stamp=Time(sec=1700000000, nanosec=123456789)), + child_frame_id="b", + ) + second = TransformStamped( + header=Header(frame_id="a", stamp=Time(sec=1700000000, nanosec=123456790)), + child_frame_id="b", + ) + buffer.add(first) + buffer.add(second) + first.header.stamp.nanosec = 0 + assert len(buffer) == 2 + stored_first = buffer.first() + assert stored_first is not None + assert stored_first.header.stamp.nanosec == 123456789 + stored_first.header.stamp.nanosec = 1 + unchanged = buffer.first() + assert unchanged is not None + assert unchanged.header.stamp.nanosec == 123456789 + latest = buffer.get() + assert latest is not None + assert latest.header.stamp.nanosec == 123456790 + + +def test_buffer_drops_out_of_order_data_older_than_retained_window(): + buffer = TBuffer(buffer_size=1) + for sec in (10, 11, 1): + buffer.add(TransformStamped(header=Header(stamp=Time(sec=sec)))) + assert len(buffer) == 2 + first = buffer.first() + assert first is not None + assert first.header.stamp.sec == 10 + + +def test_buffer_preserves_negative_stamp(): + buffer = TBuffer() + buffer.add(TransformStamped(header=Header(stamp=Time(sec=-1, nanosec=500000000)))) + result = buffer.get(time_point=-0.5, time_tolerance=0) + assert result is not None + assert result.header.stamp == Time(sec=-1, nanosec=500000000) diff --git a/dimos/protocol/tf/tf.py b/dimos/protocol/tf/tf.py index 03a4706deda..87afac6b4cc 100644 --- a/dimos/protocol/tf/tf.py +++ b/dimos/protocol/tf/tf.py @@ -15,17 +15,23 @@ # limitations under the License. from collections import deque +from copy import copy from functools import reduce +import math +import subprocess import threading import time -from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Any, Protocol, cast, runtime_checkable -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos_generated.geometry_msgs.msg import PoseStamped, TransformStamped +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage +from sortedcontainers import SortedDict # type: ignore[import-untyped] + +from dimos.msgs.geometry import compose_transforms, inverse_transform, pose_from_transform +from dimos.msgs.time import time_from_seconds, to_nanoseconds, to_seconds from dimos.types.timestamped import to_human_readable from dimos.utils.logging_config import setup_logger -from dimos.utils.timeseries.inmemory import InMemoryStore if TYPE_CHECKING: from collections.abc import Callable @@ -48,45 +54,68 @@ def get( child_frame: str, time_point: float | None = None, time_tolerance: float | None = None, - ) -> Transform | None: ... + ) -> TransformStamped | None: ... + +class TBuffer: + """A bounded time index of generated transforms, keyed by integer nanoseconds.""" -class TBuffer(InMemoryStore[Transform]): def __init__(self, buffer_size: float = 10.0) -> None: - super().__init__() + if math.isnan(buffer_size) or buffer_size <= 0: + raise ValueError("buffer_size must be positive") self.buffer_size = buffer_size + self._entries: SortedDict = SortedDict() + + def __len__(self) -> int: + return len(self._entries) + + def add(self, transform: TransformStamped) -> None: + stamp = to_nanoseconds(transform.header.stamp) + self._entries[stamp] = TransformStamped( + header=transform.header, + child_frame_id=transform.child_frame_id, + transform=transform.transform, + ) + if math.isfinite(self.buffer_size): + oldest = self._entries.peekitem(-1)[0] - round(self.buffer_size * 1_000_000_000) + while self._entries and self._entries.peekitem(0)[0] < oldest: + self._entries.popitem(0) + + def first(self) -> TransformStamped | None: + return copy(self._entries.peekitem(0)[1]) if self._entries else None - def add(self, transform: Transform) -> None: - self.save(transform) - self.prune_old(transform.ts - self.buffer_size) + def last(self) -> TransformStamped | None: + return copy(self._entries.peekitem(-1)[1]) if self._entries else None - def get(self, time_point: float | None = None, time_tolerance: float = 1.0) -> Transform | None: - """Get transform at specified time or latest if no time given.""" + def get( + self, time_point: float | None = None, time_tolerance: float = 1.0 + ) -> TransformStamped | None: if time_point is None: return self.last() - return self.find_closest(time_point, time_tolerance) + stamp = to_nanoseconds(time_from_seconds(time_point)) + if math.isnan(time_tolerance) or time_tolerance < 0: + raise ValueError("time_tolerance must be nonnegative") + index = self._entries.bisect_left(stamp) + candidates = [ + self._entries.peekitem(i) for i in (index - 1, index) if 0 <= i < len(self._entries) + ] + if not candidates: + return None + nearest, result = min(candidates, key=lambda item: (abs(item[0] - stamp), -item[0])) + if abs(nearest - stamp) / 1_000_000_000 > time_tolerance: + return None + return copy(cast("TransformStamped", result)) def __str__(self) -> str: - if len(self) == 0: + first, last = self.first(), self.last() + if first is None or last is None: return "TBuffer(empty)" - - first_item = self.first() - time_range = self.time_range() - if time_range and first_item: - start_time = to_human_readable(time_range[0]) - end_time = to_human_readable(time_range[1]) - duration = time_range[1] - time_range[0] - - frame_str = f"{first_item.frame_id} -> {first_item.child_frame_id}" - - return ( - f"TBuffer(" - f"{frame_str}, " - f"{len(self)} msgs, " - f"{duration:.2f}s [{start_time} - {end_time}])" - ) - - return f"TBuffer({len(self)} msgs)" + start, end = to_seconds(first.header.stamp), to_seconds(last.header.stamp) + return ( + f"TBuffer({first.header.frame_id} -> {first.child_frame_id}, " + f"{len(self)} msgs, {end - start:.2f}s " + f"[{to_human_readable(start)} - {to_human_readable(end)}])" + ) # stores multiple transform buffers @@ -97,10 +126,10 @@ def __init__(self, buffer_size: float = 10.0) -> None: self.buffer_size = buffer_size self._cv = threading.Condition() - def receive_transform(self, *args: Transform) -> None: + def receive_transform(self, *args: TransformStamped) -> None: with self._cv: for transform in args: - key = (transform.frame_id, transform.child_frame_id) + key = (transform.header.frame_id, transform.child_frame_id) if key not in self.buffers: self.buffers[key] = TBuffer(self.buffer_size) self.buffers[key].add(transform) @@ -134,12 +163,14 @@ def get_transform( child_frame: str, time_point: float | None = None, time_tolerance: float | None = None, - ) -> Transform | None: + ) -> TransformStamped | None: if parent_frame == child_frame: - return Transform( - frame_id=parent_frame, + return TransformStamped( + header=Header( + frame_id=parent_frame, + stamp=time_from_seconds(time_point if time_point is not None else time.time()), + ), child_frame_id=child_frame, - ts=time_point if time_point is not None else time.time(), ) # No explicit tolerance means "anything still buffered" — the buffer @@ -156,7 +187,7 @@ def get_transform( reverse_key = (child_frame, parent_frame) if reverse_key in self.buffers: transform = self.buffers[reverse_key].get(time_point, tolerance) - return transform.inverse() if transform else None + return inverse_transform(transform) if transform else None return None @@ -166,7 +197,7 @@ def _get( child_frame: str, time_point: float | None = None, time_tolerance: float | None = None, - ) -> Transform | None: + ) -> TransformStamped | None: with self._cv: simple = self.get_transform(parent_frame, child_frame, time_point, time_tolerance) @@ -180,7 +211,7 @@ def _get( if complex is None: return None - return reduce(lambda t1, t2: t1 + t2, complex) + return reduce(compose_transforms, complex) def _wait_get( self, @@ -189,7 +220,7 @@ def _wait_get( time_point: float | None, time_tolerance: float | None, forward_tolerance: float, - ) -> Transform | None: + ) -> TransformStamped | None: deadline = time.monotonic() + forward_tolerance with self._cv: while True: @@ -210,7 +241,7 @@ def get( *, forward_tolerance: float = 0.0, warn: bool = True, - ) -> Transform | None: + ) -> TransformStamped | None: result = self._get(parent_frame, child_frame, time_point, time_tolerance) if result is None and forward_tolerance > 0: result = self._wait_get( @@ -240,7 +271,7 @@ def get_pose( ) if not tf: return None - return tf.to_pose() + return pose_from_transform(tf) def get_transform_search( self, @@ -248,7 +279,7 @@ def get_transform_search( child_frame: str, time_point: float | None = None, time_tolerance: float | None = None, - ) -> list[Transform] | None: + ) -> list[TransformStamped] | None: """Search for shortest transform chain between parent and child frames using BFS.""" with self._cv: # Check if direct transform exists (already checked in get_transform, but for clarity) @@ -257,7 +288,7 @@ def get_transform_search( return [direct] # BFS to find shortest path - queue: deque[tuple[str, list[Transform]]] = deque([(parent_frame, [])]) + queue: deque[tuple[str, list[TransformStamped]]] = deque([(parent_frame, [])]) visited = {parent_frame} while queue: @@ -283,8 +314,6 @@ def get_transform_search( return None def graph(self) -> str: - import subprocess - def connection_str(connection: tuple[str, str]) -> str: (frame_from, frame_to) = connection return f"{frame_from} -> {frame_to}" @@ -335,10 +364,10 @@ def __init__(self, stream: Any | None = None, buffer_size: float = 10.0) -> None if stream is not None: self._unsubscribe = stream.subscribe(self.receive_tfmessage) - def publish(self, *transforms: Transform) -> None: + def publish(self, *transforms: TransformStamped) -> None: self.receive_transform(*transforms) if self._stream is not None: - self._stream.publish(TFMessage(*transforms)) + self._stream.publish(TFMessage(transforms=list(transforms))) def dispose(self) -> None: if self._unsubscribe is not None: diff --git a/dimos/robot/assembly/mid360_realsense_30.py b/dimos/robot/assembly/mid360_realsense_30.py index 421d361ce22..877f2c36a92 100644 --- a/dimos/robot/assembly/mid360_realsense_30.py +++ b/dimos/robot/assembly/mid360_realsense_30.py @@ -46,6 +46,9 @@ import math +from dimos_generated.geometry_msgs.msg import TransformStamped +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 + from dimos.core.coordination.blueprints import autoconnect from dimos.core.stream import In from dimos.hardware.sensors.camera.realsense.camera import RealSenseCamera @@ -53,10 +56,6 @@ from dimos.hardware.sensors.lidar.pointlio.pointlio_blueprints import mid360_for_pointlio from dimos.hardware.sensors.lidar.pointlio.recorder import PointlioRecorder from dimos.hardware.sensors.lidar.virtual_mid360.recorder import Mid360PcapRecorder -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.protocol.tf.static_tf_publisher import ( FrameSpec, StaticTfPublisher, @@ -92,7 +91,7 @@ class Mid360RealsenseStaticTf(StaticTfPublisher): """Publishes the RealSense/Mid-360 mount tree onto tf on a fixed interval.""" - def transforms(self) -> list[Transform]: + def transforms(self) -> list[TransformStamped]: return frames_to_edge_transforms(FRAMES) diff --git a/dimos/robot/deeprobotics/m20/camera.py b/dimos/robot/deeprobotics/m20/camera.py index 4e0b18957f1..c7db9a6d594 100644 --- a/dimos/robot/deeprobotics/m20/camera.py +++ b/dimos/robot/deeprobotics/m20/camera.py @@ -19,17 +19,18 @@ from typing import Any import av +from dimos_generated.foxglove_msgs.msg import CompressedVideo +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import Out -from dimos.msgs.foxglove_msgs.CompressedVideo import CompressedVideo -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.camera_info import camera_info_from_intrinsics, camera_info_with_stamp +from dimos.msgs.time import time_from_seconds from dimos.robot.deeprobotics.m20.constants import ( FRONT_CAMERA_RTSP_URL, REAR_CAMERA_RTSP_URL, @@ -44,27 +45,27 @@ # factory calibration. _FRONT_CAMERA_XYZ = (0.37646, 0.0, 0.03738) _REAR_CAMERA_XYZ = (-0.37646, 0.0, 0.03738) -_OPTICAL_ROT = Quaternion(-0.5, 0.5, -0.5, 0.5) +_OPTICAL_ROT = Quaternion(x=-0.5, y=0.5, z=-0.5, w=0.5) _CAMERA_WIDTH = 800 _CAMERA_HEIGHT = 600 _CAMERA_FOCAL_LENGTH = 607.0 * _CAMERA_WIDTH / 1280.0 -_FRONT_CAMERA_INFO = CameraInfo.from_intrinsics( +_FRONT_CAMERA_INFO = camera_info_from_intrinsics( fx=_CAMERA_FOCAL_LENGTH, fy=_CAMERA_FOCAL_LENGTH, cx=_CAMERA_WIDTH * 0.5, cy=_CAMERA_HEIGHT * 0.5, width=_CAMERA_WIDTH, height=_CAMERA_HEIGHT, - frame_id="front_camera_optical", + header=Header(frame_id="front_camera_optical"), ) -_REAR_CAMERA_INFO = CameraInfo.from_intrinsics( +_REAR_CAMERA_INFO = camera_info_from_intrinsics( fx=_CAMERA_FOCAL_LENGTH, fy=_CAMERA_FOCAL_LENGTH, cx=_CAMERA_WIDTH * 0.5, cy=_CAMERA_HEIGHT * 0.5, width=_CAMERA_WIDTH, height=_CAMERA_HEIGHT, - frame_id="rear_camera_optical", + header=Header(frame_id="rear_camera_optical"), ) @@ -115,35 +116,54 @@ def _publish_camera_metadata(self) -> None: now = time.time() self.tf.publish( TFMessage( - Transform( - translation=Vector3(*_FRONT_CAMERA_XYZ), - frame_id="base_link", - child_frame_id="front_camera_link", - ts=now, - ), - Transform( - rotation=_OPTICAL_ROT, - frame_id="front_camera_link", - child_frame_id="front_camera_optical", - ts=now, - ), - Transform( - translation=Vector3(*_REAR_CAMERA_XYZ), - rotation=Quaternion(0.0, 0.0, 1.0, 0.0), - frame_id="base_link", - child_frame_id="rear_camera_link", - ts=now, - ), - Transform( - rotation=_OPTICAL_ROT, - frame_id="rear_camera_link", - child_frame_id="rear_camera_optical", - ts=now, - ), + transforms=[ + TransformStamped( + header=Header(frame_id="base_link", stamp=time_from_seconds(now)), + child_frame_id="front_camera_link", + transform=Transform( + translation=Vector3( + x=_FRONT_CAMERA_XYZ[0], + y=_FRONT_CAMERA_XYZ[1], + z=_FRONT_CAMERA_XYZ[2], + ), + rotation=Quaternion(w=1.0), + ), + ), + TransformStamped( + header=Header( + frame_id="front_camera_link", stamp=time_from_seconds(now) + ), + child_frame_id="front_camera_optical", + transform=Transform(rotation=_OPTICAL_ROT), + ), + TransformStamped( + header=Header(frame_id="base_link", stamp=time_from_seconds(now)), + child_frame_id="rear_camera_link", + transform=Transform( + translation=Vector3( + x=_REAR_CAMERA_XYZ[0], + y=_REAR_CAMERA_XYZ[1], + z=_REAR_CAMERA_XYZ[2], + ), + rotation=Quaternion(z=1.0), + ), + ), + TransformStamped( + header=Header( + frame_id="rear_camera_link", stamp=time_from_seconds(now) + ), + child_frame_id="rear_camera_optical", + transform=Transform(rotation=_OPTICAL_ROT), + ), + ] ) ) - self.front_camera_info.publish(_FRONT_CAMERA_INFO.with_ts(now)) - self.rear_camera_info.publish(_REAR_CAMERA_INFO.with_ts(now)) + self.front_camera_info.publish( + camera_info_with_stamp(_FRONT_CAMERA_INFO, time_from_seconds(now)) + ) + self.rear_camera_info.publish( + camera_info_with_stamp(_REAR_CAMERA_INFO, time_from_seconds(now)) + ) self._stop_event.wait(1.0) def _relay(self, url: str, output: Out[CompressedVideo], camera_name: str) -> None: @@ -168,10 +188,10 @@ def _relay(self, url: str, output: Out[CompressedVideo], camera_name: str) -> No continue output.publish( CompressedVideo( - bytes(packet), + data=bytes(packet), format="h265", frame_id="", - ts=time.time(), + timestamp=time_from_seconds(time.time()), ) ) except (av.FFmpegError, IndexError, ValueError) as exc: diff --git a/dimos/robot/deeprobotics/m20/connection.py b/dimos/robot/deeprobotics/m20/connection.py index f820df24fc4..2e9c076b1c6 100644 --- a/dimos/robot/deeprobotics/m20/connection.py +++ b/dimos/robot/deeprobotics/m20/connection.py @@ -23,18 +23,15 @@ from threading import Condition, RLock from typing import Any, Literal +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.std_msgs.msg import Bool, Int32, UInt32 from pydantic import Field from reactivex.disposable import Disposable from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.std_msgs.Bool import Bool -from dimos.msgs.std_msgs.Int32 import Int32 -from dimos.msgs.std_msgs.UInt32 import UInt32 from dimos.utils.logging_config import setup_logger from dimos.utils.sequential_ids import SequentialIds @@ -139,7 +136,7 @@ def move(self, twist: Twist, duration: float = 0.0) -> bool: @rpc def stop_movement(self) -> None: """Publish an immediate zero velocity without disabling future commands.""" - self.cmd_vel.publish(Twist.zero()) + self.cmd_vel.publish(Twist()) @rpc def standup(self) -> bool: @@ -169,7 +166,7 @@ def standup(self) -> bool: def liedown(self) -> bool: """Disable velocity output and command the M20 to its Sit/prone state.""" self.stop_movement() - self.motion_state_cmd.publish(Int32(MOTION_SIT)) + self.motion_state_cmd.publish(Int32(data=MOTION_SIT)) return True @rpc @@ -221,7 +218,7 @@ def _on_gait_state(self, msg: UInt32) -> None: self._state_condition.notify_all() def _on_odometry(self, odometry: Odometry) -> None: - self.odom.publish(odometry.to_pose_stamped()) + self.odom.publish(PoseStamped(header=odometry.header, pose=odometry.pose.pose)) def _ensure_rl_control(self) -> bool: with self._lock: @@ -229,18 +226,18 @@ def _ensure_rl_control(self) -> bool: if motion_state == MOTION_RL_CONTROL: return True if motion_state != MOTION_STAND: - self.motion_state_cmd.publish(Int32(MOTION_STAND)) + self.motion_state_cmd.publish(Int32(data=MOTION_STAND)) if not self._wait_for_motion_state(MOTION_STAND, self.config.stand_timeout_s): logger.error("M20 did not confirm Stand before the transition timeout") return False - self.motion_state_cmd.publish(Int32(MOTION_RL_CONTROL)) + self.motion_state_cmd.publish(Int32(data=MOTION_RL_CONTROL)) if not self._wait_for_motion_state(MOTION_RL_CONTROL, self.config.rl_control_timeout_s): logger.error("M20 did not confirm RL Control after standing") return False return True def _set_gait_and_wait(self, gait: int) -> bool: - self.gait_cmd.publish(UInt32(gait)) + self.gait_cmd.publish(UInt32(data=gait)) return self._wait_for_gait_state(gait, self.config.gait_timeout_s) def _enter_navigation_mode(self) -> bool: diff --git a/dimos/robot/deeprobotics/m20/test_camera_messages.py b/dimos/robot/deeprobotics/m20/test_camera_messages.py new file mode 100644 index 00000000000..89a9a8694c3 --- /dev/null +++ b/dimos/robot/deeprobotics/m20/test_camera_messages.py @@ -0,0 +1,52 @@ +# Copyright 2025-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. + +from threading import Event +from types import SimpleNamespace + +from dimos_generated.sensor_msgs.msg import CameraInfo +from dimos_generated.tf2_msgs.msg import TFMessage + +from dimos.robot.deeprobotics.m20.camera import M20CameraRelay + + +class _OneIteration(Event): + def wait(self, timeout=None): + self.set() + return True + + +def test_camera_metadata_uses_generated_wire_values_without_connecting_streams(): + transforms, front, rear = [], [], [] + proxy = SimpleNamespace( + _stop_event=_OneIteration(), + tf=SimpleNamespace(publish=transforms.append), + front_camera_info=SimpleNamespace(publish=front.append), + rear_camera_info=SimpleNamespace(publish=rear.append), + ) + M20CameraRelay._publish_camera_metadata(proxy) + tf = TFMessage.decode(transforms[0].encode()) + assert len(tf.transforms) == 4 + assert [t.child_frame_id for t in tf.transforms] == [ + "front_camera_link", + "front_camera_optical", + "rear_camera_link", + "rear_camera_optical", + ] + first, second = CameraInfo.decode(front[0].encode()), CameraInfo.decode(rear[0].encode()) + assert first.header.frame_id == "front_camera_optical" + assert second.header.frame_id == "rear_camera_optical" + assert first.header.stamp == second.header.stamp == tf.transforms[0].header.stamp + assert first.width == second.width == 800 + assert first.height == second.height == 600 diff --git a/dimos/robot/deeprobotics/m20/test_connection_messages.py b/dimos/robot/deeprobotics/m20/test_connection_messages.py new file mode 100644 index 00000000000..1c60f925b38 --- /dev/null +++ b/dimos/robot/deeprobotics/m20/test_connection_messages.py @@ -0,0 +1,53 @@ +# 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. + +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseWithCovariance, + Quaternion, + Twist, + Vector3, +) +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.std_msgs.msg import Bool, Header + +from dimos.robot.deeprobotics.m20.connection import M20Connection + + +def test_generated_forwarding_preserves_readiness_pose_clock_and_zero_stop(monkeypatch): + module = M20Connection() + velocities, poses = [], [] + monkeypatch.setattr(module.cmd_vel, "publish", velocities.append) + monkeypatch.setattr(module.odom, "publish", poses.append) + command = Twist(linear=Vector3(x=0.1), angular=Vector3(z=0.2)) + try: + assert module.move(command) is False + module._on_command_ready(Bool(data=True)) + assert module.move(command) is True + source = Odometry( + header=Header(frame_id="map", stamp=Time(sec=1, nanosec=123456789)), + pose=PoseWithCovariance(pose=Pose(position=Point(x=2), orientation=Quaternion(w=1))), + ) + module._on_odometry(source) + assert poses[0].header == source.header and poses[0].pose == source.pose.pose + module.stop_movement() + assert velocities[-1] == Twist() + assert Twist.decode(velocities[0].encode()) == command + finally: + module.stop() diff --git a/dimos/robot/diy/alfred/config.py b/dimos/robot/diy/alfred/config.py index cef7bf0e7b9..09c6295cd77 100644 --- a/dimos/robot/diy/alfred/config.py +++ b/dimos/robot/diy/alfred/config.py @@ -18,9 +18,7 @@ from pathlib import Path from typing import Any -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos_generated.geometry_msgs.msg import Point, Pose DEFAULT_ADDRESS = "172.6.2.20:11323" ALFRED_URDF = Path(__file__).resolve().parent / "alfred.urdf" @@ -52,6 +50,6 @@ class AlfredConfig: d455_serial="260922302422", internal_odom_offsets={ # Mid-360 lidar: a bit forward, and a bit to the right of base center, above ground. - "mid360_link": Pose(0.20, -0.20, 0.30, *Quaternion.from_euler(Vector3(0, 0, 0))), + "mid360_link": Pose(position=Point(x=0.20, y=-0.20, z=0.30)), }, ) diff --git a/dimos/robot/diy/alfred/effector_high_level.py b/dimos/robot/diy/alfred/effector_high_level.py index aeff8a199ed..07619405d43 100644 --- a/dimos/robot/diy/alfred/effector_high_level.py +++ b/dimos/robot/diy/alfred/effector_high_level.py @@ -36,6 +36,16 @@ import time from typing import TYPE_CHECKING, Any +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseWithCovariance, + Twist, + TwistWithCovariance, + Vector3, +) +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.std_msgs.msg import Header import numpy as np from pydantic import Field, FiniteFloat @@ -44,11 +54,8 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.time import time_from_seconds from dimos.robot.diy.alfred.config import DEFAULT_ADDRESS from dimos.utils.logging_config import setup_logger @@ -171,7 +178,7 @@ async def move_velocity( self, x: float, y: float = 0.0, yaw: float = 0.0, duration: float = 0.0 ) -> str: """Move the Alfred at the given velocity for ``duration`` seconds.""" - twist = Twist(linear=Vector3(x, y, 0), angular=Vector3(0, 0, yaw)) + twist = Twist(linear=Vector3(x=x, y=y), angular=Vector3(z=yaw)) await self.move(twist, duration=duration) return f"Started moving with velocity=({x}, {y}, {yaw}) for {duration} seconds" @@ -218,21 +225,24 @@ async def _poll_wheel_odometry(self, client: portal.Client) -> None: forward = math.cos(heading) * dx + math.sin(heading) * dy left = -math.sin(heading) * dx + math.cos(heading) * dy twist = Twist( - linear=Vector3(forward / dt, left / dt, 0.0), - angular=Vector3(0.0, 0.0, turn / dt), + linear=Vector3(x=forward / dt, y=left / dt), + angular=Vector3(z=turn / dt), ) previous = (ts, x, y, yaw) self.wheel_odometry.publish( Odometry( - ts=ts, - frame_id=self.config.wheel_odom_frame_id, + header=Header( + stamp=time_from_seconds(ts), frame_id=self.config.wheel_odom_frame_id + ), child_frame_id=self.config.base_frame_id, - pose=Pose( - position=Vector3(x, y, 0.0), - orientation=Quaternion.from_euler(Vector3(0.0, 0.0, yaw)), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=x, y=y), + orientation=quaternion_from_euler(0.0, 0.0, yaw), + ) ), - twist=twist, + twist=TwistWithCovariance(twist=twist), ) ) except asyncio.CancelledError: diff --git a/dimos/robot/diy/alfred/mount_tf.py b/dimos/robot/diy/alfred/mount_tf.py index 660ffe15591..47962f957c0 100644 --- a/dimos/robot/diy/alfred/mount_tf.py +++ b/dimos/robot/diy/alfred/mount_tf.py @@ -25,14 +25,15 @@ from xml.etree import ElementTree -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos_generated.geometry_msgs.msg import Transform, TransformStamped, Vector3 +from dimos_generated.std_msgs.msg import Header + +from dimos.msgs.geometry import quaternion_from_euler from dimos.protocol.tf.static_tf_publisher import StaticTfPublisher from dimos.robot.diy.alfred.config import ALFRED_URDF -def mount_transforms() -> list[Transform]: +def mount_transforms() -> list[TransformStamped]: """One transform per fixed joint of the urdf, minus the imager frames. The drivers publish their own imager offsets from the factory extrinsics read @@ -52,11 +53,13 @@ def mount_transforms() -> list[Transform]: translation = [float(value) for value in origin.attrib["xyz"].split()] rpy = [float(value) for value in origin.attrib["rpy"].split()] transforms.append( - Transform( - translation=Vector3(*translation), - rotation=Quaternion.from_euler(Vector3(*rpy)), - frame_id=parent.attrib["link"], + TransformStamped( + header=Header(frame_id=parent.attrib["link"]), child_frame_id=child_link, + transform=Transform( + translation=Vector3(x=translation[0], y=translation[1], z=translation[2]), + rotation=quaternion_from_euler(*rpy), + ), ) ) return transforms @@ -65,5 +68,5 @@ def mount_transforms() -> list[Transform]: class AlfredMountTf(StaticTfPublisher): """Publishes Alfred's urdf mount tree onto tf on a fixed interval.""" - def transforms(self) -> list[Transform]: + def transforms(self) -> list[TransformStamped]: return mount_transforms() diff --git a/dimos/robot/diy/alfred/test_alfred.py b/dimos/robot/diy/alfred/test_alfred.py index f4b32d360a3..56025344ba8 100644 --- a/dimos/robot/diy/alfred/test_alfred.py +++ b/dimos/robot/diy/alfred/test_alfred.py @@ -26,7 +26,7 @@ def test_every_camera_frame_dim_slam_wants_is_rooted_at_base_link(): """cuVSLAM places no camera at all until every configured frame resolves against base_link, so one missing mount edge drops every image instead of degrading.""" - parents = {t.child_frame_id: t.frame_id for t in mount_transforms()} + parents = {t.child_frame_id: t.header.frame_id for t in mount_transforms()} for frame in [*IR_ENTITY_BY_FRAME, DEPTH_FRAME]: # The driver publishes the imager leaves off its own link; the mount tree owes # that link a path to base_link. diff --git a/dimos/robot/diy/alfred/test_generated_effector.py b/dimos/robot/diy/alfred/test_generated_effector.py new file mode 100644 index 00000000000..2bc95b58835 --- /dev/null +++ b/dimos/robot/diy/alfred/test_generated_effector.py @@ -0,0 +1,76 @@ +# 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. + +from types import SimpleNamespace +from unittest.mock import MagicMock + +from dimos_generated.nav_msgs.msg import Odometry +import pytest + +from dimos.msgs.geometry import yaw +from dimos.robot.diy.alfred.effector_high_level import AlfredHighLevel + + +@pytest.mark.asyncio +async def test_wheel_odometry_generated_frame_inversion_and_local_velocity( + monkeypatch: pytest.MonkeyPatch, +) -> None: + module = AlfredHighLevel() + readings = iter( + [ + {"translation": [0.0, 0.0], "rotation": 0.0}, + {"translation": [1.0, -2.0], "rotation": 0.0}, + ] + ) + client = MagicMock() + + def get_odometry(_request: object) -> SimpleNamespace: + reading = next(readings) + return SimpleNamespace(done=lambda: True, result=lambda: reading) + + client.get_odometry.side_effect = get_odometry + stamps = iter([100.0, 101.0]) + monkeypatch.setattr( + "dimos.robot.diy.alfred.effector_high_level.time", + SimpleNamespace(time=lambda: next(stamps)), + ) + received: list[Odometry] = [] + + def publish(value: Odometry) -> None: + received.append(Odometry.decode(value.encode())) + if len(received) == 2: + module._odometry_stop.set() + + module.wheel_odometry = MagicMock() + module.wheel_odometry.publish.side_effect = publish + + async def wait_or_stop(_seconds: float) -> None: + return None + + monkeypatch.setattr(module, "_wait_or_stop", wait_or_stop) + await module._poll_wheel_odometry(client) + assert len(received) == 2 + first, last = received + assert first.header.stamp.sec == 100 + assert last.header.stamp.sec == 101 + assert last.header.frame_id == "wheel_odom" + assert last.child_frame_id == "base_link" + assert (last.pose.pose.position.x, last.pose.pose.position.y) == (1.0, 2.0) + assert yaw(last.pose.pose.orientation) == 0.0 + assert (last.twist.twist.linear.x, last.twist.twist.linear.y, last.twist.twist.angular.z) == ( + 1.0, + 2.0, + 0.0, + ) + client.set_target_velocity.assert_not_called() diff --git a/dimos/robot/diy/alfred/test_mount_tf.py b/dimos/robot/diy/alfred/test_mount_tf.py new file mode 100644 index 00000000000..758c3a1a37e --- /dev/null +++ b/dimos/robot/diy/alfred/test_mount_tf.py @@ -0,0 +1,44 @@ +# 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. + +"""URDF mount extraction emits generated transforms and excludes driver-owned frames.""" + +from dimos_generated.tf2_msgs.msg import TFMessage + +from dimos.robot.diy.alfred import mount_tf + + +def test_mounts_preserve_parent_and_translation_and_skip_imager_frames(tmp_path, monkeypatch): + urdf = tmp_path / "robot.urdf" + urdf.write_text(""" + + + + + + + + + """) + monkeypatch.setattr(mount_tf, "ALFRED_URDF", urdf) + + edges = mount_tf.mount_transforms() + decoded = TFMessage.decode(TFMessage(transforms=edges).encode()) + + assert len(decoded.transforms) == 1 + edge = decoded.transforms[0] + assert (edge.header.frame_id, edge.child_frame_id) == ("base", "lidar_link") + translation = edge.transform.translation + assert (translation.x, translation.y, translation.z) == (1, 2, 3) + assert edge.transform.rotation.w == 1 diff --git a/dimos/robot/drone/camera_module.py b/dimos/robot/drone/camera_module.py index 8f2b11631f0..453a1e1ad02 100644 --- a/dimos/robot/drone/camera_module.py +++ b/dimos/robot/drone/camera_module.py @@ -21,16 +21,18 @@ import time from typing import Any -from dimos_lcm.sensor_msgs import CameraInfo +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.std_msgs.Header import Header -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.geometry import pose_from_transform +from dimos.msgs.image import image_view +from dimos.msgs.time import to_seconds from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -140,10 +142,10 @@ def _processing_loop(self) -> None: self._latest_frame = None # Get numpy array from Image - img_array = frame.data + img_array = image_view(frame) # Create header - header = Header(self.camera_frame_id) + header = Header(stamp=frame.header.stamp, frame_id=self.camera_frame_id) # Publish camera info self._publish_camera_info(header, img_array.shape) @@ -177,15 +179,14 @@ def _publish_camera_info(self, header: Header, shape: tuple[int, ...]) -> None: P = [fx, 0, cx, 0, 0, fy, cy, 0, 0, 0, 1, 0] msg = CameraInfo( - D_length=len(D), header=header, height=height, width=width, distortion_model="plumb_bob", - D=D, - K=K, - R=R, - P=P, + d=D, + k=K, + r=R, + p=P, binning_x=0, binning_y=0, ) @@ -201,17 +202,13 @@ def _publish_camera_pose(self, header: Header) -> None: transform = self.tfbuffer.get( parent_frame=self.world_frame_id, child_frame=self.camera_frame_id, - time_point=header.ts, + time_point=to_seconds(header.stamp), time_tolerance=1.0, ) if transform: - pose_msg = PoseStamped( - ts=header.ts, - frame_id=self.camera_frame_id, - position=transform.translation, - orientation=transform.rotation, - ) + pose_msg = pose_from_transform(transform) + pose_msg.header.stamp = header.stamp self.camera_pose.publish(pose_msg) except Exception as e: diff --git a/dimos/robot/drone/connection_module.py b/dimos/robot/drone/connection_module.py index e823ca8defa..e2074d287e4 100644 --- a/dimos/robot/drone/connection_module.py +++ b/dimos/robot/drone/connection_module.py @@ -21,7 +21,17 @@ import time from typing import Any -from dimos_lcm.std_msgs import String +from dimos_generated.geometry_msgs.msg import ( + Pose, + PoseStamped, + Transform, + TransformStamped, + Twist, + Vector3, +) +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header, String +from dimos_generated.tf2_msgs.msg import TFMessage from reactivex.disposable import Disposable from dimos.agents.annotation import skill @@ -30,13 +40,7 @@ from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out from dimos.mapping.models import LatLon -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.time import time_from_seconds from dimos.robot.drone.dji_video_stream import DJIDroneVideoStream from dimos.robot.drone.mavlink_connection import MavlinkConnection from dimos.utils.logging_config import setup_logger @@ -141,7 +145,7 @@ def start(self) -> None: ) # Subscribe to movement commands - self.register_disposable(Disposable(self.movecmd.subscribe(self.move))) + self.register_disposable(Disposable(self.movecmd.subscribe(self._on_move_command))) # Subscribe to Twist movement commands if self.movecmd_twist.transport: @@ -173,34 +177,32 @@ def _publish_tf(self, msg: PoseStamped) -> None: self.odom.publish(msg) # Publish base_link transform - base_link = Transform( - translation=msg.position, - rotation=msg.orientation, - frame_id="world", + position = msg.pose.position + base_link = TransformStamped( + header=msg.header, child_frame_id="base_link", - ts=msg.ts if hasattr(msg, "ts") else time.time(), + transform=Transform( + translation=Vector3(x=position.x, y=position.y, z=position.z), + rotation=msg.pose.orientation, + ), ) - - # Publish camera_link transform (camera mounted on front of drone, no gimbal factored in yet) - camera_link = Transform( - translation=Vector3(0.1, 0.0, -0.05), # 10cm forward, 5cm down - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), # No rotation relative to base - frame_id="base_link", + camera_link = TransformStamped( + header=Header(stamp=time_from_seconds(time.time()), frame_id="base_link"), child_frame_id="camera_link", - ts=time.time(), + transform=Transform(translation=Vector3(x=0.1, z=-0.05)), ) - self.tf.publish(TFMessage(base_link, camera_link)) + self.tf.publish(TFMessage(transforms=[base_link, camera_link])) def _publish_status(self, status: dict[str, Any]) -> None: """Publish drone status as JSON string.""" self._status = status - status_str = String(json.dumps(status)) + status_str = String(data=json.dumps(status)) self.status.publish(status_str) def _publish_telemetry(self, telemetry: dict[str, Any]) -> None: """Publish full telemetry as JSON string.""" - telemetry_str = String(json.dumps(telemetry)) + telemetry_str = String(data=json.dumps(telemetry)) self.telemetry.publish(telemetry_str) self._latest_telemetry = telemetry @@ -222,10 +224,8 @@ def _telemetry_loop(self) -> None: if self._odom is None: # Publish default pose default_pose = PoseStamped( - position=Vector3(0, 0, 0), - orientation=Quaternion(0, 0, 0, 1), - frame_id="world", - ts=time.time(), + header=Header(stamp=time_from_seconds(time.time()), frame_id="world"), + pose=Pose(), ) self._publish_tf(default_pose) logger.debug("Publishing default odometry") @@ -265,7 +265,7 @@ def move(self, x: float = 0.0, y: float = 0.0, z: float = 0.0, duration: float = duration: How long to move (0 = continuous) """ if self.connection: - self.connection.move(Vector3(x, y, z), duration) + self.connection.move(Vector3(x=x, y=y, z=z), duration) @skill def takeoff(self, altitude: float = 3.0) -> bool: @@ -385,7 +385,7 @@ def follow_object( duration (float, optional): How long to track for. Defaults to 120.0. """ msg = {"object_description": object_description, "duration": duration} - self.follow_object_cmd.publish(String(json.dumps(msg))) + self.follow_object_cmd.publish(String(data=json.dumps(msg))) yield "Started trying to track. First, trying to find the object." @@ -416,6 +416,10 @@ def follow_object( else: yield f"Stopped tracking '{object_description}'" + def _on_move_command(self, message: Vector3) -> None: + """Unpack generated movement vectors at the scalar skill boundary.""" + self.move(message.x, message.y, message.z) + def _on_move_twist(self, msg: Twist) -> None: """Handle Twist movement commands from tracking/navigation. diff --git a/dimos/robot/drone/dji_video_stream.py b/dimos/robot/drone/dji_video_stream.py index 92afc1ce507..54229da7480 100644 --- a/dimos/robot/drone/dji_video_stream.py +++ b/dimos/robot/drone/dji_video_stream.py @@ -23,10 +23,13 @@ import time from typing import Any +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header import numpy as np from reactivex import Observable, Subject -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_from_array +from dimos.msgs.time import time_from_seconds from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -137,7 +140,9 @@ def _capture_loop(self) -> None: frame = frame.reshape((self.height, self.width, channels)) # Create Image message (RGB format - matches GStreamer pipeline output) - img_msg = Image.from_numpy(frame, format=ImageFormat.RGB) + img_msg = image_from_array( + frame, encoding="rgb8", header=Header(stamp=time_from_seconds(time.time())) + ) # Publish self._video_subject.on_next(img_msg) @@ -217,8 +222,8 @@ def get_stream(self) -> Observable[Image]: # type: ignore[override] from dimos.utils.testing.legacy_pickle import LegacyPickleStore def _fix_format(img: Image) -> Image: - if img.format == ImageFormat.BGR: - img.format = ImageFormat.RGB + if img.encoding == "bgr8": + img.encoding = "rgb8" return img logger.info("Creating video replay stream") diff --git a/dimos/robot/drone/drone_tracking_module.py b/dimos/robot/drone/drone_tracking_module.py index ec18323ad3f..e22be472c9d 100644 --- a/dimos/robot/drone/drone_tracking_module.py +++ b/dimos/robot/drone/drone_tracking_module.py @@ -20,7 +20,9 @@ import time from typing import Any -from dimos_lcm.std_msgs import String +from dimos_generated.geometry_msgs.msg import Twist, Vector3 +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header, String import numpy as np from numpy.typing import NDArray @@ -29,9 +31,8 @@ from dimos.core.module import Module from dimos.core.stream import In, Out from dimos.models.qwen.video_query import get_bbox_from_qwen_frame -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_from_array, image_view +from dimos.msgs.time import time_from_seconds from dimos.robot.drone.drone_visual_servoing_controller import ( DroneVisualServoingController, PIDParams, @@ -114,7 +115,7 @@ def _get_latest_frame(self) -> np.ndarray[Any, np.dtype[Any]] | None: if self._latest_frame is None: return None # Convert Image to numpy array - data: np.ndarray[Any, np.dtype[Any]] = self._latest_frame.data + data: np.ndarray[Any, np.dtype[Any]] = image_view(self._latest_frame) return data @rpc @@ -270,8 +271,8 @@ def _visual_servoing_loop(self, tracker: Any, duration: float) -> None: # Publish velocity command via LCM if self.cmd_vel.transport: twist = Twist() - twist.linear = Vector3(vx, vy, 0) - twist.angular = Vector3(0, 0, 0) # No rotation for now + twist.linear = Vector3(x=vx, y=vy) + twist.angular = Vector3() # No rotation for now self.cmd_vel.publish(twist) # Publish visualization if transport is set @@ -279,7 +280,11 @@ def _visual_servoing_loop(self, tracker: Any, duration: float) -> None: overlay = self._draw_tracking_overlay( frame, (int(x), int(y), int(w), int(h)), (int(current_x), int(current_y)) ) - overlay_msg = Image.from_numpy(overlay, format=ImageFormat.BGR) + overlay_msg = image_from_array( + overlay, + encoding="bgr8", + header=Header(stamp=time_from_seconds(time.time())), + ) self.tracking_overlay.publish(overlay_msg) # Publish status @@ -304,8 +309,8 @@ def _visual_servoing_loop(self, tracker: Any, duration: float) -> None: # Stop movement by publishing zero velocity if self.cmd_vel.transport: stop_twist = Twist() - stop_twist.linear = Vector3(0, 0, 0) - stop_twist.angular = Vector3(0, 0, 0) + stop_twist.linear = Vector3() + stop_twist.angular = Vector3() self.cmd_vel.publish(stop_twist) self._tracking_active = False logger.info(f"Visual servoing loop ended after {frame_count} frames") @@ -366,7 +371,7 @@ def _publish_status(self, status: dict[str, Any]) -> None: status: Status dictionary """ if self.tracking_status.transport: - status_msg = String(json.dumps(status)) + status_msg = String(data=json.dumps(status)) self.tracking_status.publish(status_msg) def _stop_tracking(self) -> None: @@ -378,8 +383,8 @@ def _stop_tracking(self) -> None: # Send stop command via LCM if self.cmd_vel.transport: stop_twist = Twist() - stop_twist.linear = Vector3(0, 0, 0) - stop_twist.angular = Vector3(0, 0, 0) + stop_twist.linear = Vector3() + stop_twist.angular = Vector3() self.cmd_vel.publish(stop_twist) self._publish_status({"status": "stopped", "object": self._current_object}) diff --git a/dimos/robot/drone/mavlink_connection.py b/dimos/robot/drone/mavlink_connection.py index e908be4ab03..1bb15870c63 100644 --- a/dimos/robot/drone/mavlink_connection.py +++ b/dimos/robot/drone/mavlink_connection.py @@ -20,13 +20,13 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Twist, Vector3 +from dimos_generated.std_msgs.msg import Header from pymavlink import mavutil # type: ignore[import-untyped] from reactivex import Subject -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.time import time_from_seconds from dimos.utils.logging_config import setup_logger logger = setup_logger(level=logging.INFO) @@ -167,7 +167,7 @@ def _publish_odom(self) -> None: # MAVLink --> ROS conversion # MAVLink: positive pitch = nose up, positive yaw = clockwise # ROS: positive pitch = nose down, positive yaw = counter-clockwise - quaternion = Quaternion.from_euler(Vector3(roll, -pitch, -yaw)) + quaternion = quaternion_from_euler(roll, -pitch, -yaw) if not hasattr(self, "_position"): self._position = {"x": 0.0, "y": 0.0, "z": 0.0} @@ -229,10 +229,8 @@ def _publish_odom(self) -> None: ) pose = PoseStamped( - position=Vector3(self._position["x"], self._position["y"], self._position["z"]), - orientation=quaternion, - frame_id="world", - ts=current_time, + header=Header(stamp=time_from_seconds(current_time), frame_id="world"), + pose=Pose(position=Point(**self._position), orientation=quaternion), ) self._odom_subject.on_next(pose) diff --git a/dimos/robot/drone/test_drone.py b/dimos/robot/drone/test_drone.py index 2e7336c2c7d..ce30006ea4f 100644 --- a/dimos/robot/drone/test_drone.py +++ b/dimos/robot/drone/test_drone.py @@ -23,12 +23,13 @@ import unittest from unittest.mock import MagicMock, patch +from dimos_generated.geometry_msgs.msg import PoseStamped, Vector3 +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_from_array, image_view from dimos.robot.drone.connection_module import DroneConnectionModule from dimos.robot.drone.dji_video_stream import FakeDJIVideoStream @@ -85,7 +86,7 @@ def recv_side_effect(*args, **kwargs): # Verify NED to ROS conversion happened # ROS uses different conventions: positive pitch = nose down, positive yaw = counter-clockwise # So we expect sign flips in the quaternion conversion - self.assertIsNotNone(pose.orientation) + self.assertIsNotNone(pose.pose.orientation) def test_position_integration(self) -> None: """Test velocity integration for indoor flight positioning.""" @@ -221,7 +222,10 @@ def test_fake_mavlink_connection(self) -> None: def test_fake_video_stream_no_throttling(self) -> None: """Test FakeDJIVideoStream returns replay stream with format fix.""" - with patch("dimos.utils.testing.legacy_pickle.LegacyPickleStore") as mock_replay: + with ( + patch("dimos.utils.data.get_data") as mock_get_data, + patch("dimos.utils.testing.legacy_pickle.LegacyPickleStore") as mock_replay, + ): mock_stream = MagicMock() mock_replay.return_value.stream.return_value = mock_stream @@ -229,6 +233,7 @@ def test_fake_video_stream_no_throttling(self) -> None: stream.get_stream() # Verify replay store was created and stream was piped (for BGR→RGB fix) + mock_get_data.assert_called_once_with("drone") mock_replay.assert_called_once_with("drone/video") mock_replay.return_value.stream.assert_called_once() mock_stream.pipe.assert_called_once() @@ -357,16 +362,14 @@ def replay_side_effect(store_name: str): publish=lambda x: ( published_odom.append(x), print( - f"[TEST] Published odom: position=({x.position.x:.2f}, {x.position.y:.2f}, {x.position.z:.2f})" + f"[TEST] Published odom: position=({x.pose.position.x:.2f}, {x.pose.position.y:.2f}, {x.pose.position.z:.2f})" ), ) ) module.video = MagicMock( publish=lambda x: ( published_video.append(x), - print( - f"[TEST] Published video frame with shape: {x.data.shape if hasattr(x, 'data') else 'unknown'}" - ), + print(f"[TEST] Published video frame with shape: {image_view(x).shape}"), ) ) module.status = MagicMock( @@ -505,7 +508,7 @@ def test_video_stream_replay(self, mock_get_data, mock_replay) -> None: def subscribe_side_effect(callback) -> None: for frame in video_frames: - img = Image(data=frame, format=ImageFormat.BGR) + img = image_from_array(frame, encoding="bgr8", header=Header()) callback(img) received_frames.append(img) @@ -517,7 +520,7 @@ def subscribe_side_effect(callback) -> None: def piped_subscribe(callback): for frame in video_frames: - img = Image(data=frame, format=ImageFormat.RGB) # After format fix + img = image_from_array(frame, encoding="rgb8", header=Header()) # After format fix callback(img) piped_captured.append(img) @@ -537,8 +540,8 @@ def piped_subscribe(callback): self.assertEqual(len(piped_captured), 2) for _i, frame in enumerate(piped_captured): self.assertIsInstance(frame, Image) - self.assertEqual(frame.data.shape, (360, 640, 3)) - self.assertEqual(frame.format, ImageFormat.RGB) # Format should be corrected + self.assertEqual(image_view(frame).shape, (360, 640, 3)) + self.assertEqual(frame.encoding, "rgb8") # Format should be corrected class TestDroneMovementAndOdometry(unittest.TestCase): @@ -556,7 +559,7 @@ def test_movement_command_conversion(self, mock_get_data, mock_replay) -> None: # Test movement in ROS frame # ROS: X=forward, Y=left, Z=up - velocity_ros = Vector3(2.0, -1.0, 0.5) # Forward 2m/s, right 1m/s, up 0.5m/s + velocity_ros = Vector3(x=2.0, y=-1.0, z=0.5) # Forward 2m/s, right 1m/s, up 0.5m/s result = conn.move(velocity_ros, duration=1.0) self.assertTrue(result) @@ -608,8 +611,8 @@ def replay_stream_subscribe(callback) -> None: # Check odometry message odom = published_odom[0] self.assertIsInstance(odom, PoseStamped) - self.assertIsNotNone(odom.orientation) - self.assertEqual(odom.frame_id, "world") + self.assertIsNotNone(odom.pose.orientation) + self.assertEqual(odom.header.frame_id, "world") @patch("dimos.utils.testing.legacy_pickle.LegacyPickleStore") @patch("dimos.utils.data.get_data") diff --git a/dimos/robot/galaxea/r1pro/test_connection_messages.py b/dimos/robot/galaxea/r1pro/test_connection_messages.py new file mode 100644 index 00000000000..fd1670c06d8 --- /dev/null +++ b/dimos/robot/galaxea/r1pro/test_connection_messages.py @@ -0,0 +1,157 @@ +# 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. + +import math + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.dimos_msgs.msg import MotorCommandArray +from dimos_generated.geometry_msgs.msg import Twist, TwistStamped, Vector3 +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header +import pytest + +from dimos.hardware.whole_body.spec import VEL_STOP +from dimos.protocol.pubsub.impl.rospubsub import RawROS, RawROSTopic +from dimos.robot.galaxea.r1pro import connection + + +@pytest.fixture() +def robot(): + robot = connection.R1ProConnection() + try: + yield robot + finally: + robot.stop() + + +def speed(sec, nanosec, *, vx=1.0, vy=0.0, wz=0.0): + return TwistStamped( + header=Header(stamp=Time(sec=sec, nanosec=nanosec)), + twist=Twist(linear=Vector3(x=vx, y=vy), angular=Vector3(z=wz)), + ) + + +def test_integrated_pose_odometry_and_tf_preserve_nanoseconds(robot, monkeypatch): + outputs = {name: [] for name in ("odom", "odometry", "tf")} + for name, values in outputs.items(): + monkeypatch.setattr(getattr(robot, name), "publish", values.append) + robot._on_chassis_speed(speed(1700000000, 123456789), None) + robot._on_chassis_speed(speed(1700000000, 223456789, vy=0.5, wz=0.2), None) + + pose, odometry, tf = [outputs[name][0] for name in outputs] + assert pose.pose.position.x == pytest.approx(0.1) + assert pose.pose.position.y == pytest.approx(0.05) + assert pose.pose.orientation.z == pytest.approx(math.sin(0.01)) + assert odometry.pose.pose == pose.pose + assert odometry.twist.twist.linear == Vector3(x=1.0, y=0.5) + assert [edge.child_frame_id for edge in tf.transforms] == [ + robot.config.frame_id, + robot.config.lidar_frame_id, + ] + assert tf.transforms[1].header.frame_id == robot.config.frame_id + assert tf.transforms[1].transform.rotation.w == 1.0 + for message in (pose, odometry, *tf.transforms): + assert message.header.stamp == Time(sec=1700000000, nanosec=223456789) + for message in (pose, odometry, tf): + assert type(message).decode(message.encode()) == message + + +@pytest.mark.parametrize("second_stamp", [(0, 0), (-1, 0), (2, 1)]) +def test_clock_jump_does_not_move_base(robot, monkeypatch, second_stamp): + values = [] + monkeypatch.setattr(robot.odom, "publish", values.append) + robot._on_chassis_speed(speed(0, 0), None) + robot._on_chassis_speed(speed(*second_stamp), None) + assert values == [] + assert robot._odom_x == 0.0 + + +@pytest.mark.parametrize("oldest", [Time(), Time(sec=1700000000, nanosec=123456788)]) +def test_joint_snapshot_uses_oldest_exact_stamp_and_all_segments(robot, monkeypatch, oldest): + for callback, count, stamp, value in [ + (robot._on_feedback_torso, 4, oldest, 1.0), + (robot._on_feedback_left, 7, Time(sec=1700000000, nanosec=123456789), 2.0), + (robot._on_feedback_right, 7, Time(sec=1700000000, nanosec=123456790), 3.0), + ]: + callback( + JointState( + header=Header(stamp=stamp), + position=[value] * count, + velocity=[0.5] * count, + effort=[0.2] * count, + ), + None, + ) + values = [] + + def collect(message): + values.append(message) + robot._stop_event.set() + + monkeypatch.setattr(robot.motor_states, "publish", collect) + robot._publish_loop() + state = values[0] + assert list(state.position) == [1.0] * 4 + [2.0] * 7 + [3.0] * 7 + assert list(state.velocity) == [0.5] * 18 + assert list(state.name) == connection.R1PRO_UPPER_BODY_JOINTS + assert state.header.stamp == oldest + assert JointState.decode(state.encode()) == state + + +@pytest.fixture() +def command_sink(robot, monkeypatch, mocker): + ros = mocker.Mock(spec=RawROS) + monkeypatch.setattr(robot, "_ros", ros) + monkeypatch.setattr(connection, "dimos_to_ros", lambda msg, _: msg) + monkeypatch.setattr(connection, "header_now", lambda: Header(stamp=Time(sec=9, nanosec=123))) + for name in ("_cmd_torso_topic", "_cmd_left_topic", "_cmd_right_topic", "_speed_topic"): + monkeypatch.setattr(robot, name, RawROSTopic(name, JointState)) + return ros.publish + + +def motor_command(): + return MotorCommandArray( + q=list(range(18)), dq=[0.0, VEL_STOP, 0.2] * 6, kp=[0.0] * 18, kd=[0.0] * 18, tau=[0.0] * 18 + ) + + +def test_motor_command_splits_generated_sequences_and_tracking_speed(robot, command_sink): + robot._on_motor_command(motor_command()) + messages = [call.args[1] for call in command_sink.call_args_list] + assert [list(msg.position) for msg in messages] == [ + list(range(4)), + list(range(4, 11)), + list(range(11, 18)), + ] + assert list(messages[0].velocity) == [0.5, 0.5, 0.2, 0.5] + assert all(msg.header.stamp == Time(sec=9, nanosec=123) for msg in messages) + + +@pytest.mark.parametrize("field", ["q", "dq", "kp", "kd", "tau"]) +def test_malformed_motor_array_never_reaches_ros(robot, command_sink, field): + message = motor_command() + setattr(message, field, [0.0]) + robot._on_motor_command(message) + command_sink.assert_not_called() + + +def test_base_command_only_forwards_supported_axes(robot, command_sink): + robot._on_cmd_vel( + Twist(linear=Vector3(x=1.0, y=2.0, z=3.0), angular=Vector3(x=4.0, y=5.0, z=6.0)) + ) + message = command_sink.call_args.args[1] + assert message == TwistStamped( + header=Header(stamp=Time(sec=9, nanosec=123)), + twist=Twist(linear=Vector3(x=1.0, y=2.0), angular=Vector3(z=6.0)), + ) diff --git a/dimos/robot/manipulators/common/mock.py b/dimos/robot/manipulators/common/mock.py index 2420591c42a..9da13b93c64 100644 --- a/dimos/robot/manipulators/common/mock.py +++ b/dimos/robot/manipulators/common/mock.py @@ -16,11 +16,12 @@ from __future__ import annotations +from dimos_generated.sensor_msgs.msg import JointState + from dimos.control.components import HardwareComponent, HardwareType, make_joints from dimos.control.coordinator import ControlCoordinator from dimos.control.tasks.trajectory_task.trajectory_task import joint_trajectory_task from dimos.core.stream import Out -from dimos.msgs.sensor_msgs.JointState import JointState _mock_hw = HardwareComponent( hardware_id="arm", diff --git a/dimos/robot/manipulators/dual_openyam/test_blueprints.py b/dimos/robot/manipulators/dual_openyam/test_blueprints.py index 0633b071043..ab7129ed98a 100644 --- a/dimos/robot/manipulators/dual_openyam/test_blueprints.py +++ b/dimos/robot/manipulators/dual_openyam/test_blueprints.py @@ -14,6 +14,9 @@ from typing import Any, cast +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Float32, Header import pytest from pytest_mock import MockerFixture @@ -22,9 +25,6 @@ from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser from dimos.core.coordination.blueprints import Blueprint from dimos.hardware.whole_body.spec import WholeBodyAdapter -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.std_msgs.Float32 import Float32 from dimos.robot.manipulators.dual_openyam.blueprints.basic import ( DualOpenYamCoordinator, ) @@ -72,8 +72,8 @@ def test_mock_webxr_coordinator_commands_both_arms_and_grippers( DualOpenYamPinkPoseTargetSolver, "frame_poses", return_value={ - "left_grasp_frame": PoseStamped(position=[0.5, 0.2, 0.4]), - "right_grasp_frame": PoseStamped(position=[0.5, -0.2, 0.4]), + "left_grasp_frame": PoseStamped(pose=Pose(position=Point(x=0.5, y=0.2, z=0.4))), + "right_grasp_frame": PoseStamped(pose=Pose(position=Point(x=0.5, y=-0.2, z=0.4))), }, ) mocker.patch.object( @@ -91,16 +91,22 @@ def test_mock_webxr_coordinator_commands_both_arms_and_grippers( buttons = Buttons() buttons.left_grip = True buttons.right_grip = True - coordinator._dispatch("teleop_buttons", buttons) + coordinator._dispatch("teleop_buttons", buttons.to_message()) coordinator._dispatch("left_gripper_command", Float32(data=0.75)) coordinator._dispatch("right_gripper_command", Float32(data=0.25)) coordinator._dispatch( "left_cartesian_command", - PoseStamped(frame_id=DUAL_OPENYAM_WEBXR_TASK_NAME, position=[1.0, 0.0, 0.0]), + PoseStamped( + header=Header(frame_id=DUAL_OPENYAM_WEBXR_TASK_NAME), + pose=Pose(position=Point(x=1.0)), + ), ) coordinator._dispatch( "right_cartesian_command", - PoseStamped(frame_id=DUAL_OPENYAM_WEBXR_TASK_NAME, position=[-1.0, 0.0, 0.0]), + PoseStamped( + header=Header(frame_id=DUAL_OPENYAM_WEBXR_TASK_NAME), + pose=Pose(position=Point(x=-1.0)), + ), ) assert coordinator._tick_loop is not None coordinator._tick_loop._tick() diff --git a/dimos/robot/manipulators/dual_openyam/test_teleop_ik.py b/dimos/robot/manipulators/dual_openyam/test_teleop_ik.py index 103f620a797..104e24945fa 100644 --- a/dimos/robot/manipulators/dual_openyam/test_teleop_ik.py +++ b/dimos/robot/manipulators/dual_openyam/test_teleop_ik.py @@ -14,12 +14,12 @@ """Objective tests for Dual OpenYAM WebXR teleoperation.""" +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.sensor_msgs.msg import JointState import numpy as np import pytest from dimos.control.tasks.pose_target_ik import PoseTargetIKTaskConfig -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.robot.manipulators.dual_openyam.blueprints.teleop import ( _dual_openyam_webxr_task, ) @@ -76,9 +76,13 @@ def test_webxr_solver_matches_a1z_target_tracking_speed() -> None: initial = solver.frame_poses(state, _TARGET_FRAMES) targets = { frame_name: PoseStamped( - frame_id=pose.frame_id, - position=[pose.position.x, pose.position.y, pose.position.z + 0.1], - orientation=pose.orientation, + header=pose.header, + pose=Pose( + position=Point( + x=pose.pose.position.x, y=pose.pose.position.y, z=pose.pose.position.z + 0.1 + ), + orientation=pose.pose.orientation, + ), ) for frame_name, pose in initial.items() } @@ -91,7 +95,7 @@ def test_webxr_solver_matches_a1z_target_tracking_speed() -> None: current = solver.frame_poses(state, _TARGET_FRAMES) progress = np.mean( [ - current[frame_name].position.z - initial[frame_name].position.z + current[frame_name].pose.position.z - initial[frame_name].pose.position.z for frame_name in _TARGET_FRAMES ] ) diff --git a/dimos/robot/manipulators/openarm/test_openarm_teleop.py b/dimos/robot/manipulators/openarm/test_openarm_teleop.py index 5fd7bb203ca..c242f34773f 100644 --- a/dimos/robot/manipulators/openarm/test_openarm_teleop.py +++ b/dimos/robot/manipulators/openarm/test_openarm_teleop.py @@ -16,6 +16,9 @@ from typing import Any, cast +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Float32, Header import numpy as np import pytest from pytest_mock import MockerFixture @@ -28,9 +31,6 @@ from dimos.hardware.whole_body.spec import WholeBodyAdapter from dimos.manipulation.planning.kinematics.config import PinkKinematicsConfig from dimos.manipulation.planning.spec.config import RobotModelConfig -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.std_msgs.Float32 import Float32 from dimos.robot.manipulators.openarm.blueprints.basic import openarm_planner_coordinator from dimos.robot.manipulators.openarm.blueprints.teleop import ( OPENARM_WEBXR_TASK_NAME, @@ -160,8 +160,12 @@ def test_openarm_webxr_commands_both_arms_and_grippers_through_coordinator( OpenArmPinkPoseTargetSolver, "frame_poses", return_value={ - "openarm_left_grasp_frame": PoseStamped(position=[0.5, 0.2, 0.4]), - "openarm_right_grasp_frame": PoseStamped(position=[0.5, -0.2, 0.4]), + "openarm_left_grasp_frame": PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(x=0.5, y=0.2, z=0.4)) + ), + "openarm_right_grasp_frame": PoseStamped( + header=Header(frame_id=""), pose=Pose(position=Point(x=0.5, y=-0.2, z=0.4)) + ), }, ) step = mocker.patch.object( @@ -194,16 +198,22 @@ def test_openarm_webxr_commands_both_arms_and_grippers_through_coordinator( buttons.left_grip = True buttons.right_grip = True buttons.pack_analog_triggers(left=0.25, right=0.75) - coordinator._dispatch("teleop_buttons", buttons) + coordinator._dispatch("teleop_buttons", buttons.to_message()) coordinator._dispatch("left_gripper_command", Float32(data=0.75)) coordinator._dispatch("right_gripper_command", Float32(data=0.25)) coordinator._dispatch( "left_cartesian_command", - PoseStamped(frame_id=OPENARM_WEBXR_TASK_NAME, position=[1.0, 0.0, 0.0]), + PoseStamped( + header=Header(frame_id=OPENARM_WEBXR_TASK_NAME), + pose=Pose(position=Point(x=1.0, y=0.0, z=0.0)), + ), ) coordinator._dispatch( "right_cartesian_command", - PoseStamped(frame_id=OPENARM_WEBXR_TASK_NAME, position=[-1.0, 0.0, 0.0]), + PoseStamped( + header=Header(frame_id=OPENARM_WEBXR_TASK_NAME), + pose=Pose(position=Point(x=-1.0, y=0.0, z=0.0)), + ), ) assert coordinator._tick_loop is not None @@ -223,7 +233,7 @@ def test_openarm_webxr_commands_both_arms_and_grippers_through_coordinator( released = Buttons() released.right_primary = True - coordinator._dispatch("teleop_buttons", released) + coordinator._dispatch("teleop_buttons", released.to_message()) coordinator._tick_loop._tick() step.assert_called_once() finally: @@ -291,9 +301,13 @@ def test_openarm_bimanual_pink_steps_from_canonical_zero_with_bounded_updates() initial = ik.frame_poses(seed, frames) targets = { name: PoseStamped( - frame_id=pose.frame_id, - position=[pose.position.x, pose.position.y, pose.position.z + 0.01], - orientation=pose.orientation, + header=Header(frame_id=pose.header.frame_id), + pose=Pose( + position=Point( + x=pose.pose.position.x, y=pose.pose.position.y, z=pose.pose.position.z + 0.01 + ), + orientation=pose.pose.orientation, + ), ) for name, pose in initial.items() } @@ -311,12 +325,12 @@ def test_openarm_bimanual_pink_steps_from_canonical_zero_with_bounded_updates() final = ik.frame_poses(seed, frames) errors = [ np.linalg.norm( - np.array([target.position.x, target.position.y, target.position.z]) + np.array([target.pose.position.x, target.pose.position.y, target.pose.position.z]) - np.array( [ - final[name].position.x, - final[name].position.y, - final[name].position.z, + final[name].pose.position.x, + final[name].pose.position.y, + final[name].pose.position.z, ] ) ) diff --git a/dimos/robot/manipulators/openyam/test_teleop_ik.py b/dimos/robot/manipulators/openyam/test_teleop_ik.py index 945ec7d8935..fcc0cf193d3 100644 --- a/dimos/robot/manipulators/openyam/test_teleop_ik.py +++ b/dimos/robot/manipulators/openyam/test_teleop_ik.py @@ -14,13 +14,12 @@ """Objective tests for single-arm OpenYAM Quest teleoperation.""" +from dimos_generated.geometry_msgs.msg import Pose, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import JointState import pytest +from scipy.spatial.transform import Rotation from dimos.control.tasks.pose_target_ik import PinkPoseTargetSolver, PoseTargetIKTaskConfig -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.robot.manipulators.openyam.blueprints.teleop import teleop_webxr_openyam from dimos.robot.manipulators.openyam.config import OPENYAM_ARM_JOINTS, OPENYAM_HOME_JOINTS from dimos.robot.manipulators.openyam.teleop_ik import OpenYamPinkPoseTargetSolver @@ -70,10 +69,16 @@ def _orientation_target_motion( solver = _solver(solver_type) initial = JointState(name=OPENYAM_ARM_JOINTS, position=OPENYAM_HOME_JOINTS) pose = solver.frame_poses(initial, (_TARGET_FRAME,))[_TARGET_FRAME] + q = pose.pose.orientation + rotated = ( + Rotation.from_euler("xyz", [0.0, 0.2, 0.0]) * Rotation.from_quat([q.x, q.y, q.z, q.w]) + ).as_quat() target = PoseStamped( - frame_id=pose.frame_id, - position=pose.position, - orientation=Quaternion.from_euler(Vector3(0.0, 0.2, 0.0)) * pose.orientation, + header=pose.header, + pose=Pose( + position=pose.pose.position, + orientation=Quaternion(x=rotated[0], y=rotated[1], z=rotated[2], w=rotated[3]), + ), ) state = initial for _ in range(10): diff --git a/dimos/robot/unitree/b1/joystick_module.py b/dimos/robot/unitree/b1/joystick_module.py index 4040e818ecd..7c9e578895d 100644 --- a/dimos/robot/unitree/b1/joystick_module.py +++ b/dimos/robot/unitree/b1/joystick_module.py @@ -25,14 +25,14 @@ import time +from dimos_generated.geometry_msgs.msg import Twist, TwistStamped, Vector3 +from dimos_generated.std_msgs.msg import Header, Int32 + from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import Out -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.std_msgs.Int32 import Int32 +from dimos.msgs.time import time_from_seconds class JoystickModule(Module): @@ -80,10 +80,8 @@ def stop(self) -> None: # Send stop command stop_twist = Twist() stop_twist_stamped = TwistStamped( - ts=time.time(), - frame_id="base_link", - linear=stop_twist.linear, - angular=stop_twist.angular, + header=Header(stamp=time_from_seconds(time.time()), frame_id="base_link"), + twist=Twist(linear=stop_twist.linear, angular=stop_twist.angular), ) self.twist_out.publish(stop_twist_stamped) @@ -145,13 +143,13 @@ def _pygame_loop(self) -> None: self.mode_out.publish(mode_msg) # Also send zero twist stop_twist = Twist() - stop_twist.linear = Vector3(0, 0, 0) - stop_twist.angular = Vector3(0, 0, 0) + stop_twist.linear = Vector3(x=0, y=0, z=0) + stop_twist.angular = Vector3(x=0, y=0, z=0) stop_twist_stamped = TwistStamped( - ts=time.time(), - frame_id="base_link", - linear=stop_twist.linear, - angular=stop_twist.angular, + header=Header( + stamp=time_from_seconds(time.time()), frame_id="base_link" + ), + twist=Twist(linear=stop_twist.linear, angular=stop_twist.angular), ) self.twist_out.publish(stop_twist_stamped) print("EMERGENCY STOP!") @@ -164,8 +162,8 @@ def _pygame_loop(self) -> None: # Generate Twist message from held keys twist = Twist() - twist.linear = Vector3(0, 0, 0) - twist.angular = Vector3(0, 0, 0) + twist.linear = Vector3(x=0, y=0, z=0) + twist.angular = Vector3(x=0, y=0, z=0) # Apply controls based on mode if self.current_mode == 2: # WALK mode - movement control @@ -213,7 +211,8 @@ def _pygame_loop(self) -> None: twist.angular.y = -1.0 # Pitch backward twist_stamped = TwistStamped( - ts=time.time(), frame_id="base_link", linear=twist.linear, angular=twist.angular + header=Header(stamp=time_from_seconds(time.time()), frame_id="base_link"), + twist=Twist(linear=twist.linear, angular=twist.angular), ) self.twist_out.publish(twist_stamped) diff --git a/dimos/robot/unitree/b1/test_cdr_forwarding.py b/dimos/robot/unitree/b1/test_cdr_forwarding.py new file mode 100644 index 00000000000..0456fa1e8b0 --- /dev/null +++ b/dimos/robot/unitree/b1/test_cdr_forwarding.py @@ -0,0 +1,87 @@ +# Copyright 2025-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. + +from types import SimpleNamespace +from unittest.mock import Mock + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + PoseWithCovariance, + Quaternion, + Twist, + TwistStamped, + Vector3, +) +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.std_msgs.msg import Header, Int32 +import pytest + +from dimos.robot.unitree.b1.connection import B1ConnectionModule, MockB1ConnectionModule, RobotMode +from dimos.robot.unitree.b1.joystick_module import JoystickModule + + +def test_b1_odometry_forwarding_preserves_generated_pose_and_exact_header(): + source = Odometry( + header=Header(frame_id="map", stamp=Time(sec=1700000000, nanosec=123456789)), + pose=PoseWithCovariance( + pose=Pose(position=Point(x=1, y=2, z=3), orientation=Quaternion(w=1)) + ), + ) + received = [] + proxy = SimpleNamespace(odom_pose=SimpleNamespace(publish=received.append)) + B1ConnectionModule._publish_odom_pose(proxy, source) + result = PoseStamped.decode(received[0].encode()) + assert result.header == source.header + assert result.pose == source.pose.pose + + +@pytest.mark.parametrize("mode", [RobotMode.WALK, RobotMode.STAND]) +def test_b1_generated_command_mapping_preserves_mode_clamps_and_source(mode): + module = MockB1ConnectionModule() + try: + module.handle_mode(Int32.decode(Int32(data=mode).encode())) + message = TwistStamped( + header=Header(frame_id="base_link", stamp=Time(sec=1700000000, nanosec=123456789)), + twist=Twist(linear=Vector3(x=2, y=-2, z=2), angular=Vector3(x=2, y=-2, z=4)), + ) + before = message.encode() + module.handle_twist_stamped(TwistStamped.decode(before)) + command = module._current_cmd + assert module.current_mode == mode and command.mode == mode + assert command.lx == -1 and command.ly == 1 + assert command.rx == (1 if mode == RobotMode.WALK else -1) + assert command.ry == (0 if mode == RobotMode.WALK else -1) + assert message.encode() == before + assert module.socket is None + finally: + module.stop() + + +def test_b1_joystick_stop_publishes_generated_zero_without_starting_loop(monkeypatch): + module = JoystickModule() + thread = Mock() + module._thread = thread + received = [] + monkeypatch.setattr(module.twist_out, "publish", received.append) + monkeypatch.setattr("dimos.robot.unitree.b1.joystick_module.time.time", lambda: 1700000000.25) + module.stop() + assert len(received) == 1 + value = TwistStamped.decode(received[0].encode()) + assert value.header.frame_id == "base_link" + assert value.header.stamp == Time(sec=1700000000, nanosec=250000000) + assert value.twist == Twist() + thread.join.assert_called_once() diff --git a/dimos/robot/unitree/b1/unitree_b1.py b/dimos/robot/unitree/b1/unitree_b1.py index 46924f3edc0..d05c5ae3d6e 100644 --- a/dimos/robot/unitree/b1/unitree_b1.py +++ b/dimos/robot/unitree/b1/unitree_b1.py @@ -23,15 +23,15 @@ import logging import os +from dimos_generated.geometry_msgs.msg import PoseStamped, TwistStamped +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.std_msgs.msg import Int32 +from dimos_generated.tf2_msgs.msg import TFMessage + from dimos.core.coordination.module_coordinator import ModuleCoordinator from dimos.core.resource import Resource from dimos.core.transport import ROSTransport from dimos.core.transport_factory import make_transport -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.std_msgs.Int32 import Int32 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.robot.robot import Robot from dimos.robot.unitree.b1.connection import ( B1ConnectionModule, diff --git a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_coordinator.py b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_coordinator.py index 35408e812ba..5d8e5b556a0 100644 --- a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_coordinator.py +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_coordinator.py @@ -22,15 +22,15 @@ import os +from dimos_generated.dimos_msgs.msg import MotorCommandArray +from dimos_generated.sensor_msgs.msg import Imu, JointState + from dimos.control.components import HardwareComponent, HardwareType, make_humanoid_joints from dimos.control.coordinator import ControlCoordinator from dimos.control.tasks.trajectory_task.trajectory_task import joint_trajectory_task from dimos.core.coordination.blueprints import autoconnect from dimos.core.stream import Out from dimos.core.transport import LCMTransport -from dimos.msgs.sensor_msgs.Imu import Imu -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.sensor_msgs.MotorCommandArray import MotorCommandArray from dimos.robot.unitree.g1.wholebody_connection import G1WholeBodyConnection _g1_joints = make_humanoid_joints("g1") diff --git a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_record.py b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_record.py index 079b16fcc44..3ddbc64e1a5 100644 --- a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_record.py +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_record.py @@ -22,6 +22,8 @@ from typing import Any +from dimos_generated.sensor_msgs.msg import CameraInfo + from dimos.core.coordination.blueprints import autoconnect from dimos.core.coordination.module_coordinator import ModuleCoordinator from dimos.core.global_config import global_config @@ -29,11 +31,11 @@ from dimos.hardware.sensors.lidar.pointlio.module import PointLio from dimos.hardware.sensors.lidar.pointlio.pointlio_blueprints import mid360_for_pointlio from dimos.memory.module import default_recording_dir -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo from dimos.navigation.movement_manager.movement_manager import MovementManager from dimos.robot.unitree.g1.effectors.high_level.dds_sdk import G1HighLevelDdsSdk from dimos.robot.unitree.g1.g1_recorder import G1Recorder from dimos.robot.unitree.g1.g1_tf_publisher import G1TfPublisher +from dimos.visualization.rerun.message_helpers import camera_pinhole from dimos.visualization.vis_module import vis_module _RECORDING_DIR = default_recording_dir() @@ -71,18 +73,22 @@ def _static_robot_body(rr: Any) -> list[Any]: def _convert_camera_info(camera_info: CameraInfo) -> Any: - return camera_info.to_rerun( - image_topic="/world/color_image", - optical_frame="d435_color_optical_frame", - ) + return [ + ( + "/world/color_image", + camera_pinhole(camera_info, optical_frame="d435_color_optical_frame"), + ) + ] def _convert_depth_camera_info(camera_info: CameraInfo) -> Any: # Depth is aligned to color, so it shares the color optical frame. - return camera_info.to_rerun( - image_topic="/world/realsense_depth_image", - optical_frame="d435_color_optical_frame", - ) + return [ + ( + "/world/realsense_depth_image", + camera_pinhole(camera_info, optical_frame="d435_color_optical_frame"), + ) + ] _record_vis = vis_module( diff --git a/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_detection.py b/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_detection.py index ec72cc10b23..ec472fc7b26 100644 --- a/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_detection.py +++ b/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_detection.py @@ -17,13 +17,13 @@ from typing import Any +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 +from dimos_generated.vision_msgs.msg import Detection2DArray + from dimos.core.coordination.blueprints import autoconnect from dimos.core.transport import LCMTransport from dimos.hardware.sensors.camera.zed import compat as zed -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray from dimos.perception.detection.detectors.person.yolo import YoloPersonDetector from dimos.perception.detection.module3D import Detection3DModule from dimos.perception.detection.person_tracker import PersonTracker diff --git a/dimos/robot/unitree/g1/blueprints/primitive/test_unitree_g1_vis.py b/dimos/robot/unitree/g1/blueprints/primitive/test_unitree_g1_vis.py new file mode 100644 index 00000000000..8386659f8a6 --- /dev/null +++ b/dimos/robot/unitree/g1/blueprints/primitive/test_unitree_g1_vis.py @@ -0,0 +1,43 @@ +# Copyright 2025-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. + +from dimos_generated.geometry_msgs.msg import Point, PointStamped +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest + +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.robot.unitree.g1.blueprints.primitive.unitree_g1_vis import ( + _global_map_colors, + _goal_colors, + _waypoint_colors, +) + + +def test_global_map_uses_generated_xyz_and_height_colors(): + cloud = pointcloud_from_xyz(np.array([[1.0, 2.0, 0.0], [3.0, 4.0, 2.0]]), header=Header()) + result = _global_map_colors(type(cloud).decode(cloud.encode())) + np.testing.assert_array_equal( + result.positions.as_arrow_array().to_pylist(), [[1, 2, 0], [3, 4, 2]] + ) + assert result.colors.as_arrow_array().to_pylist() == [0x1E50C8FF, 0x3BDB64FF] + + +@pytest.mark.parametrize("render", [_goal_colors, _waypoint_colors]) +def test_generated_stamped_targets_keep_position_and_reject_nonfinite(render): + target = PointStamped(header=Header(frame_id="world"), point=Point(x=1, y=2, z=3)) + result = render(PointStamped.decode(target.encode())) + np.testing.assert_allclose(result.positions.as_arrow_array().to_pylist(), [[1, 2, 3.3]]) + target.point.x = float("nan") + assert render(target) is None diff --git a/dimos/robot/unitree/g1/blueprints/primitive/unitree_g1_vis.py b/dimos/robot/unitree/g1/blueprints/primitive/unitree_g1_vis.py index 8c352e5017c..236e03dd560 100644 --- a/dimos/robot/unitree/g1/blueprints/primitive/unitree_g1_vis.py +++ b/dimos/robot/unitree/g1/blueprints/primitive/unitree_g1_vis.py @@ -18,10 +18,13 @@ import math from typing import Any +from dimos_generated.geometry_msgs.msg import PointStamped +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.sensor_msgs.msg import PointCloud2 import numpy as np from dimos.core.global_config import global_config -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.pointcloud import pointcloud_xyz from dimos.robot.unitree.g1.g1_rerun import g1_costmap, g1_odometry_tf_override, g1_static_robot from dimos.visualization.vis_module import vis_module @@ -51,14 +54,17 @@ def _g1_path_colors(path: Path) -> Any: if not path.poses: return rr.LineStrips3D([]) - points = [[pose.x, pose.y, pose.z + _PATH_Z_LIFT] for pose in path.poses] + points = [ + [pose.pose.position.x, pose.pose.position.y, pose.pose.position.z + _PATH_Z_LIFT] + for pose in path.poses + ] return rr.LineStrips3D([points], colors=[_PATH_COLOR_RGBA], radii=_PATH_RADIUS_METERS) -def _global_map_colors(cloud: Any) -> Any: +def _global_map_colors(cloud: PointCloud2) -> Any: import rerun as rr - points, _ = cloud.as_numpy() + points = pointcloud_xyz(cloud) if len(points) == 0: return None @@ -76,27 +82,29 @@ def _global_map_colors(cloud: Any) -> Any: return rr.Points3D(positions=points[:, :3], colors=colors, radii=0.03) -def _waypoint_colors(waypoint: Any) -> Any: +def _waypoint_colors(waypoint: PointStamped) -> Any: import rerun as rr - if not all(math.isfinite(value) for value in (waypoint.x, waypoint.y, waypoint.z)): + if not all( + math.isfinite(value) for value in (waypoint.point.x, waypoint.point.y, waypoint.point.z) + ): return None return rr.Points3D( - positions=[[waypoint.x, waypoint.y, waypoint.z + _VIS_LIFT]], + positions=[[waypoint.point.x, waypoint.point.y, waypoint.point.z + _VIS_LIFT]], colors=[(255, 140, 0)], radii=0.22, ) -def _goal_colors(goal: Any) -> Any: +def _goal_colors(goal: PointStamped) -> Any: import rerun as rr - if not all(math.isfinite(value) for value in (goal.x, goal.y, goal.z)): + if not all(math.isfinite(value) for value in (goal.point.x, goal.point.y, goal.point.z)): return None return rr.Points3D( - positions=[[goal.x, goal.y, goal.z + _VIS_LIFT]], + positions=[[goal.point.x, goal.point.y, goal.point.z + _VIS_LIFT]], colors=[(180, 60, 220)], radii=0.3, ) diff --git a/dimos/robot/unitree/g1/config.py b/dimos/robot/unitree/g1/config.py index fae7d212e8a..4a423fc835e 100644 --- a/dimos/robot/unitree/g1/config.py +++ b/dimos/robot/unitree/g1/config.py @@ -20,9 +20,7 @@ from pathlib import Path from typing import Any -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos_generated.geometry_msgs.msg import Point, Pose @dataclass(frozen=True) @@ -43,6 +41,6 @@ class G1Config: width_clearance=0.6, internal_odom_offsets={ # Mid-360 lidar: 1.2 m above ground. - "mid360_link": Pose(0.0, 0.0, 1.2, *Quaternion.from_euler(Vector3(0, 0, 0))), + "mid360_link": Pose(position=Point(x=0.0, y=0.0, z=1.2)), }, ) diff --git a/dimos/robot/unitree/g1/connection.py b/dimos/robot/unitree/g1/connection.py index ed2c490e495..5dfc966598e 100644 --- a/dimos/robot/unitree/g1/connection.py +++ b/dimos/robot/unitree/g1/connection.py @@ -15,13 +15,13 @@ from abc import ABC, abstractmethod from typing import Any +from dimos_generated.geometry_msgs.msg import Twist from pydantic import Field from reactivex.disposable import Disposable from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In -from dimos.msgs.geometry_msgs.Twist import Twist from dimos.robot.unitree.connection import UnitreeWebRTCConnection from dimos.utils.logging_config import setup_logger diff --git a/dimos/robot/unitree/g1/connection_spec.py b/dimos/robot/unitree/g1/connection_spec.py index ed2d572df27..cfb5b93fc55 100644 --- a/dimos/robot/unitree/g1/connection_spec.py +++ b/dimos/robot/unitree/g1/connection_spec.py @@ -14,7 +14,8 @@ from typing import Any, Protocol -from dimos.msgs.geometry_msgs.Twist import Twist +from dimos_generated.geometry_msgs.msg import Twist + from dimos.spec.utils import Spec diff --git a/dimos/robot/unitree/g1/debug/demo_arrow_control.py b/dimos/robot/unitree/g1/debug/demo_arrow_control.py index 7224275ecf3..10dc3c11b18 100755 --- a/dimos/robot/unitree/g1/debug/demo_arrow_control.py +++ b/dimos/robot/unitree/g1/debug/demo_arrow_control.py @@ -20,8 +20,8 @@ import traceback from typing import Any -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos_generated.geometry_msgs.msg import Twist, Vector3 + from dimos.robot.unitree.g1.effectors.high_level.dds_sdk import G1HighLevelDdsSdk @@ -113,26 +113,39 @@ def main(stdscr: Any) -> None: # Arrow keys if key == curses.KEY_UP or key_char == "w": - twist = Twist(linear=Vector3(linear_speed, 0, 0), angular=Vector3(0, 0, 0)) + twist = Twist( + linear=Vector3(x=linear_speed, y=0, z=0), angular=Vector3(x=0, y=0, z=0) + ) action = "Moving forward..." elif key == curses.KEY_DOWN or key_char == "s": - twist = Twist(linear=Vector3(-linear_speed, 0, 0), angular=Vector3(0, 0, 0)) + twist = Twist( + linear=Vector3(x=-linear_speed, y=0, z=0), angular=Vector3(x=0, y=0, z=0) + ) action = "Moving backward..." elif key == curses.KEY_LEFT or key_char == "a": - twist = Twist(linear=Vector3(0, 0, 0), angular=Vector3(0, 0, angular_speed)) + twist = Twist( + linear=Vector3(x=0, y=0, z=0), angular=Vector3(x=0, y=0, z=angular_speed) + ) action = "Rotating left..." elif key == curses.KEY_RIGHT or key_char == "d": - twist = Twist(linear=Vector3(0, 0, 0), angular=Vector3(0, 0, -angular_speed)) + twist = Twist( + linear=Vector3(x=0, y=0, z=0), angular=Vector3(x=0, y=0, z=-angular_speed) + ) action = "Rotating right..." elif key_char == "q": - twist = Twist(linear=Vector3(0, linear_speed, 0), angular=Vector3(0, 0, 0)) + twist = Twist( + linear=Vector3(x=0, y=linear_speed, z=0), angular=Vector3(x=0, y=0, z=0) + ) action = "Strafing left..." elif key_char == "e": - twist = Twist(linear=Vector3(0, -linear_speed, 0), angular=Vector3(0, 0, 0)) + twist = Twist( + linear=Vector3(x=0, y=-linear_speed, z=0), angular=Vector3(x=0, y=0, z=0) + ) action = "Strafing right..." elif key_char == " ": conn.move( - Twist(linear=Vector3(0, 0, 0), angular=Vector3(0, 0, 0)), duration=move_duration + Twist(linear=Vector3(x=0, y=0, z=0), angular=Vector3(x=0, y=0, z=0)), + duration=move_duration, ) action = "🛑 Stopped" last_cmd_time = current_time diff --git a/dimos/robot/unitree/g1/debug/demo_arrow_control_cmd_vel.py b/dimos/robot/unitree/g1/debug/demo_arrow_control_cmd_vel.py index bd284428589..fab4326f17e 100644 --- a/dimos/robot/unitree/g1/debug/demo_arrow_control_cmd_vel.py +++ b/dimos/robot/unitree/g1/debug/demo_arrow_control_cmd_vel.py @@ -20,18 +20,17 @@ import traceback from typing import Any +from dimos_generated.geometry_msgs.msg import Twist, Vector3 import lcm from dimos.core.transport import LCMTransport -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.robot.unitree.g1.effectors.high_level.dds_sdk import G1HighLevelDdsSdk -CMD_VEL_CHANNEL = "/cmd_vel#geometry_msgs.Twist" +CMD_VEL_CHANNEL = "/cmd_vel#geometry_msgs/msg/Twist" def publish_twist(lc: lcm.LCM, twist: Twist) -> None: - lc.publish(CMD_VEL_CHANNEL, twist.lcm_encode()) + lc.publish(CMD_VEL_CHANNEL, twist.encode()) def draw_ui(stdscr: Any, state_text: str = "Not connected") -> None: @@ -116,25 +115,37 @@ def main(stdscr: Any) -> None: action = None if key == curses.KEY_UP or key_char == "w": - twist = Twist(linear=Vector3(linear_speed, 0, 0), angular=Vector3(0, 0, 0)) + twist = Twist( + linear=Vector3(x=linear_speed, y=0, z=0), angular=Vector3(x=0, y=0, z=0) + ) action = "Moving forward..." elif key == curses.KEY_DOWN or key_char == "s": - twist = Twist(linear=Vector3(-linear_speed, 0, 0), angular=Vector3(0, 0, 0)) + twist = Twist( + linear=Vector3(x=-linear_speed, y=0, z=0), angular=Vector3(x=0, y=0, z=0) + ) action = "Moving backward..." elif key == curses.KEY_LEFT or key_char == "a": - twist = Twist(linear=Vector3(0, 0, 0), angular=Vector3(0, 0, angular_speed)) + twist = Twist( + linear=Vector3(x=0, y=0, z=0), angular=Vector3(x=0, y=0, z=angular_speed) + ) action = "Rotating left..." elif key == curses.KEY_RIGHT or key_char == "d": - twist = Twist(linear=Vector3(0, 0, 0), angular=Vector3(0, 0, -angular_speed)) + twist = Twist( + linear=Vector3(x=0, y=0, z=0), angular=Vector3(x=0, y=0, z=-angular_speed) + ) action = "Rotating right..." elif key_char == "q": - twist = Twist(linear=Vector3(0, linear_speed, 0), angular=Vector3(0, 0, 0)) + twist = Twist( + linear=Vector3(x=0, y=linear_speed, z=0), angular=Vector3(x=0, y=0, z=0) + ) action = "Strafing left..." elif key_char == "e": - twist = Twist(linear=Vector3(0, -linear_speed, 0), angular=Vector3(0, 0, 0)) + twist = Twist( + linear=Vector3(x=0, y=-linear_speed, z=0), angular=Vector3(x=0, y=0, z=0) + ) action = "Strafing right..." elif key_char == " ": - stop = Twist(linear=Vector3(0, 0, 0), angular=Vector3(0, 0, 0)) + stop = Twist(linear=Vector3(x=0, y=0, z=0), angular=Vector3(x=0, y=0, z=0)) publish_twist(lc, stop) action = "Stopped" last_cmd_time = current_time @@ -160,7 +171,7 @@ def main(stdscr: Any) -> None: pass finally: draw_ui(stdscr, "Stopping...") - stop = Twist(linear=Vector3(0, 0, 0), angular=Vector3(0, 0, 0)) + stop = Twist(linear=Vector3(x=0, y=0, z=0), angular=Vector3(x=0, y=0, z=0)) publish_twist(lc, stop) time.sleep(0.5) conn.disconnect() diff --git a/dimos/robot/unitree/g1/effectors/high_level/dds_sdk.py b/dimos/robot/unitree/g1/effectors/high_level/dds_sdk.py index df62f9c0b42..72849d2f9c4 100644 --- a/dimos/robot/unitree/g1/effectors/high_level/dds_sdk.py +++ b/dimos/robot/unitree/g1/effectors/high_level/dds_sdk.py @@ -20,6 +20,7 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import Twist, Vector3 from reactivex.disposable import Disposable from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import ( # type: ignore[import-not-found] MotionSwitcherClient, @@ -37,8 +38,6 @@ from dimos.core.global_config import GlobalConfig, global_config from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.robot.unitree.g1.effectors.high_level.commands import ( ARM_API_ID, ARM_COMMANDS, @@ -297,7 +296,7 @@ def move_velocity( self, x: float, y: float = 0.0, yaw: float = 0.0, duration: float = 0.0 ) -> str: """Move the robot at the given velocity for ``duration`` seconds.""" - twist = Twist(linear=Vector3(x, y, 0), angular=Vector3(0, 0, yaw)) + twist = Twist(linear=Vector3(x=x, y=y), angular=Vector3(z=yaw)) self.move(twist, duration=duration) return f"Started moving with velocity=({x}, {y}, {yaw}) for {duration} seconds" diff --git a/dimos/robot/unitree/g1/effectors/high_level/high_level_spec.py b/dimos/robot/unitree/g1/effectors/high_level/high_level_spec.py index 01dde2ab248..146b734c527 100644 --- a/dimos/robot/unitree/g1/effectors/high_level/high_level_spec.py +++ b/dimos/robot/unitree/g1/effectors/high_level/high_level_spec.py @@ -16,8 +16,9 @@ from typing import Any, Protocol +from dimos_generated.geometry_msgs.msg import Twist + from dimos.core.stream import In -from dimos.msgs.geometry_msgs.Twist import Twist from dimos.spec.utils import Spec diff --git a/dimos/robot/unitree/g1/effectors/high_level/webrtc.py b/dimos/robot/unitree/g1/effectors/high_level/webrtc.py index 7b9743f4d02..24acc0f20db 100644 --- a/dimos/robot/unitree/g1/effectors/high_level/webrtc.py +++ b/dimos/robot/unitree/g1/effectors/high_level/webrtc.py @@ -14,6 +14,7 @@ from typing import Any +from dimos_generated.geometry_msgs.msg import Twist, Vector3 from pydantic import Field from reactivex.disposable import Disposable @@ -22,8 +23,6 @@ from dimos.core.global_config import GlobalConfig, global_config from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.robot.unitree.connection import UnitreeWebRTCConnection from dimos.robot.unitree.g1.effectors.high_level.commands import ( ARM_API_ID, @@ -116,7 +115,7 @@ def move_velocity( args = { "x": 0.5, "y": 0.0, "yaw": 0.0, "duration": 2.0 } move_velocity(**args) """ - twist = Twist(linear=Vector3(x, y, 0), angular=Vector3(0, 0, yaw)) + twist = Twist(linear=Vector3(x=x, y=y), angular=Vector3(z=yaw)) self.move(twist, duration=duration) return f"Started moving with velocity=({x}, {y}, {yaw}) for {duration} seconds" diff --git a/dimos/robot/unitree/g1/g1_recorder.py b/dimos/robot/unitree/g1/g1_recorder.py index d6f8883d676..d80152d9f44 100644 --- a/dimos/robot/unitree/g1/g1_recorder.py +++ b/dimos/robot/unitree/g1/g1_recorder.py @@ -16,18 +16,17 @@ from __future__ import annotations +from dimos_generated.sensor_msgs.msg import CameraInfo, Image from pydantic import Field from dimos.core.stream import In from dimos.hardware.sensors.lidar.pointlio.recorder import PointlioRecorder, PointlioRecorderConfig -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image class G1RecorderConfig(PointlioRecorderConfig): - # don't compress + # Preserve depth pixels; lossless compression wraps generated CDR. stream_codecs: dict[str, str] = Field( - default_factory=lambda: {"realsense_depth_image": "lz4+lcm"} + default_factory=lambda: {"realsense_depth_image": "lz4+cdr"} ) diff --git a/dimos/robot/unitree/g1/g1_tf_publisher.py b/dimos/robot/unitree/g1/g1_tf_publisher.py index 7237a76218e..aba566175ea 100644 --- a/dimos/robot/unitree/g1/g1_tf_publisher.py +++ b/dimos/robot/unitree/g1/g1_tf_publisher.py @@ -26,16 +26,17 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import Transform, TransformStamped, Vector3 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage from pydantic import Field from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import Out -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.geometry import compose_transforms, inverse_transform, quaternion_from_euler +from dimos.msgs.time import time_from_nanoseconds from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -51,25 +52,32 @@ _WAIST_ROLL_ORIGIN = (-0.0039635, 0.0, 0.044) -def torso_to_mid360() -> Transform: - """torso_link -> the Mid-360's own frame: URDF mount pitch plus the upside-down roll.""" - return Transform( - translation=Vector3(*_TORSO_MID360_XYZ), - rotation=Quaternion.from_euler(Vector3(0.0, MID360_PITCH, 0.0)) - * Quaternion.from_euler(Vector3(math.pi, 0.0, 0.0)), - frame_id="torso_link", - child_frame_id="mid360_link", +def _mount( + parent: str, + child: str, + xyz: tuple[float, float, float] = (0.0, 0.0, 0.0), + roll: float = 0.0, + pitch: float = 0.0, + yaw: float = 0.0, +) -> TransformStamped: + return TransformStamped( + header=Header(frame_id=parent), + child_frame_id=child, + transform=Transform( + translation=Vector3(x=xyz[0], y=xyz[1], z=xyz[2]), + rotation=quaternion_from_euler(roll, pitch, yaw), + ), ) -def torso_to_d435() -> Transform: +def torso_to_mid360() -> TransformStamped: + """URDF mount pitch composed after the sensor's upside-down roll.""" + return _mount("torso_link", "mid360_link", _TORSO_MID360_XYZ, roll=math.pi, pitch=MID360_PITCH) + + +def torso_to_d435() -> TransformStamped: """torso_link -> d435_link, the URDF mount.""" - return Transform( - translation=Vector3(*_TORSO_D435_XYZ), - rotation=Quaternion.from_euler(Vector3(0.0, D435_PITCH, 0.0)), - frame_id="torso_link", - child_frame_id="d435_link", - ) + return _mount("torso_link", "d435_link", _TORSO_D435_XYZ, pitch=D435_PITCH) # rt/lowstate motor indices, ordering from make_humanoid_joints("g1"). @@ -78,34 +86,21 @@ def torso_to_d435() -> Transform: _WAIST_PITCH_IDX = 14 -def base_to_torso(waist_yaw: float, waist_roll: float, waist_pitch: float) -> Transform: +def base_to_torso(waist_yaw: float, waist_roll: float, waist_pitch: float) -> TransformStamped: """base_link -> torso_link through the g1.urdf waist chain.""" - yaw = Transform( - rotation=Quaternion.from_euler(Vector3(0.0, 0.0, waist_yaw)), - frame_id="base_link", - child_frame_id="waist_yaw_link", - ) - roll = Transform( - translation=Vector3(*_WAIST_ROLL_ORIGIN), - rotation=Quaternion.from_euler(Vector3(waist_roll, 0.0, 0.0)), - frame_id="waist_yaw_link", - child_frame_id="waist_roll_link", - ) - pitch = Transform( - rotation=Quaternion.from_euler(Vector3(0.0, waist_pitch, 0.0)), - frame_id="waist_roll_link", - child_frame_id="torso_link", - ) - return yaw + roll + pitch + yaw = _mount("base_link", "waist_yaw_link", yaw=waist_yaw) + roll = _mount("waist_yaw_link", "waist_roll_link", _WAIST_ROLL_ORIGIN, roll=waist_roll) + pitch = _mount("waist_roll_link", "torso_link", pitch=waist_pitch) + return compose_transforms(compose_transforms(yaw, roll), pitch) def mount_transforms( waist_yaw: float = 0.0, waist_roll: float = 0.0, waist_pitch: float = 0.0 -) -> list[Transform]: +) -> list[TransformStamped]: """The mount tree as published: rooted at mid360_link.""" return [ - -torso_to_mid360(), - -base_to_torso(waist_yaw, waist_roll, waist_pitch), + inverse_transform(torso_to_mid360()), + inverse_transform(base_to_torso(waist_yaw, waist_roll, waist_pitch)), torso_to_d435(), ] @@ -212,8 +207,8 @@ async def _publish_loop(self) -> None: with self._waist_lock: waist_yaw, waist_roll, waist_pitch = self._waist transforms = mount_transforms(waist_yaw, waist_roll, waist_pitch) - now = time.time() + now = time.time_ns() for transform in transforms: - transform.ts = now - self.tf.publish(TFMessage(*transforms)) + transform.header.stamp = time_from_nanoseconds(now) + self.tf.publish(TFMessage(transforms=transforms)) await asyncio.sleep(period) diff --git a/dimos/robot/unitree/g1/skill_container.py b/dimos/robot/unitree/g1/skill_container.py index c3825ed29cc..68f8327fe5e 100644 --- a/dimos/robot/unitree/g1/skill_container.py +++ b/dimos/robot/unitree/g1/skill_container.py @@ -19,11 +19,11 @@ import difflib +from dimos_generated.geometry_msgs.msg import Twist, Vector3 + from dimos.agents.annotation import skill from dimos.core.core import rpc from dimos.core.module import Module -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.robot.unitree.g1.connection_spec import G1ConnectionSpec from dimos.utils.logging_config import setup_logger @@ -89,7 +89,7 @@ def move(self, x: float, y: float = 0.0, yaw: float = 0.0, duration: float = 0.0 duration: How long to move (seconds) """ - twist = Twist(linear=Vector3(x, y, 0), angular=Vector3(0, 0, yaw)) + twist = Twist(linear=Vector3(x=x, y=y), angular=Vector3(z=yaw)) self._connection.move(twist, duration=duration) return f"Started moving with velocity=({x}, {y}, {yaw}) for {duration} seconds" diff --git a/dimos/robot/unitree/g1/test_g1_manipulation_integration.py b/dimos/robot/unitree/g1/test_g1_manipulation_integration.py index 283c2d2c691..60ab74a5705 100644 --- a/dimos/robot/unitree/g1/test_g1_manipulation_integration.py +++ b/dimos/robot/unitree/g1/test_g1_manipulation_integration.py @@ -20,12 +20,13 @@ pytest.importorskip("roboplan.cartesian_planning") pytest.importorskip("roboplan.rrt") +from dimos_generated.sensor_msgs.msg import JointState + from dimos.manipulation.planning.planners import roboplan_planner as roboplan_planner_module from dimos.manipulation.planning.planners.roboplan_config import RoboPlanPlannerConfig from dimos.manipulation.planning.spec.enums import PlanningStatus from dimos.manipulation.planning.spec.validation import prepare_robot_model from dimos.manipulation.planning.world import roboplan_world as roboplan_world_module -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.robot.unitree.g1.manip_config import ( G1_LEFT_ARM_JOINTS, G1_RIGHT_ARM_JOINTS, diff --git a/dimos/robot/unitree/g1/test_g1_tf_publisher.py b/dimos/robot/unitree/g1/test_g1_tf_publisher.py index 222e26d1b51..df689179879 100644 --- a/dimos/robot/unitree/g1/test_g1_tf_publisher.py +++ b/dimos/robot/unitree/g1/test_g1_tf_publisher.py @@ -16,8 +16,10 @@ import math -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos_generated.geometry_msgs.msg import Transform, TransformStamped, Vector3 +from dimos_generated.std_msgs.msg import Header + +from dimos.msgs.geometry import inverse_transform, quaternion_euler, transform_matrix from dimos.protocol.tf.tf import MultiTBuffer from dimos.robot.unitree.g1.g1_tf_publisher import ( D435_PITCH, @@ -40,7 +42,7 @@ def _buffer( waist_yaw: float = 0.0, waist_roll: float = 0.0, waist_pitch: float = 0.0, - live: Transform | None = None, + live: TransformStamped | None = None, ) -> MultiTBuffer: buffer = MultiTBuffer() buffer.receive_transform(*mount_transforms(waist_yaw, waist_roll, waist_pitch)) @@ -53,34 +55,35 @@ def test_rest_pose_offsets_match_urdf() -> None: """The lidar sits MOUNT_Z above base_link, the offset every ground projection uses.""" leg = _buffer().get("mid360_link", "base_link") assert leg is not None - base_to_sensor = -leg - assert abs(base_to_sensor.translation.z - MOUNT_Z) < 1e-6 - assert abs(base_to_sensor.translation.x - MOUNT_X) < 1e-6 + base_to_sensor = inverse_transform(leg) + assert abs(base_to_sensor.transform.translation.z - MOUNT_Z) < 1e-6 + assert abs(base_to_sensor.transform.translation.x - MOUNT_X) < 1e-6 def test_mid360_frame_is_upside_down() -> None: """The inverted mount points the sensor's z axis at the floor in the base frame.""" leg = _buffer().get("base_link", "mid360_link") assert leg is not None - z_axis = leg.rotation.rotate_vector(Vector3(0.0, 0.0, 1.0)) - assert z_axis.z < -0.99 + z_axis = transform_matrix(leg.transform)[:3, 2] + assert z_axis[2] < -0.99 def test_flipped_level_sensor_yields_level_base() -> None: """A standing robot reports a flipped sensor pose. base_link must come out level, below it.""" - live = Transform( - translation=Vector3(0.0, 0.0, LIDAR_HEIGHT), - rotation=torso_to_mid360().rotation, - frame_id="world", + live = TransformStamped( + header=Header(frame_id="world"), child_frame_id="mid360_link", + transform=Transform( + translation=Vector3(z=LIDAR_HEIGHT), rotation=torso_to_mid360().transform.rotation + ), ) base = _buffer(live=live).get("world", "base_link") assert base is not None - euler = base.rotation.euler - assert abs(euler.x) < 1e-6 - assert abs(euler.y) < 1e-6 - assert abs(euler.z) < 1e-6 - assert math.isclose(base.translation.z, LIDAR_HEIGHT - MOUNT_Z, abs_tol=1e-6) + euler = quaternion_euler(base.transform.rotation) + assert abs(euler[0]) < 1e-6 + assert abs(euler[1]) < 1e-6 + assert abs(euler[2]) < 1e-6 + assert math.isclose(base.transform.translation.z, LIDAR_HEIGHT - MOUNT_Z, abs_tol=1e-6) def test_waist_yaw_rotates_base_link_against_the_torso() -> None: @@ -88,32 +91,32 @@ def test_waist_yaw_rotates_base_link_against_the_torso() -> None: yaw = math.pi / 4 leg = _buffer(waist_yaw=yaw).get("torso_link", "base_link") assert leg is not None - assert abs(leg.rotation.euler.z - (-yaw)) < 1e-6 + assert abs(quaternion_euler(leg.transform.rotation)[2] - (-yaw)) < 1e-6 def test_waist_pitch_rotates_base_link_against_the_torso() -> None: pitch = 0.3 leg = _buffer(waist_pitch=pitch).get("torso_link", "base_link") assert leg is not None - assert abs(leg.rotation.euler.y - (-pitch)) < 1e-6 + assert abs(quaternion_euler(leg.transform.rotation)[1] - (-pitch)) < 1e-6 def test_rest_pose_base_to_torso_matches_urdf_offsets() -> None: rest = base_to_torso(0.0, 0.0, 0.0) - assert abs(rest.translation.x - PELVIS_TORSO_X) < 1e-6 - assert abs(rest.translation.z - PELVIS_TORSO_Z) < 1e-6 - assert abs(rest.rotation.euler.x) < 1e-6 - assert abs(rest.rotation.euler.y) < 1e-6 - assert abs(rest.rotation.euler.z) < 1e-6 + assert abs(rest.transform.translation.x - PELVIS_TORSO_X) < 1e-6 + assert abs(rest.transform.translation.z - PELVIS_TORSO_Z) < 1e-6 + assert abs(quaternion_euler(rest.transform.rotation)[0]) < 1e-6 + assert abs(quaternion_euler(rest.transform.rotation)[1]) < 1e-6 + assert abs(quaternion_euler(rest.transform.rotation)[2]) < 1e-6 def test_d435_hangs_off_base_link() -> None: """The tree is rooted at mid360_link, so the camera edge is reachable by composition.""" camera = _buffer().get("base_link", "d435_link") assert camera is not None - assert abs(camera.translation.x - (PELVIS_TORSO_X + 0.0576235)) < 1e-6 - assert abs(camera.translation.z - (PELVIS_TORSO_Z + 0.42987)) < 1e-6 - assert abs(camera.rotation.euler.y - D435_PITCH) < 1e-6 + assert abs(camera.transform.translation.x - (PELVIS_TORSO_X + 0.0576235)) < 1e-6 + assert abs(camera.transform.translation.z - (PELVIS_TORSO_Z + 0.42987)) < 1e-6 + assert abs(quaternion_euler(camera.transform.rotation)[1] - D435_PITCH) < 1e-6 def test_pelvis_height_matches_config_note() -> None: diff --git a/dimos/robot/unitree/g1/test_wholebody_connection.py b/dimos/robot/unitree/g1/test_wholebody_connection.py index 35c41badabc..4a6067981e9 100644 --- a/dimos/robot/unitree/g1/test_wholebody_connection.py +++ b/dimos/robot/unitree/g1/test_wholebody_connection.py @@ -17,13 +17,16 @@ from collections.abc import Iterator from types import SimpleNamespace +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.dimos_msgs.msg import MotorCommandArray +from dimos_generated.std_msgs.msg import Header from pydantic import ValidationError import pytest -from dimos.msgs.sensor_msgs.MotorCommandArray import MotorCommandArray from dimos.robot.unitree.g1.wholebody_connection import ( _NUM_MOTOR_SLOTS, _NUM_MOTORS, + G1LowStateSnapshot, G1WholeBodyConnection, G1WholeBodyConnectionConfig, ) @@ -130,3 +133,38 @@ def test_non_finite_soft_start_is_rejected(value): # and no way back to the commanded gains. with pytest.raises(ValidationError): G1WholeBodyConnectionConfig(soft_start_seconds=value) + + +def test_generated_feedback_has_one_exact_header_and_ros_quaternion(connection, monkeypatch): + joints, imu = [], [] + monkeypatch.setattr(connection.motor_states, "publish", joints.append) + monkeypatch.setattr(connection.imu, "publish", imu.append) + header = Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id="g1_pelvis") + sample = G1LowStateSnapshot( + positions=[0.1] * 29, + velocities=[0.2] * 29, + efforts=[0.3] * 29, + quaternion=(0.8, 0.0, 0.0, 0.6), + gyroscope=(1.0, 2.0, 3.0), + accelerometer=(4.0, 5.0, 6.0), + ) + connection._publish_motor_state_and_imu(header=header, sample=sample) + assert joints[0].header == header == imu[0].header + assert list(joints[0].position) == [0.1] * 29 + assert imu[0].orientation.w == 0.8 + assert imu[0].orientation.z == 0.6 + assert imu[0].angular_velocity.x == 1.0 + assert imu[0].linear_acceleration.z == 6.0 + for message in [joints[0], imu[0]]: + assert type(message).decode(message.encode()) == message + header.stamp.nanosec = 0 + assert joints[0].header.stamp.nanosec == 123456789 + + +@pytest.mark.parametrize("field", ["dq", "kp", "kd", "tau"]) +def test_inconsistent_motor_array_is_dropped(connection, field): + publisher = _wire(connection, soft_start_seconds=0.0) + command = _command() + setattr(command, field, [0.0]) + connection._on_motor_command(command) + assert publisher.frames == [] diff --git a/dimos/robot/unitree/g1/wholebody_connection.py b/dimos/robot/unitree/g1/wholebody_connection.py index 2f185d62e8b..2139ae6c3eb 100644 --- a/dimos/robot/unitree/g1/wholebody_connection.py +++ b/dimos/robot/unitree/g1/wholebody_connection.py @@ -27,6 +27,10 @@ import time from typing import TYPE_CHECKING, Any +from dimos_generated.dimos_msgs.msg import MotorCommandArray +from dimos_generated.geometry_msgs.msg import Quaternion, Vector3 +from dimos_generated.sensor_msgs.msg import Imu, JointState +from dimos_generated.std_msgs.msg import Header from pydantic import Field from reactivex.disposable import Disposable @@ -41,11 +45,7 @@ from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out from dimos.hardware.whole_body.spec import POS_STOP, VEL_STOP -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.Imu import Imu -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.sensor_msgs.MotorCommandArray import MotorCommandArray +from dimos.msgs.time import header_now from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -69,16 +69,14 @@ def _imu_from_unitree_wxyz( gyroscope: tuple[float, float, float], accelerometer: tuple[float, float, float], *, - frame_id: str, - ts: float, + header: Header, ) -> Imu: _w, x, y, z = quaternion return Imu( - orientation=Quaternion(x, y, z, _w), - angular_velocity=Vector3(*gyroscope), - linear_acceleration=Vector3(*accelerometer), - frame_id=frame_id, - ts=ts, + orientation=Quaternion(x=x, y=y, z=z, w=_w), + angular_velocity=Vector3(x=gyroscope[0], y=gyroscope[1], z=gyroscope[2]), + linear_acceleration=Vector3(x=accelerometer[0], y=accelerometer[1], z=accelerometer[2]), + header=header, ) @@ -336,13 +334,10 @@ def _snapshot_motor_imu(self) -> G1LowStateSnapshot | None: ), ) - def _publish_motor_state_and_imu( - self, now: float, frame_id: str, sample: G1LowStateSnapshot - ) -> None: + def _publish_motor_state_and_imu(self, header: Header, sample: G1LowStateSnapshot) -> None: self.motor_states.publish( JointState( - ts=now, - frame_id=frame_id, + header=header, name=G1_JOINT_NAMES, position=sample.positions, velocity=sample.velocities, @@ -355,8 +350,7 @@ def _publish_motor_state_and_imu( sample.quaternion, sample.gyroscope, sample.accelerometer, - frame_id=frame_id, - ts=now, + header=header, ) ) @@ -369,12 +363,12 @@ def _publish_loop(self) -> None: self._drain_low_state() sample = self._snapshot_motor_imu() if sample is not None: - self._publish_motor_state_and_imu(now=time.time(), frame_id=frame_id, sample=sample) + self._publish_motor_state_and_imu(header=header_now(frame_id), sample=sample) next_tick += period sleep_for = next_tick - time.perf_counter() if sleep_for > 0: - time.sleep(sleep_for) + self._stop_event.wait(sleep_for) else: next_tick = time.perf_counter() @@ -396,8 +390,8 @@ def _soft_start_scale(self, now: float) -> float: return scale def _on_motor_command(self, msg: MotorCommandArray) -> None: - if msg.num_joints != _NUM_MOTORS: - logger.warning(f"Expected {_NUM_MOTORS} motor commands, got {msg.num_joints}; ignoring") + if any(len(values) != _NUM_MOTORS for values in (msg.q, msg.dq, msg.kp, msg.kd, msg.tau)): + logger.warning(f"Expected {_NUM_MOTORS} values in each motor-command array; ignoring") return with self._lock: diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_coordinator.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_coordinator.py index d2a4c0454a8..d43a4af6ad3 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_coordinator.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_coordinator.py @@ -22,12 +22,12 @@ from __future__ import annotations +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist + from dimos.control.components import HardwareComponent, HardwareType, make_twist_base_joints from dimos.control.coordinator import ControlCoordinator, TaskConfig from dimos.core.coordination.blueprints import autoconnect from dimos.core.transport import LCMTransport -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist from dimos.robot.unitree.go2.connection import GO2Connection _go2_joints = make_twist_base_joints("go2") diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_holonomic_benchmark.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_holonomic_benchmark.py index 632b46a5dcf..86d0fec2601 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_holonomic_benchmark.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_holonomic_benchmark.py @@ -37,11 +37,12 @@ from typing import Any +from dimos_generated.geometry_msgs.msg import PoseStamped + from dimos.control.benchmarking.benchmark import Benchmarker from dimos.core.coordination.blueprints import TransportSpec, autoconnect from dimos.core.stream import Transport from dimos.core.transport import LCMTransport -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_holonomic_controller import ( unitree_go2_holonomic_controller, ) diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_holonomic_controller.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_holonomic_controller.py index 762727023e3..f19d5b7d26c 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_holonomic_controller.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_holonomic_controller.py @@ -49,18 +49,17 @@ from __future__ import annotations +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Float32, Int8 + from dimos.control.components import HardwareComponent, HardwareType, make_twist_base_joints from dimos.control.coordinator import TaskConfig from dimos.control.path_following_coordinator import PathFollowingCoordinator from dimos.core.coordination.blueprints import autoconnect from dimos.core.stream import Out from dimos.core.transport import LCMTransport -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.std_msgs.Float32 import Float32 -from dimos.msgs.std_msgs.Int8 import Int8 from dimos.robot.unitree.go2.connection import GO2Connection from dimos.robot.unitree.keyboard_teleop import KeyboardTeleop diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_rpp_benchmark.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_rpp_benchmark.py index 965f845b232..2e0cd11ab00 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_rpp_benchmark.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_rpp_benchmark.py @@ -50,11 +50,12 @@ from typing import Any +from dimos_generated.geometry_msgs.msg import PoseStamped + from dimos.control.benchmarking.benchmark import Benchmarker from dimos.core.coordination.blueprints import TransportSpec, autoconnect from dimos.core.stream import Transport from dimos.core.transport import LCMTransport -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_rpp_controller import ( unitree_go2_rpp_controller, ) diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_rpp_controller.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_rpp_controller.py index fce12dff27f..eb24b6b4583 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_rpp_controller.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_rpp_controller.py @@ -52,18 +52,17 @@ from __future__ import annotations +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Float32, Int8 + from dimos.control.components import HardwareComponent, HardwareType, make_twist_base_joints from dimos.control.coordinator import TaskConfig from dimos.control.path_following_coordinator import PathFollowingCoordinator from dimos.core.coordination.blueprints import autoconnect from dimos.core.stream import Out from dimos.core.transport import LCMTransport -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.std_msgs.Float32 import Float32 -from dimos.msgs.std_msgs.Int8 import Int8 from dimos.robot.unitree.go2.connection import GO2Connection from dimos.robot.unitree.keyboard_teleop import KeyboardTeleop diff --git a/dimos/robot/unitree/go2/blueprints/smart/test_unitree_go2_markers.py b/dimos/robot/unitree/go2/blueprints/smart/test_unitree_go2_markers.py index bf33a2e0888..213ffb1125f 100644 --- a/dimos/robot/unitree/go2/blueprints/smart/test_unitree_go2_markers.py +++ b/dimos/robot/unitree/go2/blueprints/smart/test_unitree_go2_markers.py @@ -29,7 +29,7 @@ def test_unitree_go2_markers_uses_detector_backed_tf_stack() -> None: bp for bp in unitree_go2_markers.blueprints if bp.module is MarkerDetectionStreamModule ) assert detector.kwargs["marker_length_m"] == 0.1 - assert detector.kwargs["camera_info"].frame_id == "camera_optical" + assert detector.kwargs["camera_info"].header.frame_id == "camera_optical" assert ( unitree_go2_markers.transport_map[("detections", MarkerDetectionStreamModule)].topic.topic == "/marker_detection/detections" diff --git a/dimos/robot/unitree/go2/blueprints/smart/unitree_go2_detection.py b/dimos/robot/unitree/go2/blueprints/smart/unitree_go2_detection.py index 2d021657a2c..1c94bd5b410 100644 --- a/dimos/robot/unitree/go2/blueprints/smart/unitree_go2_detection.py +++ b/dimos/robot/unitree/go2/blueprints/smart/unitree_go2_detection.py @@ -13,11 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 +from dimos_generated.vision_msgs.msg import Detection2DArray + from dimos.core.coordination.blueprints import autoconnect from dimos.core.transport import LCMTransport -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray from dimos.perception.detection.module3D import Detection3DModule from dimos.robot.unitree.go2.blueprints.smart.unitree_go2 import unitree_go2 from dimos.robot.unitree.go2.connection import GO2Connection diff --git a/dimos/robot/unitree/go2/blueprints/smart/unitree_go2_ros.py b/dimos/robot/unitree/go2/blueprints/smart/unitree_go2_ros.py index ba35872bd98..2a05bd91680 100644 --- a/dimos/robot/unitree/go2/blueprints/smart/unitree_go2_ros.py +++ b/dimos/robot/unitree/go2/blueprints/smart/unitree_go2_ros.py @@ -13,10 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 + from dimos.core.transport import ROSTransport -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.robot.unitree.go2.blueprints.smart.unitree_go2 import unitree_go2 unitree_go2_ros = unitree_go2.transports( diff --git a/dimos/robot/unitree/go2/camera_calibration.py b/dimos/robot/unitree/go2/camera_calibration.py new file mode 100644 index 00000000000..7a79c10b353 --- /dev/null +++ b/dimos/robot/unitree/go2/camera_calibration.py @@ -0,0 +1,28 @@ +# 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. + +"""The packaged Go2 front-camera calibration, independent of its robot connection.""" + +from importlib import resources + +from dimos_generated.sensor_msgs.msg import CameraInfo +from dimos_generated.std_msgs.msg import Header + +from dimos.msgs.camera_info import camera_info_from_yaml + + +def front_camera_calibration() -> CameraInfo: + resource = resources.files("dimos.robot.unitree.go2").joinpath("front_camera_720.yaml") + with resources.as_file(resource) as path: + return camera_info_from_yaml(path, header=Header(frame_id="camera_optical")) diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index 6d95f99e75c..9ef0c858425 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -14,12 +14,22 @@ import copy from enum import Enum -from importlib import resources import sys from threading import Lock, Thread import time from typing import Any, Protocol +from dimos_generated.geometry_msgs.msg import ( + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Twist, + Vector3, +) +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage from pydantic import Field from reactivex import defer, empty, operators as ops from reactivex.abc import SchedulerBase @@ -35,16 +45,10 @@ from dimos.core.stream import In, Out from dimos.memory.replay import Replay, ReplayStream, resolve_db_path from dimos.memory.store.sqlite import SqliteStore -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.geometry import compose_transforms, transform_from_pose +from dimos.msgs.time import time_from_nanoseconds from dimos.robot.unitree.connection import UnitreeWebRTCConnection +from dimos.robot.unitree.go2.camera_calibration import front_camera_calibration from dimos.robot.unitree.type.lowstate import LowStateMsg from dimos.spec.perception import Camera, Pointcloud from dimos.utils.decorators.decorators import cached_property, simple_mcache @@ -103,16 +107,6 @@ def switch_joystick(self, enable: bool = True) -> bool: ... def publish_request(self, topic: str, data: dict) -> dict: ... # type: ignore[type-arg] -_FRONT_CAMERA_720_YAML = resources.files("dimos.robot.unitree.go2").joinpath( - "front_camera_720.yaml" -) - - -def _camera_info_static() -> CameraInfo: - with resources.as_file(_FRONT_CAMERA_720_YAML) as yaml_path: - return CameraInfo.from_yaml(str(yaml_path)) - - def _prefixed(prefix: str | None, name: str) -> str: """Apply a TF namespace prefix (ModuleConfig.frame_id_prefix) to a frame name.""" if not prefix or not name: @@ -122,16 +116,17 @@ def _prefixed(prefix: str | None, name: str) -> str: # Static camera mount chain: base_link -> camera_link -> camera_optical. # TODO we need a standardized way to specify this for all cameras in dimos -BASE_TO_OPTICAL: Transform = Transform( - translation=Vector3(0.3, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="base_link", - child_frame_id="camera_link", -) + Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(-0.5, 0.5, -0.5, 0.5), - frame_id="camera_link", - child_frame_id="camera_optical", +BASE_TO_OPTICAL = compose_transforms( + TransformStamped( + header=Header(frame_id="base_link"), + child_frame_id="camera_link", + transform=Transform(translation=Vector3(x=0.3)), + ), + TransformStamped( + header=Header(frame_id="camera_link"), + child_frame_id="camera_optical", + transform=Transform(rotation=Quaternion(x=-0.5, y=0.5, z=-0.5, w=0.5)), + ), ) @@ -352,7 +347,7 @@ class GO2Connection(Module, Camera, Pointcloud): tf: Out[TFMessage] connection: Go2ConnectionProtocol - camera_info_static: CameraInfo = _camera_info_static() + camera_info_static: CameraInfo = front_camera_calibration() _camera_info_thread: Thread | None = None _latest_lowstate: LowStateMsg | None = None @@ -380,11 +375,11 @@ def __init__(self, **kwargs: Any) -> None: if hasattr(self.connection, "camera_info_static"): self.camera_info_static = self.connection.camera_info_static - if self.config.frame_id_prefix and self.camera_info_static.frame_id: + if self.config.frame_id_prefix and self.camera_info_static.header.frame_id: # Copy so the class-level default is not mutated. self.camera_info_static = copy.copy(self.camera_info_static) - self.camera_info_static.frame_id = _prefixed( - self.config.frame_id_prefix, self.camera_info_static.frame_id + self.camera_info_static.header.frame_id = _prefixed( + self.config.frame_id_prefix, self.camera_info_static.header.frame_id ) @rpc @@ -395,7 +390,7 @@ def start(self) -> None: self.connection.start() def onimage(image: Image) -> None: - image.frame_id = _prefixed(self.config.frame_id_prefix, image.frame_id) + image.header.frame_id = _prefixed(self.config.frame_id_prefix, image.header.frame_id) self.color_image.publish(image) if self.config.lidar: @@ -451,42 +446,36 @@ def stop(self) -> None: super().stop() @classmethod - def _odom_to_tf(cls, odom: PoseStamped, prefix: str = "") -> list[Transform]: - # The odom parent frame (odom.frame_id) stays unprefixed so namespaced - # robots still hang off one shared tree root. - camera_link = Transform( - translation=Vector3(0.3, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id=_prefixed(prefix, "base_link"), - child_frame_id=_prefixed(prefix, "camera_link"), - ts=odom.ts, - ) - - camera_optical = Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(-0.5, 0.5, -0.5, 0.5), - frame_id=_prefixed(prefix, "camera_link"), - child_frame_id=_prefixed(prefix, "camera_optical"), - ts=odom.ts, - ) - + def _odom_to_tf(cls, odom: PoseStamped, prefix: str = "") -> list[TransformStamped]: + # The odometry parent stays global; only robot-local frames get a prefix. return [ - Transform.from_pose(_prefixed(prefix, "base_link"), odom), - camera_link, - camera_optical, + transform_from_pose(odom, child_frame_id=_prefixed(prefix, "base_link")), + TransformStamped( + header=Header(stamp=odom.header.stamp, frame_id=_prefixed(prefix, "base_link")), + child_frame_id=_prefixed(prefix, "camera_link"), + transform=Transform(translation=Vector3(x=0.3)), + ), + TransformStamped( + header=Header(stamp=odom.header.stamp, frame_id=_prefixed(prefix, "camera_link")), + child_frame_id=_prefixed(prefix, "camera_optical"), + transform=Transform(rotation=Quaternion(x=-0.5, y=0.5, z=-0.5, w=0.5)), + ), ] def _publish_tf(self, msg: PoseStamped) -> None: - msg.frame_id = self.config.odom_frame_id + message = copy.copy(msg) + message.header.frame_id = self.config.odom_frame_id if self.config.publish_tf: - transforms = self._odom_to_tf(msg, prefix=self.config.frame_id_prefix or "") - self.tf.publish(TFMessage(*transforms)) + transforms = self._odom_to_tf(message, prefix=self.config.frame_id_prefix or "") + self.tf.publish(TFMessage(transforms=transforms)) if self.odom.transport: - self.odom.publish(msg) + self.odom.publish(message) def publish_camera_info(self) -> None: while True: - self.camera_info.publish(self.camera_info_static.with_ts(time.time())) + message = copy.copy(self.camera_info_static) + message.header.stamp = time_from_nanoseconds(time.time_ns()) + self.camera_info.publish(message) time.sleep(1.0) @rpc diff --git a/dimos/robot/unitree/go2/dds/cli/render.py b/dimos/robot/unitree/go2/dds/cli/render.py index 7711e9cf0b7..0eae19fbb10 100644 --- a/dimos/robot/unitree/go2/dds/cli/render.py +++ b/dimos/robot/unitree/go2/dds/cli/render.py @@ -34,26 +34,42 @@ import subprocess from typing import TYPE_CHECKING, Any, cast +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + TransformStamped, + Twist, + TwistStamped, + Vector3, +) +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Header import numpy as np +from scipy.spatial.transform import Rotation import typer from dimos.memory.transform import throttle from dimos.memory.utils.progress import progress -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import compose_transforms, transform_from_matrix +from dimos.msgs.pointcloud import transform_cloud +from dimos.msgs.time import time_from_seconds from dimos.robot.unitree.go2.dds.extrinsics import LIDAR_TO_BASE from dimos.robot.unitree.go2.dds.msgs.SportModeState import SportModeState from dimos.robot.unitree.go2.dds.store import Go2McapStore +from dimos.visualization.rerun.message_helpers import ( + cloud_archetype, + image_archetype, + navigation_archetype, +) if TYPE_CHECKING: from collections.abc import Iterator + from dimos_generated.sensor_msgs.msg import Image, PointCloud2 + from dimos.memory.type.observation import Observation - from dimos.msgs.sensor_msgs.Image import Image - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 WORLD = "world" @@ -61,8 +77,8 @@ def world_accel(obs: Observation[Any]) -> np.ndarray: """IMU linear acceleration rotated into the world frame (includes gravity).""" la = obs.data.linear_acceleration - a = obs.data.orientation.rotate_vector(Vector3(la.x, la.y, la.z)) - return np.array([float(a.x), float(a.y), float(a.z)]) + q = obs.data.orientation + return np.asarray(Rotation.from_quat([q.x, q.y, q.z, q.w]).apply([la.x, la.y, la.z])) def gravity_bias( @@ -109,10 +125,11 @@ def sportmode_pose(obs: Observation[SportModeState]) -> PoseStamped: sm = obs.data w, x, y, z = (float(v) for v in sm.imu_state.quaternion) # Unitree order: wxyz return PoseStamped( - ts=obs.ts, - frame_id=WORLD, - position=[float(v) for v in sm.position], - orientation=[x, y, z, w], + header=Header(frame_id=WORLD, stamp=time_from_seconds(obs.ts)), + pose=Pose( + position=Point(x=sm.position[0], y=sm.position[1], z=sm.position[2]), + orientation=Quaternion(x=x, y=y, z=z, w=w), + ), ) @@ -127,7 +144,8 @@ def step(state: Any, obs: Observation[Any]) -> tuple[Any, TwistStamped]: if prev is not None: vel = vel + a_world * (obs.ts - prev) twist = TwistStamped( - ts=obs.ts, frame_id=WORLD, linear=vel.tolist(), angular=[0.0, 0.0, 0.0] + header=Header(frame_id=WORLD, stamp=time_from_seconds(obs.ts)), + twist=Twist(linear=Vector3(x=vel[0], y=vel[1], z=vel[2])), ) return (vel, obs.ts), twist @@ -137,20 +155,21 @@ def step(state: Any, obs: Observation[Any]) -> tuple[Any, TwistStamped]: def integrate_position(state: Any, obs: Observation[Any]) -> tuple[Any, PoseStamped]: """scan_data: TwistStamped -> PoseStamped (velocity integrated to position).""" pos, prev = state - v = obs.data.linear + v = obs.data.twist.linear if prev is not None: pos = pos + np.array([v.x, v.y, v.z]) * (obs.ts - prev) pose = PoseStamped( - ts=obs.ts, frame_id=WORLD, position=pos.tolist(), orientation=[0.0, 0.0, 0.0, 1.0] + header=Header(frame_id=WORLD, stamp=time_from_seconds(obs.ts)), + pose=Pose(position=Point(x=pos[0], y=pos[1], z=pos[2]), orientation=Quaternion(w=1)), ) return (pos, obs.ts), pose def accumulate_path(upstream: Iterator[Observation[PoseStamped]]) -> Iterator[Observation[Path]]: """transform: yield the growing nav_msgs/Path as each pose streams in.""" - path = Path(frame_id=WORLD) + path = Path(header=Header(frame_id=WORLD)) for obs in upstream: - path = path.push(obs.data) + path = Path(header=obs.data.header, poses=[*path.poses, obs.data]) yield obs.derive(data=path) @@ -160,11 +179,13 @@ def leg_odom(store: Go2McapStore, seconds: float | None) -> None: def log_pose(obs: Observation[PoseStamped]) -> None: rr.set_time("time", timestamp=obs.ts) - rr.log("world/leg_odom", obs.data.to_rerun(), rr.TransformAxes3D(axis_length=0.2)) + rr.log( + "world/leg_odom", navigation_archetype(obs.data), rr.TransformAxes3D(axis_length=0.2) + ) def log_path(obs: Observation[Path]) -> None: rr.set_time("time", timestamp=obs.ts) - rr.log("world/leg_odom_path", obs.data.to_rerun()) + rr.log("world/leg_odom_path", navigation_archetype(obs.data)) src = store.streams.sportmodestate.to_time(seconds) with progress(src.count(), "leg_odom") as bar: @@ -185,7 +206,7 @@ def imu_odom(store: Go2McapStore, seconds: float | None) -> None: def log_path(obs: Observation[Path]) -> None: rr.set_time("time", timestamp=obs.ts) - rr.log("world/imu_odom_path", obs.data.to_rerun(color=(220, 90, 90))) + rr.log("world/imu_odom_path", navigation_archetype(obs.data, color=(220, 90, 90))) gravity = gravity_bias(store) # calibrate gravity+bias on the stationary start src = store.streams.imu.to_time(seconds) @@ -207,10 +228,17 @@ def lidar(store: Go2McapStore, seconds: float | None) -> None: def log_lidar(obs: Observation[PointCloud2]) -> None: rr.set_time("time", timestamp=obs.ts) - rr.log("world/leg_odom/lidar", obs.data.to_rerun()) + rr.log("world/leg_odom/lidar", cloud_archetype(obs.data)) src = store.streams.lidar.to_time(seconds) - rr.log("world/leg_odom/lidar", LIDAR_TO_BASE.to_rerun(frameless=True), static=True) + t, q = LIDAR_TO_BASE.transform.translation, LIDAR_TO_BASE.transform.rotation + rr.log( + "world/leg_odom/lidar", + rr.Transform3D( + translation=[t.x, t.y, t.z], quaternion=rr.Quaternion(xyzw=[q.x, q.y, q.z, q.w]) + ), + static=True, + ) with progress(src.count(), "lidar") as bar: src.tap(bar).tap(log_lidar).drain() @@ -248,15 +276,19 @@ def world_lidar(store: Go2McapStore, seconds: float | None) -> None: def to_world(obs: Observation[PointCloud2]) -> PointCloud2: p, q = _interp_pose(tt, pos, quat, obs.ts) - b2w = Transform.from_pose( - WORLD, - PoseStamped(ts=obs.ts, frame_id=WORLD, position=p.tolist(), orientation=q.tolist()), + matrix = np.eye(4) + matrix[:3, :3] = Rotation.from_quat(q).as_matrix() + matrix[:3, 3] = p + b2w = TransformStamped( + header=Header(frame_id=WORLD, stamp=obs.data.header.stamp), + child_frame_id=ext.header.frame_id, + transform=transform_from_matrix(matrix), ) - return obs.data.transform(b2w.apply(ext)) # lidar -> base -> world + return transform_cloud(obs.data, compose_transforms(b2w, ext)) def log_voxels(obs: Observation[PointCloud2]) -> None: rr.set_time("time", timestamp=obs.ts) - rr.log("world/world_lidar", obs.data.to_rerun()) + rr.log("world/world_lidar", cloud_archetype(obs.data)) src = store.streams.lidar.to_time(seconds) with progress(src.count(), "world_lidar") as bar: @@ -281,7 +313,7 @@ def log_image(obs: Observation[Image]) -> None: if obs.data is None: # truncated/corrupt frame return rr.set_time("time", timestamp=obs.ts) - rr.log("world/camera", obs.data.to_rerun()) + rr.log("world/camera", image_archetype(obs.data)) src = store.streams.color_image.to_time(seconds) with progress(src.count(), "camera") as bar: diff --git a/dimos/robot/unitree/go2/dds/cli/test_render_generated.py b/dimos/robot/unitree/go2/dds/cli/test_render_generated.py new file mode 100644 index 00000000000..864a5565d67 --- /dev/null +++ b/dimos/robot/unitree/go2/dds/cli/test_render_generated.py @@ -0,0 +1,106 @@ +# 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. + +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +import math +from types import SimpleNamespace + +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Vector3 +from dimos_generated.sensor_msgs.msg import Imu +from dimos_generated.std_msgs.msg import Header +import numpy as np + +from dimos.memory.type.observation import Observation +from dimos.msgs.geometry import transform_matrix +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_xyz, transform_cloud +from dimos.msgs.time import time_from_seconds +from dimos.robot.unitree.go2.dds.cli.render import ( + accumulate_path, + integrate_position, + integrate_velocity, + sportmode_pose, + world_accel, +) +from dimos.robot.unitree.go2.dds.extrinsics import EXT_R, EXT_T, LIDAR_TO_BASE + + +def test_world_acceleration_and_integrators_keep_world_axes_and_clock(): + imu = Imu( + orientation=Quaternion(z=math.sqrt(0.5), w=math.sqrt(0.5)), + linear_acceleration=Vector3(x=1, z=9.81), + ) + observations = [Observation(ts=ts, _data=imu) for ts in (2.0, 3.0)] + np.testing.assert_allclose(world_accel(observations[0]), [0, 1, 9.81], atol=1e-12) + velocity_step = integrate_velocity(np.array([0.0, 0.0, 9.81])) + state = (np.zeros(3), None) + twists = [] + for obs in observations: + state, twist = velocity_step(state, obs) + twists.append(Observation(ts=obs.ts, _data=twist)) + assert twist.header.stamp == time_from_seconds(obs.ts) + assert twist.header.frame_id == "world" + position_state = (np.zeros(3), None) + for obs in twists: + position_state, pose = integrate_position(position_state, obs) + np.testing.assert_allclose( + [pose.pose.position.x, pose.pose.position.y, pose.pose.position.z], [0, 1, 0], atol=1e-12 + ) + assert pose.header.stamp == time_from_seconds(3.0) + + +def test_path_accumulation_keeps_previous_snapshot_and_exact_pose_headers(): + poses = [ + PoseStamped( + header=Header(frame_id="world", stamp=time_from_seconds(ts)), + pose=Pose(position=Point(x=ts), orientation=Quaternion(w=1)), + ) + for ts in (1.0, 2.0) + ] + result = list( + accumulate_path( + iter([Observation(ts=ts, _data=p) for ts, p in zip((1.0, 2.0), poses, strict=True)]) + ) + ) + assert [len(obs.data.poses) for obs in result] == [1, 2] + assert result[1].data.header == poses[1].header + assert result[0].data.poses[0].encode() == poses[0].encode() + + +def test_sportmode_wxyz_is_converted_to_generated_xyzw(): + sm = SimpleNamespace( + position=[1.0, 2.0, 3.0], imu_state=SimpleNamespace(quaternion=[0.5, -0.5, 0.5, -0.5]) + ) + pose = sportmode_pose(Observation(ts=4.25, _data=sm)) + assert pose.header.stamp == time_from_seconds(4.25) + q = pose.pose.orientation + assert (q.x, q.y, q.z, q.w) == (-0.5, 0.5, -0.5, 0.5) + assert PoseStamped.decode(pose.encode()).pose.position.y == 2.0 + + +def test_measured_lidar_mount_retains_coordinates_and_cloud_stamp(): + cloud = pointcloud_from_xyz( + np.array([[1.0, 0.0, 0.0]]), header=Header(frame_id="lidar", stamp=time_from_seconds(12.25)) + ) + original = cloud.encode() + transformed = transform_cloud(cloud, LIDAR_TO_BASE) + matrix = transform_matrix(LIDAR_TO_BASE.transform) + np.testing.assert_allclose(matrix[:3, :3], EXT_R, atol=1e-6) + np.testing.assert_allclose( + pointcloud_xyz(transformed), (matrix[:3, :3] @ [1.0, 0.0, 0.0] + EXT_T)[None, :], atol=1e-6 + ) + assert transformed.header.stamp == cloud.header.stamp + assert transformed.header.frame_id == "base_link" + assert cloud.encode() == original diff --git a/dimos/robot/unitree/go2/dds/codec.py b/dimos/robot/unitree/go2/dds/codec.py index ed645d77e04..0d3b136698a 100644 --- a/dimos/robot/unitree/go2/dds/codec.py +++ b/dimos/robot/unitree/go2/dds/codec.py @@ -31,12 +31,11 @@ import json from typing import Any, Protocol, runtime_checkable -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.Imu import Imu -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos_generated.foxglove_msgs.msg import CompressedVideo +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import CompressedImage, Imu, PointCloud2 + from dimos.robot.unitree.go2.dds import cdr, ros -from dimos.robot.unitree.go2.dds.msgs.CompressedVideo import CompressedVideo from dimos.robot.unitree.go2.dds.msgs.ControlEvent import ControlEvent from dimos.robot.unitree.go2.dds.msgs.HeightMap import HeightMap from dimos.robot.unitree.go2.dds.msgs.LowCmd import LowCmd @@ -114,7 +113,7 @@ def encode(self, msg: Any) -> bytes: "rt/utlidar/height_map_array": FnCodec(HeightMap, ros.decode_height_map), "rt/utlidar/imu": FnCodec(Imu, ros.decode_imu), "rt/utlidar/robot_odom": FnCodec(Odometry, ros.decode_odometry), - "rt/frontvideo": FnCodec(Image, ros.decode_compressed_image), + "rt/frontvideo": FnCodec(CompressedImage, ros.decode_compressed_image), "rt/frontvideo/h264": FnCodec(CompressedVideo, ros.decode_compressed_video), "rt/lowstate": CdrStructCodec(LowState), "rt/lowcmd": CdrStructCodec(LowCmd), diff --git a/dimos/robot/unitree/go2/dds/extrinsics.py b/dimos/robot/unitree/go2/dds/extrinsics.py index 8db4f797534..10d5002c601 100644 --- a/dimos/robot/unitree/go2/dds/extrinsics.py +++ b/dimos/robot/unitree/go2/dds/extrinsics.py @@ -24,11 +24,12 @@ [yaw -57.9, pitch -8.1, roll -167.3] (agrees to a few degrees). """ +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.std_msgs.msg import Header import numpy as np +from scipy.spatial.transform import Rotation -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.geometry import quaternion_from_matrix # base_link <- lidar rotation (lidar points -> base frame: p_base = EXT_R @ p_lidar + EXT_T) EXT_R = np.array( @@ -46,16 +47,21 @@ CAM_T = np.array([0.30, 0.0, 0.0], dtype=np.float64) CAM_Q = np.array([-0.5, 0.5, -0.5, 0.5], dtype=np.float64) # xyzw -# Same mounts as standard Transform msgs (typed; carry frame ids; have to_rerun). -LIDAR_TO_BASE = Transform( - translation=Vector3(EXT_T), - rotation=Quaternion.from_rotation_matrix(EXT_R), - frame_id="base_link", +# Normalize the measured six-decimal rotation with the same SciPy conversion +# as the previous numeric helper, then publish generated frame-labelled values. +LIDAR_TO_BASE = TransformStamped( + header=Header(frame_id="base_link"), child_frame_id="lidar", + transform=Transform( + translation=Vector3(x=EXT_T[0], y=EXT_T[1], z=EXT_T[2]), + rotation=quaternion_from_matrix(Rotation.from_matrix(EXT_R).as_matrix()), + ), ) -BASE_TO_CAMERA = Transform( - translation=Vector3(CAM_T), - rotation=Quaternion(CAM_Q), - frame_id="base_link", +BASE_TO_CAMERA = TransformStamped( + header=Header(frame_id="base_link"), child_frame_id="camera_optical", + transform=Transform( + translation=Vector3(x=CAM_T[0], y=CAM_T[1], z=CAM_T[2]), + rotation=Quaternion(x=CAM_Q[0], y=CAM_Q[1], z=CAM_Q[2], w=CAM_Q[3]), + ), ) diff --git a/dimos/robot/unitree/go2/dds/ros.py b/dimos/robot/unitree/go2/dds/ros.py index c9174a78c63..4d7c263a2e8 100644 --- a/dimos/robot/unitree/go2/dds/ros.py +++ b/dimos/robot/unitree/go2/dds/ros.py @@ -12,276 +12,46 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Decode the standard-ROS DDS channels into ``dimos.msgs`` types. +"""Generated ROS2 DDS codecs with explicit Unitree firmware adaptations. -The wire layouts are declared as ``__cdr_fields__`` specs and walked by the -generic :mod:`cdr` decoder (no hand-rolled byte cursors). A thin per-type -adapter maps the decoded wire struct into the dimos message — only the parts -that genuinely differ from a field copy (point-buffer reinterpretation, jpeg -decode, pose nesting) live there. +Standard payloads retain all generated fields, exact stamps and arbitrary +point layouts. The device's Imu quaternion uses wxyz and is reordered only at +this Unitree boundary. HeightMap remains a device-specific wire structure. """ -# TODO this file needs to go away, dimos/msgs are structurally the same as -# these messages here so we will write an automatic translator, temporary so -# we can iterate on go2 dds research, see if it's viable at all -# -# TODO pointcloud has timestamps and intensities, we drop those on LCM round trip -# and our pointcloud2 message doesn't support arbitrary fields per point, we need -# to implement those - -from __future__ import annotations - from dataclasses import dataclass +from dimos_generated.foxglove_msgs.msg import CompressedVideo +from dimos_generated.geometry_msgs.msg import Quaternion +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import CompressedImage, Imu, PointCloud2 import numpy as np -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.sensor_msgs.Imu import Imu -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.robot.unitree.go2.dds import cdr -from dimos.robot.unitree.go2.dds.msgs.CompressedVideo import CompressedVideo from dimos.robot.unitree.go2.dds.msgs.HeightMap import HeightMap -# Shared wire specs (Header/Time) reused by the per-message layouts below. -@dataclass -class _Time: - sec: int - nanosec: int - - __cdr_fields__ = [("sec", "i32"), ("nanosec", "u32")] - - -@dataclass -class _Header: - stamp: _Time - frame_id: str - - __cdr_fields__ = [("stamp", _Time), ("frame_id", "string")] - - -def _ts(h: _Header) -> float: - return h.stamp.sec + h.stamp.nanosec * 1e-9 - - -# sensor_msgs/Imu -@dataclass -class _ImuWire: - header: _Header - orientation: np.ndarray # f64[4] xyzw - orientation_covariance: np.ndarray # f64[9] - angular_velocity: np.ndarray # f64[3] - angular_velocity_covariance: np.ndarray - linear_acceleration: np.ndarray # f64[3] - linear_acceleration_covariance: np.ndarray - - __cdr_fields__ = [ - ("header", _Header), - ("orientation", ("array", "f64", 4)), - ("orientation_covariance", ("array", "f64", 9)), - ("angular_velocity", ("array", "f64", 3)), - ("angular_velocity_covariance", ("array", "f64", 9)), - ("linear_acceleration", ("array", "f64", 3)), - ("linear_acceleration_covariance", ("array", "f64", 9)), - ] - - def decode_imu(buf: bytes) -> Imu: - w: _ImuWire = cdr.decode(buf, _ImuWire)[0] - # Unitree fills orientation wxyz even in this sensor_msgs/Imu — reorder to xyzw - # (verified: rotating accel by this lands gravity on +z in leg-odom-stationary windows). - qw, qx, qy, qz = (float(v) for v in w.orientation) - return Imu( - orientation=Quaternion(qx, qy, qz, qw), - angular_velocity=Vector3(w.angular_velocity.tolist()), - linear_acceleration=Vector3(w.linear_acceleration.tolist()), - orientation_covariance=w.orientation_covariance.tolist(), - angular_velocity_covariance=w.angular_velocity_covariance.tolist(), - linear_acceleration_covariance=w.linear_acceleration_covariance.tolist(), - frame_id=w.header.frame_id, - ts=_ts(w.header), - ) - - -# nav_msgs/Odometry -@dataclass -class _PoseWire: - position: np.ndarray # f64[3] - orientation: np.ndarray # f64[4] xyzw - - __cdr_fields__ = [("position", ("array", "f64", 3)), ("orientation", ("array", "f64", 4))] - - -@dataclass -class _PoseWithCov: - pose: _PoseWire - covariance: np.ndarray # f64[36] - - __cdr_fields__ = [("pose", _PoseWire), ("covariance", ("array", "f64", 36))] - - -@dataclass -class _TwistWire: - linear: np.ndarray # f64[3] - angular: np.ndarray # f64[3] - - __cdr_fields__ = [("linear", ("array", "f64", 3)), ("angular", ("array", "f64", 3))] - - -@dataclass -class _TwistWithCov: - twist: _TwistWire - covariance: np.ndarray # f64[36] - - __cdr_fields__ = [("twist", _TwistWire), ("covariance", ("array", "f64", 36))] - - -@dataclass -class _OdomWire: - header: _Header - child_frame_id: str - pose: _PoseWithCov - twist: _TwistWithCov - - __cdr_fields__ = [ - ("header", _Header), - ("child_frame_id", "string"), - ("pose", _PoseWithCov), - ("twist", _TwistWithCov), - ] + message = Imu.decode(buf) + firmware = message.orientation + message.orientation = Quaternion(x=firmware.y, y=firmware.z, z=firmware.w, w=firmware.x) + return message def decode_odometry(buf: bytes) -> Odometry: - w: _OdomWire = cdr.decode(buf, _OdomWire)[0] - pose = Pose() - pose.position = Vector3(w.pose.pose.position.tolist()) - pose.orientation = Quaternion(*w.pose.pose.orientation.tolist()) - twist = Twist() - twist.linear = Vector3(w.twist.twist.linear.tolist()) - twist.angular = Vector3(w.twist.twist.angular.tolist()) - return Odometry( - ts=_ts(w.header), - frame_id=w.header.frame_id, - child_frame_id=w.child_frame_id, - pose=pose, - twist=twist, - ) - - -# sensor_msgs/PointCloud2 -@dataclass -class _PointField: - name: str - offset: int - datatype: int - count: int - - __cdr_fields__ = [ - ("name", "string"), - ("offset", "u32"), - ("datatype", "u8"), - ("count", "u32"), - ] - - -@dataclass -class _Pc2Wire: - header: _Header - height: int - width: int - fields: list[_PointField] - is_bigendian: int - point_step: int - row_step: int - data: np.ndarray # u8[] - is_dense: int - - __cdr_fields__ = [ - ("header", _Header), - ("height", "u32"), - ("width", "u32"), - ("fields", ("seq", _PointField)), - ("is_bigendian", "u8"), - ("point_step", "u32"), - ("row_step", "u32"), - ("data", ("seq", "u8")), - ("is_dense", "u8"), - ] - - -# ROS PointField datatype code -> numpy dtype -_PF_DT = {1: " PointCloud2: - w: _Pc2Wire = cdr.decode(buf, _Pc2Wire)[0] - ts, frame = _ts(w.header), w.header.frame_id - if w.point_step == 0 or w.data.size < w.point_step: - return PointCloud2.from_numpy(np.empty((0, 3), np.float32), frame, ts) - dt = np.dtype( - { - "names": [f.name for f in w.fields], - "formats": [_PF_DT[f.datatype] for f in w.fields], - "offsets": [f.offset for f in w.fields], - "itemsize": w.point_step, - } - ) - arr = w.data.view(dt) - xyz = np.stack([arr["x"], arr["y"], arr["z"]], axis=-1).astype(np.float32) - inten = arr["intensity"].astype(np.float32) if dt.names and "intensity" in dt.names else None - return PointCloud2.from_numpy(xyz, frame, ts, inten) - - -# sensor_msgs/CompressedImage -@dataclass -class _CImgWire: - header: _Header - format: str - data: np.ndarray # u8[] - - __cdr_fields__ = [ - ("header", _Header), - ("format", "string"), - ("data", ("seq", "u8")), - ] + return PointCloud2.decode(buf) -def decode_compressed_image(buf: bytes) -> Image | None: - import cv2 - - w: _CImgWire = cdr.decode(buf, _CImgWire)[0] - bgr = cv2.imdecode(w.data, cv2.IMREAD_COLOR) - if bgr is None: - return None - return Image.from_numpy(bgr, ImageFormat.BGR, w.header.frame_id, _ts(w.header)) - - -# foxglove_msgs/CompressedVideo (field order differs from CompressedImage: -# timestamp, frame_id, data, format). -@dataclass -class _CVideoWire: - stamp: _Time - frame_id: str - data: np.ndarray # u8[] - format: str - - __cdr_fields__ = [ - ("stamp", _Time), - ("frame_id", "string"), - ("data", ("seq", "u8")), - ("format", "string"), - ] +def decode_compressed_image(buf: bytes) -> CompressedImage: + return CompressedImage.decode(buf) def decode_compressed_video(buf: bytes) -> CompressedVideo: - """Decode the CDR envelope only — the encoded packet is left for H264Decoder.""" - w: _CVideoWire = cdr.decode(buf, _CVideoWire)[0] - return CompressedVideo(data=w.data, format=w.format, frame_id=w.frame_id) + return CompressedVideo.decode(buf) # unitree_go/HeightMap (rt/utlidar/height_map_array) diff --git a/dimos/robot/unitree/go2/dds/test_generated_codecs.py b/dimos/robot/unitree/go2/dds/test_generated_codecs.py new file mode 100644 index 00000000000..44930920549 --- /dev/null +++ b/dimos/robot/unitree/go2/dds/test_generated_codecs.py @@ -0,0 +1,111 @@ +# 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. + +import struct + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.foxglove_msgs.msg import CompressedVideo +from dimos_generated.geometry_msgs.msg import Quaternion +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import CompressedImage, Imu, PointCloud2, PointField +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest +from rosbags.typesys import Stores, get_typestore + +from dimos.memory.type.observation import Observation +from dimos.msgs.image import image_view +from dimos.msgs.pointcloud import pointcloud_xyz +from dimos.robot.unitree.go2.dds import ros +from dimos.robot.unitree.go2.dds.codec import GO2_CODECS +from dimos.robot.unitree.go2.dds.video import H264Decoder + + +@pytest.mark.parametrize("little_endian", [True, False]) +def test_odometry_preserves_independent_cdr_covariances_and_exact_stamp(little_endian): + source = Odometry( + header=Header(stamp=Time(sec=42, nanosec=123456789), frame_id="odom"), + child_frame_id="base_link", + ) + source.pose.pose.position.x = 1.25 + source.pose.covariance = list(range(36)) + source.twist.covariance = list(range(36, 72)) + reference = get_typestore(Stores.ROS2_JAZZY) + independent = reference.deserialize_cdr(source.encode(), source.msg_name) + payload = bytes( + reference.serialize_cdr(independent, source.msg_name, little_endian=little_endian) + ) + result = ros.decode_odometry(payload) + assert result.header.stamp.nanosec == 123456789 + assert result.header.frame_id == "odom" and result.child_frame_id == "base_link" + assert result.pose.pose.position.x == 1.25 + assert list(result.pose.covariance) == list(range(36)) + assert list(result.twist.covariance) == list(range(36, 72)) + + +def test_pointcloud_preserves_big_endian_fields_and_row_padding(): + source = PointCloud2( + header=Header(stamp=Time(sec=42, nanosec=123456789), frame_id="sensor"), + height=1, + width=1, + point_step=26, + row_step=34, + is_bigendian=True, + fields=[ + PointField(name=n, offset=i * 8, datatype=8, count=1) + for i, n in enumerate(["x", "y", "z"]) + ] + + [PointField(name="intensity", offset=24, datatype=4, count=1)], + data=struct.pack(">dddH", 1.25, -2.5, 3.0, 1234) + b"padding!", + ) + result = ros.decode_pointcloud2(source.encode()) + assert result.encode() == source.encode() + assert result.header.stamp.nanosec == 123456789 + np.testing.assert_array_equal(pointcloud_xyz(result), [[1.25, -2.5, 3.0]]) + + +def test_unitree_imu_reorders_firmware_quaternion_without_losing_metadata(): + source = Imu( + header=Header(stamp=Time(nanosec=123456789), frame_id="imu"), + orientation=Quaternion(x=0.5, y=0.1, z=0.2, w=0.3), + ) + source.orientation_covariance = list(range(9)) + result = ros.decode_imu(source.encode()) + q = result.orientation + assert (q.x, q.y, q.z, q.w) == (0.1, 0.2, 0.3, 0.5) + assert list(result.orientation_covariance) == list(range(9)) + assert result.header.stamp.nanosec == 123456789 + + +def test_compressed_image_stays_a_compressed_generated_message(): + source = CompressedImage( + header=Header(frame_id="camera"), format="jpeg", data=b"encoded-packet" + ) + result = GO2_CODECS["rt/frontvideo"].decode(source.encode()) + assert isinstance(result, CompressedImage) + assert result.encode() == source.encode() + + +def test_h264_emission_preserves_packet_stamp_and_pixel_layout(mocker): + stamp = Time(sec=42, nanosec=123456789) + packet = CompressedVideo(timestamp=stamp, frame_id="camera", format="h264", data=b"packet") + source = Observation(id=7, ts=99.0, data_type=CompressedVideo, _data=packet) + frame = mocker.Mock() + pixels = np.arange(18, dtype=np.uint8).reshape(2, 3, 3) + frame.to_ndarray.return_value = pixels + result = H264Decoder._emit(frame, source) + assert result.id == 7 and result.ts == 99.0 + assert result.data.encoding == "bgr8" and result.data.header.frame_id == "camera" + assert result.data.header.stamp.sec == 42 and result.data.header.stamp.nanosec == 123456789 + np.testing.assert_array_equal(image_view(result.data), pixels) diff --git a/dimos/robot/unitree/go2/dds/test_store.py b/dimos/robot/unitree/go2/dds/test_store.py index 296ca56e0ae..70831c198bc 100644 --- a/dimos/robot/unitree/go2/dds/test_store.py +++ b/dimos/robot/unitree/go2/dds/test_store.py @@ -21,13 +21,13 @@ import importlib.util +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import CompressedImage, Imu, PointCloud2 import numpy as np import pytest -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.Imu import Imu -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.image import image_from_compressed, image_view +from dimos.msgs.pointcloud import pointcloud_xyz from dimos.robot.unitree.go2.dds.msgs.ControlEvent import ControlEvent from dimos.robot.unitree.go2.dds.msgs.LowState import LowState from dimos.robot.unitree.go2.dds.msgs.SportModeState import SportModeState @@ -63,9 +63,9 @@ def test_lists_streams(store: Go2McapStore) -> None: def test_lidar(store: Go2McapStore) -> None: pc = store.streams.lidar.first().data - xyz = pc.points_f32() + xyz = pointcloud_xyz(pc) print( - f"\nPointCloud2: {xyz.shape[0]} pts frame={pc.frame_id!r} " + f"\nPointCloud2: {xyz.shape[0]} pts frame={pc.header.frame_id!r} " f"range=[{np.linalg.norm(xyz[:, :3], axis=1).min():.2f}, " f"{np.linalg.norm(xyz[:, :3], axis=1).max():.2f}] m" ) @@ -79,7 +79,7 @@ def test_imu(store: Go2McapStore) -> None: print( f"\nImu: |q|={np.linalg.norm([q.x, q.y, q.z, q.w]):.4f} " f"acc=({imu.linear_acceleration.x:.2f}, {imu.linear_acceleration.y:.2f}, " - f"{imu.linear_acceleration.z:.2f}) frame={imu.frame_id!r}" + f"{imu.linear_acceleration.z:.2f}) frame={imu.header.frame_id!r}" ) assert isinstance(imu, Imu) assert abs(imu.linear_acceleration.z) == pytest.approx(9.8, abs=0.5) # gravity @@ -90,7 +90,7 @@ def test_odom(store: Go2McapStore) -> None: p = odom.pose.pose.position print( f"\nOdometry: pos=({p.x:.2f}, {p.y:.2f}, {p.z:.2f}) " - f"{odom.frame_id!r} -> {odom.child_frame_id!r}" + f"{odom.header.frame_id!r} -> {odom.child_frame_id!r}" ) assert isinstance(odom, Odometry) assert odom.child_frame_id == "base_link" @@ -98,9 +98,9 @@ def test_odom(store: Go2McapStore) -> None: def test_color_image(store: Go2McapStore) -> None: img = store.streams.color_image.first().data - arr = img.as_numpy() - print(f"\nImage: {arr.shape} frame={img.frame_id!r}") - assert isinstance(img, Image) + arr = image_view(image_from_compressed(img)) + print(f"\nImage: {arr.shape} frame={img.header.frame_id!r}") + assert isinstance(img, CompressedImage) assert arr.ndim == 3 diff --git a/dimos/robot/unitree/go2/dds/video.py b/dimos/robot/unitree/go2/dds/video.py index 6b76617367c..0162a6c8eb7 100644 --- a/dimos/robot/unitree/go2/dds/video.py +++ b/dimos/robot/unitree/go2/dds/video.py @@ -36,9 +36,12 @@ from collections.abc import Iterator from typing import TYPE_CHECKING, Any +from dimos_generated.foxglove_msgs.msg import CompressedVideo +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header + from dimos.memory.transform import Transformer -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.robot.unitree.go2.dds.msgs.CompressedVideo import CompressedVideo +from dimos.msgs.image import image_from_array if TYPE_CHECKING: from dimos.memory.type.observation import Observation @@ -60,7 +63,7 @@ def __call__( continue last = obs try: - frames = decoder.decode(av.packet.Packet(packet.data.tobytes())) + frames = decoder.decode(av.packet.Packet(bytes(packet.data))) except av.error.FFmpegError: continue # P-frame with no reference yet (e.g. seeked past a keyframe) for frame in frames: @@ -75,4 +78,10 @@ def __call__( @staticmethod def _emit(frame: Any, obs: Observation[CompressedVideo]) -> Observation[Image]: bgr = frame.to_ndarray(format="bgr24") - return obs.derive(data=Image.from_numpy(bgr, ImageFormat.BGR, obs.data.frame_id, obs.ts)) + return obs.derive( + data=image_from_array( + bgr, + encoding="bgr8", + header=Header(stamp=obs.data.timestamp, frame_id=obs.data.frame_id), + ) + ) diff --git a/dimos/robot/unitree/go2/fleet_connection.py b/dimos/robot/unitree/go2/fleet_connection.py index 9f580b17c70..3e6649f0295 100644 --- a/dimos/robot/unitree/go2/fleet_connection.py +++ b/dimos/robot/unitree/go2/fleet_connection.py @@ -37,7 +37,7 @@ from typing_extensions import Self if TYPE_CHECKING: - from dimos.msgs.geometry_msgs.Twist import Twist + from dimos_generated.geometry_msgs.msg import Twist logger = setup_logger() diff --git a/dimos/robot/unitree/go2/go2_mid360_recorder.py b/dimos/robot/unitree/go2/go2_mid360_recorder.py index 5c2ab2f243d..bd1a3486749 100644 --- a/dimos/robot/unitree/go2/go2_mid360_recorder.py +++ b/dimos/robot/unitree/go2/go2_mid360_recorder.py @@ -23,12 +23,11 @@ from __future__ import annotations +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 + from dimos.core.stream import In from dimos.hardware.sensors.lidar.pointlio.recorder import PointlioRecorder -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 class Go2Mid360Recorder(PointlioRecorder): diff --git a/dimos/robot/unitree/go2/go2_mid360_static_transforms.py b/dimos/robot/unitree/go2/go2_mid360_static_transforms.py index 1158bbbc762..da411efa371 100644 --- a/dimos/robot/unitree/go2/go2_mid360_static_transforms.py +++ b/dimos/robot/unitree/go2/go2_mid360_static_transforms.py @@ -33,7 +33,9 @@ import math -from dimos.msgs.geometry_msgs.Transform import Transform +from dimos_generated.geometry_msgs.msg import TransformStamped + +from dimos.msgs.geometry import inverse_transform from dimos.protocol.tf.static_tf_publisher import ( FrameSpec, StaticTfPublisher, @@ -65,14 +67,18 @@ ] -def mount_transforms() -> list[Transform]: +def mount_transforms() -> list[TransformStamped]: """The mount tree as published: rooted at mid360_link.""" edges = {t.child_frame_id: t for t in frames_to_edge_transforms(FRAMES)} - return [-edges["mid360_link"], -edges["front_camera"], edges["camera_optical"]] + return [ + inverse_transform(edges["mid360_link"]), + inverse_transform(edges["front_camera"]), + edges["camera_optical"], + ] class Go2Mid360StaticTf(StaticTfPublisher): """Publishes the Go2/Mid-360 mount tree onto tf on a fixed interval.""" - def transforms(self) -> list[Transform]: + def transforms(self) -> list[TransformStamped]: return mount_transforms() diff --git a/dimos/robot/unitree/go2/test_go2_mid360_static_transforms.py b/dimos/robot/unitree/go2/test_go2_mid360_static_transforms.py index 7cf999cd0f0..1330ecdd685 100644 --- a/dimos/robot/unitree/go2/test_go2_mid360_static_transforms.py +++ b/dimos/robot/unitree/go2/test_go2_mid360_static_transforms.py @@ -21,6 +21,9 @@ import math +from scipy.spatial.transform import Rotation + +from dimos.msgs.geometry import inverse_transform from dimos.protocol.tf.tf import MultiTBuffer from dimos.robot.unitree.go2.go2_mid360_static_transforms import ( MID360_PITCH_DOWN, @@ -42,16 +45,17 @@ def test_mount_height_survives_the_inversion() -> None: """The lidar sits MOUNT_Z above base_link, the offset every ground projection uses.""" leg = _buffer().get("mid360_link", "base_link") assert leg is not None - base_to_sensor = -leg - assert abs(base_to_sensor.translation.z - MOUNT_Z) < 1e-6 - assert abs(base_to_sensor.translation.x - MOUNT_X) < 1e-6 + base_to_sensor = inverse_transform(leg) + assert abs(base_to_sensor.transform.translation.z - MOUNT_Z) < 1e-6 + assert abs(base_to_sensor.transform.translation.x - MOUNT_X) < 1e-6 def test_mount_pitch_survives_the_inversion() -> None: """A sign flip here steers the follower off-heading rather than failing loudly.""" leg = _buffer().get("mid360_link", "base_link") assert leg is not None - pitch = (-leg).rotation.euler.y + q = inverse_transform(leg).transform.rotation + pitch = Rotation.from_quat([q.x, q.y, q.z, q.w]).as_euler("xyz")[1] assert abs(pitch - MID360_PITCH_DOWN) < 1e-6 assert abs(math.degrees(pitch) - 60.0) < 1e-6 @@ -60,5 +64,5 @@ def test_camera_optical_hangs_off_base_link() -> None: """The tree is rooted at mid360_link, so the camera edge is reachable by composition.""" optical = _buffer().get("base_link", "camera_optical") assert optical is not None - assert abs(optical.translation.x - 0.32715) < 1e-6 - assert abs(optical.translation.z - 0.04297) < 1e-6 + assert abs(optical.transform.translation.x - 0.32715) < 1e-6 + assert abs(optical.transform.translation.z - 0.04297) < 1e-6 diff --git a/dimos/robot/unitree/mujoco_connection.py b/dimos/robot/unitree/mujoco_connection.py index 249252afdf0..83e4c83c81d 100644 --- a/dimos/robot/unitree/mujoco_connection.py +++ b/dimos/robot/unitree/mujoco_connection.py @@ -31,6 +31,9 @@ from typing import Any, TypeVar import weakref +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Twist +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header import numpy as np from numpy.typing import NDArray from reactivex import Observable, empty @@ -39,13 +42,9 @@ from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.global_config import GlobalConfig -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.robot.unitree.type.odometry import Odometry +from dimos.msgs.camera_info import camera_info_from_fov +from dimos.msgs.image import image_from_array +from dimos.msgs.time import header_now, time_from_seconds from dimos.simulation.mujoco.constants import ( LAUNCHER_PATH, LIDAR_FPS, @@ -106,12 +105,12 @@ def cleanup_on_exit( atexit.register(cleanup_on_exit) - camera_info_static: CameraInfo = CameraInfo.from_fov( - fov_deg=VIDEO_CAMERA_FOV, + camera_info_static: CameraInfo = camera_info_from_fov( + fov_degrees=VIDEO_CAMERA_FOV, width=VIDEO_WIDTH, height=VIDEO_HEIGHT, axis="vertical", - frame_id="camera_optical", + header=Header(frame_id="camera_optical"), ) def _launch(self) -> None: @@ -326,7 +325,7 @@ def get_video_frame(self) -> NDArray[Any] | None: return None - def get_odom_message(self) -> Odometry | None: + def get_odom_message(self) -> PoseStamped | None: if self.shm_data is None: return None @@ -336,13 +335,11 @@ def get_odom_message(self) -> Odometry | None: pos, quat_wxyz, timestamp = odom_data # Convert quaternion from (w,x,y,z) to (x,y,z,w) for ROS/Dimos - orientation = Quaternion(quat_wxyz[1], quat_wxyz[2], quat_wxyz[3], quat_wxyz[0]) + orientation = Quaternion(x=quat_wxyz[1], y=quat_wxyz[2], z=quat_wxyz[3], w=quat_wxyz[0]) - return Odometry( - position=Vector3(pos[0], pos[1], pos[2]), - orientation=orientation, - ts=timestamp, - frame_id="world", + return PoseStamped( + header=Header(stamp=time_from_seconds(timestamp), frame_id="world"), + pose=Pose(position=Point(x=pos[0], y=pos[1], z=pos[2]), orientation=orientation), ) return None @@ -401,15 +398,19 @@ def lidar_stream(self) -> Observable[PointCloud2]: return self._create_stream(self.get_lidar_message, LIDAR_FPS, "Lidar") @functools.cache - def odom_stream(self) -> Observable[Odometry]: + def odom_stream(self) -> Observable[PoseStamped]: return self._create_stream(self.get_odom_message, ODOM_FREQUENCY, "Odom") @functools.cache def video_stream(self) -> Observable[Image]: def get_video_as_image() -> Image | None: frame = self.get_video_frame() - # MuJoCo renderer returns RGB uint8 frames; Image.from_numpy defaults to BGR. - return Image.from_numpy(frame, format=ImageFormat.RGB) if frame is not None else None + # MuJoCo renderer returns RGB uint8 pixels. + return ( + image_from_array(frame, encoding="rgb8", header=header_now("camera_optical")) + if frame is not None + else None + ) return self._create_stream(get_video_as_image, VIDEO_FPS, "Video") diff --git a/dimos/robot/unitree/test_connection.py b/dimos/robot/unitree/test_connection.py index 8180f2d70d9..041ac8b5621 100644 --- a/dimos/robot/unitree/test_connection.py +++ b/dimos/robot/unitree/test_connection.py @@ -19,18 +19,25 @@ """ import json +import threading from typing import Any from unittest.mock import ANY, AsyncMock, MagicMock, call +from dimos_generated.geometry_msgs.msg import PoseStamped, TransformStamped, Twist, Vector3 +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 +import numpy as np import pytest +import reactivex as rx +from reactivex.scheduler import ThreadPoolScheduler from unitree_webrtc_connect.constants import DATA_CHANNEL_TYPE, RTC_TOPIC, SPORT_CMD from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.global_config import GlobalConfig -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.image import image_view +from dimos.msgs.pointcloud import pointcloud_xyz from dimos.robot.unitree import connection as conn_mod -from dimos.robot.unitree.connection import UnitreeWebRTCConnection +from dimos.robot.unitree.connection import SerializableVideoFrame, UnitreeWebRTCConnection +from dimos.robot.unitree.type.odometry import raw_odometry_msg_sample def _stub_driver(connect_exc: Exception | None = None) -> MagicMock: @@ -113,8 +120,8 @@ def test_move_api_toggle_sends_selected_wire_command( driver = _stub_driver() monkeypatch.setattr(conn_mod, "LegionConnection", MagicMock(return_value=driver)) twist = Twist( - linear=Vector3(1.5, -0.4, 0.0), - angular=Vector3(0.0, 0.0, 0.8), + linear=Vector3(x=1.5, y=-0.4), + angular=Vector3(z=0.8), ) connection = UnitreeWebRTCConnection(ip="10.0.0.99", **connection_options) @@ -165,3 +172,102 @@ def test_global_config_reads_unitree_aes_128_key_env(monkeypatch: pytest.MonkeyP """The key enters via GlobalConfig, read from the UNITREE_AES_128_KEY env var.""" monkeypatch.setenv("UNITREE_AES_128_KEY", "ee" * 16) assert GlobalConfig().unitree_aes_128_key == "ee" * 16 + + +@pytest.fixture +def sensor_scheduler(mocker): + scheduler = ThreadPoolScheduler(max_workers=2) + mocker.patch("dimos.utils.reactive.get_scheduler", return_value=scheduler) + try: + yield scheduler + finally: + scheduler.executor.shutdown(wait=True) + + +@pytest.mark.parametrize("kind", ["lidar", "video"]) +def test_sensor_streams_emit_generated_cdr_with_exact_arrival_time( + built_connection, mocker, sensor_scheduler, kind +): + connection, _driver = built_connection + mocker.patch.object(conn_mod.time, "time_ns", return_value=1700000000123456789) + values = np.array([[1.25, 2.5, 3.75]], dtype=np.float32) + pixels = np.arange(18, dtype=np.uint8).reshape(2, 3, 3) + if kind == "lidar": + raw = {"data": {"stamp": 1.0, "data": {"points": values}}} + mocker.patch.object(connection, "raw_lidar_stream", return_value=rx.just(raw)) + stream = connection.lidar_stream() + else: + frame = SerializableVideoFrame(data=pixels) + mocker.patch.object(connection, "raw_video_stream", return_value=rx.just(frame)) + stream = connection.video_stream() + received = [] + errors = [] + ready = threading.Event() + + def receive(message): + received.append(message) + ready.set() + + def fail(error): + errors.append(error) + ready.set() + + subscription = stream.subscribe(receive, fail) + try: + assert ready.wait(5), "Sensor conversion did not publish" + assert not errors + assert len(received) == 1 + message = received[0] + assert (message.header.stamp.sec, message.header.stamp.nanosec) == (1700000000, 123456789) + if kind == "lidar": + decoded = PointCloud2.decode(message.encode()) + np.testing.assert_array_equal(pointcloud_xyz(decoded), values) + assert decoded.header.frame_id == "world" + else: + decoded = Image.decode(message.encode()) + np.testing.assert_array_equal(image_view(decoded), pixels) + assert decoded.header.frame_id == "camera_optical" + finally: + subscription.dispose() + + +@pytest.mark.parametrize("as_tf", [False, True]) +def test_odometry_stream_preserves_pose_and_uses_exact_arrival_stamp( + built_connection, mocker, sensor_scheduler, as_tf +): + connection, _driver = built_connection + mocker.patch.object(conn_mod.time, "time_ns", return_value=1700000000123456789) + mocker.patch.object( + connection, "raw_odom_stream", return_value=rx.just(raw_odometry_msg_sample) + ) + stream = connection.tf_stream() if as_tf else connection.odom_stream() + received = [] + errors = [] + ready = threading.Event() + + def receive(message): + received.append(message) + ready.set() + + def fail(error): + errors.append(error) + ready.set() + + subscription = stream.subscribe(receive, fail) + try: + assert ready.wait(5), "Odometry conversion did not publish" + assert not errors + assert len(received) == 1 + cls = TransformStamped if as_tf else PoseStamped + message = cls.decode(received[0].encode()) + assert message.header.frame_id == "world" + assert (message.header.stamp.sec, message.header.stamp.nanosec) == (1700000000, 123456789) + if as_tf: + assert message.child_frame_id == "base_link" + position, rotation = message.transform.translation, message.transform.rotation + else: + position, rotation = message.pose.position, message.pose.orientation + assert (position.x, position.y, position.z) == (5.961965, -2.916958, 0.319509) + assert rotation.w == -0.242112 + finally: + subscription.dispose() diff --git a/dimos/robot/unitree/test_generated_commands.py b/dimos/robot/unitree/test_generated_commands.py new file mode 100644 index 00000000000..12ee1901377 --- /dev/null +++ b/dimos/robot/unitree/test_generated_commands.py @@ -0,0 +1,87 @@ +# 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. + +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from unittest.mock import Mock + +from dimos_generated.geometry_msgs.msg import Twist +from dimos_generated.std_msgs.msg import Float32, Int8 +import pytest + +from dimos.control.benchmarking.gate import GATE_ADVANCE, GATE_QUIT, GATE_SKIP +from dimos.robot.unitree import keyboard_teleop as keyboard +from dimos.robot.unitree.g1.skill_container import UnitreeG1SkillContainer + + +def test_keyboard_generated_operator_gates_and_slider_without_window(monkeypatch): + module = keyboard.KeyboardTeleop(disable_movement=True) + module._keys_held = set() + module._thread = Mock() + received = [] + slider = [] + motion = [] + monkeypatch.setattr(module.operator_command, "publish", received.append) + monkeypatch.setattr(module.e_max, "publish", slider.append) + monkeypatch.setattr(module.cmd_vel, "publish", motion.append) + monkeypatch.setattr(module, "_update_display", lambda message: None) + monkeypatch.setattr(keyboard.pygame, "init", Mock()) + monkeypatch.setattr(keyboard.pygame, "quit", Mock()) + monkeypatch.setattr(keyboard.pygame.display, "set_mode", lambda *args: Mock()) + monkeypatch.setattr(keyboard.pygame.display, "set_caption", Mock()) + monkeypatch.setattr(keyboard.pygame.font, "Font", lambda *args: Mock()) + monkeypatch.setattr( + keyboard.pygame.time, + "Clock", + lambda: SimpleNamespace(tick=lambda hz: module._stop_event.set()), + ) + keys = [ + keyboard.pygame.K_RETURN, + keyboard.pygame.K_k, + keyboard.pygame.K_BACKSPACE, + keyboard.pygame.K_7, + ] + events = [SimpleNamespace(type=keyboard.pygame.KEYDOWN, key=key) for key in keys] + monkeypatch.setattr(keyboard.pygame.event, "get", lambda: events) + try: + module._pygame_loop() + assert [Int8.decode(value.encode()).data for value in received] == [ + GATE_ADVANCE, + GATE_SKIP, + GATE_QUIT, + ] + assert Float32.decode(slider[0].encode()).data == pytest.approx(0.7) + assert not motion # disable_movement keeps this an operator-only interface + finally: + module.stop() + assert len(motion) == 1 and Twist.decode(motion[0].encode()) == Twist() + + +def test_g1_skill_constructs_generated_twist_at_captured_rpc_boundary(): + module = UnitreeG1SkillContainer() + gateway = Mock() + module._connection = gateway + try: + module.move(x=0.25, y=-0.125, yaw=0.5, duration=2) + gateway.move.assert_called_once() + args, kwargs = gateway.move.call_args + assert type(args[0]) is Twist + value = Twist.decode(args[0].encode()) + assert value.linear.x == 0.25 and value.linear.y == -0.125 + assert value.angular.z == 0.5 and value.linear.z == 0 + assert kwargs == {"duration": 2} + finally: + module.stop() diff --git a/dimos/robot/unitree/test_simulation_tf.py b/dimos/robot/unitree/test_simulation_tf.py new file mode 100644 index 00000000000..9ffb8d97eff --- /dev/null +++ b/dimos/robot/unitree/test_simulation_tf.py @@ -0,0 +1,87 @@ +# 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. + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.sensor_msgs.msg import CameraInfo +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage +import pytest + +from dimos.core.global_config import GlobalConfig +from dimos.robot.unitree.dimsim_connection import _odom_to_tf +from dimos.robot.unitree.g1.mujoco_sim import G1SimConnection +from dimos.robot.unitree.mujoco_connection import MujocoConnection + + +@pytest.fixture +def source_pose(): + return PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=1700000000, nanosec=123456789)), + pose=Pose(position=Point(x=1.25, y=2.5, z=0.7)), + ) + + +def test_dimsim_tf_preserves_source_pose_and_stamp(source_pose): + message = TFMessage.decode(TFMessage(transforms=_odom_to_tf(source_pose)).encode()) + assert [(edge.header.frame_id, edge.child_frame_id) for edge in message.transforms] == [ + ("world", "base_link"), + ("base_link", "camera_link"), + ("camera_link", "camera_optical"), + ("base_link", "lidar_link"), + ] + assert all(edge.header.stamp == source_pose.header.stamp for edge in message.transforms) + assert message.transforms[0].transform.translation.x == 1.25 + assert message.transforms[1].transform.translation.x == 0.3 + + +@pytest.fixture +def g1_connection(mocker): + connection = G1SimConnection(g=GlobalConfig(robot_ip="127.0.0.1")) + mocker.patch.object(connection, "connection", mocker.MagicMock(spec=MujocoConnection)) + mocker.patch.object(connection, "tf", mocker.MagicMock()) + mocker.patch.object(connection, "odom", mocker.MagicMock()) + mocker.patch.object(connection, "camera_info", mocker.MagicMock()) + try: + yield connection + finally: + connection.stop() + + +def test_g1_simulator_tf_uses_pose_time_for_all_mounts(g1_connection, source_pose): + g1_connection._publish_tf(source_pose) + message = TFMessage.decode(g1_connection.tf.publish.call_args.args[0].encode()) + assert [(edge.header.frame_id, edge.child_frame_id) for edge in message.transforms] == [ + ("world", "base_link"), + ("base_link", "camera_link"), + ("camera_link", "camera_optical"), + ("map", "world"), + ] + assert all(edge.header.stamp == source_pose.header.stamp for edge in message.transforms) + assert message.transforms[1].transform.translation.z == 0.6 + assert g1_connection.odom.publish.call_args.args[0] == source_pose + + +def test_g1_camera_info_stamps_each_publication_without_mutating_template(g1_connection, mocker): + template = CameraInfo(header=Header(frame_id="camera_optical"), width=640, height=288) + g1_connection.connection.camera_info_static = template + mocker.patch("dimos.robot.unitree.g1.mujoco_sim.time.time_ns", return_value=1700000000123456789) + mocker.patch.object( + g1_connection._stop_event, "wait", side_effect=lambda _: g1_connection._stop_event.set() + ) + g1_connection._publish_camera_info_loop() + message = CameraInfo.decode(g1_connection.camera_info.publish.call_args.args[0].encode()) + assert message.header.stamp == Time(sec=1700000000, nanosec=123456789) + assert message.header.frame_id == "camera_optical" + assert template.header.stamp == Time() diff --git a/dimos/robot/unitree/testing/tool_actors.py b/dimos/robot/unitree/testing/tool_actors.py index cd0cdc05dbf..942ceb28c2b 100644 --- a/dimos/robot/unitree/testing/tool_actors.py +++ b/dimos/robot/unitree/testing/tool_actors.py @@ -14,13 +14,13 @@ import asyncio from collections.abc import Callable +from dimos_generated.sensor_msgs.msg import PointCloud2 import pytest from dimos.core.coordination.module_coordinator import ModuleCoordinator from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.transport import LCMTransport -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.robot.unitree.type.map import Map as Mapper diff --git a/dimos/robot/unitree/testing/tool_tooling.py b/dimos/robot/unitree/testing/tool_tooling.py index fe4d76a7428..db2e48e6acb 100644 --- a/dimos/robot/unitree/testing/tool_tooling.py +++ b/dimos/robot/unitree/testing/tool_tooling.py @@ -15,14 +15,14 @@ import time from dimos.robot.unitree.type.lidar import pointcloud2_from_webrtc_lidar -from dimos.robot.unitree.type.odometry import Odometry +from dimos.robot.unitree.type.odometry import pose_from_webrtc_odometry from dimos.utils.reactive import backpressure from dimos.utils.testing.legacy_pickle import LegacyPickleStore def test_replay_all() -> None: lidar_store = LegacyPickleStore("unitree/lidar", autocast=pointcloud2_from_webrtc_lidar) - odom_store = LegacyPickleStore("unitree/odom", autocast=Odometry.from_msg) + odom_store = LegacyPickleStore("unitree/odom", autocast=pose_from_webrtc_odometry) video_store = LegacyPickleStore("unitree/video") backpressure(odom_store.stream()).subscribe(print) diff --git a/dimos/robot/unitree/type/lidar.py b/dimos/robot/unitree/type/lidar.py index c8e01962e2d..90df6536bb9 100644 --- a/dimos/robot/unitree/type/lidar.py +++ b/dimos/robot/unitree/type/lidar.py @@ -16,21 +16,21 @@ from collections.abc import Callable import time -from typing import TypedDict, TypeVar +from typing import Protocol, TypedDict, TypeVar +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header import numpy as np from reactivex import operators as ops from reactivex.observable import Observable -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.types.timestamped import Timestamped +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.utils.logging_config import setup_logger logger = setup_logger() -# Backwards compatibility alias for pickled data -LidarMessage = PointCloud2 - class RawLidarPoints(TypedDict): points: np.ndarray # Shape (N, 3) array of 3D points [x, y, z] @@ -56,34 +56,28 @@ class RawLidarMsg(TypedDict): data: RawLidarData -def pointcloud2_from_webrtc_lidar(raw_message: RawLidarMsg, ts: float | None = None) -> PointCloud2: - """Convert a raw Unitree WebRTC lidar message to PointCloud2. - - Args: - raw_message: Raw lidar message from Unitree WebRTC API - ts: Optional timestamp override. If None, uses the sensor stamp from - ``raw_message["data"]["stamp"]``. +def pointcloud2_from_webrtc_lidar( + raw_message: RawLidarMsg, *, stamp: Time | None = None +) -> PointCloud2: + """Copy decoded WebRTC XYZ points into a generated cloud in the world frame. - The sensor stamp is authoritative when valid but Unitree's publisher - occasionally re-emits a stale value on fresh scans — see - :func:`repair_stale_ts` for the downstream repair. + The raw sensor's floating-second stamp is used unless the connection supplies + an explicit arrival stamp. No Open3D object or legacy message is constructed. """ - import open3d as o3d # type: ignore[import-untyped] - data = raw_message["data"] - points = data["data"]["points"] + return pointcloud_from_xyz( + data["data"]["points"], + header=Header( + frame_id="world", stamp=stamp if stamp is not None else time_from_seconds(data["stamp"]) + ), + ) - pointcloud = o3d.geometry.PointCloud() - pointcloud.points = o3d.utility.Vector3dVector(points) - return PointCloud2( - pointcloud=pointcloud, - ts=ts if ts is not None else data["stamp"], - frame_id="world", - ) +class StampedMessage(Protocol): + header: Header -T = TypeVar("T", bound=Timestamped) +T = TypeVar("T", bound=StampedMessage) def repair_stale_ts( @@ -109,20 +103,20 @@ def _repair(item: T) -> T: nonlocal prev_good, prev_raw, n_seen, calibrated, use_system_time if use_system_time: - item.ts = now() + item.header.stamp = time_from_seconds(now()) return item if not calibrated: - if prev_raw is not None and item.ts != prev_raw: + if prev_raw is not None and to_seconds(item.header.stamp) != prev_raw: calibrated = True # lidar stamps advancing — using lidar time", - prev_raw = item.ts + prev_raw = to_seconds(item.header.stamp) n_seen += 1 - if prev_good is not None and item.ts <= prev_good: - item.ts = prev_good + default_period + if prev_good is not None and to_seconds(item.header.stamp) <= prev_good: + item.header.stamp = time_from_seconds(prev_good + default_period) - prev_good = item.ts + prev_good = to_seconds(item.header.stamp) if not calibrated and n_seen >= calibration_frames: calibrated = True diff --git a/dimos/robot/unitree/type/map.py b/dimos/robot/unitree/type/map.py index 473a6e1df21..5d8e515a365 100644 --- a/dimos/robot/unitree/type/map.py +++ b/dimos/robot/unitree/type/map.py @@ -14,10 +14,15 @@ from __future__ import annotations +from copy import deepcopy from pathlib import Path import time from typing import TYPE_CHECKING, Any +from dimos_generated.nav_msgs.msg import OccupancyGrid +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header +import numpy as np from reactivex import interval from reactivex.disposable import Disposable @@ -27,8 +32,9 @@ from dimos.mapping.pointclouds.accumulators.general import GeneralPointCloudAccumulator from dimos.mapping.pointclouds.accumulators.protocol import PointCloudAccumulator from dimos.mapping.pointclouds.occupancy import general_occupancy -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.occupancy import occupancy_from_file +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_from_xyz_rgb, pointcloud_to_open3d +from dimos.msgs.time import time_from_seconds if TYPE_CHECKING: import open3d as o3d # type: ignore[import-untyped] @@ -53,6 +59,7 @@ class Map(Module): def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) + self._latest_header = Header(frame_id="world", stamp=time_from_seconds(time.time())) self.voxel_size = self.config.voxel_size self.cost_resolution = self.config.cost_resolution self.global_publish_interval = self.config.global_publish_interval @@ -78,15 +85,22 @@ def stop(self) -> None: super().stop() def to_PointCloud2(self) -> PointCloud2: - return PointCloud2( - pointcloud=self._point_cloud_accumulator.get_point_cloud(), - ts=time.time(), - ) + cloud = self._point_cloud_accumulator.get_point_cloud() + points = np.asarray(cloud.points) + if cloud.has_colors(): + return pointcloud_from_xyz_rgb( + points, + (np.asarray(cloud.colors) * 255).astype(np.uint8), + header=deepcopy(self._latest_header), + ) + return pointcloud_from_xyz(points, header=deepcopy(self._latest_header)) # TODO: Why is this RPC? @rpc def add_frame(self, frame: PointCloud2) -> None: - self._point_cloud_accumulator.add(frame.pointcloud) + self._point_cloud_accumulator.add(pointcloud_to_open3d(frame)) + self._latest_header = deepcopy(frame.header) + self._latest_header.frame_id = "world" @property def o3d_geometry(self) -> o3d.geometry.PointCloud: @@ -106,7 +120,9 @@ def _publish(self, _: Any) -> None: if self.config.g.mujoco_global_costmap_from_occupancy: if self._preloaded_occupancy is None: path = Path(self.config.g.mujoco_global_costmap_from_occupancy) - self._preloaded_occupancy = OccupancyGrid.from_path(path) + self._preloaded_occupancy = occupancy_from_file( + path, header=deepcopy(self._latest_header) + ) occupancygrid = self._preloaded_occupancy self.global_costmap.publish(occupancygrid) diff --git a/dimos/robot/unitree/type/odometry.py b/dimos/robot/unitree/type/odometry.py index fabf800b6c1..8240d414b89 100644 --- a/dimos/robot/unitree/type/odometry.py +++ b/dimos/robot/unitree/type/odometry.py @@ -13,15 +13,16 @@ # limitations under the License. from typing import Literal, TypedDict -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.robot.unitree.type.timeseries import ( - Timestamped, +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped as GeneratedPoseStamped, + Quaternion as GeneratedQuaternion, ) -from dimos.types.timestamped import to_timestamp +from dimos_generated.std_msgs.msg import Header as GeneratedHeader -raw_odometry_msg_sample = { +raw_odometry_msg_sample: "RawOdometryMessage" = { "type": "msg", "topic": "rt/utlidar/robot_pose", "data": { @@ -73,32 +74,21 @@ class RawOdometryMessage(TypedDict): data: OdometryData -class Odometry(PoseStamped, Timestamped): # type: ignore[misc] - name = "geometry_msgs.PoseStamped" - - def __init__(self, frame_id: str = "base_link", *args, **kwargs) -> None: # type: ignore[no-untyped-def] - super().__init__(frame_id=frame_id, *args, **kwargs) # type: ignore[misc] - - @classmethod - def from_msg(cls, msg: RawOdometryMessage) -> "Odometry": - pose = msg["data"]["pose"] - - # Extract position - pos = Vector3( - pose["position"].get("x"), - pose["position"].get("y"), - pose["position"].get("z"), - ) - - rot = Quaternion( - pose["orientation"].get("x"), - pose["orientation"].get("y"), - pose["orientation"].get("z"), - pose["orientation"].get("w"), - ) - - ts = to_timestamp(msg["data"]["header"]["stamp"]) - return Odometry(position=pos, orientation=rot, ts=ts, frame_id="world") - - def __repr__(self) -> str: - return f"Odom pos({self.position}), rot({self.orientation})" +def pose_from_webrtc_odometry( + message: RawOdometryMessage, *, header: GeneratedHeader | None = None +) -> GeneratedPoseStamped: + """Copy the device's ROS-shaped pose, preserving its header unless explicitly replaced.""" + data = message["data"] + source_header = data["header"] + pose = data["pose"] + return GeneratedPoseStamped( + header=header + if header is not None + else GeneratedHeader( + stamp=Time(**source_header["stamp"]), frame_id=source_header["frame_id"] + ), + pose=Pose( + position=Point(**pose["position"]), + orientation=GeneratedQuaternion(**pose["orientation"]), + ), + ) diff --git a/dimos/robot/unitree/type/test_lidar.py b/dimos/robot/unitree/type/test_lidar.py index 631e7f680ef..e92bbf19135 100644 --- a/dimos/robot/unitree/type/test_lidar.py +++ b/dimos/robot/unitree/type/test_lidar.py @@ -16,19 +16,27 @@ import itertools from typing import cast +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest import reactivex as rx -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_xyz +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.robot.unitree.type.lidar import ( RawLidarMsg, pointcloud2_from_webrtc_lidar, repair_stale_ts, ) -from dimos.types.timestamped import Timestamped from dimos.utils.testing.replay import SensorReplay +def _stamped(seconds: float) -> PointCloud2: + return PointCloud2(header=Header(stamp=time_from_seconds(seconds))) + + @pytest.mark.self_hosted def test_init() -> None: lidar = SensorReplay("office_lidar") @@ -41,63 +49,75 @@ def test_init() -> None: def test_repair_stale_ts_extrapolates_non_monotonic_stamp() -> None: # input: two healthy frames, one stale, one healthy - raw = [Timestamped(0.000), Timestamped(0.130), Timestamped(-6.354), Timestamped(0.260)] + raw = [_stamped(0.000), _stamped(0.130), _stamped(-6.354), _stamped(0.260)] out: list[float] = [] rx.from_iterable(raw).pipe(repair_stale_ts(default_period=0.130)).subscribe( - on_next=lambda item: out.append(item.ts) + on_next=lambda item: out.append(to_seconds(item.header.stamp)) ) assert out == [0.000, 0.130, 0.260, 0.390] def test_repair_stale_ts_passes_monotonic_unchanged() -> None: - raw = [Timestamped(0.0), Timestamped(0.130), Timestamped(0.260)] + raw = [_stamped(0.0), _stamped(0.130), _stamped(0.260)] out: list[float] = [] rx.from_iterable(raw).pipe(repair_stale_ts()).subscribe( - on_next=lambda item: out.append(item.ts) + on_next=lambda item: out.append(to_seconds(item.header.stamp)) ) assert out == [0.0, 0.130, 0.260] def test_repair_stale_ts_handles_consecutive_bad_frames() -> None: # two stale-stamp frames in a row → each forward-extrapolated by default_period - raw = [Timestamped(0.0), Timestamped(-6.354), Timestamped(-6.354), Timestamped(0.500)] + raw = [_stamped(0.0), _stamped(-6.354), _stamped(-6.354), _stamped(0.500)] out: list[float] = [] rx.from_iterable(raw).pipe(repair_stale_ts(default_period=0.130)).subscribe( - on_next=lambda item: out.append(item.ts) + on_next=lambda item: out.append(to_seconds(item.header.stamp)) ) assert out == [0.0, 0.130, 0.260, 0.500] def test_repair_stale_ts_old_firmware_uses_system_time_after_calibration() -> None: - raw = [Timestamped(100.0) for _ in range(13)] + raw = [_stamped(100.0) for _ in range(13)] out: list[float] = [] rx.from_iterable(raw).pipe( repair_stale_ts(default_period=0.130, calibration_frames=10, now=lambda: 999.0) - ).subscribe(on_next=lambda item: out.append(item.ts)) + ).subscribe(on_next=lambda item: out.append(to_seconds(item.header.stamp))) assert out[0] == 100.0 assert out[1:10] == pytest.approx([100.0 + 0.130 * i for i in range(1, 10)]) assert out[10:] == [999.0, 999.0, 999.0] def test_repair_stale_ts_new_firmware_repair_persists_after_calibration() -> None: - raw = [Timestamped(i * 0.130) for i in range(11)] + [Timestamped(-6.354), Timestamped(2.0)] + raw = [_stamped(i * 0.130) for i in range(11)] + [_stamped(-6.354), _stamped(2.0)] out: list[float] = [] rx.from_iterable(raw).pipe( repair_stale_ts(default_period=0.130, calibration_frames=10, now=lambda: 999.0) - ).subscribe(on_next=lambda item: out.append(item.ts)) + ).subscribe(on_next=lambda item: out.append(to_seconds(item.header.stamp))) assert out[:11] == pytest.approx([i * 0.130 for i in range(11)]) assert out[11] == pytest.approx(10 * 0.130 + 0.130) assert out[12] == pytest.approx(2.0) def test_repair_stale_ts_calibration_boundary_one_differs() -> None: - raw = [Timestamped(5.0) for _ in range(9)] + [ - Timestamped(5.5), - Timestamped(5.6), - Timestamped(5.7), + raw = [_stamped(5.0) for _ in range(9)] + [ + _stamped(5.5), + _stamped(5.6), + _stamped(5.7), ] out: list[float] = [] rx.from_iterable(raw).pipe( repair_stale_ts(default_period=0.130, calibration_frames=10, now=lambda: 999.0) - ).subscribe(on_next=lambda item: out.append(item.ts)) + ).subscribe(on_next=lambda item: out.append(to_seconds(item.header.stamp))) assert 999.0 not in out + + +@pytest.mark.parametrize("override", [None, Time(sec=1700000000, nanosec=123456789)]) +def test_webrtc_xyz_conversion_preserves_points_and_explicit_time(override): + raw = {"data": {"stamp": -0.5, "data": {"points": np.array([[1.25, 2.5, 3.75]])}}} + cloud = pointcloud2_from_webrtc_lidar(cast("RawLidarMsg", raw), stamp=override) + decoded = PointCloud2.decode(cloud.encode()) + np.testing.assert_array_equal(pointcloud_xyz(decoded), [[1.25, 2.5, 3.75]]) + assert decoded.header.frame_id == "world" + assert decoded.header.stamp == ( + Time(sec=-1, nanosec=500000000) if override is None else override + ) diff --git a/dimos/robot/unitree/type/test_odometry.py b/dimos/robot/unitree/type/test_odometry.py index 8020684fb79..b383a18e6ee 100644 --- a/dimos/robot/unitree/type/test_odometry.py +++ b/dimos/robot/unitree/type/test_odometry.py @@ -14,9 +14,11 @@ from __future__ import annotations +from dimos_generated.geometry_msgs.msg import PoseStamped import pytest -from dimos.robot.unitree.type.odometry import Odometry +from dimos.msgs.geometry import yaw +from dimos.robot.unitree.type.odometry import pose_from_webrtc_odometry from dimos.utils.testing.replay import SensorReplay _EXPECTED_TOTAL_RAD = -4.05212 @@ -28,22 +30,24 @@ def test_dataset_size() -> None: def test_odometry_conversion_and_count() -> None: - """Each replay entry converts to :class:`Odometry` and count is correct.""" + """Each replay entry converts to :class:`PoseStamped` and count is correct.""" for raw in SensorReplay(name="raw_odometry_rotate_walk").iterate(): - odom = Odometry.from_msg(raw) + odom = pose_from_webrtc_odometry(raw) assert isinstance(raw, dict) - assert isinstance(odom, Odometry) + assert isinstance(odom, PoseStamped) def test_total_rotation_travel_iterate() -> None: total_rad = 0.0 prev_yaw: float | None = None - for odom in SensorReplay(name="raw_odometry_rotate_walk", autocast=Odometry.from_msg).iterate(): - yaw = odom.orientation.radians.z + for odom in SensorReplay( + name="raw_odometry_rotate_walk", autocast=pose_from_webrtc_odometry + ).iterate(): + angle = yaw(odom.pose.orientation) if prev_yaw is not None: - diff = yaw - prev_yaw + diff = angle - prev_yaw total_rad += diff - prev_yaw = yaw + prev_yaw = angle assert total_rad == pytest.approx(_EXPECTED_TOTAL_RAD, abs=0.001) diff --git a/dimos/robot/unitree/type/test_pose_conversion.py b/dimos/robot/unitree/type/test_pose_conversion.py new file mode 100644 index 00000000000..a8da638faed --- /dev/null +++ b/dimos/robot/unitree/type/test_pose_conversion.py @@ -0,0 +1,45 @@ +# 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. + +import copy + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.std_msgs.msg import Header +import pytest + +from dimos.robot.unitree.type.odometry import pose_from_webrtc_odometry, raw_odometry_msg_sample + + +@pytest.mark.parametrize("seconds,nanoseconds", [(-1, 987654321), (1700000000, 123456789)]) +def test_device_pose_preserves_original_ros_header(seconds, nanoseconds): + raw = copy.deepcopy(raw_odometry_msg_sample) + raw["data"]["header"]["stamp"] = {"sec": seconds, "nanosec": nanoseconds} + message = pose_from_webrtc_odometry(raw) + decoded = PoseStamped.decode(message.encode()) + assert decoded.header == Header(frame_id="odom", stamp=Time(sec=seconds, nanosec=nanoseconds)) + assert decoded.pose.position.x == 5.961965 + assert decoded.pose.position.y == -2.916958 + assert decoded.pose.position.z == 0.319509 + assert decoded.pose.orientation.w == -0.242112 + message.header.frame_id = "changed" + assert raw["data"]["header"]["frame_id"] == "odom" + + +def test_explicit_arrival_header_is_copied(): + header = Header(frame_id="world", stamp=Time(sec=5, nanosec=123)) + message = pose_from_webrtc_odometry(raw_odometry_msg_sample, header=header) + assert message.header == header + header.stamp.nanosec = 999 + assert message.header.stamp.nanosec == 123 diff --git a/dimos/robot/unitree/unitree_skill_container.py b/dimos/robot/unitree/unitree_skill_container.py index ebccee8941a..de4f6385ed4 100644 --- a/dimos/robot/unitree/unitree_skill_container.py +++ b/dimos/robot/unitree/unitree_skill_container.py @@ -19,16 +19,21 @@ import math import time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Vector3 +from dimos_generated.tf2_msgs.msg import TFMessage from unitree_webrtc_connect.constants import RTC_TOPIC from dimos.agents.annotation import skill from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import In -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.geometry import ( + pose_from_transform, + quaternion_euler, + quaternion_from_euler, + translate_pose_local, +) +from dimos.msgs.time import header_now from dimos.navigation.base import NavigationState from dimos.navigation.go2.replanning_a_star.spec import NavigationInterfaceSpec from dimos.robot.unitree.go2.connection_spec import GO2ConnectionSpec @@ -199,20 +204,24 @@ def _goal_pose( current: PoseStamped, x: float, y: float, degrees: float | None, relative: bool ) -> PoseStamped: """Where move_to sends the robot, in the world frame.""" - euler = current.orientation.to_euler() + roll, pitch, yaw = quaternion_euler(current.pose.orientation) if relative: - position = current.position + current.orientation.rotate_vector(Vector3(x, y, 0)) - yaw = euler.yaw + math.radians(degrees or 0.0) + translated = translate_pose_local(current.pose, Vector3(x=x, y=y)) + position = translated.position + heading = yaw + math.radians(degrees or 0.0) else: - position = Vector3(x, y, current.position.z) - yaw = euler.yaw if degrees is None else math.radians(degrees) - orientation = Quaternion.from_euler(Vector3(euler.roll, euler.pitch, yaw)) - return PoseStamped(position=position, orientation=orientation, frame_id="world") + position = Point(x=x, y=y, z=current.pose.position.z) + heading = yaw if degrees is None else math.radians(degrees) + orientation = quaternion_from_euler(roll, pitch, heading) + return PoseStamped( + header=header_now("world"), pose=Pose(position=position, orientation=orientation) + ) def _pose_text(pose: PoseStamped) -> str: - yaw = math.degrees(pose.orientation.to_euler().yaw) - return f"x={pose.position.x:.2f} y={pose.position.y:.2f} heading={yaw:.0f}deg" + heading = math.degrees(quaternion_euler(pose.pose.orientation)[2]) + position = pose.pose.position + return f"x={position.x:.2f} y={position.y:.2f} heading={heading:.0f}deg" class UnitreeSkillContainer(Module): @@ -265,12 +274,12 @@ def move_to( if tf is None: return "Failed to get the position of the robot." - goal = _goal_pose(tf.to_pose(), x, y, degrees, relative) + goal = _goal_pose(pose_from_transform(tf), x, y, degrees, relative) self._navigation.set_goal(goal) outcome = self._wait_for_goal() tf = self.tfbuffer.get("world", "base_link") - now = "unknown" if tf is None else _pose_text(tf.to_pose()) + now = "unknown" if tf is None else _pose_text(pose_from_transform(tf)) return f"{outcome}. Robot is at {now}; goal was {_pose_text(goal)}." def _wait_for_goal(self, timeout: float = 100.0, settle: float = 2.0) -> str: diff --git a/dimos/simulation/dimsim/cdr_schemas.py b/dimos/simulation/dimsim/cdr_schemas.py new file mode 100644 index 00000000000..44d4f858f30 --- /dev/null +++ b/dimos/simulation/dimsim/cdr_schemas.py @@ -0,0 +1,41 @@ +# 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. + +"""Export canonical generated ROS 2 schemas for the standalone DimSim bundle. + +Run ``python -m dimos.simulation.dimsim.cdr_schemas`` after changing messages. +""" + +import json + +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 + +from dimos.constants import DIMOS_PROJECT_ROOT + +SCHEMA_PATH = DIMOS_PROJECT_ROOT / "misc/DimSim/cli/bridge/cdr_schemas.ts" + + +def schema_source() -> str: + definitions = {t.msg_name: t.schema for t in (Image, PointCloud2, PoseStamped, Twist)} + return ( + "// Generated from the canonical .msg definitions; do not edit by hand.\n" + "export const schemas: Record = " + + json.dumps(definitions, indent=2) + + ";\n" + ) + + +if __name__ == "__main__": + SCHEMA_PATH.write_text(schema_source()) diff --git a/dimos/simulation/dimsim/test_cdr_schemas.py b/dimos/simulation/dimsim/test_cdr_schemas.py new file mode 100644 index 00000000000..c77bb9f8bdf --- /dev/null +++ b/dimos/simulation/dimsim/test_cdr_schemas.py @@ -0,0 +1,19 @@ +# 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. + +from dimos.simulation.dimsim.cdr_schemas import SCHEMA_PATH, schema_source + + +def test_standalone_dimsim_schemas_match_canonical_messages(): + assert SCHEMA_PATH.read_text() == schema_source() diff --git a/dimos/simulation/engines/base.py b/dimos/simulation/engines/base.py index d4b0735528d..594456e0e9a 100644 --- a/dimos/simulation/engines/base.py +++ b/dimos/simulation/engines/base.py @@ -22,7 +22,7 @@ if TYPE_CHECKING: from pathlib import Path - from dimos.msgs.sensor_msgs.JointState import JointState + from dimos_generated.sensor_msgs.msg import JointState class SimulationEngine(ABC): diff --git a/dimos/simulation/engines/mujoco_engine.py b/dimos/simulation/engines/mujoco_engine.py index b9587e01995..7ba9c24c220 100644 --- a/dimos/simulation/engines/mujoco_engine.py +++ b/dimos/simulation/engines/mujoco_engine.py @@ -40,7 +40,7 @@ from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: - from dimos.msgs.sensor_msgs.JointState import JointState + from dimos_generated.sensor_msgs.msg import JointState logger = setup_logger() @@ -725,15 +725,15 @@ def read_joint_efforts(self) -> list[float]: def write_joint_command(self, command: JointState) -> None: if command.position: self._command_mode = "position" - self._set_position_targets(command.position) + self._set_position_targets(list(command.position)) return if command.velocity: self._command_mode = "velocity" - self._set_velocity_targets(command.velocity) + self._set_velocity_targets(list(command.velocity)) return if command.effort: self._command_mode = "effort" - self._set_effort_targets(command.effort) + self._set_effort_targets(list(command.effort)) return def _set_position_targets(self, positions: list[float]) -> None: diff --git a/dimos/simulation/engines/test_mujoco_sim_module.py b/dimos/simulation/engines/test_mujoco_sim_module.py index 30761039d3e..03998366f78 100644 --- a/dimos/simulation/engines/test_mujoco_sim_module.py +++ b/dimos/simulation/engines/test_mujoco_sim_module.py @@ -20,10 +20,14 @@ from typing import Any from unittest.mock import MagicMock +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import Imu +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo +from dimos.msgs.camera_info import camera_info_from_intrinsics +from dimos.msgs.time import to_seconds from dimos.simulation.engines.mujoco_engine import CameraFrame, MujocoEngine from dimos.simulation.engines.mujoco_sim_module import MujocoSimModule, MujocoSimModuleConfig @@ -125,6 +129,27 @@ def cleanup(self) -> None: module._publish_shm_and_lcm(_FakeEngine) assert events == ["joint_state", "imu", "ready"] + pose = PoseStamped.decode(module.odom.publish.call_args.args[0].encode()) + assert pose.header.frame_id == "world" + assert (pose.pose.position.x, pose.pose.position.y, pose.pose.position.z) == ( + 0.0, + 0.0, + 0.75, + ) + assert pose.pose.orientation.w == 1.0 + imu = Imu.decode(module.imu.publish.call_args.args[0].encode()) + assert imu.orientation.w == 1.0 + assert (imu.angular_velocity.x, imu.angular_velocity.y, imu.angular_velocity.z) == ( + 0.1, + 0.2, + 0.3, + ) + assert ( + imu.linear_acceleration.x, + imu.linear_acceleration.y, + imu.linear_acceleration.z, + ) == (1.0, 2.0, 3.0) + assert imu.header.stamp.sec > 0 finally: module.stop() @@ -158,14 +183,28 @@ def disconnect(self) -> None: module._publish_tf(10.0, frame) color_tf, depth_tf, camera_link_tf = messages[-1].transforms - assert color_tf.frame_id == "link7" + assert color_tf.header.frame_id == "link7" assert color_tf.child_frame_id == "wrist_camera_color_optical_frame" - assert np.allclose(color_tf.translation.to_numpy(), [0.0, 0.0, 1.0]) - assert depth_tf.frame_id == "link7" + assert np.allclose( + [ + color_tf.transform.translation.x, + color_tf.transform.translation.y, + color_tf.transform.translation.z, + ], + [0.0, 0.0, 1.0], + ) + assert depth_tf.header.frame_id == "link7" assert depth_tf.child_frame_id == "wrist_camera_depth_optical_frame" - assert camera_link_tf.frame_id == "link7" + assert camera_link_tf.header.frame_id == "link7" assert camera_link_tf.child_frame_id == "wrist_camera_link" - assert np.allclose(camera_link_tf.translation.to_numpy(), [0.0, 0.0, 1.0]) + assert np.allclose( + [ + camera_link_tf.transform.translation.x, + camera_link_tf.transform.translation.y, + camera_link_tf.transform.translation.z, + ], + [0.0, 0.0, 1.0], + ) finally: module.stop() @@ -362,8 +401,14 @@ def test_publish_loop_stamps_messages_with_frame_timestamp() -> None: module = MujocoSimModule() try: module.config = MujocoSimModuleConfig(fps=1000) - module._camera_info_base = CameraInfo.from_intrinsics( - width=1, height=1, fx=1.0, fy=1.0, cx=0.5, cy=0.5, frame_id="wrist_camera_color_frame" + module._camera_info_base = camera_info_from_intrinsics( + width=1, + height=1, + fx=1.0, + fy=1.0, + cx=0.5, + cy=0.5, + header=Header(frame_id="wrist_camera_color_frame"), ) color: list[Any] = [] depth: list[Any] = [] @@ -377,13 +422,13 @@ def test_publish_loop_stamps_messages_with_frame_timestamp() -> None: _run_publish_loop(module, frame_ts) module._publish_camera_info() - assert [img.ts for img in color] == frame_ts - assert [img.ts for img in depth] == frame_ts - assert [msg.transforms[0].ts for msg in tf] == frame_ts + assert [to_seconds(img.header.stamp) for img in color] == frame_ts + assert [to_seconds(img.header.stamp) for img in depth] == frame_ts + assert [to_seconds(msg.transforms[0].header.stamp) for msg in tf] == frame_ts # camera_info rides the latest frame's sim clock, not wall time. - assert info[-1].ts == frame_ts[-1] - assert module.get_color_camera_info().ts == frame_ts[-1] - assert module.get_depth_camera_info().ts == frame_ts[-1] + assert to_seconds(info[-1].header.stamp) == frame_ts[-1] + assert to_seconds(module.get_color_camera_info().header.stamp) == frame_ts[-1] + assert to_seconds(module.get_depth_camera_info().header.stamp) == frame_ts[-1] finally: module.stop() @@ -392,8 +437,14 @@ def test_camera_info_falls_back_to_wall_clock_before_first_frame() -> None: module = MujocoSimModule() try: module.config = MujocoSimModuleConfig() - module._camera_info_base = CameraInfo.from_intrinsics( - width=1, height=1, fx=1.0, fy=1.0, cx=0.5, cy=0.5, frame_id="wrist_camera_color_frame" + module._camera_info_base = camera_info_from_intrinsics( + width=1, + height=1, + fx=1.0, + fy=1.0, + cx=0.5, + cy=0.5, + header=Header(frame_id="wrist_camera_color_frame"), ) assert module._latest_frame_ts is None assert module._camera_info_ts() == pytest.approx(time.time(), abs=5.0) diff --git a/dimos/simulation/habitat/nix/flake.nix b/dimos/simulation/habitat/nix/flake.nix index 89016fef741..60508eff56d 100644 --- a/dimos/simulation/habitat/nix/flake.nix +++ b/dimos/simulation/habitat/nix/flake.nix @@ -12,8 +12,8 @@ # Provides the installer, not the simulator: habitat-sim is conda-only and # headless rendering needs the host's EGL driver, so this cannot be a derivation. devShells = forAll (pkgs: { - default = pkgs.mkShellNoCC { - packages = [ pkgs.micromamba pkgs.curl pkgs.cacert ]; + default = pkgs.mkShell { + packages = [ pkgs.micromamba pkgs.curl pkgs.cacert pkgs.cmake ]; }; }); }; diff --git a/dimos/simulation/habitat/nix/install.sh b/dimos/simulation/habitat/nix/install.sh index 59e523e773b..e4450a9873a 100755 --- a/dimos/simulation/habitat/nix/install.sh +++ b/dimos/simulation/habitat/nix/install.sh @@ -16,10 +16,24 @@ if [ ! -x ./env/bin/python ]; then python=3.9 habitat-sim headless withbullet fi -# Pinned to uv.lock: the native must speak the same zenoh wire version as the -# dimos peers, or SHM payloads arrive as unreadable handles. --no-deps skips -# lcm-dimos-fork, the LCM runtime; only dimos_lcm's pure-python encoders are needed. -./env/bin/pip install --no-input --no-deps "dimos-lcm==0.1.3" +# Build standalone generated CDR messages for Habitat's Python 3.9 interpreter. +# The DimOS checkout's environment generates the package; the native environment +# builds its own extension and never imports or installs DimOS at runtime. +if [ ! -x "$ROOT/.venv/bin/python" ]; then + echo "Set up the DimOS project virtualenv before building Habitat messages." >&2 + exit 1 +fi +if [ ! -f "$ROOT/build/message-codegen/install/lib/cmake/fastcdr/fastcdr-config.cmake" ]; then + bash "$ROOT/scripts/setup_message_codegen.sh" +fi +(cd "$ROOT" && .venv/bin/python -m dimos.message_codegen.generate --package \ + --output "$OUT/messages" \ + --type geometry_msgs/msg/Twist --type nav_msgs/msg/Odometry \ + --type sensor_msgs/msg/CameraInfo --type sensor_msgs/msg/Image \ + --type sensor_msgs/msg/PointCloud2 --type tf2_msgs/msg/TFMessage) +CMAKE_PREFIX_PATH="$ROOT/build/message-codegen/install${CMAKE_PREFIX_PATH:+;$CMAKE_PREFIX_PATH}" \ + ./env/bin/pip install --no-input "$OUT/messages/python" +# Keep the native transport wire version pinned to uv.lock. ./env/bin/pip install --no-input "eclipse-zenoh==1.10.1" numpy # Annotated HM3D house, no credentials. --no-replace resumes a partial download diff --git a/dimos/simulation/mujoco/direct_cmd_vel_explorer.py b/dimos/simulation/mujoco/direct_cmd_vel_explorer.py index 73bb35a9de2..3b012c4092f 100644 --- a/dimos/simulation/mujoco/direct_cmd_vel_explorer.py +++ b/dimos/simulation/mujoco/direct_cmd_vel_explorer.py @@ -16,11 +16,11 @@ import threading from typing import TYPE_CHECKING +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist, Vector3 + from dimos.core.transport import PubSubTransport from dimos.core.transport_factory import make_transport -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.geometry import yaw if TYPE_CHECKING: from collections.abc import Callable @@ -82,13 +82,13 @@ def _drive_to(self, target_x: float, target_y: float) -> None: """Pursuit controller: steer toward the target while driving forward.""" while True: pose = self._wait_for_pose() - dx = target_x - pose.x - dy = target_y - pose.y + dx = target_x - pose.pose.position.x + dy = target_y - pose.pose.position.y distance = math.hypot(dx, dy) if distance < 0.3: break target_heading = math.atan2(dy, dx) - heading_error = self._normalize_angle(target_heading - pose.yaw) + heading_error = self._normalize_angle(target_heading - yaw(pose.pose.orientation)) # Only drive forward when roughly facing the target. if abs(heading_error) > 0.3: linear = 0.0 @@ -98,7 +98,7 @@ def _drive_to(self, target_x: float, target_y: float) -> None: assert self._cmd_vel is not None self._cmd_vel.broadcast( None, - Twist(linear=Vector3(linear, 0, 0), angular=Vector3(0, 0, angular)), + Twist(linear=Vector3(x=linear), angular=Vector3(z=angular)), ) self._stop() diff --git a/dimos/simulation/mujoco/test_generated_control_consumers.py b/dimos/simulation/mujoco/test_generated_control_consumers.py new file mode 100644 index 00000000000..bd54c2e62db --- /dev/null +++ b/dimos/simulation/mujoco/test_generated_control_consumers.py @@ -0,0 +1,70 @@ +# 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. + +from unittest.mock import MagicMock + +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Twist +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest + +from dimos.msgs.geometry import quaternion_from_euler +from dimos.simulation.mujoco.direct_cmd_vel_explorer import DirectCmdVelExplorer +from dimos.simulation.mujoco.person_on_track import PersonTrackPublisher + + +@pytest.mark.parametrize( + "heading,expected_linear,expected_angular", [(0.0, 0.8, 0.0), (1.0, 0.0, -1.5)] +) +def test_pursuit_generated_pose_preserves_steering_and_stop( + heading: float, expected_linear: float, expected_angular: float, monkeypatch: pytest.MonkeyPatch +) -> None: + explorer = DirectCmdVelExplorer() + output = MagicMock() + explorer._cmd_vel = output + poses = iter( + [ + PoseStamped( + header=Header(frame_id="world"), + pose=Pose(position=Point(), orientation=quaternion_from_euler(0.0, 0.0, heading)), + ), + PoseStamped( + header=Header(frame_id="world"), + pose=Pose(position=Point(x=1.0), orientation=Quaternion(w=1.0)), + ), + ] + ) + monkeypatch.setattr(explorer, "_wait_for_pose", lambda: next(poses)) + explorer._drive_to(1.0, 0.0) + commands = [Twist.decode(call.args[1].encode()) for call in output.broadcast.call_args_list] + assert len(commands) == 2 + assert commands[0].linear.x == expected_linear + assert commands[0].angular.z == expected_angular + assert commands[1].linear.x == 0.0 + assert commands[1].angular.z == 0.0 + + +def test_person_track_generated_pose_preserves_position_and_heading( + monkeypatch: pytest.MonkeyPatch, +) -> None: + output = MagicMock() + monkeypatch.setattr( + "dimos.simulation.mujoco.person_on_track.make_transport", lambda channel, message: output + ) + publisher = PersonTrackPublisher([(1.0, 2.0), (1.0, 3.0), (2.0, 3.0)]) + publisher.tick() + pose = Pose.decode(output.broadcast.call_args.args[1].encode()) + assert (pose.position.x, pose.position.y, pose.position.z) == (1.0, 2.0, 0.0) + assert np.isclose(pose.orientation.z, np.sin(np.pi / 4)) + assert np.isclose(pose.orientation.w, np.cos(np.pi / 4)) diff --git a/dimos/spec/mapping.py b/dimos/spec/mapping.py index f35778f40b4..93a762f687a 100644 --- a/dimos/spec/mapping.py +++ b/dimos/spec/mapping.py @@ -14,9 +14,10 @@ from typing import Protocol +from dimos_generated.nav_msgs.msg import OccupancyGrid +from dimos_generated.sensor_msgs.msg import PointCloud2 + from dimos.core.stream import Out -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 class GlobalPointcloud(Protocol): diff --git a/dimos/spec/perception.py b/dimos/spec/perception.py index 4fac65ad020..72bdae0fbd8 100644 --- a/dimos/spec/perception.py +++ b/dimos/spec/perception.py @@ -14,12 +14,10 @@ from typing import Protocol +from dimos_generated.nav_msgs.msg import Odometry as OdometryMsg +from dimos_generated.sensor_msgs.msg import CameraInfo, Image as ImageMsg, Imu, PointCloud2 + from dimos.core.stream import Out -from dimos.msgs.nav_msgs.Odometry import Odometry as OdometryMsg -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image as ImageMsg -from dimos.msgs.sensor_msgs.Imu import Imu -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 class Image(Protocol): diff --git a/dimos/teleop/hosted/arm_command.py b/dimos/teleop/hosted/arm_command.py index 47b8aec6261..04d1a272016 100644 --- a/dimos/teleop/hosted/arm_command.py +++ b/dimos/teleop/hosted/arm_command.py @@ -24,15 +24,14 @@ import time from typing import Any -from dimos_lcm.geometry_msgs import TwistStamped as LCMTwistStamped +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist, TwistStamped, Vector3 +from dimos_generated.std_msgs.msg import Float32, Header from reactivex.disposable import Disposable from dimos.control.coordinator import ControlCoordinator from dimos.core.core import rpc from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped -from dimos.msgs.std_msgs.Float32 import Float32 +from dimos.msgs.time import to_seconds from dimos.teleop.hosted.command_executor import SerializedCommandExecutor from dimos.teleop.utils.teleop_transforms import webxr_to_robot from dimos.teleop.webxr.controller_types import Hand @@ -75,7 +74,8 @@ def __init__(self, **kwargs: Any) -> None: self._last_stale_warn = 0.0 self._last_future_warn = 0.0 - self._decoders[LCMTwistStamped._get_packed_fingerprint()] = self._on_twist_bytes + self._command_types["twist"] = TwistStamped + self._decoders["twist"] = self._on_twist_bytes # No local WebSocket server — the operator connects through the broker. def _start_server(self) -> None: @@ -103,25 +103,20 @@ def stop(self) -> None: # ─── Inbound command plane (operator → robot) ───────────────────── def _on_cmd_raw(self, data: Any) -> None: - """Fingerprint-dispatch LCM bytes (Pose/Joy inherited, Twist added here).""" + """Dispatch explicit CDR command frames (pose, joy and twist).""" if isinstance(data, str): data = data.encode() - decoder = self._decoders.get(data[:8]) - if decoder is None: - return - try: - decoder(data) - except Exception: - logger.warning("cmd_raw decode failed", exc_info=True) + if isinstance(data, bytes): + self._dispatch_binary_message(data) def _on_pose_bytes(self, data: bytes) -> None: """Controller pose → robot frame; stale/future/out-of-order dropped.""" - msg = PoseStamped.lcm_decode(data) + msg = PoseStamped.decode(data) try: - hand = self._resolve_hand(msg.frame_id) + hand = self._resolve_hand(msg.header.frame_id) except ValueError: return - ts = float(msg.ts) + ts = to_seconds(msg.header.stamp) if not math.isfinite(ts): return age = time.time() - ts @@ -148,8 +143,8 @@ def _on_twist_bytes(self, data: bytes) -> None: """Browser keyboard EE-twist → coordinator's eef_twist task.""" if self._estopped: return - msg = TwistStamped.lcm_decode(data) - ts = float(msg.ts) + msg = TwistStamped.decode(data) + ts = to_seconds(msg.header.stamp) if not math.isfinite(ts): return age = time.time() - ts @@ -172,13 +167,15 @@ def _on_twist_bytes(self, data: bytes) -> None: scale = self._translation_scale self.ee_twist_command.publish( TwistStamped( - linear=[ - msg.linear.x * scale, - msg.linear.y * scale, - msg.linear.z * scale, - ], - angular=[msg.angular.x, msg.angular.y, msg.angular.z], - ts=msg.ts, + header=Header(stamp=msg.header.stamp), + twist=Twist( + linear=Vector3( + x=msg.twist.linear.x * scale, + y=msg.twist.linear.y * scale, + z=msg.twist.linear.z * scale, + ), + angular=msg.twist.angular, + ), ) ) diff --git a/dimos/teleop/hosted/camera_mux.py b/dimos/teleop/hosted/camera_mux.py index d6eeb7ae3c7..7d5dc07dd23 100644 --- a/dimos/teleop/hosted/camera_mux.py +++ b/dimos/teleop/hosted/camera_mux.py @@ -28,13 +28,16 @@ import time from typing import Any +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header import numpy as np from reactivex.disposable import Disposable from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_from_array, image_view +from dimos.msgs.time import to_nanoseconds from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -128,6 +131,8 @@ def _on_cam(self, cam: str, img: Image) -> None: def _composite(self) -> Image | None: """Selected frames → one even-sized Image; None on any error (a raise would kill the RxPY camera subscription).""" + import cv2 + with self._cam_lock: order = [c for c in self._cam_order if c in self._cam_selected] imgs = [self._cam_frames[c] for c in order if c in self._cam_frames] @@ -136,21 +141,27 @@ def _composite(self) -> Image | None: try: if len(imgs) == 1: return self._even_dims(self._stamp(self._downscale(imgs[0]))) - import cv2 - - target_h = min(im.data.shape[0] for im in imgs) + if len({im.encoding for im in imgs}) != 1: + raise ValueError("Selected camera encodings must match") + target_h = min(im.height for im in imgs) tiles = [] for im in imgs: - h, w = im.data.shape[:2] + data = image_view(im) + h, w = data.shape[:2] tiles.append( - cv2.resize(im.data, (max(1, int(w * target_h / h)), target_h)) + cv2.resize(data, (max(1, int(w * target_h / h)), target_h)) if h != target_h - else im.data + else data ) + latest = max(imgs, key=lambda im: to_nanoseconds(im.header.stamp)) return self._even_dims( self._stamp( self._downscale( - Image(data=np.hstack(tiles), format=imgs[0].format, frame_id="camera_mux") + image_from_array( + np.hstack(tiles), + encoding=imgs[0].encoding, + header=Header(stamp=latest.header.stamp, frame_id="camera_mux"), + ) ) ) ) @@ -162,28 +173,30 @@ def _composite(self) -> Image | None: def _even_dims(img: Image) -> Image: """Crop to even width/height — an odd composite crashes libx264's avcodec_open2 when a camera switch reopens the encoder.""" - data = img.data + data = image_view(img) if data.ndim < 2: return img h, w = data.shape[:2] if h % 2 == 0 and w % 2 == 0: return img data = data[: h - (h % 2), : w - (w % 2)] - return Image(data=np.ascontiguousarray(data), format=img.format, frame_id=img.frame_id) + return image_from_array( + np.ascontiguousarray(data), encoding=img.encoding, header=img.header + ) def _downscale(self, img: Image) -> Image: """Cap publish width at config.video_max_width (0 = off). Runs before _stamp so the strip's 16px cells stay decodable at the sent size.""" + import cv2 + max_w = self.config.video_max_width - if max_w <= 0 or img.data.ndim < 2: + if max_w <= 0 or image_view(img).ndim < 2: return img - h, w = img.data.shape[:2] + h, w = image_view(img).shape[:2] if w <= max_w: return img - import cv2 - - out = cv2.resize(img.data, (max_w, max(1, int(h * max_w / w)))) - return Image(data=out, format=img.format, frame_id=img.frame_id) + out = cv2.resize(image_view(img), (max_w, max(1, int(h * max_w / w)))) + return image_from_array(out, encoding=img.encoding, header=img.header) def _stamp(self, img: Image) -> Image: """Append (not overwrite) a bottom strip encoding capture time as B/W @@ -197,7 +210,7 @@ def _stamp(self, img: Image) -> Image: ] s = _STAMP_CELL_PX - data = img.data + data = image_view(img) if data.ndim < 2 or data.shape[1] < _STAMP_CELLS * s: if not self._stamp_warned: self._stamp_warned = True @@ -215,7 +228,7 @@ def _stamp(self, img: Image) -> Image: if bit: strip[:, i * s : (i + 1) * s] = 255 out = np.vstack([data, strip]) - return Image(data=out, format=img.format, frame_id=img.frame_id) + return image_from_array(out, encoding=img.encoding, header=img.header) def _set_cam_selection(self, data: bytes) -> None: """camera_select kind → filter to known cams, republish immediately so diff --git a/dimos/teleop/hosted/go2_command.py b/dimos/teleop/hosted/go2_command.py index bdac4661e7d..5ac70951f22 100644 --- a/dimos/teleop/hosted/go2_command.py +++ b/dimos/teleop/hosted/go2_command.py @@ -27,16 +27,22 @@ import time from typing import Any -from dimos_lcm.std_msgs import Bool +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Twist, + TwistStamped, +) +from dimos_generated.std_msgs.msg import Bool from reactivex.disposable import Disposable from unitree_webrtc_connect.constants import SPORT_CMD from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped +from dimos.msgs.time import header_now, to_nanoseconds from dimos.robot.unitree.go2.connection import GO2Connection from dimos.teleop.hosted.command_executor import SerializedCommandExecutor from dimos.utils.generic import finite_number @@ -100,7 +106,7 @@ def __init__(self, **kwargs: Any) -> None: self._obstacle_avoidance = True self._light = 0.0 self._posture = "StandReady" - self._last_cmd_ts = 0.0 + self._last_cmd_ns = 0 self._last_cmd_nonzero = False @rpc @@ -358,7 +364,8 @@ def _handle_nav_goal(self, msg: dict[str, Any]) -> None: self._send_ack(nonce, False) return pose = PoseStamped( - ts=time.time(), frame_id="world", position=[x, y, 0.0], orientation=[0, 0, 0, 1] + header=header_now("world"), + pose=Pose(position=Point(x=x, y=y), orientation=Quaternion(w=1)), ) try: self.goal_request.publish(pose) @@ -388,36 +395,48 @@ def _on_cmd_vel_in(self, twist: TwistStamped) -> None: """Guard raw operator drive, republish on tele_cmd_vel for MovementManager.""" if self._estopped: return - ts = float(twist.ts) - if not math.isfinite(ts): - # NaN ts passes every comparison below and would poison _last_cmd_ts. + try: + stamp_ns = to_nanoseconds(twist.header.stamp) + except ValueError: + logger.warning("dropping cmd_vel with invalid nanoseconds") return - age = time.time() - ts - if age > self.config.cmd_stale_after_sec: + age_ns = time.time_ns() - stamp_ns + if age_ns > self.config.cmd_stale_after_sec * 1_000_000_000: return - if age < 0: # future-stamped: don't advance _last_cmd_ts (would stall drive) + if age_ns < 0: # future-stamped: do not stall subsequent drive frames return - if ts <= self._last_cmd_ts: # out-of-order + if stamp_ns <= self._last_cmd_ns: return - self._last_cmd_ts = ts - - if not _all_finite(twist): + command = twist.twist + if not _all_finite(command): logger.warning("dropping non-finite cmd_vel") return + self._last_cmd_ns = stamp_ns + lin_max = self.config.max_linear_mps ang_max = self.config.max_angular_rps - twist.linear.x = _clamp(twist.linear.x, -lin_max, lin_max) - twist.linear.y = _clamp(twist.linear.y, -lin_max, lin_max) - twist.angular.z = _clamp(twist.angular.z, -ang_max, ang_max) + command.linear.x = _clamp(command.linear.x, -lin_max, lin_max) + command.linear.y = _clamp(command.linear.y, -lin_max, lin_max) + command.angular.z = _clamp(command.angular.z, -ang_max, ang_max) # Idle zeros would make MovementManager cancel the nav plan, so forward a # zero only as the release edge (prev frame moving), then stay silent. - moving = not twist.is_zero() + moving = any( + abs(value) > 1e-8 + for value in ( + command.linear.x, + command.linear.y, + command.linear.z, + command.angular.x, + command.angular.y, + command.angular.z, + ) + ) if not moving and not self._last_cmd_nonzero: return self._last_cmd_nonzero = moving - self.tele_cmd_vel.publish(Twist(linear=twist.linear, angular=twist.angular)) + self.tele_cmd_vel.publish(command) # ─── robot-authoritative state → stats module ───────────────────── diff --git a/dimos/teleop/hosted/hosted_stats.py b/dimos/teleop/hosted/hosted_stats.py index c572398b09f..3b8bdd66fdb 100644 --- a/dimos/teleop/hosted/hosted_stats.py +++ b/dimos/teleop/hosted/hosted_stats.py @@ -28,15 +28,17 @@ import time from typing import Any +from dimos_generated.dimos_msgs.msg import VideoStats +from dimos_generated.geometry_msgs.msg import TwistStamped from reactivex.disposable import Disposable from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped +from dimos.msgs.time import to_seconds from dimos.robot.unitree.go2.connection import GO2Connection from dimos.teleop.utils.stream_stats import LiveStreamStats -from dimos.teleop.utils.video_stats import VideoStats +from dimos.teleop.utils.video_stats import video_stats_from_dict from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -108,7 +110,7 @@ def _on_state_json(self, data: Any) -> None: kind = msg.get("type") if kind == "video_stats": try: - self.video_stats.publish(VideoStats.from_dict(msg)) + self.video_stats.publish(video_stats_from_dict(msg)) except (TypeError, ValueError): logger.warning("state_reliable: malformed video_stats, dropping") elif kind == "clock_report": @@ -124,10 +126,11 @@ def _on_cmd_raw(self, data: Any) -> None: if isinstance(data, str): data = data.encode() try: - cmd = TwistStamped.lcm_decode(data) + cmd = TwistStamped.decode(data) + source_time = to_seconds(cmd.header.stamp) except Exception: return - self._cmd_stats.record(cmd.ts, nbytes=len(data)) + self._cmd_stats.record(source_time, nbytes=len(data)) self.cmd_vel_stamped.publish(cmd) def _on_robot_state(self, data: Any) -> None: diff --git a/dimos/teleop/hosted/map_compress.py b/dimos/teleop/hosted/map_compress.py index 0928080aa5b..b1c64758786 100644 --- a/dimos/teleop/hosted/map_compress.py +++ b/dimos/teleop/hosted/map_compress.py @@ -27,13 +27,17 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.nav_msgs.msg import OccupancyGrid +import numpy as np from reactivex.disposable import Disposable from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid, block_max_reduce +from dimos.msgs.geometry import quaternion_euler +from dimos.msgs.occupancy import block_max_reduce, occupancy_view +from dimos.msgs.time import to_seconds from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -75,19 +79,18 @@ def stop(self) -> None: def _on_costmap(self, grid: OccupancyGrid) -> None: """Throttle → coarsen → colorize → PNG → map_out (under the 32 KB cap).""" + import cv2 + now = time.monotonic() if now - self._last_map_pub < 1.0 / self.config.map_hz: return - cells = grid.grid - if cells is None or cells.size == 0: - return - # A bad frame must drop, not kill the RxPY costmap subscription. try: - import cv2 - - res = grid.resolution + cells = occupancy_view(grid) + if cells.size == 0: + return + res = grid.info.resolution img_cells = cells if 0 < res < self.config.map_min_resolution: factor = max(1, round(self.config.map_min_resolution / res)) @@ -102,7 +105,7 @@ def _on_costmap(self, grid: OccupancyGrid) -> None: # origin lets the browser place map + robot: cell = (world_xy - origin)/res. h, w = img_cells.shape[:2] - origin = grid.origin.position + origin = grid.info.origin.position payload = { "type": "map", "fmt": "png", @@ -110,7 +113,7 @@ def _on_costmap(self, grid: OccupancyGrid) -> None: "h": int(h), "res": float(res), "origin": [float(origin.x), float(origin.y)], - "stamp": float(grid.ts), + "stamp": to_seconds(grid.header.stamp), "png_b64": png_b64, } data = json.dumps(payload, separators=(",", ":")).encode() @@ -132,10 +135,10 @@ def _on_odom(self, pose: PoseStamped) -> None: try: payload = { "type": "odom", - "x": float(pose.position.x), - "y": float(pose.position.y), - "yaw": float(pose.orientation.to_euler().yaw), - "ts": float(pose.ts), + "x": float(pose.pose.position.x), + "y": float(pose.pose.position.y), + "yaw": quaternion_euler(pose.pose.orientation)[2], + "ts": to_seconds(pose.header.stamp), } self.map_out.publish(json.dumps(payload).encode()) except Exception: @@ -146,7 +149,6 @@ def _on_odom(self, pose: PoseStamped) -> None: @staticmethod def _occupancy_to_bgra(cells: Any) -> Any: """Occupancy int8 {-1,0,1..100} → BGRA for PNG; unknown transparent.""" - import numpy as np # (B, G, R, A) — RGB reversed for OpenCV. c_unknown = (0, 0, 0, 0) diff --git a/dimos/teleop/hosted/test_arm_command.py b/dimos/teleop/hosted/test_arm_command.py index b95e953bb48..fa194ceed3c 100644 --- a/dimos/teleop/hosted/test_arm_command.py +++ b/dimos/teleop/hosted/test_arm_command.py @@ -30,14 +30,16 @@ from typing import Any from unittest.mock import MagicMock, patch +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist, TwistStamped, Vector3 +from dimos_generated.std_msgs.msg import Header import pytest from dimos.core.module import Module -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped +from dimos.msgs.time import time_from_seconds from dimos.teleop.hosted.arm_command import ArmCommandModule from dimos.teleop.webxr.controller_types import Hand, WebXRControllerState from dimos.utils.testing.waiting import wait_until +from dimos.web.relay_bridge.protocol import FrameHeader, encode_data_frame @pytest.fixture @@ -72,16 +74,40 @@ def _fake_init(self: Any, **kwargs: Any) -> None: module._cmd.stop() +def _frame(channel, message): + return encode_data_frame( + FrameHeader( + ch=channel, + seq=0, + ts=time.time(), + delivery="latest", + meta={"encoding": "cdr", "type": message.msg_name}, + ), + message.encode(), + ) + + def _pose_bytes(frame_id: str, ts: float | None = None) -> bytes: - return PoseStamped(ts=time.time() if ts is None else ts, frame_id=frame_id).lcm_encode() + return _frame( + "pose", + PoseStamped( + header=Header( + stamp=time_from_seconds(time.time() if ts is None else ts), frame_id=frame_id + ) + ), + ) def _twist_bytes(x: float = 0.1, angular_x: float = 0.0, ts: float | None = None) -> bytes: - # ts=None keeps TwistStamped's default stamp (now) — a fresh command. - kwargs = {} if ts is None else {"ts": ts} - return TwistStamped( - frame_id="eef_twist_arm", linear=[x, 0.0, 0.0], angular=[angular_x, 0.0, 0.0], **kwargs - ).lcm_encode() + return _frame( + "twist", + TwistStamped( + header=Header( + stamp=time_from_seconds(time.time() if ts is None else ts), frame_id="eef_twist_arm" + ), + twist=Twist(linear=Vector3(x=x), angular=Vector3(x=angular_x)), + ), + ) def _tick(module: ArmCommandModule) -> None: @@ -161,8 +187,8 @@ def test_twist_republished_without_task_address(module: ArmCommandModule) -> Non module._on_cmd_raw(_twist_bytes(0.2)) module.ee_twist_command.publish.assert_called_once() out = module.ee_twist_command.publish.call_args.args[0] - assert out.frame_id == "" # addressing is the port wiring, not the payload - assert out.linear.x == pytest.approx(0.2) + assert out.header.frame_id == "" # addressing is the port wiring, not the payload + assert out.twist.linear.x == pytest.approx(0.2) def test_ui_scale_disabled_is_rejected(module: ArmCommandModule) -> None: @@ -180,8 +206,8 @@ def test_ui_scale_updates_pose_and_keyboard_twist(module: ArmCommandModule) -> N assert _sent_acks(module) == [{"type": "cmd_ack", "nonce": 4, "ok": True}] assert module._translation_scale == 0.5 out = module.ee_twist_command.publish.call_args.args[0] - assert out.linear.x == pytest.approx(0.1) - assert out.angular.x == pytest.approx(0.3) + assert out.twist.linear.x == pytest.approx(0.1) + assert out.twist.angular.x == pytest.approx(0.3) def test_twist_dropped_while_estopped(module: ArmCommandModule) -> None: @@ -243,7 +269,7 @@ def test_controller_pose_publishes_on_hand_port(module: ArmCommandModule) -> Non assert not module._is_engaged[Hand.RIGHT] module.right_controller_output.publish.assert_called() out = module.right_controller_output.publish.call_args.args[0] - assert out.frame_id == "right" # handedness preserved; no task-name overwrite + assert out.header.frame_id == "right" # handedness preserved; no task-name overwrite module.left_controller_output.publish.assert_not_called() diff --git a/dimos/teleop/hosted/test_camera_mux.py b/dimos/teleop/hosted/test_camera_mux.py index c7a82fcfd67..35f8e811b63 100644 --- a/dimos/teleop/hosted/test_camera_mux.py +++ b/dimos/teleop/hosted/test_camera_mux.py @@ -24,10 +24,13 @@ import json from types import SimpleNamespace +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_from_array, image_view +from dimos.msgs.time import time_from_nanoseconds, to_nanoseconds from dimos.teleop.hosted.camera_mux import CameraMuxModule @@ -55,7 +58,7 @@ def _make(cameras: list[str], **cfg: object) -> _Mux: def _img(w: int, h: int) -> Image: - return Image(data=np.zeros((h, w, 3), dtype=np.uint8), format=ImageFormat.BGR) + return image_from_array(np.zeros((h, w, 3), dtype=np.uint8), encoding="bgr8") def _feed(mux: _Mux, cam: str, img: Image) -> None: @@ -64,7 +67,7 @@ def _feed(mux: _Mux, cam: str, img: Image) -> None: def _is_even(img: Image) -> bool: - h, w = img.data.shape[:2] + h, w = image_view(img).shape[:2] return h % 2 == 0 and w % 2 == 0 @@ -73,8 +76,8 @@ def _is_even(img: Image) -> bool: def test_even_dims_crops_odd_width_and_height() -> None: out = CameraMuxModule._even_dims(_img(641, 481)) - assert out.data.shape[:2] == (480, 640) - assert out.data.flags["C_CONTIGUOUS"] # from_ndarray needs contiguous + assert image_view(out).shape[:2] == (480, 640) + assert image_view(out).flags["C_CONTIGUOUS"] # from_ndarray needs contiguous def test_even_dims_passes_even_through_unchanged() -> None: @@ -126,8 +129,8 @@ def test_composite_returns_none_on_error_does_not_raise() -> None: mux = _make(["cam1", "cam2"]) with mux._cam_lock: mux._cam_selected = ["cam1", "cam2"] - _feed(mux, "cam1", Image(data=np.zeros((480, 640, 3), np.uint8), format=ImageFormat.BGR)) - _feed(mux, "cam2", Image(data=np.zeros((480, 640, 4), np.uint8), format=ImageFormat.BGRA)) + _feed(mux, "cam1", image_from_array(np.zeros((480, 640, 3), np.uint8), encoding="bgr8")) + _feed(mux, "cam2", image_from_array(np.zeros((480, 640, 4), np.uint8), encoding="bgra8")) assert mux._composite() is None @@ -185,7 +188,7 @@ def test_latency_stamp_fits_640_width() -> None: mux = _make(["cam1"], latency_stamp=True) _feed(mux, "cam1", _img(640, 480)) out = mux._composite() - assert out is not None and out.data.shape[:2] == (480 + 16, 640) + assert out is not None and image_view(out).shape[:2] == (480 + 16, 640) def test_fps_cap_drops_frames_within_window() -> None: @@ -199,9 +202,39 @@ def test_fps_cap_window_released_on_failed_composite() -> None: mux = _make(["cam1", "cam2"], video_max_fps=30.0) with mux._cam_lock: mux._cam_selected = ["cam1", "cam2"] - _feed(mux, "cam2", Image(data=np.zeros((480, 640, 4), np.uint8), format=ImageFormat.BGRA)) + _feed(mux, "cam2", image_from_array(np.zeros((480, 640, 4), np.uint8), encoding="bgra8")) mux._on_cam("cam1", _img(640, 480)) # BGR + BGRA hstack fails → None assert mux.published == [] _feed(mux, "cam2", _img(640, 480)) mux._on_cam("cam1", _img(640, 480)) # still inside the window — must publish assert len(mux.published) == 1 + + +def test_generated_crop_scale_stamp_preserve_zero_source_header() -> None: + mux = _make(["cam1"], video_max_width=641, latency_stamp=True) + source = image_from_array( + np.zeros((721, 1280, 3), np.uint8), + encoding="bgr8", + header=Header(frame_id="camera_optical"), + ) + mux._on_cam("cam1", Image.decode(source.encode())) + output = Image.decode(mux.published[0].encode()) + assert output.width % 2 == 0 and output.height % 2 == 0 + assert output.header.frame_id == "camera_optical" + assert to_nanoseconds(output.header.stamp) == 0 + assert image_view(output).shape[1] == 640 + + +def test_composite_preserves_latest_exact_stamp_without_aliasing() -> None: + mux = _make(["cam1", "cam2"]) + mux._cam_selected = ["cam1", "cam2"] + first, second = _img(320, 240), _img(320, 240) + first.header.stamp = time_from_nanoseconds(1700000000123456788) + second.header.stamp = time_from_nanoseconds(1700000000123456789) + _feed(mux, "cam1", first) + _feed(mux, "cam2", second) + output = mux._composite() + assert output is not None and output.header.frame_id == "camera_mux" + assert to_nanoseconds(output.header.stamp) == 1700000000123456789 + output.header.stamp.nanosec = 0 + assert second.header.stamp.nanosec == 123456789 diff --git a/dimos/teleop/hosted/test_go2_command.py b/dimos/teleop/hosted/test_go2_command.py index 17430fdb165..d604791ac92 100644 --- a/dimos/teleop/hosted/test_go2_command.py +++ b/dimos/teleop/hosted/test_go2_command.py @@ -29,10 +29,12 @@ from typing import Any from unittest.mock import MagicMock +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist, TwistStamped, Vector3 +from dimos_generated.std_msgs.msg import Header import pytest from dimos.core.module import Module -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped +from dimos.msgs.time import time_from_nanoseconds, time_from_seconds, to_nanoseconds from dimos.teleop.hosted.go2_command import ALLOWED_SPORT_CMDS, Go2CommandModule from dimos.utils.testing.waiting import wait_until @@ -62,8 +64,7 @@ def module(monkeypatch: pytest.MonkeyPatch) -> Iterator[Go2CommandModule]: def _twist(ts: float, *, vx: float = 0.3) -> TwistStamped: """A drive frame at time ``ts`` (vx=0.3 moving, vx=0 idle-joystick).""" - t = TwistStamped(ts=ts, linear=[0.0, 0.0, 0.0], angular=[0.0, 0.0, 0.0]) - t.linear.x = vx + t = TwistStamped(header=Header(stamp=time_from_seconds(ts)), twist=Twist(linear=Vector3(x=vx))) return t @@ -124,25 +125,25 @@ def test_drive_drops_stale(module: Go2CommandModule) -> None: def test_drive_drops_future(module: Go2CommandModule) -> None: module._on_cmd_vel_in(_twist(time.time() + 5.0)) module.tele_cmd_vel.publish.assert_not_called() - assert module._last_cmd_ts == 0.0 # future stamp must not poison the guard + assert module._last_cmd_ns == 0.0 # future stamp must not poison the guard def test_drive_drops_in_window_future_without_poisoning_guard(module: Go2CommandModule) -> None: # A future stamp SMALLER than cmd_stale_after_sec (clock skew) must still be - # rejected and must NOT advance _last_cmd_ts — otherwise every subsequent + # rejected and must NOT advance _last_cmd_ns — otherwise every subsequent # in-order frame would be dropped as out-of-order until wall-clock catches # up, stalling drive. Regression guard for the in-window future case. module._on_cmd_vel_in(_twist(time.time() + 0.2)) # +0.2s < 0.5s stale window module.tele_cmd_vel.publish.assert_not_called() - assert module._last_cmd_ts == 0.0 + assert module._last_cmd_ns == 0.0 # a normal fresh frame right after must still be forwarded module._on_cmd_vel_in(_twist(time.time())) module.tele_cmd_vel.publish.assert_called_once() def test_drive_drops_out_of_order(module: Go2CommandModule) -> None: - module._last_cmd_ts = time.time() - module._on_cmd_vel_in(_twist(module._last_cmd_ts - 0.1)) + module._last_cmd_ns = time.time_ns() + module._on_cmd_vel_in(_twist(module._last_cmd_ns / 1e9 - 0.1)) module.tele_cmd_vel.publish.assert_not_called() @@ -150,7 +151,7 @@ def test_drive_forwards_fresh(module: Go2CommandModule) -> None: ts = time.time() module._on_cmd_vel_in(_twist(ts)) module.tele_cmd_vel.publish.assert_called_once() - assert module._last_cmd_ts == ts + assert module._last_cmd_ns == to_nanoseconds(time_from_seconds(ts)) def test_drive_suppresses_idle_zero_stream(module: Go2CommandModule) -> None: @@ -179,17 +180,28 @@ def test_estopped_drive_is_dropped(module: Go2CommandModule) -> None: module.tele_cmd_vel.publish.assert_not_called() -def test_drive_drops_nan_timestamp(module: Go2CommandModule) -> None: - # A NaN ts passes every comparison and would poison _last_cmd_ts (ts <= NaN - # is False forever → reorder guard permanently disabled). Must be rejected. - module._on_cmd_vel_in(_twist(float("nan"))) +def test_drive_drops_invalid_nanoseconds(module: Go2CommandModule) -> None: + message = _twist(time.time()) + message.header.stamp.nanosec = 1_000_000_000 + module._on_cmd_vel_in(message) module.tele_cmd_vel.publish.assert_not_called() - assert module._last_cmd_ts == 0.0 # guard not poisoned + assert module._last_cmd_ns == 0 + + +def test_drive_orders_frames_one_nanosecond_apart(module: Go2CommandModule) -> None: + first = _twist(time.time() - 0.1) + stamp_ns = to_nanoseconds(first.header.stamp) + second = TwistStamped.decode(first.encode()) + second.header.stamp = time_from_nanoseconds(stamp_ns + 1) + module._on_cmd_vel_in(TwistStamped.decode(first.encode())) + module._on_cmd_vel_in(second) + assert module.tele_cmd_vel.publish.call_count == 2 + assert module._last_cmd_ns == stamp_ns + 1 def test_drive_drops_non_finite_velocity(module: Go2CommandModule) -> None: t = _twist(time.time()) - t.linear.x = float("inf") + t.twist.linear.x = float("inf") module._on_cmd_vel_in(t) module.tele_cmd_vel.publish.assert_not_called() @@ -197,8 +209,8 @@ def test_drive_drops_non_finite_velocity(module: Go2CommandModule) -> None: def test_drive_clamps_excessive_velocity(module: Go2CommandModule) -> None: # An untrusted operator sending huge velocities is clamped to the envelope. t = _twist(time.time()) - t.linear.x = 99.0 # way over max_linear_mps=1.5 - t.angular.z = -50.0 # way under -max_angular_rps=2.0 + t.twist.linear.x = 99.0 # way over max_linear_mps=1.5 + t.twist.angular.z = -50.0 # way under -max_angular_rps=2.0 module._on_cmd_vel_in(t) published = module.tele_cmd_vel.publish.call_args[0][0] assert published.linear.x == 1.5 @@ -377,7 +389,8 @@ def test_nav_goal_publishes_and_acks( module._handle_nav_goal({"x": 2.5, "y": -1.0, "nonce": 11}) (pose,) = module.goal_request.publish.call_args.args - assert pose.position.x == pytest.approx(2.5) + assert pose.pose.position.x == pytest.approx(2.5) + assert PoseStamped.decode(pose.encode()).header.frame_id == "world" assert acks == [(11, True)] diff --git a/dimos/teleop/hosted/test_hosted_stats.py b/dimos/teleop/hosted/test_hosted_stats.py index 1047a0c1493..78ec00b76d2 100644 --- a/dimos/teleop/hosted/test_hosted_stats.py +++ b/dimos/teleop/hosted/test_hosted_stats.py @@ -21,11 +21,12 @@ from types import SimpleNamespace from unittest.mock import MagicMock +from dimos_generated.geometry_msgs.msg import Twist, TwistStamped, Vector3 +from dimos_generated.std_msgs.msg import Header import pytest from dimos.core.module import Module -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.teleop.hosted.hosted_stats import HostedStatsModule @@ -43,7 +44,9 @@ def module(monkeypatch: pytest.MonkeyPatch) -> Iterator[HostedStatsModule]: def _cmd(vx: float = 0.3, ts: float = 123.0) -> bytes: - return TwistStamped(ts=ts, linear=Vector3(vx, 0, 0), angular=Vector3(0, 0, 0)).lcm_encode() + return TwistStamped( + header=Header(stamp=time_from_seconds(ts)), twist=Twist(linear=Vector3(x=vx)) + ).encode() def test_telemetry_without_go2(monkeypatch: pytest.MonkeyPatch) -> None: @@ -64,8 +67,8 @@ def test_cmd_raw_republishes_stamped_for_recorder(module: HostedStatsModule) -> module._on_cmd_raw(_cmd(vx=0.5, ts=42.0)) module.cmd_vel_stamped.publish.assert_called_once() published = module.cmd_vel_stamped.publish.call_args[0][0] - assert published.ts == 42.0 - assert published.linear.x == 0.5 + assert to_seconds(published.header.stamp) == 42.0 + assert published.twist.linear.x == 0.5 def test_cmd_raw_records_stats(module: HostedStatsModule) -> None: @@ -94,3 +97,10 @@ def test_state_json_ignores_foreign_kind(module: HostedStatsModule) -> None: # estop/sport/etc. on the shared state plane are owned by other modules. module._on_state_json(json.dumps({"type": "estop", "nonce": "x"}).encode()) module.video_stats.publish.assert_not_called() + + +def test_cmd_raw_rejects_invalid_ros_nanoseconds(module: HostedStatsModule) -> None: + message = TwistStamped.decode(_cmd()) + message.header.stamp.nanosec = 1_000_000_000 + module._on_cmd_raw(message.encode()) + module.cmd_vel_stamped.publish.assert_not_called() diff --git a/dimos/teleop/hosted/test_map_compress.py b/dimos/teleop/hosted/test_map_compress.py index d39f981ca88..3dde277789b 100644 --- a/dimos/teleop/hosted/test_map_compress.py +++ b/dimos/teleop/hosted/test_map_compress.py @@ -20,31 +20,29 @@ from collections.abc import Iterator import json import math -from types import SimpleNamespace from typing import Any from unittest.mock import MagicMock import cv2 +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest -from dimos.core.module import Module -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.time import time_from_seconds from dimos.teleop.hosted.map_compress import MapCompressModule @pytest.fixture -def module(monkeypatch: pytest.MonkeyPatch) -> Iterator[MapCompressModule]: - """A MapCompressModule with its throttle state initialized for real (only the - framework Module.__init__ is skipped) and its config / map_out port mocked.""" - monkeypatch.setattr(Module, "__init__", lambda self, **kwargs: None) +def module(mocker) -> Iterator[MapCompressModule]: module = MapCompressModule() - module.config = SimpleNamespace(map_hz=2.0, map_min_resolution=0.1, odom_hz=15.0) - module.map_out = MagicMock() - yield module + mocker.patch.object(module.map_out, "publish") + try: + yield module + finally: + module.stop() def _published_json(mock: MagicMock, msg_type: str) -> dict[str, Any] | None: @@ -60,8 +58,13 @@ def _published_json(mock: MagicMock, msg_type: str) -> dict[str, Any] | None: return None -def _occupancy(grid: Any) -> Any: - return OccupancyGrid(grid=np.asarray(grid, dtype=np.int8), resolution=0.1) +def _occupancy(grid: Any, resolution: float = 0.1) -> OccupancyGrid: + cells = np.asarray(grid, dtype=np.int8) + message = OccupancyGrid( + info=MapMetaData(width=cells.shape[1], height=cells.shape[0], resolution=resolution), + data=cells.ravel(), + ) + return OccupancyGrid.decode(message.encode()) def test_costmap_encodes_and_publishes_map(module: MapCompressModule) -> None: @@ -102,7 +105,7 @@ def test_block_max_preserves_obstacle_when_coarsening(module: MapCompressModule) # 0.02 m/cell → coarsen by 5× to reach 0.1. A lone obstacle must survive. cells = np.zeros((10, 10), dtype=np.int8) cells[3, 3] = 100 - module._on_costmap(OccupancyGrid(grid=cells, resolution=0.02)) + module._on_costmap(_occupancy(cells, resolution=0.02)) msg = _published_json(module.map_out, "map") assert msg is not None assert msg["res"] == pytest.approx(0.1) # coarsened 5× @@ -114,8 +117,11 @@ def test_block_max_preserves_obstacle_when_coarsening(module: MapCompressModule) def test_odom_publishes_planar_pose(module: MapCompressModule) -> None: - q = Quaternion.from_euler(Vector3(0.0, 0.0, math.pi / 2)) # yaw = 90° - pose = PoseStamped(ts=123.0, position=[1.5, -2.0, 0.3], orientation=[q.x, q.y, q.z, q.w]) + q = quaternion_from_euler(0.0, 0.0, math.pi / 2) # yaw = 90° + pose = PoseStamped( + header=Header(stamp=time_from_seconds(123)), + pose=Pose(position=Point(x=1.5, y=-2, z=0.3), orientation=q), + ) module._on_odom(pose) msg = _published_json(module.map_out, "odom") @@ -133,7 +139,7 @@ def test_empty_costmap_publishes_nothing(module: MapCompressModule) -> None: def test_odom_degenerate_quaternion_does_not_raise(module: MapCompressModule) -> None: # A zero quaternion makes to_euler() (scipy) raise; _on_odom runs inside an # RxPY subscriber, so it must drop the frame, not kill the odom stream. - pose = PoseStamped(ts=1.0, position=[0.0, 0.0, 0.0], orientation=[0.0, 0.0, 0.0, 0.0]) + pose = PoseStamped(pose=Pose(orientation=Quaternion(w=0))) module._on_odom(pose) # must not raise assert _published_json(module.map_out, "odom") is None @@ -144,3 +150,10 @@ def test_oversized_map_dropped(module: MapCompressModule) -> None: module._on_costmap(_occupancy(noise)) module.map_out.publish.assert_not_called() assert module._last_map_pub > 0 # throttle window consumed + + +def test_malformed_grid_does_not_break_next_frame(module): + module._on_costmap(OccupancyGrid(info=MapMetaData(width=2, height=2, resolution=0.1), data=[0])) + module.map_out.publish.assert_not_called() + module._on_costmap(_occupancy([[0, 100]])) + assert _published_json(module.map_out, "map")["w"] == 2 diff --git a/dimos/teleop/keyboard/keyboard_teleop_module.py b/dimos/teleop/keyboard/keyboard_teleop_module.py index 56e68f62ebf..360b7209409 100644 --- a/dimos/teleop/keyboard/keyboard_teleop_module.py +++ b/dimos/teleop/keyboard/keyboard_teleop_module.py @@ -43,12 +43,14 @@ if TYPE_CHECKING: from pygame.key import _ScancodeWrapper +from dimos_generated.geometry_msgs.msg import Twist, TwistStamped, Vector3 +from dimos_generated.std_msgs.msg import Float32 + from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import Out -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped -from dimos.msgs.std_msgs.Float32 import Float32 +from dimos.msgs.time import header_now from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -235,7 +237,15 @@ def _publish_twist( linear: TwistVector = (0.0, 0.0, 0.0), angular: TwistVector = (0.0, 0.0, 0.0), ) -> None: - self.ee_twist_command.publish(TwistStamped(linear=list(linear), angular=list(angular))) + self.ee_twist_command.publish( + TwistStamped( + header=header_now(), + twist=Twist( + linear=Vector3(x=linear[0], y=linear[1], z=linear[2]), + angular=Vector3(x=angular[0], y=angular[1], z=angular[2]), + ), + ) + ) def _publish_gripper_command(self, *, opening: float) -> None: """Publish a changed normalized gripper opening.""" diff --git a/dimos/teleop/keyboard/test_keyboard_teleop_module.py b/dimos/teleop/keyboard/test_keyboard_teleop_module.py index f04845006e6..1a1ba1a680c 100644 --- a/dimos/teleop/keyboard/test_keyboard_teleop_module.py +++ b/dimos/teleop/keyboard/test_keyboard_teleop_module.py @@ -16,9 +16,9 @@ from collections.abc import Iterator +from dimos_generated.geometry_msgs.msg import TwistStamped import pytest -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped import dimos.teleop.keyboard.keyboard_teleop_module as keyboard_mod from dimos.teleop.keyboard.keyboard_teleop_module import ( KeyboardTeleopModule, @@ -50,11 +50,11 @@ def test_publish_twist_emits_unaddressed_twist_stamped( module._publish_twist(linear=(0.1, 0.2, 0.3), angular=(0.4, 0.5, 0.6)) - msg = publish.call_args.args[0] + msg = TwistStamped.decode(publish.call_args.args[0].encode()) assert isinstance(msg, TwistStamped) - assert msg.frame_id == "" # no task-name address in the payload - assert [msg.linear.x, msg.linear.y, msg.linear.z] == [0.1, 0.2, 0.3] - assert [msg.angular.x, msg.angular.y, msg.angular.z] == [0.4, 0.5, 0.6] + assert msg.header.frame_id == "" # no task-name address in the payload + assert [msg.twist.linear.x, msg.twist.linear.y, msg.twist.linear.z] == [0.1, 0.2, 0.3] + assert [msg.twist.angular.x, msg.twist.angular.y, msg.twist.angular.z] == [0.4, 0.5, 0.6] def test_publish_twist_defaults_to_zero_twist(module: KeyboardTeleopModule, mocker) -> None: @@ -62,9 +62,9 @@ def test_publish_twist_defaults_to_zero_twist(module: KeyboardTeleopModule, mock module._publish_twist() - msg = publish.call_args.args[0] - assert [msg.linear.x, msg.linear.y, msg.linear.z] == [0.0, 0.0, 0.0] - assert [msg.angular.x, msg.angular.y, msg.angular.z] == [0.0, 0.0, 0.0] + msg = TwistStamped.decode(publish.call_args.args[0].encode()) + assert [msg.twist.linear.x, msg.twist.linear.y, msg.twist.linear.z] == [0.0, 0.0, 0.0] + assert [msg.twist.angular.x, msg.twist.angular.y, msg.twist.angular.z] == [0.0, 0.0, 0.0] def test_twist_from_keys_maps_translation_keys_to_eef_linear_twist() -> None: @@ -98,9 +98,9 @@ def test_final_key_release_publishes_zero_velocity(module: KeyboardTeleopModule, assert held == set() assert publish.call_count == 1 - msg = publish.call_args.args[0] - assert [msg.linear.x, msg.linear.y, msg.linear.z] == [0.0, 0.0, 0.0] - assert [msg.angular.x, msg.angular.y, msg.angular.z] == [0.0, 0.0, 0.0] + msg = TwistStamped.decode(publish.call_args.args[0].encode()) + assert [msg.twist.linear.x, msg.twist.linear.y, msg.twist.linear.z] == [0.0, 0.0, 0.0] + assert [msg.twist.angular.x, msg.twist.angular.y, msg.twist.angular.z] == [0.0, 0.0, 0.0] def test_keyup_preserves_remaining_motion_key(module: KeyboardTeleopModule, mocker) -> None: @@ -112,8 +112,8 @@ def test_keyup_preserves_remaining_motion_key(module: KeyboardTeleopModule, mock assert held == {keyboard_mod.pygame.K_a} assert publish.call_count == 1 - msg = publish.call_args.args[0] - assert [msg.linear.x, msg.linear.y, msg.linear.z] == [0.0, 0.05, 0.0] + msg = TwistStamped.decode(publish.call_args.args[0].encode()) + assert [msg.twist.linear.x, msg.twist.linear.y, msg.twist.linear.z] == [0.0, 0.05, 0.0] def test_keyup_publishes_directly_without_timeout_wait( diff --git a/dimos/teleop/phone/phone_extensions.py b/dimos/teleop/phone/phone_extensions.py index e46bab239a2..1839defb728 100644 --- a/dimos/teleop/phone/phone_extensions.py +++ b/dimos/teleop/phone/phone_extensions.py @@ -19,11 +19,10 @@ - SimplePhoneTeleop: Filters to ground robot axes and outputs cmd_vel: Out[Twist] """ +from dimos_generated.geometry_msgs.msg import Twist, TwistStamped, Vector3 + from dimos.core.core import rpc from dimos.core.stream import Out -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped -from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.teleop.phone.phone_teleop_module import PhoneTeleopModule @@ -48,7 +47,7 @@ def stop(self) -> None: def _publish_msg(self, output_msg: TwistStamped) -> None: self.cmd_vel.publish( Twist( - linear=Vector3(x=output_msg.linear.x, y=output_msg.linear.y, z=0.0), - angular=Vector3(x=0.0, y=0.0, z=output_msg.linear.z), + linear=Vector3(x=output_msg.twist.linear.x, y=output_msg.twist.linear.y, z=0.0), + angular=Vector3(x=0.0, y=0.0, z=output_msg.twist.linear.z), ) ) diff --git a/dimos/teleop/phone/phone_teleop_module.py b/dimos/teleop/phone/phone_teleop_module.py index 241e0538895..2417bcf529c 100644 --- a/dimos/teleop/phone/phone_teleop_module.py +++ b/dimos/teleop/phone/phone_teleop_module.py @@ -27,8 +27,8 @@ import time from typing import Any -from dimos_lcm.geometry_msgs import TwistStamped as LCMTwistStamped -from dimos_lcm.std_msgs import Bool as LCMBool +from dimos_generated.geometry_msgs.msg import Twist, TwistStamped, Vector3 +from dimos_generated.std_msgs.msg import Bool, Header from fastapi import WebSocket, WebSocketDisconnect from fastapi.responses import HTMLResponse from fastapi.staticfiles import StaticFiles @@ -37,11 +37,12 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import Out -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.std_msgs.Bool import Bool from dimos.utils.logging_config import setup_logger +from dimos.web.relay_bridge.protocol import ( + ProtocolError, + decode_data_frame, + peek_data_frame_lengths, +) from dimos.web.robot_web_interface import RobotWebInterface logger = setup_logger() @@ -88,11 +89,11 @@ def __init__(self, **kwargs: Any) -> None: self._web_server = RobotWebInterface(port=self.config.server_port) self._web_server_thread: threading.Thread | None = None - # Fingerprint-based message dispatch table - self._decoders: dict[bytes, Any] = { - LCMTwistStamped._get_packed_fingerprint(): self._on_sensors_bytes, - LCMBool._get_packed_fingerprint(): self._on_button_bytes, + self._command_types: dict[str, type[TwistStamped] | type[Bool]] = { + "sensors": TwistStamped, + "button": Bool, } + self._decoders = {"sensors": self._on_sensors_bytes, "button": self._on_button_bytes} self._setup_routes() @@ -104,6 +105,13 @@ async def teleop_index() -> HTMLResponse: index_path = STATIC_DIR / "index.html" return HTMLResponse(content=index_path.read_text()) + @self._web_server.app.get("/teleop/schema") + async def command_schemas() -> dict[str, dict[str, str]]: + return { + channel: {"type": message_type.msg_name, "definition": message_type.schema} + for channel, message_type in self._command_types.items() + } + if STATIC_DIR.is_dir(): self._web_server.app.mount( "/static", StaticFiles(directory=str(STATIC_DIR)), name="teleop_static" @@ -116,17 +124,33 @@ async def websocket_endpoint(ws: WebSocket) -> None: try: while True: data = await ws.receive_bytes() - fingerprint = data[:8] - decoder = self._decoders.get(fingerprint) - if decoder: - decoder(data) - else: - logger.warning(f"Unknown message fingerprint: {fingerprint.hex()}") + self._dispatch_binary_message(data) except WebSocketDisconnect: logger.info("Phone client disconnected") except Exception: logger.exception("WebSocket error") + def _dispatch_binary_message(self, data: bytes) -> bool: + """Route a complete CDR frame by explicit channel and schema identity.""" + try: + lengths = peek_data_frame_lengths(data) + if lengths is None or lengths[2] != len(data): + return False + frame = decode_data_frame(data) + message_type = self._command_types.get(frame.header.ch) + metadata = frame.header.meta or {} + if ( + message_type is None + or metadata.get("encoding") != "cdr" + or metadata.get("type") != message_type.msg_name + ): + return False + self._decoders[frame.header.ch](frame.payload) + except (ProtocolError, ValueError): + logger.warning("Dropping malformed phone command frame") + return False + return True + @rpc def start(self) -> None: super().start() @@ -156,14 +180,14 @@ def _disengage(self) -> None: logger.info("Phone teleop disengaged") def _on_sensors_bytes(self, data: bytes) -> None: - """Decode raw LCM bytes into TwistStamped and update sensor state.""" - msg = TwistStamped.lcm_decode(data) + """Decode generated CDR bytes into TwistStamped and update sensor state.""" + msg = TwistStamped.decode(data) with self._lock: self._current_sensors = msg def _on_button_bytes(self, data: bytes) -> None: - """Decode raw LCM bytes into Bool and update button state.""" - msg = Bool.lcm_decode(data) + """Decode generated CDR bytes into Bool and update button state.""" + msg = Bool.decode(data) with self._lock: self._teleop_button = bool(msg.data) @@ -250,7 +274,13 @@ def _get_output_twist(self) -> TwistStamped | None: if current is None or initial is None: return None - delta: Twist = Twist(current) - Twist(initial) + delta = Twist( + linear=Vector3( + x=current.twist.linear.x - initial.twist.linear.x, + y=current.twist.linear.y - initial.twist.linear.y, + z=current.twist.linear.z - initial.twist.linear.z, + ) + ) # Handle yaw wraparound (linear.z = yaw, 0-360 degrees) d_yaw = delta.linear.z @@ -261,17 +291,18 @@ def _get_output_twist(self) -> TwistStamped | None: cfg = self.config return TwistStamped( - ts=current.ts, - frame_id="phone", - linear=Vector3( - x=-delta.linear.y * cfg.linear_gain, # pitch forward -> drive forward - y=-delta.linear.x * cfg.linear_gain, # roll right -> strafe right - z=d_yaw * cfg.linear_gain, # yaw delta - ), - angular=Vector3( - x=current.angular.x * cfg.angular_gain, - y=current.angular.y * cfg.angular_gain, - z=current.angular.z * cfg.angular_gain, + header=Header(stamp=current.header.stamp, frame_id="phone"), + twist=Twist( + linear=Vector3( + x=-delta.linear.y * cfg.linear_gain, # pitch forward -> drive forward + y=-delta.linear.x * cfg.linear_gain, # roll right -> strafe right + z=d_yaw * cfg.linear_gain, # yaw delta + ), + angular=Vector3( + x=current.twist.angular.x * cfg.angular_gain, + y=current.twist.angular.y * cfg.angular_gain, + z=current.twist.angular.z * cfg.angular_gain, + ), ), ) diff --git a/dimos/teleop/phone/test_phone_cdr.py b/dimos/teleop/phone/test_phone_cdr.py new file mode 100644 index 00000000000..e426374ffc1 --- /dev/null +++ b/dimos/teleop/phone/test_phone_cdr.py @@ -0,0 +1,121 @@ +# Copyright 2025-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. + +"""Exercise schema advertisement and CDR command routing without physical sensors.""" + +from collections.abc import Iterator + +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Twist, TwistStamped, Vector3 +from dimos_generated.std_msgs.msg import Bool, Header +from fastapi.testclient import TestClient +import pytest + +from dimos.msgs.protocol import DimosMsg +from dimos.teleop.phone.phone_extensions import SimplePhoneTeleop +from dimos.teleop.phone.phone_teleop_module import PhoneTeleopModule +from dimos.web.relay_bridge.protocol import FrameHeader, encode_data_frame + + +def frame(channel: str, message: DimosMsg, *, type_name: str | None = None) -> bytes: + return encode_data_frame( + FrameHeader( + ch=channel, + seq=1, + ts=0, + delivery="latest", + meta={"type": type_name or message.msg_name, "encoding": "cdr"}, + ), + message.encode(), + ) + + +@pytest.fixture +def module() -> Iterator[PhoneTeleopModule]: + instance = PhoneTeleopModule() + try: + yield instance + finally: + instance.stop() + + +def test_websocket_and_schema_endpoint_use_generated_contract(module: PhoneTeleopModule) -> None: + with TestClient(module._web_server.app) as client: + schemas = client.get("/teleop/schema").json() + assert schemas["sensors"] == { + "type": TwistStamped.msg_name, + "definition": TwistStamped.schema, + } + assert schemas["button"] == {"type": Bool.msg_name, "definition": Bool.schema} + with client.websocket_connect("/ws") as ws: + ws.send_bytes(frame("button", Bool(data=True))) + ws.send_bytes(frame("sensors", TwistStamped(header=Header(frame_id="phone")))) + assert module._teleop_button is True + assert module._current_sensors.header.frame_id == "phone" + + +@pytest.mark.parametrize( + "wire", + [ + b"", + Bool(data=True).encode(), + frame("unknown", Bool(data=True)), + frame("button", Bool(data=True), type_name="std_msgs/msg/Int8"), + frame("sensors", Bool(data=True)), + frame("button", Bool(data=True)) + b"trailing", + ], +) +def test_wrong_identity_or_legacy_unframed_bytes_are_rejected( + module: PhoneTeleopModule, wire: bytes +) -> None: + assert module._dispatch_binary_message(wire) is False + assert module._teleop_button is False + assert module._current_sensors is None + + +def test_generated_control_math_preserves_exact_stamp_and_wraps_yaw( + module: PhoneTeleopModule, +) -> None: + initial = TwistStamped(twist=Twist(linear=Vector3(z=350))) + current = TwistStamped( + header=Header(stamp=Time(sec=1700000000, nanosec=123456789)), + twist=Twist(linear=Vector3(x=15, y=-30, z=10), angular=Vector3(z=60)), + ) + module._on_sensors_bytes(initial.encode()) + assert module._engage() + module._on_sensors_bytes(current.encode()) + output = module._get_output_twist() + decoded = TwistStamped.decode(output.encode()) + assert decoded.header.stamp.nanosec == 123456789 + assert decoded.header.frame_id == "phone" + assert decoded.twist.linear.x == 1.0 + assert decoded.twist.linear.y == -0.5 + assert decoded.twist.linear.z == pytest.approx(20 / 30) + assert decoded.twist.angular.z == 2.0 + + +def test_ground_robot_extension_publishes_generated_twist() -> None: + module = SimplePhoneTeleop() + outputs: list[Twist] = [] + unsubscribe = module.cmd_vel.subscribe(outputs.append) + try: + module._publish_msg(TwistStamped(twist=Twist(linear=Vector3(x=1, y=2, z=3)))) + result = Twist.decode(outputs[0].encode()) + assert result.linear.x == 1 + assert result.linear.y == 2 + assert result.linear.z == 0 + assert result.angular.z == 3 + finally: + unsubscribe() + module.stop() diff --git a/dimos/teleop/phone/web/static/index.html b/dimos/teleop/phone/web/static/index.html index 5bad9c727e9..12bdcd0c45a 100644 --- a/dimos/teleop/phone/web/static/index.html +++ b/dimos/teleop/phone/web/static/index.html @@ -165,7 +165,27 @@

Sensors

+ + diff --git a/web/sdk/fixture/demo_cdr.ts b/web/sdk/fixture/demo_cdr.ts new file mode 100644 index 00000000000..6e13812c927 --- /dev/null +++ b/web/sdk/fixture/demo_cdr.ts @@ -0,0 +1,59 @@ +import frames from "../../shared/fixtures/cdr_frames.json"; +import { createDecoderRegistry } from "@dimos/sdk"; +import type { ChannelSpec } from "@dimos/shared"; + +const registry = createDecoderRegistry(); +const messages = document.querySelector("#message")!; +const endian = document.querySelector("#endian")!; +for (const [index, vector] of frames.vectors.entries()) { + messages.add(new Option(vector.name, String(index))); +} +function render(): void { + const vector = frames.vectors[Number(messages.value)]; + const channel: ChannelSpec = { + ch: vector.name, + dir: "rx", + encoding: vector.encoding, + delivery: "reliable", + maxHz: 1, + params: { cdr: vector.schema }, + publish: "none", + requiredScope: null, + }; + const encoded = endian.value === "big" ? vector.big_endian_b64 : vector.payload_b64; + const payload = Uint8Array.from(atob(encoded), (c) => c.charCodeAt(0)); + const decoded = registry.resolve(channel)!(payload, { + ch: vector.name, + seq: 1, + ts: 0, + delivery: "reliable", + }); + document.querySelector("#status")!.textContent = + `Decoded ${payload.length} CDR bytes · ${endian.value} endian`; + document.querySelector("#type")!.textContent = vector.schema.type; + document.querySelector("#schema")!.textContent = vector.schema.definition; + document.querySelector("#value")!.textContent = JSON.stringify( + decoded.value, + (_key, value) => + typeof value === "bigint" + ? `${value}n` + : ArrayBuffer.isView(value) + ? Array.from(value as unknown as ArrayLike) + : value, + 2, + ); + const canvas = document.querySelector("#image")!; + canvas.hidden = vector.name !== "image"; + if (!canvas.hidden) { + const image = decoded.value as { data: Uint8Array }; + const rgba = new Uint8ClampedArray(8); + for (let i = 0; i < 2; i++) { + rgba.set(image.data.subarray(i * 3, i * 3 + 3), i * 4); + rgba[i * 4 + 3] = 255; + } + canvas.getContext("2d")!.putImageData(new ImageData(rgba, 2, 1), 0, 0); + } +} +messages.onchange = render; +endian.onchange = render; +render(); diff --git a/web/sdk/package.json b/web/sdk/package.json index d3521895cda..a67fa16b3e4 100644 --- a/web/sdk/package.json +++ b/web/sdk/package.json @@ -21,5 +21,9 @@ "typescript": "~5.9.2", "vite": "^7.1.0", "vitest": "^3.2.4" + }, + "dependencies": { + "@foxglove/rosmsg": "5.0.5", + "@foxglove/rosmsg2-serialization": "3.1.2" } } diff --git a/web/sdk/src/decoders/cdr.test.ts b/web/sdk/src/decoders/cdr.test.ts new file mode 100644 index 00000000000..822efef3366 --- /dev/null +++ b/web/sdk/src/decoders/cdr.test.ts @@ -0,0 +1,45 @@ +import { describe, expect, it } from "vitest"; +import frames from "../../../shared/fixtures/cdr_frames.json"; +import { cdrDecoder, cdrDecoderFor, cdrPreview } from "./cdr.ts"; +const HEADER = { ch: "message", seq: 1, ts: 0, delivery: "latest" as const }; +const bytes = (value: string) => Uint8Array.from(atob(value), (c) => c.charCodeAt(0)); +function plain(value: unknown): unknown { + if (typeof value === "bigint") return String(value); + if (ArrayBuffer.isView(value)) return Array.from(value as unknown as ArrayLike, plain); + if (Array.isArray(value)) return value.map(plain); + if (value !== null && typeof value === "object") { + return Object.fromEntries( + Object.entries(value).map(([key, item]) => [key, plain(item)]), + ); + } + return value; +} +describe("generated CDR browser decoding", () => { + for (const vector of frames.vectors) { + it(`decodes ${vector.name} in both byte orders from its embedded schema`, () => { + const decode = cdrDecoder(vector.schema); + for (const encoded of [vector.payload_b64, vector.big_endian_b64]) { + expect(plain(decode(bytes(encoded), HEADER).value)).toEqual(vector.value); + } + }); + } + it("rejects truncation, invalid encapsulation and trailing bytes", () => { + const vector = frames.vectors[0]; + const decode = cdrDecoder(vector.schema); + const payload = bytes(vector.payload_b64); + expect(() => decode(payload.subarray(0, payload.length - 1), HEADER)).toThrow(); + expect(() => decode(new Uint8Array([...payload, 0]), HEADER)).toThrow(/trailing/); + payload[0] = 255; + expect(() => decode(payload, HEADER)).toThrow(/header/); + }); + it("declines missing and malformed schemas", () => { + expect(cdrDecoderFor(undefined)).toBeNull(); + expect(cdrDecoderFor({ type: "p/msg/T", definition: "not a definition !" })).toBeNull(); + expect(cdrDecoderFor({ type: "p.T", definition: "float64 x" })).toBeNull(); + }); + it("bounds previews of large images and arrays", () => { + const text = cdrPreview({ data: new Uint8Array(1000000), label: "x".repeat(10000) }); + expect(text.length).toBeLessThan(550); + expect(text).toContain("Uint8Array(1000000)"); + }); +}); diff --git a/web/sdk/src/decoders/cdr.ts b/web/sdk/src/decoders/cdr.ts new file mode 100644 index 00000000000..10c244c780d --- /dev/null +++ b/web/sdk/src/decoders/cdr.ts @@ -0,0 +1,106 @@ +// Decode generated messages from the manifest's complete ROS2 definition. +import { parse } from "@foxglove/rosmsg"; +import { MessageReader } from "@foxglove/rosmsg2-serialization"; +import type { Decoder } from "./index.ts"; + +export interface CdrSchema { + type: string; + definition: string; +} +export const CDR_ENCODING_RE = /\.cdr\.v1$/; +export const CDR_PREVIEW_MAX_CHARS = 512; +const PREVIEW_MAX_ITEMS = 8; +const PREVIEW_MAX_STRING = 64; + +export function isCdrSchema(value: unknown): value is CdrSchema { + if (value === null || typeof value !== "object") return false; + const schema = value as Record; + return typeof schema.type === "string" && + /^[A-Za-z][A-Za-z0-9_]*\/msg\/[A-Za-z][A-Za-z0-9_]*$/.test(schema.type) && + typeof schema.definition === "string"; +} + +export function cdrDecoder(schema: CdrSchema): Decoder { + const definitions = parse(schema.definition, { ros2: true }); + const reader = new MessageReader(definitions); + return (payload) => { + // Generated codecs use encapsulated XCDR1 with no trailing bytes. + if ( + payload.length < 4 || payload[0] !== 0 || payload[1] > 1 || + payload[2] !== 0 || payload[3] !== 0 + ) throw new Error("invalid XCDR1 header"); + const value = reader.readMessage(payload); + if (reader.lastReadByteLength() !== payload.byteLength) throw new Error("trailing CDR bytes"); + return { value, preview: cdrPreview(value) }; + }; +} + +export function cdrDecoderFor(value: unknown): Decoder | null { + if (!isCdrSchema(value)) return null; + try { + return cdrDecoder(value); + } catch { + return null; + } +} + +/** Variable-length arrays require an explicit panel subscription in the cockpit. */ +export function cdrHasSequences(schema: CdrSchema): boolean { + return parse(schema.definition, { ros2: true }).some((definition) => + definition.definitions.some((field) => field.isArray && field.arrayLength === undefined) + ); +} + +export function cdrPreview(value: unknown, budget = CDR_PREVIEW_MAX_CHARS): string { + const parts: string[] = []; + let len = 0; + let full = false; + const push = (s: string) => { + parts.push(s); + len += s.length; + if (len >= budget) full = true; + }; + const items = (n: number, at: (i: number) => unknown) => { + for (let i = 0; i < n && i < PREVIEW_MAX_ITEMS && !full; i++) { + if (i > 0) push(", "); + walk(at(i)); + } + if (n > PREVIEW_MAX_ITEMS && !full) push(", ..."); + push("]"); + }; + const walk = (v: unknown): void => { + if (full) return; + if (v === null || v === undefined) { + push("null"); + } else if (typeof v === "string") { + push( + JSON.stringify(v.length > PREVIEW_MAX_STRING ? `${v.slice(0, PREVIEW_MAX_STRING)}...` : v), + ); + } else if (typeof v === "number" || typeof v === "boolean" || typeof v === "bigint") { + push(String(v)); + } else if (ArrayBuffer.isView(v)) { + const a = v as unknown as ArrayLike; + push(`${v.constructor.name}(${a.length})[`); + items(a.length, (i) => a[i]); + } else if (Array.isArray(v)) { + push("["); + items(v.length, (i) => v[i]); + } else if (typeof v === "object") { + push("{"); + let first = true; + for (const [key, item] of Object.entries(v)) { + if (full) break; + if (!first) push(", "); + first = false; + push(`${key}: `); + walk(item); + } + push("}"); + } else { + push(String(v)); + } + }; + walk(value); + const text = parts.join(""); + return full ? `${text.slice(0, budget)} ... (truncated)` : text; +} diff --git a/web/sdk/src/decoders/lcm.test.ts b/web/sdk/src/decoders/lcm.test.ts deleted file mode 100644 index e490de9577f..00000000000 --- a/web/sdk/src/decoders/lcm.test.ts +++ /dev/null @@ -1,374 +0,0 @@ -import { describe, expect, it } from "vitest"; -import type { FrameHeader } from "@dimos/shared"; -import lcmFrames from "../../../shared/fixtures/lcm_frames.json"; -import { - compileLcmDecoder, - isLcmSchema, - LCM_PREVIEW_MAX_CHARS, - lcmDecoder, - lcmDecoderFor, - LcmOversizedError, - lcmPreview, - type LcmSchema, - type LcmValue, - MAX_LCM_ARRAY_ELEMENTS, -} from "./lcm.ts"; - -const HEADER: FrameHeader = { ch: "p", seq: 1, ts: 0, delivery: "reliable" }; - -interface LcmVector { - name: string; - encoding: string; - schema: LcmSchema; - payload_b64: string; - value: unknown; -} - -// Python is the reference (dimos/web/relay_bridge/gen_lcm_fixtures.py): -// pytest re-encodes these, vitest decodes them. -const VECTORS = (lcmFrames as unknown as { vectors: LcmVector[] }).vectors; - -function b64ToBytes(b64: string): Uint8Array { - return Uint8Array.from(atob(b64), (c) => c.charCodeAt(0)); -} - -/** The fixture's JSON form of a decoded value: byte views as base64, other - * typed arrays as plain arrays, bigints as decimal strings. */ -function plain(v: unknown): unknown { - if (v instanceof Uint8Array) return btoa(String.fromCharCode(...v)); - if (typeof v === "bigint") return v.toString(); - if (v instanceof BigInt64Array) return Array.from(v, (x) => x.toString()); - if (ArrayBuffer.isView(v)) return Array.from(v as unknown as ArrayLike); - if (Array.isArray(v)) return v.map(plain); - if (typeof v === "object" && v !== null) { - return Object.fromEntries(Object.entries(v).map(([k, x]) => [k, plain(x)])); - } - return v; -} - -function vector(name: string): LcmVector { - const vec = VECTORS.find((v) => v.name === name); - if (vec === undefined) throw new Error(`no fixture vector ${name}`); - return vec; -} - -function decodeVector(name: string): { payload: Uint8Array; value: LcmValue } { - const vec = vector(name); - const payload = b64ToBytes(vec.payload_b64); - return { payload, value: compileLcmDecoder(vec.schema)(payload) }; -} - -// A hand-built schema for the error cases: x, then n strings and n floats. -const POINT: LcmSchema = { - type: "t.P", - fp: "0011223344556677", - structs: { - "t.P": [["x", "double", null], ["n", "int32_t", null], ["xs", "float", ["n"]], [ - "tags", - "string", - ["n"], - ]], - }, -}; - -/** A zeroed frame of `size` bytes that starts with the fingerprint. */ -function stamped(size: number, fp = POINT.fp): Uint8Array { - const bytes = new Uint8Array(size); - for (let i = 0; i < 8; i++) bytes[i] = parseInt(fp.slice(i * 2, i * 2 + 2), 16); - return bytes; -} - -/** A t.P frame: fingerprint, x = 1.5, n, two floats, two "a" strings. */ -function pointFrame(fp = POINT.fp, n = 2, strLen = 2): Uint8Array { - const bytes = stamped(8 + 8 + 4 + 8 + 12, fp); - const view = new DataView(bytes.buffer, 8); - view.setFloat64(0, 1.5, false); - view.setInt32(8, n, false); - view.setFloat32(12, 0.5, false); - view.setFloat32(16, -0.25, false); - for (const at of [20, 26]) { - view.setInt32(at, strLen, false); - bytes[8 + at + 4] = 0x61; // "a" - bytes[8 + at + 5] = 0; - } - return bytes; -} - -describe("lcm golden vectors", () => { - for (const vec of VECTORS) { - it(`decodes ${vec.name} (${vec.encoding})`, () => { - const decode = compileLcmDecoder(vec.schema); - expect(plain(decode(b64ToBytes(vec.payload_b64)))).toEqual(vec.value); - }); - } -}); - -describe("lcm value shape", () => { - it("keeps the LCM field names in wire order, count fields included", () => { - expect(Object.keys(decodeVector("laser_scan").value)).toEqual([ - "ranges_length", - "intensities_length", - "header", - "angle_min", - "angle_max", - "angle_increment", - "time_increment", - "scan_time", - "range_min", - "range_max", - "ranges", - "intensities", - ]); - }); - - it("views byte[] and int8_t[] into the frame, copies other arrays as typed arrays", () => { - const cloud = decodeVector("point_cloud"); - expect(cloud.value.data).toBeInstanceOf(Uint8Array); - expect((cloud.value.data as Uint8Array).buffer).toBe(cloud.payload.buffer); - expect(cloud.value.is_dense).toBe(true); - expect(cloud.value.fields).toHaveLength(3); - const grid = decodeVector("occupancy_grid"); - expect(grid.value.data).toBeInstanceOf(Int8Array); - expect((grid.value.data as Int8Array).buffer).toBe(grid.payload.buffer); - expect(Array.from(grid.value.data as Int8Array)).toEqual([0, 100, -1, 50, 0, 7]); - const scan = decodeVector("laser_scan").value; - expect(scan.ranges).toBeInstanceOf(Float32Array); - expect(scan.intensities).toBeInstanceOf(Float32Array); - expect((scan.intensities as Float32Array).length).toBe(0); - expect(decodeVector("imu").value.orientation_covariance).toBeInstanceOf(Float64Array); - expect(decodeVector("joint_state").value.name).toEqual(["shoulder", "elbow", "wrist_é"]); - }); - - it("decodes int64_t to a bigint, losslessly", () => { - expect(decodeVector("int64").value.data).toBe(-9007199254740993n); - }); -}); - -describe("lcm decoder errors", () => { - const decode = compileLcmDecoder(POINT); - - it("decodes the hand-built frame", () => { - expect(plain(decode(pointFrame()))).toEqual({ - x: 1.5, - n: 2, - xs: [0.5, -0.25], - tags: ["a", "a"], - }); - }); - - it("rejects the wrong fingerprint and a frame shorter than one", () => { - expect(() => decode(pointFrame("0011223344556678"))).toThrow(/fingerprint mismatch for t\.P/); - expect(() => decode(new Uint8Array(4))).toThrow(/shorter than the fingerprint/); - }); - - it("rejects truncated, overlong and malformed frames", () => { - expect(() => decode(pointFrame().subarray(0, 24))).toThrow(); - expect(() => decode(pointFrame().subarray(0, 8 + 8 + 4 + 8 + 4))).toThrow(/bad string length/); - const trailing = new Uint8Array(pointFrame().byteLength + 1); - trailing.set(pointFrame()); - expect(() => decode(trailing)).toThrow(/trailing bytes: 1/); - expect(() => decode(pointFrame(POINT.fp, 2, 0))).toThrow(/bad string length 0/); - }); - - it("rejects negative and oversized counts before allocating", () => { - expect(() => decode(pointFrame(POINT.fp, -1))).toThrow(/bad array length -1/); - expect(() => decode(pointFrame(POINT.fp, 1_000_000))).toThrow(/overruns the payload/); - }); - - it("refuses schemas it cannot serve at compile time", () => { - const withRoot = (rows: unknown): LcmSchema => - ({ ...POINT, structs: { "t.P": rows } }) as LcmSchema; - // A type that is neither a primitive nor a struct key reads as a missing struct. - expect(() => compileLcmDecoder(withRoot([["x", "int128_t", null]]))).toThrow( - /schema has no struct int128_t/, - ); - expect(() => compileLcmDecoder(withRoot([["h", "t.H", null]]))).toThrow( - /schema has no struct t\.H/, - ); - expect(() => compileLcmDecoder(withRoot([["m", "double", [3, 3]]]))).toThrow( - /multi-dimensional arrays are not supported \(t\.P\.m\)/, - ); - // A computed __proto__ key would set the record's prototype in browsers. - expect(() => compileLcmDecoder(withRoot([["__proto__", "double", null]]))).toThrow( - /field __proto__ is not supported \(t\.P\)/, - ); - }); - - it("rejects recursive schemas at compile time", () => { - const direct: LcmSchema = { ...POINT, structs: { "t.P": [["self", "t.P", null]] } }; - expect(() => compileLcmDecoder(direct)).toThrow(/recursive struct t\.P/); - // Through an array: n = 1 at every level would nest until the stack overflows. - const indirect: LcmSchema = { - ...POINT, - structs: { - "t.P": [["n", "int32_t", null], ["kids", "t.K", ["n"]]], - "t.K": [["parent", "t.P", null]], - }, - }; - expect(() => compileLcmDecoder(indirect)).toThrow(/recursive struct t\.P/); - expect(lcmDecoderFor(direct)).toBeNull(); - }); - - it("keeps a leading BOM, as Python's decoder does", () => { - const text = new TextEncoder().encode("\ufeffrobot"); - const bytes = stamped(8 + 4 + text.length + 1); - new DataView(bytes.buffer).setInt32(8, text.length + 1, false); - bytes.set(text, 12); - const decode = compileLcmDecoder({ ...POINT, structs: { "t.P": [["s", "string", null]] } }); - expect(decode(bytes).s).toBe("\ufeffrobot"); - }); -}); - -describe("lcm schema gate and decoder factory", () => { - it("accepts the fixture schemas and rejects the rest", () => { - for (const vec of VECTORS) expect(isLcmSchema(vec.schema)).toBe(true); - for ( - const bad of [ - null, - undefined, - "t.P", - {}, - { type: "t.P", fp: "zz", structs: {} }, - { type: "t.P", fp: POINT.fp, structs: null }, - { type: 1, fp: POINT.fp, structs: {} }, - ] - ) { - expect(isLcmSchema(bad)).toBe(false); - } - }); - - it("returns null for an unusable schema and a previewing decoder otherwise", () => { - expect(lcmDecoderFor(undefined)).toBeNull(); - expect(lcmDecoderFor({ ...POINT, structs: {} })).toBeNull(); - expect(() => lcmDecoder({ ...POINT, structs: {} })).toThrow(/schema has no struct t\.P/); - const decoder = lcmDecoderFor(POINT); - expect(decoder).not.toBeNull(); - const decoded = decoder!(pointFrame(), HEADER); - expect(plain(decoded.value)).toEqual({ x: 1.5, n: 2, xs: [0.5, -0.25], tags: ["a", "a"] }); - expect(decoded.preview).toBe( - '{x: 1.5, n: 2, xs: Float32Array(2)[0.5, -0.25], tags: ["a", "a"]}', - ); - }); -}); - -// Two struct arrays and a string array, each with its own count field. -const CLOUD: LcmSchema = { - type: "t.C", - fp: POINT.fp, - structs: { - "t.C": [ - ["n", "int32_t", null], - ["m", "int32_t", null], - ["k", "int32_t", null], - ["a", "t.X", ["n"]], - ["b", "t.X", ["m"]], - ["tags", "string", ["k"]], - ], - "t.X": [["x", "float", null]], - }, -}; - -function cloudFrame(n: number, m: number, k: number): Uint8Array { - const bytes = stamped(8 + 12 + 4 * (n + m) + 6 * k); - const view = new DataView(bytes.buffer, 8); - view.setInt32(0, n, false); - view.setInt32(4, m, false); - view.setInt32(8, k, false); - // Floats stay zero; every tag is "a". - for (let i = 0, at = 12 + 4 * (n + m); i < k; i++, at += 6) { - view.setInt32(at, 2, false); - bytes[8 + at + 4] = 0x61; - } - return bytes; -} - -// n booleans behind their count: the one array kind that is a plain JS array. -const FLAGS: LcmSchema = { - type: "t.F", - fp: POINT.fp, - structs: { "t.F": [["n", "int32_t", null], ["flags", "boolean", ["n"]]] }, -}; - -function flagsFrame(n: number): Uint8Array { - const bytes = stamped(8 + 4 + n); - new DataView(bytes.buffer).setInt32(8, n, false); - bytes.fill(1, 12); - return bytes; -} - -describe("lcm element budget", () => { - const decode = compileLcmDecoder(CLOUD); - const MAX = MAX_LCM_ARRAY_ELEMENTS; - - it("decodes a frame at the budget", () => { - expect((decode(cloudFrame(MAX, 0, 0)).a as unknown[]).length).toBe(MAX); - expect((decode(cloudFrame(0, 0, MAX)).tags as string[]).length).toBe(MAX); - }); - - it("reports a frame over the budget as oversized instead of expanding it", () => { - expect(() => decode(cloudFrame(MAX + 1, 0, 0))).toThrow(LcmOversizedError); - expect(() => decode(cloudFrame(0, 0, MAX + 1))).toThrow(/over 100000 array elements \(tags\)/); - // The budget is per frame: sibling arrays add up. - expect(() => decode(cloudFrame(MAX / 2, MAX / 2 + 1, 0))).toThrow(/\(b\)/); - const frame = cloudFrame(MAX + 1, 0, 0); - expect(lcmDecoderFor(CLOUD)!(frame, HEADER)).toEqual({ - value: undefined, - preview: `(oversized lcm message: over 100000 array elements (a), ${frame.byteLength} B)`, - }); - }); - - it("charges boolean arrays against the budget too", () => { - const flags = compileLcmDecoder(FLAGS); - const value = flags(flagsFrame(MAX)).flags as boolean[]; - expect(value).toHaveLength(MAX); - expect(value[MAX - 1]).toBe(true); - expect(() => flags(flagsFrame(MAX + 1))).toThrow(/over 100000 array elements \(flags\)/); - }); - - it("takes a page's own budget", () => { - const small = lcmDecoder(CLOUD, { maxArrayElements: 10 }); - expect((small(cloudFrame(10, 0, 0), HEADER).value as LcmValue).a).toHaveLength(10); - expect(small(cloudFrame(11, 0, 0), HEADER).preview).toMatch(/over 10 array elements \(a\)/); - const big = compileLcmDecoder(CLOUD, { maxArrayElements: MAX + 1 }); - expect(big(cloudFrame(MAX + 1, 0, 0)).a).toHaveLength(MAX + 1); - }); -}); - -describe("lcmPreview", () => { - it("prints nested fields, typed arrays and bigints", () => { - const value = { - a: 1, - s: "x", - b: true, - big: 5n, - arr: new Float64Array([1.5, 2]), - nested: { q: null }, - }; - expect(lcmPreview(value)).toBe( - '{a: 1, s: "x", b: true, big: 5, arr: Float64Array(2)[1.5, 2], nested: {q: null}}', - ); - }); - - it("shows the head of long arrays and cuts long strings", () => { - expect(lcmPreview({ r: new Float32Array(20) })).toBe( - "{r: Float32Array(20)[0, 0, 0, 0, 0, 0, 0, 0, ...]}", - ); - expect(lcmPreview(["y".repeat(100)])).toBe(`["${"y".repeat(64)}..."]`); - }); - - it("stays short on a huge array and stops at its budget on a wide message", () => { - const poses = Array.from({ length: 100_000 }, (_, i) => ({ x: i, y: -i })); - const long = lcmPreview({ poses_length: poses.length, poses }); - expect(long.startsWith("{poses_length: 100000, poses: [{x: 0, y: 0}, {x: 1, y: -1}, ")).toBe( - true, - ); - expect(long.endsWith(", ...]}")).toBe(true); - expect(long.length).toBeLessThan(200); - const wide = Object.fromEntries( - Array.from({ length: 200 }, (_, i) => [`field_${i}`, "v".repeat(20)]), - ); - const text = lcmPreview(wide); - expect(text.endsWith(" ... (truncated)")).toBe(true); - expect(text.length).toBeLessThanOrEqual(LCM_PREVIEW_MAX_CHARS + 20); - }); -}); diff --git a/web/sdk/src/decoders/lcm.ts b/web/sdk/src/decoders/lcm.ts deleted file mode 100644 index 3c926a87dba..00000000000 --- a/web/sdk/src/decoders/lcm.ts +++ /dev/null @@ -1,424 +0,0 @@ -// Schema-driven LCM decoder for `.lcm.v1` channels. The robot ships -// each channel's schema (exported from the generated dimos_lcm class) in the -// manifest record's params.lcm; compileLcmDecoder turns it into closures once -// per adopted manifest, so any DimOS message decodes with no generated code. -// Values are plain objects with the LCM fields in wire order (including the -// *_length count fields); byte[] and int8_t[] are zero-copy views into the -// frame, other primitive arrays are fresh typed arrays, int64_t is a bigint. -// Decoding runs on the ingest path, so a frame that would expand into more -// struct, string or boolean array elements than the budget (a page's own, -// else MAX_LCM_ARRAY_ELEMENTS) is reported as oversized instead (the json.v1 -// decoder's approach to a payload above its cap). - -import type { Decoder } from "./index.ts"; - -export type LcmDim = number | string; -export type LcmField = [name: string, type: string, dims: LcmDim[] | null]; -export interface LcmSchema { - type: string; - /** 16 hex chars: the big-endian packed fingerprint every frame starts with. */ - fp: string; - structs: Record; -} -export type LcmValue = Record; - -export const LCM_ENCODING_RE = /\.lcm\.v1$/; -// Struct, string and boolean array elements one frame may expand into JS -// values (typed arrays and views cost what the payload weighs and need no -// budget). A Path with 1000 poses is 1000; a page that needs more registers -// lcmDecoder(schema, { maxArrayElements }) for that encoding. -export const MAX_LCM_ARRAY_ELEMENTS = 100_000; - -export interface LcmDecoderOptions { - /** Per-frame budget of struct, string and boolean array elements. */ - maxArrayElements?: number; -} - -/** Thrown mid-decode by a compiled decoder for a frame over its element budget. */ -export class LcmOversizedError extends Error {} - -interface Budget { - max: number; - left: number; -} -// Longest preview the raw-value UI mounts for an LCM channel. -export const LCM_PREVIEW_MAX_CHARS = 512; -const PREVIEW_MAX_ITEMS = 8; -const PREVIEW_MAX_STRING = 64; - -type FieldDecoder = (view: DataView, bytes: Uint8Array, off: number, out: LcmValue) => number; - -// fatal: corrupted bytes must fail decode, not U+FFFD their way onto screen. -// ignoreBOM: a leading U+FEFF is part of the string (Python's decoder keeps it). -const utf8 = new TextDecoder("utf-8", { fatal: true, ignoreBOM: true }); -const SIZE: Record = { - int8_t: 1, - int16_t: 2, - int32_t: 4, - int64_t: 8, - float: 4, - double: 8, - boolean: 1, - byte: 1, -}; - -/** Shape check for a manifest's params.lcm; the compiler validates the rows. */ -export function isLcmSchema(x: unknown): x is LcmSchema { - if (typeof x !== "object" || x === null) return false; - const s = x as Record; - return typeof s.type === "string" && typeof s.fp === "string" && - /^[0-9a-f]{16}$/.test(s.fp) && typeof s.structs === "object" && s.structs !== null; -} - -function readString(view: DataView, bytes: Uint8Array, off: number): [string, number] { - const len = view.getInt32(off, false); // counts the trailing NUL - off += 4; - if (len < 1 || off + len > bytes.byteLength) { - throw new Error(`bad string length ${len} at ${off}`); - } - return [utf8.decode(bytes.subarray(off, off + len - 1)), off + len]; -} - -function scalar(name: string, type: string): FieldDecoder { - switch (type) { - case "double": - return (v, _b, off, out) => { - out[name] = v.getFloat64(off, false); - return off + 8; - }; - case "float": - return (v, _b, off, out) => { - out[name] = v.getFloat32(off, false); - return off + 4; - }; - case "int32_t": - return (v, _b, off, out) => { - out[name] = v.getInt32(off, false); - return off + 4; - }; - case "int16_t": - return (v, _b, off, out) => { - out[name] = v.getInt16(off, false); - return off + 2; - }; - case "int8_t": - return (v, _b, off, out) => { - out[name] = v.getInt8(off); - return off + 1; - }; - case "int64_t": - return (v, _b, off, out) => { - out[name] = v.getBigInt64(off, false); - return off + 8; - }; - case "byte": - return (v, _b, off, out) => { - out[name] = v.getUint8(off); - return off + 1; - }; - case "boolean": - return (v, _b, off, out) => { - out[name] = v.getUint8(off) !== 0; - return off + 1; - }; - case "string": - return (v, b, off, out) => { - const [text, next] = readString(v, b, off); - out[name] = text; - return next; - }; - default: - throw new Error(`unknown primitive ${type}`); - } -} - -/** Array length: the fixed size, or the value of the count field already - * decoded into the same struct (LCM puts the count before the array). */ -function count(dim: LcmDim, out: LcmValue): number { - const n = typeof dim === "number" ? dim : out[dim]; - if (typeof n !== "number" || n < 0 || !Number.isInteger(n)) { - throw new Error(`bad array length ${String(n)}`); - } - return n; -} - -function reserve(budget: Budget, name: string, n: number): void { - if (n > budget.left) { - throw new LcmOversizedError(`over ${budget.max} array elements (${name})`); - } - budget.left -= n; -} - -function checked(name: string, type: string, size: number) { - return (bytes: Uint8Array, off: number, n: number) => { - if (off + n * size > bytes.byteLength) { - throw new Error(`${name}: ${n} x ${type} overruns the payload`); - } - }; -} - -function primitiveArray(name: string, type: string, dim: LcmDim, budget: Budget): FieldDecoder { - const check = checked(name, type, SIZE[type] ?? 1); - switch (type) { - case "byte": - return (_v, b, off, out) => { - const n = count(dim, out); - check(b, off, n); - out[name] = b.subarray(off, off + n); - return off + n; - }; - case "int8_t": - return (_v, b, off, out) => { - const n = count(dim, out); - check(b, off, n); - out[name] = new Int8Array(b.buffer, b.byteOffset + off, n); - return off + n; - }; - case "boolean": - return (_v, b, off, out) => { - const n = count(dim, out); - check(b, off, n); - reserve(budget, name, n); - const a = new Array(n); - for (let i = 0; i < n; i++) a[i] = b[off + i] !== 0; - out[name] = a; - return off + n; - }; - case "float": - return (v, b, off, out) => { - const n = count(dim, out); - check(b, off, n); - const a = new Float32Array(n); - for (let i = 0; i < n; i++, off += 4) a[i] = v.getFloat32(off, false); - out[name] = a; - return off; - }; - case "double": - return (v, b, off, out) => { - const n = count(dim, out); - check(b, off, n); - const a = new Float64Array(n); - for (let i = 0; i < n; i++, off += 8) a[i] = v.getFloat64(off, false); - out[name] = a; - return off; - }; - case "int16_t": - return (v, b, off, out) => { - const n = count(dim, out); - check(b, off, n); - const a = new Int16Array(n); - for (let i = 0; i < n; i++, off += 2) a[i] = v.getInt16(off, false); - out[name] = a; - return off; - }; - case "int32_t": - return (v, b, off, out) => { - const n = count(dim, out); - check(b, off, n); - const a = new Int32Array(n); - for (let i = 0; i < n; i++, off += 4) a[i] = v.getInt32(off, false); - out[name] = a; - return off; - }; - case "int64_t": - return (v, b, off, out) => { - const n = count(dim, out); - check(b, off, n); - const a = new BigInt64Array(n); - for (let i = 0; i < n; i++, off += 8) a[i] = v.getBigInt64(off, false); - out[name] = a; - return off; - }; - case "string": - return (v, b, off, out) => { - const n = count(dim, out); - check(b, off, n); // every string is at least 5 bytes: no giant Array from a bad count - reserve(budget, name, n); - const a = new Array(n); - for (let i = 0; i < n; i++) [a[i], off] = readString(v, b, off); - out[name] = a; - return off; - }; - default: - throw new Error(`unknown primitive ${type}`); - } -} - -function struct(name: string, fields: FieldDecoder[]): FieldDecoder { - return (v, b, off, out) => { - const o: LcmValue = {}; - for (let i = 0; i < fields.length; i++) off = fields[i](v, b, off, o); - out[name] = o; - return off; - }; -} - -function structArray( - name: string, - fields: FieldDecoder[], - dim: LcmDim, - budget: Budget, -): FieldDecoder { - const check = checked(name, "struct", 1); - return (v, b, off, out) => { - const n = count(dim, out); - check(b, off, n); // a corrupt count must not allocate a giant Array first - reserve(budget, name, n); - const a = new Array(n); - for (let j = 0; j < n; j++) { - const o: LcmValue = {}; - for (let i = 0; i < fields.length; i++) off = fields[i](v, b, off, o); - a[j] = o; - } - out[name] = a; - return off; - }; -} - -/** Compile a schema into a decoder of whole frames (fingerprint included). - * Throws on a schema it cannot serve: unknown primitives, missing or - * recursive structs, multi-dimensional arrays, a field named __proto__. The - * decoder throws LcmOversizedError, before expanding the array that crosses - * it, for a frame over the element budget. */ -export function compileLcmDecoder( - schema: LcmSchema, - { maxArrayElements = MAX_LCM_ARRAY_ELEMENTS }: LcmDecoderOptions = {}, -): (payload: Uint8Array) => LcmValue { - // null: the struct is being compiled, so a reference to it is a cycle. - const compiled = new Map(); - const budget: Budget = { max: maxArrayElements, left: 0 }; - const structDecoders = (name: string): FieldDecoder[] => { - const done = compiled.get(name); - if (done) return done; - if (done === null) throw new Error(`recursive struct ${name}`); - const rows = schema.structs[name]; - if (!Array.isArray(rows)) throw new Error(`schema has no struct ${name}`); - compiled.set(name, null); - const fields: FieldDecoder[] = []; - for (const [fname, ftype, dims] of rows) { - // A computed __proto__ key sets a record's prototype in browsers. - if (fname === "__proto__") throw new Error(`field __proto__ is not supported (${name})`); - const prim = ftype === "string" || Object.hasOwn(SIZE, ftype); - if (dims === null) { - fields.push(prim ? scalar(fname, ftype) : struct(fname, structDecoders(ftype))); - } else if (dims.length === 1) { - fields.push( - prim - ? primitiveArray(fname, ftype, dims[0], budget) - : structArray(fname, structDecoders(ftype), dims[0], budget), - ); - } else { - throw new Error(`multi-dimensional arrays are not supported (${name}.${fname})`); - } - } - compiled.set(name, fields); - return fields; - }; - const root = structDecoders(schema.type); - const fp = Uint8Array.from( - { length: 8 }, - (_, i) => parseInt(schema.fp.slice(i * 2, i * 2 + 2), 16), - ); - return (payload: Uint8Array): LcmValue => { - if (payload.byteLength < 8) throw new Error("payload shorter than the fingerprint"); - for (let i = 0; i < 8; i++) { - if (payload[i] !== fp[i]) throw new Error(`fingerprint mismatch for ${schema.type}`); - } - const view = new DataView(payload.buffer, payload.byteOffset, payload.byteLength); - const out: LcmValue = {}; - budget.left = maxArrayElements; - let off = 8; - for (let i = 0; i < root.length; i++) off = root[i](view, payload, off, out); - if (off !== payload.byteLength) throw new Error(`trailing bytes: ${payload.byteLength - off}`); - return out; - }; -} - -/** Bounded text form of a decoded value for the raw channel table: the first - * elements of arrays and typed arrays, cut strings, bigint digits. Cost is - * proportional to the budget, never to the message. */ -export function lcmPreview(value: unknown, budget = LCM_PREVIEW_MAX_CHARS): string { - const parts: string[] = []; - let len = 0; - let full = false; - const push = (s: string) => { - parts.push(s); - len += s.length; - if (len >= budget) full = true; - }; - const items = (n: number, at: (i: number) => unknown) => { - for (let i = 0; i < n && i < PREVIEW_MAX_ITEMS && !full; i++) { - if (i > 0) push(", "); - walk(at(i)); - } - if (n > PREVIEW_MAX_ITEMS && !full) push(", ..."); - push("]"); - }; - const walk = (v: unknown): void => { - if (full) return; - if (v === null || v === undefined) { - push("null"); - } else if (typeof v === "string") { - push( - JSON.stringify(v.length > PREVIEW_MAX_STRING ? `${v.slice(0, PREVIEW_MAX_STRING)}...` : v), - ); - } else if (typeof v === "number" || typeof v === "boolean" || typeof v === "bigint") { - push(String(v)); - } else if (ArrayBuffer.isView(v)) { - const a = v as unknown as ArrayLike; - push(`${v.constructor.name}(${a.length})[`); - items(a.length, (i) => a[i]); - } else if (Array.isArray(v)) { - push("["); - items(v.length, (i) => v[i]); - } else if (typeof v === "object") { - push("{"); - let first = true; - for (const [key, item] of Object.entries(v)) { - if (full) break; - if (!first) push(", "); - first = false; - push(`${key}: `); - walk(item); - } - push("}"); - } else { - push(String(v)); - } - }; - walk(value); - const text = parts.join(""); - return full ? `${text.slice(0, budget)} ... (truncated)` : text; -} - -/** The Decoder for a schema (a manifest channel's params.lcm): the value and - * its preview, or an oversized preview for a frame over the budget. Throws on - * a schema compileLcmDecoder cannot serve. A page that needs a bigger budget - * registers this under the channel's encoding. */ -export function lcmDecoder(schema: LcmSchema, opts?: LcmDecoderOptions): Decoder { - const decode = compileLcmDecoder(schema, opts); - return (payload) => { - let value: LcmValue; - try { - value = decode(payload); - } catch (e) { - if (!(e instanceof LcmOversizedError)) throw e; - return { - value: undefined, - preview: `(oversized lcm message: ${e.message}, ${payload.byteLength} B)`, - }; - } - return { value, preview: lcmPreview(value) }; - }; -} - -/** lcmDecoder for a channel's params.lcm, or null when the schema is missing - * or unusable (the channel then renders as one without a decoder). */ -export function lcmDecoderFor(lcm: unknown): Decoder | null { - if (!isLcmSchema(lcm)) return null; - try { - return lcmDecoder(lcm); - } catch { - return null; - } -} diff --git a/web/sdk/src/session.test.ts b/web/sdk/src/session.test.ts index f80715177e1..0c4a8cf76a6 100644 --- a/web/sdk/src/session.test.ts +++ b/web/sdk/src/session.test.ts @@ -17,13 +17,13 @@ import { type Session, } from "./session.ts"; import { PublishError } from "./errors.ts"; -import lcmFrames from "../../shared/fixtures/lcm_frames.json"; +import cdrFrames from "../../shared/fixtures/cdr_frames.json"; // The pose_stamped golden vector (Python-generated): the schema a robot puts -// in params.lcm plus the exact frame bytes. -const POSE_LCM = (lcmFrames as { vectors: { schema: unknown; payload_b64: string }[] }) +// in params.cdr plus the exact frame bytes. +const POSE_CDR = (cdrFrames as { vectors: { schema: unknown; payload_b64: string }[] }) .vectors[0]; -const poseLcmPayload = () => Uint8Array.from(atob(POSE_LCM.payload_b64), (c) => c.charCodeAt(0)); +const poseCdrPayload = () => Uint8Array.from(atob(POSE_CDR.payload_b64), (c) => c.charCodeAt(0)); import { FakeRelayEnd, INFO, @@ -761,51 +761,51 @@ describe("Session over a fake WebTransport", () => { expect(slot.preview).toBe("2 points"); }); - it("decodes a *.lcm.v1 channel with the schema from the manifest", async () => { + it("decodes a *.cdr.v1 channel with the schema from the manifest", async () => { const { relay, handle } = start(); await goLive(relay, handle, ROBOT_A, [ spec({ - ch: "lcm_pose", - encoding: "geometry_msgs.PoseStamped.lcm.v1", - params: { lcm: POSE_LCM.schema }, + ch: "cdr_pose", + encoding: "geometry_msgs/msg/PoseStamped.cdr.v1", + params: { cdr: POSE_CDR.schema }, }), ]); - handle.subscribe("lcm_pose", () => {}); - await until(() => relay.subs().includes("lcm_pose"), "sub"); - relay.pushRaw(1, poseLcmPayload(), "lcm_pose"); - await until(() => handle.store.get("lcm_pose") !== null, "frame"); - const slot = handle.store.get("lcm_pose")!; + handle.subscribe("cdr_pose", () => {}); + await until(() => relay.subs().includes("cdr_pose"), "sub"); + relay.pushRaw(1, poseCdrPayload(), "cdr_pose"); + await until(() => handle.store.get("cdr_pose") !== null, "frame"); + const slot = handle.store.get("cdr_pose")!; const value = slot.value as { pose: { position: { x: number } }; header: { frame_id: string } }; expect(value.pose.position.x).toBe(1.5); expect(value.header.frame_id).toBe("map"); expect(slot.preview).toContain("position: {x: 1.5, y: -2.5, z: 0.25}"); - // Another fingerprint is a decode error; the last good value stays. - const foreign = poseLcmPayload(); + // An invalid encapsulation header is a decode error; the last good value stays. + const foreign = poseCdrPayload(); foreign[0] ^= 0xff; - relay.pushRaw(2, foreign, "lcm_pose"); + relay.pushRaw(2, foreign, "cdr_pose"); await until(() => { handle.store.publishUi(); - return handle.store.getUiSnapshot("lcm_pose").stats.frames === 2; + return handle.store.getUiSnapshot("cdr_pose").stats.frames === 2; }, "bad frame counted"); - expect(handle.store.get("lcm_pose")!.seq).toBe(1); - expect(handle.store.getUiSnapshot("lcm_pose").stats.decodeErrors).toBe(1); + expect(handle.store.get("cdr_pose")!.seq).toBe(1); + expect(handle.store.getUiSnapshot("cdr_pose").stats.decodeErrors).toBe(1); }); - it("counts frames of an lcm channel without a usable schema, storing no value", async () => { + it("counts frames of an cdr channel without a usable schema, storing no value", async () => { const { relay, handle } = start(); await goLive(relay, handle, ROBOT_A, [ - spec({ ch: "lcm_pose", encoding: "geometry_msgs.PoseStamped.lcm.v1", params: {} }), + spec({ ch: "cdr_pose", encoding: "geometry_msgs/msg/PoseStamped.cdr.v1", params: {} }), ]); - handle.subscribe("lcm_pose", () => {}); - await until(() => relay.subs().includes("lcm_pose"), "sub"); - relay.pushRaw(1, poseLcmPayload(), "lcm_pose"); + handle.subscribe("cdr_pose", () => {}); + await until(() => relay.subs().includes("cdr_pose"), "sub"); + relay.pushRaw(1, poseCdrPayload(), "cdr_pose"); await until(() => { handle.store.publishUi(); - return handle.store.getUiSnapshot("lcm_pose").stats.frames === 1; + return handle.store.getUiSnapshot("cdr_pose").stats.frames === 1; }, "frame counted"); - expect(handle.store.get("lcm_pose")?.value).toBeUndefined(); - expect(handle.store.getUiSnapshot("lcm_pose").stats.decodeErrors).toBe(0); + expect(handle.store.get("cdr_pose")?.value).toBeUndefined(); + expect(handle.store.getUiSnapshot("cdr_pose").stats.decodeErrors).toBe(0); }); it("retries the watch when the robot reappears after unknown_robot", async () => { diff --git a/web/shared/fixtures/cdr_frames.json b/web/shared/fixtures/cdr_frames.json new file mode 100644 index 00000000000..f3fb8e1839b --- /dev/null +++ b/web/shared/fixtures/cdr_frames.json @@ -0,0 +1,236 @@ +{ + "vectors": [ + { + "name": "pose_stamped", + "encoding": "geometry_msgs/msg/PoseStamped.cdr.v1", + "schema": { + "type": "geometry_msgs/msg/PoseStamped", + "definition": "# A Pose with reference coordinate frame and timestamp\n\nstd_msgs/Header header\nPose pose\n================================================================================\nMSG: builtin_interfaces/Time\n# This message communicates ROS Time defined here:\n# https://design.ros2.org/articles/clock_and_time.html\n\n# The seconds component, valid over all int32 values.\nint32 sec\n\n# The nanoseconds component, valid in the range [0, 1e9), to be added to the seconds component. \n# e.g.\n# The time -1.7 seconds is represented as {sec: -2, nanosec: 3e8}\n# The time 1.7 seconds is represented as {sec: 1, nanosec: 7e8}\nuint32 nanosec\n================================================================================\nMSG: geometry_msgs/Point\n# This contains the position of a point in free space\nfloat64 x\nfloat64 y\nfloat64 z\n================================================================================\nMSG: geometry_msgs/Pose\n# A representation of pose in free space, composed of position and orientation.\n\nPoint position\nQuaternion orientation\n================================================================================\nMSG: geometry_msgs/Quaternion\n# This represents an orientation in free space in quaternion form.\n\nfloat64 x 0\nfloat64 y 0\nfloat64 z 0\nfloat64 w 1\n================================================================================\nMSG: std_msgs/Header\n# Standard metadata for higher-level stamped data types.\n# This is generally used to communicate timestamped data\n# in a particular coordinate frame.\n\n# Two-integer timestamp that is expressed as seconds and nanoseconds.\nbuiltin_interfaces/Time stamp\n\n# Transform frame with which this data is associated.\nstring frame_id\n" + }, + "value": { + "header": { + "stamp": { + "sec": 1757400000, + "nanosec": 500000123 + }, + "frame_id": "map" + }, + "pose": { + "position": { + "x": 1.5, + "y": -2.5, + "z": 0.25 + }, + "orientation": { + "x": 0.0, + "y": 0.0, + "z": 0.6, + "w": 0.8 + } + } + }, + "payload_b64": "AAEAAMDLv2h7Zc0dBAAAAG1hcAAAAAAAAAD4PwAAAAAAAATAAAAAAAAA0D8AAAAAAAAAAAAAAAAAAAAAMzMzMzMz4z+amZmZmZnpPw==", + "big_endian_b64": "AAAAAGi/y8AdzWV7AAAABG1hcAA/+AAAAAAAAMAEAAAAAAAAP9AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP+MzMzMzMzM/6ZmZmZmZmg==" + }, + { + "name": "custom_segments", + "encoding": "dimos_msgs/msg/LineSegments3D.cdr.v1", + "schema": { + "type": "dimos_msgs/msg/LineSegments3D", + "definition": "# Copyright 2026 Dimensional Inc.\n# SPDX-License-Identifier: Apache-2.0\n\n# Weighted segments in a common coordinate frame.\nstd_msgs/Header header\nLineSegment3D[] segments\n================================================================================\nMSG: builtin_interfaces/Time\n# This message communicates ROS Time defined here:\n# https://design.ros2.org/articles/clock_and_time.html\n\n# The seconds component, valid over all int32 values.\nint32 sec\n\n# The nanoseconds component, valid in the range [0, 1e9), to be added to the seconds component. \n# e.g.\n# The time -1.7 seconds is represented as {sec: -2, nanosec: 3e8}\n# The time 1.7 seconds is represented as {sec: 1, nanosec: 7e8}\nuint32 nanosec\n================================================================================\nMSG: dimos_msgs/LineSegment3D\n# Copyright 2026 Dimensional Inc.\n# SPDX-License-Identifier: Apache-2.0\n\ngeometry_msgs/Point start\ngeometry_msgs/Point end\nfloat64 weight 1.0\n================================================================================\nMSG: geometry_msgs/Point\n# This contains the position of a point in free space\nfloat64 x\nfloat64 y\nfloat64 z\n================================================================================\nMSG: std_msgs/Header\n# Standard metadata for higher-level stamped data types.\n# This is generally used to communicate timestamped data\n# in a particular coordinate frame.\n\n# Two-integer timestamp that is expressed as seconds and nanoseconds.\nbuiltin_interfaces/Time stamp\n\n# Transform frame with which this data is associated.\nstring frame_id\n" + }, + "value": { + "header": { + "stamp": { + "sec": 1757400000, + "nanosec": 500000123 + }, + "frame_id": "map" + }, + "segments": [ + { + "start": { + "x": 1.0, + "y": 0.0, + "z": 0.0 + }, + "end": { + "x": 0.0, + "y": 2.0, + "z": 0.0 + }, + "weight": 4.0 + } + ] + }, + "payload_b64": "AAEAAMDLv2h7Zc0dBAAAAG1hcAABAAAAAAAAAAAAAAAAAPA/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAABBA", + "big_endian_b64": "AAAAAGi/y8AdzWV7AAAABG1hcAAAAAABAAAAAD/wAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAEAQAAAAAAAA" + }, + { + "name": "image", + "encoding": "sensor_msgs/msg/Image.cdr.v1", + "schema": { + "type": "sensor_msgs/msg/Image", + "definition": "# This message contains an uncompressed image\n# (0, 0) is at top-left corner of image\n\nstd_msgs/Header header # Header timestamp should be acquisition time of image\n # Header frame_id should be optical frame of camera\n # origin of frame should be optical center of cameara\n # +x should point to the right in the image\n # +y should point down in the image\n # +z should point into to plane of the image\n # If the frame_id here and the frame_id of the CameraInfo\n # message associated with the image conflict\n # the behavior is undefined\n\nuint32 height # image height, that is, number of rows\nuint32 width # image width, that is, number of columns\n\n# The legal values for encoding are in file include/sensor_msgs/image_encodings.hpp\n# If you want to standardize a new string format, join\n# ros-users@lists.ros.org and send an email proposing a new encoding.\n\nstring encoding # Encoding of pixels -- channel meaning, ordering, size\n # taken from the list of strings in include/sensor_msgs/image_encodings.hpp\n\nuint8 is_bigendian # is this data bigendian?\nuint32 step # Full row length in bytes\nuint8[] data # actual matrix data, size is (step * rows)\n================================================================================\nMSG: builtin_interfaces/Time\n# This message communicates ROS Time defined here:\n# https://design.ros2.org/articles/clock_and_time.html\n\n# The seconds component, valid over all int32 values.\nint32 sec\n\n# The nanoseconds component, valid in the range [0, 1e9), to be added to the seconds component. \n# e.g.\n# The time -1.7 seconds is represented as {sec: -2, nanosec: 3e8}\n# The time 1.7 seconds is represented as {sec: 1, nanosec: 7e8}\nuint32 nanosec\n================================================================================\nMSG: std_msgs/Header\n# Standard metadata for higher-level stamped data types.\n# This is generally used to communicate timestamped data\n# in a particular coordinate frame.\n\n# Two-integer timestamp that is expressed as seconds and nanoseconds.\nbuiltin_interfaces/Time stamp\n\n# Transform frame with which this data is associated.\nstring frame_id\n" + }, + "value": { + "header": { + "stamp": { + "sec": 1757400000, + "nanosec": 500000123 + }, + "frame_id": "map" + }, + "height": 1, + "width": 2, + "encoding": "rgb8", + "is_bigendian": 0, + "step": 6, + "data": [ + 0, + 1, + 2, + 253, + 254, + 255 + ] + }, + "payload_b64": "AAEAAMDLv2h7Zc0dBAAAAG1hcAABAAAAAgAAAAUAAAByZ2I4AAAAAAYAAAAGAAAAAAEC/f7/", + "big_endian_b64": "AAAAAGi/y8AdzWV7AAAABG1hcAAAAAABAAAAAgAAAAVyZ2I4AAAAAAAAAAYAAAAGAAEC/f7/" + }, + { + "name": "imu", + "encoding": "sensor_msgs/msg/Imu.cdr.v1", + "schema": { + "type": "sensor_msgs/msg/Imu", + "definition": "# This is a message to hold data from an IMU (Inertial Measurement Unit)\n#\n# Accelerations should be in m/s^2 (not in g's), and rotational velocity should be in rad/sec\n#\n# If the covariance of the measurement is known, it should be filled in (if all you know is the\n# variance of each measurement, e.g. from the datasheet, just put those along the diagonal)\n# A covariance matrix of all zeros will be interpreted as \"covariance unknown\", and to use the\n# data a covariance will have to be assumed or gotten from some other source\n#\n# If you have no estimate for one of the data elements (e.g. your IMU doesn't produce an\n# orientation estimate), please set element 0 of the associated covariance matrix to -1\n# If you are interpreting this message, please check for a value of -1 in the first element of each\n# covariance matrix, and disregard the associated estimate.\n\nstd_msgs/Header header\n\ngeometry_msgs/Quaternion orientation\nfloat64[9] orientation_covariance # Row major about x, y, z axes\n\ngeometry_msgs/Vector3 angular_velocity\nfloat64[9] angular_velocity_covariance # Row major about x, y, z axes\n\ngeometry_msgs/Vector3 linear_acceleration\nfloat64[9] linear_acceleration_covariance # Row major x, y z\n================================================================================\nMSG: builtin_interfaces/Time\n# This message communicates ROS Time defined here:\n# https://design.ros2.org/articles/clock_and_time.html\n\n# The seconds component, valid over all int32 values.\nint32 sec\n\n# The nanoseconds component, valid in the range [0, 1e9), to be added to the seconds component. \n# e.g.\n# The time -1.7 seconds is represented as {sec: -2, nanosec: 3e8}\n# The time 1.7 seconds is represented as {sec: 1, nanosec: 7e8}\nuint32 nanosec\n================================================================================\nMSG: geometry_msgs/Quaternion\n# This represents an orientation in free space in quaternion form.\n\nfloat64 x 0\nfloat64 y 0\nfloat64 z 0\nfloat64 w 1\n================================================================================\nMSG: geometry_msgs/Vector3\n# This represents a vector in free space.\n\n# This is semantically different than a point.\n# A vector is always anchored at the origin.\n# When a transform is applied to a vector, only the rotational component is applied.\n\nfloat64 x\nfloat64 y\nfloat64 z\n================================================================================\nMSG: std_msgs/Header\n# Standard metadata for higher-level stamped data types.\n# This is generally used to communicate timestamped data\n# in a particular coordinate frame.\n\n# Two-integer timestamp that is expressed as seconds and nanoseconds.\nbuiltin_interfaces/Time stamp\n\n# Transform frame with which this data is associated.\nstring frame_id\n" + }, + "value": { + "header": { + "stamp": { + "sec": 1757400000, + "nanosec": 500000123 + }, + "frame_id": "map" + }, + "orientation": { + "x": 0.0, + "y": 0.0, + "z": 0.0, + "w": 1.0 + }, + "orientation_covariance": [ + 0.0, + 1.0, + 2.0, + 3.0, + 4.0, + 5.0, + 6.0, + 7.0, + 8.0 + ], + "angular_velocity": { + "x": 0.0, + "y": 0.0, + "z": 0.0 + }, + "angular_velocity_covariance": [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ], + "linear_acceleration": { + "x": 0.0, + "y": 0.0, + "z": 0.0 + }, + "linear_acceleration_covariance": [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ] + }, + "payload_b64": "AAEAAMDLv2h7Zc0dBAAAAG1hcAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAADwPwAAAAAAAAAAAAAAAAAA8D8AAAAAAAAAQAAAAAAAAAhAAAAAAAAAEEAAAAAAAAAUQAAAAAAAABhAAAAAAAAAHEAAAAAAAAAgQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA==", + "big_endian_b64": "AAAAAGi/y8AdzWV7AAAABG1hcAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA/8AAAAAAAAAAAAAAAAAAAP/AAAAAAAABAAAAAAAAAAEAIAAAAAAAAQBAAAAAAAABAFAAAAAAAAEAYAAAAAAAAQBwAAAAAAABAIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA==" + }, + { + "name": "joints", + "encoding": "sensor_msgs/msg/JointState.cdr.v1", + "schema": { + "type": "sensor_msgs/msg/JointState", + "definition": "# This is a message that holds data to describe the state of a set of torque controlled joints.\n#\n# The state of each joint (revolute or prismatic) is defined by:\n# * the position of the joint (rad or m),\n# * the velocity of the joint (rad/s or m/s) and\n# * the effort that is applied in the joint (Nm or N).\n#\n# Each joint is uniquely identified by its name\n# The header specifies the time at which the joint states were recorded. All the joint states\n# in one message have to be recorded at the same time.\n#\n# This message consists of a multiple arrays, one for each part of the joint state.\n# The goal is to make each of the fields optional. When e.g. your joints have no\n# effort associated with them, you can leave the effort array empty.\n#\n# All arrays in this message should have the same size, or be empty.\n# This is the only way to uniquely associate the joint name with the correct\n# states.\n\nstd_msgs/Header header\n\nstring[] name\nfloat64[] position\nfloat64[] velocity\nfloat64[] effort\n================================================================================\nMSG: builtin_interfaces/Time\n# This message communicates ROS Time defined here:\n# https://design.ros2.org/articles/clock_and_time.html\n\n# The seconds component, valid over all int32 values.\nint32 sec\n\n# The nanoseconds component, valid in the range [0, 1e9), to be added to the seconds component. \n# e.g.\n# The time -1.7 seconds is represented as {sec: -2, nanosec: 3e8}\n# The time 1.7 seconds is represented as {sec: 1, nanosec: 7e8}\nuint32 nanosec\n================================================================================\nMSG: std_msgs/Header\n# Standard metadata for higher-level stamped data types.\n# This is generally used to communicate timestamped data\n# in a particular coordinate frame.\n\n# Two-integer timestamp that is expressed as seconds and nanoseconds.\nbuiltin_interfaces/Time stamp\n\n# Transform frame with which this data is associated.\nstring frame_id\n" + }, + "value": { + "header": { + "stamp": { + "sec": 1757400000, + "nanosec": 500000123 + }, + "frame_id": "map" + }, + "name": [ + "joint_a", + "joint_b" + ], + "position": [ + 1.5, + -2.5 + ], + "velocity": [], + "effort": [] + }, + "payload_b64": "AAEAAMDLv2h7Zc0dBAAAAG1hcAACAAAACAAAAGpvaW50X2EACAAAAGpvaW50X2IAAgAAAAAAAAAAAPg/AAAAAAAABMAAAAAAAAAAAA==", + "big_endian_b64": "AAAAAGi/y8AdzWV7AAAABG1hcAAAAAACAAAACGpvaW50X2EAAAAACGpvaW50X2IAAAAAAj/4AAAAAAAAwAQAAAAAAAAAAAAAAAAAAA==" + }, + { + "name": "int64", + "encoding": "std_msgs/msg/Int64.cdr.v1", + "schema": { + "type": "std_msgs/msg/Int64", + "definition": "# This was originally provided as an example message.\n# It is deprecated as of Foxy\n# It is recommended to create your own semantically meaningful message.\n# However if you would like to continue using this please use the equivalent in example_msgs.\n\nint64 data\n" + }, + "value": { + "data": "-9007199254740993" + }, + "payload_b64": "AAEAAP///////9//", + "big_endian_b64": "AAAAAP/f////////" + }, + { + "name": "empty", + "encoding": "std_msgs/msg/Empty.cdr.v1", + "schema": { + "type": "std_msgs/msg/Empty", + "definition": "\n" + }, + "value": {}, + "payload_b64": "AAEAAAA=", + "big_endian_b64": "AAAAAAA=" + }, + { + "name": "string", + "encoding": "std_msgs/msg/String.cdr.v1", + "schema": { + "type": "std_msgs/msg/String", + "definition": "# This was originally provided as an example message.\n# It is deprecated as of Foxy\n# It is recommended to create your own semantically meaningful message.\n# However if you would like to continue using this please use the equivalent in example_msgs.\n\nstring data\n" + }, + "value": { + "data": "robot 🦾 café" + }, + "payload_b64": "AAEAABEAAAByb2JvdCDwn6a+IGNhZsOpAA==", + "big_endian_b64": "AAAAAAAAABFyb2JvdCDwn6a+IGNhZsOpAA==" + } + ] +} diff --git a/web/shared/fixtures/costmap_frames.json b/web/shared/fixtures/costmap_frames.json index e9ffa3d305c..9a48edb53ee 100644 --- a/web/shared/fixtures/costmap_frames.json +++ b/web/shared/fixtures/costmap_frames.json @@ -5,7 +5,7 @@ "meta": { "w": 5, "h": 4, - "res": 0.05, + "res": 0.05000000074505806, "origin": [ -1.25, 2.5, @@ -20,7 +20,7 @@ "meta": { "w": 3, "h": 3, - "res": 0.1, + "res": 0.10000000149011612, "origin": [ 1.5, -0.75, diff --git a/web/shared/fixtures/lcm_frames.json b/web/shared/fixtures/lcm_frames.json deleted file mode 100644 index 692e42df4c0..00000000000 --- a/web/shared/fixtures/lcm_frames.json +++ /dev/null @@ -1,1114 +0,0 @@ -{ - "vectors": [ - { - "name": "pose_stamped", - "encoding": "geometry_msgs.PoseStamped.lcm.v1", - "schema": { - "type": "geometry_msgs.PoseStamped", - "fp": "6a82696458c279a0", - "structs": { - "geometry_msgs.PoseStamped": [ - [ - "header", - "std_msgs.Header", - null - ], - [ - "pose", - "geometry_msgs.Pose", - null - ] - ], - "std_msgs.Header": [ - [ - "seq", - "int32_t", - null - ], - [ - "stamp", - "std_msgs.Time", - null - ], - [ - "frame_id", - "string", - null - ] - ], - "std_msgs.Time": [ - [ - "sec", - "int32_t", - null - ], - [ - "nsec", - "int32_t", - null - ] - ], - "geometry_msgs.Pose": [ - [ - "position", - "geometry_msgs.Point", - null - ], - [ - "orientation", - "geometry_msgs.Quaternion", - null - ] - ], - "geometry_msgs.Point": [ - [ - "x", - "double", - null - ], - [ - "y", - "double", - null - ], - [ - "z", - "double", - null - ] - ], - "geometry_msgs.Quaternion": [ - [ - "x", - "double", - null - ], - [ - "y", - "double", - null - ], - [ - "z", - "double", - null - ], - [ - "w", - "double", - null - ] - ] - } - }, - "payload_b64": "aoJpZFjCeaAAAAAAaL/LwB3NZQAAAAAEbWFwAD/4AAAAAAAAwAQAAAAAAAA/0AAAAAAAAAAAAAAAAAAAAAAAAAAAAAA/4zMzMzMzMz/pmZmZmZma", - "value": { - "header": { - "seq": 0, - "stamp": { - "sec": 1757400000, - "nsec": 500000000 - }, - "frame_id": "map" - }, - "pose": { - "position": { - "x": 1.5, - "y": -2.5, - "z": 0.25 - }, - "orientation": { - "x": 0.0, - "y": 0.0, - "z": 0.6, - "w": 0.8 - } - } - } - }, - { - "name": "laser_scan", - "encoding": "sensor_msgs.LaserScan.lcm.v1", - "schema": { - "type": "sensor_msgs.LaserScan", - "fp": "8de5754e00f25f00", - "structs": { - "sensor_msgs.LaserScan": [ - [ - "ranges_length", - "int32_t", - null - ], - [ - "intensities_length", - "int32_t", - null - ], - [ - "header", - "std_msgs.Header", - null - ], - [ - "angle_min", - "float", - null - ], - [ - "angle_max", - "float", - null - ], - [ - "angle_increment", - "float", - null - ], - [ - "time_increment", - "float", - null - ], - [ - "scan_time", - "float", - null - ], - [ - "range_min", - "float", - null - ], - [ - "range_max", - "float", - null - ], - [ - "ranges", - "float", - [ - "ranges_length" - ] - ], - [ - "intensities", - "float", - [ - "intensities_length" - ] - ] - ], - "std_msgs.Header": [ - [ - "seq", - "int32_t", - null - ], - [ - "stamp", - "std_msgs.Time", - null - ], - [ - "frame_id", - "string", - null - ] - ], - "std_msgs.Time": [ - [ - "sec", - "int32_t", - null - ], - [ - "nsec", - "int32_t", - null - ] - ] - } - }, - "payload_b64": "jeV1TgDyXwAAAAAFAAAAAAAAAAdov8vAHc1lAAAAAAZsYXNlcgC/wAAAP8AAAD9AAAAAAAAAPczMzT6AAABB8AAAPwAAAD+gAABAMAAAQfAAAEFIAAA=", - "value": { - "ranges_length": 5, - "intensities_length": 0, - "header": { - "seq": 7, - "stamp": { - "sec": 1757400000, - "nsec": 500000000 - }, - "frame_id": "laser" - }, - "angle_min": -1.5, - "angle_max": 1.5, - "angle_increment": 0.75, - "time_increment": 0.0, - "scan_time": 0.10000000149011612, - "range_min": 0.25, - "range_max": 30.0, - "ranges": [ - 0.5, - 1.25, - 2.75, - 30.0, - 12.5 - ], - "intensities": [] - } - }, - { - "name": "imu", - "encoding": "sensor_msgs.Imu.lcm.v1", - "schema": { - "type": "sensor_msgs.Imu", - "fp": "31ab205c8dd57aa8", - "structs": { - "sensor_msgs.Imu": [ - [ - "header", - "std_msgs.Header", - null - ], - [ - "orientation", - "geometry_msgs.Quaternion", - null - ], - [ - "orientation_covariance", - "double", - [ - 9 - ] - ], - [ - "angular_velocity", - "geometry_msgs.Vector3", - null - ], - [ - "angular_velocity_covariance", - "double", - [ - 9 - ] - ], - [ - "linear_acceleration", - "geometry_msgs.Vector3", - null - ], - [ - "linear_acceleration_covariance", - "double", - [ - 9 - ] - ] - ], - "std_msgs.Header": [ - [ - "seq", - "int32_t", - null - ], - [ - "stamp", - "std_msgs.Time", - null - ], - [ - "frame_id", - "string", - null - ] - ], - "std_msgs.Time": [ - [ - "sec", - "int32_t", - null - ], - [ - "nsec", - "int32_t", - null - ] - ], - "geometry_msgs.Quaternion": [ - [ - "x", - "double", - null - ], - [ - "y", - "double", - null - ], - [ - "z", - "double", - null - ], - [ - "w", - "double", - null - ] - ], - "geometry_msgs.Vector3": [ - [ - "x", - "double", - null - ], - [ - "y", - "double", - null - ], - [ - "z", - "double", - null - ] - ] - } - }, - "payload_b64": "MasgXI3VeqgAAAAAaL/LwB3NZQAAAAAEaW11AAAAAAAAAAAAAAAAAAAAAAA/4zMzMzMzMz/pmZmZmZmaP4R64UeuFHs/hHrhR64Uez+EeuFHrhR7P4R64UeuFHs/hHrhR64Uez+EeuFHrhR7P4R64UeuFHs/hHrhR64Uez+EeuFHrhR7P7mZmZmZmZq/yZmZmZmZmj/TMzMzMzMzAAAAAAAAAAA/8AAAAAAAAEAAAAAAAAAAQAgAAAAAAABAEAAAAAAAAEAUAAAAAAAAQBgAAAAAAABAHAAAAAAAAEAgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAEAjAAAAAAAAP564UeuFHrg/nrhR64UeuD+euFHrhR64P564UeuFHrg/nrhR64UeuD+euFHrhR64P564UeuFHrg/nrhR64UeuD+euFHrhR64", - "value": { - "header": { - "seq": 0, - "stamp": { - "sec": 1757400000, - "nsec": 500000000 - }, - "frame_id": "imu" - }, - "orientation": { - "x": 0.0, - "y": 0.0, - "z": 0.6, - "w": 0.8 - }, - "orientation_covariance": [ - 0.01, - 0.01, - 0.01, - 0.01, - 0.01, - 0.01, - 0.01, - 0.01, - 0.01 - ], - "angular_velocity": { - "x": 0.1, - "y": -0.2, - "z": 0.3 - }, - "angular_velocity_covariance": [ - 0.0, - 1.0, - 2.0, - 3.0, - 4.0, - 5.0, - 6.0, - 7.0, - 8.0 - ], - "linear_acceleration": { - "x": 0.0, - "y": 0.0, - "z": 9.5 - }, - "linear_acceleration_covariance": [ - 0.03, - 0.03, - 0.03, - 0.03, - 0.03, - 0.03, - 0.03, - 0.03, - 0.03 - ] - } - }, - { - "name": "occupancy_grid", - "encoding": "nav_msgs.OccupancyGrid.lcm.v1", - "schema": { - "type": "nav_msgs.OccupancyGrid", - "fp": "e7dfd179cdfc3b65", - "structs": { - "nav_msgs.OccupancyGrid": [ - [ - "data_length", - "int32_t", - null - ], - [ - "header", - "std_msgs.Header", - null - ], - [ - "info", - "nav_msgs.MapMetaData", - null - ], - [ - "data", - "int8_t", - [ - "data_length" - ] - ] - ], - "std_msgs.Header": [ - [ - "seq", - "int32_t", - null - ], - [ - "stamp", - "std_msgs.Time", - null - ], - [ - "frame_id", - "string", - null - ] - ], - "std_msgs.Time": [ - [ - "sec", - "int32_t", - null - ], - [ - "nsec", - "int32_t", - null - ] - ], - "nav_msgs.MapMetaData": [ - [ - "map_load_time", - "std_msgs.Time", - null - ], - [ - "resolution", - "float", - null - ], - [ - "width", - "int32_t", - null - ], - [ - "height", - "int32_t", - null - ], - [ - "origin", - "geometry_msgs.Pose", - null - ] - ], - "geometry_msgs.Pose": [ - [ - "position", - "geometry_msgs.Point", - null - ], - [ - "orientation", - "geometry_msgs.Quaternion", - null - ] - ], - "geometry_msgs.Point": [ - [ - "x", - "double", - null - ], - [ - "y", - "double", - null - ], - [ - "z", - "double", - null - ] - ], - "geometry_msgs.Quaternion": [ - [ - "x", - "double", - null - ], - [ - "y", - "double", - null - ], - [ - "z", - "double", - null - ], - [ - "w", - "double", - null - ] - ] - } - }, - "payload_b64": "59/Rec38O2UAAAAGAAAAAGi/y8AdzWUAAAAABG1hcABov8vAHc1lAD1MzM0AAAADAAAAAsAUAAAAAAAAwBQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD/wAAAAAAAAAGT/MgAH", - "value": { - "data_length": 6, - "header": { - "seq": 0, - "stamp": { - "sec": 1757400000, - "nsec": 500000000 - }, - "frame_id": "map" - }, - "info": { - "map_load_time": { - "sec": 1757400000, - "nsec": 500000000 - }, - "resolution": 0.05000000074505806, - "width": 3, - "height": 2, - "origin": { - "position": { - "x": -5.0, - "y": -5.0, - "z": 0.0 - }, - "orientation": { - "x": 0.0, - "y": 0.0, - "z": 0.0, - "w": 1.0 - } - } - }, - "data": [ - 0, - 100, - -1, - 50, - 0, - 7 - ] - } - }, - { - "name": "joint_state", - "encoding": "sensor_msgs.JointState.lcm.v1", - "schema": { - "type": "sensor_msgs.JointState", - "fp": "ccefa8a492b967e3", - "structs": { - "sensor_msgs.JointState": [ - [ - "name_length", - "int32_t", - null - ], - [ - "position_length", - "int32_t", - null - ], - [ - "velocity_length", - "int32_t", - null - ], - [ - "effort_length", - "int32_t", - null - ], - [ - "header", - "std_msgs.Header", - null - ], - [ - "name", - "string", - [ - "name_length" - ] - ], - [ - "position", - "double", - [ - "position_length" - ] - ], - [ - "velocity", - "double", - [ - "velocity_length" - ] - ], - [ - "effort", - "double", - [ - "effort_length" - ] - ] - ], - "std_msgs.Header": [ - [ - "seq", - "int32_t", - null - ], - [ - "stamp", - "std_msgs.Time", - null - ], - [ - "frame_id", - "string", - null - ] - ], - "std_msgs.Time": [ - [ - "sec", - "int32_t", - null - ], - [ - "nsec", - "int32_t", - null - ] - ] - } - }, - "payload_b64": "zO+opJK5Z+MAAAADAAAAAwAAAAMAAAADAAAAAGi/y8AdzWUAAAAABGFybQAAAAAJc2hvdWxkZXIAAAAABmVsYm93AAAAAAl3cmlzdF/DqQA/uZmZmZmZmr/JmZmZmZmaP9MzMzMzMzM/hHrhR64Uez+UeuFHrhR7P564UeuFHrg/8AAAAAAAAEAAAAAAAAAAQAgAAAAAAAA=", - "value": { - "name_length": 3, - "position_length": 3, - "velocity_length": 3, - "effort_length": 3, - "header": { - "seq": 0, - "stamp": { - "sec": 1757400000, - "nsec": 500000000 - }, - "frame_id": "arm" - }, - "name": [ - "shoulder", - "elbow", - "wrist_\u00e9" - ], - "position": [ - 0.1, - -0.2, - 0.3 - ], - "velocity": [ - 0.01, - 0.02, - 0.03 - ], - "effort": [ - 1.0, - 2.0, - 3.0 - ] - } - }, - { - "name": "point_cloud", - "encoding": "sensor_msgs.PointCloud2.lcm.v1", - "schema": { - "type": "sensor_msgs.PointCloud2", - "fp": "f5eb3da1c2853175", - "structs": { - "sensor_msgs.PointCloud2": [ - [ - "fields_length", - "int32_t", - null - ], - [ - "data_length", - "int32_t", - null - ], - [ - "header", - "std_msgs.Header", - null - ], - [ - "height", - "int32_t", - null - ], - [ - "width", - "int32_t", - null - ], - [ - "fields", - "sensor_msgs.PointField", - [ - "fields_length" - ] - ], - [ - "is_bigendian", - "boolean", - null - ], - [ - "point_step", - "int32_t", - null - ], - [ - "row_step", - "int32_t", - null - ], - [ - "data", - "byte", - [ - "data_length" - ] - ], - [ - "is_dense", - "boolean", - null - ] - ], - "std_msgs.Header": [ - [ - "seq", - "int32_t", - null - ], - [ - "stamp", - "std_msgs.Time", - null - ], - [ - "frame_id", - "string", - null - ] - ], - "std_msgs.Time": [ - [ - "sec", - "int32_t", - null - ], - [ - "nsec", - "int32_t", - null - ] - ], - "sensor_msgs.PointField": [ - [ - "name", - "string", - null - ], - [ - "offset", - "int32_t", - null - ], - [ - "datatype", - "byte", - null - ], - [ - "count", - "int32_t", - null - ] - ] - } - }, - "payload_b64": "9es9ocKFMXUAAAADAAAAGAAAAABov8vAHc1lAAAAAAZsaWRhcgAAAAABAAAAAgAAAAJ4AAAAAAAHAAAAAQAAAAJ5AAAAAAQHAAAAAQAAAAJ6AAAAAAgHAAAAAQAAAAAMAAAAGAAAgD8AAABAAABAQAAAgMAAALBAAADIQAE=", - "value": { - "fields_length": 3, - "data_length": 24, - "header": { - "seq": 0, - "stamp": { - "sec": 1757400000, - "nsec": 500000000 - }, - "frame_id": "lidar" - }, - "height": 1, - "width": 2, - "fields": [ - { - "name": "x", - "offset": 0, - "datatype": 7, - "count": 1 - }, - { - "name": "y", - "offset": 4, - "datatype": 7, - "count": 1 - }, - { - "name": "z", - "offset": 8, - "datatype": 7, - "count": 1 - } - ], - "is_bigendian": false, - "point_step": 12, - "row_step": 24, - "data": "AACAPwAAAEAAAEBAAACAwAAAsEAAAMhA", - "is_dense": true - } - }, - { - "name": "tf_message", - "encoding": "tf2_msgs.TFMessage.lcm.v1", - "schema": { - "type": "tf2_msgs.TFMessage", - "fp": "c2b8a1c33a8923ec", - "structs": { - "tf2_msgs.TFMessage": [ - [ - "transforms_length", - "int32_t", - null - ], - [ - "transforms", - "geometry_msgs.TransformStamped", - [ - "transforms_length" - ] - ] - ], - "geometry_msgs.TransformStamped": [ - [ - "header", - "std_msgs.Header", - null - ], - [ - "child_frame_id", - "string", - null - ], - [ - "transform", - "geometry_msgs.Transform", - null - ] - ], - "std_msgs.Header": [ - [ - "seq", - "int32_t", - null - ], - [ - "stamp", - "std_msgs.Time", - null - ], - [ - "frame_id", - "string", - null - ] - ], - "std_msgs.Time": [ - [ - "sec", - "int32_t", - null - ], - [ - "nsec", - "int32_t", - null - ] - ], - "geometry_msgs.Transform": [ - [ - "translation", - "geometry_msgs.Vector3", - null - ], - [ - "rotation", - "geometry_msgs.Quaternion", - null - ] - ], - "geometry_msgs.Vector3": [ - [ - "x", - "double", - null - ], - [ - "y", - "double", - null - ], - [ - "z", - "double", - null - ] - ], - "geometry_msgs.Quaternion": [ - [ - "x", - "double", - null - ], - [ - "y", - "double", - null - ], - [ - "z", - "double", - null - ], - [ - "w", - "double", - null - ] - ] - } - }, - "payload_b64": "wrihwzqJI+wAAAACAAAAAWi/y8AdzWUAAAAABG1hcAAAAAAKYmFzZV9saW5rAD+5mZmZmZmaAAAAAAAAAAA/4AAAAAAAAAAAAAAAAAAAAAAAAAAAAAA/4zMzMzMzMz/pmZmZmZmaAAAAAWi/y8AdzWUAAAAACmJhc2VfbGluawAAAAAHY2FtZXJhAL/wAAAAAAAAQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD/wAAAAAAAA", - "value": { - "transforms_length": 2, - "transforms": [ - { - "header": { - "seq": 1, - "stamp": { - "sec": 1757400000, - "nsec": 500000000 - }, - "frame_id": "map" - }, - "child_frame_id": "base_link", - "transform": { - "translation": { - "x": 0.1, - "y": 0.0, - "z": 0.5 - }, - "rotation": { - "x": 0.0, - "y": 0.0, - "z": 0.6, - "w": 0.8 - } - } - }, - { - "header": { - "seq": 1, - "stamp": { - "sec": 1757400000, - "nsec": 500000000 - }, - "frame_id": "base_link" - }, - "child_frame_id": "camera", - "transform": { - "translation": { - "x": -1.0, - "y": 2.0, - "z": 0.0 - }, - "rotation": { - "x": 0.0, - "y": 0.0, - "z": 0.0, - "w": 1.0 - } - } - } - ] - } - }, - { - "name": "int64", - "encoding": "std_msgs.Int64.lcm.v1", - "schema": { - "type": "std_msgs.Int64", - "fp": "2cbcf9fdee91023e", - "structs": { - "std_msgs.Int64": [ - [ - "data", - "int64_t", - null - ] - ] - } - }, - "payload_b64": "LLz5/e6RAj7/3////////w==", - "value": { - "data": "-9007199254740993" - } - } - ] -} From 58e39e9a7a8607b10bbfdcde8a45d101b7191787 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Thu, 1 Oct 2026 06:22:27 +0000 Subject: [PATCH 02/11] [autofix.ci] apply automated fixes --- dimos/models/vl/test_base.py | 1 - dimos/models/vl/tool_base.py | 2 +- dimos/models/vl/tool_vlm.py | 2 +- 3 files changed, 2 insertions(+), 3 deletions(-) diff --git a/dimos/models/vl/test_base.py b/dimos/models/vl/test_base.py index 5d00e5395a0..3b07021fe07 100644 --- a/dimos/models/vl/test_base.py +++ b/dimos/models/vl/test_base.py @@ -17,7 +17,6 @@ import pytest from dimos.models.vl.qwen import QwenVlModel -from dimos_generated.sensor_msgs.msg import Image from dimos.msgs.image import image_from_file from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.utils.data import get_data diff --git a/dimos/models/vl/tool_base.py b/dimos/models/vl/tool_base.py index 2cd64850e00..f34c954b8c9 100644 --- a/dimos/models/vl/tool_base.py +++ b/dimos/models/vl/tool_base.py @@ -12,12 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.sensor_msgs.msg import Image import pytest from dimos.core.transport import LCMTransport from dimos.models.vl.moondream import MoondreamVlModel from dimos.models.vl.qwen import QwenVlModel -from dimos_generated.sensor_msgs.msg import Image from dimos.msgs.image import image_from_file from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.utils.data import get_data diff --git a/dimos/models/vl/tool_vlm.py b/dimos/models/vl/tool_vlm.py index a8b99dd445a..b1c111d1b4f 100644 --- a/dimos/models/vl/tool_vlm.py +++ b/dimos/models/vl/tool_vlm.py @@ -16,10 +16,10 @@ import time from typing import TYPE_CHECKING +from dimos_generated.sensor_msgs.msg import Image import pytest from dimos.models.vl.moondream import MoondreamVlModel -from dimos_generated.sensor_msgs.msg import Image from dimos.utils.testing.replay import TimedSensorReplay if TYPE_CHECKING: From 05d4302761edf2be8e22a089abde86be3fcff69b Mon Sep 17 00:00:00 2001 From: cc Date: Tue, 6 Oct 2026 10:24:29 -0700 Subject: [PATCH 03/11] fix: select recorded map streams and accumulate poseless world clouds --- dimos/cli/commands/map.py | 4 +- dimos/mapping/cli/replay.py | 73 +++---- dimos/mapping/cli/replay_marker.py | 24 ++- dimos/mapping/cli/streams.py | 49 +++++ dimos/mapping/cli/test_stream_selection.py | 209 +++++++++++++++++++++ 5 files changed, 320 insertions(+), 39 deletions(-) create mode 100644 dimos/mapping/cli/streams.py create mode 100644 dimos/mapping/cli/test_stream_selection.py diff --git a/dimos/cli/commands/map.py b/dimos/cli/commands/map.py index fee305185cc..ad378e4a41b 100644 --- a/dimos/cli/commands/map.py +++ b/dimos/cli/commands/map.py @@ -23,7 +23,9 @@ from dimos.mapping.cli.replay_marker import main as _map_replay_marker_main from dimos.mapping.cli.view import main as _map_view_main -map_app = typer.Typer(help="Voxel-map tools over recorded sqlite datasets") +map_app = typer.Typer( + help="Voxel-map tools over recorded datasets; rename/pose-fill require SQLite" +) map_app.command("global")(_map_main) map_app.command("rename")(_map_rename_main) map_app.command("pose-fill")(_map_pose_fill_main) diff --git a/dimos/mapping/cli/replay.py b/dimos/mapping/cli/replay.py index 6a93394cbe8..45d299eeb2d 100644 --- a/dimos/mapping/cli/replay.py +++ b/dimos/mapping/cli/replay.py @@ -17,8 +17,7 @@ Lidar clouds are assumed to be in world frame and logged directly under their entity path (no parent transform). Entities written: -- ``world/lidar`` — Go2 L1 per-frame point cloud -- ``world/fastlio_lidar`` — fastlio_lidar raw cloud (if present) +- ``world/lidar`` — the selected per-frame point cloud - ``world/_voxels`` — growing voxel map, one per PointCloud2 stream (``--map``) - ``world/_map`` — single static voxel map, one per PointCloud2 stream (``--map-final``) - ``world/fastlio`` — fastlio_odometry pose axis (if present) @@ -37,6 +36,7 @@ import typer +from dimos.mapping.cli.streams import select_stream from dimos.memory.utils.progress import progress if TYPE_CHECKING: @@ -122,7 +122,13 @@ def _log_path( def main( - dataset: str = typer.Argument(..., help="Dataset .db: bare name (cwd or data/) or path"), + dataset: str = typer.Argument(..., help="Dataset .db or .mcap: bare name or path"), + lidar_stream: str | None = typer.Option( + None, "--lidar", help="PointCloud2 stream; auto-select only a unique candidate" + ), + image_stream: str | None = typer.Option( + None, "--image", help="Image stream; auto-select a unique candidate, omit if absent" + ), out: Path | None = typer.Option( None, "--out", help="Output .rrd path (default: ./.rrd)" ), @@ -157,7 +163,7 @@ def main( map_source: list[str] = typer.Option( [], "--map-source", - help="PointCloud2 stream(s) to map; repeatable. Default: all PointCloud2 streams", + help="PointCloud2 stream(s) to map; repeatable. Default: the selected --lidar stream", ), map_carve_columns: bool = typer.Option( False, @@ -203,10 +209,13 @@ def main( cam_info = front_camera_calibration() with store: - # Resolve which streams to voxelize: all PointCloud2 streams, or the - # explicit --map-source subset. Validate up front so typos fail fast. - pc_streams = [n for n, t in stream_payload_types(store).items() if t is PointCloud2] - map_sources = list(map_source) or pc_streams + # Resolve roles and the explicit --map-source subset before opening output. + types = stream_payload_types(store) + selected_lidar = select_stream(types, PointCloud2, lidar_stream, "--lidar") + selected_image = select_stream(types, Image, image_stream, "--image", required=False) + assert selected_lidar is not None + pc_streams = [n for n, t in types.items() if t is PointCloud2] + map_sources = list(map_source) or [selected_lidar] if (map or map_final) and (bad := [s for s in map_sources if s not in pc_streams]): raise typer.BadParameter(f"--map-source: not PointCloud2 stream(s): {', '.join(bad)}") @@ -234,16 +243,10 @@ def main( def clipped(name: str, ptype: type[Any]) -> Stream[Any]: return store.stream(name, ptype).from_time(seek or None).to_time(duration) - lidar = clipped("lidar", PointCloud2) - color_image = clipped("color_image", Image) - has_livox = "fastlio_lidar" in store.streams - livox = clipped("fastlio_lidar", PointCloud2) if has_livox else None + lidar = clipped(selected_lidar, PointCloud2) # Per-frame raw clouds. _log_clouds(" lidar", lidar, "world/lidar", voxel, point_mode) - if livox is not None: - _log_clouds("fastlio_lidar", livox, "world/fastlio_lidar", voxel, point_mode) - # Accumulated voxel maps over the selected PointCloud2 streams. # --map logs a growing map per stream; --map-final logs one static map # per stream. --map-carve-columns clears the Z column under each surface @@ -336,25 +339,27 @@ def clipped(name: str, ptype: type[Any]) -> Stream[Any]: color=(0, 200, 100), # green ) - # Pass 2: camera pose + image per color_image. - cam_pipeline = ( - color_image.transform(throttle(1.0 / camera_hz)) if camera_hz > 0 else color_image - ) - n_img = cam_pipeline.count() - with progress(n_img, " color_image") as bar: - for img_obs in cam_pipeline: - bar(img_obs) - rr.set_time(TIMELINE, timestamp=img_obs.ts) - if img_obs.pose_tuple is not None: - x, y, z, qx, qy, qz, qw = img_obs.pose_tuple - rr.log( - "world/camera", - rr.Transform3D( - translation=[x, y, z], quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw]) - ), - ) - rr.log("world/camera/image", image_archetype(img_obs.data)) - + if selected_image is not None: + color_image = clipped(selected_image, Image) + # Pass 2: camera pose + image per color_image. + cam_pipeline = ( + color_image.transform(throttle(1.0 / camera_hz)) if camera_hz > 0 else color_image + ) + n_img = cam_pipeline.count() + with progress(n_img, " color_image") as bar: + for img_obs in cam_pipeline: + bar(img_obs) + rr.set_time(TIMELINE, timestamp=img_obs.ts) + if img_obs.pose_tuple is not None: + x, y, z, qx, qy, qz, qw = img_obs.pose_tuple + rr.log( + "world/camera", + rr.Transform3D( + translation=[x, y, z], + quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw]), + ), + ) + rr.log("world/camera/image", image_archetype(img_obs.data)) print(f"wrote {out}") if no_gui: print(f"open with: rerun {out}") diff --git a/dimos/mapping/cli/replay_marker.py b/dimos/mapping/cli/replay_marker.py index 3525ea7c69f..b7b10bdd235 100644 --- a/dimos/mapping/cli/replay_marker.py +++ b/dimos/mapping/cli/replay_marker.py @@ -31,6 +31,8 @@ import typer +from dimos.mapping.cli.streams import select_stream + if TYPE_CHECKING: from dimos.memory.stream import Stream @@ -39,7 +41,13 @@ def main( - dataset: str = typer.Argument(..., help="Dataset .db: bare name (cwd or data/) or path"), + dataset: str = typer.Argument(..., help="Dataset .db or .mcap: bare name or path"), + lidar_stream: str | None = typer.Option( + None, "--lidar", help="PointCloud2 stream; auto-select only a unique candidate" + ), + image_stream: str | None = typer.Option( + None, "--image", help="Image stream; auto-select only a unique candidate" + ), out: Path | None = typer.Option( None, "--out", help="Output .rrd path (default: ./.rrd)" ), @@ -66,7 +74,7 @@ def main( from dimos_generated.sensor_msgs.msg import Image, PointCloud2 import rerun as rr - from dimos.memory.cli.dataset import open_store, resolve_dataset + from dimos.memory.cli.dataset import open_store, resolve_dataset, stream_payload_types from dimos.memory.transform import QualityWindow, SpeedLimit from dimos.memory.vis.color import Color from dimos.msgs.image import image_sharpness @@ -76,6 +84,14 @@ def main( db_path = resolve_dataset(dataset) store = open_store(db_path) + try: + types = stream_payload_types(store) + selected_lidar = select_stream(types, PointCloud2, lidar_stream, "--lidar") + selected_image = select_stream(types, Image, image_stream, "--image") + assert selected_lidar is not None and selected_image is not None + except Exception: + store.stop() + raise if out is None: out = Path.cwd() / f"{db_path.stem}.rrd" cam_info = front_camera_calibration() @@ -90,8 +106,8 @@ def main( rr.log("world/camera", pinhole, static=True) with store: - color_image = store.stream("color_image", Image).from_time(seek or None).to_time(duration) - lidar = store.stream("lidar", PointCloud2).from_time(seek or None).to_time(duration) + color_image = store.stream(selected_image, Image).from_time(seek or None).to_time(duration) + lidar = store.stream(selected_lidar, PointCloud2).from_time(seek or None).to_time(duration) # Pass 1: robot base pose over time (from lidar.pose). for lidar_obs in lidar: diff --git a/dimos/mapping/cli/streams.py b/dimos/mapping/cli/streams.py new file mode 100644 index 00000000000..e25c5be214d --- /dev/null +++ b/dimos/mapping/cli/streams.py @@ -0,0 +1,49 @@ +# 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. + +"""Resolve recorded stream roles without depending on robot-specific names.""" + +from collections.abc import Mapping + +import typer + + +def select_stream( + streams: Mapping[str, type], + compatible: type | tuple[type, ...], + selected: str | None, + option: str, + *, + required: bool = True, +) -> str | None: + """Accept an explicit compatible stream or the sole compatible candidate.""" + types = compatible if isinstance(compatible, tuple) else (compatible,) + candidates = sorted(name for name, kind in streams.items() if issubclass(kind, types)) + if selected is not None: + if selected not in candidates: + raise typer.BadParameter( + f"{selected!r} is not a compatible stream. Candidates: {', '.join(candidates) or 'none'}", + param_hint=option, + ) + return selected + if len(candidates) == 1: + return candidates[0] + if len(candidates) > 1: + raise typer.BadParameter( + f"Multiple compatible streams: {', '.join(candidates)}. Select one with {option}.", + param_hint=option, + ) + if required: + raise typer.BadParameter("No compatible stream found", param_hint=option) + return None diff --git a/dimos/mapping/cli/test_stream_selection.py b/dimos/mapping/cli/test_stream_selection.py new file mode 100644 index 00000000000..b5e01b3494c --- /dev/null +++ b/dimos/mapping/cli/test_stream_selection.py @@ -0,0 +1,209 @@ +# 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. + +"""Select recorded streams by compatible type, with explicit ambiguity errors.""" + +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage +import numpy as np +import pytest +from typer.testing import CliRunner + +from dimos.cli.commands.map import map_app +from dimos.mapping.cli.streams import select_stream +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.protocol.cdr_mcap import CdrMcapWriter + + +def recording(path, clouds, images=(), frame="world", transforms=()): + cloud = pointcloud_from_xyz(np.array([[1.0, 2.0, 3.0]]), header=Header(frame_id=frame)) + image = Image( + header=Header(frame_id="camera"), + height=1, + width=1, + encoding="rgb8", + step=3, + data=[255, 0, 0], + ) + with CdrMcapWriter(path) as writer: + for name, value in ( + [(n, cloud) for n in clouds] + + [(n, image) for n in images] + + ([("tf", TFMessage(transforms=list(transforms)))] if transforms else []) + ): + writer.write( + name, + value.encode(), + schema_name=value.msg_name, + schema=value.schema, + log_time_ns=1_000_000_000, + publish_time_ns=1_000_000_000, + ) + + +def test_unique_type_and_optional_missing_stream(): + assert ( + select_stream({"front_cloud": PointCloud2, "camera": Image}, PointCloud2, None, "--lidar") + == "front_cloud" + ) + assert ( + select_stream({"front_cloud": PointCloud2}, Image, None, "--image", required=False) is None + ) + + +@pytest.mark.parametrize("command", ["global", "replay", "replay-marker"]) +def test_multiple_clouds_require_selection_even_when_one_is_named_lidar(tmp_path, command): + source, output = tmp_path / "source.mcap", tmp_path / "out.rrd" + recording(source, ["lidar", "second_cloud"], ["camera"]) + result = CliRunner().invoke(map_app, [command, str(source), "--no-gui", "--out", str(output)]) + assert result.exit_code == 2, result.output + assert "Multiple compatible streams: lidar, second_cloud" in result.output + assert "--lidar" in result.output + assert not output.exists() + + +def test_cloud_only_replay_selects_unique_arbitrary_name(tmp_path): + source, output = tmp_path / "source.mcap", tmp_path / "out.rrd" + recording(source, ["front_cloud"]) + result = CliRunner().invoke(map_app, ["replay", str(source), "--no-gui", "--out", str(output)]) + assert result.exit_code == 0, result.output + assert "[1/1]" in result.output + assert output.stat().st_size > 0 + + +def test_explicit_cloud_and_image_disambiguate_actual_replay(tmp_path): + source, output = tmp_path / "source.mcap", tmp_path / "out.rrd" + recording(source, ["front", "rear"], ["left_camera", "right_camera"]) + result = CliRunner().invoke( + map_app, + [ + "replay", + str(source), + "--lidar", + "rear", + "--image", + "right_camera", + "--no-gui", + "--out", + str(output), + ], + ) + assert result.exit_code == 0, result.output + assert output.stat().st_size > 0 + + +@pytest.mark.parametrize("selected", ["missing", "camera"]) +def test_invalid_explicit_cloud_fails_before_output(tmp_path, selected): + source, output = tmp_path / "source.mcap", tmp_path / "out.rrd" + recording(source, ["front"], ["camera"]) + result = CliRunner().invoke( + map_app, ["replay", str(source), "--lidar", selected, "--no-gui", "--out", str(output)] + ) + assert result.exit_code == 2, result.output + assert "not a compatible stream" in result.output + assert not output.exists() + + +def test_multiple_images_require_selection(tmp_path): + source, output = tmp_path / "source.mcap", tmp_path / "out.rrd" + recording(source, ["front"], ["left", "right"]) + result = CliRunner().invoke(map_app, ["replay", str(source), "--no-gui", "--out", str(output)]) + assert result.exit_code == 2, result.output + assert "Multiple compatible streams: left, right" in result.output + assert "--image" in result.output + assert not output.exists() + + +@pytest.mark.parametrize("command,extra", [("global", ["--markers"]), ("replay-marker", [])]) +def test_marker_commands_reject_ambiguous_images_before_output(tmp_path, command, extra): + source, output = tmp_path / "source.mcap", tmp_path / "out.rrd" + recording(source, ["cloud"], ["left", "right"]) + result = CliRunner().invoke( + map_app, [command, str(source), *extra, "--no-gui", "--out", str(output)] + ) + assert result.exit_code == 2, result.output + assert "Multiple compatible streams: left, right" in result.output + assert not output.exists() + + +def test_world_cloud_global_needs_no_pose(tmp_path): + source, output = tmp_path / "source.mcap", tmp_path / "out.rrd" + recording(source, ["cloud"]) + result = CliRunner().invoke( + map_app, + [ + "global", + str(source), + "--device", + "CPU:0", + "--block-count", + "100", + "--no-gui", + "--out", + str(output), + ], + ) + assert result.exit_code == 0, result.output + assert "kept [1/1]" in result.output + assert output.stat().st_size > 0 + + +@pytest.mark.parametrize("option", [["--pgo-tol", "0.3"], ["--pgo"]]) +def test_world_cloud_trajectory_options_require_pose(tmp_path, option): + source, output = tmp_path / "source.mcap", tmp_path / "out.rrd" + recording(source, ["cloud"]) + result = CliRunner().invoke( + map_app, + ["global", str(source), "--device", "CPU:0", "--no-gui", "--out", str(output), *option], + ) + assert result.exit_code == 2, result.output + assert "trajectory" in result.output + assert not output.exists() + + +@pytest.mark.parametrize("with_tf", [False, True]) +def test_sensor_cloud_requires_real_registration_tf(tmp_path, with_tf): + source, output = tmp_path / "source.mcap", tmp_path / "out.rrd" + tf = TransformStamped( + header=Header(frame_id="world"), + child_frame_id="sensor", + transform=Transform(translation=Vector3(x=2.0), rotation=Quaternion(w=1.0)), + ) + recording(source, ["cloud"], frame="sensor", transforms=[tf] if with_tf else []) + result = CliRunner().invoke( + map_app, + [ + "global", + str(source), + "--frame", + "world", + "--pgo-tol", + "0.3", + "--device", + "CPU:0", + "--block-count", + "100", + "--no-gui", + "--out", + str(output), + ], + ) + assert result.exit_code == (0 if with_tf else 2), result.output + assert output.exists() == with_tf + if with_tf: + assert "kept [1/1]" in result.output + else: + assert "cannot register" in result.output From 7ef82a893353599a742d3c0fb34cbc748809202d Mon Sep 17 00:00:00 2001 From: cc Date: Fri, 9 Oct 2026 18:54:31 -0700 Subject: [PATCH 04/11] fix(runtime): preserve region clocks and reconcile generated consumers --- dimos/control/tick_loop.py | 1 + dimos/imitation/dataprep/core.py | 1 + dimos/memory/codecs/test_json.py | 6 +++--- .../navigation/global_planner/mls_planner/utils/plan_rrd.py | 1 + .../blueprints/manipulation/test_r1pro_planar_preview.py | 2 ++ dimos/teleop/webxr/test_module.py | 1 + 6 files changed, 9 insertions(+), 3 deletions(-) diff --git a/dimos/control/tick_loop.py b/dimos/control/tick_loop.py index 7de179586db..d164513b6cb 100644 --- a/dimos/control/tick_loop.py +++ b/dimos/control/tick_loop.py @@ -42,6 +42,7 @@ from dimos.msgs.geometry_msgs.Transform import Transform from dimos.msgs.sensor_msgs.JointState import JointState from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.time import time_from_seconds from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: diff --git a/dimos/imitation/dataprep/core.py b/dimos/imitation/dataprep/core.py index 859a20deb14..5e882b63297 100644 --- a/dimos/imitation/dataprep/core.py +++ b/dimos/imitation/dataprep/core.py @@ -57,6 +57,7 @@ validate_source_kinds as validate_source_kinds, ) from dimos.memory.store.mcap import McapStore +from dimos.msgs.image import image_from_compressed, image_view from dimos.msgs.sensor_msgs.JointState import JointState from dimos.msgs.std_msgs.String import String diff --git a/dimos/memory/codecs/test_json.py b/dimos/memory/codecs/test_json.py index 57ee2c22461..2964a52aa0c 100644 --- a/dimos/memory/codecs/test_json.py +++ b/dimos/memory/codecs/test_json.py @@ -12,19 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.std_msgs.msg import String import pytest from dimos.memory.codecs.base import codec_from_id from dimos.memory.codecs.json import JsonCodec -from dimos.msgs.std_msgs.String import String @pytest.mark.parametrize( "text", ['{"label":"拿起积木 🦾","counter":9007199254740993}', "[1,2,null]", '"label"'] ) def test_json_storage_preserves_text_without_a_transport_envelope(text): - codec = codec_from_id("json", "dimos.msgs.std_msgs.String.String") - encoded = codec.encode(String(text)) + codec = codec_from_id("json", f"{String.__module__}.{String.__qualname__}") + encoded = codec.encode(String(data=text)) assert encoded == text.encode("utf-8") assert codec.decode(encoded).data == text diff --git a/dimos/navigation/global_planner/mls_planner/utils/plan_rrd.py b/dimos/navigation/global_planner/mls_planner/utils/plan_rrd.py index 97b269bdaf5..a771b546a13 100644 --- a/dimos/navigation/global_planner/mls_planner/utils/plan_rrd.py +++ b/dimos/navigation/global_planner/mls_planner/utils/plan_rrd.py @@ -40,6 +40,7 @@ from dimos.mapping.ray_tracing.viz import log_loaded_map, voxel_map_points from dimos.memory.store.sqlite import SqliteStore from dimos.memory.tf import StreamTF, tf_stream +from dimos.memory.transform import FnTransformer from dimos.memory.type.observation import Observation from dimos.memory.vis.utils import DEFAULT_RENDER_VOXEL, default_render_voxel from dimos.msgs.geometry_msgs.Transform import Transform diff --git a/dimos/robot/galaxea/r1pro/blueprints/manipulation/test_r1pro_planar_preview.py b/dimos/robot/galaxea/r1pro/blueprints/manipulation/test_r1pro_planar_preview.py index 743fd91118a..a699e65d072 100644 --- a/dimos/robot/galaxea/r1pro/blueprints/manipulation/test_r1pro_planar_preview.py +++ b/dimos/robot/galaxea/r1pro/blueprints/manipulation/test_r1pro_planar_preview.py @@ -17,6 +17,8 @@ import os import platform +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import numpy as np import pytest from dimos.control.components import HardwareType, make_twist_base_joints diff --git a/dimos/teleop/webxr/test_module.py b/dimos/teleop/webxr/test_module.py index fa611934fc8..8e49963736b 100644 --- a/dimos/teleop/webxr/test_module.py +++ b/dimos/teleop/webxr/test_module.py @@ -28,6 +28,7 @@ from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.Joy import Joy from dimos.msgs.std_msgs.String import String +from dimos.msgs.time import time_from_seconds from dimos.teleop.webxr.body_tracking import BodyTrackingSnapshot from dimos.teleop.webxr.controller_types import ( Buttons, From 4ff7b556f6aec8cfd13f000a6f987170e8a8626b Mon Sep 17 00:00:00 2001 From: cc Date: Fri, 9 Oct 2026 19:07:28 -0700 Subject: [PATCH 05/11] fix: reconcile root native builds with generated CDR packages --- .../hardware/sensors/camera/v4l2/rust/Cargo.toml | 4 ++-- .../sensors/camera/v4l2/rust/src/module.rs | 16 ++++++++-------- native/rust/dimos-module/src/pointcloud.rs | 2 +- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/dimos/hardware/sensors/camera/v4l2/rust/Cargo.toml b/dimos/hardware/sensors/camera/v4l2/rust/Cargo.toml index 6b517352683..40540bf7eac 100644 --- a/dimos/hardware/sensors/camera/v4l2/rust/Cargo.toml +++ b/dimos/hardware/sensors/camera/v4l2/rust/Cargo.toml @@ -18,11 +18,11 @@ path = "dimos_v4l2_camera::module::V4L2Camera" python = "dimos.hardware.sensors.camera.v4l2.module:V4L2Camera" [package.metadata.dimos.module.v4l2_camera.outputs] -jpeg_out = "sensor_msgs.CompressedImage" +jpeg_out = "sensor_msgs/msg/CompressedImage" [dependencies] dimos-module = { path = "../../../../../../native/rust/dimos-module" } -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +dimos-generated-messages = { path = "../../../../../../packages/dimos-generated/src/dimos_generated_schemas/package/rust" } tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal"] } serde = { version = "1", features = ["derive"] } tracing = "0.1" diff --git a/dimos/hardware/sensors/camera/v4l2/rust/src/module.rs b/dimos/hardware/sensors/camera/v4l2/rust/src/module.rs index 9bf25ae712a..0478b312df0 100644 --- a/dimos/hardware/sensors/camera/v4l2/rust/src/module.rs +++ b/dimos/hardware/sensors/camera/v4l2/rust/src/module.rs @@ -17,9 +17,11 @@ use std::time::{Duration, Instant}; -use dimos_module::{native_config, Module, Output}; -use lcm_msgs::sensor_msgs::CompressedImage; -use lcm_msgs::std_msgs::{Header, Time}; +use dimos_generated_messages::{ + builtin_interfaces::msg::time::Time, sensor_msgs::msg::compressed_image::CompressedImage, + std_msgs::msg::header::Header, +}; +use dimos_module::{cdr, native_config, Module, Output}; use tracing::{info, warn}; use crate::capture::{fourcc, Capture}; @@ -56,7 +58,7 @@ pub struct Config { #[derive(Module)] #[module(name = "v4l2_camera", setup = start)] pub struct V4L2Camera { - #[output(encode = CompressedImage::encode)] + #[output(encode = cdr::encode)] jpeg_out: Output, #[config] @@ -115,7 +117,7 @@ fn pump( ) { let mut clock = CaptureClock::new(driver_clocks(), wall_s); let (width, height) = (config.width as usize, config.height as usize); - let (mut timeouts, mut seq) = (0, 0i32); + let mut timeouts = 0; let (mut frames, mut encode_s, mut window) = (0u32, 0.0, Instant::now()); loop { let frame = match capture.next(FRAME_TIMEOUT_MS) { @@ -154,17 +156,15 @@ fn pump( }; let message = CompressedImage { header: Header { - seq, stamp: Time { sec: stamp.trunc() as i32, - nsec: (stamp.fract() * 1e9) as i32, + nanosec: (stamp.fract() * 1e9) as u32, }, frame_id: config.frame_id.clone(), }, format: "jpeg".into(), data, }; - seq = seq.wrapping_add(1); runtime.block_on(output.publish(&message)).ok(); frames += 1; if window.elapsed() >= STATS_EVERY { diff --git a/native/rust/dimos-module/src/pointcloud.rs b/native/rust/dimos-module/src/pointcloud.rs index 31b8ef0aee6..58a07e04fb2 100644 --- a/native/rust/dimos-module/src/pointcloud.rs +++ b/native/rust/dimos-module/src/pointcloud.rs @@ -113,7 +113,7 @@ mod tests { let field = |name: &str, off: i32| PointField { name: name.into(), offset: off, - datatype: PointField::FLOAT32 as u8, + datatype: PointField::FLOAT32, count: 1, }; PointCloud2 { From 0b42934c7874feaf88098533385386423ab0512b Mon Sep 17 00:00:00 2001 From: cc Date: Fri, 9 Oct 2026 19:36:44 -0700 Subject: [PATCH 06/11] fix(cdr): reconcile current-main consumers and native modules --- dimos/agents/skills/point_nav.py | 27 +-- dimos/core/native_module.py | 19 ++- .../hardware/sensors/lidar/pointlio/module.py | 5 +- dimos/mapping/pointclouds/occupancy.py | 1 + dimos/mapping/ray_tracing/utils/loaded_map.py | 9 +- dimos/mapping/relocalization/lidar/replay.py | 59 ++++--- dimos/memory/type/recording.py | 9 +- dimos/navigation/go2/loop_closure/module.py | 28 ++-- dimos/navigation/go2/loop_closure/pgo_map.py | 29 +++- dimos/navigation/local_planner/module.py | 77 ++++++--- dimos/navigation/local_planner/native.py | 7 +- dimos/navigation/local_planner/profile.py | 16 +- .../navigation/local_planner/rust/Cargo.lock | 111 +++++++++++-- .../navigation/local_planner/rust/Cargo.toml | 12 +- .../local_planner/rust/src/module/msg.rs | 34 ++-- .../local_planner/rust/src/module/planner.rs | 11 +- .../local_planner/rust/src/module/tf_pose.rs | 14 +- dimos/navigation/local_planner/search/base.py | 21 +-- .../navigation/local_planner/search/target.py | 41 +++-- dimos/navigation/local_planner/test_module.py | 123 ++++++++++---- .../navigation/local_planner/test_profile.py | 45 ++--- dimos/navigation/local_planner/test_viz.py | 48 ++++-- dimos/navigation/local_planner/viz.py | 25 +-- dimos/navigation/test_tf_pose.py | 26 +-- dimos/navigation/tf_pose.py | 12 +- .../trajectory_follower/fancy/controller.py | 9 +- .../trajectory_follower/fancy/laws/hinted.py | 24 +-- .../trajectory_follower/fancy/laws/seed.py | 21 ++- .../trajectory_follower/fancy/module.py | 43 +++-- .../trajectory_follower/fancy/native.py | 9 +- .../trajectory_follower/fancy/rust/Cargo.lock | 113 +++++++++++-- .../trajectory_follower/fancy/rust/Cargo.toml | 12 +- .../fancy/rust/src/module/follower.rs | 13 +- .../fancy/test_controller.py | 33 ++-- .../trajectory_follower/fancy/test_module.py | 72 ++++++-- .../fancy/test_rust_parity.py | 32 ++-- dimos/perception/depth2depth_cloud/module.py | 6 +- .../depth2depth_cloud/rust/Cargo.lock | 97 +++++++++-- .../depth2depth_cloud/rust/Cargo.toml | 12 +- .../depth2depth_cloud/rust/flake.nix | 3 + .../depth2depth_cloud/rust/src/module.rs | 43 +++-- .../depth2depth_cloud/rust/src/undistort.rs | 40 +++-- dimos/perception/localize/support_plane.py | 8 +- dimos/perception/localize/test_rig.py | 63 +++++++ dimos/perception/localize/types.py | 2 +- dimos/robot/unitree/go2/base.py | 17 +- dimos/robot/unitree/go2/dds/module.py | 15 +- dimos/robot/unitree/go2/dds/rust/Cargo.lock | 108 +++++++++--- dimos/robot/unitree/go2/dds/rust/Cargo.toml | 30 ++-- .../robot/unitree/go2/dds/rust/src/module.rs | 157 ++++++------------ .../unitree/go2/dds/rust/tests/config.json | 25 ++- .../dds/rust/vendor/cyclonedds-sys/Cargo.toml | 4 +- .../dds/rust/vendor/cyclonedds-sys/build.rs | 3 - dimos/robot/unitree/mujoco_connection.py | 5 +- .../engines/test_mujoco_sim_module.py | 13 +- 55 files changed, 1238 insertions(+), 603 deletions(-) create mode 100644 dimos/perception/localize/test_rig.py diff --git a/dimos/agents/skills/point_nav.py b/dimos/agents/skills/point_nav.py index 271158bd07e..280b1e223a9 100644 --- a/dimos/agents/skills/point_nav.py +++ b/dimos/agents/skills/point_nav.py @@ -16,15 +16,15 @@ from collections.abc import AsyncIterator import math +from dimos_generated.geometry_msgs.msg import Point, PointStamped, PoseStamped +from dimos_generated.tf2_msgs.msg import TFMessage from dimos_lcm.std_msgs import Bool from dimos.agents.annotation import skill from dimos.agents.capabilities import CAP_MOVEMENT from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PointStamped import PointStamped -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.time import header_now class PointNavSkillContainerConfig(ModuleConfig): @@ -92,19 +92,23 @@ async def go_to(self, x: float, y: float, timeout_s: float = 90.0) -> str: # TODO: fix this in the planner. It drops a goal the robot is already within its # goal_tolerance of without publishing anything, so no goal_reached would ever end this # go_to. Until the planner announces such a goal, repeat its check here. - distance = math.dist((x, y), (pose.x, pose.y)) + distance = math.dist((x, y), (pose.pose.position.x, pose.pose.position.y)) if distance < self.config.arrival_radius_m: self.stop_tool("go_to") return ( - f"Already there: robot at ({pose.x:.2f}, {pose.y:.2f}), " + f"Already there: robot at ({pose.pose.position.x:.2f}, {pose.pose.position.y:.2f}), " f"{distance:.2f} m from the target." ) # One x, y can have walkable surfaces at several heights, e.g. downstairs and upstairs. # The planner picks the one nearest the goal's z, so send the height of the floor under # the robot. - floor_z = pose.z - self.config.base_height_m - self.goal.publish(PointStamped(x, y, floor_z, frame_id=self.config.world_frame)) + floor_z = pose.pose.position.z - self.config.base_height_m + self.goal.publish( + PointStamped( + header=header_now(self.config.world_frame), point=Point(x=x, y=y, z=floor_z) + ) + ) self._target = (x, y) self._timeout = asyncio.get_running_loop().call_later( timeout_s, self._finish, f"Gave up after {timeout_s:g}s" @@ -118,7 +122,10 @@ def _finish(self, outcome: str) -> None: self._timeout = None # A NaN goal cancels the planner's goal self.goal.publish( - PointStamped(math.nan, math.nan, math.nan, frame_id=self.config.world_frame) + PointStamped( + header=header_now(self.config.world_frame), + point=Point(x=math.nan, y=math.nan, z=math.nan), + ) ) # Tell the BasicPathFollower to stop self.stop_movement.publish(Bool(True)) @@ -126,8 +133,8 @@ def _finish(self, outcome: str) -> None: if pose is None: where = "robot position unknown" else: - distance = math.dist(self._target, (pose.x, pose.y)) - where = f"robot at ({pose.x:.2f}, {pose.y:.2f}), {distance:.2f} m from the target" + distance = math.dist(self._target, (pose.pose.position.x, pose.pose.position.y)) + where = f"robot at ({pose.pose.position.x:.2f}, {pose.pose.position.y:.2f}), {distance:.2f} m from the target" self.tool_update("go_to", f"{outcome}; {where}.") self.stop_tool("go_to") diff --git a/dimos/core/native_module.py b/dimos/core/native_module.py index ff98bcf2d90..f709ddb3519 100644 --- a/dimos/core/native_module.py +++ b/dimos/core/native_module.py @@ -41,6 +41,7 @@ class MyCppModule(NativeModule): from __future__ import annotations +from dataclasses import asdict, is_dataclass import enum import functools import json @@ -53,6 +54,7 @@ class MyCppModule(NativeModule): import time from typing import IO, Any +import numpy as np from pydantic import Field, model_validator from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT @@ -88,6 +90,21 @@ def _set_process_to_die_when_parent_dies() -> None: logger = setup_logger() +def _json_values(value: Any) -> Any: + """Convert generated value fields to JSON without changing their schema names.""" + if is_dataclass(value) and not isinstance(value, type): + return _json_values(asdict(value)) + if isinstance(value, np.ndarray): + return value.tolist() + if isinstance(value, np.generic): + return value.item() + if isinstance(value, dict): + return {key: _json_values(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_json_values(item) for item in value] + return value + + class LogFormat(enum.Enum): TEXT = "text" JSON = "json" @@ -202,7 +219,7 @@ def to_config_dict(self) -> dict[str, Any]: # An opted-in base field is sent even when None, so the native struct # reports a null it can name rather than a field that looks unset. return { - k: v + k: _json_values(v) for k, v in self.model_dump().items() if k not in ignore_fields and (v is not None or k in self.base_fields) } diff --git a/dimos/hardware/sensors/lidar/pointlio/module.py b/dimos/hardware/sensors/lidar/pointlio/module.py index ff80d1fba4b..f7c609438dc 100644 --- a/dimos/hardware/sensors/lidar/pointlio/module.py +++ b/dimos/hardware/sensors/lidar/pointlio/module.py @@ -36,14 +36,13 @@ from typing import TYPE_CHECKING, Literal +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import Imu, PointCloud2 from pydantic import Field from dimos.core.core import rpc from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.sensor_msgs.Imu import Imu -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.spec import perception diff --git a/dimos/mapping/pointclouds/occupancy.py b/dimos/mapping/pointclouds/occupancy.py index d723b901543..21a1c8f75c6 100644 --- a/dimos/mapping/pointclouds/occupancy.py +++ b/dimos/mapping/pointclouds/occupancy.py @@ -18,6 +18,7 @@ from typing import TYPE_CHECKING, Any, Protocol, TypeVar import numpy as np +from numpy.typing import NDArray from scipy import ndimage from dimos.msgs.geometry_msgs.Pose import Pose diff --git a/dimos/mapping/ray_tracing/utils/loaded_map.py b/dimos/mapping/ray_tracing/utils/loaded_map.py index 2ed096246c9..8269cd62560 100644 --- a/dimos/mapping/ray_tracing/utils/loaded_map.py +++ b/dimos/mapping/ray_tracing/utils/loaded_map.py @@ -16,13 +16,14 @@ from __future__ import annotations +from dimos_generated.sensor_msgs.msg import PointCloud2 import numpy as np from numpy.typing import NDArray from dimos.memory.store.sqlite import SqliteStore from dimos.memory.tf import StreamTF from dimos.memory.type.observation import Observation -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_xyz, transform_cloud LOADED_MAP_STREAM = "loaded_map" @@ -41,10 +42,10 @@ def place_loaded_map( loaded_map: Observation[PointCloud2], tf: StreamTF, world_frame: str, ts: float ) -> NDArray[np.float32]: """World-frame points of the loaded map, placed by the transform at ts.""" - placement = tf.get(world_frame, loaded_map.data.frame_id, time_point=ts) + placement = tf.get(world_frame, loaded_map.data.header.frame_id, time_point=ts) if placement is None: raise RuntimeError( - f"no {world_frame}->{loaded_map.data.frame_id} transform at ts={ts:.3f} " + f"no {world_frame}->{loaded_map.data.header.frame_id} transform at ts={ts:.3f} " "to place the loaded map" ) - return loaded_map.data.transform(placement).points_f32() + return pointcloud_xyz(transform_cloud(loaded_map.data, placement)).astype(np.float32) diff --git a/dimos/mapping/relocalization/lidar/replay.py b/dimos/mapping/relocalization/lidar/replay.py index 7845e0a3b23..a63b79a9b6d 100644 --- a/dimos/mapping/relocalization/lidar/replay.py +++ b/dimos/mapping/relocalization/lidar/replay.py @@ -22,8 +22,14 @@ import math from pathlib import Path -from typing import TYPE_CHECKING, NamedTuple - +from typing import NamedTuple + +from dimos_generated.geometry_msgs.msg import TransformStamped as Transform +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage +from dimos_message_build.registry import decode as cdr_decode +import numpy as np import typer from dimos.mapping.ray_tracing.module import RayTracingVoxelMapConfig @@ -34,12 +40,10 @@ from dimos.mapping.relocalization.module import yaw_deg from dimos.memory.store.sqlite import SqliteStore from dimos.memory.tf import StreamTF -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2, register_colormap_annotation -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.pointcloud import pointcloud_to_open3d, pointcloud_xyz, transform_cloud +from dimos.msgs.time import time_from_seconds from dimos.utils.data import resolve_named_path - -if TYPE_CHECKING: - from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.visualization.rerun.message_helpers import register_colormap_annotation # rerun is imported inside the functions that log: it is heavy and only this command loads it. @@ -68,7 +72,8 @@ class Replay(NamedTuple): def fix_error(fix: Transform, recorded: Transform) -> tuple[float, float]: """Yaw in degrees and translation in meters between two fixes of the same frame pair.""" dyaw = (yaw_deg(fix) - yaw_deg(recorded) + 180.0) % 360.0 - 180.0 - return dyaw, (fix.translation - recorded.translation).length() + a, b = fix.transform.translation, recorded.transform.translation + return dyaw, math.dist((a.x, a.y, a.z), (b.x, b.y, b.z)) def recorded_fix(store: SqliteStore, world_frame: str, map_frame: str) -> Transform | None: @@ -77,7 +82,7 @@ def recorded_fix(store: SqliteStore, world_frame: str, map_frame: str) -> Transf return None for obs in store.stream("tf", TFMessage).order_by("ts"): for tf in obs.data.transforms: - if (tf.frame_id, tf.child_frame_id) == (world_frame, map_frame): + if (tf.header.frame_id, tf.child_frame_id) == (world_frame, map_frame): return tf return None @@ -131,7 +136,7 @@ def replay( ray = RayTraceMap(voxel_size=voxel_size) frames = lidar.transform(pose_from_tf(tf, world_frame)).transform(ray) - relocalizer = LidarRelocalizer(premap.pointcloud, PRESETS[preset]) + relocalizer = LidarRelocalizer(pointcloud_to_open3d(premap), PRESETS[preset]) recorded = recorded_fix(store, world_frame, MAP_FRAME) attempts: list[ReplayAttempt] = [] @@ -147,11 +152,18 @@ def replay( if obs.ts >= stop_at: break rr.set_time(TIMELINE, timestamp=obs.ts) - rr.log("world/local_map", voxel_map_points(obs.data.points_f32(), voxel_size)) - if fix is not None or obs.ts < next_attempt or len(obs.data) < min_local_points: + rr.log( + "world/local_map", + voxel_map_points(pointcloud_xyz(obs.data).astype(np.float32), voxel_size), + ) + if ( + fix is not None + or obs.ts < next_attempt + or obs.data.width * obs.data.height < min_local_points + ): continue next_attempt = obs.ts + reloc_interval - fix_attempt = relocalizer.attempt(obs.data.pointcloud, world_frame, MAP_FRAME) + fix_attempt = relocalizer.attempt(pointcloud_to_open3d(obs.data), world_frame, MAP_FRAME) attempts.append(ReplayAttempt(obs.ts - t0, fix_attempt.result.fitness, fix_attempt.fix)) rr.log(FITNESS_SERIES, rr.Scalars(fix_attempt.result.fitness)) _print_attempt(attempts[-1], recorded) @@ -159,12 +171,12 @@ def replay( continue fix, fix_ts = fix_attempt.fix, obs.ts stop_at = obs.ts + after_s - log_loaded_map(premap.transform(fix).points_f32()) + log_loaded_map(pointcloud_xyz(transform_cloud(premap, fix)).astype(np.float32)) if recorded is not None: rr.log( "world/recorded_map", rr.Points3D( - premap.transform(recorded).points_f32(), + pointcloud_xyz(transform_cloud(premap, recorded)).astype(np.float32), colors=[RECORDED_MAP_COLOR], radii=PREMAP_POINT_RADIUS, ), @@ -180,9 +192,12 @@ def write_loaded_map( if stream.count() > 0: print(f"{LOADED_MAP_STREAM} already has {stream.count()} messages, leaving it") return False - placed = PointCloud2(pointcloud=premap.transform(fix).pointcloud, frame_id=world_frame, ts=ts) + placed = transform_cloud(premap, fix) + placed.header = Header(stamp=time_from_seconds(ts), frame_id=world_frame) stream.append(placed, ts=ts) - print(f"wrote {len(placed)} placed premap points to {LOADED_MAP_STREAM} at {ts:.3f}") + print( + f"wrote {placed.width * placed.height} placed premap points to {LOADED_MAP_STREAM} at {ts:.3f}" + ) return True @@ -190,7 +205,7 @@ def _print_attempt(attempt: ReplayAttempt, recorded: Transform | None) -> None: if attempt.fix is None: print(f"{attempt.t_s:.1f}s refused fitness={attempt.fitness:.3f}") return - t = attempt.fix.translation + t = attempt.fix.transform.translation line = ( f"{attempt.t_s:.1f}s fitness={attempt.fitness:.3f} " f"t=({t.x:.2f}, {t.y:.2f}, {t.z:.2f}) yaw={yaw_deg(attempt.fix):.1f}deg" @@ -203,7 +218,7 @@ def _print_attempt(attempt: ReplayAttempt, recorded: Transform | None) -> None: def main( recording: str = typer.Argument(..., help="Recording .db: bare name (cwd or data/) or path"), - premap: str = typer.Option(..., "--premap", help="Premap .pc2.lcm: bare name or path"), + premap: str = typer.Option(..., "--premap", help="Premap .pc2.cdr: bare name or path"), lidar: str = typer.Option("lidar", "--lidar", help="Lidar stream in the recording"), world_frame: str = typer.Option("odom", "--world-frame", help="Frame the live map is built in"), preset: str = typer.Option(DEFAULT_PRESET, "--preset", help=f"One of {sorted(PRESETS)}"), @@ -232,9 +247,11 @@ def main( ) -> None: """Replay a recording through the relocalizer and compare its fix with the robot's.""" db_path = resolve_named_path(recording, ".db") - premap_cloud = PointCloud2.lcm_decode(resolve_named_path(premap, ".pc2.lcm").read_bytes()) + premap_cloud = cdr_decode(resolve_named_path(premap, ".pc2.cdr").read_bytes(), PointCloud2) _init_recording(db_path.stem, out) - print(f"relocalizing {db_path.name} against {premap} ({len(premap_cloud)} points, {preset})") + print( + f"relocalizing {db_path.name} against {premap} ({premap_cloud.width * premap_cloud.height} points, {preset})" + ) store = SqliteStore(path=str(db_path)) with store: result = replay( diff --git a/dimos/memory/type/recording.py b/dimos/memory/type/recording.py index 489325030c6..e66679a9ef0 100644 --- a/dimos/memory/type/recording.py +++ b/dimos/memory/type/recording.py @@ -14,11 +14,4 @@ """Policy for recording to an existing artifact.""" -import enum - - -class OnExisting(str, enum.Enum): - OVERWRITE = "overwrite" - ERROR = "error" - BACKUP = "backup" - APPEND = "append" +from dimos.memory.recording_policy import OnExisting as OnExisting diff --git a/dimos/navigation/go2/loop_closure/module.py b/dimos/navigation/go2/loop_closure/module.py index fc72ecd34a1..0a34b1845e7 100644 --- a/dimos/navigation/go2/loop_closure/module.py +++ b/dimos/navigation/go2/loop_closure/module.py @@ -16,18 +16,16 @@ import time +from dimos_generated.dimos_msgs.msg import LineSegment3D, LineSegments3D +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseArray, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Float32, Header from reactivex.disposable import Disposable from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseArray import PoseArray -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.LineSegments3D import LineSegments3D -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.std_msgs.Float32 import Float32 -from dimos.msgs.std_msgs.Header import Header +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.navigation.go2.loop_closure.pgo import PGOConfig from dimos.navigation.go2.loop_closure.pgo_map import LIVE_PGO, PGOMap from dimos.robot.unitree.type.lidar import repair_stale_ts @@ -95,7 +93,7 @@ def _on_odom(self, pose: PoseStamped) -> None: def _on_lidar(self, cloud: PointCloud2) -> None: assert self._map is not None t0, loops = time.perf_counter(), self._map.n_loops - rebuilt = self._map.add(cloud, self._pose) + rebuilt = self._map.add(cloud, self._pose.pose if self._pose is not None else None) self._frames += 1 emit = rebuilt or self._frames % self.config.emit_every == 0 global_map = self._map.global_map() if emit else None @@ -105,7 +103,7 @@ def _on_lidar(self, cloud: PointCloud2) -> None: self.global_map.publish(global_map) if self._map.n_keyframes != self._graph_keyframes: self._graph_keyframes = self._map.n_keyframes - self._publish_graph(cloud.ts) + self._publish_graph(to_seconds(cloud.header.stamp)) if self.config.telemetry: self.mapper_frame_ms.publish(Float32(frame_ms)) if self._map.n_loops != loops: @@ -119,14 +117,20 @@ def _publish_graph(self, ts: float) -> None: positions, quats = self._map.keyframe_poses() self.pgo_keyframes.publish( PoseArray( - Header(ts, frame_id), + Header(stamp=time_from_seconds(ts), frame_id=frame_id), [ - Pose(*xyz, *quat) + Pose(position=Point(*xyz), orientation=Quaternion(*quat)) for xyz, quat in zip(positions.tolist(), quats.tolist(), strict=True) ], ) ) segments, scores = self._map.loop_segments() self.pgo_loops.publish( - LineSegments3D(ts=ts, frame_id=frame_id, segments=segments, weights=scores) + LineSegments3D( + header=Header(stamp=time_from_seconds(ts), frame_id=frame_id), + segments=[ + LineSegment3D(start=Point(*start), end=Point(*end), weight=float(score)) + for (start, end), score in zip(segments.tolist(), scores.tolist(), strict=True) + ], + ) ) diff --git a/dimos/navigation/go2/loop_closure/pgo_map.py b/dimos/navigation/go2/loop_closure/pgo_map.py index 3ab5abcd876..d08b13c047b 100644 --- a/dimos/navigation/go2/loop_closure/pgo_map.py +++ b/dimos/navigation/go2/loop_closure/pgo_map.py @@ -31,11 +31,13 @@ import sys import time +from dimos_generated.geometry_msgs.msg import Pose +from dimos_generated.sensor_msgs.msg import PointCloud2 import numpy as np from dimos.mapping.voxels.grid import VoxelGrid -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_xyz +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.navigation.go2.loop_closure.pgo import PGOConfig, PoseGraph, _PGOState, _pose_to_pose3 from dimos.utils.logging_config import setup_logger @@ -79,18 +81,24 @@ def n_loops(self) -> int: def add(self, cloud: PointCloud2, pose: Pose | None) -> bool: """Insert a world-frame lidar frame taken at odom `pose`. True if the map was rebuilt.""" self._grid.add_frame(cloud) - if pose is not None and not (pose.position.is_zero() or pose.orientation.is_zero()): + if pose is not None and not ( + all(getattr(pose.position, axis) == 0 for axis in ("x", "y", "z")) + or all(getattr(pose.orientation, axis) == 0 for axis in ("x", "y", "z", "w")) + ): t0, loops = time.perf_counter(), self.n_loops - self._pgo.process(_pose_to_pose3(pose), cloud.ts, cloud) + self._pgo.process(_pose_to_pose3(pose), to_seconds(cloud.header.stamp), cloud) if self.n_loops != loops: self.loop_ms = (time.perf_counter() - t0) * 1e3 # frame time, not wall time, so replay and tests pace the same as the robot - if self.n_loops == self._placed_loops or cloud.ts - self._last_rebuild_ts < self._cooldown: + if ( + self.n_loops == self._placed_loops + or to_seconds(cloud.header.stamp) - self._last_rebuild_ts < self._cooldown + ): return False t0 = time.perf_counter() self._grid.reproject(self.graph().place) - self._placed_loops, self._last_rebuild_ts = self.n_loops, cloud.ts + self._placed_loops, self._last_rebuild_ts = self.n_loops, to_seconds(cloud.header.stamp) self.rebuild_ms = (time.perf_counter() - t0) * 1e3 logger.info( "PGO map rebuilt", @@ -144,7 +152,9 @@ def main(dataset: str) -> None: def frames() -> Iterator[tuple[PointCloud2, Pose | None]]: for obs in lidar: cloud = obs.data - cloud.ts = obs.ts # old recordings carry Unitree's stale lidar stamps + cloud.header.stamp = time_from_seconds( + obs.ts + ) # old recordings carry Unitree's stale lidar stamps yield cloud, obs.pose live = PGOMap() @@ -163,7 +173,10 @@ def frames() -> Iterator[tuple[PointCloud2, Pose | None]]: for cloud, _ in frames(): gold.add_frame(graph.place(cloud)) - got, want = live.global_map().points_f32(), gold.get_global_pointcloud2().points_f32() + got, want = ( + pointcloud_xyz(live.global_map()).astype(np.float32), + pointcloud_xyz(gold.get_global_pointcloud2()).astype(np.float32), + ) missing = 1 - _within_one_voxel(want, got, 0.05) extra = 1 - _within_one_voxel(got, want, 0.05) steps = np.array(step_ms) diff --git a/dimos/navigation/local_planner/module.py b/dimos/navigation/local_planner/module.py index e61d41ba530..a3178a33266 100644 --- a/dimos/navigation/local_planner/module.py +++ b/dimos/navigation/local_planner/module.py @@ -26,6 +26,11 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage import numpy as np from numpy.typing import NDArray from pydantic import Field, ImportString @@ -35,13 +40,9 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import IO, In, Out -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.geometry import quaternion_from_euler, yaw +from dimos.msgs.pointcloud import pointcloud_xyz +from dimos.msgs.time import time_from_seconds from dimos.navigation import spec from dimos.navigation.embodiment.base import Embodiment from dimos.navigation.embodiment.go2 import GO2 @@ -70,7 +71,7 @@ def annotate( ) -> Path: """The planner's route as the follower's path: stamped, grounded, precision profile in the timestamps.""" nav = stamped(ref, ts=ts, frame_id=frame_id, ground_z=ground_z) - xy = np.array([[p.position.x, p.position.y] for p in ref.poses]).reshape(-1, 2) + xy = np.array([[p.pose.position.x, p.pose.position.y] for p in ref.poses]).reshape(-1, 2) clearance = path_clearance(xy, obstacles, emb.width / 2.0) return encode_precision(nav, clearance, emb, t0=ts) @@ -79,16 +80,20 @@ def stamped(ref: Path, ts: float = 0.0, frame_id: str = "odom", ground_z: float """The route stamped with time, frame and ground: the search's z = 0 is not the floor, `ground_z` is.""" poses = [ PoseStamped( - ts=ts, - frame_id=frame_id, - position=Vector3(p.position.x, p.position.y, ground_z), - orientation=Quaternion( - p.orientation.x, p.orientation.y, p.orientation.z, p.orientation.w + header=Header(stamp=time_from_seconds(ts), frame_id=frame_id), + pose=Pose( + position=Point(x=p.pose.position.x, y=p.pose.position.y, z=ground_z), + orientation=Quaternion( + p.pose.orientation.x, + p.pose.orientation.y, + p.pose.orientation.z, + p.pose.orientation.w, + ), ), ) for p in ref.poses ] - return Path(ts=ts, frame_id=frame_id, poses=poses) + return Path(header=Header(stamp=time_from_seconds(ts), frame_id=frame_id), poses=poses) def carrot_along( @@ -221,7 +226,7 @@ def _on_local_map(self, msg: PointCloud2) -> None: def _on_planner_path(self, msg: Path) -> None: # an empty path is MLS finding no route: no carrot, hold the last local plan - xy = np.array([[p.position.x, p.position.y] for p in msg.poses]).reshape(-1, 2) + xy = np.array([[p.pose.position.x, p.pose.position.y] for p in msg.poses]).reshape(-1, 2) with self._lock: self._global_xy = xy if len(xy) else None @@ -252,14 +257,18 @@ def tick(self) -> None: self._stale = False logger.info("local_map is live again, resuming planning") # the gate reads the carrot, not the array (replan_due) - goal = carrot_along(global_xy, (pose.x, pose.y), self.config.goal_lookahead_m) + goal = carrot_along( + global_xy, + (pose.pose.position.x, pose.pose.position.y), + self.config.goal_lookahead_m, + ) if self.due(cloud_seq, goal): if self.retask(goal): # a new task: warm start, hysteresis and the incumbent are about the old one self._episode.reset() self._incumbent = None # the base rides emb.base_height above the surface - ground_z = pose.position.z - self._emb.base_height + ground_z = pose.pose.position.z - self._emb.base_height if self.plan_once(cloud, pose, goal, ground_z): self._planned = (cloud_seq, goal) @@ -282,7 +291,14 @@ def clear(self) -> None: self._incumbent = None self._episode.reset() logger.info("planner_path is empty, clearing the local plan") - self.path.publish(Path(ts=time.time(), frame_id=self.config.world_frame, poses=[])) + self.path.publish( + Path( + header=Header( + stamp=time_from_seconds(time.time()), frame_id=self.config.world_frame + ), + poses=[], + ) + ) def hold(self, pose: PoseStamped, age: float) -> None: """Refuse the way the planner does: a single-pose stub reads as "stop".""" @@ -295,13 +311,19 @@ def hold(self, pose: PoseStamped, age: float) -> None: max_map_age_s=self.config.max_map_age_s, ) ts = time.time() + header = Header(stamp=time_from_seconds(ts), frame_id=self.config.world_frame) stub = PoseStamped( - ts=ts, - frame_id=self.config.world_frame, - position=Vector3(pose.x, pose.y, pose.position.z - self._emb.base_height), - orientation=Quaternion.from_euler(Vector3(0.0, 0.0, pose.yaw)), + header=header, + pose=Pose( + position=Point( + x=pose.pose.position.x, + y=pose.pose.position.y, + z=pose.pose.position.z - self._emb.base_height, + ), + orientation=quaternion_from_euler(0.0, 0.0, yaw(pose.pose.orientation)), + ), ) - held = Path(ts=ts, frame_id=self.config.world_frame, poses=[stub]) + held = Path(header=header, poses=[stub]) self._published = True self.path.publish(held) @@ -314,13 +336,16 @@ def plan_once( ) -> bool: """Plan and publish. False when the search raised and nothing went out.""" # the model decides the obstacles here; the follower's room hint is measured off the same points - raw = cloud.points_f32() + raw = pointcloud_xyz(cloud).astype(np.float32) pts = hard_points(self._model, raw, ground_z) try: ref = self._episode.plan( pts[:, :2], - pose, - Pose(goal[0], goal[1], 0.0), + pose.pose, + Pose( + position=Point(x=goal[0], y=goal[1], z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), self._incumbent, ground=ground_points(raw, ground_z), unseen_cost=self.config.unseen_cost, diff --git a/dimos/navigation/local_planner/native.py b/dimos/navigation/local_planner/native.py index c0a32db4959..b8d0de36042 100644 --- a/dimos/navigation/local_planner/native.py +++ b/dimos/navigation/local_planner/native.py @@ -22,11 +22,12 @@ from typing import TYPE_CHECKING, Any +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.tf2_msgs.msg import TFMessage + from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.stream import IO, In, Out -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.navigation import spec from dimos.navigation.embodiment.base import Embodiment from dimos.navigation.local_planner.module import LocalPlannerConfig diff --git a/dimos/navigation/local_planner/profile.py b/dimos/navigation/local_planner/profile.py index 585248cc6d6..87382ca52b1 100644 --- a/dimos/navigation/local_planner/profile.py +++ b/dimos/navigation/local_planner/profile.py @@ -28,10 +28,12 @@ from typing import Any +from dimos_generated.nav_msgs.msg import Path import numpy as np from numpy.typing import NDArray -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import yaw +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.navigation.embodiment.base import Embodiment # The governor curve is the embodiment's (max_speed, min_speed, speed_clearance, @@ -49,8 +51,8 @@ def governor_speed(clearance: NDArray[np.floating[Any]], emb: Embodiment) -> NDA def _segments(path: Path) -> tuple[NDArray[np.float64], NDArray[np.float64]]: - xy = np.array([[p.position.x, p.position.y] for p in path.poses]).reshape(-1, 2) - yaws = np.array([p.yaw for p in path.poses]) + xy = np.array([[p.pose.position.x, p.pose.position.y] for p in path.poses]).reshape(-1, 2) + yaws = np.array([yaw(p.pose.orientation) for p in path.poses]) return xy, yaws @@ -68,7 +70,7 @@ def encode_precision( xy, yaws = _segments(path) v = governor_speed(clearance, emb) if len(clearance) == n else np.full(n, emb.max_speed) t = t0 - path.poses[0].ts = t + path.poses[0].header.stamp = time_from_seconds(t) for i in range(1, n): ds = float(np.linalg.norm(xy[i] - xy[i - 1])) if ds < _FAN_EPS: @@ -76,8 +78,8 @@ def encode_precision( t += dyaw / emb.max_yaw_rate else: t += ds / float(min(v[i - 1], v[i])) - path.poses[i].ts = t - path.ts = t0 + path.poses[i].header.stamp = time_from_seconds(t) + path.header.stamp = time_from_seconds(t0) return path @@ -90,7 +92,7 @@ def decode_ceilings(path: Path, lo: float, hi: float) -> NDArray[np.float64] | N n = len(path.poses) if n < 2: return None - ts = np.array([p.ts for p in path.poses]) + ts = np.array([to_seconds(p.header.stamp) for p in path.poses]) dt = np.diff(ts) if np.any(dt < 0) or not np.any(dt > 0): return None diff --git a/dimos/navigation/local_planner/rust/Cargo.lock b/dimos/navigation/local_planner/rust/Cargo.lock index 7bfb7576bef..ba85d0d4e11 100644 --- a/dimos/navigation/local_planner/rust/Cargo.lock +++ b/dimos/navigation/local_planner/rust/Cargo.lock @@ -226,6 +226,20 @@ name = "bytemuck" version = "1.25.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a1f896587b6f2c069c73d2f0913e2d590c3990285cd2f0b6aa02b786b4c679c" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.6", +] [[package]] name = "byteorder" @@ -580,9 +594,19 @@ dependencies = [ ] [[package]] -name = "dimos-lcm" +name = "dimos-generated-messages" +version = "0.1.0" +dependencies = [ + "heapless", + "re_cdr", + "ros2msg", + "serde", + "serde-big-array", +] + +[[package]] +name = "dimos-lcm-transport" version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" dependencies = [ "byteorder", "socket2 0.5.10", @@ -593,9 +617,9 @@ dependencies = [ name = "dimos-local-planner" version = "0.1.0" dependencies = [ + "dimos-generated-messages", "dimos-local-planner", "dimos-module", - "lcm-msgs", "numpy", "pyo3", "serde", @@ -610,11 +634,12 @@ dependencies = [ name = "dimos-module" version = "0.1.0" dependencies = [ - "dimos-lcm", + "dimos-generated-messages", + "dimos-lcm-transport", "dimos-module-macros", - "lcm-msgs", "nalgebra", "rayon", + "re_cdr", "serde", "serde_json", "tokio", @@ -966,6 +991,15 @@ version = "0.33.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e762b9b188634fc508d4ec54b62ea3d352c43fd431d18d05d46f559f217f4725" +[[package]] +name = "hash32" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47d60b12902ba28e2730cd37e95b8c9223af2808df9e902d4df49588d1470606" +dependencies = [ + "byteorder", +] + [[package]] name = "hashbrown" version = "0.12.3" @@ -1008,6 +1042,17 @@ version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +[[package]] +name = "heapless" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bfb9eb618601c89945a70e254898da93b13be0388091d42117462b265bb3fad" +dependencies = [ + "hash32", + "serde", + "stable_deref_trait", +] + [[package]] name = "heck" version = "0.5.0" @@ -1413,14 +1458,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" -[[package]] -name = "lcm-msgs" -version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" -dependencies = [ - "byteorder", -] - [[package]] name = "libc" version = "0.2.189" @@ -1997,6 +2034,16 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + [[package]] name = "proc-macro-crate" version = "3.5.0" @@ -2297,6 +2344,18 @@ dependencies = [ "yasna", ] +[[package]] +name = "re_cdr" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bba10c21433b597481f366b978cdf2617e56785b38cb0dcc3833c6b8ad7e061a" +dependencies = [ + "bytemuck", + "byteorder", + "serde", + "thiserror", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -2412,6 +2471,23 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "ros2msg" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fe291be5cb6d4b8e6bf9c9de80c40bd1cb09ecac50deafa2fdcae0027e04988" +dependencies = [ + "heck", + "pest", + "pest_derive", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "syn 2.0.119", + "thiserror", +] + [[package]] name = "rustc-hash" version = "2.1.3" @@ -2652,6 +2728,15 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "serde-big-array" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11fc7cc2c76d73e0f27ee52abbd64eec84d46f370c88371120433196934e4b7f" +dependencies = [ + "serde", +] + [[package]] name = "serde_core" version = "1.0.229" diff --git a/dimos/navigation/local_planner/rust/Cargo.toml b/dimos/navigation/local_planner/rust/Cargo.toml index 445f9a75654..6c753eeb247 100644 --- a/dimos/navigation/local_planner/rust/Cargo.toml +++ b/dimos/navigation/local_planner/rust/Cargo.toml @@ -22,7 +22,7 @@ serde_json = { version = "1", features = ["float_roundtrip"] } pyo3 = { version = "0.25", features = ["extension-module", "abi3-py310"], optional = true } numpy = { version = "0.25", optional = true } dimos-module = { path = "../../../../native/rust/dimos-module", optional = true } -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen", optional = true } +dimos-generated-messages = { path = "../../../../packages/dimos-generated/src/dimos_generated_schemas/package/rust", optional = true } tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal", "time"], optional = true } tracing = { version = "0.1", optional = true } tracing-subscriber = { version = "0.3", features = ["env-filter"], optional = true } @@ -35,7 +35,7 @@ validator = { version = "0.21", features = ["derive"], optional = true } # `module` the native module that runs on the robot. They never build together. default = [] python = ["dep:pyo3", "dep:numpy"] -module = ["dep:dimos-module", "dep:lcm-msgs", "dep:tokio", "dep:tracing", "dep:tracing-subscriber", "dep:validator"] +module = ["dep:dimos-module", "dep:dimos-generated-messages", "dep:tokio", "dep:tracing", "dep:tracing-subscriber", "dep:validator"] # `Emb::fixture()`, a test body for this crate and its dependants. test-bodies = [] @@ -48,12 +48,12 @@ python = "dimos.navigation.local_planner.native:LocalPlannerNative" threads = 2 [package.metadata.dimos.module.local_planner.inputs] -local_map = "sensor_msgs.PointCloud2" -planner_path = "nav_msgs.Path" -tf = "tf2_msgs.TFMessage" +local_map = "sensor_msgs/msg/PointCloud2" +planner_path = "nav_msgs/msg/Path" +tf = "tf2_msgs/msg/TFMessage" [package.metadata.dimos.module.local_planner.outputs] -path = "nav_msgs.Path" +path = "nav_msgs/msg/Path" [profile.release] lto = "thin" diff --git a/dimos/navigation/local_planner/rust/src/module/msg.rs b/dimos/navigation/local_planner/rust/src/module/msg.rs index a722f1c9cd3..a97dc65b837 100644 --- a/dimos/navigation/local_planner/rust/src/module/msg.rs +++ b/dimos/navigation/local_planner/rust/src/module/msg.rs @@ -17,10 +17,15 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use dimos_generated_messages::builtin_interfaces::msg::time::Time; +use dimos_generated_messages::geometry_msgs::msg::point::Point; +use dimos_generated_messages::geometry_msgs::msg::pose::Pose; +use dimos_generated_messages::geometry_msgs::msg::pose_stamped::PoseStamped; +use dimos_generated_messages::geometry_msgs::msg::quaternion::Quaternion; +use dimos_generated_messages::geometry_msgs::msg::vector3::Vector3; +use dimos_generated_messages::nav_msgs::msg::path::Path; +use dimos_generated_messages::std_msgs::msg::header::Header; use dimos_module::{error_throttled, Output}; -use lcm_msgs::geometry_msgs::{Point, Pose, PoseStamped, Quaternion, Vector3}; -use lcm_msgs::nav_msgs::Path; -use lcm_msgs::std_msgs::{Header, Time}; /// A plan waypoint as the laws and the planner see it: `(x, y, yaw)`. pub type State = [f64; 3]; @@ -45,19 +50,19 @@ pub fn quat_of_yaw(yaw: f64) -> Quaternion { /// Seconds since the unix epoch as one float, the python `PoseStamped.ts`. pub fn secs_of(t: &Time) -> f64 { - t.sec as f64 + t.nsec as f64 * 1e-9 + t.sec as f64 + t.nanosec as f64 * 1e-9 } /// A float second count as a `Time`, matching `Header.__init__`'s truncation -/// (`sec = int(ts)`, `nsec = int((ts - sec) * 1e9)`). +/// (`sec = int(ts)`, `nanosec = int((ts - sec) * 1e9)`). pub fn time_of_secs(ts: f64) -> Time { if !ts.is_finite() { - return Time { sec: 0, nsec: 0 }; + return Time { sec: 0, nanosec: 0 }; } - let sec = ts.trunc(); + let sec = ts.floor(); Time { sec: sec.clamp(i32::MIN as f64, i32::MAX as f64) as i32, - nsec: ((ts - sec) * 1e9) as i32, + nanosec: ((ts - sec) * 1e9) as u32, } } @@ -71,7 +76,6 @@ pub fn now_secs() -> f64 { pub fn header(frame_id: &str, ts: f64) -> Header { Header { - seq: 0, stamp: time_of_secs(ts), frame_id: frame_id.into(), } @@ -136,8 +140,12 @@ pub fn path_stamps(path: &Path) -> Vec { } /// A body-frame twist, the only shape either module publishes. -pub fn twist(vx: f64, vy: f64, wz: f64) -> lcm_msgs::geometry_msgs::Twist { - lcm_msgs::geometry_msgs::Twist { +pub fn twist( + vx: f64, + vy: f64, + wz: f64, +) -> dimos_generated_messages::geometry_msgs::msg::twist::Twist { + dimos_generated_messages::geometry_msgs::msg::twist::Twist { linear: Vector3 { x: vx, y: vy, @@ -191,8 +199,8 @@ mod tests { #[test] fn a_non_finite_stamp_does_not_wrap_around() { // a saturating `as i32` would give i32::MAX seconds; zero is the honest "no time" - assert_eq!(time_of_secs(f64::NAN), Time { sec: 0, nsec: 0 }); - assert_eq!(time_of_secs(f64::INFINITY), Time { sec: 0, nsec: 0 }); + assert_eq!(time_of_secs(f64::NAN), Time { sec: 0, nanosec: 0 }); + assert_eq!(time_of_secs(f64::INFINITY), Time { sec: 0, nanosec: 0 }); } #[test] diff --git a/dimos/navigation/local_planner/rust/src/module/planner.rs b/dimos/navigation/local_planner/rust/src/module/planner.rs index 0364abf0a71..c0728804aa2 100644 --- a/dimos/navigation/local_planner/rust/src/module/planner.rs +++ b/dimos/navigation/local_planner/rust/src/module/planner.rs @@ -21,9 +21,10 @@ use std::time::{Duration, Instant}; use crate::planner::{plan_explored, Emb, COMMIT_MARGIN}; use crate::{clearance, stamps}; +use dimos_generated_messages::nav_msgs::msg::path::Path; +use dimos_generated_messages::sensor_msgs::msg::point_cloud2::PointCloud2; +use dimos_module::cdr; use dimos_module::{debug_throttled, native_config, warn_throttled, Input, Module, Output, Tf}; -use lcm_msgs::nav_msgs::Path; -use lcm_msgs::sensor_msgs::PointCloud2; use tracing::{debug, info, warn}; use validator::ValidationError; @@ -95,13 +96,13 @@ struct Shared { #[derive(Module)] #[module(name = "local_planner", setup = spawn_worker, teardown = stop_worker)] pub struct LocalPlanner { - #[input(decode = PointCloud2::decode, handler = on_local_map)] + #[input(decode = cdr::decode, handler = on_local_map)] local_map: Input, - #[input(decode = Path::decode, handler = on_planner_path)] + #[input(decode = cdr::decode, handler = on_planner_path)] planner_path: Input, - #[output(encode = Path::encode)] + #[output(encode = cdr::encode)] path: Output, #[config] diff --git a/dimos/navigation/local_planner/rust/src/module/tf_pose.rs b/dimos/navigation/local_planner/rust/src/module/tf_pose.rs index 6351ee970b6..9240e1864ca 100644 --- a/dimos/navigation/local_planner/rust/src/module/tf_pose.rs +++ b/dimos/navigation/local_planner/rust/src/module/tf_pose.rs @@ -86,12 +86,14 @@ fn edge_of(t: Transform) -> (f64, Isometry3) { pub fn state_of(iso: &Isometry3) -> State { let q = iso.rotation.quaternion(); - let yaw = yaw_of(&lcm_msgs::geometry_msgs::Quaternion { - x: q.i, - y: q.j, - z: q.k, - w: q.w, - }); + let yaw = yaw_of( + &dimos_generated_messages::geometry_msgs::msg::quaternion::Quaternion { + x: q.i, + y: q.j, + z: q.k, + w: q.w, + }, + ); [iso.translation.x, iso.translation.y, yaw] } diff --git a/dimos/navigation/local_planner/search/base.py b/dimos/navigation/local_planner/search/base.py index 7203d19446a..99f83187b65 100644 --- a/dimos/navigation/local_planner/search/base.py +++ b/dimos/navigation/local_planner/search/base.py @@ -40,14 +40,14 @@ import math from typing import Any, Protocol +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Header import numpy as np from numpy.typing import NDArray -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import quaternion_from_euler, yaw +from dimos.msgs.time import time_from_seconds RESOLUTION = 0.1 # waypoint spacing of the published route (m) # Largest yaw change one published waypoint may command (rad). A consumer that @@ -81,7 +81,7 @@ def states_of(path: Path | None) -> NDArray[np.float64] | None: if path is None or not path.poses: return None return np.array( - [[p.position.x, p.position.y, p.orientation.euler[2]] for p in path.poses] + [[p.pose.position.x, p.pose.position.y, yaw(p.pose.orientation)] for p in path.poses] ).reshape(-1, 3) @@ -104,8 +104,9 @@ def densify_states(states: NDArray[np.float64], res: float) -> list[NDArray[np.f def pose_stamped(x: float, y: float, yaw: float) -> PoseStamped: return PoseStamped( - ts=0.0, # deterministic: nothing here reads a stamp, and caches get pickled - frame_id="world", - position=[float(x), float(y), 0.0], - orientation=Quaternion.from_euler(Vector3(0.0, 0.0, float(yaw))), + header=Header(stamp=time_from_seconds(0.0), frame_id="world"), + pose=Pose( + position=Point(x=float(x), y=float(y), z=0.0), + orientation=quaternion_from_euler(0.0, 0.0, float(yaw)), + ), ) diff --git a/dimos/navigation/local_planner/search/target.py b/dimos/navigation/local_planner/search/target.py index a7a8bdaffde..5e86bbc5cf2 100644 --- a/dimos/navigation/local_planner/search/target.py +++ b/dimos/navigation/local_planner/search/target.py @@ -19,11 +19,14 @@ from typing import Any +from dimos_generated.geometry_msgs.msg import Pose +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Header import numpy as np from numpy.typing import NDArray -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import yaw +from dimos.msgs.time import time_from_seconds from dimos.navigation.embodiment.base import Embodiment from dimos.navigation.local_planner.search.base import ( RESOLUTION, @@ -67,8 +70,12 @@ def plan( # what ships. Port se2.py when parity on it is wanted. band = np.asarray(obstacles, dtype=float).reshape(-1, 2) - xs = [pose.x, goal.x] + ([] if not len(band) else [band[:, 0].min(), band[:, 0].max()]) - ys = [pose.y, goal.y] + ([] if not len(band) else [band[:, 1].min(), band[:, 1].max()]) + xs = [pose.position.x, goal.position.x] + ( + [] if not len(band) else [band[:, 0].min(), band[:, 0].max()] + ) + ys = [pose.position.y, goal.position.y] + ( + [] if not len(band) else [band[:, 1].min(), band[:, 1].max()] + ) # Anchored on the world lattice: a new return may add rows, never move a sample. x0, y0 = anchor(min(xs) - PAD), anchor(min(ys) - PAD) x1, y1 = max(xs) + PAD, max(ys) + PAD @@ -79,17 +86,21 @@ def plan( states = se2_search( grid, (x0, y0, x1, y1), - (pose.x, pose.y, pose.yaw), - (goal.x, goal.y), + (pose.position.x, pose.position.y, yaw(pose.orientation)), + (goal.position.x, goal.position.y), self._emb, self._emb.precision, incumbent=states_of(incumbent), ) if states is None: - return Path(ts=0.0, frame_id="world", poses=[pose_stamped(pose.x, pose.y, pose.yaw)]) + return Path( + header=Header(stamp=time_from_seconds(0.0), frame_id="world"), + poses=[pose_stamped(pose.position.x, pose.position.y, yaw(pose.orientation))], + ) dense = densify_states(states, self._res) return Path( - ts=0.0, frame_id="world", poses=[pose_stamped(x, y, yaw) for x, y, yaw in dense] + header=Header(stamp=time_from_seconds(0.0), frame_id="world"), + poses=[pose_stamped(x, y, yaw) for x, y, yaw in dense], ) @@ -120,8 +131,8 @@ def plan( inc = states_of(incumbent) out = self._mod.plan( pts, - (pose.x, pose.y, pose.yaw), - (goal.x, goal.y), + (pose.position.x, pose.position.y, yaw(pose.orientation)), + (goal.position.x, goal.position.y), self._emb, self._res, None if inc is None else np.ascontiguousarray(inc, dtype=np.float64), @@ -132,8 +143,14 @@ def plan( unseen_cost, ) if out is None or not len(out): - return Path(ts=0.0, frame_id="world", poses=[pose_stamped(pose.x, pose.y, pose.yaw)]) - return Path(ts=0.0, frame_id="world", poses=[pose_stamped(x, y, yaw) for x, y, yaw in out]) + return Path( + header=Header(stamp=time_from_seconds(0.0), frame_id="world"), + poses=[pose_stamped(pose.position.x, pose.position.y, yaw(pose.orientation))], + ) + return Path( + header=Header(stamp=time_from_seconds(0.0), frame_id="world"), + poses=[pose_stamped(x, y, yaw) for x, y, yaw in out], + ) def make_py(emb: Embodiment, resolution: float = RESOLUTION, **_: Any) -> TargetEpisode: diff --git a/dimos/navigation/local_planner/test_module.py b/dimos/navigation/local_planner/test_module.py index 90ca6ffa461..cdfce5d6e09 100644 --- a/dimos/navigation/local_planner/test_module.py +++ b/dimos/navigation/local_planner/test_module.py @@ -15,18 +15,26 @@ from dataclasses import replace import math +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Header import numpy as np from numpy.typing import NDArray import pytest from dimos.core.module import ModuleConfig -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.core.transport_factory import rpc_backend +from dimos.msgs.geometry import quaternion_from_euler, yaw +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import time_from_seconds from dimos.navigation.embodiment.base import Embodiment from dimos.navigation.embodiment.go2 import GO2 from dimos.navigation.local_planner import module as planner_module @@ -48,6 +56,12 @@ _BUILT: list[LocalPlanner] = [] +@pytest.fixture(autouse=True) +def isolate_rpc_boundary(mocker): + for method in ("start", "serve_module_rpc", "stop"): + mocker.patch.object(rpc_backend(), method) + + @pytest.fixture(autouse=True) def _stop_modules(): yield @@ -63,18 +77,26 @@ def _planner(**config) -> LocalPlanner: def test_stamped_preserves_positions_and_yaw(): ref = Path( - frame_id="world", poses=[pose_stamped(0.0, 0.0, 0.0), pose_stamped(1.0, 2.0, math.pi / 2)] + header=Header(stamp=time_from_seconds(0.0), frame_id="world"), + poses=[pose_stamped(0.0, 0.0, 0.0), pose_stamped(1.0, 2.0, math.pi / 2)], ) nav = stamped(ref, ts=3.0, frame_id="odom") - assert nav.frame_id == "odom" + assert nav.header.frame_id == "odom" assert len(nav.poses) == 2 - assert nav.poses[1].position.x == 1.0 - assert nav.poses[1].position.y == 2.0 - assert abs(nav.poses[1].yaw - math.pi / 2) < 1e-9 + assert nav.poses[1].pose.position.x == 1.0 + assert nav.poses[1].pose.position.y == 2.0 + assert abs(yaw(nav.poses[1].pose.orientation) - math.pi / 2) < 1e-9 def test_stamped_empty(): - assert len(stamped(Path(frame_id="world", poses=[])).poses) == 0 + assert ( + len( + stamped( + Path(header=Header(stamp=time_from_seconds(0.0), frame_id="world"), poses=[]) + ).poses + ) + == 0 + ) def test_carrot_walks_arc_from_closest_waypoint(): @@ -108,20 +130,23 @@ def test_hold_publishes_single_pose_stub_at_the_current_pose(): planner, published = _holding_planner() planner.hold( PoseStamped( - position=(1.5, -2.0, 0.0), orientation=Quaternion.from_euler(Vector3(0, 0, math.pi / 2)) + header=Header(stamp=time_from_seconds(0.0), frame_id=""), + pose=Pose( + position=Point(1.5, -2.0, 0.0), orientation=quaternion_from_euler(0, 0, math.pi / 2) + ), ), age=7.0, ) assert len(published) == 1 path = published[0] - assert path.frame_id == "odom" + assert path.header.frame_id == "odom" # a single pose is the planner's refusal shape assert len(path.poses) == 1 - assert path.poses[0].position.x == 1.5 - assert path.poses[0].position.y == -2.0 - assert abs(path.poses[0].yaw - math.pi / 2) < 1e-9 + assert path.poses[0].pose.position.x == 1.5 + assert path.poses[0].pose.position.y == -2.0 + assert abs(yaw(path.poses[0].pose.orientation) - math.pi / 2) < 1e-9 # the stub stands on the ground like every planned pose - assert abs(path.poses[0].position.z - (0.0 - GO2.base_height)) < 1e-9 + assert abs(path.poses[0].pose.position.z - (0.0 - GO2.base_height)) < 1e-9 def test_a_route_that_vanished_clears_the_plan_once(): @@ -129,17 +154,29 @@ def test_a_route_that_vanished_clears_the_plan_once(): # a route that never existed is "waiting", not a clear planner.clear() assert published == [] - planner.hold(PoseStamped(position=(1.5, -2.0, 0.0)), age=7.0) + planner.hold( + PoseStamped( + header=Header(stamp=time_from_seconds(0.0), frame_id=""), + pose=Pose(position=Point(1.5, -2.0, 0.0), orientation=Quaternion(0.0, 0.0, 0.0, 1.0)), + ), + age=7.0, + ) planner.clear() planner.clear() assert len(published) == 2 - assert published[1].poses == [] and published[1].frame_id == "odom" + assert published[1].poses == [] and published[1].header.frame_id == "odom" assert planner._planned is None and planner._incumbent is None def test_hold_stub_stops_the_controller(): planner, published = _holding_planner() - planner.hold(PoseStamped(position=(1.5, -2.0, 0.0)), age=7.0) + planner.hold( + PoseStamped( + header=Header(stamp=time_from_seconds(0.0), frame_id=""), + pose=Pose(position=Point(1.5, -2.0, 0.0), orientation=Quaternion(0.0, 0.0, 0.0, 1.0)), + ), + age=7.0, + ) pose = pose_stamped(1.5, -2.0, 0.0) twist = PursuitController().update(pose, published[0], t=0.0) assert (twist.linear.x, twist.linear.y, twist.angular.z) == (0.0, 0.0, 0.0) @@ -152,7 +189,15 @@ def test_hold_warns_once_per_stale_episode(monkeypatch): ) planner, _published = _holding_planner() for _ in range(3): - planner.hold(PoseStamped(position=(0.0, 0.0, 0.0)), age=7.0) + planner.hold( + PoseStamped( + header=Header(stamp=time_from_seconds(0.0), frame_id=""), + pose=Pose( + position=Point(0.0, 0.0, 0.0), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + ), + ), + age=7.0, + ) assert len(warnings) == 1 @@ -169,17 +214,25 @@ def test_a_stale_pose_is_a_missing_pose(monkeypatch): tf, clock = MultiTBuffer(), _Clock() planner._pose_src = TfPose(tf, "base_link", planner.config.max_map_age_s, clock=clock) tf.receive_transform( - Transform( - translation=Vector3(0.0, 0.0, 0.3), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="odom", + TransformStamped( + header=Header(stamp=time_from_seconds(5.0), frame_id="odom"), child_frame_id="base_link", - ts=5.0, + transform=Transform( + translation=Vector3(0.0, 0.0, 0.3), rotation=Quaternion(0.0, 0.0, 0.0, 1.0) + ), + ) + ) + planner._on_local_map( + pointcloud_from_xyz( + np.zeros((0, 3), np.float32), + header=Header(stamp=time_from_seconds(0.0), frame_id="odom"), ) ) - planner._on_local_map(PointCloud2.from_numpy(np.zeros((0, 3), np.float32), frame_id="odom")) planner._on_planner_path( - Path(frame_id="odom", poses=[pose_stamped(0.0, 0.0, 0.0), pose_stamped(3.0, 0.0, 0.0)]) + Path( + header=Header(stamp=time_from_seconds(0.0), frame_id="odom"), + poses=[pose_stamped(0.0, 0.0, 0.0), pose_stamped(3.0, 0.0, 0.0)], + ) ) planner.tick() assert len(published) == 1 @@ -230,7 +283,7 @@ def test_the_band_rides_the_body_not_the_map_origin(): def test_a_map_with_a_non_finite_return_still_plans(): planner, published = _holding_planner() room = np.concatenate([_room(-0.28), np.array([[np.nan, 0.0, -0.08]], dtype=np.float32)]) - cloud = PointCloud2.from_numpy(room, frame_id="odom") + cloud = pointcloud_from_xyz(room, header=Header(stamp=time_from_seconds(0.0), frame_id="odom")) planner._on_local_map(cloud) assert planner.plan_once(cloud, pose_stamped(0.0, 0.0, 0.0), (0.5, 0.0), ground_z=-0.28) assert len(published) == 1 and len(published[0].poses) >= 2 @@ -249,10 +302,14 @@ def _detour(emb: Embodiment, cloud: NDArray[np.float64], ground_z: float) -> flo hard = hard_points(load_model("body_band", emb), cloud, ground_z) episode = make_py(emb) episode.reset() - path = episode.plan(hard[:, :2], Pose(0.0, 0.0, 0.0), Pose(3.0, 0.0, 0.0)) + path = episode.plan( + hard[:, :2], + Pose(position=Point(0.0, 0.0, 0.0), orientation=Quaternion(0.0, 0.0, 0.0, 1.0)), + Pose(position=Point(3.0, 0.0, 0.0), orientation=Quaternion(0.0, 0.0, 0.0, 1.0)), + ) if len(path.poses) < 2: return math.inf # a refusal is the strongest form of "it saw the wall" - return max(abs(p.position.y) for p in path.poses) + return max(abs(p.pose.position.y) for p in path.poses) def test_a_tall_body_plans_around_what_the_old_band_cut_off(): diff --git a/dimos/navigation/local_planner/test_profile.py b/dimos/navigation/local_planner/test_profile.py index d987ab31757..35b4edf974a 100644 --- a/dimos/navigation/local_planner/test_profile.py +++ b/dimos/navigation/local_planner/test_profile.py @@ -15,12 +15,14 @@ import itertools import math +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Header import numpy as np +import pytest -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.navigation.embodiment.go2 import GO2 from dimos.navigation.local_planner.profile import ( decode_ceilings, @@ -34,15 +36,16 @@ def _pose(x: float, y: float, yaw: float = 0.0, ts: float = 0.0) -> PoseStamped: return PoseStamped( - ts=ts, - frame_id="odom", - position=Vector3(x, y, 0.0), - orientation=Quaternion.from_euler(Vector3(0, 0, yaw)), + header=Header(stamp=time_from_seconds(ts), frame_id="odom"), + pose=Pose(position=Point(x, y, 0.0), orientation=quaternion_from_euler(0, 0, yaw)), ) def _path(n: int = 21, step: float = 0.1) -> Path: - return Path(frame_id="odom", poses=[_pose(i * step, 0.0, ts=1.0) for i in range(n)]) + return Path( + header=Header(stamp=time_from_seconds(0.0), frame_id="odom"), + poses=[_pose(i * step, 0.0, ts=1.0) for i in range(n)], + ) def test_governor_curve_endpoints() -> None: @@ -61,7 +64,7 @@ def test_roundtrip_recovers_governor_speeds() -> None: v = governor_speed(clearance, GO2) # segment ceiling = min of endpoint speeds; waypoint i carries segment i-1 expected = np.minimum(v[:-1], v[1:]) - np.testing.assert_allclose(ceilings[1:], expected, rtol=1e-9) + np.testing.assert_allclose(ceilings[1:], expected, rtol=1e-8) assert ceilings[0] == ceilings[1] @@ -73,7 +76,8 @@ def test_tight_zone_reads_slow() -> None: ceilings = decode_ceilings(path, MIN_SPEED, MAX_SPEED) assert ceilings is not None assert ceilings[10] == np.clip(MIN_SPEED, MIN_SPEED, MAX_SPEED) - assert ceilings[3] == MAX_SPEED + # ROS time quantizes each waypoint to nanoseconds. + assert ceilings[3] == pytest.approx(MAX_SPEED, abs=1e-8) def test_unstamped_path_decodes_none() -> None: @@ -83,7 +87,7 @@ def test_unstamped_path_decodes_none() -> None: def test_non_monotone_decodes_none() -> None: path = _path() encode_precision(path, np.full(len(path.poses), 1.0), GO2) - path.poses[5].ts = 999.0 + path.poses[5].header.stamp = time_from_seconds(999.0) assert decode_ceilings(path, MIN_SPEED, MAX_SPEED) is None @@ -91,29 +95,32 @@ def test_fan_keeps_timeline_monotone_and_inherits_ceiling() -> None: poses = [_pose(0.0, 0.0, 0.0), _pose(0.1, 0.0, 0.0)] poses += [_pose(0.1, 0.0, yaw) for yaw in (0.3, 0.6, 0.9)] # rotation in place poses += [_pose(0.2, 0.0, 0.9)] - path = Path(frame_id="odom", poses=poses) + path = Path(header=Header(stamp=time_from_seconds(0.0), frame_id="odom"), poses=poses) encode_precision(path, np.full(len(poses), 1.0), GO2) - ts = [p.ts for p in path.poses] + ts = [to_seconds(p.header.stamp) for p in path.poses] assert all(b > a for a, b in itertools.pairwise(ts)) ceilings = decode_ceilings(path, MIN_SPEED, MAX_SPEED) assert ceilings is not None - assert ceilings[3] == MAX_SPEED # fan inherits, does not read as slow motion + # ROS time quantizes each waypoint to nanoseconds. + assert ceilings[3] == pytest.approx( + MAX_SPEED, abs=1e-8 + ) # fan inherits, does not read as slow motion def test_decode_clips_to_safe_band() -> None: # a foreign producer stamping absurdly fast segments cannot raise the cap path = _path(n=3) for i, p in enumerate(path.poses): - p.ts = i * 1e-4 # implies 1000 m/s + p.header.stamp = time_from_seconds(i * 1e-4) # implies 1000 m/s ceilings = decode_ceilings(path, MIN_SPEED, MAX_SPEED) assert ceilings is not None assert float(np.max(ceilings)) <= MAX_SPEED def test_empty_and_single_pose() -> None: - empty = Path(frame_id="odom", poses=[]) + empty = Path(header=Header(stamp=time_from_seconds(0.0), frame_id="odom"), poses=[]) assert decode_ceilings(encode_precision(empty, np.zeros(0), GO2), MIN_SPEED, MAX_SPEED) is None - single = Path(frame_id="odom", poses=[_pose(0, 0)]) + single = Path(header=Header(stamp=time_from_seconds(0.0), frame_id="odom"), poses=[_pose(0, 0)]) assert decode_ceilings(encode_precision(single, np.zeros(1), GO2), MIN_SPEED, MAX_SPEED) is None @@ -126,4 +133,4 @@ def test_stall_never_reads_as_schedule() -> None: ca, cb = decode_ceilings(a, MIN_SPEED, MAX_SPEED), decode_ceilings(b, MIN_SPEED, MAX_SPEED) assert ca is not None and cb is not None np.testing.assert_allclose(ca, cb) - assert math.isclose(b.poses[0].ts, 12345.6) + assert math.isclose(to_seconds(b.poses[0].header.stamp), 12345.6) diff --git a/dimos/navigation/local_planner/test_viz.py b/dimos/navigation/local_planner/test_viz.py index e0826a68607..99c3a9d2197 100644 --- a/dimos/navigation/local_planner/test_viz.py +++ b/dimos/navigation/local_planner/test_viz.py @@ -17,13 +17,14 @@ from dataclasses import replace import math +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Header import numpy as np from numpy.typing import NDArray -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.time import time_from_seconds from dimos.navigation.embodiment.go2 import GO2 from dimos.navigation.local_planner.profile import encode_precision from dimos.navigation.local_planner.viz import ( @@ -36,14 +37,16 @@ def _pose(x: float, y: float, yaw: float = 0.0) -> PoseStamped: return PoseStamped( - frame_id="odom", - position=Vector3(x, y, 0.0), - orientation=Quaternion.from_euler(Vector3(0, 0, yaw)), + header=Header(stamp=time_from_seconds(0.0), frame_id="odom"), + pose=Pose(position=Point(x, y, 0.0), orientation=quaternion_from_euler(0, 0, yaw)), ) def _plan(n: int = 40, yaw: float = 0.0) -> Path: - return Path(frame_id="odom", poses=[_pose(k * 0.1, 0.0, yaw) for k in range(n)]) + return Path( + header=Header(stamp=time_from_seconds(0.0), frame_id="odom"), + poses=[_pose(k * 0.1, 0.0, yaw) for k in range(n)], + ) def _rgba(boxes) -> NDArray[np.float64]: @@ -84,25 +87,31 @@ def test_colour_tracks_the_precision_the_planner_stamped() -> None: def test_an_unstamped_plan_draws_one_flat_colour() -> None: """No pinch is claimed where no precision was stamped. - Freshly built poses carry wall-clock stamps microseconds apart, so the - decoder clips the absurd speed to cruise rather than returning None. It - must not manufacture a varying profile out of noise: every box comes out - the same colour. + Canonical fixtures carry equal stamps until precision is encoded. + The renderer must not infer clearance from that unstamped path. """ plan = _plan() # never passed through encode_precision room = plan_clearance(plan, GO2) - assert room is not None and len(np.unique(np.round(room, 6))) == 1 + assert room is None colors = _rgba(render_body(plan)) assert len({tuple(c[:3]) for c in colors}) == 1 def test_the_veto_stub_is_drawn_rather_than_blanked() -> None: """An empty viewport reads as a dead module; a refusal must look refused.""" - veto = render_body(Path(frame_id="odom", poses=[_pose(1.0, 2.0, 0.5)])) + veto = render_body( + Path( + header=Header(stamp=time_from_seconds(0.0), frame_id="odom"), + poses=[_pose(1.0, 2.0, 0.5)], + ) + ) assert veto is not None assert len(veto.centers.pa_array) == 1 # an EMPTY path is different: hold the last good picture - assert render_body(Path(frame_id="odom", poses=[])) is None + assert ( + render_body(Path(header=Header(stamp=time_from_seconds(0.0), frame_id="odom"), poses=[])) + is None + ) def test_the_box_sits_on_the_body_not_on_the_pose_point() -> None: @@ -123,8 +132,13 @@ def test_the_plan_pins_its_frame_and_draws_its_body_off_the_one_path_topic() -> assert by_entity == {"world/path": rr.Transform3D, "world/path/body": rr.Boxes3D} assert drawn[0] == ("world/path", rr.Transform3D(parent_frame="tf#/odom")) # a refusal stub is drawn too, so it is not mistaken for a dead module - veto = render_plan(Path(frame_id="odom", poses=[_pose(1.0, 2.0)])) + veto = render_plan( + Path(header=Header(stamp=time_from_seconds(0.0), frame_id="odom"), poses=[_pose(1.0, 2.0)]) + ) assert veto is not None and veto[-1][0] == "world/path/body" # an empty path is "no route": hold the last picture rather than blank it - assert render_plan(Path(frame_id="odom", poses=[])) is None + assert ( + render_plan(Path(header=Header(stamp=time_from_seconds(0.0), frame_id="odom"), poses=[])) + is None + ) assert set(motion_visual_override(body_dilate_m=-0.03)) == {"world/path"} diff --git a/dimos/navigation/local_planner/viz.py b/dimos/navigation/local_planner/viz.py index ea2ca332b1b..7d85bcf254c 100644 --- a/dimos/navigation/local_planner/viz.py +++ b/dimos/navigation/local_planner/viz.py @@ -23,11 +23,12 @@ from functools import partial from typing import TYPE_CHECKING, Any +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.nav_msgs.msg import Path import numpy as np from numpy.typing import NDArray -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import yaw from dimos.navigation.embodiment.base import Embodiment from dimos.navigation.embodiment.go2 import GO2 from dimos.navigation.local_planner.profile import ( @@ -46,10 +47,10 @@ def _body_centre(pose: PoseStamped, center_off: float) -> tuple[float, float]: """The body centre for a pose point, offset along the pose's own heading.""" - yaw = pose.yaw + heading = yaw(pose.pose.orientation) return ( - pose.position.x + center_off * np.cos(yaw), - pose.position.y + center_off * np.sin(yaw), + pose.pose.position.x + center_off * np.cos(heading), + pose.pose.position.y + center_off * np.sin(heading), ) @@ -82,16 +83,18 @@ def render_body( p = msg.poses[0] cx, cy = _body_centre(p, center_off) return rr.Boxes3D( - centers=[[cx, cy, p.position.z + height / 2.0 + _Z_LIFT]], + centers=[[cx, cy, p.pose.position.z + height / 2.0 + _Z_LIFT]], half_sizes=[[length / 2.0, width / 2.0, height / 2.0]], - rotation_axis_angles=[rr.RotationAxisAngle([0.0, 0.0, 1.0], rr.Angle(rad=p.yaw))], + rotation_axis_angles=[ + rr.RotationAxisAngle([0.0, 0.0, 1.0], rr.Angle(rad=yaw(p.pose.orientation))) + ], colors=[[255, 60, 60, 200]], radii=[line_radius], fill_mode="majorwireframe", labels=["VETO: no safe route"], ) - xy = np.array([[p.position.x, p.position.y] for p in msg.poses]).reshape(-1, 2) + xy = np.array([[p.pose.position.x, p.pose.position.y] for p in msg.poses]).reshape(-1, 2) seg = np.linalg.norm(np.diff(xy, axis=0), axis=1) arcs = np.concatenate([[0.0], np.cumsum(seg)]) # one box per stride_m of arc, plus the last pose so the goal end is drawn @@ -103,8 +106,8 @@ def render_body( for i in picks: p = msg.poses[int(i)] cx, cy = _body_centre(p, center_off) - centers.append([cx, cy, p.position.z + height / 2.0 + _Z_LIFT]) - angles.append(rr.RotationAxisAngle([0.0, 0.0, 1.0], rr.Angle(rad=p.yaw))) + centers.append([cx, cy, p.pose.position.z + height / 2.0 + _Z_LIFT]) + angles.append(rr.RotationAxisAngle([0.0, 0.0, 1.0], rr.Angle(rad=yaw(p.pose.orientation)))) if clear is None: colors.append([100, 160, 255, 160]) # unstamped: the plan line's own blue else: @@ -139,7 +142,7 @@ def render_plan( return None out: RerunMulti = [ # the bridge only pins single-archetype entities to their tf frame; the boxes are the plan - (entity, rr.Transform3D(parent_frame=f"tf#/{msg.frame_id}")), + (entity, rr.Transform3D(parent_frame=f"tf#/{msg.header.frame_id}")), ] boxes = render_body(msg, emb, stride_m, line_radius) if boxes is not None: diff --git a/dimos/navigation/test_tf_pose.py b/dimos/navigation/test_tf_pose.py index 97a2bca2055..2458d425b9b 100644 --- a/dimos/navigation/test_tf_pose.py +++ b/dimos/navigation/test_tf_pose.py @@ -12,9 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.std_msgs.msg import Header + +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.navigation.tf_pose import TfPose from dimos.protocol.tf.tf import MultiTBuffer @@ -34,7 +35,7 @@ def get( time_tolerance: float | None = None, *, forward_tolerance: float = 0.0, - ) -> Transform | None: + ) -> TransformStamped | None: self.gets += 1 return super().get( parent_frame, @@ -53,13 +54,11 @@ def __call__(self) -> float: return self.t -def _body_at(x: float, ts: float) -> Transform: - return Transform( - translation=Vector3(x, 0.0, 0.3), - rotation=IDENTITY, - frame_id="odom", +def _body_at(x: float, ts: float) -> TransformStamped: + return TransformStamped( + header=Header(stamp=time_from_seconds(ts), frame_id="odom"), child_frame_id="base_link", - ts=ts, + transform=Transform(translation=Vector3(x, 0.0, 0.3), rotation=IDENTITY), ) @@ -68,7 +67,12 @@ def test_tf_pose_is_the_world_to_base_edge(): tf.receive_transform(_body_at(1.0, ts=5.0)) pose = TfPose(tf, "base_link", max_age_s=2.5, clock=clock).get("odom") assert pose is not None - assert (pose.position.x, pose.position.z, pose.ts, pose.frame_id) == (1.0, 0.3, 5.0, "odom") + assert ( + pose.pose.position.x, + pose.pose.position.z, + to_seconds(pose.header.stamp), + pose.header.frame_id, + ) == (1.0, 0.3, 5.0, "odom") def test_tf_pose_goes_stale_when_the_stamp_stops_advancing(): diff --git a/dimos/navigation/tf_pose.py b/dimos/navigation/tf_pose.py index a1a2f08e45d..c874f716812 100644 --- a/dimos/navigation/tf_pose.py +++ b/dimos/navigation/tf_pose.py @@ -19,10 +19,14 @@ import time from typing import TYPE_CHECKING +from dimos.msgs.geometry import pose_from_transform +from dimos.msgs.time import to_seconds + if TYPE_CHECKING: from collections.abc import Callable - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped + from dimos_generated.geometry_msgs.msg import PoseStamped + from dimos.protocol.tf.tf import TFLookup @@ -59,8 +63,8 @@ def get(self, world_frame: str) -> PoseStamped | None: self._next_lookup = now + self.RETRY_PERIOD_S self._seen = None return None - if self._seen is None or self._seen[0] != tf.ts: - self._seen = (tf.ts, now) + if self._seen is None or self._seen[0] != to_seconds(tf.header.stamp): + self._seen = (to_seconds(tf.header.stamp), now) if now - self._seen[1] > self.max_age_s: return None - return tf.to_pose(ts=tf.ts) + return pose_from_transform(tf) diff --git a/dimos/navigation/trajectory_follower/fancy/controller.py b/dimos/navigation/trajectory_follower/fancy/controller.py index eef34feb27f..6e07dda0e6e 100644 --- a/dimos/navigation/trajectory_follower/fancy/controller.py +++ b/dimos/navigation/trajectory_follower/fancy/controller.py @@ -27,12 +27,12 @@ import math from typing import Any, Protocol +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist +from dimos_generated.nav_msgs.msg import Path import numpy as np from numpy.typing import NDArray -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import yaw from dimos.protocol.service.spec import BaseConfig @@ -87,6 +87,7 @@ def path_xy_yaw(path: Path) -> NDArray[np.float64]: """The plan as a contiguous (N, 3) float64 array, the rust marshalling.""" return np.ascontiguousarray( np.array( - [[p.position.x, p.position.y, p.yaw] for p in path.poses], dtype=np.float64 + [[p.pose.position.x, p.pose.position.y, yaw(p.pose.orientation)] for p in path.poses], + dtype=np.float64, ).reshape(-1, 3) ) diff --git a/dimos/navigation/trajectory_follower/fancy/laws/hinted.py b/dimos/navigation/trajectory_follower/fancy/laws/hinted.py index f7d02853bcf..dddfcc30d83 100644 --- a/dimos/navigation/trajectory_follower/fancy/laws/hinted.py +++ b/dimos/navigation/trajectory_follower/fancy/laws/hinted.py @@ -22,13 +22,13 @@ import math from typing import Any +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist, Vector3 +from dimos_generated.nav_msgs.msg import Path import numpy as np from numpy.typing import NDArray -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import yaw +from dimos.msgs.time import to_seconds from dimos.navigation.embodiment.base import Embodiment from dimos.navigation.embodiment.go2 import GO2 from dimos.navigation.local_planner.profile import decode_ceilings @@ -76,12 +76,12 @@ def update( self, pose: PoseStamped, path: Path, t: float, clearance: NDArray[np.float64] | None = None ) -> Twist: cfg, emb = self.config, self.emb - if len(path) < 2: + if len(path.poses) < 2: # empty path or single-pose veto stub: the planner says stop return Twist(Vector3(0, 0, 0), Vector3(0, 0, 0)) - xy = np.array([[p.position.x, p.position.y] for p in path.poses]) - yaws = np.array([p.yaw for p in path.poses]) - px, py, pyaw = pose.position.x, pose.position.y, pose.yaw + xy = np.array([[p.pose.position.x, p.pose.position.y] for p in path.poses]) + yaws = np.array([yaw(p.pose.orientation) for p in path.poses]) + px, py, pyaw = pose.pose.position.x, pose.pose.position.y, yaw(pose.pose.orientation) n = len(xy) seg = np.linalg.norm(np.diff(xy, axis=0), axis=1) if n > 1 else np.zeros(1) @@ -198,10 +198,14 @@ def update( clr = None if clearance is None else np.ascontiguousarray(clearance, dtype=np.float64) # the law reads only stamp deltas: a precision profile, not a schedule ts = np.ascontiguousarray( - np.array([p.ts for p in path.poses], dtype=np.float64).reshape(-1) + np.array([to_seconds(p.header.stamp) for p in path.poses], dtype=np.float64).reshape(-1) ) vx, vy, wz = self._mod.update_hinted( - (float(pose.position.x), float(pose.position.y), float(pose.yaw)), + ( + float(pose.pose.position.x), + float(pose.pose.position.y), + float(yaw(pose.pose.orientation)), + ), path_xy_yaw(path), clr, ts, diff --git a/dimos/navigation/trajectory_follower/fancy/laws/seed.py b/dimos/navigation/trajectory_follower/fancy/laws/seed.py index 14072fe6311..edcce8c7958 100644 --- a/dimos/navigation/trajectory_follower/fancy/laws/seed.py +++ b/dimos/navigation/trajectory_follower/fancy/laws/seed.py @@ -22,13 +22,12 @@ import math from typing import Any +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist, Vector3 +from dimos_generated.nav_msgs.msg import Path import numpy as np from numpy.typing import NDArray -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import yaw from dimos.navigation.embodiment.base import Embodiment from dimos.navigation.embodiment.go2 import GO2 from dimos.navigation.trajectory_follower.fancy.controller import ( @@ -67,12 +66,12 @@ def update( self, pose: PoseStamped, path: Path, t: float, clearance: NDArray[np.float64] | None = None ) -> Twist: cfg, emb = self.config, self.emb - if len(path) < 2: + if len(path.poses) < 2: # empty path or a single-pose veto stub: hold return Twist(Vector3(0, 0, 0), Vector3(0, 0, 0)) - xy = np.array([[p.position.x, p.position.y] for p in path.poses]) - yaws = np.array([p.yaw for p in path.poses]) - px, py, pyaw = pose.position.x, pose.position.y, pose.yaw + xy = np.array([[p.pose.position.x, p.pose.position.y] for p in path.poses]) + yaws = np.array([yaw(p.pose.orientation) for p in path.poses]) + px, py, pyaw = pose.pose.position.x, pose.pose.position.y, yaw(pose.pose.orientation) seg = np.linalg.norm(np.diff(xy, axis=0), axis=1) if len(xy) > 1 else np.zeros(1) arcs = np.concatenate([[0.0], np.cumsum(seg)]) @@ -141,7 +140,11 @@ def update( ) -> Twist: clr = None if clearance is None else np.ascontiguousarray(clearance, dtype=np.float64) vx, vy, wz = self._mod.update_seed( - (float(pose.position.x), float(pose.position.y), float(pose.yaw)), + ( + float(pose.pose.position.x), + float(pose.pose.position.y), + float(yaw(pose.pose.orientation)), + ), path_xy_yaw(path), clr, self._emb, diff --git a/dimos/navigation/trajectory_follower/fancy/module.py b/dimos/navigation/trajectory_follower/fancy/module.py index 43b78873249..a62099e6156 100644 --- a/dimos/navigation/trajectory_follower/fancy/module.py +++ b/dimos/navigation/trajectory_follower/fancy/module.py @@ -26,6 +26,10 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist, Vector3 +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Bool +from dimos_generated.tf2_msgs.msg import TFMessage from pydantic import ImportString from reactivex.disposable import Disposable @@ -33,11 +37,6 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import IO, In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.std_msgs.Bool import Bool -from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.navigation import spec from dimos.navigation.embodiment.base import Embodiment from dimos.navigation.embodiment.go2 import GO2 @@ -84,7 +83,7 @@ class TrajectoryFollowerConfig(ModuleConfig): goal_tolerance: float = 0.20 # planar distance that counts as arrival (m) # The planner's own body: the law decodes the path stamps with its governor band. embodiment: Embodiment = GO2 - # The pose is the `path.frame_id -> base_frame` edge on tf, read each tick. + # The pose is the `path.header.frame_id -> base_frame` edge on tf, read each tick. base_frame: str = "base_link" # Deadman: zero the twist once the held path is this old, measured from arrival; must clear the ~1 Hz replan cadence. max_path_age_s: float = 2.5 @@ -129,7 +128,9 @@ def stop(self) -> None: self._stop_event.set() if self._thread is not None: self._thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) - self.nav_cmd_vel.publish(Twist()) + self.nav_cmd_vel.publish( + Twist(linear=Vector3(0.0, 0.0, 0.0), angular=Vector3(0.0, 0.0, 0.0)) + ) super().stop() def _on_path(self, msg: Path) -> None: @@ -144,10 +145,14 @@ def _on_path(self, msg: Path) -> None: self._path_at = time.monotonic() if len(msg.poses) >= 2: # a single-pose stub is a refusal, not an arrival target - self._latch.set_goal((msg.poses[-1].position.x, msg.poses[-1].position.y)) + self._latch.set_goal( + (msg.poses[-1].pose.position.x, msg.poses[-1].pose.position.y) + ) if not msg.poses: # now, not a control period later - self.nav_cmd_vel.publish(Twist()) + self.nav_cmd_vel.publish( + Twist(linear=Vector3(0.0, 0.0, 0.0), angular=Vector3(0.0, 0.0, 0.0)) + ) def _control_loop(self) -> None: period = 1.0 / self.config.control_frequency @@ -165,29 +170,37 @@ def tick(self) -> None: age = None if self._path_at is None else now - self._path_at pose = None if path is not None and self._pose_src is not None: - pose = self._pose_src.get(path.frame_id) + pose = self._pose_src.get(path.header.frame_id) if path is not None and pose is not None: assert age is not None self.step(pose, path, age) elif path is not None: # a plan with no live pose under it: the deadman on the pose - self.nav_cmd_vel.publish(Twist()) + self.nav_cmd_vel.publish( + Twist(linear=Vector3(0.0, 0.0, 0.0), angular=Vector3(0.0, 0.0, 0.0)) + ) def step(self, pose: PoseStamped, path: Path, age: float) -> None: """One control tick against a path that arrived `age` seconds ago.""" # the deadman outranks arrival: a goal reached against an unrefreshed plan is a coincidence if age > self.config.max_path_age_s: - self.nav_cmd_vel.publish(Twist()) + self.nav_cmd_vel.publish( + Twist(linear=Vector3(0.0, 0.0, 0.0), angular=Vector3(0.0, 0.0, 0.0)) + ) return - xy = (pose.position.x, pose.position.y) + xy = (pose.pose.position.x, pose.pose.position.y) with self._lock: arrived, reached = self._latch.arrive(xy), self._latch.reached if arrived: - self.nav_cmd_vel.publish(Twist()) + self.nav_cmd_vel.publish( + Twist(linear=Vector3(0.0, 0.0, 0.0), angular=Vector3(0.0, 0.0, 0.0)) + ) self.goal_reached.publish(Bool(True)) logger.info("goal reached", x=round(xy[0], 2), y=round(xy[1], 2)) return if reached: - self.nav_cmd_vel.publish(Twist()) + self.nav_cmd_vel.publish( + Twist(linear=Vector3(0.0, 0.0, 0.0), angular=Vector3(0.0, 0.0, 0.0)) + ) return self.nav_cmd_vel.publish(self._controller.update(pose, path, time.monotonic())) diff --git a/dimos/navigation/trajectory_follower/fancy/native.py b/dimos/navigation/trajectory_follower/fancy/native.py index 70ef137c84d..9ecbcf7e2e0 100644 --- a/dimos/navigation/trajectory_follower/fancy/native.py +++ b/dimos/navigation/trajectory_follower/fancy/native.py @@ -27,12 +27,13 @@ from typing import TYPE_CHECKING, Any +from dimos_generated.geometry_msgs.msg import Twist +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Bool +from dimos_generated.tf2_msgs.msg import TFMessage + from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.stream import IO, In, Out -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.std_msgs.Bool import Bool -from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.navigation import spec from dimos.navigation.embodiment.base import Embodiment from dimos.navigation.trajectory_follower.fancy.module import TrajectoryFollowerConfig diff --git a/dimos/navigation/trajectory_follower/fancy/rust/Cargo.lock b/dimos/navigation/trajectory_follower/fancy/rust/Cargo.lock index 82335d66aa3..7c233769417 100644 --- a/dimos/navigation/trajectory_follower/fancy/rust/Cargo.lock +++ b/dimos/navigation/trajectory_follower/fancy/rust/Cargo.lock @@ -226,6 +226,20 @@ name = "bytemuck" version = "1.25.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a1f896587b6f2c069c73d2f0913e2d590c3990285cd2f0b6aa02b786b4c679c" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.6", +] [[package]] name = "byteorder" @@ -580,9 +594,19 @@ dependencies = [ ] [[package]] -name = "dimos-lcm" +name = "dimos-generated-messages" +version = "0.1.0" +dependencies = [ + "heapless", + "re_cdr", + "ros2msg", + "serde", + "serde-big-array", +] + +[[package]] +name = "dimos-lcm-transport" version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" dependencies = [ "byteorder", "socket2 0.5.10", @@ -593,8 +617,8 @@ dependencies = [ name = "dimos-local-planner" version = "0.1.0" dependencies = [ + "dimos-generated-messages", "dimos-module", - "lcm-msgs", "serde", "serde_json", "tokio", @@ -607,11 +631,12 @@ dependencies = [ name = "dimos-module" version = "0.1.0" dependencies = [ - "dimos-lcm", + "dimos-generated-messages", + "dimos-lcm-transport", "dimos-module-macros", - "lcm-msgs", "nalgebra", "rayon", + "re_cdr", "serde", "serde_json", "tokio", @@ -636,10 +661,10 @@ dependencies = [ name = "dimos-trajectory-follower" version = "0.1.0" dependencies = [ + "dimos-generated-messages", "dimos-local-planner", "dimos-module", "dimos-trajectory-follower", - "lcm-msgs", "numpy", "pyo3", "serde", @@ -981,6 +1006,15 @@ version = "0.33.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e762b9b188634fc508d4ec54b62ea3d352c43fd431d18d05d46f559f217f4725" +[[package]] +name = "hash32" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47d60b12902ba28e2730cd37e95b8c9223af2808df9e902d4df49588d1470606" +dependencies = [ + "byteorder", +] + [[package]] name = "hashbrown" version = "0.12.3" @@ -1023,6 +1057,17 @@ version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +[[package]] +name = "heapless" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bfb9eb618601c89945a70e254898da93b13be0388091d42117462b265bb3fad" +dependencies = [ + "hash32", + "serde", + "stable_deref_trait", +] + [[package]] name = "heck" version = "0.5.0" @@ -1428,14 +1473,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" -[[package]] -name = "lcm-msgs" -version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" -dependencies = [ - "byteorder", -] - [[package]] name = "libc" version = "0.2.189" @@ -2012,6 +2049,16 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + [[package]] name = "proc-macro-crate" version = "3.5.0" @@ -2312,6 +2359,18 @@ dependencies = [ "yasna", ] +[[package]] +name = "re_cdr" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bba10c21433b597481f366b978cdf2617e56785b38cb0dcc3833c6b8ad7e061a" +dependencies = [ + "bytemuck", + "byteorder", + "serde", + "thiserror", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -2427,6 +2486,23 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "ros2msg" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fe291be5cb6d4b8e6bf9c9de80c40bd1cb09ecac50deafa2fdcae0027e04988" +dependencies = [ + "heck", + "pest", + "pest_derive", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "syn 2.0.119", + "thiserror", +] + [[package]] name = "rustc-hash" version = "2.1.3" @@ -2667,6 +2743,15 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "serde-big-array" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11fc7cc2c76d73e0f27ee52abbd64eec84d46f370c88371120433196934e4b7f" +dependencies = [ + "serde", +] + [[package]] name = "serde_core" version = "1.0.229" diff --git a/dimos/navigation/trajectory_follower/fancy/rust/Cargo.toml b/dimos/navigation/trajectory_follower/fancy/rust/Cargo.toml index 03ec5c9ba16..032706f59e9 100644 --- a/dimos/navigation/trajectory_follower/fancy/rust/Cargo.toml +++ b/dimos/navigation/trajectory_follower/fancy/rust/Cargo.toml @@ -21,7 +21,7 @@ required-features = ["module"] dimos-local-planner = { path = "../../../local_planner/rust", default-features = false } serde = { version = "1", features = ["derive"], optional = true } dimos-module = { path = "../../../../../native/rust/dimos-module", optional = true } -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen", optional = true } +dimos-generated-messages = { path = "../../../../../packages/dimos-generated/src/dimos_generated_schemas/package/rust", optional = true } tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal", "time"], optional = true } tracing = { version = "0.1", optional = true } tracing-subscriber = { version = "0.3", features = ["env-filter"], optional = true } @@ -35,7 +35,7 @@ serde_json = { version = "1", features = ["float_roundtrip"], optional = true } # Python; maturin builds with defaults and gets the extension module. default = [] python = ["dep:pyo3", "dep:numpy", "dep:serde_json"] -module = ["dimos-local-planner/module", "dep:serde", "dep:dimos-module", "dep:lcm-msgs", "dep:tokio", "dep:tracing", "dep:tracing-subscriber", "dep:validator"] +module = ["dimos-local-planner/module", "dep:serde", "dep:dimos-module", "dep:dimos-generated-messages", "dep:tokio", "dep:tracing", "dep:tracing-subscriber", "dep:validator"] [dev-dependencies] dimos-local-planner = { path = "../../../local_planner/rust", default-features = false, features = ["test-bodies"] } @@ -48,12 +48,12 @@ python = "dimos.navigation.trajectory_follower.fancy.native:TrajectoryFollowerNa threads = 2 [package.metadata.dimos.module.trajectory_follower.inputs] -path = "nav_msgs.Path" -tf = "tf2_msgs.TFMessage" +path = "nav_msgs/msg/Path" +tf = "tf2_msgs/msg/TFMessage" [package.metadata.dimos.module.trajectory_follower.outputs] -nav_cmd_vel = "geometry_msgs.Twist" -goal_reached = "std_msgs.Bool" +nav_cmd_vel = "geometry_msgs/msg/Twist" +goal_reached = "std_msgs/msg/Bool" [profile.release] lto = "thin" diff --git a/dimos/navigation/trajectory_follower/fancy/rust/src/module/follower.rs b/dimos/navigation/trajectory_follower/fancy/rust/src/module/follower.rs index 2a3effa75ae..51a5448a264 100644 --- a/dimos/navigation/trajectory_follower/fancy/rust/src/module/follower.rs +++ b/dimos/navigation/trajectory_follower/fancy/rust/src/module/follower.rs @@ -20,11 +20,12 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use crate::laws::hinted::update as hinted_update; +use dimos_generated_messages::geometry_msgs::msg::twist::Twist; +use dimos_generated_messages::nav_msgs::msg::path::Path; +use dimos_generated_messages::std_msgs::msg::bool::Bool; use dimos_local_planner::planner::Emb; +use dimos_module::cdr; use dimos_module::{native_config, Input, Module, Output, Tf}; -use lcm_msgs::geometry_msgs::Twist; -use lcm_msgs::nav_msgs::Path; -use lcm_msgs::std_msgs::Bool; // serde derives come in through #[native_config] use tracing::info; @@ -118,13 +119,13 @@ struct Shared { #[derive(Module)] #[module(name = "trajectory_follower", setup = spawn_worker, teardown = stop_worker)] pub struct TrajectoryFollower { - #[input(decode = Path::decode, handler = on_path)] + #[input(decode = cdr::decode, handler = on_path)] path: Input, - #[output(encode = Twist::encode)] + #[output(encode = cdr::encode)] nav_cmd_vel: Output, - #[output(encode = Bool::encode)] + #[output(encode = cdr::encode)] goal_reached: Output, #[config] diff --git a/dimos/navigation/trajectory_follower/fancy/test_controller.py b/dimos/navigation/trajectory_follower/fancy/test_controller.py index 167029c2c8f..5af770baf3b 100644 --- a/dimos/navigation/trajectory_follower/fancy/test_controller.py +++ b/dimos/navigation/trajectory_follower/fancy/test_controller.py @@ -15,14 +15,15 @@ from dataclasses import replace import math +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Header import numpy as np from pydantic import ValidationError import pytest -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.time import time_from_seconds from dimos.navigation.embodiment.go2 import GO2 from dimos.navigation.trajectory_follower.fancy.controller import ControllerConfig from dimos.navigation.trajectory_follower.fancy.laws.seed import PursuitController @@ -30,21 +31,23 @@ def _pose(x: float, y: float, yaw: float = 0.0) -> PoseStamped: return PoseStamped( - frame_id="world", - position=Vector3(x, y, 0.0), - orientation=Quaternion.from_euler(Vector3(0, 0, yaw)), + header=Header(stamp=time_from_seconds(0.0), frame_id="world"), + pose=Pose(position=Point(x, y, 0.0), orientation=quaternion_from_euler(0, 0, yaw)), ) def _straight_path(n: int = 40, step: float = 0.1) -> Path: - return Path(frame_id="world", poses=[_pose(i * step, 0.0) for i in range(n)]) + return Path( + header=Header(stamp=time_from_seconds(0.0), frame_id="world"), + poses=[_pose(i * step, 0.0) for i in range(n)], + ) def _fan_path() -> Path: """Rotate in place at the origin: yaw steps with zero displacement.""" yaws = [0.0, 0.3, 0.6, 0.9, 1.2, 1.5] poses = [_pose(0.0, 0.0, y) for y in yaws] + [_pose(0.1, 0.0, 1.5), _pose(1.0, 0.1, 1.5)] - return Path(frame_id="world", poses=poses) + return Path(header=Header(stamp=time_from_seconds(0.0), frame_id="world"), poses=poses) def test_on_path_drives_forward() -> None: @@ -84,7 +87,11 @@ def test_fan_done_resumes_translation() -> None: def test_empty_path_stops() -> None: - tw = PursuitController().update(_pose(0.0, 0.0), Path(frame_id="world", poses=[]), 0.0) + tw = PursuitController().update( + _pose(0.0, 0.0), + Path(header=Header(stamp=time_from_seconds(0.0), frame_id="world"), poses=[]), + 0.0, + ) assert tw.linear.x == 0.0 and tw.angular.z == 0.0 @@ -105,7 +112,7 @@ def test_config_rejects_unknown_fields() -> None: def test_governor_creeps_in_tight_room() -> None: path = _straight_path() - tight = np.full(len(path), 0.06) # barely above the precision floor + tight = np.full(len(path.poses), 0.06) # barely above the precision floor tw = PursuitController().update(_pose(0.0, 0.0), path, 0.0, clearance=tight) # within the bottom twentieth of the governor's band, whatever the ceiling is creep = GO2.min_speed + 0.05 * (GO2.max_speed - GO2.min_speed) @@ -114,7 +121,7 @@ def test_governor_creeps_in_tight_room() -> None: def test_governor_full_speed_in_open_room() -> None: path = _straight_path() - wide = np.full(len(path), 1.0) + wide = np.full(len(path.poses), 1.0) tw_open = PursuitController().update(_pose(-1.0, 0.0), path, 0.0, clearance=wide) tw_bare = PursuitController().update(_pose(-1.0, 0.0), path, 0.0) assert math.hypot(tw_open.linear.x, tw_open.linear.y) == pytest.approx( @@ -124,7 +131,7 @@ def test_governor_full_speed_in_open_room() -> None: def test_governor_reads_room_ahead_not_behind() -> None: path = _straight_path() # 4 m of path - clear = np.full(len(path), 1.0) + clear = np.full(len(path.poses), 1.0) clear[:5] = 0.06 # tight patch already behind the lookahead window tw = PursuitController().update(_pose(1.5, 0.0), path, 0.0, clearance=clear) assert math.hypot(tw.linear.x, tw.linear.y) > GO2.min_speed + 0.1 diff --git a/dimos/navigation/trajectory_follower/fancy/test_module.py b/dimos/navigation/trajectory_follower/fancy/test_module.py index 24885fade9a..a59eed9d680 100644 --- a/dimos/navigation/trajectory_follower/fancy/test_module.py +++ b/dimos/navigation/trajectory_follower/fancy/test_module.py @@ -13,15 +13,23 @@ # limitations under the License. +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from dimos.core.module import ModuleConfig -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path +from dimos.core.transport_factory import rpc_backend +from dimos.msgs.time import time_from_seconds from dimos.navigation.local_planner.obstacles import path_clearance from dimos.navigation.tf_pose import TfPose from dimos.navigation.trajectory_follower.fancy.laws import seed @@ -37,6 +45,12 @@ _BUILT: list[TrajectoryFollower] = [] +@pytest.fixture(autouse=True) +def isolate_rpc_boundary(mocker): + for method in ("start", "serve_module_rpc", "stop"): + mocker.patch.object(rpc_backend(), method) + + @pytest.fixture(autouse=True) def _stop_modules(): yield @@ -91,16 +105,28 @@ def _follower(**config): def _base_at(z: float) -> PoseStamped: """The tf-resolved base pose; the ground sits emb.base_height under it.""" - return PoseStamped(ts=0.0, frame_id="odom", position=Vector3(0.0, 0.0, z)) + return PoseStamped( + header=Header(stamp=time_from_seconds(0.0), frame_id="odom"), + pose=Pose(position=Point(0.0, 0.0, z), orientation=Quaternion(0.0, 0.0, 0.0, 1.0)), + ) def _straight_path(end_x: float = 0.5) -> Path: return Path( - ts=0.0, - frame_id="odom", + header=Header(stamp=time_from_seconds(0.0), frame_id="odom"), poses=[ - PoseStamped(ts=0.0, frame_id="odom", position=Vector3(0.0, 0.0, 0.0)), - PoseStamped(ts=0.0, frame_id="odom", position=Vector3(end_x, 0.0, 0.0)), + PoseStamped( + header=Header(stamp=time_from_seconds(0.0), frame_id="odom"), + pose=Pose( + position=Point(0.0, 0.0, 0.0), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + ), + ), + PoseStamped( + header=Header(stamp=time_from_seconds(0.0), frame_id="odom"), + pose=Pose( + position=Point(end_x, 0.0, 0.0), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + ), + ), ], ) @@ -154,7 +180,10 @@ def test_a_stale_path_outranks_an_arrival(): out, reached = _driven(follower), _reached(follower) path = _straight_path(0.5) follower._on_path(path) - at_goal = PoseStamped(ts=0.0, frame_id="odom", position=Vector3(0.5, 0.0, 0.01)) + at_goal = PoseStamped( + header=Header(stamp=time_from_seconds(0.0), frame_id="odom"), + pose=Pose(position=Point(0.5, 0.0, 0.01), orientation=Quaternion(0.0, 0.0, 0.0, 1.0)), + ) follower.step(at_goal, path, age=9.0) assert out[-1].linear.x == 0.0 assert not reached @@ -175,12 +204,12 @@ def _on_tf(follower: TrajectoryFollower, frame: str = "odom") -> tuple[MultiTBuf tf, clock = MultiTBuffer(), _Clock() follower._pose_src = TfPose(tf, "base_link", follower.config.max_path_age_s, clock=clock) tf.receive_transform( - Transform( - translation=Vector3(0.0, 0.0, 0.01), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id=frame, + TransformStamped( + header=Header(stamp=time_from_seconds(5.0), frame_id=frame), child_frame_id="base_link", - ts=5.0, + transform=Transform( + translation=Vector3(0.0, 0.0, 0.01), rotation=Quaternion(0.0, 0.0, 0.0, 1.0) + ), ) ) return tf, clock @@ -193,7 +222,12 @@ def test_the_pose_comes_off_tf_in_the_path_frame(heard): follower._on_path(_straight_path(5.0)) # an odom plan; the edge on tf is map -> base_link follower.tick() assert (out[-1].linear.x, out[-1].linear.y, out[-1].angular.z) == (0.0, 0.0, 0.0) - follower._on_path(Path(ts=0.0, frame_id="map", poses=_straight_path(5.0).poses)) + follower._on_path( + Path( + header=Header(stamp=time_from_seconds(0.0), frame_id="map"), + poses=_straight_path(5.0).poses, + ) + ) clock.t += TfPose.RETRY_PERIOD_S # the miss parked lookups for a period follower.tick() assert out[-1].linear.x > 0.0 @@ -218,7 +252,9 @@ def test_an_empty_path_drops_the_plan_and_zeroes_the_twist(): follower._on_path(path) follower.step(pose, path, age=0.0) assert out[-1].linear.x > 0.0 - follower._on_path(Path(ts=0.0, frame_id=path.frame_id, poses=[])) + follower._on_path( + Path(header=Header(stamp=time_from_seconds(0.0), frame_id=path.header.frame_id), poses=[]) + ) assert (out[-1].linear.x, out[-1].linear.y, out[-1].angular.z) == (0.0, 0.0, 0.0) assert follower._path is None diff --git a/dimos/navigation/trajectory_follower/fancy/test_rust_parity.py b/dimos/navigation/trajectory_follower/fancy/test_rust_parity.py index 650a7b952c1..164f16ef8b2 100644 --- a/dimos/navigation/trajectory_follower/fancy/test_rust_parity.py +++ b/dimos/navigation/trajectory_follower/fancy/test_rust_parity.py @@ -21,15 +21,19 @@ from dataclasses import replace import math +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.time import time_from_seconds, to_seconds + pytest.importorskip("dimos_trajectory_follower") -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.nav_msgs.msg import Path + from dimos.navigation.embodiment.go2 import GO2 from dimos.navigation.local_planner.obstacles import path_clearance as follower_clearance from dimos.navigation.local_planner.profile import encode_precision @@ -51,14 +55,16 @@ def _pose(x: float, y: float, yaw: float = 0.0) -> PoseStamped: return PoseStamped( - frame_id="world", - position=Vector3(x, y, 0.0), - orientation=Quaternion.from_euler(Vector3(0, 0, yaw)), + header=Header(stamp=time_from_seconds(0.0), frame_id="world"), + pose=Pose(position=Point(x, y, 0.0), orientation=quaternion_from_euler(0, 0, yaw)), ) def _path(states: list[tuple[float, float, float]]) -> Path: - return Path(frame_id="world", poses=[_pose(*s) for s in states]) + return Path( + header=Header(stamp=time_from_seconds(0.0), frame_id="world"), + poses=[_pose(*s) for s in states], + ) def _straight(rng: np.random.Generator) -> list[tuple[float, float, float]]: @@ -141,7 +147,9 @@ def _cases(seed: int = 20260802, n: int = CASES): # type: ignore[no-untyped-def encode_precision(path, enc, GO2, t0=float(srng.uniform(0.0, 1e9))) if k % 8 == 3: for q in path.poses: - q.ts = 5.0 # flat: not the dialect, both must ignore it + q.header.stamp = time_from_seconds( + 5.0 + ) # flat: not the dialect, both must ignore it # non-default gains and plant every fifth case: the params tuple order is load-bearing emb = GO2 @@ -199,9 +207,9 @@ def test_rust_matches_python(law: str) -> None: seen["held"] += 1 if abs(a[2]) >= emb.max_yaw_rate - 1e-12 or math.hypot(a[0], a[1]) >= top - 1e-12: seen["clamped"] += 1 - if case[3] is not None and len(case[3]) == len(case[2]): + if case[3] is not None and len(case[3]) == len(case[2].poses): seen["governed"] += 1 - if len(case[2]) > 2 and math.hypot(a[0], a[1]) < 1e-3 and abs(a[2]) > 1e-3: + if len(case[2].poses) > 2 and math.hypot(a[0], a[1]) < 1e-3 and abs(a[2]) > 1e-3: seen["fan"] += 1 # the sweep is only worth its tolerance if it actually reached the branches assert all(v > 0 for v in seen.values()), f"unexercised branches: {seen}" @@ -243,7 +251,7 @@ def test_encode_precision_matches_python() -> None: clr = np.asarray([] if clearance is None else clearance, dtype=np.float64) # a t0 well off zero: an offset that cancels in the diff would hide a base-stamp divergence t0 = 1754212345.75 - want = [p.ts for p in encode_precision(path, clr, GO2, t0=t0).poses] + want = [to_seconds(p.header.stamp) for p in encode_precision(path, clr, GO2, t0=t0).poses] got = rs.encode_precision(path_xy_yaw(path), clr if len(clr) else None, t0, GOVERNOR) assert len(got) == len(want) for k, (x, y) in enumerate(zip(want, got, strict=True)): diff --git a/dimos/perception/depth2depth_cloud/module.py b/dimos/perception/depth2depth_cloud/module.py index b078a32ac89..26a26eb726f 100644 --- a/dimos/perception/depth2depth_cloud/module.py +++ b/dimos/perception/depth2depth_cloud/module.py @@ -19,14 +19,12 @@ from pathlib import Path from typing import TYPE_CHECKING +from dimos_generated.sensor_msgs.msg import CameraInfo, CompressedImage, PointCloud2 +from dimos_generated.tf2_msgs.msg import TFMessage from pydantic import Field, field_validator, model_validator from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.CompressedImage import CompressedImage -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage def has_nvidia_gpu() -> bool: diff --git a/dimos/perception/depth2depth_cloud/rust/Cargo.lock b/dimos/perception/depth2depth_cloud/rust/Cargo.lock index 6064e9c4a45..b9e3d7a3ed6 100644 --- a/dimos/perception/depth2depth_cloud/rust/Cargo.lock +++ b/dimos/perception/depth2depth_cloud/rust/Cargo.lock @@ -831,8 +831,8 @@ name = "dimos-depth2depth-cloud" version = "0.1.0" dependencies = [ "depth2depth", + "dimos-generated-messages", "dimos-module", - "lcm-msgs", "nalgebra", "rayon", "serde", @@ -843,9 +843,19 @@ dependencies = [ ] [[package]] -name = "dimos-lcm" +name = "dimos-generated-messages" +version = "0.1.0" +dependencies = [ + "heapless", + "re_cdr", + "ros2msg", + "serde", + "serde-big-array", +] + +[[package]] +name = "dimos-lcm-transport" version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" dependencies = [ "byteorder", "socket2 0.5.10", @@ -856,11 +866,12 @@ dependencies = [ name = "dimos-module" version = "0.1.0" dependencies = [ - "dimos-lcm", + "dimos-generated-messages", + "dimos-lcm-transport", "dimos-module-macros", - "lcm-msgs", "nalgebra", "rayon", + "re_cdr", "serde", "serde_json", "tokio", @@ -1614,6 +1625,15 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "hash32" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47d60b12902ba28e2730cd37e95b8c9223af2808df9e902d4df49588d1470606" +dependencies = [ + "byteorder", +] + [[package]] name = "hashbrown" version = "0.12.3" @@ -1658,6 +1678,17 @@ version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +[[package]] +name = "heapless" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bfb9eb618601c89945a70e254898da93b13be0388091d42117462b265bb3fad" +dependencies = [ + "hash32", + "serde", + "stable_deref_trait", +] + [[package]] name = "heck" version = "0.5.0" @@ -2073,14 +2104,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" -[[package]] -name = "lcm-msgs" -version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" -dependencies = [ - "byteorder", -] - [[package]] name = "libc" version = "0.2.189" @@ -2826,6 +2849,16 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + [[package]] name = "proc-macro-crate" version = "3.5.0" @@ -3160,6 +3193,18 @@ dependencies = [ "yasna", ] +[[package]] +name = "re_cdr" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bba10c21433b597481f366b978cdf2617e56785b38cb0dcc3833c6b8ad7e061a" +dependencies = [ + "bytemuck", + "byteorder", + "serde", + "thiserror 2.0.20", +] + [[package]] name = "reborrow" version = "0.5.5" @@ -3282,6 +3327,23 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "ros2msg" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fe291be5cb6d4b8e6bf9c9de80c40bd1cb09ecac50deafa2fdcae0027e04988" +dependencies = [ + "heck", + "pest", + "pest_derive", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "syn 2.0.119", + "thiserror 2.0.20", +] + [[package]] name = "rustc-hash" version = "2.1.3" @@ -3564,6 +3626,15 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "serde-big-array" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11fc7cc2c76d73e0f27ee52abbd64eec84d46f370c88371120433196934e4b7f" +dependencies = [ + "serde", +] + [[package]] name = "serde_core" version = "1.0.229" diff --git a/dimos/perception/depth2depth_cloud/rust/Cargo.toml b/dimos/perception/depth2depth_cloud/rust/Cargo.toml index 2d44f1e431c..cb6d5b6a4d8 100644 --- a/dimos/perception/depth2depth_cloud/rust/Cargo.toml +++ b/dimos/perception/depth2depth_cloud/rust/Cargo.toml @@ -23,17 +23,17 @@ path = "dimos_depth2depth_cloud::module::Depth2DepthCloud" python = "dimos.perception.depth2depth_cloud.module:Depth2DepthCloud" [package.metadata.dimos.module.depth2depth_cloud.inputs] -image = "sensor_msgs.CompressedImage" -camera_info = "sensor_msgs.CameraInfo" -lidar = "sensor_msgs.PointCloud2" -tf = "tf2_msgs.TFMessage" +image = "sensor_msgs/msg/CompressedImage" +camera_info = "sensor_msgs/msg/CameraInfo" +lidar = "sensor_msgs/msg/PointCloud2" +tf = "tf2_msgs/msg/TFMessage" [package.metadata.dimos.module.depth2depth_cloud.outputs] -depth_cloud = "sensor_msgs.PointCloud2" +depth_cloud = "sensor_msgs/msg/PointCloud2" [dependencies] dimos-module = { path = "../../../../native/rust/dimos-module" } -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +dimos-generated-messages = { path = "../../../../packages/dimos-generated/src/dimos_generated_schemas/package/rust" } tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal"] } serde = { version = "1", features = ["derive"] } tracing = "0.1" diff --git a/dimos/perception/depth2depth_cloud/rust/flake.nix b/dimos/perception/depth2depth_cloud/rust/flake.nix index 6eef7889cb5..a504cd98b3f 100644 --- a/dimos/perception/depth2depth_cloud/rust/flake.nix +++ b/dimos/perception/depth2depth_cloud/rust/flake.nix @@ -39,6 +39,9 @@ mkdir -p $out/native/rust cp -r ${../../../../native/rust/dimos-module} $out/native/rust/dimos-module cp -r ${../../../../native/rust/dimos-module-macros} $out/native/rust/dimos-module-macros + cp -r ${../../../../native/rust/dimos-lcm-transport} $out/native/rust/dimos-lcm-transport + mkdir -p $out/packages/dimos-generated/src/dimos_generated_schemas/package + cp -r ${../../../../packages/dimos-generated/src/dimos_generated_schemas/package/rust} $out/packages/dimos-generated/src/dimos_generated_schemas/package/rust ''; generatedCargoNix = crate2nix.tools.${system}.generatedCargoNix { diff --git a/dimos/perception/depth2depth_cloud/rust/src/module.rs b/dimos/perception/depth2depth_cloud/rust/src/module.rs index 3e40d6498c4..000f9068b84 100644 --- a/dimos/perception/depth2depth_cloud/rust/src/module.rs +++ b/dimos/perception/depth2depth_cloud/rust/src/module.rs @@ -31,10 +31,15 @@ use std::time::{Duration, Instant}; use depth2depth::{ Calibration, CalibrationConfig, CloudOptions, Config as ModelConfig, Depth2Depth, }; +use dimos_generated_messages::builtin_interfaces::msg::time::Time; +use dimos_generated_messages::sensor_msgs::msg::camera_info::CameraInfo; +use dimos_generated_messages::sensor_msgs::msg::compressed_image::CompressedImage; +use dimos_generated_messages::sensor_msgs::msg::point_cloud2::PointCloud2; +use dimos_generated_messages::sensor_msgs::msg::point_field::PointField; +use dimos_generated_messages::std_msgs::msg::header::Header; +use dimos_module::cdr; use dimos_module::pointcloud::extract_xyz; use dimos_module::{native_config, warn_throttled, Input, Module, Output, Tf}; -use lcm_msgs::sensor_msgs::{CameraInfo, CompressedImage, PointCloud2, PointField}; -use lcm_msgs::std_msgs::{Header, Time}; use nalgebra::{Isometry3, Point3}; use tokio::sync::Notify; use tokio::task::JoinHandle; @@ -122,19 +127,19 @@ pub struct Config { #[derive(Module)] #[module(name = "depth2depth_cloud", setup = start, teardown = stop)] pub struct Depth2DepthCloud { - #[input(decode = CompressedImage::decode, handler = on_image)] + #[input(decode = cdr::decode, handler = on_image)] image: Input, - #[input(decode = CameraInfo::decode, handler = on_camera_info)] + #[input(decode = cdr::decode, handler = on_camera_info)] camera_info: Input, - #[input(decode = PointCloud2::decode, handler = on_lidar)] + #[input(decode = cdr::decode, handler = on_lidar)] lidar: Input, #[tf] tf: Tf, - #[output(encode = PointCloud2::encode)] + #[output(encode = cdr::encode)] depth_cloud: Output, #[config] @@ -418,8 +423,8 @@ impl Worker { }; let decoded = Instant::now(); if !undistort.as_ref().is_some_and(|(held, held_width, _)| { - held.K == info.K - && held.D == info.D + held.k == info.k + && held.d == info.d && held.distortion_model == info.distortion_model && held.width == info.width && *held_width == width @@ -712,10 +717,10 @@ fn decode_rgb(image: &CompressedImage, scale: usize) -> Option<(Vec, usize, } fn make_cloud(points: &[[f32; 3]], source: &Header, frame_id: String) -> PointCloud2 { - let field = |name: &str, offset: i32| PointField { + let field = |name: &str, offset: u32| PointField { name: name.into(), offset, - datatype: PointField::FLOAT32 as u8, + datatype: PointField::FLOAT32, count: 1, }; let data: Vec = points @@ -725,16 +730,15 @@ fn make_cloud(points: &[[f32; 3]], source: &Header, frame_id: String) -> PointCl .collect(); PointCloud2 { header: Header { - seq: source.seq, stamp: source.stamp.clone(), frame_id, }, height: 1, - width: points.len() as i32, + width: points.len() as u32, fields: vec![field("x", 0), field("y", 4), field("z", 8)], is_bigendian: false, point_step: 12, - row_step: 12 * points.len() as i32, + row_step: 12 * points.len() as u32, data, is_dense: true, } @@ -754,7 +758,7 @@ fn or_derived(configured: usize, derived: usize) -> usize { } fn seconds(stamp: &Time) -> f64 { - stamp.sec as f64 + stamp.nsec as f64 * 1e-9 + stamp.sec as f64 + stamp.nanosec as f64 * 1e-9 } /// Config, then calibration, then image frame, since drivers stamp frames nobody publishes a transform for. @@ -845,8 +849,12 @@ mod tests { #[test] fn a_non_jpeg_format_is_refused_rather_than_fed_to_the_decoder() { let png = CompressedImage { + header: Header { + stamp: Time { sec: 0, nanosec: 0 }, + frame_id: String::new(), + }, format: "png".into(), - ..Default::default() + data: Vec::new(), }; assert!(decode_rgb(&png, 1).is_none()); } @@ -854,8 +862,7 @@ mod tests { #[test] fn the_cloud_carries_the_points_and_the_source_stamp() { let header = Header { - seq: 7, - stamp: Time { sec: 3, nsec: 4 }, + stamp: Time { sec: 3, nanosec: 4 }, frame_id: "image".into(), }; let cloud = make_cloud( @@ -867,7 +874,7 @@ mod tests { ( cloud.width, cloud.header.frame_id.as_str(), - cloud.header.stamp.nsec + cloud.header.stamp.nanosec ), (2, "camera", 4) ); diff --git a/dimos/perception/depth2depth_cloud/rust/src/undistort.rs b/dimos/perception/depth2depth_cloud/rust/src/undistort.rs index 41751feff60..23f1132e9ab 100644 --- a/dimos/perception/depth2depth_cloud/rust/src/undistort.rs +++ b/dimos/perception/depth2depth_cloud/rust/src/undistort.rs @@ -16,7 +16,7 @@ //! projection both assume straight lines stay straight. use depth2depth::Pinhole; -use lcm_msgs::sensor_msgs::CameraInfo; +use dimos_generated_messages::sensor_msgs::msg::camera_info::CameraInfo; use rayon::prelude::*; /// Pinhole + distortion, as it arrives in a `CameraInfo`. @@ -56,21 +56,21 @@ impl Lens { /// The lens as seen in an image `scale` times the calibration's resolution; None without usable intrinsics /// or for a distortion model this cannot undistort. pub fn from_info(info: &CameraInfo, scale: f64) -> Option { - let (fx, fy) = (info.K[0], info.K[4]); + let (fx, fy) = (info.k[0], info.k[4]); if !(fx > 0.0 && fy > 0.0 && fx.is_finite() && fy.is_finite()) { return None; } let model = Model::from_name(&info.distortion_model)?; // Eight coefficients mean the rational model whatever `distortion_model` claims; some drivers mislabel it plumb_bob. let mut distortion = [0.0; 8]; - for (slot, value) in distortion.iter_mut().zip(&info.D) { + for (slot, value) in distortion.iter_mut().zip(&info.d) { *slot = *value; } Some(Self { fx: fx * scale, fy: fy * scale, - cx: (info.K[2] + 0.5) * scale - 0.5, - cy: (info.K[5] + 0.5) * scale - 0.5, + cx: (info.k[2] + 0.5) * scale - 0.5, + cy: (info.k[5] + 0.5) * scale - 0.5, model, distortion, }) @@ -185,9 +185,27 @@ mod tests { CameraInfo { width: 1920, height: 1536, - K: [1000.0, 0.0, 959.5, 0.0, 1000.0, 767.5, 0.0, 0.0, 1.0], - D: coefficients.to_vec(), - ..Default::default() + k: [1000.0, 0.0, 959.5, 0.0, 1000.0, 767.5, 0.0, 0.0, 1.0], + d: coefficients.to_vec(), + header: dimos_generated_messages::std_msgs::msg::header::Header { + stamp: dimos_generated_messages::builtin_interfaces::msg::time::Time { + sec: 0, + nanosec: 0, + }, + frame_id: String::new(), + }, + distortion_model: "plumb_bob".into(), + r: [0.; 9], + p: [0.; 12], + binning_x: 0, + binning_y: 0, + roi: dimos_generated_messages::sensor_msgs::msg::region_of_interest::RegionOfInterest { + x_offset: 0, + y_offset: 0, + width: 0, + height: 0, + do_rectify: false, + }, } } @@ -227,7 +245,7 @@ mod tests { let mut head = info(&[ -0.6792, -0.638, 0.0002, -0.0002, -0.0302, -0.2868, -0.9987, -0.1752, ]); - head.K = [1012.59, 0.0, 962.21, 0.0, 1012.13, 765.67, 0.0, 0.0, 1.0]; + head.k = [1012.59, 0.0, 962.21, 0.0, 1012.13, 765.67, 0.0, 0.0, 1.0]; let lens = Lens::from_info(&head, 1.0).unwrap(); for ((x, y), (u, v)) in [ ((-0.9467, -0.7570), (300.77, 237.5)), @@ -246,7 +264,7 @@ mod tests { // The Go2 front camera (front_camera_720.yaml); OpenCV's fisheye::distortPoints puts these rays here. let mut go2 = info(&[-0.0730943, -0.0234114, -0.0069306, 0.0092387]); go2.distortion_model = "equidistant".into(); - go2.K = [ + go2.k = [ 797.4756, 0.0, 643.5352, 0.0, 796.4872, 349.2784, 0.0, 0.0, 1.0, ]; let lens = Lens::from_info(&go2, 1.0).unwrap(); @@ -274,7 +292,7 @@ mod tests { #[test] fn zero_focal_length_is_refused() { let mut bad = info(&[0.0; 5]); - bad.K[0] = 0.0; + bad.k[0] = 0.0; assert!(Lens::from_info(&bad, 1.0).is_none()); } diff --git a/dimos/perception/localize/support_plane.py b/dimos/perception/localize/support_plane.py index bf28c061393..c751724dd62 100644 --- a/dimos/perception/localize/support_plane.py +++ b/dimos/perception/localize/support_plane.py @@ -21,6 +21,7 @@ import numpy as np +from dimos.msgs.pointcloud import pointcloud_xyz, voxel_downsample_cloud from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -59,14 +60,11 @@ def fit_support_plane( for obs in keyframes[:: max(1, len(keyframes) // 5)][:5]: cloud = rig.backdrop(obs.ts, frame) if cloud is not None: - clouds.append(cloud.voxel_downsample(0.01)) + clouds.append(pointcloud_xyz(voxel_downsample_cloud(cloud, 0.01))) if not clouds: return None - merged = clouds[0] - for cloud in clouds[1:]: - merged = merged + cloud - points = np.asarray(merged.pointcloud.points) + points = np.concatenate(clouds) if len(points) < 500: return None diff --git a/dimos/perception/localize/test_rig.py b/dimos/perception/localize/test_rig.py new file mode 100644 index 00000000000..d7524681ecd --- /dev/null +++ b/dimos/perception/localize/test_rig.py @@ -0,0 +1,63 @@ +# 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. + +"""Canonical pose frames and camera extrinsics used by localization.""" + +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest + +from dimos.memory.type.observation import Observation +from dimos.msgs.geometry import transform_matrix +from dimos.msgs.time import time_from_seconds, to_seconds +from dimos.perception.localize.rig import Rig +from dimos.robot.unitree.go2.camera_calibration import front_camera_calibration + + +@pytest.mark.parametrize("query, expected_x", [(1.0, 2.0), (1.5, 3.0)]) +def test_pose_interpolation_preserves_world_frame_for_camera_composition(mocker, query, expected_x): + poses = mocker.Mock() + poses.at.return_value = [ + Observation(ts=1.0, pose=(2.0, 0.0, 0.0)), + Observation(ts=2.0, pose=(4.0, 0.0, 0.0)), + ] + mount = TransformStamped( + header=Header(stamp=time_from_seconds(0), frame_id="base_link"), + child_frame_id="camera_optical", + transform=Transform( + translation=Vector3(x=0.5, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + rig = Rig( + cameras={"camera_optical": front_camera_calibration()}, + color=None, + world_frame="world", + poses=poses, + mounts={"camera_optical": mount}, + ) + + pose = rig.pose_at(query) + camera = rig.world_to_optical(query) + + assert pose.header.frame_id == "world" + assert pose.pose.position.x == expected_x + assert to_seconds(pose.header.stamp) == query + assert camera.header.frame_id == "camera_optical" + assert camera.child_frame_id == "world" + np.testing.assert_allclose( + transform_matrix(camera.transform)[:3, 3], [-expected_x - 0.5, 0.0, 0.0] + ) + assert to_seconds(camera.header.stamp) == query diff --git a/dimos/perception/localize/types.py b/dimos/perception/localize/types.py index 4c41c8d9650..501c25c5dc3 100644 --- a/dimos/perception/localize/types.py +++ b/dimos/perception/localize/types.py @@ -20,7 +20,7 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + from dimos_generated.sensor_msgs.msg import PointCloud2 @dataclass diff --git a/dimos/robot/unitree/go2/base.py b/dimos/robot/unitree/go2/base.py index e47188e9878..b4280e35bc4 100644 --- a/dimos/robot/unitree/go2/base.py +++ b/dimos/robot/unitree/go2/base.py @@ -25,21 +25,20 @@ import threading from typing import Any, Literal +from dimos_generated.foxglove_msgs.msg import CompressedVideo +from dimos_generated.geometry_msgs.msg import Twist +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import CameraInfo, PointCloud2 +from dimos_generated.std_msgs.msg import String from pydantic import Field, field_validator from reactivex.disposable import Disposable from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import In, Out -from dimos.msgs.foxglove_msgs.CompressedVideo import CompressedVideo -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.std_msgs.String import String from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.protocol.tf.static_tf_publisher import StaticTfPublisherConfig -from dimos.robot.unitree.go2.connection import _camera_info_static +from dimos.robot.unitree.go2.camera_calibration import front_camera_calibration from dimos.robot.unitree.go2.go2_mid360_static_transforms import ( CAMERA_XYZ, MID360_MOUNT_PRESETS, @@ -55,7 +54,7 @@ class Go2BaseConfig(StaticTfPublisherConfig): # yaw) tuple or a name from MID360_MOUNT_PRESETS. mid360_mount: tuple[float, float, float] | str = MID360_MOUNT_PRESETS["SF"] # The front camera's calibration, published at camera_info_hz. - camera_info: CameraInfo = Field(default_factory=_camera_info_static) + camera_info: CameraInfo = Field(default_factory=front_camera_calibration) camera_info_hz: float = Field(default=1.0, gt=0.0) # The frame the live odometry moves; the mount edges above it are inverted so it # never gets two parents. GO2DDS publishes its own odom edge only for base_link. @@ -107,7 +106,7 @@ def _startup_pose(self) -> None: @rpc def send_command(self, verb: str) -> None: """Fire an action verb at the robot side.""" - self.command.transport.publish(String(verb)) + self.command.transport.publish(String(data=verb)) @rpc def sport_command(self, api_id: int) -> None: diff --git a/dimos/robot/unitree/go2/dds/module.py b/dimos/robot/unitree/go2/dds/module.py index f1d704c934c..4116e846856 100644 --- a/dimos/robot/unitree/go2/dds/module.py +++ b/dimos/robot/unitree/go2/dds/module.py @@ -24,15 +24,18 @@ import time from typing import TYPE_CHECKING, Literal +from dimos_generated.sensor_msgs.msg import ( + BatteryState, + CompressedImage, + Imu, + JointState, + Joy, + PointCloud2, +) + from dimos.core.core import rpc from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.stream import Out -from dimos.msgs.sensor_msgs.BatteryState import BatteryState -from dimos.msgs.sensor_msgs.CompressedImage import CompressedImage -from dimos.msgs.sensor_msgs.Imu import Imu -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.sensor_msgs.Joy import Joy -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.robot.unitree.go2.base import Go2Base, Go2BaseConfig diff --git a/dimos/robot/unitree/go2/dds/rust/Cargo.lock b/dimos/robot/unitree/go2/dds/rust/Cargo.lock index 7a909e8a8f5..60ff855efd3 100644 --- a/dimos/robot/unitree/go2/dds/rust/Cargo.lock +++ b/dimos/robot/unitree/go2/dds/rust/Cargo.lock @@ -166,9 +166,9 @@ checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" [[package]] name = "bindgen" -version = "0.70.1" +version = "0.72.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f49d8fed880d473ea71efb9bf597651e77201bdd4893efe54c9e5d65ae04ce6f" +checksum = "993776b509cfb49c750f11b8f07a46fa23e0a1386ffc01fb1e7d343efc387895" dependencies = [ "bitflags 2.13.2", "cexpr", @@ -179,7 +179,7 @@ dependencies = [ "proc-macro2", "quote", "regex", - "rustc-hash 1.1.0", + "rustc-hash", "shlex 1.3.0", "syn 2.0.119", ] @@ -252,6 +252,20 @@ name = "bytemuck" version = "1.25.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a1f896587b6f2c069c73d2f0913e2d590c3990285cd2f0b6aa02b786b4c679c" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.6", +] [[package]] name = "byteorder" @@ -660,14 +674,25 @@ dependencies = [ "subtle", ] +[[package]] +name = "dimos-generated-messages" +version = "0.1.0" +dependencies = [ + "heapless", + "re_cdr", + "ros2msg", + "serde", + "serde-big-array", +] + [[package]] name = "dimos-go2-dds" version = "0.1.0" dependencies = [ "cdr", "cyclonedds-rs", + "dimos-generated-messages", "dimos-module", - "lcm-msgs", "serde", "serde-big-array", "serde_json", @@ -677,9 +702,8 @@ dependencies = [ ] [[package]] -name = "dimos-lcm" +name = "dimos-lcm-transport" version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" dependencies = [ "byteorder", "socket2 0.5.10", @@ -690,11 +714,12 @@ dependencies = [ name = "dimos-module" version = "0.1.0" dependencies = [ - "dimos-lcm", + "dimos-generated-messages", + "dimos-lcm-transport", "dimos-module-macros", - "lcm-msgs", "nalgebra", "rayon", + "re_cdr", "serde", "serde_json", "tokio", @@ -1062,6 +1087,15 @@ version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" +[[package]] +name = "hash32" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47d60b12902ba28e2730cd37e95b8c9223af2808df9e902d4df49588d1470606" +dependencies = [ + "byteorder", +] + [[package]] name = "hashbrown" version = "0.12.3" @@ -1104,6 +1138,23 @@ version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +[[package]] +name = "heapless" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bfb9eb618601c89945a70e254898da93b13be0388091d42117462b265bb3fad" +dependencies = [ + "hash32", + "serde", + "stable_deref_trait", +] + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + [[package]] name = "hermit-abi" version = "0.5.3" @@ -1503,14 +1554,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" -[[package]] -name = "lcm-msgs" -version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" -dependencies = [ - "byteorder", -] - [[package]] name = "libc" version = "0.2.189" @@ -2162,7 +2205,7 @@ dependencies = [ "pin-project-lite", "quinn-proto", "quinn-udp", - "rustc-hash 2.1.3", + "rustc-hash", "rustls", "socket2 0.6.5", "thiserror 2.0.20", @@ -2184,7 +2227,7 @@ dependencies = [ "rand 0.10.2", "rand_pcg", "ring", - "rustc-hash 2.1.3", + "rustc-hash", "rustls", "rustls-pki-types", "rustls-platform-verifier", @@ -2335,6 +2378,18 @@ dependencies = [ "yasna", ] +[[package]] +name = "re_cdr" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bba10c21433b597481f366b978cdf2617e56785b38cb0dcc3833c6b8ad7e061a" +dependencies = [ + "bytemuck", + "byteorder", + "serde", + "thiserror 2.0.20", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -2451,10 +2506,21 @@ dependencies = [ ] [[package]] -name = "rustc-hash" -version = "1.1.0" +name = "ros2msg" +version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" +checksum = "2fe291be5cb6d4b8e6bf9c9de80c40bd1cb09ecac50deafa2fdcae0027e04988" +dependencies = [ + "heck", + "pest", + "pest_derive", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "syn 2.0.119", + "thiserror 2.0.20", +] [[package]] name = "rustc-hash" diff --git a/dimos/robot/unitree/go2/dds/rust/Cargo.toml b/dimos/robot/unitree/go2/dds/rust/Cargo.toml index 1b153ed99e2..629a0033ea9 100644 --- a/dimos/robot/unitree/go2/dds/rust/Cargo.toml +++ b/dimos/robot/unitree/go2/dds/rust/Cargo.toml @@ -17,7 +17,7 @@ path = "src/bin/go2_dds.rs" # Our forks, vendored: the robot has no GitHub access. See vendor/*/Cargo.toml. cyclonedds-rs = { path = "vendor/cyclonedds-rs", default-features = false } dimos-module = { path = "../../../../../../native/rust/dimos-module" } -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +dimos-generated-messages = { path = "../../../../../../packages/dimos-generated/src/dimos_generated_schemas/package/rust" } serde = { version = "1", features = ["derive"] } serde-big-array = "0.5" tokio = { version = "1", features = ["rt-multi-thread", "macros"] } @@ -33,22 +33,22 @@ path = "dimos_go2_dds::module::Go2Dds" python = "dimos.robot.unitree.go2.dds.module:GO2DDS" [package.metadata.dimos.module.go2_dds.inputs] -cmd_vel = "geometry_msgs.Twist" -command = "std_msgs.String" -tf = "tf2_msgs.TFMessage" +cmd_vel = "geometry_msgs/msg/Twist" +command = "std_msgs/msg/String" +tf = "tf2_msgs/msg/TFMessage" [package.metadata.dimos.module.go2_dds.outputs] -odometry = "nav_msgs.Odometry" -lidar = "sensor_msgs.PointCloud2" -video = "foxglove_msgs.CompressedVideo" -image = "sensor_msgs.CompressedImage" -lidar_raw = "sensor_msgs.PointCloud2" -lidar_imu = "sensor_msgs.Imu" -joint_state = "sensor_msgs.JointState" -imu = "sensor_msgs.Imu" -battery = "sensor_msgs.BatteryState" -joy = "sensor_msgs.Joy" -camera_info = "sensor_msgs.CameraInfo" +odometry = "nav_msgs/msg/Odometry" +lidar = "sensor_msgs/msg/PointCloud2" +video = "foxglove_msgs/msg/CompressedVideo" +image = "sensor_msgs/msg/CompressedImage" +lidar_raw = "sensor_msgs/msg/PointCloud2" +lidar_imu = "sensor_msgs/msg/Imu" +joint_state = "sensor_msgs/msg/JointState" +imu = "sensor_msgs/msg/Imu" +battery = "sensor_msgs/msg/BatteryState" +joy = "sensor_msgs/msg/Joy" +camera_info = "sensor_msgs/msg/CameraInfo" [profile.release] lto = "thin" diff --git a/dimos/robot/unitree/go2/dds/rust/src/module.rs b/dimos/robot/unitree/go2/dds/rust/src/module.rs index 8b950be18b6..eda7db8a6c4 100644 --- a/dimos/robot/unitree/go2/dds/rust/src/module.rs +++ b/dimos/robot/unitree/go2/dds/rust/src/module.rs @@ -26,18 +26,30 @@ use std::sync::Arc; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use cyclonedds_rs::{DdsWriter, SampleBuffer}; +use dimos_generated_messages::builtin_interfaces::msg::time::Time as VideoTime; +use dimos_generated_messages::builtin_interfaces::msg::time::Time; +use dimos_generated_messages::foxglove_msgs::msg::compressed_video::CompressedVideo; +use dimos_generated_messages::geometry_msgs::msg::point::Point; +use dimos_generated_messages::geometry_msgs::msg::pose::Pose; +use dimos_generated_messages::geometry_msgs::msg::pose_with_covariance::PoseWithCovariance; +use dimos_generated_messages::geometry_msgs::msg::quaternion::Quaternion; +use dimos_generated_messages::geometry_msgs::msg::twist::Twist; +use dimos_generated_messages::geometry_msgs::msg::twist_with_covariance::TwistWithCovariance; +use dimos_generated_messages::geometry_msgs::msg::vector3::Vector3; +use dimos_generated_messages::nav_msgs::msg::odometry::Odometry; +use dimos_generated_messages::sensor_msgs::msg::battery_state::BatteryState; +use dimos_generated_messages::sensor_msgs::msg::camera_info::CameraInfo; +use dimos_generated_messages::sensor_msgs::msg::compressed_image::CompressedImage; +use dimos_generated_messages::sensor_msgs::msg::imu::Imu; +use dimos_generated_messages::sensor_msgs::msg::joint_state::JointState; +use dimos_generated_messages::sensor_msgs::msg::joy::Joy; +use dimos_generated_messages::sensor_msgs::msg::point_cloud2::PointCloud2; +use dimos_generated_messages::sensor_msgs::msg::point_field::PointField; +use dimos_generated_messages::std_msgs::msg as std_msgs; +use dimos_generated_messages::std_msgs::msg::header::Header; +use dimos_module::cdr; use dimos_module::nalgebra::{Isometry3, Quaternion as NQuaternion, Translation3, UnitQuaternion}; use dimos_module::{native_config, Input, Module, Output, Tf, Transform}; -use lcm_msgs::builtin_interfaces::Time as VideoTime; -use lcm_msgs::foxglove_msgs::CompressedVideo; -use lcm_msgs::geometry_msgs::{ - Point, Pose, PoseWithCovariance, Quaternion, Twist, TwistWithCovariance, Vector3, -}; -use lcm_msgs::nav_msgs::Odometry; -use lcm_msgs::sensor_msgs::{ - BatteryState, CameraInfo, CompressedImage, Imu, JointState, Joy, PointCloud2, PointField, -}; -use lcm_msgs::std_msgs::{self, Header, Time}; use tokio::runtime::Handle; use tracing::{info, warn}; @@ -93,7 +105,6 @@ pub struct Config { /// front_camera -> mid360_link, fixed-axis rpy in degrees. pub mid360_mount: [f64; 3], /// The front camera's calibration, published at `camera_info_hz`. - #[serde(with = "camera_info_dict")] pub camera_info: CameraInfo, /// Mount tree rate on `tf`. #[validate(range(exclusive_min = 0.0))] @@ -122,66 +133,6 @@ pub struct Config { pub deadman_ms: u64, } -/// python's `CameraInfo` config dict <-> the message; the stamp is set per publish. -mod camera_info_dict { - use lcm_msgs::sensor_msgs::CameraInfo; - use lcm_msgs::std_msgs::Header; - use serde::{Deserialize, Deserializer, Serialize, Serializer}; - - #[derive(Serialize, Deserialize)] - #[allow(non_snake_case)] - struct Dict { - frame_id: String, - height: i32, - width: i32, - distortion_model: String, - D: Vec, - K: [f64; 9], - R: [f64; 9], - P: [f64; 12], - #[serde(default)] - binning_x: i32, - #[serde(default)] - binning_y: i32, - } - - pub fn deserialize<'de, De: Deserializer<'de>>(d: De) -> Result { - let c = Dict::deserialize(d)?; - Ok(CameraInfo { - header: Header { - frame_id: c.frame_id, - ..Default::default() - }, - height: c.height, - width: c.width, - distortion_model: c.distortion_model, - D: c.D, - K: c.K, - R: c.R, - P: c.P, - binning_x: c.binning_x, - binning_y: c.binning_y, - ..Default::default() - }) - } - - pub fn serialize(c: &CameraInfo, s: S) -> Result { - Dict { - frame_id: c.header.frame_id.clone(), - height: c.height, - width: c.width, - distortion_model: c.distortion_model.clone(), - D: c.D.clone(), - K: c.K, - R: c.R, - P: c.P, - binning_x: c.binning_x, - binning_y: c.binning_y, - } - .serialize(s) - } -} - fn known_encoding(v: &str) -> Result<(), validator::ValidationError> { match v { "h264" | "jpeg" => Ok(()), @@ -298,43 +249,43 @@ pub fn parse_verb(verb: &str) -> Option> { #[derive(Module)] #[module(name = "go2_dds", setup = start, teardown = stop)] pub struct Go2Dds { - #[input(decode = Twist::decode, handler = on_cmd_vel)] + #[input(decode = cdr::decode, handler = on_cmd_vel)] cmd_vel: Input, - #[input(decode = std_msgs::String::decode, handler = on_command)] - command: Input, + #[input(decode = cdr::decode, handler = on_command)] + command: Input, - #[output(encode = Odometry::encode)] + #[output(encode = cdr::encode)] odometry: Output, - #[output(encode = PointCloud2::encode)] + #[output(encode = cdr::encode)] lidar: Output, - #[output(encode = CompressedVideo::encode)] + #[output(encode = cdr::encode)] video: Output, - #[output(encode = CompressedImage::encode)] + #[output(encode = cdr::encode)] image: Output, - #[output(encode = PointCloud2::encode)] + #[output(encode = cdr::encode)] lidar_raw: Output, - #[output(encode = Imu::encode)] + #[output(encode = cdr::encode)] lidar_imu: Output, - #[output(encode = JointState::encode)] + #[output(encode = cdr::encode)] joint_state: Output, - #[output(encode = Imu::encode)] + #[output(encode = cdr::encode)] imu: Output, - #[output(encode = BatteryState::encode)] + #[output(encode = cdr::encode)] battery: Output, - #[output(encode = Joy::encode)] + #[output(encode = cdr::encode)] joy: Output, - #[output(encode = CameraInfo::encode)] + #[output(encode = cdr::encode)] camera_info: Output, #[config] @@ -414,7 +365,7 @@ impl Go2Dds { }); } - async fn on_command(&mut self, msg: std_msgs::String) { + async fn on_command(&mut self, msg: std_msgs::string::String) { match parse_verb(&msg.data) { Some(cmds) => { info!(verb = %msg.data, ?cmds, "command"); @@ -446,7 +397,7 @@ fn time_of_secs(ts: f64) -> Time { let sec = ts.trunc(); Time { sec: sec as i32, - nsec: ((ts - sec) * 1e9) as i32, + nanosec: ((ts - sec) * 1e9) as u32, } } @@ -497,7 +448,6 @@ impl Cadence { fn header(frame_id: &str, ts: f64) -> Header { Header { - seq: 0, stamp: time_of_secs(ts), frame_id: frame_id.into(), } @@ -580,9 +530,9 @@ pub fn battery(s: &types::LowState, ts: f64) -> BatteryState { capacity: f32::NAN, design_capacity: f32::NAN, percentage: b.soc as f32 / 100.0, - power_supply_status: status as u8, - power_supply_health: BatteryState::POWER_SUPPLY_HEALTH_UNKNOWN as u8, - power_supply_technology: BatteryState::POWER_SUPPLY_TECHNOLOGY_LION as u8, + power_supply_status: status, + power_supply_health: BatteryState::POWER_SUPPLY_HEALTH_UNKNOWN, + power_supply_technology: BatteryState::POWER_SUPPLY_TECHNOLOGY_LION, present: true, cell_voltage: b .cell_vol @@ -612,7 +562,6 @@ pub fn odometry(s: &types::Odometry, ts: f64) -> (Odometry, Transform) { let (l, a) = (&s.twist.twist.linear, &s.twist.twist.angular); let odom = Odometry { header: Header { - seq: 0, stamp: time_of_secs(ts), frame_id: s.header.frame_id.clone(), }, @@ -660,25 +609,24 @@ pub fn odometry(s: &types::Odometry, ts: f64) -> (Odometry, Transform) { pub fn pointcloud(s: &types::PointCloud2, ts: f64) -> PointCloud2 { PointCloud2 { header: Header { - seq: 0, stamp: time_of_secs(ts), frame_id: s.header.frame_id.clone(), }, - height: s.height as i32, - width: s.width as i32, + height: s.height, + width: s.width, fields: s .fields .iter() .map(|f| PointField { name: f.name.clone(), - offset: f.offset as i32, + offset: f.offset, datatype: f.datatype, - count: f.count as i32, + count: f.count, }) .collect(), is_bigendian: s.is_bigendian, - point_step: s.point_step as i32, - row_step: s.row_step as i32, + point_step: s.point_step, + row_step: s.row_step, data: s.data.clone(), is_dense: s.is_dense, } @@ -918,7 +866,6 @@ impl DdsLoop { if let Some(data) = newest_jpeg(replies) { let msg = CompressedImage { header: Header { - seq: 0, stamp: time_of_secs(now_secs()), frame_id: CAMERA_FRAME.to_string(), }, @@ -999,7 +946,7 @@ impl VideoLoop { let msg = CompressedVideo { timestamp: VideoTime { sec: (ts_ns / 1_000_000_000) as i32, - nanosec: (ts_ns % 1_000_000_000) as i32, + nanosec: (ts_ns % 1_000_000_000) as u32, }, frame_id: CAMERA_FRAME.to_string(), data: au.data, @@ -1115,7 +1062,7 @@ mod tests { assert_eq!(b.cell_voltage, vec![4.1]); assert_eq!( b.power_supply_status, - BatteryState::POWER_SUPPLY_STATUS_DISCHARGING as u8 + BatteryState::POWER_SUPPLY_STATUS_DISCHARGING ); let remote = types::WirelessController { lx: 0.5, @@ -1152,8 +1099,8 @@ mod tests { let info = config().camera_info; assert_eq!((info.width, info.height), (1280, 720)); assert_eq!(info.header.frame_id, CAMERA_FRAME); - assert_eq!(info.K[0], 797.4756164864929); - assert_eq!(info.D.len(), 4); + assert_eq!(info.k[0], 797.4756164864929); + assert_eq!(info.d.len(), 4); } #[test] @@ -1207,7 +1154,7 @@ mod tests { odom.header.stamp, Time { sec: 1234, - nsec: 500_000_000 + nanosec: 500_000_000 } ); assert_eq!( diff --git a/dimos/robot/unitree/go2/dds/rust/tests/config.json b/dimos/robot/unitree/go2/dds/rust/tests/config.json index 7452be70e3e..4635f979c12 100644 --- a/dimos/robot/unitree/go2/dds/rust/tests/config.json +++ b/dimos/robot/unitree/go2/dds/rust/tests/config.json @@ -16,17 +16,23 @@ 0.0 ], "camera_info": { - "frame_id": "camera_optical", + "header": { + "stamp": { + "sec": 0, + "nanosec": 0 + }, + "frame_id": "camera_optical" + }, "height": 720, "width": 1280, "distortion_model": "equidistant", - "D": [ + "d": [ -0.07309428880537933, -0.02341140740909078, -0.0069305931780026956, 0.009238684474464793 ], - "K": [ + "k": [ 797.4756164864929, 0.0, 643.5352167821186, @@ -37,7 +43,7 @@ 0.0, 1.0 ], - "R": [ + "r": [ 1.0, 0.0, 0.0, @@ -48,7 +54,7 @@ 0.0, 1.0 ], - "P": [ + "p": [ 797.4756164864929, 0.0, 643.5352167821186, @@ -63,7 +69,14 @@ 0.0 ], "binning_x": 0, - "binning_y": 0 + "binning_y": 0, + "roi": { + "x_offset": 0, + "y_offset": 0, + "height": 0, + "width": 0, + "do_rectify": false + } }, "camera_info_hz": 1.0, "tf_root": "base_link", diff --git a/dimos/robot/unitree/go2/dds/rust/vendor/cyclonedds-sys/Cargo.toml b/dimos/robot/unitree/go2/dds/rust/vendor/cyclonedds-sys/Cargo.toml index 2f766bdd822..5699315789e 100644 --- a/dimos/robot/unitree/go2/dds/rust/vendor/cyclonedds-sys/Cargo.toml +++ b/dimos/robot/unitree/go2/dds/rust/vendor/cyclonedds-sys/Cargo.toml @@ -18,8 +18,8 @@ shm = [] default = [] [build-dependencies] -# 0.55 chokes on anonymous unions with clang >= 16; 0.70 names them. -bindgen = "0.70" +# 0.72 retains complete struct definitions with current Clang, including Clang 22. +bindgen = "0.72" glob = "0.3" cc = "1.0" diff --git a/dimos/robot/unitree/go2/dds/rust/vendor/cyclonedds-sys/build.rs b/dimos/robot/unitree/go2/dds/rust/vendor/cyclonedds-sys/build.rs index 7bee2362b1c..04013ed46f9 100644 --- a/dimos/robot/unitree/go2/dds/rust/vendor/cyclonedds-sys/build.rs +++ b/dimos/robot/unitree/go2/dds/rust/vendor/cyclonedds-sys/build.rs @@ -648,9 +648,6 @@ mod build { } pub fn main() { - for (key, value) in env::vars() { - println!("{}: {}", key, value); - } let mut headerloc = find_cyclonedds().unwrap(); if let Some(iceoryx_headers) = find_iceoryx("v2.0.2") { diff --git a/dimos/robot/unitree/mujoco_connection.py b/dimos/robot/unitree/mujoco_connection.py index 83e4c83c81d..0478cb14e0e 100644 --- a/dimos/robot/unitree/mujoco_connection.py +++ b/dimos/robot/unitree/mujoco_connection.py @@ -44,6 +44,7 @@ from dimos.core.global_config import GlobalConfig from dimos.msgs.camera_info import camera_info_from_fov from dimos.msgs.image import image_from_array +from dimos.msgs.pointcloud import pointcloud_from_xyz from dimos.msgs.time import header_now, time_from_seconds from dimos.simulation.mujoco.constants import ( LAUNCHER_PATH, @@ -352,7 +353,9 @@ def get_lidar_message(self) -> PointCloud2 | None: if seq > self._last_lidar_seq and lidar is not None: self._last_lidar_seq = seq points, ts = lidar - return PointCloud2.from_numpy(points, frame_id="world", timestamp=ts) + return pointcloud_from_xyz( + points, header=Header(stamp=time_from_seconds(ts), frame_id="world") + ) return None diff --git a/dimos/simulation/engines/test_mujoco_sim_module.py b/dimos/simulation/engines/test_mujoco_sim_module.py index 03998366f78..27bf4598fb5 100644 --- a/dimos/simulation/engines/test_mujoco_sim_module.py +++ b/dimos/simulation/engines/test_mujoco_sim_module.py @@ -244,9 +244,16 @@ def disconnect(self) -> None: ] assert children[3:] == ["apple"] apple = messages[-1].transforms[3] - assert apple.frame_id == "world" - assert apple.ts == pytest.approx(time.time(), abs=5.0) - assert np.allclose(apple.translation.to_numpy(), [0.4, 0.08, 0.17]) + assert apple.header.frame_id == "world" + assert to_seconds(apple.header.stamp) == pytest.approx(time.time(), abs=5.0) + assert np.allclose( + [ + apple.transform.translation.x, + apple.transform.translation.y, + apple.transform.translation.z, + ], + [0.4, 0.08, 0.17], + ) finally: module.stop() From 4e6383d77a0f1ba819f3290451a0251133aa45f3 Mon Sep 17 00:00:00 2001 From: cc Date: Fri, 9 Oct 2026 20:01:56 -0700 Subject: [PATCH 07/11] fix(cdr): repair native and executable example regressions --- dimos/cli/bake/test_bake_e2e.py | 3 +- dimos/e2e_tests/test_map3d_browser.py | 9 +++--- docs/usage/evals.md | 3 +- docs/usage/lcm.md | 20 ++++++------ docs/usage/transforms.md | 16 +++++++--- native/cpp/tests/lidar_demo.cpp | 44 +++++++++++++++++++++++++++ 6 files changed, 75 insertions(+), 20 deletions(-) create mode 100644 native/cpp/tests/lidar_demo.cpp diff --git a/dimos/cli/bake/test_bake_e2e.py b/dimos/cli/bake/test_bake_e2e.py index 8c7c5ace192..054402e3feb 100644 --- a/dimos/cli/bake/test_bake_e2e.py +++ b/dimos/cli/bake/test_bake_e2e.py @@ -27,6 +27,7 @@ import subprocess import time +from dimos_generated.dimos_msgs.msg import RegionPointCloud2 from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 from dimos_generated.sensor_msgs.msg import PointCloud2 from dimos_generated.std_msgs.msg import Header @@ -148,7 +149,7 @@ def test_the_host_publishes_its_outputs_and_hides_the_suppressed_hop(baked, zeno transports = [] try: for name, msg_type in ( - ("surface_map", PointCloud2), + ("surface_map", RegionPointCloud2), ("global_map", PointCloud2), ("local_map", PointCloud2), ): diff --git a/dimos/e2e_tests/test_map3d_browser.py b/dimos/e2e_tests/test_map3d_browser.py index 74c92329b4e..4c8ba598997 100644 --- a/dimos/e2e_tests/test_map3d_browser.py +++ b/dimos/e2e_tests/test_map3d_browser.py @@ -25,10 +25,12 @@ from collections.abc import Callable import cv2 +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import time_from_seconds from dimos.web.cockpit import Map3D, cockpit pytest.importorskip("playwright") @@ -42,10 +44,9 @@ # ~2.6 km, so the frame only shows if the clip planes follow the fit, and the # voxels are still ~4 px wide from there. Their heights put one voxel at each # end of the height ramp and one in the middle. -CLOUD = PointCloud2.from_numpy( +CLOUD = pointcloud_from_xyz( np.array([[-1000.0, 0.0, 0.0], [0.0, 0.0, 16.0], [1000.0, 0.0, 32.0]], np.float32), - frame_id="world", - timestamp=1.0, + header=Header(stamp=time_from_seconds(1.0), frame_id="world"), ) diff --git a/docs/usage/evals.md b/docs/usage/evals.md index f1a26f1b256..7772866b0d8 100644 --- a/docs/usage/evals.md +++ b/docs/usage/evals.md @@ -246,7 +246,8 @@ odom = store.stream("odom", PoseStamped) for i in range(20): odom.append( PoseStamped(header=Header(stamp=time_from_seconds(1000.0+i),frame_id="world"), - pose=Pose(position=Point(x=float(i),y=2.5))), + pose=Pose(position=Point(x=float(i), y=2.5, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0))), ts=1000.0 + i, ) store.stop() diff --git a/docs/usage/lcm.md b/docs/usage/lcm.md index bc8322e9740..3739723f46e 100644 --- a/docs/usage/lcm.md +++ b/docs/usage/lcm.md @@ -6,8 +6,8 @@ CDR messages. LCM and Zenoh use the same message bytes; choosing a transport does not change a message's schema or codec. Message definitions are ROS2 `.msg` files. The standalone generator produces -Python, C++ and Rust value types without requiring ROS. Python bindings share -the C++ Fast CDR implementation; native Rust types use `re_cdr`. +Python, C++ and Rust value types without requiring ROS. Python classes/codecs use rosbags; C++ uses upstream ROSIDL/Fast CDR and +native Rust types use `re_cdr`. See [add and use a message](/docs/development/messages.md) for the complete local-message user story and the three language examples. @@ -15,12 +15,13 @@ local-message user story and the three language examples. ```python session=cdr_transport_demo ansi=false from dimos_generated.geometry_msgs.msg import Vector3 +from dimos_message_build.registry import encode, decode message = Vector3(x=1.0, y=2.0, z=3.0) -payload = message.encode() -decoded = Vector3.decode(payload) +payload = encode(message) +decoded = decode(payload, Vector3) assert (decoded.x, decoded.y, decoded.z) == (1.0, 2.0, 3.0) -print(message.msg_name) +print(message.__msgtype__) print(f"Decoded: x={decoded.x}, y={decoded.y}, z={decoded.z}") ``` @@ -54,12 +55,13 @@ dot=32.0 ```python session=cdr_transport_demo ansi=false from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.std_msgs.msg import Header from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_xyz points = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32) -cloud = pointcloud_from_xyz(points, header=Header(frame_id="camera")) -roundtrip = PointCloud2.decode(cloud.encode()) +cloud = pointcloud_from_xyz(points, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="camera")) +roundtrip = decode(encode(cloud), PointCloud2) np.testing.assert_array_equal(pointcloud_xyz(roundtrip), points) print(f"PointCloud: {roundtrip.width * roundtrip.height} points") print(f"Frame: {roundtrip.header.frame_id}") @@ -104,8 +106,8 @@ finally: In-process value: 1.0, 2.0, 3.0 ``` -For inter-process exchange, encode with `message.encode()` and decode using the -known generated type. Separate Python-object serialization paths retain their +For inter-process exchange, use the registry’s `encode(message)` and +`decode(payload, MessageType)` functions. Separate Python-object serialization paths retain their own contracts; CDR is the typed-message representation. ## Available types and custom packages diff --git a/docs/usage/transforms.md b/docs/usage/transforms.md index 8ed211c2145..2ef58541608 100644 --- a/docs/usage/transforms.md +++ b/docs/usage/transforms.md @@ -77,14 +77,15 @@ A generated `Transform` contains translation and rotation. `TransformStamped` ad ```python from dimos_generated.geometry_msgs.msg import TransformStamped +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.std_msgs.msg import Header from dimos.msgs.geometry import compose_transforms, inverse_transform, transform_matrix from dimos.msgs.time import time_from_seconds from dimos_generated.geometry_msgs.msg import Quaternion from dimos_generated.geometry_msgs.msg import Transform from dimos_generated.geometry_msgs.msg import Vector3 -camera_transform = TransformStamped(header=Header(frame_id='base_link'), child_frame_id='camera_link', transform=Transform(translation=Vector3(x=0.5, y=0.0, z=0.3), rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0))) -print(camera_transform) +camera_transform = TransformStamped(header=Header(stamp=time_from_seconds(0.0), frame_id='base_link'), child_frame_id='camera_link', transform=Transform(translation=Vector3(x=0.5, y=0.0, z=0.3), rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0))) +print(camera_transform.header.frame_id, camera_transform.child_frame_id) ``` ```results @@ -97,14 +98,15 @@ Transforms can be composed and inverted: ```python from dimos_generated.geometry_msgs.msg import TransformStamped +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.std_msgs.msg import Header from dimos.msgs.geometry import compose_transforms, inverse_transform, transform_matrix from dimos.msgs.time import time_from_seconds from dimos_generated.geometry_msgs.msg import Quaternion from dimos_generated.geometry_msgs.msg import Transform from dimos_generated.geometry_msgs.msg import Vector3 -t1 = TransformStamped(header=Header(frame_id='base_link'), child_frame_id='camera_link', transform=Transform(translation=Vector3(x=1.0, y=0.0, z=0.0), rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0))) -t2 = TransformStamped(header=Header(frame_id='camera_link'), child_frame_id='end_effector', transform=Transform(translation=Vector3(x=0.0, y=0.5, z=0.0), rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0))) +t1 = TransformStamped(header=Header(stamp=time_from_seconds(0.0), frame_id='base_link'), child_frame_id='camera_link', transform=Transform(translation=Vector3(x=1.0, y=0.0, z=0.0), rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0))) +t2 = TransformStamped(header=Header(stamp=Time(sec=0, nanosec=0), frame_id='camera_link'), child_frame_id='end_effector', transform=Transform(translation=Vector3(x=0.0, y=0.5, z=0.0), rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0))) t3 = compose_transforms(t1, t2) print(f'Composed: {t3.header.frame_id} -> {t3.child_frame_id}') print(f'Translation: ({t3.transform.translation.x}, {t3.transform.translation.y}, {t3.transform.translation.z})') @@ -124,6 +126,7 @@ For integration with libraries like NumPy or OpenCV: ```python from dimos_generated.geometry_msgs.msg import TransformStamped +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.std_msgs.msg import Header from dimos.msgs.geometry import compose_transforms, inverse_transform, transform_matrix from dimos.msgs.time import time_from_seconds @@ -202,6 +205,7 @@ This example demonstrates how multiple modules publish and receive transforms. T ```python skip ansi=false from dimos_generated.geometry_msgs.msg import TransformStamped +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.std_msgs.msg import Header from dimos.msgs.geometry import compose_transforms, inverse_transform, transform_matrix from dimos.msgs.time import time_from_seconds @@ -366,6 +370,7 @@ text "CameraModule" italic at ((CL.x + CO.x)/2, CL.s.y - 0.25in) import time from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.std_msgs.msg import Header from dimos.msgs.time import time_from_nanoseconds from dimos.protocol.tf.tf import MultiTBuffer @@ -381,7 +386,8 @@ for i in range(5): ), child_frame_id="camera_link", transform=Transform( - translation=Vector3(x=float(i)), rotation=Quaternion(w=1.0) + translation=Vector3(x=float(i), y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) ), ) tf.receive_transform(t) diff --git a/native/cpp/tests/lidar_demo.cpp b/native/cpp/tests/lidar_demo.cpp new file mode 100644 index 00000000000..346eb8ed23a --- /dev/null +++ b/native/cpp/tests/lidar_demo.cpp @@ -0,0 +1,44 @@ +// Copyright 2026 Dimensional Inc. +// SPDX-License-Identifier: Apache-2.0 + +#include +#include +#include +#include +#include + +#include "../../../dimos/hardware/sensors/lidar/common/point_cloud_utils.hpp" + +int main() { + auto cloud = dimos::make_xyzi_cloud("lidar", 1.25, 2); + for (int i = 0; i < 2; ++i) { + float* point = dimos::xyzi_point(cloud, i); + point[0] = static_cast(i + 1); + point[1] = 2; + point[2] = 3; + point[3] = 0.5; + } + for (bool little : {true, false}) { + auto bytes = dimos::native::cdr_encode(cloud, little); + auto decoded = dimos::native::cdr_decode(bytes.data(), bytes.size()); + assert(decoded == cloud); + assert(decoded.header.stamp.sec == 1 && decoded.header.stamp.nanosec == 250000000); + assert(decoded.width == 2 && decoded.fields.size() == 4 && decoded.data.size() == 32); + std::cout << (little ? "little" : "big") << "-endian CDR: " + << decoded.width << " XYZI points, frame=" << decoded.header.frame_id + << ", stamp=" << decoded.header.stamp.sec << "s + " + << decoded.header.stamp.nanosec << "ns, intensity=" + << dimos::xyzi_point(decoded, 0)[3] << '\n'; + } + auto negative = dimos::time_from_seconds(-0.25); + assert(negative.sec == -1 && negative.nanosec == 750000000); + bool rejected = false; + try { dimos::time_from_seconds(std::numeric_limits::infinity()); } + catch (const std::invalid_argument&) { rejected = true; } + assert(rejected); + rejected = false; + try { dimos::make_xyzi_cloud("lidar", 0, -1); } + catch (const std::invalid_argument&) { rejected = true; } + assert(rejected); + std::cout << "Negative timestamps normalize; invalid timestamps and negative point counts fail.\n"; +} From 14aa539a1306b9d49d95d5bcf7237c515a37808c Mon Sep 17 00:00:00 2001 From: cc Date: Fri, 9 Oct 2026 20:09:10 -0700 Subject: [PATCH 08/11] fix(cdr): adapt current-main dataset and SDK regressions --- dimos/agents/skills/point_nav.py | 2 +- dimos/control/tasks/test_pose_target_ik.py | 20 ++- dimos/imitation/dataprep/build.py | 9 +- dimos/imitation/dataprep/test_mcap_source.py | 82 +++++----- .../mls_planner/utils/test_plan_rrd.py | 20 ++- dimos/web/test_lcm_codec.py | 150 ------------------ examples/message-project/rust/Cargo.lock | 96 +++++------ 7 files changed, 136 insertions(+), 243 deletions(-) delete mode 100644 dimos/web/test_lcm_codec.py diff --git a/dimos/agents/skills/point_nav.py b/dimos/agents/skills/point_nav.py index 280b1e223a9..6f261988418 100644 --- a/dimos/agents/skills/point_nav.py +++ b/dimos/agents/skills/point_nav.py @@ -17,8 +17,8 @@ import math from dimos_generated.geometry_msgs.msg import Point, PointStamped, PoseStamped +from dimos_generated.std_msgs.msg import Bool from dimos_generated.tf2_msgs.msg import TFMessage -from dimos_lcm.std_msgs import Bool from dimos.agents.annotation import skill from dimos.agents.capabilities import CAP_MOVEMENT diff --git a/dimos/control/tasks/test_pose_target_ik.py b/dimos/control/tasks/test_pose_target_ik.py index fac8b268b6c..3ea70f586ce 100644 --- a/dimos/control/tasks/test_pose_target_ik.py +++ b/dimos/control/tasks/test_pose_target_ik.py @@ -373,14 +373,28 @@ def test_current_frame_poses_uses_live_coordinator_seed(mocker: MockerFixture) - def test_measured_frame_poses_follow_measured_joints_not_the_target(mocker: MockerFixture) -> None: solver = _solver(mocker) - target = PoseStamped(frame_id="world", position=(9.0, 9.0, 9.0)) # commanded, never reached + target = PoseStamped( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id="world"), + pose=Pose( + position=Point(x=9.0, y=9.0, z=9.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + ) # commanded, never reached task = _Task(_config(), solver, _snapshot(targets={"tool": target})) poses = task.measured_frame_poses(_state(positions={"arm/a": 0.2, "arm/b": 0.3})) - assert poses == {"tool": PoseStamped(frame_id="base")} # FK result, not the target + assert poses == { + "tool": PoseStamped( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id="base"), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + } # FK result, not the target seed, frames = solver.frame_poses.call_args.args - assert seed.position == [0.2, 0.3] and list(frames) == ["tool"] + np.testing.assert_array_equal(seed.position, [0.2, 0.3]) + assert list(frames) == ["tool"] def test_measured_frame_poses_are_empty_without_full_feedback(mocker: MockerFixture) -> None: diff --git a/dimos/imitation/dataprep/build.py b/dimos/imitation/dataprep/build.py index 55809b134d3..48ebe634008 100644 --- a/dimos/imitation/dataprep/build.py +++ b/dimos/imitation/dataprep/build.py @@ -81,7 +81,14 @@ def _recording_codecs(path: Path) -> dict[str, StreamCodec]: return codecs for channel in summary.channels.values(): payload_type = channel.metadata.get("dimos.payload_type") - if payload_type and channel.message_encoding in {"jpeg", "lcm", "lz4+lcm"}: + if payload_type and channel.message_encoding in { + "cdr", + "lz4+cdr", + "json", + "jpeg", + "lcm", + "lz4+lcm", + }: try: codecs[channel.topic] = cast( "StreamCodec", codec_from_id(channel.message_encoding, payload_type) diff --git a/dimos/imitation/dataprep/test_mcap_source.py b/dimos/imitation/dataprep/test_mcap_source.py index 21ac99e07ac..fe2b21031b0 100644 --- a/dimos/imitation/dataprep/test_mcap_source.py +++ b/dimos/imitation/dataprep/test_mcap_source.py @@ -21,6 +21,9 @@ from typing import Any import weakref +from dimos_generated.sensor_msgs.msg import Image, JointState +from dimos_generated.std_msgs.msg import Header, String +from dimos_message_build.registry import encode as cdr_encode, schema as message_schema from mcap.writer import Writer as McapWriter import numpy as np import pytest @@ -35,24 +38,29 @@ Sample, SyncConfig, ) -from dimos.memory.codecs.jpeg import JpegCodec -from dimos.memory.codecs.lcm import LcmCodec +from dimos.memory.codecs.cdr import CdrCodec from dimos.memory.codecs.lz4 import Lz4Codec -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.std_msgs.String import String +from dimos.msgs.image import image_from_array +from dimos.msgs.time import time_from_seconds def _register_channel( writer: McapWriter, name: str, payload_type: type[Any], - message_encoding: str = "lcm", + message_encoding: str = "cdr", ) -> int: + schema_id = 0 + if message_encoding in {"cdr", "lz4+cdr"}: + schema_id = writer.register_schema( + name=payload_type.__msgtype__, + encoding="ros2msg", + data=message_schema(payload_type.__msgtype__).encode(), + ) return writer.register_channel( topic=name, message_encoding=message_encoding, - schema_id=0, + schema_id=schema_id, metadata={ "dimos.payload_type": f"{payload_type.__module__}.{payload_type.__qualname__}", "dimos.observation_time": "publish_time", @@ -68,7 +76,7 @@ def _write_message(writer: McapWriter, channel_id: int, ts: float, message: Any) publish_time=timestamp_ns, data=message.to_json().encode("utf-8") if isinstance(message, EpisodeStatus) - else (message if isinstance(message, bytes) else message.lcm_encode()), + else (message if isinstance(message, bytes) else cdr_encode(message)), ) @@ -77,7 +85,7 @@ def _write_collection(path: Path) -> None: writer = McapWriter(output) writer.start(profile="dimos", library="test") channels = { - "color_image": _register_channel(writer, "color_image", Image, "jpeg"), + "color_image": _register_channel(writer, "color_image", Image), "coordinator_joint_state": _register_channel( writer, "coordinator_joint_state", JointState ), @@ -105,13 +113,10 @@ def _write_collection(path: Path) -> None: writer, channels["color_image"], ts, - JpegCodec().encode( - Image( - ts=ts, - frame_id="wrist_camera_link", - format=ImageFormat.RGB, - data=np.full((8, 8, 3), index, dtype=np.uint8), - ) + image_from_array( + np.full((8, 8, 3), index, dtype=np.uint8), + encoding="rgb8", + header=Header(stamp=time_from_seconds(ts), frame_id="wrist_camera_link"), ), ) _write_message( @@ -119,12 +124,11 @@ def _write_collection(path: Path) -> None: channels["coordinator_joint_state"], ts, JointState( - ts=ts, - frame_id="coordinator", + header=Header(stamp=time_from_seconds(ts), frame_id="coordinator"), name=["shoulder", "wrist"], - position=[float(index), float(index + 1)], - velocity=[0.0, 0.0], - effort=[0.0, 0.0], + position=np.array([float(index), float(index + 1)]), + velocity=np.zeros(2), + effort=np.zeros(2), ), ) _write_message( @@ -132,12 +136,11 @@ def _write_collection(path: Path) -> None: channels["applied_joint_position_command"], ts, JointState( - ts=ts, - frame_id="coordinator", + header=Header(stamp=time_from_seconds(ts), frame_id="coordinator"), name=["shoulder", "wrist"], - position=[float(index) + 0.25, float(index) + 1.25], - velocity=[], - effort=[], + position=np.array([float(index) + 0.25, float(index) + 1.25]), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ) _write_message( @@ -163,7 +166,7 @@ def _config(source: Path, output: Path) -> DataPrepConfig: observation={ "observation.images.wrist": FeatureSpec( stream="color_image", - field="data", + field=None, dtype="video", shape=(8, 8, 3), names=["height", "width", "channels"], @@ -223,12 +226,18 @@ def writer(samples: Iterator[Sample], selected_output: OutputConfig) -> Path: np.testing.assert_array_equal(received[0].action["action"], [0.25, 1.25]) -@pytest.mark.parametrize("encoding", ["lcm", "lz4+lcm"]) +@pytest.mark.parametrize("encoding", ["cdr", "lz4+cdr"]) def test_recording_metadata_decodes_typed_messages(tmp_path: Path, encoding: str) -> None: path = tmp_path / "recording.mcap" - message = JointState(ts=12.5, name=["arm/joint1"], position=[0.25]) - codec = LcmCodec(JointState) - payload = Lz4Codec(codec).encode(message) if encoding == "lz4+lcm" else codec.encode(message) + message = JointState( + header=Header(stamp=time_from_seconds(12.5), frame_id="arm"), + name=["arm/joint1"], + position=np.array([0.25]), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + codec = CdrCodec(JointState) + payload = Lz4Codec(codec).encode(message) if encoding == "lz4+cdr" else codec.encode(message) with path.open("wb") as file: writer = McapWriter(file) writer.start() @@ -239,7 +248,7 @@ def test_recording_metadata_decodes_typed_messages(tmp_path: Path, encoding: str with _open_recording(path) as store: observation = store.stream("measured").first() assert observation.ts == 11.5 - assert observation.data.lcm_encode() == message.lcm_encode() + assert cdr_encode(observation.data) == cdr_encode(message) def test_missing_message_package_reports_the_recorded_type(tmp_path: Path) -> None: @@ -249,7 +258,7 @@ def test_missing_message_package_reports_the_recorded_type(tmp_path: Path) -> No writer.start() channel = writer.register_channel( topic="custom_state", - message_encoding="lcm", + message_encoding="cdr", schema_id=0, metadata={"dimos.payload_type": "missing_recording_package.State"}, ) @@ -266,14 +275,15 @@ def test_sample_emission_does_not_keep_every_decoded_camera_frame(tmp_path, mock config = _config(recording, tmp_path / "dataset") config.observation["camera_alias"] = config.observation["observation.images.wrist"] decoded = [] - decode = JpegCodec.decode + decode = CdrCodec.decode def track_decode(codec, payload): image = decode(codec, payload) - decoded.append(weakref.ref(image.data)) + if isinstance(image, Image): + decoded.append(weakref.ref(image.data)) return image - mocker.patch.object(JpegCodec, "decode", autospec=True, side_effect=track_decode) + mocker.patch.object(CdrCodec, "decode", autospec=True, side_effect=track_decode) with _open_recording(recording) as store: episode = extract_episodes(store, config.episodes)[0] samples = iter_episode_samples( diff --git a/dimos/navigation/global_planner/mls_planner/utils/test_plan_rrd.py b/dimos/navigation/global_planner/mls_planner/utils/test_plan_rrd.py index a93e026a10a..de962eb5af0 100644 --- a/dimos/navigation/global_planner/mls_planner/utils/test_plan_rrd.py +++ b/dimos/navigation/global_planner/mls_planner/utils/test_plan_rrd.py @@ -22,8 +22,7 @@ from dimos.mapping.ray_tracing.transformer import RayTraceMap from dimos.memory.tf import StreamTF from dimos.memory.type.observation import Observation -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_from_xyz from dimos.navigation.global_planner.mls_planner.mls_planner import MLSPlanner from dimos.navigation.global_planner.mls_planner.utils.plan_rrd import Seeding, SeedStage @@ -40,9 +39,22 @@ def test_seeding_lands_one_region_per_step_and_matches_a_full_rebuild( ) -> None: # Three floor patches, each in its own region cell. cloud = np.vstack([_patch(0.0), _patch(8.0), _patch(16.0)]) - loaded_map = Observation(id=0, ts=5.0, _data=PointCloud2.from_numpy(cloud, frame_id="map")) + loaded_map = Observation( + id=0, + ts=5.0, + _data=pointcloud_from_xyz( + cloud, header=Header(frame_id="map", stamp=Time(sec=5, nanosec=0)) + ), + ) tf_lookup = mocker.create_autospec(StreamTF, instance=True) - tf_lookup.get.return_value = Transform(frame_id="odom", child_frame_id="map") + tf_lookup.get.return_value = TransformStamped( + header=Header(frame_id="odom", stamp=Time(sec=5, nanosec=0)), + child_frame_id="map", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) # No support gate, so every seeded voxel reaches the planner. ray = RayTraceMap(voxel_size=VOXEL_SIZE, support_min=0) planner = MLSPlanner(voxel_size=VOXEL_SIZE, robot_height=0.5) diff --git a/dimos/web/test_lcm_codec.py b/dimos/web/test_lcm_codec.py deleted file mode 100644 index d9895063005..00000000000 --- a/dimos/web/test_lcm_codec.py +++ /dev/null @@ -1,150 +0,0 @@ -# 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. - -"""Generic .lcm.v1 codec: the default rule, schema export, the -encoder's fingerprint guard, and the golden vectors shared with -web/sdk/src/decoders/lcm.test.ts.""" - -import base64 -import json -import subprocess -import sys -from typing import Any - -from dimos_lcm import ( - geometry_msgs as lcm_geometry_msgs, - nav_msgs as lcm_nav_msgs, - sensor_msgs as lcm_sensor_msgs, - std_msgs as lcm_std_msgs, - tf2_msgs as lcm_tf2_msgs, - visualization_msgs as lcm_visualization_msgs, -) -import pytest - -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.MotorCommandArray import MotorCommandArray -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.web.lcm_codec import ( - check_lcm_params, - default_encoding, - encode_lcm_v1, - export_schema, - schema_class_for, -) -from dimos.web.relay_bridge.gen_lcm_fixtures import build_messages, walk_value -from dimos.web.relay_bridge.locate import find_web_dir - -with open(find_web_dir() / "shared" / "fixtures" / "lcm_frames.json") as f: - VECTORS: list[dict[str, Any]] = json.load(f)["vectors"] -MESSAGES = dict(build_messages()) - - -class _Bare: - msg_name = "Bare" - - def lcm_encode(self) -> bytes: - return b"" - - -def test_default_encoding() -> None: - assert default_encoding(PoseStamped, "rx") == "geometry_msgs.PoseStamped.lcm.v1" - assert default_encoding(Transform, "rx") == "tf2_msgs.TFMessage.lcm.v1" - assert default_encoding(PoseStamped, "tx") == "json.v1" - assert default_encoding(dict, "rx") == "json.v1" - # A generated class is named by its module path: its msg_name may be bare. - assert lcm_visualization_msgs.MarkerArray.msg_name == "MarkerArray" - assert ( - default_encoding(lcm_visualization_msgs.MarkerArray, "rx") - == "visualization_msgs.MarkerArray.lcm.v1" - ) - with pytest.raises(ValueError, match=r"sensor_msgs\.Image has no default.*jpeg\.v1"): - default_encoding(Image, "rx") - - -def test_schema_class_for_shapes() -> None: - assert schema_class_for(Twist) is lcm_geometry_msgs.Twist # direct subclass - assert schema_class_for(Odometry) is lcm_nav_msgs.Odometry # composition - # PoseStamped inherits Pose's fingerprint method; msg_name decides. - assert schema_class_for(PoseStamped) is lcm_geometry_msgs.PoseStamped - assert schema_class_for(Transform) is lcm_tf2_msgs.TFMessage - assert schema_class_for(lcm_std_msgs.Int64) is lcm_std_msgs.Int64 - markers = lcm_visualization_msgs.MarkerArray - assert schema_class_for(markers) is markers - - -@pytest.mark.parametrize( - ("message_type", "match"), - [ - (JointTrajectory, "declares its own LCM fingerprint"), - (MotorCommandArray, r"dimos_lcm.*@web_encoder"), - (dict, "not a DimOS message"), - (_Bare, r"must be '\.', got 'Bare'"), - ], - ids=["foreign_fingerprint", "no_schema", "no_lcm_encode", "unqualified_name"], -) -def test_schema_class_for_rejections(message_type: type, match: str) -> None: - with pytest.raises(ValueError, match=match): - schema_class_for(message_type) - - -def test_export_schema_dims() -> None: - scan = export_schema(lcm_sensor_msgs.LaserScan)["structs"]["sensor_msgs.LaserScan"] - assert ["ranges", "float", ["ranges_length"]] in scan - imu = export_schema(lcm_sensor_msgs.Imu)["structs"]["sensor_msgs.Imu"] - assert ["orientation_covariance", "double", [9]] in imu - assert export_schema(lcm_std_msgs.Empty)["structs"] == {"std_msgs.Empty": []} - - -def test_encode_lcm_v1_guards_the_fingerprint() -> None: - params = {"lcm": export_schema(lcm_geometry_msgs.PoseStamped)} - check_lcm_params(params) - msg = PoseStamped(ts=1.0, position=[1.0, 2.0, 3.0], orientation=[0.0, 0.0, 0.0, 1.0]) - assert encode_lcm_v1(msg, params) == msg.lcm_encode() - with pytest.raises(ValueError, match=r"Pose\.lcm_encode\(\) wrote fingerprint .* expects 6a82"): - encode_lcm_v1(Pose(1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 1.0), params) - for bad in ({}, {"lcm": {"type": "x", "fp": "zz", "structs": {}}}, {"lcm": {"type": "x"}}): - with pytest.raises(ValueError, match=r"need params\['lcm'\]"): - check_lcm_params(bad) - - -@pytest.mark.parametrize("vec", VECTORS, ids=lambda v: v["name"]) -def test_golden_vector(vec: dict[str, Any]) -> None: - msg = MESSAGES[vec["name"]] - generated = schema_class_for(type(msg)) - assert vec["encoding"] == default_encoding(type(msg), "rx") - assert export_schema(generated) == vec["schema"] - payload = base64.b64decode(vec["payload_b64"]) - assert encode_lcm_v1(msg, {"lcm": vec["schema"]}) == payload - assert walk_value(generated.lcm_decode(payload), generated) == vec["value"] - - -def test_golden_vectors_cover_every_message() -> None: - assert [v["name"] for v in VECTORS] == list(MESSAGES) - - -def test_import_stays_light() -> None: - # cockpit.py and codecs.py import this module at module scope; the LCM - # classes load only when a schema is resolved. - code = ( - "import sys; import dimos.web.lcm_codec; " - "assert 'numpy' not in sys.modules; " - "assert 'dimos_lcm' not in sys.modules; " - "assert 'dimos.web.relay_bridge.relay_bridge_module' not in sys.modules" - ) - subprocess.run([sys.executable, "-c", code], check=True) diff --git a/examples/message-project/rust/Cargo.lock b/examples/message-project/rust/Cargo.lock index 8ce3ef42073..65d0d347351 100644 --- a/examples/message-project/rust/Cargo.lock +++ b/examples/message-project/rust/Cargo.lock @@ -3847,9 +3847,9 @@ dependencies = [ [[package]] name = "zenoh" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f7ad636296bfdb8a86bbbc9bae1b6a57f0e77b715769f99885ca2ca2368c73a" +checksum = "453ba5d28a1197653aae4bb024fd74a9d5c7051d5a19a77f3a43d83f40c22584" dependencies = [ "ahash", "arc-swap", @@ -3899,18 +3899,18 @@ dependencies = [ [[package]] name = "zenoh-buffers" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a81e52197e18bc75587c566df732726aee6e0a1c8779d7a2ff609f461f6b1a07" +checksum = "555a8169c9888fc5571f1b38d2843a63db786621e153abad1b075c40b20e83b4" dependencies = [ "zenoh-collections", ] [[package]] name = "zenoh-codec" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "31b792ce46854f252af31869030224c2da6bd9d72a1cd2812e29b457444a07b1" +checksum = "bc959f727893eab9b66ec051c8f4a2d6f66b4ee70bf23f2ed244aa754705e854" dependencies = [ "rand 0.8.8", "tracing", @@ -3922,18 +3922,18 @@ dependencies = [ [[package]] name = "zenoh-collections" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42b897fae0ddce178e3e58d3bb9ca131dd2a301dd869ebded37d9e1f3784aac0" +checksum = "0c05abd8dcfd2239a8dfb5b79bd4ed85094c6bb6d6b991d3529bff364cc48e29" dependencies = [ "ahash", ] [[package]] name = "zenoh-config" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b39224e3029571e0572f88590eb9078bcd60dedc6381ad1a5c8ad5fb7a474bf" +checksum = "59799d556b4fac79886cb0a9c90f3bc48d0894816b99c8a90a3509b1809a02dd" dependencies = [ "json5", "nonempty-collections", @@ -3957,9 +3957,9 @@ dependencies = [ [[package]] name = "zenoh-core" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9335e8f6509fcbec68a073372c13a69ba90c3a284a86427257607d75586b6344" +checksum = "975073c51353c1c33e81a98370bf3560863540913e5f30091923f5a328fc46d4" dependencies = [ "lazy_static", "tokio", @@ -3969,9 +3969,9 @@ dependencies = [ [[package]] name = "zenoh-crypto" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b496e8c984260c4d80ce58b42956081dc8656f0b875766e5ae0f8e39510fdee" +checksum = "e281a36cfd351ba9232a24899d946a782d0ff3ac0c3d759f87f3e7cbe09a4553" dependencies = [ "aes", "hmac", @@ -3983,9 +3983,9 @@ dependencies = [ [[package]] name = "zenoh-keyexpr" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "af891f8d9f52c3e8617712b65b08fa9ae02f1f1b57e2739ff194bb65df1255df" +checksum = "02c6fcb7f18846b4e9b097c61ed7c656ec8d7858ab54a12d5f5025fe1514caae" dependencies = [ "getrandom 0.2.17", "hashbrown 0.16.1", @@ -3999,9 +3999,9 @@ dependencies = [ [[package]] name = "zenoh-link" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04f796c9df0ef7736d140363507498d8c5d0acf621c158113f122aa992979d3d" +checksum = "c55c4fbc923328f0b62b1b492723ff53b757a36421f70218e652f0aa899cde6c" dependencies = [ "zenoh-config", "zenoh-link-commons", @@ -4013,9 +4013,9 @@ dependencies = [ [[package]] name = "zenoh-link-commons" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4dd110109f4f89450491bd8942152020f632eb86d8fac10349a186185f654873" +checksum = "6d52520b4c22d05173c8621e3cbe18d838a17d74477891ed9c59a7aeedaa325d" dependencies = [ "async-trait", "base64 0.22.1", @@ -4050,9 +4050,9 @@ dependencies = [ [[package]] name = "zenoh-link-quic_datagram" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97b2cddd28955877b7eceedfaf09c801aac31ab8f0cac32c936f31358f5af133" +checksum = "1af3abf41c719c82ca356864b5ddeeb7c77cb29fe777776b40b31eb4d19de288" dependencies = [ "async-trait", "rustls-webpki", @@ -4067,9 +4067,9 @@ dependencies = [ [[package]] name = "zenoh-link-tcp" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4801cc6fb608b5069f3ea6b82484c5f727019764e90e14838fe6930a9dffd9a1" +checksum = "4d4702d55fe9f17cbd1d3b601901f8ba0a693a7ad1ce685a1ea4c60f775552b7" dependencies = [ "async-trait", "socket2 0.5.10", @@ -4085,9 +4085,9 @@ dependencies = [ [[package]] name = "zenoh-link-udp" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6785a0b6f0b94131c002a8358aaef3eb99f5ab0e0cd444f28d1f1fac38bd4091" +checksum = "b140e7b65715fe8b29b706b2b9fc3009fc6181346931c8926d290113fe38df4b" dependencies = [ "async-trait", "libc", @@ -4108,9 +4108,9 @@ dependencies = [ [[package]] name = "zenoh-macros" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c813ccde166cf0527402d46615257e0d152140e66f8ecd685721f9399251407" +checksum = "f483b343fd79b6c1bd7d79f130b5c6c208e7fef243a4d7966a6bf9fdecbd70f8" dependencies = [ "proc-macro2", "quote", @@ -4120,9 +4120,9 @@ dependencies = [ [[package]] name = "zenoh-plugin-trait" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "95cd9d698675a5b2f3e958a35df56f6d60f0388669158055b436c0cbe5357ad8" +checksum = "7d57f491f63afa7413d68ee68e4de2bc9ce361c4ac89a76d5a0a0e3ca0624e21" dependencies = [ "git-version", "libloading", @@ -4138,9 +4138,9 @@ dependencies = [ [[package]] name = "zenoh-protocol" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce5505e2569421d3cb6792d4b3d9a6589d04add3bf34c0a7c37c03c2573b79ad" +checksum = "4492192986df044627f33c8eeb8aa237fb4d3bbd2d8cc68bc7eee16748c2fb61" dependencies = [ "const_format", "rand 0.8.8", @@ -4154,18 +4154,18 @@ dependencies = [ [[package]] name = "zenoh-result" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cf23203517d6dfe04393c42053093c7cffd9e385015e21c85dafaff9a342cccf" +checksum = "8d6d0ca24186d9cfbf14f2d08b8b0a09afac690fe5cd68899d8f896f9247a3dd" dependencies = [ "anyhow", ] [[package]] name = "zenoh-runtime" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ddf72c7310a090b0b22496649274661de8a17ca0d0674af8f72f184a3ed78664" +checksum = "91e8ea6ed9fb3300d02bc303b3202a21a758aa5d7a1e7dcc38610428910af9f9" dependencies = [ "lazy_static", "ron", @@ -4178,9 +4178,9 @@ dependencies = [ [[package]] name = "zenoh-shm" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d8f49ff9d36be3fdeada23103c5ba10d6e7159269bd8b07fb744448005fb5734" +checksum = "1600eef8e5b1148585b020985bd4b7851d86b585dc90dd306343127e552c591f" dependencies = [ "advisory-lock", "async-trait", @@ -4208,9 +4208,9 @@ dependencies = [ [[package]] name = "zenoh-stats" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f610d551359e0d8864138d4c4c11ed8e3c9f6991717f4eef198ab323b4ff01f4" +checksum = "b2f2972f429c8686e6d4963649b96e1f74bee0826fc80d398292913d0bb852fe" dependencies = [ "ahash", "prometheus-client", @@ -4222,9 +4222,9 @@ dependencies = [ [[package]] name = "zenoh-sync" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9792a8fa9b0f30ebabd3e6b278441178b6832462e959388aa18f92a93527935b" +checksum = "d57ed2156b61df6e877b03fb13dd72227d3d286a7fa689120109bad289ddf12b" dependencies = [ "arc-swap", "event-listener", @@ -4237,9 +4237,9 @@ dependencies = [ [[package]] name = "zenoh-task" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fbb419757b0c6e824bebf34959302e6edf5a7caea4e725de88007a56f3ba4d4d" +checksum = "887115c0f81c72ff8c9b99e6fccd9e1c3829d83e16b8b063174786b9fba08810" dependencies = [ "futures", "tokio", @@ -4251,9 +4251,9 @@ dependencies = [ [[package]] name = "zenoh-transport" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fc431ec1c066d7047c601ae4b4d4a359882c5464cd9c4ff3fe905ff6e10017af" +checksum = "72bae2e0aca56ac6431bd6c37c6ac3fcc9367446cbb356bd44137cc01d1493ee" dependencies = [ "async-trait", "crossbeam-utils", @@ -4289,9 +4289,9 @@ dependencies = [ [[package]] name = "zenoh-util" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3f99b13044636f453bcb49ca520db4b1400c8e444fdf578e63a1aad0ed9004e0" +checksum = "43a473dcc62026b9d9589acc22c62b04cb7dab4070bba5543d2afac741ddf947" dependencies = [ "async-trait", "const_format", From 7c3d411ebae0ee1aaa95de72db9a807f485f7821 Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 15 Oct 2026 00:17:11 -0700 Subject: [PATCH 09/11] refactor(cdr): migrate runtime callers onto published generator --- Cargo.toml | 1 + dimos/agents/mcp/test_generated_image.py | 3 +- dimos/agents/skills/navigation.py | 3 +- dimos/agents/skills/person_follow.py | 27 +- .../skills/test_generated_navigation.py | 14 +- .../skills/test_generated_person_follow.py | 35 +- .../skills/test_unitree_skill_container.py | 3 + .../skills/visual_servoing/bbox_navigation.py | 4 +- .../visual_servoing/detection_navigation.py | 4 +- .../test_visual_servoing_2d.py | 56 ++ .../visual_servoing/visual_servoing_2d.py | 10 +- dimos/agents/test_generated_vlm.py | 5 +- dimos/agents/typesafe/navigation.py | 46 +- dimos/agents/typesafe/test_navigation.py | 68 +- dimos/agents/typesafe/test_world_state.py | 102 ++- dimos/agents/typesafe/world_state.py | 152 +++- dimos/cli/bake/test_bake_e2e.py | 3 +- dimos/cli/cameracalibrate/cameracalibrate.py | 10 +- .../cameracalibrate/test_cameracalibrate.py | 65 +- dimos/cli/hardware/g1.py | 12 +- dimos/cli/spy/test_spy.py | 6 +- dimos/cli/test_topic.py | 27 +- dimos/cli/topic.py | 4 +- dimos/cloud/preview.py | 101 ++- dimos/cloud/test_preview.py | 111 ++- dimos/conftest.py | 39 +- dimos/control/benchmarking/benchmark.py | 10 +- dimos/control/benchmarking/paths.py | 36 +- dimos/control/benchmarking/score.py | 4 +- dimos/control/benchmarking/test_benchmark.py | 3 +- dimos/control/contract/test_presets.py | 5 +- dimos/control/contract/test_validate.py | 37 +- dimos/control/contract/validate.py | 15 +- dimos/control/coordinator.py | 21 +- .../cartesian_ik_task/cartesian_ik_task.py | 8 +- .../test_cartesian_ik_task.py | 66 +- .../eef_twist_task/test_eef_twist_task.py | 27 +- .../holonomic_pose_follower_task.py | 3 +- .../test_holonomic_pose_follower_task.py | 19 +- .../test_progress_reference.py | 35 +- .../path_follower_task/path_follower_task.py | 48 +- .../test_path_follower_task.py | 3 +- .../test_planar_base_trajectory_task.py | 22 +- dimos/control/tasks/pose_target_ik.py | 19 +- .../test_rpp_path_follower_task.py | 7 +- .../tasks/teleop_ik_task/teleop_ik_task.py | 7 +- .../teleop_ik_task/test_teleop_ik_task.py | 28 +- dimos/control/tasks/test_pose_target_ik.py | 167 ++++- .../trajectory_task/test_trajectory_task.py | 115 ++- .../tasks/trajectory_task/trajectory_task.py | 26 +- .../tasks/velocity_task/test_velocity_task.py | 30 +- .../tasks/velocity_task/velocity_task.py | 2 +- dimos/control/test_control.py | 248 +++++-- dimos/control/test_coordinator_commands.py | 17 +- dimos/control/test_coordinator_routing.py | 282 +++++++- dimos/control/test_frame_pose_tf.py | 26 +- dimos/control/test_operator_hold.py | 26 +- dimos/control/test_per_robot_joint_states.py | 15 +- dimos/control/tick_loop.py | 55 +- .../coordination/blueprint_config/values.py | 7 +- .../coordination/test_module_coordinator.py | 25 +- dimos/core/native_module.py | 35 +- dimos/core/test_cdr_transport.py | 30 +- dimos/core/test_native_module.py | 58 +- dimos/core/test_transport.py | 2 +- dimos/core/test_transport_factory.py | 6 +- dimos/core/test_zenoh_transport.py | 28 +- dimos/core/transport.py | 181 ++++- dimos/core/transport_factory.py | 2 +- dimos/e2e_tests/cdr_replay_fixture.py | 6 +- dimos/e2e_tests/conftest.py | 21 +- dimos/e2e_tests/dim_sim_client.py | 13 +- dimos/e2e_tests/lcm_spy.py | 24 +- dimos/e2e_tests/test_cdr_spy.py | 31 +- dimos/e2e_tests/test_control_coordinator.py | 87 ++- .../e2e_tests/test_custom_channel_browser.py | 21 +- dimos/e2e_tests/test_dimsim_typesafe_agent.py | 6 +- dimos/e2e_tests/test_dimsim_walk_forward.py | 5 +- dimos/e2e_tests/test_lcm_channel_browser.py | 4 +- .../test_manipulation_planning_groups.py | 8 +- dimos/e2e_tests/test_map_click_browser.py | 11 +- dimos/e2e_tests/test_rtc_video_browser.py | 8 +- dimos/evals/agents/test_no_dimos.py | 5 +- dimos/evals/conftest.py | 8 +- dimos/evals/environments/habitat.py | 20 +- dimos/evals/environments/image_file.py | 16 +- .../environments/lib/benchmark_scenes.py | 4 +- .../evals/environments/lib/recorded_poses.py | 11 +- dimos/evals/environments/mujoco_sim.py | 7 +- dimos/evals/environments/test_dimsim.py | 10 +- dimos/evals/environments/test_habitat.py | 48 +- dimos/evals/environments/test_mujoco_sim.py | 83 ++- dimos/evals/scorers.py | 37 +- dimos/evals/suites/dimsim_house.py | 2 +- dimos/evals/suites/lib/generate.py | 7 +- dimos/evals/test_evals.py | 7 +- dimos/evals/test_mem2_wiring.py | 7 +- dimos/evals/vqa/primitives/test_edge_tam.py | 24 +- dimos/evals/vqa/test_contracts.py | 4 +- dimos/evals/vqa/test_pointcloud_frame.py | 29 +- .../isolated_python/test_generated_example.py | 5 +- dimos/experimental/memory/README.md | 76 +- dimos/experimental/memory/rust/Cargo.toml | 2 +- dimos/experimental/memory/rust/flake.nix | 50 +- .../experimental/memory/rust/src/decoding.rs | 301 ++++---- .../experimental/memory/rust/src/encoding.rs | 85 +-- dimos/experimental/memory/rust/src/lib.rs | 372 ++++++---- .../memory/rust/src/store/mcap.rs | 160 +---- .../experimental/memory/rust_cli_recorder.py | 4 +- dimos/experimental/memory/rust_recorder.py | 77 +- .../experimental/memory/test_rust_recorder.py | 104 ++- .../memory/test_rust_recorder_e2e.py | 110 +-- .../robot/bosdyn/spot/effectors/high_level.py | 7 +- .../bosdyn/spot/effectors/test_high_level.py | 5 +- .../robot/bosdyn/spot/recorder.py | 7 +- dimos/experimental/robot/bosdyn/spot/rerun.py | 35 +- .../robot/bosdyn/spot/test_utils.py | 15 +- dimos/experimental/robot/bosdyn/spot/utils.py | 10 +- dimos/experimental/security_demo/conftest.py | 20 +- .../security_demo/security_module.py | 33 +- .../security_demo/test_security_module.py | 24 +- .../experimental/world_belief/test_absence.py | 23 +- .../world_belief/test_generated_recall.py | 6 +- .../world_belief/worldbelief_recorder.py | 6 +- .../world_belief/xarm6_blueprint.py | 6 +- .../hardware/drive_trains/flowbase/config.py | 4 +- .../drive_trains/transport/test_adapter.py | 34 +- .../camera/gstreamer/test_generated_frames.py | 3 +- dimos/hardware/sensors/camera/module.py | 21 +- .../sensors/camera/realsense/blueprints.py | 10 +- .../sensors/camera/realsense/camera.py | 4 + .../sensors/camera/realsense/rust/Cargo.lock | 35 + .../sensors/camera/realsense/rust/Cargo.toml | 2 +- .../sensors/camera/realsense/rust/src/main.rs | 11 +- dimos/hardware/sensors/camera/test_module.py | 9 +- dimos/hardware/sensors/camera/test_webcam.py | 26 +- dimos/hardware/sensors/camera/v4l2/module.py | 2 +- .../sensors/camera/v4l2/rust/Cargo.toml | 2 +- dimos/hardware/sensors/camera/webcam.py | 22 +- dimos/hardware/sensors/camera/zed/camera.py | 25 +- .../sensors/camera/zed/test_generated_tf.py | 16 +- .../lidar/common/point_cloud_utils.hpp | 6 +- .../sensors/lidar/fastlio2/cpp/main.cpp | 4 +- .../hardware/sensors/lidar/fastlio2/module.py | 23 +- .../sensors/lidar/fastlio2/recorder.py | 7 +- .../sensors/lidar/livox/rust/Cargo.toml | 2 +- .../sensors/lidar/livox/rust/src/module.rs | 10 +- .../hardware/sensors/lidar/livox/test_e2e.py | 45 +- .../hardware/sensors/lidar/pointlio/module.py | 2 +- .../sensors/lidar/pointlio/recorder.py | 7 +- .../sensors/lidar/pointlio/rust/Cargo.toml | 12 +- .../sensors/lidar/pointlio/rust/src/module.rs | 156 +++-- .../sensors/lidar/test_generated_recorders.py | 42 +- .../hardware/whole_body/transport/adapter.py | 11 +- .../whole_body/transport/test_adapter.py | 15 +- dimos/hosted/fragment.py | 4 +- dimos/hosted/fragment_compiler.py | 2 +- dimos/hosted/test_fragment_compiler.py | 6 +- dimos/imitation/collection/episode_monitor.py | 6 +- dimos/imitation/collection/recorder.py | 8 +- dimos/imitation/collection/test_blueprint.py | 2 +- .../collection/test_episode_monitor.py | 13 +- dimos/imitation/dataprep/core.py | 30 +- dimos/imitation/dataprep/test_core.py | 214 +++++- .../dataprep/test_generated_images.py | 37 +- dimos/imitation/dataprep/test_mcap_source.py | 2 +- dimos/imitation/dataprep/test_schema.py | 2 +- dimos/imitation/test_datacollection_e2e.py | 28 +- dimos/manipulation/conftest.py | 8 +- .../cartesian_motion_controller.py | 5 +- .../test_generated_legacy_controllers.py | 44 +- .../joint_trajectory_controller.py | 5 +- dimos/manipulation/execution_manager.py | 26 +- .../grasping/grasp_gen_x/test_module.py | 31 +- .../manipulation/grasping/heuristic_grasp.py | 3 +- .../grasping/test_heuristic_grasp.py | 15 +- dimos/manipulation/manipulation_module.py | 167 +++-- dimos/manipulation/manipulation_skills.py | 40 +- dimos/manipulation/pick_and_place_module.py | 75 +- .../planning/groups/test_planning_groups.py | 40 +- dimos/manipulation/planning/groups/utils.py | 23 +- .../kinematics/drake_optimization_ik.py | 18 +- .../planning/kinematics/jacobian_ik.py | 66 +- .../planning/kinematics/pink_ik.py | 18 +- .../test_drake_optimization_ik_selection.py | 33 +- .../kinematics/test_jacobian_ik_selection.py | 15 +- .../planning/kinematics/test_pink_ik.py | 495 +++++++------ .../planning/monitor/robot_state_monitor.py | 31 +- .../monitor/test_robot_state_monitor.py | 56 +- .../planning/monitor/test_world_monitor.py | 344 +++++++-- .../planning/monitor/world_monitor.py | 74 +- .../planning/planners/roboplan_planner.py | 46 +- .../planning/planners/rrt_planner.py | 29 +- .../planning/planners/selected_joint_space.py | 34 +- .../planners/test_rrt_planner_selection.py | 211 +++++- dimos/manipulation/planning/spec/config.py | 14 +- .../planning/spec/test_joint_space.py | 22 +- .../joint_trajectory_generator.py | 15 +- .../trajectory_generator/parametrizer.py | 7 +- .../roboplan_toppra_parametrizer.py | 8 +- .../trajectory_generator/test_parametrizer.py | 114 ++- .../test_roboplan_toppra_parametrizer.py | 32 +- .../test_simple_parametrizer.py | 55 +- .../manipulation/planning/utils/path_utils.py | 32 +- .../planning/utils/point_cloud_self_filter.py | 36 +- .../utils/roboplan_point_cloud_self_filter.py | 17 +- .../utils/test_point_cloud_self_filter.py | 55 +- .../test_roboplan_point_cloud_self_filter.py | 207 ++++-- .../planning/world/drake_world.py | 19 +- .../planning/world/roboplan_world.py | 11 +- .../world/test_drake_world_planning_groups.py | 95 ++- dimos/manipulation/sdk.py | 11 +- dimos/manipulation/test_execution_manager.py | 43 +- .../test_generated_plan_materialization.py | 74 +- .../manipulation/test_manipulation_module.py | 101 ++- .../test_manipulation_monitor_preview.py | 107 ++- .../test_manipulation_primitives.py | 69 +- dimos/manipulation/test_manipulation_spec.py | 49 +- dimos/manipulation/test_manipulation_unit.py | 437 +++++++++--- .../manipulation/test_pick_and_place_unit.py | 113 ++- dimos/manipulation/test_plan_execution.py | 83 ++- dimos/manipulation/test_planning_factory.py | 7 +- dimos/manipulation/test_roboplan.py | 616 +++++++++++++--- .../manipulation/test_roboplan_integration.py | 178 ++++- dimos/manipulation/test_sdk.py | 45 +- dimos/manipulation/visualization/operator.py | 18 +- .../visualization/test_factory.py | 59 +- .../visualization/test_operator.py | 162 ++++- dimos/manipulation/visualization/viser/gui.py | 80 ++- .../visualization/viser/test_gui.py | 83 ++- .../viser/test_scene_grasp_proposals.py | 5 +- .../viser/test_scene_obstacles.py | 11 +- .../viser/test_viser_visualization.py | 43 +- .../viser/test_visualizer_lifecycle.py | 32 +- dimos/mapping/cli/map.py | 265 ++++--- dimos/mapping/cli/test_pgo_accumulate.py | 12 +- dimos/mapping/cli/test_pose_fill.py | 43 +- dimos/mapping/cli/test_stream_selection.py | 29 +- dimos/mapping/cli/test_view.py | 7 +- dimos/mapping/cli/view.py | 3 +- dimos/mapping/dim_slam/rust/Cargo.lock | 35 + dimos/mapping/dim_slam/rust/Cargo.toml | 2 +- dimos/mapping/dim_slam/rust/flake.nix | 7 +- dimos/mapping/dim_slam/rust/src/main.rs | 6 +- .../mapping/dim_slam/rust/src/msg_convert.rs | 11 +- dimos/mapping/experimental/mesh.py | 50 +- dimos/mapping/experimental/mesh_rrd.py | 18 +- dimos/mapping/experimental/test_mesh.py | 6 +- dimos/mapping/occupancy/conftest.py | 15 +- dimos/mapping/occupancy/gradient.py | 12 +- dimos/mapping/occupancy/operations.py | 8 +- dimos/mapping/occupancy/path_resampling.py | 15 +- .../test_generated_visualizations.py | 47 +- dimos/mapping/occupancy/test_path_mask.py | 19 +- .../mapping/occupancy/test_path_resampling.py | 39 +- dimos/mapping/odometry_hist.py | 4 +- dimos/mapping/pointclouds/demo.py | 6 +- dimos/mapping/pointclouds/occupancy.py | 115 +-- .../pointclouds/test_generated_occupancy.py | 17 +- dimos/mapping/pointclouds/test_occupancy.py | 9 +- dimos/mapping/ray_tracing/module.py | 15 +- dimos/mapping/ray_tracing/rust/Cargo.toml | 22 +- dimos/mapping/ray_tracing/rust/src/module.rs | 163 ++--- dimos/mapping/ray_tracing/test_transformer.py | 44 +- dimos/mapping/ray_tracing/test_viz.py | 16 +- dimos/mapping/ray_tracing/transformer.py | 24 +- .../ray_tracing/utils/test_loaded_map.py | 31 +- dimos/mapping/ray_tracing/viz.py | 16 +- .../relocalization/lidar/relocalize.py | 17 +- dimos/mapping/relocalization/lidar/replay.py | 2 +- .../relocalization/lidar/test_replay.py | 56 +- dimos/mapping/relocalization/lidar/tune.py | 54 +- dimos/mapping/relocalization/module.py | 48 +- dimos/mapping/relocalization/test_module.py | 34 +- dimos/mapping/test_generated_costmapper.py | 10 +- dimos/mapping/test_generated_odometry_hist.py | 38 +- dimos/mapping/test_generated_voxels.py | 21 +- dimos/mapping/test_voxels.py | 16 +- dimos/mapping/voxels/grid.py | 58 +- dimos/mapping/voxels/impl/packed.py | 20 +- dimos/mapping/voxels/test_grid.py | 3 +- dimos/mapping/voxels/test_stamped.py | 14 +- dimos/memory/cli/render.py | 2 +- dimos/memory/cli/test_render.py | 12 +- dimos/memory/codecs/README.md | 79 +-- dimos/memory/codecs/base.py | 68 +- dimos/memory/codecs/cdr.py | 14 +- dimos/memory/codecs/json.py | 4 +- dimos/memory/codecs/test_codecs.py | 345 ++++----- dimos/memory/demo_data.py | 4 +- dimos/memory/module.py | 24 +- dimos/memory/semantic_search.py | 10 +- dimos/memory/store/mcap.py | 92 +-- dimos/memory/store/test_mcap.py | 71 +- dimos/memory/store/test_mcap_wire.py | 71 +- dimos/memory/tap.py | 14 +- dimos/memory/test_replay_module.py | 62 +- dimos/memory/test_tap.py | 190 ++++- dimos/memory/tool_e2e.py | 10 +- dimos/memory/type/observation.py | 2 +- dimos/memory/type/test_generated_query.py | 46 +- dimos/memory/vis/space/rerun.py | 5 +- .../vis/space/test_generated_geometry.py | 52 +- dimos/memory/vis/space/test_space.py | 183 ++++- dimos/message_codegen/test_builtin_catalog.py | 2 +- dimos/message_codegen/test_builtin_install.py | 2 +- .../test_builtin_preparation.py | 6 +- dimos/models/embedding/siglip2.py | 9 +- dimos/models/vl/test_generated_images.py | 3 +- dimos/msgs/camera_info.py | 28 +- dimos/msgs/contacts.py | 20 + dimos/msgs/control.py | 65 ++ dimos/msgs/control_msgs/test_ControlValues.py | 182 ----- .../foxglove_msgs/test_CompressedVideo.py | 25 +- dimos/msgs/geometry.py | 25 +- dimos/msgs/geometry_msgs/test_PointStamped.py | 10 +- dimos/msgs/geometry_msgs/test_Pose.py | 109 ++- dimos/msgs/geometry_msgs/test_PoseStamped.py | 15 +- .../geometry_msgs/test_PoseWithCovariance.py | 53 +- .../test_PoseWithCovarianceStamped.py | 32 +- dimos/msgs/geometry_msgs/test_Quaternion.py | 17 +- dimos/msgs/geometry_msgs/test_Transform.py | 90 ++- dimos/msgs/geometry_msgs/test_Twist.py | 36 +- dimos/msgs/geometry_msgs/test_TwistStamped.py | 15 +- .../geometry_msgs/test_TwistWithCovariance.py | 41 +- .../test_TwistWithCovarianceStamped.py | 26 +- dimos/msgs/geometry_msgs/test_Vector3.py | 41 +- .../msgs/geometry_msgs/test_WrenchStamped.py | 42 +- dimos/msgs/geometry_msgs/tool_publish.py | 5 +- dimos/msgs/helpers.py | 2 +- dimos/msgs/image.py | 13 +- .../manipulation_msgs/test_grasp_candidate.py | 30 +- dimos/msgs/mesh.py | 62 ++ dimos/msgs/nav_msgs/test_LineSegments3D.py | 126 ++-- dimos/msgs/nav_msgs/test_OccupancyGrid.py | 117 +++- dimos/msgs/nav_msgs/test_Odometry.py | 183 ++++- dimos/msgs/nav_msgs/test_Path.py | 56 +- dimos/msgs/occupancy.py | 13 +- dimos/msgs/pointcloud.py | 68 +- dimos/msgs/protocol.py | 16 +- dimos/msgs/sensor_msgs/test_CameraInfo.py | 127 ++-- .../msgs/sensor_msgs/test_CompressedImage.py | 28 +- dimos/msgs/sensor_msgs/test_Joy.py | 121 ++-- dimos/msgs/sensor_msgs/test_PointCloud2.py | 259 ++++--- dimos/msgs/sensor_msgs/test_image.py | 12 +- dimos/msgs/sim_msgs/Contacts.py | 54 -- dimos/msgs/sim_msgs/test_Contacts.py | 26 +- dimos/msgs/std_msgs/test_header.py | 10 +- dimos/msgs/test_camera_info.py | 27 +- dimos/msgs/test_control.py | 114 +++ dimos/msgs/test_generated_projection.py | 72 +- dimos/msgs/test_generated_views.py | 94 ++- dimos/msgs/test_geometry.py | 73 +- dimos/msgs/test_mesh.py | 37 + dimos/msgs/test_msg_construction.py | 232 ++++-- dimos/msgs/test_occupancy.py | 90 ++- dimos/msgs/test_pointcloud.py | 136 +++- dimos/msgs/test_pointcloud_colors.py | 31 +- dimos/msgs/test_regions.py | 101 +++ dimos/msgs/test_time.py | 3 +- dimos/msgs/test_trajectory.py | 44 +- dimos/msgs/tf2_msgs/test_TFMessage.py | 44 +- dimos/msgs/tf2_msgs/test_TFMessage_lcmpub.py | 6 +- dimos/msgs/time.py | 12 + .../vision_msgs/test_detection3d_rerun.py | 30 +- .../dannav/geometry/path_distancer.py | 5 +- .../holonomic_tc/holonomic_path_controller.py | 14 +- .../holonomic_tracking_controller.py | 4 +- .../dannav/holonomic_tc/module.py | 17 +- .../holonomic_tc/test_dan_holonomic_tc.py | 14 +- .../dannav/local_planner/module.py | 19 +- .../local_planner/test_dan_local_planner.py | 58 +- .../test_generated_frontiers.py | 118 ++++ .../wavefront_frontier_goal_selector.py | 2 +- .../experimental/patrolling/module.py | 6 +- .../patrolling/routers/base_patrol_router.py | 6 +- .../routers/coverage_patrol_router.py | 24 +- .../routers/frontier_patrol_router.py | 15 +- .../routers/random_patrol_router.py | 17 +- .../patrolling/routers/visitation_history.py | 2 +- .../patrolling/test_create_patrol_router.py | 37 +- .../patrolling/test_generated_patrol.py | 68 ++ .../experimental/patrolling/utilities.py | 4 +- .../mls_planner/mls_planner_native.py | 22 +- .../mls_planner/rust/Cargo.toml | 24 +- .../mls_planner/rust/src/module.rs | 300 ++++---- .../mls_planner/test_start_relay.py | 71 +- .../mls_planner/test_transformer.py | 35 +- .../global_planner/mls_planner/test_viz.py | 113 ++- .../global_planner/mls_planner/transformer.py | 35 +- .../mls_planner/utils/plan_rrd.py | 95 ++- .../mls_planner/utils/test_plan_rrd.py | 64 ++ .../global_planner/mls_planner/viz.py | 103 ++- dimos/navigation/global_planner/test_viz.py | 69 +- dimos/navigation/global_planner/viz.py | 15 +- dimos/navigation/go2/loop_closure/eval.py | 27 +- dimos/navigation/go2/loop_closure/pgo.py | 222 ++++-- dimos/navigation/go2/loop_closure/pgo_auto.py | 4 +- dimos/navigation/go2/loop_closure/test_pgo.py | 180 +++-- .../go2/loop_closure/test_pgo_auto.py | 108 +++ .../go2/loop_closure/test_pgo_generated.py | 177 +++++ .../go2/loop_closure/test_pgo_map.py | 55 +- .../go2/replanning_a_star/controllers.py | 10 +- .../go2/replanning_a_star/global_planner.py | 7 +- .../go2/replanning_a_star/local_planner.py | 12 +- .../go2/replanning_a_star/min_cost_astar.py | 9 +- .../go2/replanning_a_star/module.py | 4 +- .../test_generated_navigation.py | 216 ++++++ .../test_generated_pipeline.py | 130 ++++ .../replanning_a_star/test_global_planner.py | 26 +- .../replanning_a_star/test_goal_validator.py | 35 +- .../replanning_a_star/test_min_cost_astar.py | 76 +- .../navigation/local_planner/rust/Cargo.toml | 2 +- .../movement_manager/test_movement_manager.py | 11 +- dimos/navigation/spec.py | 10 +- .../test_generated_bbox_navigation.py | 46 +- .../trajectory_follower/basic/module.py | 41 +- .../trajectory_follower/basic/test_module.py | 88 ++- .../trajectory_follower/fancy/rust/Cargo.toml | 2 +- .../depth2depth_cloud/rust/Cargo.toml | 2 +- .../depth2depth_cloud/rust/flake.nix | 6 +- dimos/perception/detection/conftest.py | 7 +- .../detection/detectors/grounding_dino.py | 8 +- dimos/perception/detection/detectors/owlv2.py | 12 +- dimos/perception/detection/module2D.py | 6 +- dimos/perception/detection/module3D.py | 2 +- dimos/perception/detection/person_tracker.py | 3 +- .../detection/reid/test_generated_module.py | 7 +- .../detection/type/detection2d/bbox.py | 11 +- .../detection/type/detection2d/point.py | 11 +- .../type/detection2d/test_generated_bbox.py | 11 +- .../detection/type/detection3d/base.py | 13 +- .../detection/type/detection3d/bbox.py | 27 +- .../type/detection3d/imageDetections3DPC.py | 5 +- .../detection/type/detection3d/pointcloud.py | 185 +++-- .../type/detection3d/pointcloud_filters.py | 11 +- .../detection3d/test_generated_filters.py | 34 +- .../detection3d/test_generated_projection.py | 81 ++- .../detection3d/test_imageDetections3D.py | 7 +- .../detection/type/detection3d/test_marker.py | 3 +- .../type/detection3d/test_pointcloud.py | 92 ++- dimos/perception/experimental/moduleDB.py | 2 +- dimos/perception/experimental/object.py | 28 +- .../perception/experimental/object_tracker.py | 72 +- .../experimental/object_tracker_2d.py | 10 +- .../experimental/object_tracker_3d.py | 44 +- .../experimental/perceive_loop_skill.py | 15 +- .../temporal_memory/temporal_memory.py | 31 +- .../test_temporal_memory_module.py | 69 +- .../temporal_memory/window_analyzer.py | 3 +- .../experimental/test_generated_tracker.py | 17 +- .../experimental/test_object_generated.py | 41 +- ...test_object_scene_registration_temporal.py | 15 +- .../experimental/test_spatial_memory.py | 17 +- .../fiducial/blueprints/desk_marker_tf.py | 8 +- dimos/perception/fiducial/marker_detect.py | 2 +- dimos/perception/fiducial/marker_pose.py | 2 +- dimos/perception/fiducial/marker_tf_module.py | 5 +- .../fiducial/test_fixture_verification.py | 6 +- .../fiducial/test_generated_marker_detect.py | 27 +- dimos/perception/fiducial/test_helpers.py | 14 +- dimos/perception/fiducial/test_marker_pose.py | 41 +- .../fiducial/test_marker_tf_integration.py | 26 +- .../fiducial/test_marker_tf_module.py | 31 +- dimos/perception/localize/association.py | 22 +- dimos/perception/localize/dandetect.py | 2 +- .../localize/demo_blueprints/xarm_feed.py | 6 +- dimos/perception/localize/localize.py | 22 +- dimos/perception/localize/module.py | 66 +- dimos/perception/localize/rig.py | 76 +- dimos/perception/localize/support_plane.py | 3 +- dimos/perception/localize/test_rig.py | 115 ++- .../localize/tools/tool_localize.py | 84 ++- dimos/porcelain/test_remote_module_source.py | 2 +- .../pubsub/benchmark/test_replay_bench.py | 9 +- dimos/protocol/pubsub/benchmark/testdata.py | 72 +- .../pubsub/benchmark/tool_replay_bench.py | 4 +- dimos/protocol/pubsub/encoders.py | 10 +- dimos/protocol/pubsub/impl/jpeg_lcm.py | 5 +- dimos/protocol/pubsub/impl/jpeg_shm.py | 41 ++ dimos/protocol/pubsub/impl/lcmpubsub.py | 2 +- .../pubsub/impl/rospubsub_conversion.py | 18 +- dimos/protocol/pubsub/impl/shmpubsub.py | 353 ++++++++++ dimos/protocol/pubsub/impl/test_jpeg_cdr.py | 23 +- dimos/protocol/pubsub/impl/test_lcmpubsub.py | 22 +- dimos/protocol/pubsub/impl/test_rospubsub.py | 57 +- .../pubsub/impl/test_rospubsub_conversion.py | 72 +- dimos/protocol/pubsub/impl/test_shmpubsub.py | 210 ++++++ .../impl/webrtc/providers/test_video_track.py | 20 +- .../providers/test_video_track_generated.py | 5 +- .../impl/webrtc/providers/video_track.py | 35 +- .../pubsub/impl/webrtc/test_transport.py | 29 +- .../pubsub/impl/webrtc/tool_blueprint_e2e.py | 11 +- .../pubsub/impl/webrtc/tool_broker_e2e.py | 27 +- dimos/protocol/pubsub/impl/zenohpubsub.py | 2 +- dimos/protocol/pubsub/registry.py | 25 +- dimos/protocol/pubsub/shm/ipc_factory.py | 249 +++++++ dimos/protocol/pubsub/shm/test_ipc_factory.py | 90 +++ dimos/protocol/pubsub/test_pattern_sub.py | 6 +- dimos/protocol/pubsub/test_registry.py | 65 +- dimos/protocol/service/spec.py | 10 + dimos/protocol/service/test_lcmservice.py | 2 +- dimos/protocol/tf/static_tf_publisher.py | 3 +- dimos/protocol/tf/test_static_tf_publisher.py | 3 +- dimos/protocol/tf/test_tf.py | 193 ++++- dimos/protocol/tf/tf.py | 26 +- .../m20/blueprints/m20_kronknav.py | 20 +- dimos/robot/deeprobotics/m20/camera.py | 20 +- dimos/robot/deeprobotics/m20/connection.py | 6 +- .../deeprobotics/m20/test_camera_messages.py | 8 +- .../m20/test_connection_messages.py | 26 +- dimos/robot/diy/alfred/config.py | 7 +- dimos/robot/diy/alfred/effector_high_level.py | 19 +- dimos/robot/diy/alfred/mount_tf.py | 3 +- .../diy/alfred/test_generated_effector.py | 3 +- dimos/robot/diy/alfred/test_mount_tf.py | 3 +- dimos/robot/drone/camera_module.py | 12 +- dimos/robot/drone/connection_module.py | 12 +- dimos/robot/drone/dji_video_stream.py | 4 +- dimos/robot/drone/drone_tracking_module.py | 26 +- dimos/robot/drone/test_drone.py | 9 +- .../blueprints/basic/r1pro_coordinator.py | 20 +- .../manipulation/test_r1pro_planar_preview.py | 24 +- dimos/robot/galaxea/r1pro/connection.py | 177 +++-- dimos/robot/galaxea/r1pro/head_cameras.py | 27 +- dimos/robot/galaxea/r1pro/lio.py | 62 +- .../galaxea/r1pro/test_connection_messages.py | 51 +- .../robot/galaxea/r1pro/test_head_cameras.py | 15 +- .../robot/manipulators/common/coordinators.py | 6 +- .../dual_openyam/test_blueprints.py | 42 +- .../dual_openyam/test_teleop_ik.py | 12 +- .../openarm/test_openarm_teleop.py | 59 +- .../manipulators/openyam/test_teleop_ik.py | 19 +- .../manipulators/xarm/blueprints/grasp.py | 42 +- dimos/robot/manipulators/xarm/config.py | 23 +- dimos/robot/raw_robot_bridge.py | 113 ++- dimos/robot/test_raw_robot_bridge.py | 129 +++- dimos/robot/unitree/b1/connection.py | 18 +- dimos/robot/unitree/b1/joystick_module.py | 19 +- dimos/robot/unitree/b1/test_cdr_forwarding.py | 29 +- dimos/robot/unitree/b1/test_connection.py | 72 +- dimos/robot/unitree/connection.py | 66 +- dimos/robot/unitree/dimsim_connection.py | 91 ++- .../blueprints/basic/unitree_g1_groot_wbc.py | 25 +- .../g1/blueprints/basic/unitree_g1_teleop.py | 8 +- .../primitive/test_unitree_g1_vis.py | 15 +- .../primitive/unitree_g1_primitive_no_nav.py | 41 +- dimos/robot/unitree/g1/config.py | 6 +- .../g1/debug/demo_arrow_control_cmd_vel.py | 3 +- .../g1/effectors/high_level/dds_sdk.py | 2 +- .../unitree/g1/effectors/high_level/webrtc.py | 2 +- dimos/robot/unitree/g1/g1_tf_publisher.py | 3 +- dimos/robot/unitree/g1/mujoco_sim.py | 96 ++- dimos/robot/unitree/g1/skill_container.py | 2 +- .../g1/test_g1_manipulation_integration.py | 23 +- dimos/robot/unitree/g1/test_g1_teleop.py | 23 +- .../robot/unitree/g1/test_g1_tf_publisher.py | 6 +- .../unitree/g1/test_wholebody_connection.py | 29 +- .../robot/unitree/g1/wholebody_connection.py | 17 +- dimos/robot/unitree/go2/base.py | 2 +- .../unitree/go2/blueprints/basic/go2_sim.py | 3 +- .../basic/test_unitree_go2_joystick_record.py | 23 +- .../go2/blueprints/basic/unitree_go2_basic.py | 13 +- .../go2/blueprints/smart/unitree_go2.py | 12 +- .../unitree/go2/blueprints/teleop_webxr.py | 3 +- .../go2/blueprints/test_tf_topology.py | 11 +- dimos/robot/unitree/go2/camera_calibration.py | 5 +- dimos/robot/unitree/go2/connection.py | 31 +- dimos/robot/unitree/go2/dds/cli/render.py | 12 +- .../go2/dds/cli/test_render_generated.py | 23 +- dimos/robot/unitree/go2/dds/extrinsics.py | 5 +- dimos/robot/unitree/go2/dds/ros.py | 11 +- dimos/robot/unitree/go2/dds/rust/Cargo.toml | 2 +- .../robot/unitree/go2/dds/test_blueprints.py | 16 +- .../unitree/go2/dds/test_generated_codecs.py | 65 +- dimos/robot/unitree/go2/test_connection.py | 107 ++- .../unitree/go2/test_replay_benchmark.py | 5 +- dimos/robot/unitree/keyboard_teleop.py | 63 +- dimos/robot/unitree/mujoco_connection.py | 3 +- dimos/robot/unitree/test_connection.py | 11 +- .../robot/unitree/test_generated_commands.py | 13 +- dimos/robot/unitree/test_mujoco_connection.py | 27 +- dimos/robot/unitree/test_simulation_tf.py | 34 +- dimos/robot/unitree/type/odometry.py | 3 +- dimos/robot/unitree/type/test_lidar.py | 15 +- .../unitree/type/test_pose_conversion.py | 3 +- .../robot/unitree/unitree_skill_container.py | 2 +- dimos/simulation/dimsim/cdr_schemas.py | 5 +- dimos/simulation/engines/mujoco_engine.py | 6 +- dimos/simulation/engines/mujoco_sim_module.py | 257 ++++--- .../engines/test_mujoco_sim_module.py | 10 +- dimos/simulation/go2_sim/test_world.py | 8 +- dimos/simulation/go2_sim/world.py | 115 ++- dimos/simulation/habitat/connection.py | 12 +- dimos/simulation/habitat/nix/install.sh | 2 +- dimos/simulation/habitat/server.py | 223 +++--- dimos/simulation/habitat/test_server.py | 164 ++++- .../mujoco/direct_cmd_vel_explorer.py | 8 +- dimos/simulation/mujoco/model.py | 8 +- dimos/simulation/mujoco/person_on_track.py | 2 +- .../test_generated_control_consumers.py | 22 +- dimos/simulation/scenes/procedural.py | 2 +- dimos/teleop/hosted/arm_command.py | 7 +- dimos/teleop/hosted/blueprints/cloudflare.py | 6 +- dimos/teleop/hosted/camera_mux.py | 11 +- dimos/teleop/hosted/go2_command.py | 6 +- dimos/teleop/hosted/hosted_stats.py | 3 +- dimos/teleop/hosted/test_arm_command.py | 25 +- dimos/teleop/hosted/test_camera_mux.py | 8 +- dimos/teleop/hosted/test_go2_command.py | 12 +- dimos/teleop/hosted/test_hosted_stats.py | 14 +- dimos/teleop/hosted/test_map_compress.py | 61 +- .../keyboard/test_keyboard_teleop_module.py | 9 +- dimos/teleop/phone/phone_teleop_module.py | 15 +- dimos/teleop/phone/test_phone_cdr.py | 57 +- dimos/teleop/utils/test_report.py | 53 +- dimos/teleop/utils/test_video_stats.py | 14 +- dimos/teleop/webxr/extensions.py | 10 +- dimos/teleop/webxr/module.py | 120 +++- dimos/teleop/webxr/test_module.py | 166 +++-- dimos/utils/test_transform_utils.py | 28 +- dimos/utils/test_twist.py | 8 +- dimos/utils/testing/test_moment.py | 7 +- dimos/utils/testing/test_sqlite_moment.py | 3 +- dimos/visualization/rerun/bridge.py | 50 +- .../visualization/rerun/test_bridge_topics.py | 32 +- .../rerun/test_detection3d_bridge.py | 42 +- .../rerun/test_generated_camera_bridge.py | 45 +- .../rerun/test_generated_cloud_bridge.py | 35 +- .../rerun/test_generated_entity_bridge.py | 12 +- .../rerun/test_generated_navigation_bridge.py | 48 +- .../rerun/test_generated_occupancy_bridge.py | 52 +- .../rerun/test_generated_tf_bridge.py | 19 +- .../rerun/test_websocket_server.py | 13 +- dimos/visualization/rerun/websocket_server.py | 4 +- dimos/web/cdr_codec.py | 12 +- dimos/web/cockpit.py | 26 +- dimos/web/codecs.py | 4 +- dimos/web/relay_bridge/builtin_codecs.py | 45 +- dimos/web/relay_bridge/gen_cdr_fixtures.py | 59 +- .../web/relay_bridge/gen_costmap_fixtures.py | 16 +- dimos/web/relay_bridge/gen_voxel_fixtures.py | 10 +- dimos/web/relay_bridge/relay_bridge_module.py | 7 +- dimos/web/relay_bridge/rtc_publisher.py | 2 +- dimos/web/relay_bridge/teleop.py | 10 +- .../web/relay_bridge/test_costmap_encoding.py | 63 +- dimos/web/relay_bridge/test_map_codecs.py | 22 +- .../test_relay_bridge_authoring.py | 38 +- .../web/relay_bridge/test_relay_bridge_e2e.py | 105 ++- .../relay_bridge/test_relay_bridge_module.py | 187 +++-- .../web/relay_bridge/test_relay_bridge_rtc.py | 8 +- dimos/web/relay_bridge/test_rtc_publisher.py | 6 +- dimos/web/relay_bridge/test_voxel_encoding.py | 9 +- dimos/web/test_cdr_codec.py | 26 +- dimos/web/test_cockpit.py | 145 ++-- dimos/web/test_codecs.py | 4 +- .../test_websocket_vis_module.py | 42 +- .../web/websocket_vis/websocket_vis_module.py | 80 ++- docs/development/message-source-packages.md | 2 +- docs/development/messages-in-repository.md | 4 +- docs/usage/lcm.md | 2 +- docs/usage/transports/index.md | 96 ++- docs/web/bridge.md | 16 +- docs/web/cockpit.md | 4 +- docs/web/protocol.md | 4 +- examples/native-modules/cpp/src/cdr_relay.cpp | 4 +- examples/native-modules/cpp/src/ping.cpp | 3 +- examples/native-modules/cpp/src/pong.cpp | 3 +- examples/native-modules/rust/Cargo.toml | 2 +- examples/native-modules/rust/src/cdr_relay.rs | 4 +- examples/native-modules/rust/src/ping.rs | 2 +- examples/native-modules/rust/src/pong.rs | 2 +- examples/native-modules/rust_ping_pong.py | 3 +- examples/native-modules/rust_tf.py | 63 +- .../optimize_candidates_child.py | 5 +- .../optimize_patrol_router_child.py | 5 +- native/cpp/dimos_nativeConfig.cmake | 7 + native/cpp/include/dimos/native/module.hpp | 4 +- native/cpp/messages.nix | 83 ++- native/cpp/source-archives.json | 38 + native/cpp/tests/CMakeLists.txt | 4 +- native/cpp/tests/test_cdr_codec.cpp | 5 +- native/cpp/tests/test_module.cpp | 1 + native/message-source.nix | 53 ++ native/rust/dimos-lcm-transport/Cargo.toml | 11 + native/rust/dimos-lcm-transport/README.md | 37 + .../dimos-lcm-transport/examples/interop.rs | 46 ++ native/rust/dimos-lcm-transport/src/lib.rs | 41 ++ .../rust/dimos-lcm-transport/src/transport.rs | 661 ++++++++++++++++++ native/rust/dimos-module-macros/src/lib.rs | 8 +- native/rust/dimos-module/Cargo.toml | 3 +- native/rust/dimos-module/src/cdr.rs | 82 ++- native/rust/dimos-module/src/pointcloud.rs | 307 +++++--- native/rust/dimos-module/src/tf.rs | 87 +-- native/rust/dimos-module/src/time.rs | 4 +- packages/dimos-generated/catalog.json | 4 +- .../custom/dimos_msgs/msg/ControlValues.msg | 13 + .../custom/dimos_msgs/msg/TriangleMesh.msg | 6 + scripts/install_native_messages.sh | 31 +- scripts/package_native_sdk.py | 5 +- scripts/test_external_native_messages.sh | 2 +- scripts/test_message_user_story.sh | 8 +- web/cockpit/src/subscriptions.ts | 11 +- web/deno.lock | 31 + web/sdk/src/decoders/decoders.test.ts | 26 +- web/sdk/src/decoders/index.ts | 20 +- web/sdk/src/index.ts | 9 +- 707 files changed, 21982 insertions(+), 7937 deletions(-) create mode 100644 dimos/agents/skills/visual_servoing/test_visual_servoing_2d.py create mode 100644 dimos/msgs/contacts.py create mode 100644 dimos/msgs/control.py delete mode 100644 dimos/msgs/control_msgs/test_ControlValues.py create mode 100644 dimos/msgs/mesh.py delete mode 100644 dimos/msgs/sim_msgs/Contacts.py create mode 100644 dimos/msgs/test_control.py create mode 100644 dimos/msgs/test_mesh.py create mode 100644 dimos/msgs/test_regions.py create mode 100644 dimos/navigation/experimental/frontier_exploration/test_generated_frontiers.py create mode 100644 dimos/navigation/experimental/patrolling/test_generated_patrol.py create mode 100644 dimos/navigation/go2/loop_closure/test_pgo_auto.py create mode 100644 dimos/navigation/go2/loop_closure/test_pgo_generated.py create mode 100644 dimos/navigation/go2/replanning_a_star/test_generated_navigation.py create mode 100644 dimos/navigation/go2/replanning_a_star/test_generated_pipeline.py create mode 100644 dimos/protocol/pubsub/impl/jpeg_shm.py create mode 100644 dimos/protocol/pubsub/impl/shmpubsub.py create mode 100644 dimos/protocol/pubsub/impl/test_shmpubsub.py create mode 100644 dimos/protocol/pubsub/shm/ipc_factory.py create mode 100644 dimos/protocol/pubsub/shm/test_ipc_factory.py create mode 100644 native/cpp/dimos_nativeConfig.cmake create mode 100644 native/cpp/source-archives.json create mode 100644 native/message-source.nix create mode 100644 native/rust/dimos-lcm-transport/Cargo.toml create mode 100644 native/rust/dimos-lcm-transport/README.md create mode 100644 native/rust/dimos-lcm-transport/examples/interop.rs create mode 100644 native/rust/dimos-lcm-transport/src/lib.rs create mode 100644 native/rust/dimos-lcm-transport/src/transport.rs create mode 100644 packages/dimos-generated/custom/dimos_msgs/msg/ControlValues.msg create mode 100644 packages/dimos-generated/custom/dimos_msgs/msg/TriangleMesh.msg diff --git a/Cargo.toml b/Cargo.toml index ac8f8656ccf..c23d93a945b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,7 @@ [workspace] resolver = "2" exclude = [ + "packages/dimos-generated/build/message-source/rust", "examples/message-project/sdk/rust/dimos-module", "examples/message-project/sdk/rust/dimos-module-macros", ] diff --git a/dimos/agents/mcp/test_generated_image.py b/dimos/agents/mcp/test_generated_image.py index 3fcc5442df5..457002e0d91 100644 --- a/dimos/agents/mcp/test_generated_image.py +++ b/dimos/agents/mcp/test_generated_image.py @@ -18,6 +18,7 @@ import base64 import cv2 +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.sensor_msgs.msg import Image import numpy as np @@ -34,7 +35,7 @@ def test_generated_image_mcp_result_reaches_model_history(mocker): handle_request( {"method": "tools/call", "id": 1, "params": {"name": "observe"}}, [], - {"observe": lambda: Image.decode(original.encode())}, + {"observe": lambda: cdr_decode(cdr_encode(original), Image)}, ) ) assert response is not None diff --git a/dimos/agents/skills/navigation.py b/dimos/agents/skills/navigation.py index ef379788be3..b3a16795253 100644 --- a/dimos/agents/skills/navigation.py +++ b/dimos/agents/skills/navigation.py @@ -289,6 +289,7 @@ def _get_goal_pose_from_result(self, result: dict[str, Any]) -> PoseStamped | No return PoseStamped( header=header_now("map"), pose=Pose( - position=Point(x=pos_x, y=pos_y), orientation=quaternion_from_euler(0, 0, theta) + position=Point(x=pos_x, y=pos_y, z=0.0), + orientation=quaternion_from_euler(0, 0, theta), ), ) diff --git a/dimos/agents/skills/person_follow.py b/dimos/agents/skills/person_follow.py index b1ccefca1d3..4639f86c185 100644 --- a/dimos/agents/skills/person_follow.py +++ b/dimos/agents/skills/person_follow.py @@ -17,6 +17,9 @@ import time from typing import TYPE_CHECKING, Any +from dimos_generated.geometry_msgs.msg import Twist, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 +from dimos_generated.tf2_msgs.msg import TFMessage import numpy as np from reactivex.disposable import Disposable @@ -32,11 +35,7 @@ from dimos.models.qwen.bbox import BBox from dimos.models.vl.base import VlModel from dimos.models.vl.create import create -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.image import image_from_array from dimos.utils.logging_config import setup_logger from dimos.utils.turbojpeg import get_turbojpeg @@ -206,7 +205,9 @@ def stop_following(self) -> str: """ self._stop_following() - self.cmd_vel.publish(Twist.zero()) + self.cmd_vel.publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) if self._thread is not None: self._thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) @@ -249,7 +250,9 @@ def _follow_person( ) if len(initial_detections) == 0: - self.cmd_vel.publish(Twist.zero()) + self.cmd_vel.publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) return f"EdgeTAM failed to segment '{query}'." logger.info(f"EdgeTAM initialized with {len(initial_detections)} detections") @@ -281,7 +284,9 @@ def _follow_loop(self, tracker: "EdgeTAMProcessor", query: str) -> None: detections = tracker.process_image(latest_image) if len(detections) == 0: - self.cmd_vel.publish(Twist.zero()) + self.cmd_vel.publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) lost_count += 1 if lost_count > self._max_lost_frames: @@ -323,7 +328,9 @@ def _stop_following(self) -> None: self._should_stop.set() def _send_stop_reason(self, query: str, reason: str) -> None: - self.cmd_vel.publish(Twist.zero()) + self.cmd_vel.publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) self.tool_update( "follow_person", f"Person follow stopped for '{query}'. Reason: {reason}.", @@ -334,4 +341,4 @@ def _send_stop_reason(self, query: str, reason: str) -> None: def _decode_base64_image(b64: str) -> Image: bgr_array = get_turbojpeg().decode(base64.b64decode(b64)) - return Image(data=bgr_array, format=ImageFormat.BGR) + return image_from_array(bgr_array, encoding="bgr8") diff --git a/dimos/agents/skills/test_generated_navigation.py b/dimos/agents/skills/test_generated_navigation.py index 49551742803..52503d7eba6 100644 --- a/dimos/agents/skills/test_generated_navigation.py +++ b/dimos/agents/skills/test_generated_navigation.py @@ -16,7 +16,10 @@ import math +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.std_msgs.msg import Header import pytest from dimos.agents.skills.navigation import NavigationSkillContainer @@ -41,9 +44,10 @@ def test_tagged_location_preserves_euler_orientation_through_cdr_goal(skill): value, memory, navigator = skill angles = (0.2, -0.3, math.pi / 2) pose = PoseStamped( - pose=Pose(position=Point(x=1, y=2, z=3), orientation=quaternion_from_euler(*angles)) + pose=Pose(position=Point(x=1, y=2, z=3), orientation=quaternion_from_euler(*angles)), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) - value._on_odom(PoseStamped.decode(pose.encode())) + value._on_odom(cdr_decode(cdr_encode(pose), PoseStamped)) assert "Tagged 'desk'" in value.tag_location("desk") tagged = memory.tag_location.call_args.args[0] assert tagged.position == (1, 2, 3) @@ -51,7 +55,7 @@ def test_tagged_location_preserves_euler_orientation_through_cdr_goal(skill): memory.query_tagged_location.return_value = tagged assert "Found a tagged location" in value._navigate_by_tagged_location("desk") goal = navigator.set_goal.call_args.args[0] - decoded = PoseStamped.decode(goal.encode()) + decoded = cdr_decode(cdr_encode(goal), PoseStamped) assert decoded.header.frame_id == "map" assert decoded.pose.position == pose.pose.position assert quaternion_euler(decoded.pose.orientation) == pytest.approx(angles) @@ -63,8 +67,8 @@ def test_semantic_result_builds_generated_map_goal(skill): goal = value._get_goal_pose_from_result(result) assert goal is not None value._navigate_to(goal, "Found desk") - decoded = PoseStamped.decode(navigator.set_goal.call_args.args[0].encode()) + decoded = cdr_decode(cdr_encode(navigator.set_goal.call_args.args[0]), PoseStamped) assert decoded.header.frame_id == "map" - assert decoded.pose.position == Point(x=4, y=-2) + assert decoded.pose.position == Point(x=4, y=-2, z=0.0) assert quaternion_euler(decoded.pose.orientation) == pytest.approx((0, 0, math.pi / 2)) assert value._get_goal_pose_from_result({"distance": 0.9}) is None diff --git a/dimos/agents/skills/test_generated_person_follow.py b/dimos/agents/skills/test_generated_person_follow.py index 8122ba5ea19..67b065c37cc 100644 --- a/dimos/agents/skills/test_generated_person_follow.py +++ b/dimos/agents/skills/test_generated_person_follow.py @@ -15,8 +15,11 @@ from types import SimpleNamespace from unittest.mock import MagicMock -from dimos_generated.geometry_msgs.msg import Twist -from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Twist, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, RegionOfInterest +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest @@ -30,12 +33,24 @@ @pytest.mark.parametrize("lost", [False, True]) def test_generated_person_follow_loop_stops_cleanly(monkeypatch, lost): monkeypatch.setattr(person_follow, "create", lambda _: MagicMock()) - camera = CameraInfo(width=640, height=480, k=[500, 0, 320, 0, 500, 240, 0, 0, 1]) + camera = CameraInfo( + width=640, + height=480, + k=np.array([500, 0, 320, 0, 500, 240, 0, 0, 1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + distortion_model="", + d=np.array([], dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), + ) module = person_follow.PersonFollowSkillContainer( camera_info=camera, g=GlobalConfig(simulation="") ) image = image_from_array(np.zeros((480, 640, 3), dtype=np.uint8), encoding="rgb8") - image = Image.decode(image.encode()) + image = cdr_decode(cdr_encode(image), Image) detection = Detection2DBBox( bbox=(270, 100, 370, 300), track_id=1, @@ -53,7 +68,7 @@ def test_generated_person_follow_loop_stops_cleanly(monkeypatch, lost): outputs = [] def receive(value): - outputs.append(Twist.decode(value.encode())) + outputs.append(cdr_decode(cdr_encode(value), Twist)) if not lost: module._should_stop.set() @@ -70,9 +85,15 @@ def receive(value): thread.join(timeout=2) assert not thread.is_alive() assert outputs - assert outputs[-1] == Twist() + assert outputs[-1] == Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) if lost: - assert all(value == Twist() for value in outputs) + assert all( + value + == Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + for value in outputs + ) else: assert outputs[0].linear.x == 0.5 assert outputs[0].angular.z == 0 diff --git a/dimos/agents/skills/test_unitree_skill_container.py b/dimos/agents/skills/test_unitree_skill_container.py index 278c71eb370..00d3003a9f0 100644 --- a/dimos/agents/skills/test_unitree_skill_container.py +++ b/dimos/agents/skills/test_unitree_skill_container.py @@ -16,7 +16,9 @@ import math from typing import Any +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.std_msgs.msg import Header from langchain_core.messages import HumanMessage import pytest @@ -85,6 +87,7 @@ def _pose(x: float, y: float, yaw_deg: float) -> PoseStamped: position=Point(x=x, y=y, z=0.3), orientation=quaternion_from_euler(0, 0, math.radians(yaw_deg)), ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) diff --git a/dimos/agents/skills/visual_servoing/bbox_navigation.py b/dimos/agents/skills/visual_servoing/bbox_navigation.py index de56e99ba18..166c681a79b 100644 --- a/dimos/agents/skills/visual_servoing/bbox_navigation.py +++ b/dimos/agents/skills/visual_servoing/bbox_navigation.py @@ -70,7 +70,9 @@ def _on_detection(self, det: Detection2DArray) -> None: ) goal = PoseStamped( header=det.header, - pose=Pose(position=Point(x=z, y=-x, z=-y), orientation=Quaternion(w=1)), + pose=Pose( + position=Point(x=z, y=-x, z=-y), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), ) logger.debug( f"BBox center: ({center_x:.1f}, {center_y:.1f}) → " diff --git a/dimos/agents/skills/visual_servoing/detection_navigation.py b/dimos/agents/skills/visual_servoing/detection_navigation.py index 16b66c692c8..9b99911534c 100644 --- a/dimos/agents/skills/visual_servoing/detection_navigation.py +++ b/dimos/agents/skills/visual_servoing/detection_navigation.py @@ -205,6 +205,6 @@ def _compute_twist_from_3d( ) return Twist( - linear=Vector3(x=linear_x), - angular=Vector3(z=angular_z), + linear=Vector3(x=linear_x, y=0.0, z=0.0), + angular=Vector3(z=angular_z, x=0.0, y=0.0), ) diff --git a/dimos/agents/skills/visual_servoing/test_visual_servoing_2d.py b/dimos/agents/skills/visual_servoing/test_visual_servoing_2d.py new file mode 100644 index 00000000000..0d81b35e575 --- /dev/null +++ b/dimos/agents/skills/visual_servoing/test_visual_servoing_2d.py @@ -0,0 +1,56 @@ +# Copyright 2025-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. + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Twist +from dimos_generated.sensor_msgs.msg import CameraInfo, RegionOfInterest +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest + +from dimos.agents.skills.visual_servoing.visual_servoing_2d import VisualServoing2D + + +@pytest.mark.parametrize( + "bbox, linear, angular", + [ + ((270, 0, 370, 100), 0.5, 0), + ((245, 0, 395, 100), 0, 0), + ((120, 0, 520, 100), -0.3, 0), + ((370, 0, 470, 100), 0.36, -0.2), + ((320, 0, 320, 100), 0, 0), + ], +) +def test_generated_camera_to_twist(bbox, linear, angular): + camera = CameraInfo( + width=640, + height=480, + k=np.array([500, 0, 320, 0, 500, 240, 0, 0, 1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + distortion_model="", + d=np.array([], dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), + ) + controller = VisualServoing2D(cdr_decode(cdr_encode(camera), CameraInfo)) + output = controller.compute_twist(bbox, 640) + decoded = cdr_decode(cdr_encode(output), Twist) + assert decoded.linear.x == pytest.approx(linear) + assert decoded.angular.z == pytest.approx(angular) + assert decoded.linear.y == decoded.linear.z == 0 + assert decoded.angular.x == decoded.angular.y == 0 diff --git a/dimos/agents/skills/visual_servoing/visual_servoing_2d.py b/dimos/agents/skills/visual_servoing/visual_servoing_2d.py index fe209e7f95b..005cf4a3182 100644 --- a/dimos/agents/skills/visual_servoing/visual_servoing_2d.py +++ b/dimos/agents/skills/visual_servoing/visual_servoing_2d.py @@ -89,7 +89,7 @@ def compute_twist( estimated_distance = self._estimate_distance(bbox) if estimated_distance is None: - return Twist() + return Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) # Calculate distance error (positive = too far, need to move forward) distance_error = estimated_distance - self._target_distance @@ -118,8 +118,8 @@ def compute_twist( linear_x = max(linear_x, 0.1) return Twist( - linear=Vector3(x=linear_x), - angular=Vector3(z=angular_z), + linear=Vector3(x=linear_x, y=0.0, z=0.0), + angular=Vector3(z=angular_z, x=0.0, y=0.0), ) def _get_normalized_x(self, pixel_x: float) -> float: @@ -135,7 +135,7 @@ def _get_normalized_x(self, pixel_x: float) -> float: """ fx = self._camera_info.k[0] # focal length x cx = self._camera_info.k[2] # optical center x - return (pixel_x - cx) / fx + return float((pixel_x - cx) / fx) def _estimate_distance(self, bbox: tuple[float, float, float, float]) -> float | None: """Estimate distance to object based on bounding box size and camera intrinsics. @@ -162,4 +162,4 @@ def _estimate_distance(self, bbox: tuple[float, float, float, float]) -> float | fx = self._camera_info.k[0] # focal length x in pixels estimated_distance = (self._assumed_object_width * fx) / bbox_width - return estimated_distance + return float(estimated_distance) diff --git a/dimos/agents/test_generated_vlm.py b/dimos/agents/test_generated_vlm.py index 2a456afca18..cac403becfb 100644 --- a/dimos/agents/test_generated_vlm.py +++ b/dimos/agents/test_generated_vlm.py @@ -17,6 +17,7 @@ import base64 import cv2 +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.sensor_msgs.msg import Image from langchain_core.messages import AIMessage, HumanMessage import numpy as np @@ -44,7 +45,7 @@ def test_generated_image_reaches_model_as_jpeg(agent, use_stream): pixels = np.zeros((12, 16, 3), dtype=np.uint8) pixels[..., 0] = 255 original = image_from_array(pixels, encoding="rgb8") - image = Image.decode(original.encode()) + image = cdr_decode(cdr_encode(original), Image) answers = [] unsubscribe = value.answer_stream.subscribe(answers.append) try: @@ -63,7 +64,7 @@ def test_generated_image_reaches_model_as_jpeg(agent, use_stream): ) assert decoded.shape == (12, 16, 3) assert decoded[0, 0].tolist() == [0, 0, 254] - assert image.encode() == original.encode() + assert cdr_encode(image) == cdr_encode(original) finally: unsubscribe() diff --git a/dimos/agents/typesafe/navigation.py b/dimos/agents/typesafe/navigation.py index a3fd9d242be..2d8ac35715a 100644 --- a/dimos/agents/typesafe/navigation.py +++ b/dimos/agents/typesafe/navigation.py @@ -22,6 +22,11 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import Point, PointStamped, PoseStamped, Twist, Vector3 +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.vision_msgs.msg import Detection2DArray, Detection3DArray from langchain_core.messages import AIMessage, HumanMessage from langchain_core.messages.base import BaseMessage from reactivex import Observable, operators as ops @@ -53,14 +58,7 @@ from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.core import rpc from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PointStamped import PointStamped -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray -from dimos.types.timestamped import Timestamped +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.utils.logging_config import setup_logger from dimos.utils.reactive import LatestReader, getter_streaming @@ -125,17 +123,15 @@ def trigger(self) -> Observable[object]: @staticmethod def _pose_of(o: Odometry) -> PoseStamped: - return PoseStamped( - ts=o.ts, frame_id=o.frame_id, position=o.position, orientation=o.orientation - ) + return PoseStamped(header=o.header, pose=o.pose.pose) def _fresh(self, name: str) -> Any: """Latest message on `name`, or None when absent or older than `stale_s`.""" try: - msg: Timestamped = self._latest[name]() + msg = self._latest[name]() except TimeoutError: return None - return msg if time.time() - msg.ts <= self.config.stale_s else None + return msg if time.time() - to_seconds(msg.header.stamp) <= self.config.stale_s else None def state(self, trigger: object) -> WorldState | None: with self._lock: @@ -207,7 +203,9 @@ def stop(self) -> None: self._publisher.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) for reader in self._latest.values(): reader.dispose() - self.cmd_vel.publish(Twist.zero()) + self.cmd_vel.publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) super().stop() @rpc @@ -229,7 +227,9 @@ def _change_goal(self, goal: str | None) -> None: self._zero_since = None self._motion = "idle" if goal is None: - self.cmd_vel.publish(Twist.zero()) + self.cmd_vel.publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) self.agent_idle.publish(goal is None) def _gave_up(self, gen: int) -> bool: @@ -299,7 +299,12 @@ def _resolve_goal(self, x: float, y: float, gen: int) -> None: prev = self._goal_xy self._goal_xy = (x, y) if prev is None or abs(prev[0] - x) > 0.1 or abs(prev[1] - y) > 0.1: - self.goal.publish(PointStamped(x, y, 0.0, frame_id="world")) + self.goal.publish( + PointStamped( + header=Header(stamp=time_from_seconds(time.time()), frame_id="world"), + point=Point(x=x, y=y, z=0.0), + ) + ) def _set_target(self, target: Vec3, *, immediate: bool = False, gen: int | None = None) -> None: now = time.monotonic() @@ -310,7 +315,9 @@ def _set_target(self, target: Vec3, *, immediate: bool = False, gen: int | None self._decided_at = now if immediate: self._current = ZERO - self.cmd_vel.publish(Twist.zero()) + self.cmd_vel.publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) self._zero_since = (self._zero_since or now) if target == ZERO else None def _publish_loop(self) -> None: @@ -334,6 +341,9 @@ def _publish_loop(self) -> None: self._current = (nxt[0], nxt[1], nxt[2]) if self._goal is not None or any(cur): self.cmd_vel.publish( - Twist(linear=(nxt[0], nxt[1], 0.0), angular=(0.0, 0.0, nxt[2])) + Twist( + linear=Vector3(x=nxt[0], y=nxt[1], z=0.0), + angular=Vector3(x=0.0, y=0.0, z=nxt[2]), + ) ) self._stop_event.wait(dt) diff --git a/dimos/agents/typesafe/test_navigation.py b/dimos/agents/typesafe/test_navigation.py index 85abc252728..5b156acf7bd 100644 --- a/dimos/agents/typesafe/test_navigation.py +++ b/dimos/agents/typesafe/test_navigation.py @@ -15,19 +15,26 @@ import threading import time +from dimos_generated.geometry_msgs.msg import ( + PointStamped, + PoseStamped, + PoseWithCovariance, + Twist, + TwistWithCovariance, + Vector3, +) +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from reactivex.scheduler import ThreadPoolScheduler from dimos.agents.typesafe.navigation import TypeSafeNavigationAgent from dimos.agents.typesafe.test_drive import answers -from dimos.agents.typesafe.test_world_state import det2d, det3d +from dimos.agents.typesafe.test_world_state import det2d, det3d, pose from dimos.agents.typesafe.types import Answers, Question from dimos.core.transport import LCMTransport, pLCMTransport -from dimos.msgs.geometry_msgs.PointStamped import PointStamped -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.msgs.time import time_from_seconds class FakeSystemOne: @@ -84,7 +91,7 @@ def rig(monkeypatch: pytest.MonkeyPatch) -> Iterator[Rig]: def odom(a: TypeSafeNavigationAgent, x: float = 0.0) -> None: - a.odom.transport.publish(PoseStamped(position=(x, 0, 0.4), frame_id="world")) + a.odom.transport.publish(pose(x, 0, 0)) def scene(a: TypeSafeNavigationAgent, robot_x: float = 0.0) -> None: @@ -132,7 +139,9 @@ def test_drives_then_stops_and_clears(rig: Rig) -> None: assert until(lambda: moving(twists)) fake.answers = answers(stop=0.95) assert until(lambda: (odom(a), a.current_goal() is None)[1]) - assert twists[-1].is_zero() + assert twists[-1] == Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) def test_arrival_by_distance(rig: Rig) -> None: @@ -151,12 +160,33 @@ def test_odometry_feeds_pose(rig: Rig) -> None: a.detections_3d.transport.publish(det3d("chair", 3.0, 0.0)) time.sleep(0.05) a.odometry.transport.publish( - Odometry(ts=time.time(), frame_id="world", pose=Pose(position=(2.8, 0, 0.4))) + Odometry( + child_frame_id="", + twist=TwistWithCovariance( + twist=Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ), + covariance=np.zeros(36), + ), + header=Header(stamp=time_from_seconds(time.time()), frame_id="world"), + pose=PoseWithCovariance(pose=pose(2.8, 0, 0).pose, covariance=np.zeros(36)), + ) ) assert until( lambda: ( a.odometry.transport.publish( - Odometry(ts=time.time(), frame_id="world", pose=Pose(position=(2.8, 0, 0.4))) + Odometry( + child_frame_id="", + twist=TwistWithCovariance( + twist=Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), + angular=Vector3(x=0.0, y=0.0, z=0.0), + ), + covariance=np.zeros(36), + ), + header=Header(stamp=time_from_seconds(time.time()), frame_id="world"), + pose=PoseWithCovariance(pose=pose(2.8, 0, 0).pose, covariance=np.zeros(36)), + ) ), a.current_goal() is None, )[1] @@ -172,7 +202,11 @@ def test_inference_failure_zeroes_target(rig: Rig) -> None: assert until(lambda: moving(twists)) fake.fail = True odom(a) - assert until(lambda: twists[-1].is_zero(), timeout=3.0) + assert until( + lambda: twists[-1] + == Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + timeout=3.0, + ) def test_deadman_zeroes_when_inference_stalls(rig: Rig) -> None: @@ -184,7 +218,11 @@ def test_deadman_zeroes_when_inference_stalls(rig: Rig) -> None: assert until(lambda: moving(twists)) fake.gate.clear() odom(a) - assert until(lambda: twists[-1].is_zero(), timeout=3.0) + assert until( + lambda: twists[-1] + == Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + timeout=3.0, + ) fake.gate.set() @@ -196,7 +234,7 @@ def test_goal_point_published_and_used_when_target_leaves_view(rig: Rig) -> None a.set_goal("go to the chair") scene(a) assert until(lambda: moving(twists) and bool(goals)) - assert (goals[0].x, goals[0].y, goals[0].frame_id) == (3.0, 0.0, "world") + assert (goals[0].point.x, goals[0].point.y, goals[0].header.frame_id) == (3.0, 0.0, "world") # detections go stale, odom keeps coming: the latched goal point keeps it driving a.config.stale_s = 0.2 time.sleep(0.3) @@ -235,7 +273,9 @@ def test_lost_detections_give_up(rig: Rig) -> None: assert until(lambda: moving(twists)) a.config.stale_s = 0.0 # every detection is now stale assert until(lambda: (odom(a), a.current_goal() is None)[1]) - assert twists[-1].is_zero() + assert twists[-1] == Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) def test_timeout_clears_only_the_goal_that_timed_out(rig: Rig) -> None: diff --git a/dimos/agents/typesafe/test_world_state.py b/dimos/agents/typesafe/test_world_state.py index 6fe7fc8e053..62fea7d2e00 100644 --- a/dimos/agents/typesafe/test_world_state.py +++ b/dimos/agents/typesafe/test_world_state.py @@ -14,54 +14,85 @@ import math import time -from dimos_lcm.vision_msgs import ( +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + PoseWithCovariance, + Quaternion, + Vector3, +) +from dimos_generated.std_msgs.msg import Header +from dimos_generated.vision_msgs.msg import ( BoundingBox2D, BoundingBox3D, Detection2D, + Detection2DArray, Detection3D, + Detection3DArray, ObjectHypothesis, ObjectHypothesisWithPose, + Point2D, Pose2D, ) import numpy as np -from dimos.agents.typesafe.world_state import bearing_word, build_world_state, distance_word -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.std_msgs.Header import Header -from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.agents.typesafe.world_state import ( + bearing_word, + build_world_state, + distance_word, + pose_json, +) +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import time_from_seconds def pose(x: float, y: float, yaw_deg: float) -> PoseStamped: - p = PoseStamped(position=(x, y, 0.4), frame_id="world") - p.orientation = Quaternion.from_euler(Vector3(0, 0, math.radians(yaw_deg))) + p = PoseStamped( + header=Header(stamp=time_from_seconds(time.time()), frame_id="world"), + pose=Pose( + position=Point(x=x, y=y, z=0.4), + orientation=quaternion_from_euler(0, 0, math.radians(yaw_deg)), + ), + ) return p def det3d(label: str, x: float, y: float, ts: float | None = None) -> Detection3DArray: ts = time.time() if ts is None else ts - d = Detection3D() - d.header = Header(ts, "world") - d.results = [ObjectHypothesisWithPose(hypothesis=ObjectHypothesis(class_id=label, score=0.9))] - d.results_length = 1 - d.bbox = BoundingBox3D(center=Pose(position=(x, y, 0.3)), size=Vector3(0.5, 0.5, 0.6)) - return Detection3DArray(detections_length=1, header=Header(ts, "world"), detections=[d]) + header = Header(stamp=time_from_seconds(ts), frame_id="world") + center = Pose( + position=Point(x=x, y=y, z=0.3), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) + result = ObjectHypothesisWithPose( + hypothesis=ObjectHypothesis(class_id=label, score=0.9), + pose=PoseWithCovariance(pose=center, covariance=np.zeros(36)), + ) + d = Detection3D( + header=header, + results=[result], + bbox=BoundingBox3D(center=center, size=Vector3(x=0.5, y=0.5, z=0.6)), + id="", + ) + return Detection3DArray(header=header, detections=[d]) def det2d(label: str, cx: float, w: float, h: float, ts: float | None = None) -> Detection2DArray: ts = time.time() if ts is None else ts - d = Detection2D() - d.header = Header(ts, "camera") - d.results = [ObjectHypothesisWithPose(hypothesis=ObjectHypothesis(class_id=label, score=0.8))] - d.results_length = 1 - center = Pose2D() - center.position.x, center.position.y = cx, 360.0 - d.bbox = BoundingBox2D(center=center, size_x=w, size_y=h) - return Detection2DArray(detections_length=1, header=Header(ts, "camera"), detections=[d]) + header = Header(stamp=time_from_seconds(ts), frame_id="camera") + center = Pose2D(position=Point2D(x=cx, y=360.0), theta=0.0) + result = ObjectHypothesisWithPose( + hypothesis=ObjectHypothesis(class_id=label, score=0.8), + pose=PoseWithCovariance(pose=pose(0, 0, 0).pose, covariance=np.zeros(36)), + ) + d = Detection2D( + header=header, + results=[result], + bbox=BoundingBox2D(center=center, size_x=w, size_y=h), + id="", + ) + return Detection2DArray(header=header, detections=[d]) def _state(**kw): # type: ignore[no-untyped-def] @@ -78,12 +109,12 @@ def _state(**kw): # type: ignore[no-untyped-def] def test_pose_to_json() -> None: - assert pose(1, 2, 90).to_json() == { + assert pose_json(pose(1, 2, 90)) == { "position": {"x": 1.0, "y": 2.0, "z": 0.4}, "yaw_deg": 90.0, "heading": "north", } - assert pose(0, 0, -45).to_json()["heading"] == "south_east" + assert pose_json(pose(0, 0, -45))["heading"] == "south_east" def test_words() -> None: @@ -118,7 +149,12 @@ def test_room_sectors_in_robot_frame() -> None: pts = np.array( [[1.0, 1.3, 0.5], [3.0, 1.0, 0.5], [1.0, -2.0, 0.5], [1.0, 1.0, 5.0]] ) # ahead 0.3, right 2, behind 3, above - room = _state(pose=p, lidar=PointCloud2.from_numpy(pts, frame_id="world"))["room"] + room = _state( + pose=p, + lidar=pointcloud_from_xyz( + pts, header=Header(stamp=time_from_seconds(0.0), frame_id="world") + ), + )["room"] assert room["ahead"] == {"clear_m": 0.3, "state": "blocked"} assert room["right"] == {"clear_m": 2.0, "state": "clear"} assert (room["behind"]["clear_m"], room["left"]["clear_m"]) == (3.0, 5.0) @@ -136,3 +172,11 @@ def test_goal_point_relative_to_pose() -> None: "distance_m": 2.0, } assert "goal_point" not in _state() + + +def test_objects_3d_preserve_marker_identity() -> None: + detections = det3d(" chair ", 3.0, 0.0) + detections.detections[0].id = " 42 " + assert _state(detections_3d=detections)["objects"][0]["label"] == "chair id=42" + detections.detections[0].results = [] + assert _state(detections_3d=detections)["objects"][0]["label"] == "id=42" diff --git a/dimos/agents/typesafe/world_state.py b/dimos/agents/typesafe/world_state.py index ad3c79f77bc..553c3f678b4 100644 --- a/dimos/agents/typesafe/world_state.py +++ b/dimos/agents/typesafe/world_state.py @@ -18,6 +18,7 @@ import math from typing import TYPE_CHECKING, TypedDict +import numpy as np from typing_extensions import NotRequired from dimos.agents.typesafe.constants import ( @@ -27,15 +28,86 @@ SECTOR_NAMES, SIZE_WORDS_2D, ) -from dimos.msgs.geometry_msgs.PoseStamped import PoseJson, XyzJson -from dimos.msgs.sensor_msgs.PointCloud2 import SectorJson -from dimos.msgs.vision_msgs.Detection2DArray import BBoxJson +from dimos.msgs.geometry import yaw +from dimos.msgs.pointcloud import pointcloud_xyz if TYPE_CHECKING: - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 - from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray - from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray + from dimos_generated.geometry_msgs.msg import PoseStamped + from dimos_generated.sensor_msgs.msg import PointCloud2 + from dimos_generated.vision_msgs.msg import Detection2DArray, Detection3DArray + + +class XyzJson(TypedDict): + x: float + y: float + z: float + + +class PoseJson(TypedDict): + position: XyzJson + yaw_deg: float + heading: str + + +class BBoxJson(TypedDict): + cx: float + cy: float + w: float + h: float + + +class SectorJson(TypedDict): + clear_m: float + state: str + + +def pose_json(pose: PoseStamped) -> PoseJson: + position = pose.pose.position + degrees = math.degrees(yaw(pose.pose.orientation)) + headings = ( + "east", + "north_east", + "north", + "north_west", + "west", + "south_west", + "south", + "south_east", + ) + return { + "position": { + "x": round(position.x, 2), + "y": round(position.y, 2), + "z": round(position.z, 2), + }, + "yaw_deg": round(degrees, 1), + "heading": headings[round(degrees / 45.0) % 8], + } + + +def room_sectors( + cloud: PointCloud2, pose: PoseStamped, z_min: float, z_max: float, max_range: float +) -> dict[str, SectorJson]: + position = pose.pose.position + points = pointcloud_xyz(cloud) - np.array([position.x, position.y, position.z]) + c, s = np.cos(-yaw(pose.pose.orientation)), np.sin(-yaw(pose.pose.orientation)) + points = np.column_stack( + (c * points[:, 0] - s * points[:, 1], s * points[:, 0] + c * points[:, 1], points[:, 2]) + ) + radii = np.hypot(points[:, 0], points[:, 1]) + keep = (points[:, 2] > z_min) & (points[:, 2] < z_max) & (radii > 0.05) & (radii < max_range) + points, radii = points[keep], radii[keep] + bins = np.round( + np.arctan2(points[:, 1], points[:, 0]) / (2 * np.pi / len(SECTOR_NAMES)) + ).astype(int) % len(SECTOR_NAMES) + out: dict[str, SectorJson] = {} + for index, name in enumerate(SECTOR_NAMES): + clear_m = float(radii[bins == index].min()) if (bins == index).any() else max_range + out[name] = { + "clear_m": round(clear_m, 2), + "state": "blocked" if clear_m < 0.5 else "tight" if clear_m < 1.0 else "clear", + } + return out class ObjectState(TypedDict): @@ -84,21 +156,42 @@ def distance_word(d: float) -> str: def relative(pose: PoseStamped, x: float, y: float) -> tuple[str, float, str, float]: """(bearing word, bearing deg, distance word, distance m) of a world XY from the pose.""" - dx, dy = x - pose.x, y - pose.y - dist, rel = math.hypot(dx, dy), math.atan2(dy, dx) - pose.yaw + dx, dy = x - pose.pose.position.x, y - pose.pose.position.y + dist, rel = math.hypot(dx, dy), math.atan2(dy, dx) - yaw(pose.pose.orientation) deg = round(math.degrees(math.atan2(math.sin(rel), math.cos(rel))), 1) return bearing_word(rel), deg, distance_word(dist), round(dist, 2) def _objects_3d(dets: Detection3DArray, pose: PoseStamped) -> list[ObjectState]: out: list[ObjectState] = [] - for d in dets.to_json(): - bearing, deg, distance, dist = relative(pose, d["position"]["x"], d["position"]["y"]) + for detection in dets.detections: + center = detection.bbox.center.position + position: XyzJson = { + "x": round(center.x, 2), + "y": round(center.y, 2), + "z": round(center.z, 2), + } + label = next( + ( + result.hypothesis.class_id + for result in detection.results + if result.hypothesis.class_id + ), + "", + ) + marker_id = detection.id.strip() + label = label.strip() + if marker_id: + label = f"{label} id={marker_id}" if label else f"id={marker_id}" + score = round( + max((result.hypothesis.score for result in detection.results), default=0.0), 2 + ) + bearing, deg, distance, dist = relative(pose, position["x"], position["y"]) out.append( { - "label": d["label"], - "score": d["score"], - "position": d["position"], + "label": label, + "score": score, + "position": position, "bearing": bearing, "bearing_deg": deg, "distance": distance, @@ -111,13 +204,30 @@ def _objects_3d(dets: Detection3DArray, pose: PoseStamped) -> list[ObjectState]: def _objects_2d(dets: Detection2DArray, image_size: tuple[int, int]) -> list[ObjectState]: w, h = image_size out: list[ObjectState] = [] - for d in dets.to_json(): - b = d["bbox"] + for detection in dets.detections: + center = detection.bbox.center.position + b: BBoxJson = { + "cx": round(center.x, 1), + "cy": round(center.y, 1), + "w": round(detection.bbox.size_x, 1), + "h": round(detection.bbox.size_y, 1), + } + label = next( + ( + result.hypothesis.class_id + for result in detection.results + if result.hypothesis.class_id + ), + "", + ) + score = round( + max((result.hypothesis.score for result in detection.results), default=0.0), 2 + ) area = b["w"] * b["h"] / (w * h) out.append( { - "label": d["label"], - "score": d["score"], + "label": label, + "score": score, "bbox": b, "bearing": BEARING_WORDS_2D[min(4, int(5 * b["cx"] / w))], "size": next(word for limit, word in SIZE_WORDS_2D if area > limit), @@ -146,7 +256,7 @@ def build_world_state( objects = [] state: WorldState = { "goal": goal, - "robot": {**pose.to_json(), "motion": motion}, + "robot": {**pose_json(pose), "motion": motion}, "objects": objects, } if goal_xy is not None: @@ -161,7 +271,5 @@ def build_world_state( } if lidar is not None: z_min, z_max, max_range = lidar_band - state["room"] = lidar.to_json( - pose, sectors=SECTOR_NAMES, z_min=z_min, z_max=z_max, max_range=max_range - ) + state["room"] = room_sectors(lidar, pose, z_min, z_max, max_range) return state diff --git a/dimos/cli/bake/test_bake_e2e.py b/dimos/cli/bake/test_bake_e2e.py index 054402e3feb..cee6467ed28 100644 --- a/dimos/cli/bake/test_bake_e2e.py +++ b/dimos/cli/bake/test_bake_e2e.py @@ -178,7 +178,8 @@ def test_the_host_publishes_its_outputs_and_hides_the_suppressed_hop(baked, zeno header=Header(stamp=time_from_nanoseconds(stamp_ns), frame_id="odom"), child_frame_id="lidar", transform=Transform( - translation=Vector3(z=SENSOR_Z), rotation=Quaternion(w=1.0) + translation=Vector3(z=SENSOR_Z, x=0.0, y=0.0), + rotation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), ), ) ] diff --git a/dimos/cli/cameracalibrate/cameracalibrate.py b/dimos/cli/cameracalibrate/cameracalibrate.py index 1189aad776c..120aa6c841f 100644 --- a/dimos/cli/cameracalibrate/cameracalibrate.py +++ b/dimos/cli/cameracalibrate/cameracalibrate.py @@ -30,10 +30,13 @@ # Use setdefault so an explicit OPENCV_OPENCL_RUNTIME from the environment still wins. os.environ.setdefault("OPENCV_OPENCL_RUNTIME", "disabled") +from dimos_generated.sensor_msgs.msg import Image import numpy as np import typer import yaml +from dimos.msgs.image import image_to_bgr + _IMAGE_EXTS = frozenset({".png", ".jpg", ".jpeg"}) @@ -102,7 +105,7 @@ def write_camera_info_yaml( ) -> None: """Write ROS-style CameraInfo YAML loadable by dimos CameraInfo helpers. - The emitted schema is accepted by ``CameraInfo.from_yaml``. + The emitted schema is accepted by ``dimos.msgs.camera_info.camera_info_from_yaml``. """ k = np.asarray(K, dtype=np.float64).reshape(3, 3) d = np.asarray(D, dtype=np.float64).ravel() @@ -544,11 +547,10 @@ def _capture_frames_from_topic( ``topic_uri`` follows the pubsub registry format ``":"``, e.g. ``"zenoh:dimos/color_image"`` or ``"jpeg_lcm:/color_image"``. The publisher must - emit ``sensor_msgs.Image`` messages; ``Image.to_opencv()`` normalizes the + emit ``sensor_msgs.Image`` messages; ``image_to_bgr()`` normalizes the payload to BGR before detection. Raises ``RuntimeError`` if no frames arrive within ``timeout_sec``. """ - from dimos.msgs.sensor_msgs.Image import Image from dimos.protocol.pubsub.registry import subscribe_pubsub_uri if target_count < 1: @@ -560,7 +562,7 @@ def _capture_frames_from_topic( def _on_image(msg: Any) -> None: try: - arr = msg.to_opencv() + arr = image_to_bgr(msg) except (AttributeError, ValueError): return with lock: diff --git a/dimos/cli/cameracalibrate/test_cameracalibrate.py b/dimos/cli/cameracalibrate/test_cameracalibrate.py index 26501b509b9..6f076de3f14 100644 --- a/dimos/cli/cameracalibrate/test_cameracalibrate.py +++ b/dimos/cli/cameracalibrate/test_cameracalibrate.py @@ -17,6 +17,9 @@ from unittest.mock import MagicMock import cv2 +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from typer.testing import CliRunner @@ -31,8 +34,8 @@ load_frames_from_folder, write_camera_info_yaml, ) -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo as DimosCameraInfo -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.camera_info import camera_info_from_yaml, intrinsic_matrix +from dimos.msgs.image import image_from_array def _synthetic_chessboard_gray( @@ -158,15 +161,17 @@ def test_cli_folder_with_synthetic_images_writes_yaml_preview_and_camera_info( assert preview_image is not None assert preview_image.shape == (480, 640, 3) - dimos_info = DimosCameraInfo.from_yaml(str(out)) + dimos_info = camera_info_from_yaml( + out, header=Header(frame_id="camera_optical", stamp=Time(sec=0, nanosec=0)) + ) assert dimos_info.width == 640 assert dimos_info.height == 480 assert dimos_info.distortion_model == "plumb_bob" - assert dimos_info.frame_id == "camera_optical" - assert dimos_info.get_K_matrix().shape == (3, 3) - assert dimos_info.get_D_coeffs().shape == (5,) - assert dimos_info.get_R_matrix().shape == (3, 3) - assert dimos_info.get_P_matrix().shape == (3, 4) + assert dimos_info.header.frame_id == "camera_optical" + assert intrinsic_matrix(dimos_info).shape == (3, 3) + assert np.asarray(dimos_info.d).shape == (5,) + assert np.asarray(dimos_info.r).reshape(3, 3).shape == (3, 3) + assert np.asarray(dimos_info.p).reshape(3, 4).shape == (3, 4) def test_cli_help_lists_cameracalibrate_flags() -> None: @@ -228,7 +233,9 @@ def test_cli_folder_writes_only_explicit_yaml_and_prints_rms(tmp_path: Path) -> assert out.exists() preview = tmp_path / "camera_info.preview.png" assert not preview.exists() - dimos_info = DimosCameraInfo.from_yaml(str(out)) + dimos_info = camera_info_from_yaml( + out, header=Header(frame_id="camera_optical", stamp=Time(sec=0, nanosec=0)) + ) assert dimos_info.width == 640 assert dimos_info.height == 480 assert dimos_info.distortion_model == "plumb_bob" @@ -464,7 +471,7 @@ def test_capture_frames_from_topic_mocked_space_fills_target( cols, rows = 9, 6 gray = _synthetic_chessboard_gray(640, 480, cols, rows, square_px=40) bgr = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR) - image_msg = Image.from_opencv(bgr) + image_msg = image_from_array(bgr, encoding="bgr8") record: dict[str, object] = {} monkeypatch.setattr( @@ -525,7 +532,7 @@ def test_cli_topic_source_uri_passes_through( cols, rows = 9, 6 gray = _synthetic_chessboard_gray(640, 480, cols, rows, square_px=40) bgr = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR) - image_msg = Image.from_opencv(bgr) + image_msg = image_from_array(bgr, encoding="bgr8") record: dict[str, object] = {} monkeypatch.setattr( @@ -801,16 +808,18 @@ def test_write_camera_info_yaml_round_trip_matches_k_d_size_and_model(tmp_path: D=D, distortion_model="plumb_bob", ) - dimos_info = DimosCameraInfo.from_yaml(path) + dimos_info = camera_info_from_yaml( + path, header=Header(frame_id="camera_optical", stamp=Time(sec=0, nanosec=0)) + ) assert dimos_info.width == 640 assert dimos_info.height == 480 assert dimos_info.distortion_model == "plumb_bob" - assert np.allclose(dimos_info.get_K_matrix(), K) - assert np.allclose(dimos_info.get_D_coeffs(), D) + assert np.allclose(intrinsic_matrix(dimos_info), K) + assert np.allclose(np.asarray(dimos_info.d), D) def test_write_camera_info_yaml_round_trip(tmp_path: Path) -> None: - """YAML written by ``write_camera_info_yaml`` round-trips through ``CameraInfo.from_yaml``.""" + """YAML written by ``write_camera_info_yaml`` round-trips through ``camera_info_from_yaml``.""" K = np.array([[600.0, 0.5, 400.0], [0.0, 605.0, 300.5], [0.0, 0.0, 1.0]], dtype=np.float64) D = np.array([-0.12, 0.08, 0.002, -0.001, 0.0], dtype=np.float64) R = np.array([[0.999, -0.01, 0.0], [0.01, 0.999, 0.0], [0.0, 0.0, 1.0]], dtype=np.float64) @@ -830,15 +839,17 @@ def test_write_camera_info_yaml_round_trip(tmp_path: Path) -> None: P=P, distortion_model="plumb_bob", ) - dimos_info = DimosCameraInfo.from_yaml(path) + dimos_info = camera_info_from_yaml( + path, header=Header(frame_id="camera_optical", stamp=Time(sec=0, nanosec=0)) + ) assert dimos_info.width == 800 assert dimos_info.height == 600 assert dimos_info.distortion_model == "plumb_bob" - assert dimos_info.frame_id == "camera_optical" - assert np.allclose(dimos_info.get_K_matrix(), K) - assert np.allclose(dimos_info.get_D_coeffs(), D) - assert np.allclose(dimos_info.get_R_matrix(), R) - assert np.allclose(dimos_info.get_P_matrix(), P) + assert dimos_info.header.frame_id == "camera_optical" + assert np.allclose(intrinsic_matrix(dimos_info), K) + assert np.allclose(np.asarray(dimos_info.d), D) + assert np.allclose(np.asarray(dimos_info.r).reshape(3, 3), R) + assert np.allclose(np.asarray(dimos_info.p).reshape(3, 4), P) def test_write_camera_info_yaml_custom_r_p_and_distortion_model(tmp_path: Path) -> None: @@ -858,11 +869,13 @@ def test_write_camera_info_yaml_custom_r_p_and_distortion_model(tmp_path: Path) P=P, distortion_model="rational_polynomial", ) - dimos_info = DimosCameraInfo.from_yaml(path) + dimos_info = camera_info_from_yaml( + path, header=Header(frame_id="camera_optical", stamp=Time(sec=0, nanosec=0)) + ) assert dimos_info.width == 320 assert dimos_info.height == 240 assert dimos_info.distortion_model == "rational_polynomial" - assert np.allclose(dimos_info.get_K_matrix(), K) - assert np.allclose(dimos_info.get_D_coeffs(), D) - assert np.allclose(dimos_info.get_R_matrix(), R) - assert np.allclose(dimos_info.get_P_matrix(), P) + assert np.allclose(intrinsic_matrix(dimos_info), K) + assert np.allclose(np.asarray(dimos_info.d), D) + assert np.allclose(np.asarray(dimos_info.r).reshape(3, 3), R) + assert np.allclose(np.asarray(dimos_info.p).reshape(3, 4), P) diff --git a/dimos/cli/hardware/g1.py b/dimos/cli/hardware/g1.py index b41ab602078..9c33e9fd8a3 100644 --- a/dimos/cli/hardware/g1.py +++ b/dimos/cli/hardware/g1.py @@ -19,7 +19,10 @@ import time from typing import Any, NoReturn, Protocol, TypeGuard +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header +import numpy as np import typer from dimos.control.tasks.trajectory_task.trajectory_task import JOINT_TRAJECTORY_TASK_NAME @@ -163,7 +166,14 @@ def _execute_ready_pose( _require_armed_and_enabled(coordinator) _require_teleop_disengaged(coordinator) targets = { - group: JointState(position=list(positions)) for group, positions in G1_READY_JOINTS.items() + group: JointState( + position=np.asarray(list(positions), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + for group, positions in G1_READY_JOINTS.items() } planned = manipulation.plan_to_joints(targets, speed_scale=G1_READY_SPEED_SCALE) if not planned.succeeded: diff --git a/dimos/cli/spy/test_spy.py b/dimos/cli/spy/test_spy.py index 11fe0d18dea..cf6fc2c2ce7 100644 --- a/dimos/cli/spy/test_spy.py +++ b/dimos/cli/spy/test_spy.py @@ -26,6 +26,7 @@ import time from typing import Any +from dimos_cdrgen.registry import encode as cdr_encode from dimos_generated.geometry_msgs.msg import Vector3 import pytest @@ -45,7 +46,7 @@ from dimos.protocol.service.zenohservice import ZenohSessionPool VEC = Vector3(x=1.0, y=2.0, z=3.0) -VEC_BYTES = VEC.encode() +VEC_BYTES = cdr_encode(VEC) # TopicStats: pure, deterministic (injected timestamps, no sleeps) @@ -201,7 +202,8 @@ def _assert_no_decode(monkeypatch): def boom(*a, **k): raise AssertionError("spy decoded a payload on the hot path") - monkeypatch.setattr(Vector3, "decode", staticmethod(boom)) + monkeypatch.setattr("dimos_cdrgen.registry.decode", boom) + monkeypatch.setattr("dimos.protocol.pubsub.encoders.cdr_decode", boom) @pytest.fixture diff --git a/dimos/cli/test_topic.py b/dimos/cli/test_topic.py index ed7340bb1da..fa8789c8c6a 100644 --- a/dimos/cli/test_topic.py +++ b/dimos/cli/test_topic.py @@ -12,9 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_cdrgen.registry import encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Twist, Vector3 -from dimos_generated.sensor_msgs.msg import CameraInfo +from dimos_generated.sensor_msgs.msg import CameraInfo, RegionOfInterest from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from dimos.cli.topic import _build_eval_context, _decode_typed_lcm_message, topic_send @@ -30,12 +33,19 @@ def test_decode_typed_lcm_message_resolves_message_submodule() -> None: width=1920, height=1080, distortion_model="plumb_bob", - header=Header(frame_id="camera_optical"), + header=Header(frame_id="camera_optical", stamp=Time(sec=0, nanosec=0)), + d=np.array([], dtype=np.float64), + k=np.zeros(9, dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), ) decoded = _decode_typed_lcm_message( "/camera_info#sensor_msgs/msg/CameraInfo", - msg.encode(), + cdr_encode(msg), ) assert isinstance(decoded, CameraInfo) @@ -51,7 +61,7 @@ def test_build_eval_context_maps_message_names_to_classes() -> None: for name in ("Twist", "Vector3", "PoseStamped"): cls = context[name] assert isinstance(cls, type) - assert cls.__name__ == name + assert cls.__msgtype__.rsplit("/", 1)[1] == name def test_topic_send_delivers_over_lcm(monkeypatch: pytest.MonkeyPatch, lcm_url: str) -> None: @@ -74,11 +84,16 @@ def capturing_make_transport(name: str, msg_type: type) -> PubSubTransport[objec bus.subscribe(Topic(topic="/test_topic_send", msg_type=Twist), collector) try: - topic_send("/test_topic_send", "Twist(linear=Vector3(x=0.5), angular=Vector3())") + topic_send( + "/test_topic_send", + "Twist(linear=Vector3(x=0.5, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0))", + ) collector.wait() finally: for transport in transports: transport.stop() bus.stop() - assert collector.results[0][0] == Twist(linear=Vector3(x=0.5), angular=Vector3()) + assert collector.results[0][0] == Twist( + linear=Vector3(x=0.5, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) diff --git a/dimos/cli/topic.py b/dimos/cli/topic.py index d9896ad58b2..bff7921689b 100644 --- a/dimos/cli/topic.py +++ b/dimos/cli/topic.py @@ -20,12 +20,12 @@ from types import SimpleNamespace from typing import cast +from dimos_cdrgen.registry import decode as cdr_decode, message_types import typer from dimos.core.global_config import global_config from dimos.core.transport import PubSubTransport from dimos.core.transport_factory import make_transport, transport_topic -from dimos.message_codegen.registry import message_types from dimos.msgs.helpers import resolve_msg_type from dimos.protocol.pubsub.impl.lcmpubsub import LCMPubSubBase, Topic from dimos.protocol.pubsub.impl.zenohpubsub import Zenoh @@ -48,7 +48,7 @@ def _decode_typed_lcm_message(channel: str, data: bytes) -> object: cls = resolve_msg_type(msg_name) if cls is None: raise ValueError(f"Could not resolve message type from channel: {channel}") - return cls.decode(data) + return cdr_decode(data, cls.__msgtype__) def _listen_forever(listening_msg: str, on_stop: Callable[[], None] = lambda: None) -> None: diff --git a/dimos/cloud/preview.py b/dimos/cloud/preview.py index de8dfcd4119..4f32197d0d5 100644 --- a/dimos/cloud/preview.py +++ b/dimos/cloud/preview.py @@ -31,6 +31,16 @@ except ImportError: # pragma: no cover HAS_AV = False +import cv2 +from dimos_generated.geometry_msgs.msg import ( + PoseStamped, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.sensor_msgs.msg import Image, Joy, PointCloud2 +from dimos_generated.std_msgs.msg import Header + from dimos.cloud.constants import ( PREVIEW_BAND, PREVIEW_FORMAT, @@ -48,11 +58,16 @@ TIMELAPSE_MAX_S, WORLD_FRAMES, ) -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.Joy import Joy -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.geometry import yaw +from dimos.msgs.image import image_brightness, image_resize_to_fit, image_to_bgr, image_to_jpeg +from dimos.msgs.pointcloud import ( + concatenate_clouds, + pointcloud_xyz, + select_points, + transform_cloud, + voxel_downsample_cloud, +) +from dimos.msgs.time import time_from_seconds if TYPE_CHECKING: from dimos.memory.store.base import Store @@ -76,22 +91,26 @@ def _evenly(n: int, k: int) -> set[int]: def _fit(pc: PointCloud2, points: int) -> PointCloud2: """Coarser voxels until it fits: keeps the shape better than random sampling.""" - n = len(pc.points_f32()) + n = len(pointcloud_xyz(pc)) return ( - pc if n <= points else pc.voxel_downsample(PREVIEW_MAP_VOXEL * float(np.sqrt(n / points))) + pc + if n <= points + else voxel_downsample_cloud(pc, PREVIEW_MAP_VOXEL * float(np.sqrt(n / points))) ) def _pack(pc: PointCloud2, origin: np.ndarray) -> str: """The browser's point format: little-endian int16 triplets of PREVIEW_SCALE around origin.""" - q = np.clip(np.round((pc.points_f32() - origin) / PREVIEW_SCALE), -32767, 32767).astype(" dict[str, Any] | None: """`dimos-spatial-preview-v2`, or None for a recording without LiDAR, camera or poses.""" lidar, camera = _stream(store, PointCloud2, "lidar"), _stream(store, Image, "color") - odom = _stream(store, Pose, "odom") + odom = _stream(store, PoseStamped, "odom") joy = _stream(store, Joy, "joy") poses, scans, world, ends = [], [], None, [] if lidar is not None: @@ -104,19 +123,35 @@ def build(store: Store) -> dict[str, Any] | None: if i not in merged: continue pc = obs.data - if pc.frame_id not in WORLD_FRAMES: # sensor frame: place it with its own pose + if pc.header.frame_id not in WORLD_FRAMES: # sensor frame: place it with its own pose if obs.pose is None: continue - pc = pc.transform(Transform.from_pose(WORLD_FRAMES[0], obs.pose)) - pc = pc.voxel_downsample(PREVIEW_MAP_VOXEL) + pc = transform_cloud( + pc, + TransformStamped( + header=Header(stamp=time_from_seconds(obs.ts), frame_id=WORLD_FRAMES[0]), + child_frame_id=pc.header.frame_id, + transform=Transform( + translation=Vector3( + x=obs.pose.position.x, y=obs.pose.position.y, z=obs.pose.position.z + ), + rotation=obs.pose.orientation, + ), + ), + ) + pc = voxel_downsample_cloud(pc, PREVIEW_MAP_VOXEL) if i in shown: scans.append((obs.ts, pc)) - world = pc if world is None else (world + pc).voxel_downsample(PREVIEW_MAP_VOXEL) + world = ( + pc + if world is None + else voxel_downsample_cloud(concatenate_clouds(world, pc), PREVIEW_MAP_VOXEL) + ) ends = ends[:1] + ends[-1:] if odom is not None and odom.count() > len( poses ): # denser than the LiDAR poses, or the only path - poses = [(o.ts, o.data) for o in odom] + poses = [(o.ts, o.data.pose) for o in odom] sticks = [] if joy is not None: picked = _evenly(joy.count(), PREVIEW_JOY_SAMPLES) @@ -132,13 +167,31 @@ def build(store: Store) -> dict[str, Any] | None: z = float(np.median([p.position.z for _, p in poses])) if poses else 0.0 lo, hi = z + PREVIEW_BAND[0], z + PREVIEW_BAND[1] if world is not None: - world = _fit(world.filter_by_height(lo, hi), PREVIEW_MAP_POINTS) - scans = [(t, _fit(pc.filter_by_height(lo, hi), PREVIEW_SCAN_POINTS)) for t, pc in scans] + world = _fit( + select_points( + world, (pointcloud_xyz(world)[:, 2] > lo) & (pointcloud_xyz(world)[:, 2] < hi) + ), + PREVIEW_MAP_POINTS, + ) + scans = [ + ( + t, + _fit( + select_points( + pc, (pointcloud_xyz(pc)[:, 2] > lo) & (pointcloud_xyz(pc)[:, 2] < hi) + ), + PREVIEW_SCAN_POINTS, + ), + ) + for t, pc in scans + ] t0, t1 = min(ends), max(ends) - traj = [[t - t0, p.position.x, p.position.y, p.position.z, p.yaw] for t, p in poses] + traj = [ + [t - t0, p.position.x, p.position.y, p.position.z, yaw(p.orientation)] for t, p in poses + ] traj = [traj[i] for i in sorted(_evenly(len(traj), 3000))] - pts = world.points_f32().astype(np.float64) if world is not None else np.zeros((0, 3)) + pts = pointcloud_xyz(world).astype(np.float64) if world is not None else np.zeros((0, 3)) if not len(pts): # no map: frame the path pts = np.array([r[1:4] for r in traj] or [[0.0, 0.0, z]]) origin = np.round(pts.mean(axis=0), 2) @@ -163,13 +216,15 @@ def build(store: Store) -> dict[str, Any] | None: "camera": [ { "t": round(o.ts - t0, 3), - "jpeg": o.data.to_base64( - 70, max_width=PREVIEW_THUMB_PX, max_height=PREVIEW_THUMB_PX - ), + "jpeg": base64.b64encode( + image_to_jpeg( + image_resize_to_fit(o.data, PREVIEW_THUMB_PX, PREVIEW_THUMB_PX)[0], 70 + ) + ).decode(), } for o in shots ], - "thumb": int(np.argmax([o.data.brightness for o in shots])) if shots else None, + "thumb": int(np.argmax([image_brightness(o.data) for o in shots])) if shots else None, "joy": [ [round(o.ts - t0, 3), [round(a, 2) for a in o.data.axes], list(o.data.buttons)] for o in sticks @@ -203,7 +258,7 @@ def timelapse(store: Store, out: Path) -> dict[str, Any] | None: video.options = {"crf": str(TIMELAPSE_CRF), "preset": "veryfast"} frame = av.VideoFrame.from_ndarray( np.ascontiguousarray( - img.resize(video.width, video.height).to_bgr().as_numpy()[:, :, :3] + cv2.resize(image_to_bgr(img), (video.width, video.height)), dtype=np.uint8 ), format="bgr24", ) diff --git a/dimos/cloud/test_preview.py b/dimos/cloud/test_preview.py index 622efbb8ceb..38d2dd2e07a 100644 --- a/dimos/cloud/test_preview.py +++ b/dimos/cloud/test_preview.py @@ -18,17 +18,28 @@ from types import SimpleNamespace from typing import Any, cast +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import Image, Joy, PointCloud2 +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from dimos.cloud import preview from dimos.cloud.constants import PREVIEW_SCALE from dimos.memory.store.base import Store -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.Joy import Joy -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.image import image_from_array +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import time_from_seconds + + +def make_pose( + x: float, y: float, z: float, qx: float = 0.0, qy: float = 0.0, qz: float = 0.0, qw: float = 1.0 +) -> Pose: + return Pose(position=Point(x=x, y=y, z=z), orientation=Quaternion(x=qx, y=qy, z=qz, w=qw)) + + +def make_image(array: np.ndarray) -> Image: + return image_from_array(array, encoding="rgb8" if array.ndim == 3 else "16UC1") class Stream: @@ -58,17 +69,22 @@ def points(doc: dict[str, Any], b64: str) -> np.ndarray: def cloud(t: float, *xyz: tuple[float, float, float], frame: str = "world") -> PointCloud2: - return PointCloud2.from_numpy(np.array(xyz, dtype=np.float32), frame_id=frame, timestamp=t) + return pointcloud_from_xyz( + np.array(xyz, dtype=np.float32), header=Header(stamp=time_from_seconds(t), frame_id=frame) + ) def test_world_frame_recording() -> None: lidar = [ - (100.0 + i, cloud(100.0 + i, (i, 1, 0.5), (i, -1, 0.5), (i, 0, 9.0)), Pose(i * 0.5, 0, 0.3)) + ( + 100.0 + i, + cloud(100.0 + i, (i, 1, 0.5), (i, -1, 0.5), (i, 0, 9.0)), + make_pose(i * 0.5, 0, 0.3), + ) for i in range(10) ] frames = [ - (100.0 + i, Image.from_numpy(np.full((48, 64, 3), 25 * i, np.uint8)), None) - for i in range(10) + (100.0 + i, make_image(np.full((48, 64, 3), 25 * i, np.uint8)), None) for i in range(10) ] mapper = [(99.0, cloud(99.0, (50, 50, 0.5)), None)] # listed first, not a scan doc = preview.build(store(global_map=mapper, lidar=lidar, color_image=frames)) @@ -81,7 +97,7 @@ def test_world_frame_recording() -> None: def test_sensor_frame_scans_use_their_own_pose() -> None: - mount = Pose( + mount = make_pose( 5.0, 0.0, 0.3, 0.0, 0.0, float(np.sin(np.pi / 4)), float(np.cos(np.pi / 4)) ) # 90 deg yaw lidar = [ @@ -96,10 +112,9 @@ def test_sensor_frame_scans_use_their_own_pose() -> None: def test_timelapse(tmp_path: Path) -> None: frames = [ - (100.0 + i, Image.from_numpy(np.full((48, 64, 3), 8 * i, np.uint8)), None) - for i in range(30) + (100.0 + i, make_image(np.full((48, 64, 3), 8 * i, np.uint8)), None) for i in range(30) ] - depth = [(100.0, Image.from_numpy(np.zeros((48, 64), np.uint16)), None)] # listed first + depth = [(100.0, make_image(np.zeros((48, 64), np.uint16)), None)] # listed first meta = preview.timelapse(store(depth_image=depth, color_image=frames), tmp_path / "t.mp4") assert ( meta is not None and meta["speed"] == 1.0 and meta["duration_s"] == 29 @@ -109,10 +124,20 @@ def test_timelapse(tmp_path: Path) -> None: def test_path_from_denser_odometry() -> None: odom = [ - (100.0 + i, PoseStamped(position=[i, 2, 0.3], frame_id="world"), None) for i in range(5) + ( + 100.0 + i, + PoseStamped( + header=Header(stamp=time_from_seconds(0.0), frame_id="world"), + pose=make_pose(i, 2, 0.3), + ), + None, + ) + for i in range(5) ] - frames = [(101.0, Image.from_numpy(np.zeros((48, 64, 3), np.uint8)), None)] - lidar = [(102.0, cloud(102.0, (0, 0, 0.3)), Pose(9, 9, 0.3))] # one posed scan: odom is denser + frames = [(101.0, make_image(np.zeros((48, 64, 3), np.uint8)), None)] + lidar = [ + (102.0, cloud(102.0, (0, 0, 0.3)), make_pose(9, 9, 0.3)) + ] # one posed scan: odom is denser doc = preview.build(store(color_image=frames, odom=odom, lidar=lidar)) assert doc is not None and len(doc["scans"]) == 1 and doc["duration_s"] == 4 assert [r[1] for r in doc["trajectory"]] == [0, 1, 2, 3, 4] and len(doc["camera"]) == 1 @@ -120,11 +145,35 @@ def test_path_from_denser_odometry() -> None: def test_joystick_samples_ride_along() -> None: odom = [ - (100.0 + i, PoseStamped(position=[i, 0, 0.3], frame_id="world"), None) for i in range(3) + ( + 100.0 + i, + PoseStamped( + header=Header(stamp=time_from_seconds(0.0), frame_id="world"), + pose=make_pose(i, 0, 0.3), + ), + None, + ) + for i in range(3) ] sticks = [ - (100.5, Joy(axes=[0.0, 0.4567, -1.0, 0.0], buttons=[1, 0, 0]), None), - (101.5, Joy(axes=[0.0, 0.0, 0.0, 0.0], buttons=[0, 0, 1]), None), + ( + 100.5, + Joy( + header=Header(stamp=time_from_seconds(0.0), frame_id=""), + axes=np.array([0.0, 0.4567, -1.0, 0.0], dtype=np.float32), + buttons=np.array([1, 0, 0], dtype=np.int32), + ), + None, + ), + ( + 101.5, + Joy( + header=Header(stamp=time_from_seconds(0.0), frame_id=""), + axes=np.array([0.0, 0.0, 0.0, 0.0], dtype=np.float32), + buttons=np.array([0, 0, 1], dtype=np.int32), + ), + None, + ), ] doc = preview.build(store(odom=odom, joystick=sticks)) assert doc is not None and doc["streams"]["joystick"] == {"name": "joystick", "count": 2} @@ -138,8 +187,28 @@ def test_joystick_samples_ride_along() -> None: def test_joystick_is_capped_evenly(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(preview, "PREVIEW_JOY_SAMPLES", 3) - odom = [(100.0, PoseStamped(position=[0, 0, 0.3], frame_id="world"), None)] - sticks = [(100.0 + i, Joy(axes=[i / 10], buttons=[]), None) for i in range(10)] + odom = [ + ( + 100.0, + PoseStamped( + header=Header(stamp=time_from_seconds(0.0), frame_id="world"), + pose=make_pose(0, 0, 0.3), + ), + None, + ) + ] + sticks = [ + ( + 100.0 + i, + Joy( + header=Header(stamp=time_from_seconds(0.0), frame_id=""), + axes=np.array([i / 10], dtype=np.float32), + buttons=np.array([], dtype=np.int32), + ), + None, + ) + for i in range(10) + ] doc = preview.build(store(odom=odom, joystick=sticks)) assert doc is not None and doc["streams"]["joystick"]["count"] == 10 assert [r[0] for r in doc["joy"]] == [0, 4, 9] and [r[1] for r in doc["joy"]] == [ diff --git a/dimos/conftest.py b/dimos/conftest.py index 7216e75a340..d8e155aa6f5 100644 --- a/dimos/conftest.py +++ b/dimos/conftest.py @@ -20,6 +20,7 @@ import pathlib import platform import shutil +import sys import tempfile import threading import time @@ -84,6 +85,7 @@ def _lcm_bucket(seed: str) -> int: if soft < target: resource.setrlimit(resource.RLIMIT_NOFILE, (target, hard)) +from dimos_cdrgen.registry import message_types from dotenv import dotenv_values import pytest import tqdm @@ -171,6 +173,25 @@ def pytest_sessionstart(session): _arm_crash_dumps() +@pytest.hookimpl(optionalhook=True) +def pytest_testnodedown(node, error): + """Expose a crashed worker's dump before a later job timeout loses it.""" + if error is None: + return + base = pathlib.Path( + os.environ.get("DIMOS_CRASH_DIR") + or os.path.join(os.environ.get("RUNNER_TEMP") or tempfile.gettempdir(), "pytest-crash") + ) + for path in sorted(base.glob(f"crash-{node.gateway.id}-*.log")): + try: + dump = path.read_text(errors="replace") + except OSError as exc: + print(f"Cannot read worker crash diagnostic {path.name}: {exc}", flush=True) + continue + if dump: + print(f"\nWorker crash diagnostic {path.name}:\n{dump}", flush=True) + + def pytest_ignore_collect(collection_path: pathlib.Path) -> bool | None: # Nested Python projects own their dependencies and test invocation. if collection_path.is_dir() and (collection_path / "pyproject.toml").is_file(): @@ -228,6 +249,18 @@ def pytest_configure(config): ) +@pytest.fixture(autouse=True) +def _restore_reference_types_module(): + # Reference typestores replace rosbags' process-global class module. Keep + # their test-local identities from invalidating canonical message pickles. + message_types() + previous = sys.modules["rosbags.usertypes"] + try: + yield + finally: + sys.modules["rosbags.usertypes"] = previous + + def _global_config_guard(): from dimos.core.global_config import global_config @@ -274,8 +307,10 @@ def mcp_url(mcp_port: int) -> str: @pytest.fixture(scope="session") def lcm_url() -> str: - """The LCM bus URL pinned for this xdist worker (or the default).""" - return os.environ.get("LCM_DEFAULT_URL", "udpm://239.255.76.67:7667?ttl=0") + """The worker's LCM bus URL, with an optional test-only provider override.""" + return os.environ.get("DIMOS_TEST_LCM_URL") or os.environ.get( + "LCM_DEFAULT_URL", "udpm://239.255.76.67:7667?ttl=0" + ) @pytest.fixture diff --git a/dimos/control/benchmarking/benchmark.py b/dimos/control/benchmarking/benchmark.py index 9de58d99463..1c467e8e752 100644 --- a/dimos/control/benchmarking/benchmark.py +++ b/dimos/control/benchmarking/benchmark.py @@ -51,6 +51,7 @@ import time from typing import Any, Literal +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Twist from dimos_generated.nav_msgs.msg import Path as NavPath from dimos_generated.std_msgs.msg import Float32, Header, Int8 @@ -126,7 +127,7 @@ def shift_path_to_start_at_pose(path: NavPath, start_pose: PoseStamped) -> NavPa rx, ry = p.pose.position.x - px0, p.pose.position.y - py0 new.append( PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=Point(x=sx + rx * cd - ry * sd, y=sy + rx * sd + ry * cd, z=0.0), orientation=quaternion_from_euler( @@ -135,7 +136,7 @@ def shift_path_to_start_at_pose(path: NavPath, start_pose: PoseStamped) -> NavPa ), ) ) - return NavPath(poses=new) + return NavPath(poses=new, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) # Recording (flat per-run JSON) + round-trip @@ -202,14 +203,15 @@ def reference_path(self) -> NavPath: return NavPath( poses=[ PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=Point(x=x, y=y, z=0.0), orientation=quaternion_from_euler(0.0, 0.0, yaw), ), ) for x, y, yaw in self.reference - ] + ], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) diff --git a/dimos/control/benchmarking/paths.py b/dimos/control/benchmarking/paths.py index 0138ec9248e..9efa2df7b82 100644 --- a/dimos/control/benchmarking/paths.py +++ b/dimos/control/benchmarking/paths.py @@ -22,6 +22,7 @@ import math +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point as GeoPoint, Pose, @@ -48,19 +49,19 @@ def _xy_to_path(executed_xy: list[tuple[float, float]]) -> Path: """Wrap (x, y) tuples in a nav_msgs.Path so memory Polyline can render them.""" poses = [ PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=GeoPoint(x=x, y=y, z=0.0), orientation=quaternion_from_euler(0.0, 0.0, 0.0) ), ) for x, y in executed_xy ] - return Path(poses=poses) + return Path(poses=poses, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) def _pose(x: float, y: float, yaw: float) -> PoseStamped: return PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=GeoPoint(x=x, y=y, z=0.0), orientation=quaternion_from_euler(0.0, 0.0, yaw) ), @@ -80,7 +81,7 @@ def _path_from_xy(xs: list[float], ys: list[float]) -> Path: dy = ys[i] - ys[i - 1] yaw = math.atan2(dy, dx) poses.append(_pose(xs[i], ys[i], yaw)) - return Path(poses=poses) + return Path(poses=poses, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) # Path generators @@ -205,7 +206,10 @@ def straight_rotate( ) -> Path: """Translate straight along +x while the commanded yaw ramps 0 -> yaw_end.""" n = round(length / step) - return Path(poses=[_pose(i * step, 0.0, yaw_end * i / n) for i in range(n + 1)]) + return Path( + poses=[_pose(i * step, 0.0, yaw_end * i / n) for i in range(n + 1)], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) def strafe_line(length: float = 2.0, step: float = 0.05) -> Path: @@ -216,7 +220,10 @@ def strafe_line(length: float = 2.0, step: float = 0.05) -> Path: turning wrong: it must be executed as pure lateral translation. """ n = round(length / step) - return Path(poses=[_pose(0.0, i * step, 0.0) for i in range(n + 1)]) + return Path( + poses=[_pose(0.0, i * step, 0.0) for i in range(n + 1)], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) def circle_offset_heading( @@ -235,7 +242,7 @@ def circle_offset_heading( y = radius * (1.0 - math.cos(theta)) tangent = math.atan2(math.sin(theta), math.cos(theta)) poses.append(_pose(x, y, tangent + offset)) - return Path(poses=poses) + return Path(poses=poses, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) def hold_heading(path: Path, yaw: float = 0.0) -> Path: @@ -379,7 +386,7 @@ def sidestep_1m(distance: float = 1.0, n_points: int = 20) -> Path: for i in range(n_points + 1): a = i / n_points poses.append(_pose(0.0, a * distance, 0.0)) - return Path(poses=poses) + return Path(poses=poses, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) def short_battery() -> dict[str, Path]: @@ -465,8 +472,8 @@ def trajectory_to_svg( sp.add(Polyline(msg=_xy_to_path(executed_xy), color=_EXE_COLOR, width=_EXE_WIDTH)) sx, sy = executed_xy[0] ex, ey = executed_xy[-1] - sp.add(Point(msg=GeoPoint(x=sx, y=sy), color=_START_COLOR, radius=_MARKER_RADIUS)) - sp.add(Point(msg=GeoPoint(x=ex, y=ey), color=_END_COLOR, radius=_MARKER_RADIUS)) + sp.add(Point(msg=GeoPoint(x=sx, y=sy, z=0.0), color=_START_COLOR, radius=_MARKER_RADIUS)) + sp.add(Point(msg=GeoPoint(x=ex, y=ey, z=0.0), color=_END_COLOR, radius=_MARKER_RADIUS)) return sp.to_svg() @@ -512,7 +519,7 @@ def multi_trajectory_to_svg( if xy: sp.add(Polyline(msg=_xy_to_path(xy), color=color, width=_EXE_WIDTH)) sx, sy = xy[0] - sp.add(Point(msg=GeoPoint(x=sx, y=sy), color=color, radius=_MARKER_RADIUS * 0.7)) + sp.add(Point(msg=GeoPoint(x=sx, y=sy, z=0.0), color=color, radius=_MARKER_RADIUS * 0.7)) # Legend row (world coords below the plot). ly = y_min - 0.4 - i * 0.25 sp.add( @@ -520,20 +527,21 @@ def multi_trajectory_to_svg( msg=Path( poses=[ PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=GeoPoint(x=x_min, y=ly, z=0.0), orientation=quaternion_from_euler(0, 0, 0), ), ), PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=GeoPoint(x=x_min + 0.4, y=ly, z=0.0), orientation=quaternion_from_euler(0, 0, 0), ), ), - ] + ], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ), color=color, width=_EXE_WIDTH, diff --git a/dimos/control/benchmarking/score.py b/dimos/control/benchmarking/score.py index d17ee0036be..8f13b9b8590 100644 --- a/dimos/control/benchmarking/score.py +++ b/dimos/control/benchmarking/score.py @@ -60,7 +60,9 @@ def _executed_from_recording(rec: RunRecording) -> ExecutedTrajectory: cmd_twist=Twist( linear=Vector3(x=cvx, y=cvy, z=0.0), angular=Vector3(x=0.0, y=0.0, z=cwz) ), - actual_twist=Twist(), + actual_twist=Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ), ) ) return ExecutedTrajectory(ticks=ticks, arrived=rec.arrived) diff --git a/dimos/control/benchmarking/test_benchmark.py b/dimos/control/benchmarking/test_benchmark.py index 958d3537d90..49f991ce28d 100644 --- a/dimos/control/benchmarking/test_benchmark.py +++ b/dimos/control/benchmarking/test_benchmark.py @@ -21,6 +21,7 @@ from types import SimpleNamespace +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Twist, Vector3 from dimos_generated.std_msgs.msg import Header import pytest @@ -53,7 +54,7 @@ def _pose(x=0.0, y=0.0, yaw=0.0) -> PoseStamped: return PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=Point(x=x, y=y, z=0.0), orientation=quaternion_from_euler(0.0, 0.0, yaw) ), diff --git a/dimos/control/contract/test_presets.py b/dimos/control/contract/test_presets.py index 9550cf75b5c..f9724c89e46 100644 --- a/dimos/control/contract/test_presets.py +++ b/dimos/control/contract/test_presets.py @@ -31,6 +31,7 @@ from pathlib import Path import pickle +from dimos_generated.dimos_msgs.msg import ControlValues import pytest from dimos.control.contract.conftest import ARM_JOINTS, G1_JOINTS @@ -50,7 +51,7 @@ validate_command, validate_description, ) -from dimos.msgs.control_msgs.ControlValues import ControlValues +from dimos.msgs.control import control_values from dimos.robot.assets.model import RobotModel CHASSIS_LIMITS = { @@ -123,7 +124,7 @@ def preset_chassis() -> ControlDescription: def command(values: dict[str, float]) -> ControlValues: - return ControlValues( + return control_values( source="coordinator", epoch=0, sequence=1, diff --git a/dimos/control/contract/test_validate.py b/dimos/control/contract/test_validate.py index 42863ca2ba3..c635caad3cc 100644 --- a/dimos/control/contract/test_validate.py +++ b/dimos/control/contract/test_validate.py @@ -24,6 +24,7 @@ import pickle from types import MappingProxyType +from dimos_generated.dimos_msgs.msg import ControlValues import pytest from dimos.control.contract.description import ( @@ -42,7 +43,7 @@ validate_description, validate_state, ) -from dimos.msgs.control_msgs.ControlValues import ControlValues +from dimos.msgs.control import control_values def errors_of(desc: ControlDescription) -> list[str]: @@ -56,7 +57,7 @@ def errors_of(desc: ControlDescription) -> list[str]: def command(source: str, values: dict[str, float], *, sequence: int = 1) -> ControlValues: """A command frame carrying ``values``.""" - return ControlValues( + return control_values( source=source, source_ts=0.0, epoch=0, @@ -69,7 +70,7 @@ def command(source: str, values: dict[str, float], *, sequence: int = 1) -> Cont def full_state(desc: ControlDescription) -> ControlValues: """A complete, well-formed state frame for ``desc``.""" keys = list(desc.state_keys()) - return ControlValues( + return control_values( source=desc.source, source_ts=0.0, epoch=0, @@ -266,12 +267,12 @@ def test_frame_rejected_error_pickles_and_reads_well() -> None: def test_mismatched_array_lengths_are_caught_by_both_validators( xarm: ControlDescription, ) -> None: - """ControlValues forbids this at construction, so only a mutated frame gets here. + """The control_values helper forbids this at construction, so only a mutated frame gets here. Both paths still check it, because validate_command is specified never to raise and a ragged zip would break that promise. """ - frame = ControlValues("arm", 0.0, 0, 1, ["arm/joint1/position"], [0.1]) + frame = control_values("arm", 0.0, 0, 1, ["arm/joint1/position"], [0.1]) frame.values = [0.1, 0.2] with pytest.raises(FrameRejectedError, match="shape"): @@ -302,7 +303,7 @@ def test_state_includes_the_imu(g1: ControlDescription) -> None: def test_state_from_the_wrong_source(xarm: ControlDescription) -> None: keys = list(xarm.state_keys()) - frame = ControlValues("somebody_else", 0.0, 0, 1, keys, [0.0] * len(keys)) + frame = control_values("somebody_else", 0.0, 0, 1, keys, [0.0] * len(keys)) with pytest.raises(FrameRejectedError, match="source"): validate_state(xarm, frame) @@ -311,7 +312,7 @@ def test_state_from_the_wrong_source(xarm: ControlDescription) -> None: def test_state_missing_a_key(xarm: ControlDescription) -> None: """A state frame is a full report: a partial one is a bug, not a sparse update.""" keys = list(xarm.state_keys())[:-1] - frame = ControlValues("arm", 0.0, 0, 1, keys, [0.0] * len(keys)) + frame = control_values("arm", 0.0, 0, 1, keys, [0.0] * len(keys)) with pytest.raises(FrameRejectedError, match="missing_key"): validate_state(xarm, frame) @@ -319,7 +320,7 @@ def test_state_missing_a_key(xarm: ControlDescription) -> None: def test_state_with_an_unknown_key(xarm: ControlDescription) -> None: keys = [*xarm.state_keys(), "arm/joint1/surprise"] - frame = ControlValues("arm", 0.0, 0, 1, keys, [0.0] * len(keys)) + frame = control_values("arm", 0.0, 0, 1, keys, [0.0] * len(keys)) with pytest.raises(FrameRejectedError, match="unknown_key"): validate_state(xarm, frame) @@ -327,7 +328,7 @@ def test_state_with_an_unknown_key(xarm: ControlDescription) -> None: def test_state_with_a_duplicate_key(xarm: ControlDescription) -> None: keys = [*xarm.state_keys(), xarm.state_keys()[0]] - frame = ControlValues("arm", 0.0, 0, 1, keys, [0.0] * len(keys)) + frame = control_values("arm", 0.0, 0, 1, keys, [0.0] * len(keys)) with pytest.raises(FrameRejectedError, match="duplicate_key"): validate_state(xarm, frame) @@ -401,7 +402,7 @@ def test_command_rejects_an_undeclared_key(xarm: ControlDescription) -> None: def test_command_rejects_duplicate_keys(xarm: ControlDescription) -> None: - frame = ControlValues( + frame = control_values( "c", 0.0, 0, 1, ["arm/joint1/position", "arm/joint1/position"], [0.1, 0.2] ) @@ -518,3 +519,19 @@ def test_a_base_may_be_six_dof() -> None: # someone will copy, so the units have to be physically right. assert all(drone.unit_of(f"drone/body/{a}") is Unit.M_PER_S for a in linear) assert all(drone.unit_of(f"drone/body/{a}") is Unit.RAD_PER_S for a in angular) + + +@pytest.mark.parametrize("bad", [float("nan"), float("inf"), float("-inf")]) +def test_direct_wire_values_reject_non_finite_input_at_both_boundaries( + xarm: ControlDescription, bad: float +) -> None: + state = full_state(xarm) + state.values[0] = bad + with pytest.raises(FrameRejectedError, match="non_finite"): + validate_state(xarm, state) + + frame = command("c", {"arm/joint1/position": 0.1}) + frame.values[0] = bad + result = validate_command(xarm, frame, last_sequence=None) + assert isinstance(result, Rejected) + assert result.reason == "non_finite" diff --git a/dimos/control/contract/validate.py b/dimos/control/contract/validate.py index b44da42e584..d8a1b5b04a4 100644 --- a/dimos/control/contract/validate.py +++ b/dimos/control/contract/validate.py @@ -38,10 +38,11 @@ from dataclasses import dataclass import math +from dimos_generated.dimos_msgs.msg import ControlValues + from dimos.control.contract.description import ControlDescription, ResourceKind from dimos.control.contract.keys import SEPARATOR, is_valid_segment from dimos.control.contract.sequence import is_newer -from dimos.msgs.control_msgs.ControlValues import ControlValues class DescriptionError(Exception): @@ -194,6 +195,11 @@ def validate_state(desc: ControlDescription, frame: ControlValues) -> dict[str, "shape", f"{len(frame.interface_names)} names against {len(frame.values)} values", ) + for name, value in zip(frame.interface_names, frame.values, strict=True): + if not math.isfinite(value): + raise FrameRejectedError( + "non_finite", f"non-finite value for interface {name!r}: {value!r}" + ) if frame.source != desc.source: raise FrameRejectedError( "source", f"frame is from {frame.source!r}, expected {desc.source!r}" @@ -243,13 +249,14 @@ def validate_command( Returns: The values to send, or the reason they were refused. """ - # ControlValues enforces equal lengths on construction, so this only fires - # for a frame built by hand or mutated after the fact. It is still checked - # here because this path is specified never to raise. + # Generated wire values are validated here before any command is applied. if len(frame.interface_names) != len(frame.values): return Rejected( "shape", f"{len(frame.interface_names)} names against {len(frame.values)} values" ) + for name, value in zip(frame.interface_names, frame.values, strict=True): + if not math.isfinite(value): + return Rejected("non_finite", f"non-finite value for interface {name!r}: {value!r}") if not is_newer(frame.sequence, last_sequence): return Rejected( "sequence", f"frame sequence {frame.sequence} not newer than {last_sequence}" diff --git a/dimos/control/coordinator.py b/dimos/control/coordinator.py index dcb7db3cc93..aab1f3533ba 100644 --- a/dimos/control/coordinator.py +++ b/dimos/control/coordinator.py @@ -32,6 +32,14 @@ import time from typing import TYPE_CHECKING, Any +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Twist +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Float32, Header +from dimos_generated.tf2_msgs.msg import TFMessage +from dimos_generated.trajectory_msgs.msg import JointTrajectory +import numpy as np + from dimos.control.components import ( TWIST_SUFFIX_MAP, HardwareComponent, @@ -64,11 +72,6 @@ ) from dimos.hardware.manipulators.spec import ManipulatorAdapter from dimos.hardware.whole_body.spec import WholeBodyAdapter -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.std_msgs.Float32 import Float32 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -709,7 +712,13 @@ def _map_twist_to_base_joints(self, msg: Twist) -> None: velocities.append(value) if names: - joint_state = JointState(name=names, velocity=velocities) + joint_state = JointState( + name=names, + velocity=np.asarray(velocities, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + position=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) self._dispatch("joint_command", joint_state) @rpc diff --git a/dimos/control/tasks/cartesian_ik_task/cartesian_ik_task.py b/dimos/control/tasks/cartesian_ik_task/cartesian_ik_task.py index 33784e372dc..a9f776ecd14 100644 --- a/dimos/control/tasks/cartesian_ik_task/cartesian_ik_task.py +++ b/dimos/control/tasks/cartesian_ik_task/cartesian_ik_task.py @@ -22,7 +22,9 @@ from typing import TYPE_CHECKING import attrs +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Pose, PoseStamped +from dimos_generated.std_msgs.msg import Header from dimos.control.task import CoordinatorState from dimos.control.tasks.pose_target_ik import ( @@ -72,7 +74,11 @@ def is_active(self) -> bool: def on_cartesian_command(self, pose: Pose | PoseStamped, t_now: float) -> bool: """Accept an absolute target pose and activate tracking.""" - target = copy.deepcopy(pose) if isinstance(pose, PoseStamped) else PoseStamped(pose=pose) + target = ( + copy.deepcopy(pose) + if isinstance(pose, PoseStamped) + else PoseStamped(pose=pose, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) + ) with self._lock: self._target_pose = target self._last_update_time = t_now diff --git a/dimos/control/tasks/cartesian_ik_task/test_cartesian_ik_task.py b/dimos/control/tasks/cartesian_ik_task/test_cartesian_ik_task.py index fdf21be6b34..b82d0b1643c 100644 --- a/dimos/control/tasks/cartesian_ik_task/test_cartesian_ik_task.py +++ b/dimos/control/tasks/cartesian_ik_task/test_cartesian_ik_task.py @@ -16,9 +16,11 @@ from pathlib import Path -from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from pytest_mock import MockerFixture @@ -55,10 +57,19 @@ def test_cartesian_leaf_maps_absolute_pose_to_configured_frame( mocker: MockerFixture, ) -> None: solver = mocker.Mock(spec=PinkPoseTargetSolver) - solver.step.return_value = JointState(name=["arm/joint"], position=[0.01]) + solver.step.return_value = JointState( + name=["arm/joint"], + position=np.array([0.01], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) task = CartesianIKTask("cartesian", _config(), solver=solver) target = PoseStamped( - header=Header(frame_id="world"), pose=Pose(position=Point(x=0.4, y=0.2, z=0.1)) + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.4, y=0.2, z=0.1), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), ) task.on_cartesian_command(target, t_now=2.0) @@ -77,7 +88,16 @@ def test_cartesian_leaf_maps_absolute_pose_to_configured_frame( def test_cartesian_leaf_clears_after_timeout(mocker: MockerFixture) -> None: solver = mocker.Mock(spec=PinkPoseTargetSolver) task = CartesianIKTask("cartesian", _config(), solver=solver) - task.on_cartesian_command(PoseStamped(header=Header(frame_id=""), pose=Pose()), t_now=1.0) + task.on_cartesian_command( + PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + t_now=1.0, + ) output = task.compute( CoordinatorState( @@ -94,19 +114,49 @@ def test_cartesian_leaf_clears_after_timeout(mocker: MockerFixture) -> None: def test_cartesian_clear_reseeds_command_from_feedback(mocker: MockerFixture) -> None: solver = mocker.Mock(spec=PinkPoseTargetSolver) - solver.step.return_value = JointState(name=["arm/joint"], position=[0.1]) + solver.step.return_value = JointState( + name=["arm/joint"], + position=np.array([0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) task = CartesianIKTask("cartesian", _config(), solver=solver) state = CoordinatorState( joints=JointStateSnapshot(joint_positions={"arm/joint": 0.0}), t_now=1.0, dt=0.01, ) - task.on_cartesian_command(PoseStamped(header=Header(frame_id=""), pose=Pose()), t_now=1.0) + task.on_cartesian_command( + PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + t_now=1.0, + ) assert task.compute(state) is not None task.clear() - task.on_cartesian_command(PoseStamped(header=Header(frame_id=""), pose=Pose()), t_now=1.1) - solver.step.return_value = JointState(name=["arm/joint"], position=[0.01]) + task.on_cartesian_command( + PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + t_now=1.1, + ) + solver.step.return_value = JointState( + name=["arm/joint"], + position=np.array([0.01], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) assert task.compute(state) is not None assert solver.reset.call_count == 1 diff --git a/dimos/control/tasks/eef_twist_task/test_eef_twist_task.py b/dimos/control/tasks/eef_twist_task/test_eef_twist_task.py index 7e69037f904..7cc246e7e0f 100644 --- a/dimos/control/tasks/eef_twist_task/test_eef_twist_task.py +++ b/dimos/control/tasks/eef_twist_task/test_eef_twist_task.py @@ -14,7 +14,16 @@ from pathlib import Path -from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Twist, TwistStamped, Vector3 +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Twist, + TwistStamped, + Vector3, +) from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header import numpy as np @@ -42,12 +51,19 @@ def _solver(mocker: MockerFixture) -> PinkPoseTargetSolver: solver = mocker.Mock(spec=PinkPoseTargetSolver) solver.frame_poses.return_value = { "tool": PoseStamped( - header=Header(frame_id="world"), pose=Pose(position=Point(x=0.0, y=0.0, z=0.0)) + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) } solver.step.return_value = JointState( name=["arm/joint1", "arm/joint2"], - position=[0.01, 0.02], + position=np.array([0.01, 0.02], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) return solver @@ -94,7 +110,10 @@ def _state(t_now: float = 1.0, *, dt: float = 0.01) -> CoordinatorState: def _twist(x: float = 0.1) -> TwistStamped: - return TwistStamped(header=Header(frame_id="tool"), twist=Twist(linear=Vector3(x=x))) + return TwistStamped( + header=Header(frame_id="tool", stamp=Time(sec=0, nanosec=0)), + twist=Twist(linear=Vector3(x=x, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + ) def test_twist_integrates_target_and_uses_shared_persistent_command( diff --git a/dimos/control/tasks/holonomic_pose_follower_task/holonomic_pose_follower_task.py b/dimos/control/tasks/holonomic_pose_follower_task/holonomic_pose_follower_task.py index e953e0ff094..3ebdec81817 100644 --- a/dimos/control/tasks/holonomic_pose_follower_task/holonomic_pose_follower_task.py +++ b/dimos/control/tasks/holonomic_pose_follower_task/holonomic_pose_follower_task.py @@ -26,6 +26,7 @@ from pathlib import Path as _FsPath from typing import Any, Literal +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped from dimos_generated.nav_msgs.msg import Path from dimos_generated.std_msgs.msg import Float32, Header @@ -79,7 +80,7 @@ def _clamp(v: float, limit: float) -> float: def _pose_stamped(pose: tuple[float, float, float]) -> PoseStamped: return PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=Point(x=pose[0], y=pose[1], z=0.0), orientation=quaternion_from_euler(0.0, 0.0, pose[2]), diff --git a/dimos/control/tasks/holonomic_pose_follower_task/test_holonomic_pose_follower_task.py b/dimos/control/tasks/holonomic_pose_follower_task/test_holonomic_pose_follower_task.py index 1f8d17f1bfa..6d6e72c22ce 100644 --- a/dimos/control/tasks/holonomic_pose_follower_task/test_holonomic_pose_follower_task.py +++ b/dimos/control/tasks/holonomic_pose_follower_task/test_holonomic_pose_follower_task.py @@ -21,6 +21,7 @@ import math from types import SimpleNamespace +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Twist, Vector3 from dimos_generated.nav_msgs.msg import Path from dimos_generated.std_msgs.msg import Float32, Header @@ -54,7 +55,7 @@ def _pose(x=0.0, y=0.0, yaw=0.0): return PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=Point(x=x, y=y, z=0.0), orientation=quaternion_from_euler(0.0, 0.0, yaw) ), @@ -109,7 +110,9 @@ def _run_closed_loop(task, path, max_ticks=2400, on_tick=None): cmd_twist=Twist( linear=Vector3(x=vx, y=vy, z=0.0), angular=Vector3(x=0.0, y=0.0, z=wz) ), - actual_twist=Twist(), + actual_twist=Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ), ) ) if on_tick is not None: @@ -167,7 +170,10 @@ def test_set_speed_refused_while_active(): def test_rejects_pure_rotation_path(): task = _task() - path = Path(poses=[_pose(0, 0, 0.0), _pose(0, 0, 1.0)]) + path = Path( + poses=[_pose(0, 0, 0.0), _pose(0, 0, 1.0)], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) assert not task.start_path(path, _pose()) @@ -307,7 +313,10 @@ def test_replan_reprojects_without_reramp(): assert speed_before > 0.3 # Replan: a new strafe path starting at the robot (as a planner would emit). - path_b = Path(poses=[_pose(plant.x, plant.y + 2.0 * i / 40, 0.0) for i in range(41)]) + path_b = Path( + poses=[_pose(plant.x, plant.y + 2.0 * i / 40, 0.0) for i in range(41)], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) task.start_path(path_b, _pose(plant.x, plant.y, plant.yaw)) speeds = [] for k in range(30, 40): @@ -355,7 +364,7 @@ def test_speed_regulator_slows_for_fast_yaw_stretch(): poses = [_pose(i * 0.1, 0.0, 0.0) for i in range(11)] poses += [_pose(1.0 + i * 0.1, 0.0, (math.pi / 2) * i / 6) for i in range(1, 7)] poses += [_pose(1.6 + i * 0.1, 0.0, math.pi / 2) for i in range(1, 11)] - path = Path(poses=poses) + path = Path(poses=poses, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) task = _task(speed=0.8) v_by_x: dict[float, float] = {} diff --git a/dimos/control/tasks/holonomic_pose_follower_task/test_progress_reference.py b/dimos/control/tasks/holonomic_pose_follower_task/test_progress_reference.py index d23fca46d50..b6365d79c26 100644 --- a/dimos/control/tasks/holonomic_pose_follower_task/test_progress_reference.py +++ b/dimos/control/tasks/holonomic_pose_follower_task/test_progress_reference.py @@ -19,6 +19,7 @@ import math +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped from dimos_generated.nav_msgs.msg import Path from dimos_generated.std_msgs.msg import Header @@ -32,7 +33,7 @@ def _pose(x, y, yaw): return PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=Point(x=x, y=y, z=0.0), orientation=quaternion_from_euler(0.0, 0.0, yaw) ), @@ -40,7 +41,10 @@ def _pose(x, y, yaw): def _straight_rotate(length=2.0, yaw_end=math.pi / 2, n=40): - return Path(poses=[_pose(length * i / n, 0.0, yaw_end * i / n) for i in range(n + 1)]) + return Path( + poses=[_pose(length * i / n, 0.0, yaw_end * i / n) for i in range(n + 1)], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) def _circle(radius=1.0, n=100): @@ -48,14 +52,21 @@ def _circle(radius=1.0, n=100): for i in range(n + 1): th = 2 * math.pi * i / n poses.append(_pose(radius * math.sin(th), radius * (1 - math.cos(th)), 0.0)) - return Path(poses=poses) + return Path(poses=poses, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) def test_rejects_degenerate_paths(): with pytest.raises(ValueError): - ProgressPathReference(Path(poses=[_pose(0, 0, 0)])) + ProgressPathReference( + Path(poses=[_pose(0, 0, 0)], header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) + ) with pytest.raises(ValueError): - ProgressPathReference(Path(poses=[_pose(0, 0, 0), _pose(0, 0, 1.0)])) + ProgressPathReference( + Path( + poses=[_pose(0, 0, 0), _pose(0, 0, 1.0)], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + ) def test_arc_length_and_end_pose(): @@ -86,7 +97,10 @@ def test_yaw_interpolation_across_pi_wrap(): # Commanded yaw sweeps 170deg -> 190deg (i.e. -170deg): interpolation must # go THROUGH pi, not backwards through 0. a, b = math.radians(170), math.radians(190) - path = Path(poses=[_pose(i * 0.1, 0.0, a + (b - a) * i / 10) for i in range(11)]) + path = Path( + poses=[_pose(i * 0.1, 0.0, a + (b - a) * i / 10) for i in range(11)], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) ref = ProgressPathReference(path) mid = ref.sample(ref.length / 2) assert abs(mid.yaw) == pytest.approx(math.pi, abs=1e-6) @@ -132,7 +146,10 @@ def test_advance_allows_bounded_backslide(): def test_coincident_waypoints_keep_later_yaw(): # Duplicate position with a different yaw = stop-and-rotate hint; the # later yaw survives as the target at that station. - path = Path(poses=[_pose(0, 0, 0), _pose(1.0, 0, 0), _pose(1.0, 0, 1.0), _pose(2.0, 0, 1.0)]) + path = Path( + poses=[_pose(0, 0, 0), _pose(1.0, 0, 0), _pose(1.0, 0, 1.0), _pose(2.0, 0, 1.0)], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) ref = ProgressPathReference(path) assert ref.length == pytest.approx(2.0) assert ref.sample(1.0).yaw == pytest.approx(1.0) @@ -142,7 +159,9 @@ def test_max_rates_ahead_sees_upcoming_yaw_demand(): # Flat commanded yaw for 1 m, then a fast 90deg twist over 0.5 m. poses = [_pose(i * 0.1, 0.0, 0.0) for i in range(11)] poses += [_pose(1.0 + i * 0.1, 0.0, (math.pi / 2) * i / 5) for i in range(1, 6)] - ref = ProgressPathReference(Path(poses=poses)) + ref = ProgressPathReference( + Path(poses=poses, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) + ) dyaw_flat, _ = ref.max_rates_ahead(0.0, 0.5) dyaw_ahead, _ = ref.max_rates_ahead(0.6, 0.5) assert dyaw_flat == pytest.approx(0.0, abs=1e-9) diff --git a/dimos/control/tasks/path_follower_task/path_follower_task.py b/dimos/control/tasks/path_follower_task/path_follower_task.py index 7bbe432f43a..a9b1bd4f3ff 100644 --- a/dimos/control/tasks/path_follower_task/path_follower_task.py +++ b/dimos/control/tasks/path_follower_task/path_follower_task.py @@ -29,6 +29,9 @@ from dataclasses import dataclass, field, replace from typing import TYPE_CHECKING, Any, Literal +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Float32, Header import numpy as np from dimos.control.benchmarking.velocity_profile import ( @@ -52,11 +55,8 @@ VelocityTrackingPID, ) from dimos.core.global_config import global_config as _gc -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.std_msgs.Float32 import Float32 +from dimos.msgs.geometry import quaternion_euler, quaternion_from_euler +from dimos.msgs.time import time_from_seconds from dimos.navigation.go2.replanning_a_star.controllers import PController from dimos.navigation.go2.replanning_a_star.path_distancer import PathDistancer from dimos.protocol.service.spec import BaseConfig @@ -228,9 +228,11 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None: self.start_path( path, PoseStamped( - ts=state.t_now, - position=Vector3(float(px), float(py), 0.0), - orientation=Quaternion.from_euler(Vector3(0.0, 0.0, float(pyaw))), + header=Header(frame_id="", stamp=time_from_seconds(state.t_now)), + pose=Pose( + position=Point(x=float(px), y=float(py), z=0.0), + orientation=quaternion_from_euler(0.0, 0.0, float(pyaw)), + ), ), ) if not self._running(): @@ -246,9 +248,11 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None: yaw = pos.get(self._joint_names_list[2]) if x is not None and y is not None and yaw is not None: self._current_odom = PoseStamped( - ts=state.t_now, - position=Vector3(float(x), float(y), 0.0), - orientation=Quaternion.from_euler(Vector3(0.0, 0.0, float(yaw))), + header=Header(frame_id="", stamp=time_from_seconds(state.t_now)), + pose=Pose( + position=Point(x=float(x), y=float(y), z=0.0), + orientation=quaternion_from_euler(0.0, 0.0, float(yaw)), + ), ) if self._current_odom is None: return None @@ -283,8 +287,8 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None: # PathSpeedCap.speed_limit_at — min over the next ~8 waypoints so # braking starts BEFORE a corner rather than at it. if self._velocity_profile is not None and self._velocity_profile_pts is not None: - x = self._current_odom.position.x - y = self._current_odom.position.y + x = self._current_odom.pose.position.x + y = self._current_odom.pose.position.y i = int(np.argmin(np.sum((self._velocity_profile_pts - np.array([x, y])) ** 2, axis=1))) j = min(len(self._velocity_profile), i + 8) vlim = float(np.min(self._velocity_profile[i:j])) @@ -294,7 +298,7 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None: vx, vy, wz = vx * k, vy * k, wz * k elif self._profile_cap is not None: vx, vy, wz = self._profile_cap.cap( - self._current_odom.position.x, self._current_odom.position.y, vx, vy, wz + self._current_odom.pose.position.x, self._current_odom.pose.position.y, vx, vy, wz ) # Inner-loop gain compensation (mutually exclusive - PI wins if both @@ -340,8 +344,8 @@ def on_preempted(self, by_task: str, joints: frozenset[str]) -> None: def _step_initial_rotation(self) -> tuple[float, float, float]: assert self._path is not None and self._current_odom is not None - first_yaw = self._path.poses[0].orientation.euler[2] - robot_yaw = self._current_odom.orientation.euler[2] + first_yaw = quaternion_euler(self._path.poses[0].pose.orientation)[2] + robot_yaw = quaternion_euler(self._current_odom.pose.orientation)[2] yaw_err = angle_diff(first_yaw, robot_yaw) if abs(yaw_err) < self._config.orientation_tolerance: @@ -364,7 +368,7 @@ def _windowed_closest(self, pos: np.ndarray, window: int = 20) -> int: best_idx = lo best_d_sq = float("inf") for i in range(lo, hi): - p = self._path.poses[i].position + p = self._path.poses[i].pose.position d_sq = (p.x - pos[0]) ** 2 + (p.y - pos[1]) ** 2 if d_sq < best_d_sq: best_d_sq = d_sq @@ -376,7 +380,7 @@ def _step_path_following(self) -> tuple[float, float, float]: assert self._distancer is not None assert self._current_odom is not None - pos = np.array([self._current_odom.position.x, self._current_odom.position.y]) + pos = np.array([self._current_odom.pose.position.x, self._current_odom.pose.position.y]) closest = self._windowed_closest(pos) if closest > self._max_progress_idx: @@ -410,8 +414,8 @@ def _step_path_following(self) -> tuple[float, float, float]: def _step_final_rotation(self) -> tuple[float, float, float]: assert self._path is not None and self._current_odom is not None - goal_yaw = self._path.poses[-1].orientation.euler[2] - robot_yaw = self._current_odom.orientation.euler[2] + goal_yaw = quaternion_euler(self._path.poses[-1].pose.orientation)[2] + robot_yaw = quaternion_euler(self._current_odom.pose.orientation)[2] yaw_err = angle_diff(goal_yaw, robot_yaw) if abs(yaw_err) < self._config.orientation_tolerance: @@ -529,8 +533,8 @@ def start_path( self._velocity_profile = None self._velocity_profile_pts = None - first_yaw = path.poses[0].orientation.euler[2] - robot_yaw = current_odom.orientation.euler[2] + first_yaw = quaternion_euler(path.poses[0].pose.orientation)[2] + robot_yaw = quaternion_euler(current_odom.pose.orientation)[2] yaw_err = angle_diff(first_yaw, robot_yaw) self._controller.reset_yaw_error(yaw_err) diff --git a/dimos/control/tasks/path_follower_task/test_path_follower_task.py b/dimos/control/tasks/path_follower_task/test_path_follower_task.py index 8e21d67dee7..9c0a9122744 100644 --- a/dimos/control/tasks/path_follower_task/test_path_follower_task.py +++ b/dimos/control/tasks/path_follower_task/test_path_follower_task.py @@ -20,6 +20,7 @@ import math +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped from dimos_generated.std_msgs.msg import Header @@ -55,7 +56,7 @@ def _task(**overrides) -> PathFollowerTask: def _start_aligned(task: PathFollowerTask, path) -> None: odom = PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=Point(x=path.poses[0].pose.position.x, y=path.poses[0].pose.position.y, z=0.0), orientation=quaternion_from_euler( diff --git a/dimos/control/tasks/planar_base_trajectory_task/test_planar_base_trajectory_task.py b/dimos/control/tasks/planar_base_trajectory_task/test_planar_base_trajectory_task.py index 4ac516a1b47..57960c59a25 100644 --- a/dimos/control/tasks/planar_base_trajectory_task/test_planar_base_trajectory_task.py +++ b/dimos/control/tasks/planar_base_trajectory_task/test_planar_base_trajectory_task.py @@ -14,7 +14,9 @@ import math +from dimos_generated.builtin_interfaces.msg import Duration from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import numpy as np import pytest from dimos.control.task import CoordinatorState, JointStateSnapshot @@ -60,8 +62,13 @@ def _line(start, goal, duration, n=60): points = [ JointTrajectoryPoint( time_from_start=duration_from_seconds(duration * i / (n - 1)), - positions=[s + (g - s) * i / (n - 1) for s, g in zip(start, goal, strict=True)], - velocities=velocity, + positions=np.asarray( + [s + (g - s) * i / (n - 1) for s, g in zip(start, goal, strict=True)], + dtype=np.float64, + ), + velocities=np.asarray(velocity, dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) for i in range(n) ] @@ -159,7 +166,16 @@ def test_rejects_trajectories_it_cannot_follow(): never_ends.points[-1] = point two_columns = JointTrajectory( header=header_now(), - points=[JointTrajectoryPoint(positions=[0.0, 0.0], velocities=[0.0, 0.0])], + points=[ + JointTrajectoryPoint( + positions=np.array([0.0, 0.0], dtype=np.float64), + velocities=np.array([0.0, 0.0], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + time_from_start=Duration(sec=0, nanosec=0), + ) + ], + joint_names=[], ) diagonal_too_fast = _line((0.0, 0.0, 0.0), (1.0, 1.0, 0.0), duration=1.0) diff --git a/dimos/control/tasks/pose_target_ik.py b/dimos/control/tasks/pose_target_ik.py index f7d9ee52698..5fa31fdaf51 100644 --- a/dimos/control/tasks/pose_target_ik.py +++ b/dimos/control/tasks/pose_target_ik.py @@ -27,6 +27,7 @@ from typing import TYPE_CHECKING import attrs +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Pose, PoseStamped from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header @@ -443,7 +444,13 @@ def _step_frame_targets( command_limit_margin=command_limit_margin, ) return _StreamingStepResult( - command=JointState(name=list(joint_names), position=command_positions.tolist()), + command=JointState( + name=list(joint_names), + position=np.asarray(command_positions.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), bounded_increment=bounded_increment, ) @@ -515,7 +522,7 @@ def _frame_poses( for frame_name, frame_context in context.frames.items(): pose = pose_from_matrix(base_world @ self._current_frame_matrix(frame_context, q)) poses[frame_name] = PoseStamped( - header=Header(frame_id=robot_model.base_link), + header=Header(frame_id=robot_model.base_link, stamp=Time(sec=0, nanosec=0)), pose=Pose(position=pose.position, orientation=pose.orientation), ) return poses @@ -783,7 +790,13 @@ def _measured_joint_state(self, state: CoordinatorState) -> JointState | None: if position is None: return None positions.append(position) - return JointState(name=list(self._joint_names), position=positions) + return JointState( + name=list(self._joint_names), + position=np.asarray(positions, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) def _reset_command_state(self) -> None: """Discard the active command trajectory and clear a tracking fault.""" diff --git a/dimos/control/tasks/rpp_path_follower_task/test_rpp_path_follower_task.py b/dimos/control/tasks/rpp_path_follower_task/test_rpp_path_follower_task.py index d65870a14fc..d503ec1ce72 100644 --- a/dimos/control/tasks/rpp_path_follower_task/test_rpp_path_follower_task.py +++ b/dimos/control/tasks/rpp_path_follower_task/test_rpp_path_follower_task.py @@ -19,6 +19,7 @@ from types import SimpleNamespace +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.std_msgs.msg import Header import pytest @@ -84,7 +85,7 @@ def _task(artifact_path: str, **params) -> RPPPathFollowerTask: def _odom(x=0.0, y=0.0, yaw=0.0) -> PoseStamped: return PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=Point(x=x, y=y, z=0.0), orientation=quaternion_from_euler(0.0, 0.0, yaw) ), @@ -246,7 +247,7 @@ def test_drives_to_arrival_forward_only(artifact_path): def _ident_path(points): """Path with identity orientation on every pose (what MLSPlannerNative emits).""" return Path( - header=Header(frame_id="odom"), + header=Header(frame_id="odom", stamp=Time(sec=0, nanosec=0)), poses=[ PoseStamped( header=Header(frame_id="", stamp=time_from_seconds(0.0)), @@ -281,7 +282,7 @@ def test_positions_unchanged_by_synthesis(): def test_path_with_real_headings_is_untouched(): # A planner that already provides per-pose orientations must be left alone. oriented = Path( - header=Header(frame_id="odom"), + header=Header(frame_id="odom", stamp=Time(sec=0, nanosec=0)), poses=[ PoseStamped( header=Header(frame_id="", stamp=time_from_seconds(0.0)), diff --git a/dimos/control/tasks/teleop_ik_task/teleop_ik_task.py b/dimos/control/tasks/teleop_ik_task/teleop_ik_task.py index 752a2ca78cf..6ee4548488b 100644 --- a/dimos/control/tasks/teleop_ik_task/teleop_ik_task.py +++ b/dimos/control/tasks/teleop_ik_task/teleop_ik_task.py @@ -24,6 +24,7 @@ from typing import TYPE_CHECKING import attrs +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Pose, PoseStamped from dimos_generated.std_msgs.msg import Header, UInt32 @@ -175,7 +176,11 @@ def _on_controller_pose( ) -> bool: if hand not in self._bindings: return False - sample = copy.deepcopy(pose) if isinstance(pose, PoseStamped) else PoseStamped(pose=pose) + sample = ( + copy.deepcopy(pose) + if isinstance(pose, PoseStamped) + else PoseStamped(pose=pose, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) + ) with self._lock: if self._session_state is _SessionState.ESTOPPED: return False diff --git a/dimos/control/tasks/teleop_ik_task/test_teleop_ik_task.py b/dimos/control/tasks/teleop_ik_task/test_teleop_ik_task.py index a7bf6309e41..081ffe65d9a 100644 --- a/dimos/control/tasks/teleop_ik_task/test_teleop_ik_task.py +++ b/dimos/control/tasks/teleop_ik_task/test_teleop_ik_task.py @@ -17,9 +17,11 @@ from pathlib import Path from typing import cast +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header, UInt32 +import numpy as np import pytest from pytest_mock import MockerFixture @@ -68,14 +70,14 @@ def _solver(mocker: MockerFixture) -> PinkPoseTargetSolver: solver = mocker.Mock(spec=PinkPoseTargetSolver) solver.frame_poses.return_value = { "left_tool": PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=Point(x=1.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), ), ), "right_tool": PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=Point(x=-1.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), @@ -83,7 +85,11 @@ def _solver(mocker: MockerFixture) -> PinkPoseTargetSolver: ), } solver.step.return_value = JointState( - name=["robot/left", "robot/right"], position=[0.01, -0.01] + name=["robot/left", "robot/right"], + position=np.array([0.01, -0.01], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) return solver @@ -124,7 +130,7 @@ def test_face_buttons_do_not_engage_arm_teleop(mocker: MockerFixture) -> None: def _pose(x: float) -> PoseStamped: return PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=Point(x=x, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) ), @@ -230,14 +236,24 @@ def test_deadman_reengagement_reseeds_command_from_feedback( ) task.on_teleop_buttons(_buttons(left=True), 1.0) task.on_left_cartesian_command(_pose(0.1), 1.0) - solver.step.return_value = JointState(name=["robot/left", "robot/right"], position=[0.1, -0.1]) + solver.step.return_value = JointState( + name=["robot/left", "robot/right"], + position=np.array([0.1, -0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) assert task.compute(_state()) is not None task.on_teleop_buttons(_buttons(), 1.1) task.on_teleop_buttons(_buttons(left=True), 1.2) task.on_left_cartesian_command(_pose(0.2), 1.2) solver.step.return_value = JointState( - name=["robot/left", "robot/right"], position=[0.01, -0.01] + name=["robot/left", "robot/right"], + position=np.array([0.01, -0.01], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) assert task.compute(_state(1.2)) is not None diff --git a/dimos/control/tasks/test_pose_target_ik.py b/dimos/control/tasks/test_pose_target_ik.py index 3ea70f586ce..5eb05cac30f 100644 --- a/dimos/control/tasks/test_pose_target_ik.py +++ b/dimos/control/tasks/test_pose_target_ik.py @@ -16,7 +16,8 @@ from pathlib import Path -from dimos_generated.geometry_msgs.msg import Pose, PoseStamped +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header import numpy as np @@ -96,10 +97,20 @@ def _config( def _solver(mocker: MockerFixture, positions: list[float] | None = None) -> PinkPoseTargetSolver: solver = mocker.Mock(spec=PinkPoseTargetSolver) solver.step.return_value = JointState( - name=["arm/a", "arm/b"], position=positions or [0.01, -0.01] + name=["arm/a", "arm/b"], + position=np.asarray(positions or [0.01, -0.01], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) solver.frame_poses.return_value = { - "tool": PoseStamped(header=Header(frame_id="base"), pose=Pose()) + "tool": PoseStamped( + header=Header(frame_id="base", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) } return solver @@ -119,7 +130,16 @@ def _snapshot( extra_joint_positions: dict[str, float] | None = None, ) -> FrameTargetSnapshot: return FrameTargetSnapshot( - targets=targets or {"tool": PoseStamped(header=Header(frame_id="world"), pose=Pose())}, + targets=targets + or { + "tool": PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + }, last_update_time=last_update_time, extra_joint_positions=extra_joint_positions or {}, ) @@ -142,24 +162,59 @@ def advance(**kwargs: object) -> _StreamingStepResult: return _StreamingStepResult( command=JointState( name=list(command.name), - position=[position + 0.1 for position in command.position], + position=np.asarray( + [position + 0.1 for position in command.position], dtype=np.float64 + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), bounded_increment=np.array([0.1, 0.1]), ) step = mocker.patch.object(solver, "_step_frame_targets", side_effect=advance) - targets = {"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())} + targets = { + "tool": PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + } - assert solver.step(targets, JointState(name=["arm/a", "arm/b"], position=[0.0, 0.0]), 0.01) assert solver.step( targets, - JointState(name=["arm/a", "arm/b"], position=[-0.3, -0.3]), + JointState( + name=["arm/a", "arm/b"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + 0.01, + ) + assert solver.step( + targets, + JointState( + name=["arm/a", "arm/b"], + position=np.array([-0.3, -0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), 0.01, ) - assert step.call_args_list[0].kwargs["command_state"].position == [0.0, 0.0] - assert step.call_args_list[1].kwargs["command_state"].position == [0.1, 0.1] - assert step.call_args_list[1].kwargs["measured_state"].position == [-0.3, -0.3] + np.testing.assert_array_equal( + step.call_args_list[0].kwargs["command_state"].position, [0.0, 0.0] + ) + np.testing.assert_array_equal( + step.call_args_list[1].kwargs["command_state"].position, [0.1, 0.1] + ) + np.testing.assert_array_equal( + step.call_args_list[1].kwargs["measured_state"].position, [-0.3, -0.3] + ) assert step.call_args_list[1].kwargs["command_increment_history"][0] == pytest.approx( [0.1, 0.1] ) @@ -174,26 +229,64 @@ def test_pose_target_solver_reset_reseeds_from_feedback(mocker: MockerFixture) - "_step_frame_targets", side_effect=[ _StreamingStepResult( - command=JointState(name=["arm/a", "arm/b"], position=[0.1, 0.1]), + command=JointState( + name=["arm/a", "arm/b"], + position=np.array([0.1, 0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), bounded_increment=np.array([0.1, 0.1]), ), _StreamingStepResult( - command=JointState(name=["arm/a", "arm/b"], position=[-0.2, -0.2]), + command=JointState( + name=["arm/a", "arm/b"], + position=np.array([-0.2, -0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), bounded_increment=np.array([0.1, 0.1]), ), ], ) - targets = {"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())} + targets = { + "tool": PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + } - assert solver.step(targets, JointState(name=["arm/a", "arm/b"], position=[0.0, 0.0]), 0.01) + assert solver.step( + targets, + JointState( + name=["arm/a", "arm/b"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + 0.01, + ) solver.reset() assert solver.step( targets, - JointState(name=["arm/a", "arm/b"], position=[-0.3, -0.3]), + JointState( + name=["arm/a", "arm/b"], + position=np.array([-0.3, -0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), 0.01, ) - assert step.call_args_list[1].kwargs["command_state"].position == [-0.3, -0.3] + np.testing.assert_array_equal( + step.call_args_list[1].kwargs["command_state"].position, [-0.3, -0.3] + ) assert step.call_args_list[1].kwargs["command_increment_history"] == () @@ -205,15 +298,35 @@ def test_pose_target_solver_reset_during_step_discards_command_and_filter_histor def reset_during_step(**_kwargs: object) -> _StreamingStepResult: solver.reset() return _StreamingStepResult( - command=JointState(name=["arm/a", "arm/b"], position=[0.1, 0.1]), + command=JointState( + name=["arm/a", "arm/b"], + position=np.array([0.1, 0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), bounded_increment=np.array([0.1, 0.1]), ) mocker.patch.object(solver, "_step_frame_targets", side_effect=reset_during_step) result = solver.step( - {"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())}, - JointState(name=["arm/a", "arm/b"], position=[0.0, 0.0]), + { + "tool": PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + }, + JointState( + name=["arm/a", "arm/b"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), 0.01, ) @@ -256,7 +369,7 @@ def test_compute_calls_one_pink_step_and_preserves_output_order( assert output is not None assert output.joint_names == ["arm/a", "arm/b", "arm/gripper"] - assert output.positions == [0.02, -0.03, 0.4] + np.testing.assert_array_equal(output.positions, [0.02, -0.03, 0.4]) solver.step.assert_called_once() assert solver.step.call_args.args[2] == 0.01 assert task.claim().joints == frozenset({"arm/a", "arm/b", "arm/gripper"}) @@ -365,10 +478,18 @@ def test_current_frame_poses_uses_live_coordinator_seed(mocker: MockerFixture) - poses = task.current_frame_poses(_state(positions={"arm/a": 0.2, "arm/b": 0.3}), ["tool"]) - assert poses == {"tool": PoseStamped(header=Header(frame_id="base"), pose=Pose())} + assert poses == { + "tool": PoseStamped( + header=Header(frame_id="base", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + } seed = solver.frame_poses.call_args.args[0] assert seed.name == ["arm/a", "arm/b"] - assert seed.position == [0.2, 0.3] + np.testing.assert_array_equal(seed.position, [0.2, 0.3]) def test_measured_frame_poses_follow_measured_joints_not_the_target(mocker: MockerFixture) -> None: diff --git a/dimos/control/tasks/trajectory_task/test_trajectory_task.py b/dimos/control/tasks/trajectory_task/test_trajectory_task.py index cf126d5e73f..a111f647574 100644 --- a/dimos/control/tasks/trajectory_task/test_trajectory_task.py +++ b/dimos/control/tasks/trajectory_task/test_trajectory_task.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos_generated.builtin_interfaces.msg import Duration +from dimos_generated.builtin_interfaces.msg import Duration, Time +from dimos_generated.std_msgs.msg import Header from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import numpy as np import pytest from dimos.control.task import ( @@ -39,7 +41,15 @@ def test_jtt_does_not_pull_joint_back_after_another_task_moves_it(single_point, first = JointTrajectory( header=header_now(), joint_names=["joint"], - points=[JointTrajectoryPoint(positions=[first_target])], + points=[ + JointTrajectoryPoint( + positions=np.array([first_target], dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + time_from_start=Duration(sec=0, nanosec=0), + ) + ], ) assert task.execute(first, {}).status is TrajectoryExecutionStatus.ACCEPTED output = task.compute(state) @@ -54,11 +64,31 @@ def test_jtt_does_not_pull_joint_back_after_another_task_moves_it(single_point, joints=JointStateSnapshot(joint_positions={"joint": -1.0}), t_now=1.0, dt=0.01 ) assert task.compute(state) is None - points = [JointTrajectoryPoint(positions=[-1.2])] + points = [ + JointTrajectoryPoint( + positions=np.array([-1.2], dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + time_from_start=Duration(sec=0, nanosec=0), + ) + ] if not single_point: points = [ - JointTrajectoryPoint(positions=[-1.0]), - JointTrajectoryPoint(positions=[-1.2], time_from_start=duration_from_seconds(0.1)), + JointTrajectoryPoint( + positions=np.array([-1.0], dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + time_from_start=Duration(sec=0, nanosec=0), + ), + JointTrajectoryPoint( + positions=np.array([-1.2], dtype=np.float64), + time_from_start=duration_from_seconds(0.1), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ] assert ( task.execute( @@ -90,7 +120,15 @@ def test_completed_joint_reanchors_while_other_joint_keeps_command_continuity(): initial = JointTrajectory( header=header_now(), joint_names=["finished", "running"], - points=[JointTrajectoryPoint(positions=[0.1, 1.0])], + points=[ + JointTrajectoryPoint( + positions=np.array([0.1, 1.0], dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + time_from_start=Duration(sec=0, nanosec=0), + ) + ], ) assert task.execute(initial, {}).status is TrajectoryExecutionStatus.ACCEPTED output = task.compute(CoordinatorState(joints=measured, t_now=0.1, dt=0.1)) @@ -101,8 +139,20 @@ def test_completed_joint_reanchors_while_other_joint_keeps_command_continuity(): header=header_now(), joint_names=["finished"], points=[ - JointTrajectoryPoint(positions=[-1.0]), - JointTrajectoryPoint(positions=[0.1], time_from_start=duration_from_seconds(2.0)), + JointTrajectoryPoint( + positions=np.array([-1.0], dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + time_from_start=Duration(sec=0, nanosec=0), + ), + JointTrajectoryPoint( + positions=np.array([0.1], dtype=np.float64), + time_from_start=duration_from_seconds(2.0), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ], ) assert ( @@ -124,7 +174,17 @@ def test_completed_joint_reanchors_while_other_joint_keeps_command_continuity(): def test_completed_trajectory_does_not_bypass_start_validation(positions, expected): task = JointTrajectoryTask(JointTrajectoryTaskConfig(joint_names=["joint"])) initial = JointTrajectory( - header=header_now(), joint_names=["joint"], points=[JointTrajectoryPoint(positions=[0.1])] + header=header_now(), + joint_names=["joint"], + points=[ + JointTrajectoryPoint( + positions=np.array([0.1], dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + time_from_start=Duration(sec=0, nanosec=0), + ) + ], ) assert task.execute(initial, {}).status is TrajectoryExecutionStatus.ACCEPTED task.compute( @@ -136,8 +196,20 @@ def test_completed_trajectory_does_not_bypass_start_validation(positions, expect header=header_now(), joint_names=["joint"], points=[ - JointTrajectoryPoint(positions=[0.1]), - JointTrajectoryPoint(positions=[1.0], time_from_start=duration_from_seconds(2.0)), + JointTrajectoryPoint( + positions=np.array([0.1], dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + time_from_start=Duration(sec=0, nanosec=0), + ), + JointTrajectoryPoint( + positions=np.array([1.0], dtype=np.float64), + time_from_start=duration_from_seconds(2.0), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ], ) assert task.execute(trajectory, positions).status is expected @@ -150,14 +222,29 @@ def test_jtt_validates_normalized_durations_without_float_ordering(nanosec, acce trajectory = JointTrajectory( joint_names=["joint"], points=[ - JointTrajectoryPoint(positions=[0.0]), JointTrajectoryPoint( - positions=[0.1], time_from_start=Duration(sec=100_000_000, nanosec=nanosec) + positions=np.array([0.0], dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + time_from_start=Duration(sec=0, nanosec=0), + ), + JointTrajectoryPoint( + positions=np.array([0.1], dtype=np.float64), + time_from_start=Duration(sec=100_000_000, nanosec=nanosec), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), JointTrajectoryPoint( - positions=[0.2], time_from_start=Duration(sec=100_000_000, nanosec=1) + positions=np.array([0.2], dtype=np.float64), + time_from_start=Duration(sec=100_000_000, nanosec=1), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) result = task.execute(trajectory, {"joint": 0.0}) expected = ( diff --git a/dimos/control/tasks/trajectory_task/trajectory_task.py b/dimos/control/tasks/trajectory_task/trajectory_task.py index 50697d2da2a..821b07e23bd 100644 --- a/dimos/control/tasks/trajectory_task/trajectory_task.py +++ b/dimos/control/tasks/trajectory_task/trajectory_task.py @@ -28,9 +28,11 @@ import math from typing import TYPE_CHECKING, Annotated, Any +from dimos_generated.builtin_interfaces.msg import Duration from dimos_generated.dimos_msgs.msg import TrajectoryStatus from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import numpy as np from pydantic import BeforeValidator, ConfigDict, Field from pydantic.dataclasses import dataclass as pydantic_dataclass @@ -232,7 +234,7 @@ def claim(self) -> ResourceClaim: def on_joint_command(self, msg: JointState, t_now: float) -> bool: """Convert streamed joint positions into a velocity-bounded trajectory target.""" - if not msg.position or len(msg.name) != len(msg.position): + if not len(msg.position) or len(msg.name) != len(msg.position): return False selected = [ (name, position) @@ -244,7 +246,15 @@ def on_joint_command(self, msg: JointState, t_now: float) -> bool: trajectory = JointTrajectory( header=header_now(), joint_names=[name for name, _ in selected], - points=[JointTrajectoryPoint(positions=[position for _, position in selected])], + points=[ + JointTrajectoryPoint( + positions=np.asarray([position for _, position in selected], dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + time_from_start=Duration(sec=0, nanosec=0), + ) + ], ) return self.execute(trajectory, {}).status is TrajectoryExecutionStatus.ACCEPTED @@ -454,10 +464,12 @@ def execute( points=[ JointTrajectoryPoint( time_from_start=trajectory.points[0].time_from_start, - positions=first_positions, - velocities=trajectory.points[0].velocities, - accelerations=trajectory.points[0].accelerations, - effort=trajectory.points[0].effort, + positions=np.asarray(first_positions, dtype=np.float64), + velocities=np.asarray(trajectory.points[0].velocities, dtype=np.float64), + accelerations=np.asarray( + trajectory.points[0].accelerations, dtype=np.float64 + ), + effort=np.asarray(trajectory.points[0].effort, dtype=np.float64), ), *list(trajectory.points)[1:], ], @@ -554,6 +566,7 @@ def get_status(self, t_now: float) -> TrajectoryStatus: progress=progress, time_elapsed=duration_from_seconds(elapsed), time_remaining=duration_from_seconds(remaining), + error="", ) completed = self._state == TrajectoryState.COMPLETED return TrajectoryStatus( @@ -564,6 +577,7 @@ def get_status(self, t_now: float) -> TrajectoryStatus: self._last_duration if completed else self._last_elapsed ), time_remaining=duration_from_seconds(0.0), + error="", ) diff --git a/dimos/control/tasks/velocity_task/test_velocity_task.py b/dimos/control/tasks/velocity_task/test_velocity_task.py index 1a4fbaa9bdf..58e0dc02b2d 100644 --- a/dimos/control/tasks/velocity_task/test_velocity_task.py +++ b/dimos/control/tasks/velocity_task/test_velocity_task.py @@ -16,7 +16,10 @@ from __future__ import annotations +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header +import numpy as np from dimos.control.task import CoordinatorState, JointStateSnapshot from dimos.control.tasks.velocity_task.velocity_task import ( @@ -31,7 +34,16 @@ def _task() -> JointVelocityTask: def test_on_joint_command_sets_velocities() -> None: task = _task() - assert task.on_joint_command(JointState(name=["a/j1", "a/j2"], velocity=[0.3, -0.1]), 1.0) + assert task.on_joint_command( + JointState( + name=["a/j1", "a/j2"], + velocity=np.array([0.3, -0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + position=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + 1.0, + ) out = task.compute(CoordinatorState(joints=JointStateSnapshot(), t_now=1.0)) assert out is not None assert out.velocities == [0.3, -0.1] @@ -41,8 +53,20 @@ def test_on_joint_command_ignores_position_bearing_messages() -> None: # Mirrors the coordinator's if-position-elif-velocity split: a message # carrying positions must never drive the velocity task. task = _task() - both = JointState(name=["a/j1", "a/j2"], position=[0.1, 0.2], velocity=[0.3, -0.1]) + both = JointState( + name=["a/j1", "a/j2"], + position=np.array([0.1, 0.2], dtype=np.float64), + velocity=np.array([0.3, -0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + effort=np.array([], dtype=np.float64), + ) assert not task.on_joint_command(both, 1.0) - positions = JointState(name=["a/j1", "a/j2"], position=[0.1, 0.2]) + positions = JointState( + name=["a/j1", "a/j2"], + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) assert not task.on_joint_command(positions, 1.0) assert not task.is_active() diff --git a/dimos/control/tasks/velocity_task/velocity_task.py b/dimos/control/tasks/velocity_task/velocity_task.py index d88071be47d..2a1eced26c1 100644 --- a/dimos/control/tasks/velocity_task/velocity_task.py +++ b/dimos/control/tasks/velocity_task/velocity_task.py @@ -242,7 +242,7 @@ def on_joint_command(self, msg: JointState, t_now: float) -> bool: """Uniform stream handler: digest the velocity half of a joint_command.""" # Position-bearing messages belong to the servo half (the coordinator # routes position XOR velocity; positions win when both are present). - if msg.position or not msg.velocity: + if len(msg.position) or not len(msg.velocity): return False return self.set_velocities_by_name(dict(zip(msg.name, msg.velocity, strict=True)), t_now) diff --git a/dimos/control/test_control.py b/dimos/control/test_control.py index 2bda2c44147..250264216a7 100644 --- a/dimos/control/test_control.py +++ b/dimos/control/test_control.py @@ -23,9 +23,12 @@ from typing import Any from unittest.mock import MagicMock +from dimos_generated.builtin_interfaces.msg import Duration, Time from dimos_generated.geometry_msgs.msg import Twist, TwistStamped, Vector3 from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import numpy as np import pytest from dimos.control._control_test_helpers import RecordingTask @@ -118,14 +121,18 @@ def simple_trajectory(): joint_names=["arm/joint1", "arm/joint2", "arm/joint3"], points=[ JointTrajectoryPoint( - positions=[0.0, 0.0, 0.0], - velocities=[0.0, 0.0, 0.0], + positions=np.array([0.0, 0.0, 0.0], dtype=np.float64), + velocities=np.array([0.0, 0.0, 0.0], dtype=np.float64), time_from_start=duration_from_seconds(0.0), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), JointTrajectoryPoint( - positions=[1.0, 0.5, 0.25], - velocities=[0.0, 0.0, 0.0], + positions=np.array([1.0, 0.5, 0.25], dtype=np.float64), + velocities=np.array([0.0, 0.0, 0.0], dtype=np.float64), time_from_start=duration_from_seconds(1.0), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ], ) @@ -269,9 +276,9 @@ def test_read_state(self, connected_hardware): assert "arm/joint1" in state assert len(state) == 6 joint_state = state["arm/joint1"] - assert joint_state.position == 0.0 - assert joint_state.velocity == 0.0 - assert joint_state.effort == 0.0 + np.testing.assert_array_equal(joint_state.position, 0.0) + np.testing.assert_array_equal(joint_state.velocity, 0.0) + np.testing.assert_array_equal(joint_state.effort, 0.0) def test_write_command(self, connected_hardware, mock_adapter): commands = { @@ -391,7 +398,7 @@ def test_map_twist_to_base_joints_routes_planar_twist_via_joint_command( assert stream == "joint_command" assert isinstance(joint_state, JointState) assert joint_state.name == ["base/vx", "base/vy", "base/wz"] - assert joint_state.velocity == [1.0, 2.0, 3.0] + np.testing.assert_array_equal(joint_state.velocity, [1.0, 2.0, 3.0]) def test_reset_runtime_state_calls_task_hooks(self, make_coordinator): class ResettableTask(BaseControlTask): @@ -532,9 +539,7 @@ def test_execute_and_cancel_without_trajectory_task_are_semantic(self, make_coor coordinator = make_coordinator() execute_result = coordinator.execute_trajectory( - JointTrajectory( - header=header_now(), - ) + JointTrajectory(header=header_now(), joint_names=[], points=[]) ) cancel_result = coordinator.cancel_trajectory() @@ -601,7 +606,10 @@ def test_joint_command_handler_converts_claimed_positions(self, trajectory_task) accepted = trajectory_task.on_joint_command( JointState( name=["arm/joint2", "other/joint"], - position=[0.25, 9.0], + position=np.array([0.25, 9.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), t_now=1.0, ) @@ -609,14 +617,32 @@ def test_joint_command_handler_converts_claimed_positions(self, trajectory_task) assert accepted assert trajectory_task._trajectory is not None assert trajectory_task._trajectory.joint_names == ["arm/joint2"] - assert trajectory_task._trajectory.points[-1].positions == [0.25] + np.testing.assert_array_equal(trajectory_task._trajectory.points[-1].positions, [0.25]) @pytest.mark.parametrize( "command", [ - JointState(name=["arm/joint1"], velocity=[0.5]), - JointState(name=["arm/joint1", "arm/joint2"], position=[0.5]), - JointState(name=["other/joint"], position=[0.5]), + JointState( + name=["arm/joint1"], + velocity=np.array([0.5], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + position=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=["arm/joint1", "arm/joint2"], + position=np.array([0.5], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=["other/joint"], + position=np.array([0.5], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ], ) def test_joint_command_handler_ignores_non_position_commands(self, trajectory_task, command): @@ -643,14 +669,18 @@ def test_execute_partial_subset_and_claims_full_configuration(self, trajectory_t joint_names=["arm/joint2", "arm/joint3"], points=[ JointTrajectoryPoint( - positions=[0.0, 0.0], - velocities=[0.0, 0.0], + positions=np.array([0.0, 0.0], dtype=np.float64), + velocities=np.array([0.0, 0.0], dtype=np.float64), time_from_start=duration_from_seconds(0.0), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), JointTrajectoryPoint( - positions=[0.5, 1.0], - velocities=[0.0, 0.0], + positions=np.array([0.5, 1.0], dtype=np.float64), + velocities=np.array([0.0, 0.0], dtype=np.float64), time_from_start=duration_from_seconds(1.0), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ], ) @@ -701,7 +731,11 @@ def test_execute_accepts_start_position_at_tolerance(self, trajectory_task, simp joint_names=[], points=[ JointTrajectoryPoint( - time_from_start=duration_from_seconds(0.0), positions=[], velocities=[] + time_from_start=duration_from_seconds(0.0), + positions=np.array([], dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) ], ), @@ -711,8 +745,10 @@ def test_execute_accepts_start_position_at_tolerance(self, trajectory_task, simp points=[ JointTrajectoryPoint( time_from_start=duration_from_seconds(0.0), - positions=[0.0, 0.0], - velocities=[0.0, 0.0], + positions=np.array([0.0, 0.0], dtype=np.float64), + velocities=np.array([0.0, 0.0], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) ], ), @@ -722,8 +758,10 @@ def test_execute_accepts_start_position_at_tolerance(self, trajectory_task, simp points=[ JointTrajectoryPoint( time_from_start=duration_from_seconds(0.0), - positions=[0.0], - velocities=[0.0], + positions=np.array([0.0], dtype=np.float64), + velocities=np.array([0.0], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) ], ), @@ -733,7 +771,11 @@ def test_execute_accepts_start_position_at_tolerance(self, trajectory_task, simp joint_names=["arm/joint1"], points=[ JointTrajectoryPoint( - time_from_start=duration_from_seconds(0.0), positions=[], velocities=[0.0] + time_from_start=duration_from_seconds(0.0), + positions=np.array([], dtype=np.float64), + velocities=np.array([0.0], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) ], ), @@ -743,8 +785,10 @@ def test_execute_accepts_start_position_at_tolerance(self, trajectory_task, simp points=[ JointTrajectoryPoint( time_from_start=duration_from_seconds(0.0), - positions=[float("nan")], - velocities=[0.0], + positions=np.array([float("nan")], dtype=np.float64), + velocities=np.array([0.0], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) ], ), @@ -754,8 +798,10 @@ def test_execute_accepts_start_position_at_tolerance(self, trajectory_task, simp points=[ JointTrajectoryPoint( time_from_start=duration_from_seconds(0.1), - positions=[0.0], - velocities=[0.0], + positions=np.array([0.0], dtype=np.float64), + velocities=np.array([0.0], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) ], ), @@ -765,13 +811,17 @@ def test_execute_accepts_start_position_at_tolerance(self, trajectory_task, simp points=[ JointTrajectoryPoint( time_from_start=duration_from_seconds(0.0), - positions=[0.0], - velocities=[0.0], + positions=np.array([0.0], dtype=np.float64), + velocities=np.array([0.0], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), JointTrajectoryPoint( time_from_start=duration_from_seconds(0.0), - positions=[1.0], - velocities=[0.0], + positions=np.array([1.0], dtype=np.float64), + velocities=np.array([0.0], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ], ), @@ -795,10 +845,18 @@ def test_compute_emits_active_subset_and_retains_final_target(self, trajectory_t joint_names=["arm/joint2"], points=[ JointTrajectoryPoint( - positions=[0.0], velocities=[0.0], time_from_start=duration_from_seconds(0.0) + positions=np.array([0.0], dtype=np.float64), + velocities=np.array([0.0], dtype=np.float64), + time_from_start=duration_from_seconds(0.0), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), JointTrajectoryPoint( - positions=[1.0], velocities=[0.0], time_from_start=duration_from_seconds(1.0) + positions=np.array([1.0], dtype=np.float64), + velocities=np.array([0.0], dtype=np.float64), + time_from_start=duration_from_seconds(1.0), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ], ) @@ -827,10 +885,18 @@ def test_replacement_reset_and_cancel_clear_active_subset(self, trajectory_task) joint_names=["arm/joint1"], points=[ JointTrajectoryPoint( - positions=[0.0], velocities=[0.0], time_from_start=duration_from_seconds(0.0) + positions=np.array([0.0], dtype=np.float64), + velocities=np.array([0.0], dtype=np.float64), + time_from_start=duration_from_seconds(0.0), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), JointTrajectoryPoint( - positions=[1.0], velocities=[0.0], time_from_start=duration_from_seconds(1.0) + positions=np.array([1.0], dtype=np.float64), + velocities=np.array([0.0], dtype=np.float64), + time_from_start=duration_from_seconds(1.0), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ], ) @@ -839,10 +905,18 @@ def test_replacement_reset_and_cancel_clear_active_subset(self, trajectory_task) joint_names=["arm/joint3"], points=[ JointTrajectoryPoint( - positions=[2.0], velocities=[0.0], time_from_start=duration_from_seconds(0.0) + positions=np.array([2.0], dtype=np.float64), + velocities=np.array([0.0], dtype=np.float64), + time_from_start=duration_from_seconds(0.0), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), JointTrajectoryPoint( - positions=[3.0], velocities=[0.0], time_from_start=duration_from_seconds(1.0) + positions=np.array([3.0], dtype=np.float64), + velocities=np.array([0.0], dtype=np.float64), + time_from_start=duration_from_seconds(1.0), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ], ) @@ -915,7 +989,7 @@ def test_trajectory_completes(self, trajectory_task, simple_trajectory, coordina # On completion, returns final position (not None) to hold at goal assert output is not None - assert output.positions == [1.0, 0.5, 0.25] # Final trajectory point + np.testing.assert_array_equal(output.positions, [1.0, 0.5, 0.25]) # Final trajectory point assert not trajectory_task.is_active() assert trajectory_task.get_state() == TrajectoryState.COMPLETED @@ -929,7 +1003,15 @@ def test_one_point_stream_target_is_velocity_bounded(self): target = JointTrajectory( header=header_now(), joint_names=["arm/joint1"], - points=[JointTrajectoryPoint(positions=[1.0])], + points=[ + JointTrajectoryPoint( + positions=np.array([1.0], dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + time_from_start=Duration(sec=0, nanosec=0), + ) + ], ) state = JointStateSnapshot(joint_positions={"arm/joint1": 0.0}) @@ -973,17 +1055,41 @@ def test_replacement_anchors_at_bounded_command_and_preserves_other_joints(self) first = JointTrajectory( header=header_now(), joint_names=["arm/joint1"], - points=[JointTrajectoryPoint(positions=[1.0])], + points=[ + JointTrajectoryPoint( + positions=np.array([1.0], dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + time_from_start=Duration(sec=0, nanosec=0), + ) + ], ) other = JointTrajectory( header=header_now(), joint_names=["arm/joint2"], - points=[JointTrajectoryPoint(positions=[-1.0])], + points=[ + JointTrajectoryPoint( + positions=np.array([-1.0], dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + time_from_start=Duration(sec=0, nanosec=0), + ) + ], ) replacement = JointTrajectory( header=header_now(), joint_names=["arm/joint1"], - points=[JointTrajectoryPoint(positions=[-1.0])], + points=[ + JointTrajectoryPoint( + positions=np.array([-1.0], dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + time_from_start=Duration(sec=0, nanosec=0), + ) + ], ) task.execute(first, {}) @@ -1009,11 +1115,19 @@ def test_partial_replacement_status_includes_untouched_joint_run(self): header=header_now(), joint_names=["arm/joint1", "arm/joint2"], points=[ - JointTrajectoryPoint(positions=[0.0, 0.0], velocities=[0.0, 0.0]), JointTrajectoryPoint( - positions=[10.0, 10.0], - velocities=[0.0, 0.0], + positions=np.array([0.0, 0.0], dtype=np.float64), + velocities=np.array([0.0, 0.0], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + time_from_start=Duration(sec=0, nanosec=0), + ), + JointTrajectoryPoint( + positions=np.array([10.0, 10.0], dtype=np.float64), + velocities=np.array([0.0, 0.0], dtype=np.float64), time_from_start=duration_from_seconds(10.0), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ], ) @@ -1021,9 +1135,19 @@ def test_partial_replacement_status_includes_untouched_joint_run(self): header=header_now(), joint_names=["arm/joint1"], points=[ - JointTrajectoryPoint(positions=[2.5], velocities=[0.0]), JointTrajectoryPoint( - positions=[3.5], velocities=[0.0], time_from_start=duration_from_seconds(1.0) + positions=np.array([2.5], dtype=np.float64), + velocities=np.array([0.0], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + time_from_start=Duration(sec=0, nanosec=0), + ), + JointTrajectoryPoint( + positions=np.array([3.5], dtype=np.float64), + velocities=np.array([0.0], dtype=np.float64), + time_from_start=duration_from_seconds(1.0), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ], ) @@ -1272,14 +1396,18 @@ def test_partial_trajectory_and_gripper_command_share_hardware_write(self, mocke joint_names=["arm/joint1", "arm/joint2"], points=[ JointTrajectoryPoint( - positions=[0.0, 0.0], - velocities=[0.0, 0.0], + positions=np.array([0.0, 0.0], dtype=np.float64), + velocities=np.array([0.0, 0.0], dtype=np.float64), time_from_start=duration_from_seconds(0.0), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), JointTrajectoryPoint( - positions=[0.5, 0.5], - velocities=[0.0, 0.0], + positions=np.array([0.5, 0.5], dtype=np.float64), + velocities=np.array([0.0, 0.0], dtype=np.float64), time_from_start=duration_from_seconds(1.0), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ], ) @@ -1422,14 +1550,18 @@ def test_full_trajectory_execution(self, mock_adapter, wait_until): joint_names=[f"arm/joint{i + 1}" for i in range(6)], points=[ JointTrajectoryPoint( - positions=[0.0] * 6, - velocities=[0.0] * 6, + positions=np.asarray([0.0] * 6, dtype=np.float64), + velocities=np.asarray([0.0] * 6, dtype=np.float64), time_from_start=duration_from_seconds(0.0), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), JointTrajectoryPoint( - positions=[0.5] * 6, - velocities=[0.0] * 6, + positions=np.asarray([0.5] * 6, dtype=np.float64), + velocities=np.asarray([0.0] * 6, dtype=np.float64), time_from_start=duration_from_seconds(0.5), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ], ) diff --git a/dimos/control/test_coordinator_commands.py b/dimos/control/test_coordinator_commands.py index ec21eee326b..26606d868ee 100644 --- a/dimos/control/test_coordinator_commands.py +++ b/dimos/control/test_coordinator_commands.py @@ -32,6 +32,7 @@ from typing import Any from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import numpy as np import pytest import dimos.control.coordinator as coord_mod @@ -129,8 +130,20 @@ def _trajectory() -> JointTrajectory: header=header_now(), joint_names=["arm/joint1", "arm/joint2"], points=[ - JointTrajectoryPoint(time_from_start=duration_from_seconds(0.0), positions=[0.0, 0.0]), - JointTrajectoryPoint(time_from_start=duration_from_seconds(1.0), positions=[0.1, 0.2]), + JointTrajectoryPoint( + time_from_start=duration_from_seconds(0.0), + positions=np.array([0.0, 0.0], dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointTrajectoryPoint( + time_from_start=duration_from_seconds(1.0), + positions=np.array([0.1, 0.2], dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ], ) diff --git a/dimos/control/test_coordinator_routing.py b/dimos/control/test_coordinator_routing.py index a2f7e9d6eda..181e4272fac 100644 --- a/dimos/control/test_coordinator_routing.py +++ b/dimos/control/test_coordinator_routing.py @@ -28,9 +28,19 @@ import threading from typing import Any -from dimos_generated.geometry_msgs.msg import Pose, PoseStamped, Twist, TwistStamped, Vector3 +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Twist, + TwistStamped, + Vector3, +) from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from dimos.control._control_test_helpers import RecordingTask @@ -167,19 +177,38 @@ def test_position_only_updates_trajectory_task_once(self, make_coordinator, mock trajectory = coordinator.get_task(JOINT_TRAJECTORY_TASK_NAME) execute = mocker.spy(trajectory, "execute") - taps["joint_command"].emit(JointState(name=ARM_JOINTS, position=[0.1, 0.2])) + taps["joint_command"].emit( + JointState( + name=ARM_JOINTS, + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) assert execute.call_count == 1 - assert trajectory._trajectory.points[-1].positions == [ - 0.1, - 0.2, - ] + np.testing.assert_array_equal( + trajectory._trajectory.points[-1].positions, + [ + 0.1, + 0.2, + ], + ) assert coordinator.get_task("vel1")._velocities is None def test_velocity_only_updates_velocity_task(self, make_coordinator): coordinator, taps = _streaming_coordinator(make_coordinator) - taps["joint_command"].emit(JointState(name=ARM_JOINTS, velocity=[0.5, 0.6])) + taps["joint_command"].emit( + JointState( + name=ARM_JOINTS, + velocity=np.array([0.5, 0.6], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + position=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) assert coordinator.get_task("vel1")._velocities == [0.5, 0.6] assert coordinator.get_task(JOINT_TRAJECTORY_TASK_NAME)._trajectory is None @@ -188,21 +217,36 @@ def test_position_wins_when_both_present(self, make_coordinator): coordinator, taps = _streaming_coordinator(make_coordinator) taps["joint_command"].emit( - JointState(name=ARM_JOINTS, position=[0.1, 0.2], velocity=[0.5, 0.6]) + JointState( + name=ARM_JOINTS, + position=np.array([0.1, 0.2], dtype=np.float64), + velocity=np.array([0.5, 0.6], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + effort=np.array([], dtype=np.float64), + ) ) - assert coordinator.get_task(JOINT_TRAJECTORY_TASK_NAME)._trajectory.points[ - -1 - ].positions == [ - 0.1, - 0.2, - ] + np.testing.assert_array_equal( + coordinator.get_task(JOINT_TRAJECTORY_TASK_NAME)._trajectory.points[-1].positions, + [ + 0.1, + 0.2, + ], + ) assert coordinator.get_task("vel1")._velocities is None def test_unclaimed_joints_route_to_nobody(self, make_coordinator): coordinator, taps = _streaming_coordinator(make_coordinator) - taps["joint_command"].emit(JointState(name=["other/joint9"], position=[1.0])) + taps["joint_command"].emit( + JointState( + name=["other/joint9"], + position=np.array([1.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) assert coordinator.get_task(JOINT_TRAJECTORY_TASK_NAME)._trajectory is None assert coordinator.get_task("vel1")._velocities is None @@ -210,7 +254,15 @@ def test_unclaimed_joints_route_to_nobody(self, make_coordinator): def test_empty_message_routes_to_nobody(self, make_coordinator): coordinator, taps = _streaming_coordinator(make_coordinator) - taps["joint_command"].emit(JointState(name=[], position=[])) + taps["joint_command"].emit( + JointState( + name=[], + position=np.array([], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) assert coordinator.get_task(JOINT_TRAJECTORY_TASK_NAME)._trajectory is None assert coordinator.get_task("vel1")._velocities is None @@ -254,14 +306,30 @@ def test_dual_arm_ports_isolate_left_from_right(self, make_coordinator): ) coordinator.start() - taps["left_cartesian"].emit(PoseStamped(header=Header(frame_id=""), pose=Pose())) + taps["left_cartesian"].emit( + PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ) left = coordinator.get_task("cartesian_left") right = coordinator.get_task("cartesian_right") assert len(left.cartesian_calls) == 1 assert right.cartesian_calls == [] - taps["right_cartesian"].emit(PoseStamped(header=Header(frame_id=""), pose=Pose())) + taps["right_cartesian"].emit( + PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ) assert len(left.cartesian_calls) == 1 assert len(right.cartesian_calls) == 1 @@ -274,7 +342,15 @@ def test_cartesian_name_match_needs_no_stream_bind(self, make_coordinator): ) coordinator.start() - taps["cartesian_command"].emit(PoseStamped(header=Header(frame_id=""), pose=Pose())) + taps["cartesian_command"].emit( + PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ) calls = coordinator.get_task("cart").cartesian_calls assert len(calls) == 1 @@ -293,9 +369,23 @@ def test_frame_id_is_not_consulted(self, make_coordinator, mocker): warn = mocker.patch.object(coord_mod.logger, "warning") taps["cartesian_command"].emit( - PoseStamped(header=Header(frame_id="some_other_task"), pose=Pose()) + PoseStamped( + header=Header(frame_id="some_other_task", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ) + taps["cartesian_command"].emit( + PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) ) - taps["cartesian_command"].emit(PoseStamped(header=Header(frame_id=""), pose=Pose())) calls = coordinator.get_task("cart").cartesian_calls assert len(calls) == 2 @@ -678,7 +768,15 @@ def test_bare_add_task_gets_no_stream_routing(self, make_coordinator): bare = CardlessStreamTask("bare", frozenset(ARM_JOINTS)) coordinator.add_task(bare) - taps["joint_command"].emit(JointState(name=ARM_JOINTS, position=[0.1, 0.2])) + taps["joint_command"].emit( + JointState( + name=ARM_JOINTS, + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) assert coordinator.get_task(JOINT_TRAJECTORY_TASK_NAME)._trajectory is not None assert bare.position_targets == [] @@ -688,8 +786,24 @@ def test_remove_task_prunes_its_routes(self, make_coordinator): trajectory = coordinator.get_task(JOINT_TRAJECTORY_TASK_NAME) assert coordinator.remove_task(JOINT_TRAJECTORY_TASK_NAME) - taps["joint_command"].emit(JointState(name=ARM_JOINTS, position=[0.1, 0.2])) - taps["joint_command"].emit(JointState(name=ARM_JOINTS, velocity=[0.5, 0.6])) + taps["joint_command"].emit( + JointState( + name=ARM_JOINTS, + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) + taps["joint_command"].emit( + JointState( + name=ARM_JOINTS, + velocity=np.array([0.5, 0.6], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + position=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) assert trajectory._trajectory is None assert coordinator.get_task("vel1")._velocities == [0.5, 0.6] @@ -703,9 +817,17 @@ def test_runtime_add_task_with_type_activates_routing(self, make_coordinator): assert coordinator.add_task(task, task_type="trajectory") assert taps["joint_command"].subscribed - taps["joint_command"].emit(JointState(name=ARM_JOINTS, position=[0.3, 0.4])) + taps["joint_command"].emit( + JointState( + name=ARM_JOINTS, + position=np.array([0.3, 0.4], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) assert task._trajectory is not None - assert task._trajectory.points[-1].positions == [0.3, 0.4] + np.testing.assert_array_equal(task._trajectory.points[-1].positions, [0.3, 0.4]) def test_runtime_registered_card_routes_with_zero_coordinator_edits( self, make_coordinator, probe_card_type @@ -715,7 +837,15 @@ def test_runtime_registered_card_routes_with_zero_coordinator_edits( probe = ProbeTask("probe1", frozenset(ARM_JOINTS)) assert coordinator.add_task(probe, task_type=probe_card_type) - taps["joint_command"].emit(JointState(name=ARM_JOINTS, position=[0.1, 0.2])) + taps["joint_command"].emit( + JointState( + name=ARM_JOINTS, + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) assert len(probe.probe_commands) == 1 msg, t_now = probe.probe_commands[0] @@ -732,11 +862,35 @@ def test_claim_overlap_gate_blocks_non_overlapping_and_empty( probe = ProbeTask("probe1", frozenset(ARM_JOINTS)) assert coordinator.add_task(probe, task_type=probe_card_type) - taps["joint_command"].emit(JointState(name=["other/joint9"], position=[1.0])) - taps["joint_command"].emit(JointState(name=[], position=[])) + taps["joint_command"].emit( + JointState( + name=["other/joint9"], + position=np.array([1.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) + taps["joint_command"].emit( + JointState( + name=[], + position=np.array([], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) assert probe.probe_commands == [] - taps["joint_command"].emit(JointState(name=ARM_JOINTS, position=[0.1, 0.2])) + taps["joint_command"].emit( + JointState( + name=ARM_JOINTS, + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) assert len(probe.probe_commands) == 1 def test_dispatch_isolates_raising_handler_from_siblings( @@ -751,7 +905,15 @@ def test_dispatch_isolates_raising_handler_from_siblings( # raiser is first in the route list; its exception must neither abort # delivery to recorder nor propagate out of the port callback (emit). - taps["joint_command"].emit(JointState(name=ARM_JOINTS, position=[0.1, 0.2])) + taps["joint_command"].emit( + JointState( + name=ARM_JOINTS, + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) assert len(recorder.probe_commands) == 1 @@ -852,7 +1014,13 @@ def test_subclass_port_routes_with_zero_coordinator_edits( assert coordinator.add_task(probe, task_type=card) assert taps["custom_in"].subscribed - msg = JointState(name=ARM_JOINTS, position=[0.1, 0.2]) + msg = JointState( + name=ARM_JOINTS, + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) taps["custom_in"].emit(msg) assert len(probe.probe_commands) == 1 @@ -866,7 +1034,15 @@ def test_direct_routing_delivers_without_a_claim_gate(self, make_coordinator, re assert coordinator.add_task(probe, task_type=card) # Joints the task does not claim, and no frame_id: neither gate applies. - taps["custom_in"].emit(JointState(name=["other/joint9"], position=[1.0])) + taps["custom_in"].emit( + JointState( + name=["other/joint9"], + position=np.array([1.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) assert len(probe.probe_commands) == 1 @@ -923,12 +1099,28 @@ def test_two_instances_read_separate_ports(self, make_coordinator, register_card assert coordinator.add_task(a, task_type=card, stream_bind={"sensor_in": "a_in"}) assert coordinator.add_task(b, task_type=card, stream_bind={"sensor_in": "b_in"}) - taps["a_in"].emit(JointState(name=ARM_JOINTS, position=[0.1, 0.2])) + taps["a_in"].emit( + JointState( + name=ARM_JOINTS, + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) assert len(a.probe_commands) == 1 assert b.probe_commands == [] - taps["b_in"].emit(JointState(name=ARM_JOINTS, position=[0.3, 0.4])) + taps["b_in"].emit( + JointState( + name=ARM_JOINTS, + position=np.array([0.3, 0.4], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) assert len(a.probe_commands) == 1 assert len(b.probe_commands) == 1 @@ -961,7 +1153,15 @@ def test_stream_bind_from_task_config(self, make_coordinator, register_card): ) coordinator.start() - taps["b_in"].emit(JointState(name=ARM_JOINTS, position=[0.1, 0.2])) + taps["b_in"].emit( + JointState( + name=ARM_JOINTS, + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) assert coordinator.get_task("a").probe_commands == [] assert len(coordinator.get_task("b").probe_commands) == 1 @@ -1059,6 +1259,14 @@ def test_direct_cross_talk_warns_naming_both_tasks( assert "first" in logged and "second" in logged assert "stream_bind" in logged - taps["custom_in"].emit(JointState(name=ARM_JOINTS, position=[0.1, 0.2])) + taps["custom_in"].emit( + JointState( + name=ARM_JOINTS, + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) assert len(first.probe_commands) == 1 assert len(second.probe_commands) == 1 diff --git a/dimos/control/test_frame_pose_tf.py b/dimos/control/test_frame_pose_tf.py index 09dfc2fe16c..4b9f3ebce70 100644 --- a/dimos/control/test_frame_pose_tf.py +++ b/dimos/control/test_frame_pose_tf.py @@ -19,6 +19,10 @@ import threading from unittest.mock import MagicMock +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage import pytest from dimos.control.components import HardwareComponent, HardwareType, make_joints @@ -27,8 +31,6 @@ from dimos.control.task import BaseControlTask, CoordinatorState, ResourceClaim from dimos.control.tick_loop import TickLoop from dimos.hardware.manipulators.spec import ManipulatorAdapter -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.robot.manipulators.common.coordinators import ArmTwistCoordinator JOINTS = make_joints("arm", 2) @@ -53,7 +55,15 @@ def on_preempted(self, by_task: str, joints: frozenset[str]) -> None: def measured_frame_poses(self, state: CoordinatorState) -> dict[str, PoseStamped]: x = state.joints.get_position(JOINTS[0]) - return {"link_tcp": PoseStamped(frame_id="link_base", position=(x, 0.0, 0.5))} + return { + "link_tcp": PoseStamped( + header=Header(frame_id="link_base", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=x, y=0.0, z=0.5), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + } def test_measured_frame_poses_are_published_as_world_tf_at_a_limited_rate() -> None: @@ -81,9 +91,13 @@ def test_measured_frame_poses_are_published_as_world_tf_at_a_limited_rate() -> N (message,) = published (transform,) = message.transforms - assert (transform.frame_id, transform.child_frame_id) == ("world", "link_tcp") - assert transform.translation.to_numpy().tolist() == [0.25, 0.0, 0.5] - assert transform.ts == joint_states.call_args_list[0].args[0].ts + assert (transform.header.frame_id, transform.child_frame_id) == ("world", "link_tcp") + assert [ + transform.transform.translation.x, + transform.transform.translation.y, + transform.transform.translation.z, + ] == [0.25, 0.0, 0.5] + assert transform.header.stamp == joint_states.call_args_list[0].args[0].header.stamp def test_only_coordinators_declaring_tf_can_publish_frame_poses() -> None: diff --git a/dimos/control/test_operator_hold.py b/dimos/control/test_operator_hold.py index 1b9b8c74289..42231daa2f2 100644 --- a/dimos/control/test_operator_hold.py +++ b/dimos/control/test_operator_hold.py @@ -24,6 +24,10 @@ from collections.abc import Iterator from unittest.mock import MagicMock +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.std_msgs.msg import Header +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import numpy as np import pytest from dimos.control.components import ( @@ -55,9 +59,8 @@ from dimos.control.tick_loop import TickLoop from dimos.hardware.drive_trains.spec import TwistBaseAdapter from dimos.hardware.manipulators.spec import ManipulatorAdapter -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint -from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState +from dimos.msgs.time import duration_from_seconds +from dimos.msgs.trajectory import TrajectoryState ARM_JOINTS = make_joints("arm", 3) ARM_POSITIONS = [0.1, 0.2, 0.3] @@ -95,10 +98,23 @@ def on_preempted(self, by_task: str, joints: frozenset[str]) -> None: def _trajectory() -> JointTrajectory: """One second from the arm's current pose to +0.5 rad on every joint.""" return JointTrajectory( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), joint_names=list(ARM_JOINTS), points=[ - TrajectoryPoint(positions=list(ARM_POSITIONS), time_from_start=0.0), - TrajectoryPoint(positions=[p + 0.5 for p in ARM_POSITIONS], time_from_start=1.0), + JointTrajectoryPoint( + positions=np.asarray(ARM_POSITIONS, dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + time_from_start=duration_from_seconds(0.0), + ), + JointTrajectoryPoint( + positions=np.asarray([p + 0.5 for p in ARM_POSITIONS], dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + time_from_start=duration_from_seconds(1.0), + ), ], ) diff --git a/dimos/control/test_per_robot_joint_states.py b/dimos/control/test_per_robot_joint_states.py index ca68cc8de15..84d3e8e9d1b 100644 --- a/dimos/control/test_per_robot_joint_states.py +++ b/dimos/control/test_per_robot_joint_states.py @@ -21,7 +21,10 @@ from typing import Any from unittest.mock import MagicMock +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from dimos.control.components import HardwareComponent, HardwareType, make_joints @@ -361,11 +364,19 @@ def test_subclass_serves_its_own_input_and_outputs_together( right = OutTap(coordinator.right_arm_joints) coordinator.start() - commands.emit(JointState(name=LEFT_JOINTS, position=[0.4, 0.5])) + commands.emit( + JointState( + name=LEFT_JOINTS, + position=np.array([0.4, 0.5], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) trajectory = coordinator.get_task("joint_trajectory")._trajectory assert trajectory is not None - assert trajectory.points[-1].positions == [0.4, 0.5] + np.testing.assert_array_equal(trajectory.points[-1].positions, [0.4, 0.5]) wait_until(lambda: bool(left.count and right.count), timeout=5.0) assert list(left.latest().name) == LEFT_JOINTS assert list(right.latest().name) == RIGHT_JOINTS diff --git a/dimos/control/tick_loop.py b/dimos/control/tick_loop.py index d164513b6cb..49493455625 100644 --- a/dimos/control/tick_loop.py +++ b/dimos/control/tick_loop.py @@ -31,6 +31,12 @@ import time from typing import TYPE_CHECKING, NamedTuple +from dimos_generated.geometry_msgs.msg import Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage +import numpy as np + from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.control.task import ( ControlTask, @@ -39,20 +45,18 @@ JointStateSnapshot, ResourceClaim, ) -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.msgs.time import time_from_seconds from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: from collections.abc import Callable + from dimos_generated.geometry_msgs.msg import PoseStamped + from dimos.control.components import HardwareId, JointName, JointState as JointReading, TaskName from dimos.control.hardware_interface import ConnectedHardware from dimos.hardware.manipulators.spec import ControlMode from dimos.hardware.whole_body.spec import IMUState - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped logger = setup_logger() @@ -449,12 +453,15 @@ def _publish_joint_state(self, snapshot: JointStateSnapshot) -> None: """Publish aggregated JointState for external consumers.""" names = list(snapshot.joint_positions.keys()) msg = JointState( - ts=snapshot.timestamp, - frame_id=self._frame_id, + header=Header(stamp=time_from_seconds(snapshot.timestamp), frame_id=self._frame_id), name=names, - position=[snapshot.joint_positions[n] for n in names], - velocity=[snapshot.joint_velocities.get(n, 0.0) for n in names], - effort=[snapshot.joint_efforts.get(n, 0.0) for n in names], + position=np.asarray([snapshot.joint_positions[n] for n in names], dtype=np.float64), + velocity=np.asarray( + [snapshot.joint_velocities.get(n, 0.0) for n in names], dtype=np.float64 + ), + effort=np.asarray( + [snapshot.joint_efforts.get(n, 0.0) for n in names], dtype=np.float64 + ), ) if self._publish_callback: self._publish_callback(msg) @@ -471,18 +478,21 @@ def _publish_frame_poses(self, state: CoordinatorState) -> None: ts = state.joints.timestamp self._publish_tf_callback( TFMessage( - *( - Transform( - translation=pose.position, - rotation=pose.orientation, - # IK solvers return world coordinates (base pose applied), but - # stamp them with the base link's name. - frame_id="world", + transforms=[ + TransformStamped( + header=Header(stamp=time_from_seconds(ts), frame_id="world"), child_frame_id=frame, - ts=ts, + transform=Transform( + translation=Vector3( + x=pose.pose.position.x, + y=pose.pose.position.y, + z=pose.pose.position.z, + ), + rotation=pose.pose.orientation, + ), ) for frame, pose in poses.items() - ) + ] ) ) @@ -498,12 +508,11 @@ def _publish_robot_joint_states( for hw_id, state in per_hardware.items(): names = list(state.keys()) msg = JointState( - ts=timestamp, - frame_id=hw_id, + header=Header(stamp=time_from_seconds(timestamp), frame_id=hw_id), name=names, - position=[state[n].position for n in names], - velocity=[state[n].velocity for n in names], - effort=[state[n].effort for n in names], + position=np.asarray([state[n].position for n in names], dtype=np.float64), + velocity=np.asarray([state[n].velocity for n in names], dtype=np.float64), + effort=np.asarray([state[n].effort for n in names], dtype=np.float64), ) try: publish(hw_id, msg) diff --git a/dimos/core/coordination/blueprint_config/values.py b/dimos/core/coordination/blueprint_config/values.py index d1c354ffd4a..8130d7d2bcb 100644 --- a/dimos/core/coordination/blueprint_config/values.py +++ b/dimos/core/coordination/blueprint_config/values.py @@ -20,7 +20,7 @@ import copy from dataclasses import fields, is_dataclass from types import MappingProxyType -from typing import Any, cast +from typing import Any, ClassVar, cast, get_origin from pydantic import BaseModel @@ -61,6 +61,11 @@ def plain(value: Any, *, exclude_unset: bool = False) -> Any: return { info.name: plain(getattr(value, info.name), exclude_unset=exclude_unset) for info in fields(value) + if get_origin(info.type) is not ClassVar + and not ( + isinstance(info.type, str) + and info.type.startswith(("ClassVar[", "typing.ClassVar[")) + ) } if isinstance(value, Mapping): return { diff --git a/dimos/core/coordination/test_module_coordinator.py b/dimos/core/coordination/test_module_coordinator.py index d7b374c1426..c6945f179bc 100644 --- a/dimos/core/coordination/test_module_coordinator.py +++ b/dimos/core/coordination/test_module_coordinator.py @@ -18,7 +18,8 @@ from types import MappingProxyType from typing import Any, Protocol -from dimos_generated.geometry_msgs.msg import TransformStamped +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 from dimos_generated.sensor_msgs.msg import Image from dimos_generated.std_msgs.msg import Header from dimos_generated.tf2_msgs.msg import TFMessage @@ -1166,7 +1167,16 @@ class IoTfPublisher(Module): def send(self, child: str) -> None: self.tf.publish( TFMessage( - transforms=[TransformStamped(header=Header(frame_id="world"), child_frame_id=child)] + transforms=[ + TransformStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + child_frame_id=child, + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ] ) ) @@ -1188,7 +1198,16 @@ async def handle_tf(self, msg: TFMessage) -> None: def send(self, child: str) -> None: self.tf.publish( TFMessage( - transforms=[TransformStamped(header=Header(frame_id="world"), child_frame_id=child)] + transforms=[ + TransformStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + child_frame_id=child, + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ] ) ) diff --git a/dimos/core/native_module.py b/dimos/core/native_module.py index f709ddb3519..602794f105e 100644 --- a/dimos/core/native_module.py +++ b/dimos/core/native_module.py @@ -44,6 +44,7 @@ class MyCppModule(NativeModule): from dataclasses import asdict, is_dataclass import enum import functools +from importlib.util import find_spec import json import os from pathlib import Path @@ -54,6 +55,12 @@ class MyCppModule(NativeModule): import time from typing import IO, Any +from dimos_cdrgen.native_build import ( + build_environment, + installed_prefixes, + prepare_cpp, + stage_module, +) import numpy as np from pydantic import Field, model_validator @@ -151,6 +158,8 @@ class NativeModuleConfig(ModuleConfig): executable: str source_dir: str | None = None build_command: str | None = None + # Opt-in CMake source preparation; Rust/Nix modules keep their own build commands. + cmake_message_packages: list[str] = Field(default_factory=list) extra_args: list[str] = Field(default_factory=list) extra_env: dict[str, str] = Field(default_factory=dict) # Session settings for this module alone, e.g. opening it as the zenoh router @@ -558,11 +567,35 @@ def _maybe_build(self) -> None: build_command=self.config.build_command, ) build_start = time.perf_counter() + build_env = {**os.environ, **self.config.extra_env} + if self.config.cmake_message_packages: + prefixes: list[str] = [] + for package in self.config.cmake_message_packages: + provider = find_spec(package + "_schemas") + if provider is None or provider.origin is None: + raise ValueError(f"Missing installed message source package: {package}") + prefix = prepare_cpp(Path(provider.origin).parent / "package") + prefixes.extend( + value for value in installed_prefixes(prefix) if value not in prefixes + ) + package_path = Path(__file__).resolve().parents[1] + sdk = package_path / "_native" / "cpp" + if not (sdk / "CMakeLists.txt").is_file(): + sdk = package_path.parent / "native" / "cpp" + if not (sdk / "CMakeLists.txt").is_file(): + raise FileNotFoundError("DimOS native C++ source resources are missing") + prefixes.append(str(sdk)) + if self._cwd is None: + raise ValueError("CMake source preparation requires resolved source_dir") + source, exe = stage_module(Path(self._cwd), exe) + self._cwd = str(source) + self._executable = str(exe) + build_env = build_environment([Path(prefix) for prefix in prefixes], build_env) proc = subprocess.Popen( self.config.build_command, shell=True, cwd=self._cwd, - env={**os.environ, **self.config.extra_env}, + env=build_env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, ) diff --git a/dimos/core/test_cdr_transport.py b/dimos/core/test_cdr_transport.py index ed175046001..00e7cfce9cc 100644 --- a/dimos/core/test_cdr_transport.py +++ b/dimos/core/test_cdr_transport.py @@ -18,12 +18,15 @@ import threading import uuid +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from dimos.core.transport import LCMTransport, SHMTransport, ZenohTransport +from dimos.protocol.pubsub.impl.lcmpubsub import LCM from dimos.protocol.service.zenohservice import ZenohSessionPool @@ -37,9 +40,24 @@ def transport_pair(request, lcm_url): def make(message_type): topic = f"dimos/cdr/{uuid.uuid4().hex[:8]}" transports = [] + # Native memq queues belong to a handle, which permits only one + # dispatcher. Share its pubsub/service for an in-process roundtrip; + # UDP endpoints retain independent native handles and dispatchers. + memory_bus = ( + LCM(url=lcm_url) + if request.param == "lcm" and lcm_url.startswith("memq://") + else None + ) for _ in range(2): if request.param == "lcm": - transport = LCMTransport(topic, message_type, url=lcm_url) + transport = LCMTransport( + topic, + message_type, + url=lcm_url, + lcm=memory_bus.l if memory_bus is not None else None, + ) + if memory_bus is not None: + transport.lcm = memory_bus elif request.param == "zenoh": transport = ZenohTransport( topic, @@ -78,7 +96,15 @@ def on_message(message): def test_large_generated_image_preserves_pixels_and_header(transport_pair, retry_until): publisher, subscriber = transport_pair(Image) pixels = np.arange(640 * 480 * 3, dtype=np.uint8) - message = Image(width=640, height=480, encoding="rgb8", step=1920, data=pixels) + message = Image( + width=640, + height=480, + encoding="rgb8", + step=1920, + data=np.asarray(pixels, dtype=np.uint8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + is_bigendian=0, + ) message.header.frame_id = "camera" message.header.stamp.sec = 1_700_000_000 message.header.stamp.nanosec = 123_456_789 diff --git a/dimos/core/test_native_module.py b/dimos/core/test_native_module.py index 52dc3e83dca..ba1a374f8b8 100644 --- a/dimos/core/test_native_module.py +++ b/dimos/core/test_native_module.py @@ -23,11 +23,16 @@ from io import BytesIO import json from pathlib import Path +import shlex +import sys import time from types import SimpleNamespace from typing import Any from unittest.mock import Mock, patch +from dimos_generated.geometry_msgs.msg import Twist +from dimos_generated.sensor_msgs.msg import Imu, PointCloud2 +from dimos_generated.tf2_msgs.msg import TFMessage from pydantic import ValidationError import pytest @@ -42,10 +47,6 @@ from dimos.core.transport import LCMTransport, ZenohTransport from dimos.core.transport_factory import make_transport, rpc_backend, transport_topic from dimos.experimental.memory.rust_recorder import RustRecorderConfig -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.sensor_msgs.Imu import Imu -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.protocol.pubsub.impl.zenohpubsub import QOS_NEVER_DROP, Topic as ZenohTopic from dimos.protocol.service import zenohservice from dimos.protocol.service.zenohservice import ZenohConfig @@ -174,8 +175,8 @@ def test_manual(dimos_cluster: ModuleCoordinator, args_file: str) -> None: native_module.stop() assert read_json_file(args_file) == { - "cmd_vel": "/cmd_vel#geometry_msgs.Twist", - "pointcloud": "/my/custom/lidar#sensor_msgs.PointCloud2", + "cmd_vel": "/cmd_vel#geometry_msgs/msg/Twist", + "pointcloud": "/my/custom/lidar#sensor_msgs/msg/PointCloud2", "output_file": args_file, "some_param": "2.5", } @@ -190,8 +191,8 @@ def test_io_port_topic_reaches_the_native_process() -> None: module.set_transport("tf", transports[1]) assert module._collect_topics() == { - "cmd_vel": "/cmd_vel#geometry_msgs.Twist", - "tf": "/tf#tf2_msgs.TFMessage", + "cmd_vel": "/cmd_vel#geometry_msgs/msg/Twist", + "tf": "/tf#tf2_msgs/msg/TFMessage", } finally: module.stop() @@ -207,7 +208,7 @@ def test_tf_topic_comes_from_the_declared_port_only() -> None: try: module.set_transport("cmd_vel", transport) - assert module._collect_topics() == {"cmd_vel": "/cmd_vel#geometry_msgs.Twist"} + assert module._collect_topics() == {"cmd_vel": "/cmd_vel#geometry_msgs/msg/Twist"} finally: module.stop() with contextlib.suppress(Exception): @@ -520,6 +521,45 @@ def test_json_mode_malformed_falls_back_to_plain_text() -> None: assert calls[0][1] == "not json at all" +def test_cmake_preparation_uses_writable_cache_and_existing_build_hook(tmp_path, monkeypatch): + source = tmp_path / "sources" + source.mkdir() + (source / "build_it.py").write_text( + "import os\nfrom pathlib import Path\n" + "Path('build').mkdir(exist_ok=True)\n" + "Path('build/app').write_text(os.environ['CMAKE_PREFIX_PATH'])\n" + ) + source.chmod(0o555) + prefix = tmp_path / "prepared" + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache")) + prepare = Mock(return_value=prefix) + monkeypatch.setattr(native_module_mod, "prepare_cpp", prepare) + monkeypatch.setattr(native_module_mod, "installed_prefixes", lambda root: [str(root)]) + backend = rpc_backend() + monkeypatch.setattr(backend, "start", Mock()) + monkeypatch.setattr(backend, "serve_module_rpc", Mock()) + monkeypatch.setattr(backend, "stop", Mock()) + monkeypatch.setattr(native_module_mod, "get_project_root", lambda: tmp_path) + module = StubBuildModule( + source_dir="sources", + executable="build/app", + build_command=shlex.join([sys.executable, "build_it.py"]), + cmake_message_packages=["dimos_generated"], + ) + try: + module._prepare_native() + executable = Path(module._executable) + assert executable.is_relative_to(tmp_path / "cache") + assert str(prefix) in executable.read_text() + assert not (source / "build").exists() + assert prepare.call_count == 1 + module._prepare_native() + assert prepare.call_count == 1 + finally: + module.stop() + source.chmod(0o755) + + @pytest.mark.parametrize( "module_path,class_name,executable", [ diff --git a/dimos/core/test_transport.py b/dimos/core/test_transport.py index b25292a3b2b..d2805cbd7c3 100644 --- a/dimos/core/test_transport.py +++ b/dimos/core/test_transport.py @@ -14,6 +14,7 @@ import pickle +from dimos_generated.sensor_msgs.msg import Image import pytest from dimos.core.transport import ( @@ -21,7 +22,6 @@ LCMTransport, pLCMTransport, ) -from dimos.msgs.sensor_msgs.Image import Image @pytest.mark.parametrize("transport_class", [pLCMTransport, LCMTransport, JpegLcmTransport]) diff --git a/dimos/core/test_transport_factory.py b/dimos/core/test_transport_factory.py index 42504371c96..177f0bf096e 100644 --- a/dimos/core/test_transport_factory.py +++ b/dimos/core/test_transport_factory.py @@ -14,6 +14,8 @@ from __future__ import annotations +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 import pytest from dimos.core.global_config import GlobalConfig @@ -31,10 +33,6 @@ session_config, transport_topic, ) -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.protocol.pubsub.impl.zenohpubsub import QOS_LATEST_WINS, QOS_NEVER_DROP from dimos.protocol.rpc.pubsubrpc import LCMRPC from dimos.protocol.rpc.zenohrpc import ZenohRPC diff --git a/dimos/core/test_zenoh_transport.py b/dimos/core/test_zenoh_transport.py index c47fe6e772f..f96b8af7c0c 100644 --- a/dimos/core/test_zenoh_transport.py +++ b/dimos/core/test_zenoh_transport.py @@ -19,7 +19,10 @@ from types import SimpleNamespace from typing import cast +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Vector3 as TypedMsg from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header import numpy as np from pydantic import ValidationError import pytest @@ -48,21 +51,6 @@ from dimos.protocol.service.zenohservice import ZenohSessionPool -class TypedMsg: - """A generated-message codec contract for transport selection tests.""" - - msg_name = "test_msgs/msg/TypedMsg" - schema = "" - - @staticmethod - def encode() -> bytes: - return b"" - - @staticmethod - def decode(data: bytes): - return TypedMsg() - - class UntypedMsg: """A Python-only message without a generated codec contract.""" @@ -196,7 +184,15 @@ def test_zenoh_transport_broadcast_and_subscribe(retry_until, session_pool, coll t.start() t.subscribe(collector.callback) - test_img = Image(width=2, height=2, step=6, encoding="rgb8", data=np.zeros(12, dtype=np.uint8)) + test_img = Image( + width=2, + height=2, + step=6, + encoding="rgb8", + data=np.zeros(12, dtype=np.uint8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + is_bigendian=0, + ) retry_until(collector.event, lambda: t.broadcast(None, test_img)) assert isinstance(collector.received[0], Image) t.stop() diff --git a/dimos/core/transport.py b/dimos/core/transport.py index ffef78418dd..0a1c9a42d05 100644 --- a/dimos/core/transport.py +++ b/dimos/core/transport.py @@ -25,11 +25,13 @@ cast, ) +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode import numpy as np from dimos.core.stream import Stream, Transport from dimos.msgs.protocol import DimosMsg from dimos.protocol.pubsub.impl.lcmpubsub import LCM, PickleLCM, Topic as LCMTopic +from dimos.protocol.pubsub.impl.shmpubsub import BytesSharedMemory, PickleSharedMemory from dimos.protocol.pubsub.impl.webrtc.providers.broker import BrokerConfig from dimos.protocol.pubsub.impl.webrtc.providers.spec import AudioProvider, ProviderConfig from dimos.protocol.pubsub.impl.webrtc.webrtcpubsub import WebRTCPubSub @@ -41,6 +43,13 @@ from dimos.stream.audio.base import AudioEvent from dimos.utils import colors from dimos.utils.logging_config import setup_logger +from dimos.web.relay_bridge.protocol import ( + FrameHeader, + ProtocolError, + decode_data_frame, + encode_data_frame, + peek_data_frame_lengths, +) logger = setup_logger() @@ -153,7 +162,7 @@ def stop(self) -> None: def __reduce__(self): # type: ignore[no-untyped-def] return ( partial(LCMTransport, **self.lcm.config.model_dump(exclude={"lcm"})), - (self.topic.topic, self.topic.lcm_type), + (self.topic.topic, self.topic.msg_type), ) def broadcast(self, _, msg) -> None: # type: ignore[no-untyped-def] @@ -172,17 +181,22 @@ def subscribe( class JpegLcmTransport(LCMTransport): # type: ignore[type-arg] def __init__(self, topic: str, type: type, **kwargs) -> None: # type: ignore[no-untyped-def] + from dimos_generated.sensor_msgs.msg import CompressedImage, Image + from dimos.protocol.pubsub.impl.jpeg_lcm import ( JpegLCM, ) # ~330ms: deferred to avoid pulling in Image/cv2/rerun + if type is not Image: + raise TypeError("JpegLcmTransport requires generated sensor_msgs/Image") + self.image_type = type self.lcm = JpegLCM(**kwargs) # type: ignore[assignment] - super().__init__(topic, type) + super().__init__(topic, CompressedImage) def __reduce__(self): # type: ignore[no-untyped-def] return ( partial(JpegLcmTransport, **self.lcm.config.model_dump(exclude={"lcm"})), - (self.topic.topic, self.topic.lcm_type), + (self.topic.topic, self.image_type), ) def start(self) -> None: @@ -194,6 +208,122 @@ def stop(self) -> None: self._started = False +class pSHMTransport(PubSubTransport[T]): + _started: bool = False + + def __init__(self, topic: str, **kwargs) -> None: # type: ignore[no-untyped-def] + super().__init__(topic) + self.shm = PickleSharedMemory(**kwargs) + + def __reduce__(self): # type: ignore[no-untyped-def] + return ( + partial(pSHMTransport, default_capacity=self.shm.config.default_capacity), + (self.topic,), + ) + + def broadcast(self, _, msg) -> None: # type: ignore[no-untyped-def] + if not self._started: + self.start() + + self.shm.publish(self.topic, msg) + + def subscribe( + self, callback: Callable[[T], Any], selfstream: Stream[T] | None = None + ) -> Callable[[], None]: + if not self._started: + self.start() + return self.shm.subscribe(self.topic, lambda msg, topic: callback(msg)) + + def start(self) -> None: + self.shm.start() + self._started = True + + def stop(self) -> None: + self.shm.stop() + self._started = False + + +class SHMTransport(PubSubTransport[T]): + _started: bool = False + + def __init__(self, topic: str, msg_type: type[DimosMsg] | None = None, **kwargs: Any) -> None: + self._base_topic = topic + self._msg_type = msg_type + super().__init__(str(LCMTopic(topic, msg_type))) + self.shm = BytesSharedMemory(**kwargs) + + def __reduce__(self): # type: ignore[no-untyped-def] + return ( + partial(SHMTransport, default_capacity=self.shm.config.default_capacity), + (self._base_topic, self._msg_type), + ) + + def broadcast(self, _, msg) -> None: # type: ignore[no-untyped-def] + if not self._started: + self.start() + + payload = cdr_encode(msg) if self._msg_type is not None else msg + self.shm.publish(self.topic, payload) + + def subscribe( + self, callback: Callable[[T], Any], selfstream: Stream[T] | None = None + ) -> Callable[[], None]: + if not self._started: + self.start() + msg_type = self._msg_type + if msg_type is None: + return self.shm.subscribe(self.topic, lambda msg, topic: callback(cast("T", msg))) + return self.shm.subscribe( + self.topic, + lambda msg, topic: callback(cast("T", cdr_decode(msg, msg_type.__msgtype__))), + ) + + def start(self) -> None: + self.shm.start() + self._started = True + + def stop(self) -> None: + self.shm.stop() + self._started = False + + +class JpegShmTransport(PubSubTransport[T]): + _started: bool = False + + def __init__(self, topic: str, quality: int = 75, **kwargs) -> None: # type: ignore[no-untyped-def] + super().__init__(topic) + from dimos.protocol.pubsub.impl.jpeg_shm import ( + JpegSharedMemory, + ) # deferred to avoid pulling in Image/cv2/rerun + + self.shm = JpegSharedMemory(quality=quality, **kwargs) + self.quality = quality + + def __reduce__(self): # type: ignore[no-untyped-def] + return (JpegShmTransport, (self.topic, self.quality)) + + def broadcast(self, _, msg) -> None: # type: ignore[no-untyped-def] + if not self._started: + self.start() + + self.shm.publish(self.topic, msg) + + def subscribe( + self, callback: Callable[[T], Any], selfstream: Stream[T] | None = None + ) -> Callable[[], None]: + if not self._started: + self.start() + return self.shm.subscribe(self.topic, lambda msg, topic: callback(msg)) + + def start(self) -> None: + self.shm.start() + self._started = True + + def stop(self) -> None: + self.shm.stop() + self._started = False + + class ROSTransport(PubSubTransport[DimosMsg]): _ros: DimosROS | None = None @@ -296,8 +426,8 @@ class WebRTCTransport(PubSubTransport[M]): also be used directly with an explicit ``config``. Two modes: * **Raw bytes** (``msg_type=None``): messages pass through as ``bytes``. - * **Typed LCM** (``msg_type=SomeMsg``): LCM-encoded on ``broadcast()``, - LCM-decoded on ``subscribe()`` (foreign types on the shared channel are skipped) — so multiple + * **Typed CDR** (``msg_type=SomeMsg``): CDR in the shared browser data-frame + format, with channel and type metadata checked before decoding. Multiple transports sharing one multiplexed DataChannel each receive only their own message type. @@ -324,6 +454,7 @@ def __init__( self._msg_type = msg_type self._config = config or self._config_cls(**config_kwargs) self._pubsub: WebRTCPubSub | None = None + self._frame_sequence = 0 # Guards first-use init: concurrent subscribe()/broadcast() must not # construct two WebRTCPubSub wrappers (one would silently orphan any # subscribe_all state). Never pickled — __reduce__ rebuilds via the @@ -337,7 +468,25 @@ def broadcast(self, _: Stream[M] | None, msg: M) -> None: if not self._started: self.start() assert self._pubsub is not None - data = msg.lcm_encode() if self._msg_type is not None else msg + data: bytes | M + if self._msg_type is not None: + if msg.__msgtype__ != self._msg_type.__msgtype__: + raise ValueError("message type does not match declared WebRTC channel type") + with self._init_lock: + sequence = self._frame_sequence + self._frame_sequence += 1 + data = encode_data_frame( + FrameHeader( + ch=self.topic, + seq=sequence, + ts=time.time(), + delivery="latest", + meta={"type": self._msg_type.__msgtype__, "encoding": "cdr"}, + ), + cdr_encode(msg), + ) + else: + data = msg self._pubsub.publish(self.topic, data) # type: ignore[arg-type] def subscribe( @@ -351,12 +500,20 @@ def subscribe( msg_type = self._msg_type def _typed_cb(data: bytes, _topic: str) -> None: - # The channel is multiplexed (e.g. the browser sends Twists and - # Poses on cmd_unreliable); lcm_decode verifies the wire - # fingerprint and raises on other types — skip those. try: - msg = msg_type.lcm_decode(data) - except ValueError: + lengths = peek_data_frame_lengths(data) + if lengths is None or lengths[2] != len(data): + return + frame = decode_data_frame(data) + meta = frame.header.meta or {} + if ( + frame.header.ch != self.topic + or meta.get("encoding") != "cdr" + or meta.get("type") != msg_type.__msgtype__ + ): + return + msg = cdr_decode(frame.payload, msg_type.__msgtype__) + except (ProtocolError, ValueError): return callback(msg) # type: ignore[arg-type] @@ -503,7 +660,7 @@ class CloudflareAudioTransport(WebRTCAudioTransport): class ZenohTransport(PubSubTransport[T]): - """Zenoh transport with LCM encoding for typed DimosMsg. + """Zenoh transport with CDR encoding for typed DimosMsg. Accepts either a plain topic string plus message type, or a full `ZenohTopic` carrying per-topic settings: `ZenohTransport(ZenohTopic("bla", diff --git a/dimos/core/transport_factory.py b/dimos/core/transport_factory.py index 074df0a92ae..ffad47240d7 100644 --- a/dimos/core/transport_factory.py +++ b/dimos/core/transport_factory.py @@ -92,7 +92,7 @@ def default_zenoh_qos_for(name: str, msg_name: str) -> ZenohQoS | None: def default_zenoh_qos(name: str, msg_type: type | None = None) -> ZenohQoS | None: """Default publisher QoS for a logical channel; None = zenoh defaults.""" - return default_zenoh_qos_for(name, getattr(msg_type, "msg_name", "")) + return default_zenoh_qos_for(name, getattr(msg_type, "__msgtype__", "")) def make_transport( diff --git a/dimos/e2e_tests/cdr_replay_fixture.py b/dimos/e2e_tests/cdr_replay_fixture.py index cb953a948cd..6e2dd09fb5c 100644 --- a/dimos/e2e_tests/cdr_replay_fixture.py +++ b/dimos/e2e_tests/cdr_replay_fixture.py @@ -46,8 +46,10 @@ def write_go2_cdr_replay(path: Path, *, duration_s: int = 180) -> None: pose = PoseStamped( header=header, pose=Pose( - position=Point(x=math.cos(angle), y=math.sin(angle)), - orientation=Quaternion(z=math.sin(angle / 2), w=math.cos(angle / 2)), + position=Point(x=math.cos(angle), y=math.sin(angle), z=0.0), + orientation=Quaternion( + z=math.sin(angle / 2), w=math.cos(angle / 2), x=0.0, y=0.0 + ), ), ) pixels = np.full((240, 320, 3), index % 256, dtype=np.uint8) diff --git a/dimos/e2e_tests/conftest.py b/dimos/e2e_tests/conftest.py index 95e9153c35d..2c8a9f66e17 100644 --- a/dimos/e2e_tests/conftest.py +++ b/dimos/e2e_tests/conftest.py @@ -15,21 +15,24 @@ from __future__ import annotations from collections.abc import Callable, Generator, Iterator +from pathlib import Path import subprocess import sys import threading import time from typing import TYPE_CHECKING, Any +from dimos_cdrgen.registry import decode as cdr_decode +from dimos_generated.geometry_msgs.msg import PoseStamped import pytest from dimos.core.coordination.coordinator_rpc import CoordinatorRPC from dimos.core.global_config import global_config from dimos.core.transport_factory import make_transport +from dimos.e2e_tests.cdr_replay_fixture import write_go2_cdr_replay from dimos.e2e_tests.dim_sim_client import DimSimClient from dimos.e2e_tests.dimos_cli_call import DimosCliCall from dimos.e2e_tests.lcm_spy import LcmSpy -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.simulation.mujoco.direct_cmd_vel_explorer import DirectCmdVelExplorer if TYPE_CHECKING: @@ -212,7 +215,7 @@ def dim_sim(): @pytest.fixture def spawn_wall_on_pose(lcm_spy: LcmSpy, dim_sim: DimSimClient): """Spawn a dim_sim wall when the robot's /odom comes within `threshold` metres of `point`.""" - odom_topic = "/odom#geometry_msgs.PoseStamped" + odom_topic = "/odom#geometry_msgs/msg/PoseStamped" stop_event = threading.Event() workers: list[threading.Thread] = [] errors: list[BaseException] = [] @@ -225,9 +228,9 @@ def spawn(*, point, threshold, wall) -> None: def on_odom(data): if triggered.is_set(): return - pose = PoseStamped.lcm_decode(data) - dx = pose.x - px - dy = pose.y - py + pose = cdr_decode(data, PoseStamped) + dx = pose.pose.position.x - px + dy = pose.pose.position.y - py if dx * dx + dy * dy < threshold_sq: triggered.set() @@ -289,3 +292,11 @@ def explore() -> None: direct_cmd_vel_explorer.follow_points(points) return explore + + +@pytest.fixture(scope="session") +def cdr_go2_replay_db(tmp_path_factory: pytest.TempPathFactory) -> Path: + """Current generated-message replay, independent of retired LFS codecs.""" + path = tmp_path_factory.mktemp("go2-cdr") / "replay.db" + write_go2_cdr_replay(path) + return path diff --git a/dimos/e2e_tests/dim_sim_client.py b/dimos/e2e_tests/dim_sim_client.py index d3ece964fd0..7f5f0d21d6f 100644 --- a/dimos/e2e_tests/dim_sim_client.py +++ b/dimos/e2e_tests/dim_sim_client.py @@ -12,10 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.std_msgs.msg import Header + from dimos.core.global_config import global_config from dimos.core.transport import PubSubTransport from dimos.core.transport_factory import make_transport -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.simulation.dimsim.scene_client import SceneClient @@ -56,8 +59,10 @@ def add_wall(self, x1: float, y1: float, x2: float, y2: float) -> None: def publish_goal(self, x: float, y: float) -> None: self._goal_request.publish( PoseStamped( - position=(x, y, 0), - orientation=(0, 0, 0, 1), - frame_id="world", + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=x, y=y, z=0.0), + orientation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), + ), ) ) diff --git a/dimos/e2e_tests/lcm_spy.py b/dimos/e2e_tests/lcm_spy.py index 14875fef524..6d6b919a5d3 100644 --- a/dimos/e2e_tests/lcm_spy.py +++ b/dimos/e2e_tests/lcm_spy.py @@ -19,19 +19,30 @@ import threading from typing import Any +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.geometry_msgs.msg import PoseStamped + from dimos.core.global_config import global_config from dimos.core.transport_factory import transport_topic -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.protocol import DimosMsg from dimos.protocol.pubsub.impl.lcmpubsub import LCMPubSubBase, Topic from dimos.protocol.pubsub.impl.zenohpubsub import Topic as ZenohTopic, ZenohPubSubBase from dimos.utils.testing.waiting import wait_until +def _wire_topic(channel: str) -> ZenohTopic: + """The active backend's topic for an LCM-style channel ('/odom#pkg.Msg'). + + ZenohTopic is an LCM Topic plus zenoh's key expression, so both buses take it. + """ + parsed = Topic.from_channel_str(transport_topic(channel)) + return ZenohTopic(parsed.topic, parsed.msg_type) + + class LcmSpy: """Sniffs every message on the active transport's bus. - Topics are named the LCM way (`/odom#geometry_msgs.PoseStamped`); + Topics are named the LCM way (`/odom#geometry_msgs/msg/PoseStamped`); `transport_topic` maps them to the running backend's wire name. """ @@ -79,6 +90,9 @@ def _on_message(self, data: bytes, topic: Topic) -> None: for listener in listeners: listener(data) + def publish(self, topic: str, msg: Any) -> None: + self._bus.publish(_wire_topic(topic), cdr_encode(msg)) + def save_topic(self, topic: str) -> None: with self._saved_topics_lock: self._saved_topics[transport_topic(topic)] = topic @@ -121,7 +135,7 @@ def wait_for_message_result( event = threading.Event() def listener(msg: bytes) -> None: - data = type.lcm_decode(msg) + data = cdr_decode(msg, type) if predicate(data): event.set() @@ -136,12 +150,12 @@ def wait_until_odom_position( self, x: float, y: float, threshold: float = 1, timeout: float = 60 ) -> None: def predicate(msg: PoseStamped) -> bool: - pos = msg.position + pos = msg.pose.position distance = math.sqrt((pos.x - x) ** 2 + (pos.y - y) ** 2) return distance < threshold self.wait_for_message_result( - "/odom#geometry_msgs.PoseStamped", + "/odom#geometry_msgs/msg/PoseStamped", PoseStamped, predicate, f"Failed to get to position x={x}, y={y}", diff --git a/dimos/e2e_tests/test_cdr_spy.py b/dimos/e2e_tests/test_cdr_spy.py index 4c72e0b163d..347d7acd0ad 100644 --- a/dimos/e2e_tests/test_cdr_spy.py +++ b/dimos/e2e_tests/test_cdr_spy.py @@ -12,7 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_cdrgen.registry import decode as cdr_decode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.std_msgs.msg import Header from dimos.core.global_config import global_config @@ -23,11 +25,16 @@ def test_spy_publishes_generated_cdr_on_qualified_channel(mocker, monkeypatch): monkeypatch.setattr(global_config, "transport", "lcm") bus = mocker.patch("dimos.e2e_tests.lcm_spy.LCMPubSubBase").return_value spy = LcmSpy() - message = PoseStamped(header=Header(frame_id="map"), pose=Pose(position=Point(x=2, y=3))) + message = PoseStamped( + header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=2, y=3, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + ) spy.publish("/goal_request#geometry_msgs/msg/PoseStamped", message) topic, payload = bus.publish.call_args.args assert str(topic) == "/goal_request#geometry_msgs/msg/PoseStamped" - decoded = PoseStamped.decode(payload) + decoded = cdr_decode(payload, PoseStamped) assert (decoded.pose.position.x, decoded.pose.position.y, decoded.header.frame_id) == ( 2, 3, @@ -47,5 +54,19 @@ def test_spy_odometry_predicate_uses_nested_pose(mocker, monkeypatch): PoseStamped, 60, ) - assert predicate(PoseStamped(pose=Pose(position=Point(x=2, y=3)))) - assert not predicate(PoseStamped(pose=Pose(position=Point(x=3, y=3)))) + assert predicate( + PoseStamped( + pose=Pose( + position=Point(x=2, y=3, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + ) + assert not predicate( + PoseStamped( + pose=Pose( + position=Point(x=3, y=3, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + ) diff --git a/dimos/e2e_tests/test_control_coordinator.py b/dimos/e2e_tests/test_control_coordinator.py index 8b3c3ec663a..3adc89ad27c 100644 --- a/dimos/e2e_tests/test_control_coordinator.py +++ b/dimos/e2e_tests/test_control_coordinator.py @@ -22,6 +22,10 @@ import platform import time +from dimos_cdrgen.registry import decode as cdr_decode +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import numpy as np import pytest from dimos.control.coordinator import ControlCoordinator @@ -31,10 +35,8 @@ TrajectoryExecutionStatus, ) from dimos.core.rpc_client import RPCClient -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint -from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState +from dimos.msgs.time import duration_from_seconds, header_now +from dimos.msgs.trajectory import TrajectoryState pytestmark = pytest.mark.self_hosted @@ -47,7 +49,7 @@ def test_coordinator_starts_and_responds_to_rpc( ) -> None: """Test that coordinator starts and responds to RPC queries.""" # Save topics we care about (topic names carry the type suffix) - joint_state_topic = "/coordinator_joint_state#sensor_msgs.JointState" + joint_state_topic = "/coordinator_joint_state#sensor_msgs/msg/JointState" lcm_spy.save_topic(joint_state_topic) # Start the mock coordinator blueprint @@ -83,14 +85,14 @@ def test_coordinator_executes_trajectory( ) -> None: """Test that coordinator executes a trajectory via RPC.""" # Save topics - lcm_spy.save_topic("/coordinator_joint_state#sensor_msgs.JointState") + lcm_spy.save_topic("/coordinator_joint_state#sensor_msgs/msg/JointState") # Start coordinator start_blueprint("coordinator-mock") wait_for_system_ready() # Wait for it to be ready - lcm_spy.wait_for_saved_topic("/coordinator_joint_state#sensor_msgs.JointState") + lcm_spy.wait_for_saved_topic("/coordinator_joint_state#sensor_msgs/msg/JointState") # Create RPC client client = RPCClient(None, ControlCoordinator) @@ -101,17 +103,22 @@ def test_coordinator_executes_trajectory( # Create a simple trajectory trajectory = JointTrajectory( + header=header_now(), joint_names=[f"arm/joint{i + 1}" for i in range(7)], points=[ - TrajectoryPoint( - time_from_start=0.0, - positions=[0.0] * 7, - velocities=[0.0] * 7, + JointTrajectoryPoint( + time_from_start=duration_from_seconds(0.0), + positions=np.asarray([0.0] * 7, dtype=np.float64), + velocities=np.asarray([0.0] * 7, dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), - TrajectoryPoint( - time_from_start=0.5, - positions=[0.1] * 7, - velocities=[0.0] * 7, + JointTrajectoryPoint( + time_from_start=duration_from_seconds(0.5), + positions=np.asarray([0.1] * 7, dtype=np.float64), + velocities=np.asarray([0.0] * 7, dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ], ) @@ -139,7 +146,7 @@ def test_coordinator_joint_state_published( self, lcm_spy, start_blueprint, wait_for_system_ready ) -> None: """Test that joint state messages are published at expected rate.""" - joint_state_topic = "/coordinator_joint_state#sensor_msgs.JointState" + joint_state_topic = "/coordinator_joint_state#sensor_msgs/msg/JointState" lcm_spy.save_topic(joint_state_topic) # Start coordinator @@ -170,7 +177,7 @@ def test_coordinator_joint_state_published( with lcm_spy._messages_lock: raw_msg = lcm_spy.messages[joint_state_topic][0] - joint_state = JointState.lcm_decode(raw_msg) + joint_state = cdr_decode(raw_msg, JointState) assert len(joint_state.name) == 7 assert len(joint_state.position) == 7 assert "arm/joint1" in joint_state.name @@ -179,28 +186,33 @@ def test_coordinator_cancel_trajectory( self, lcm_spy, start_blueprint, wait_for_system_ready ) -> None: """Test that a running trajectory can be cancelled.""" - lcm_spy.save_topic("/coordinator_joint_state#sensor_msgs.JointState") + lcm_spy.save_topic("/coordinator_joint_state#sensor_msgs/msg/JointState") # Start coordinator start_blueprint("coordinator-mock") wait_for_system_ready() - lcm_spy.wait_for_saved_topic("/coordinator_joint_state#sensor_msgs.JointState") + lcm_spy.wait_for_saved_topic("/coordinator_joint_state#sensor_msgs/msg/JointState") client = RPCClient(None, ControlCoordinator) try: # Create a long trajectory (5 seconds) trajectory = JointTrajectory( + header=header_now(), joint_names=[f"arm/joint{i + 1}" for i in range(7)], points=[ - TrajectoryPoint( - time_from_start=0.0, - positions=[0.0] * 7, - velocities=[0.0] * 7, + JointTrajectoryPoint( + time_from_start=duration_from_seconds(0.0), + positions=np.asarray([0.0] * 7, dtype=np.float64), + velocities=np.asarray([0.0] * 7, dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), - TrajectoryPoint( - time_from_start=5.0, - positions=[1.0] * 7, - velocities=[0.0] * 7, + JointTrajectoryPoint( + time_from_start=duration_from_seconds(5.0), + positions=np.asarray([1.0] * 7, dtype=np.float64), + velocities=np.asarray([0.0] * 7, dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ], ) @@ -225,12 +237,12 @@ def test_dual_arm_coordinator( self, lcm_spy, start_blueprint, wait_until, wait_for_system_ready ) -> None: """Test dual-arm coordinator moving both arms with one combined trajectory.""" - lcm_spy.save_topic("/coordinator_joint_state#sensor_msgs.JointState") + lcm_spy.save_topic("/coordinator_joint_state#sensor_msgs/msg/JointState") # Start dual-arm mock coordinator start_blueprint("coordinator-dual-mock") wait_for_system_ready() - lcm_spy.wait_for_saved_topic("/coordinator_joint_state#sensor_msgs.JointState") + lcm_spy.wait_for_saved_topic("/coordinator_joint_state#sensor_msgs/msg/JointState") client = RPCClient(None, ControlCoordinator) try: @@ -247,11 +259,24 @@ def test_dual_arm_coordinator( # One trajectory moving the left arm (7 joints) and right arm # (6 joints) together trajectory = JointTrajectory( + header=header_now(), joint_names=[f"left_arm/joint{i + 1}" for i in range(7)] + [f"right_arm/joint{i + 1}" for i in range(6)], points=[ - TrajectoryPoint(time_from_start=0.0, positions=[0.0] * 13), - TrajectoryPoint(time_from_start=0.5, positions=[0.2] * 7 + [0.3] * 6), + JointTrajectoryPoint( + time_from_start=duration_from_seconds(0.0), + positions=np.asarray([0.0] * 13, dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointTrajectoryPoint( + time_from_start=duration_from_seconds(0.5), + positions=np.asarray([0.2] * 7 + [0.3] * 6, dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ], ) diff --git a/dimos/e2e_tests/test_custom_channel_browser.py b/dimos/e2e_tests/test_custom_channel_browser.py index 2ff83eb292a..29586fd7812 100644 --- a/dimos/e2e_tests/test_custom_channel_browser.py +++ b/dimos/e2e_tests/test_custom_channel_browser.py @@ -53,16 +53,25 @@ header=Header(stamp=time_from_seconds(42.5), frame_id="world"), poses=[ PoseStamped( - header=Header(stamp=time_from_seconds(42.5)), - pose=Pose(position=Point(x=1.5, y=-2.5), orientation=Quaternion(w=1)), + header=Header(stamp=time_from_seconds(42.5), frame_id=""), + pose=Pose( + position=Point(x=1.5, y=-2.5, z=0.0), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), ), PoseStamped( - header=Header(stamp=time_from_seconds(42.5)), - pose=Pose(position=Point(x=2.0, y=0.25), orientation=Quaternion(w=1)), + header=Header(stamp=time_from_seconds(42.5), frame_id=""), + pose=Pose( + position=Point(x=2.0, y=0.25, z=0.0), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), ), PoseStamped( - header=Header(stamp=time_from_seconds(42.5)), - pose=Pose(position=Point(x=3.5, y=4.0), orientation=Quaternion(w=1)), + header=Header(stamp=time_from_seconds(42.5), frame_id=""), + pose=Pose( + position=Point(x=3.5, y=4.0, z=0.0), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), ), ], ) diff --git a/dimos/e2e_tests/test_dimsim_typesafe_agent.py b/dimos/e2e_tests/test_dimsim_typesafe_agent.py index 0506fbe73a4..afaf850f9b7 100644 --- a/dimos/e2e_tests/test_dimsim_typesafe_agent.py +++ b/dimos/e2e_tests/test_dimsim_typesafe_agent.py @@ -22,12 +22,12 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.vision_msgs.msg import Detection3DArray import pytest from dimos.agents.typesafe.test_world_state import det3d from dimos.core.transport_factory import make_transport -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray CHAIR = (4.0, 2.0) @@ -69,7 +69,7 @@ def test_go_to_the_chair( wait_for_system_ready(timeout=1200.0) # Server physics (odom/lidar) starts only after the browser ships its snapshot. lcm_spy.wait_for_message_result( - "/odom#geometry_msgs.PoseStamped", PoseStamped, lambda _m: True, "no odom", timeout=600 + "/odom#geometry_msgs/msg/PoseStamped", PoseStamped, lambda _m: True, "no odom", timeout=600 ) dim_sim.set_agent_position(1.0, 2.0) diff --git a/dimos/e2e_tests/test_dimsim_walk_forward.py b/dimos/e2e_tests/test_dimsim_walk_forward.py index fe4a955adca..3fc5f08d191 100644 --- a/dimos/e2e_tests/test_dimsim_walk_forward.py +++ b/dimos/e2e_tests/test_dimsim_walk_forward.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.geometry_msgs.msg import PoseStamped import pytest -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - @pytest.mark.self_hosted_large def test_walk_forward( @@ -35,7 +34,7 @@ def test_walk_forward( # The coordinator can be ready while DimSim is still loading its scene. # Teleports sent before physics starts are discarded by the bridge. lcm_spy.wait_for_message_result( - "/odom#geometry_msgs.PoseStamped", + "/odom#geometry_msgs/msg/PoseStamped", PoseStamped, predicate=lambda _: True, fail_message="DimSim did not publish initial odometry", diff --git a/dimos/e2e_tests/test_lcm_channel_browser.py b/dimos/e2e_tests/test_lcm_channel_browser.py index 37a8d03c80d..12c325b4a28 100644 --- a/dimos/e2e_tests/test_lcm_channel_browser.py +++ b/dimos/e2e_tests/test_lcm_channel_browser.py @@ -39,7 +39,9 @@ POSE = PoseStamped( header=Header(stamp=time_from_seconds(42.5), frame_id="map"), - pose=Pose(position=Point(x=1.5, y=-2.5, z=0.25), orientation=Quaternion(w=1)), + pose=Pose( + position=Point(x=1.5, y=-2.5, z=0.25), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), ) # A page whose SDK and relay live on another origin, reading the decoded CDR diff --git a/dimos/e2e_tests/test_manipulation_planning_groups.py b/dimos/e2e_tests/test_manipulation_planning_groups.py index 004e0018575..bf40844eaeb 100644 --- a/dimos/e2e_tests/test_manipulation_planning_groups.py +++ b/dimos/e2e_tests/test_manipulation_planning_groups.py @@ -23,7 +23,10 @@ from collections.abc import Callable import time +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from dimos.control.coordinator import ControlCoordinator @@ -140,7 +143,10 @@ def _offset_target(snapshot: ManipulationSnapshot, group_id: str, delta: float) assert current is not None return JointState( name=current.name, - position=[position + delta for position in current.position], + position=np.asarray([position + delta for position in current.position], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) diff --git a/dimos/e2e_tests/test_map_click_browser.py b/dimos/e2e_tests/test_map_click_browser.py index 0695f848b6f..d3b0e46aedc 100644 --- a/dimos/e2e_tests/test_map_click_browser.py +++ b/dimos/e2e_tests/test_map_click_browser.py @@ -18,6 +18,7 @@ import threading from typing import NamedTuple +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, PointStamped, Pose, PoseStamped, Quaternion from dimos_generated.nav_msgs.msg import Path from dimos_generated.std_msgs.msg import Header @@ -52,8 +53,10 @@ class MapBridge(NamedTuple): def _path(*xy: tuple[float, float]) -> Path: poses = [ PoseStamped( - header=Header(stamp=time_from_seconds(1.0)), - pose=Pose(position=Point(x=x, y=y), orientation=Quaternion(w=1)), + header=Header(stamp=time_from_seconds(1.0), frame_id=""), + pose=Pose( + position=Point(x=x, y=y, z=0.0), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), ) for x, y in xy ] @@ -141,5 +144,7 @@ def test_click_to_goal_and_cancel(map_bridge: MapBridge, chromium_page: Page) -> cancel.click() assert wait_until(lambda: stops == [True], timeout=15.0) - map_bridge.path.publish(Path()) + map_bridge.path.publish( + Path(header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), poses=[]) + ) expect(cancel).to_have_count(0, timeout=30_000) diff --git a/dimos/e2e_tests/test_rtc_video_browser.py b/dimos/e2e_tests/test_rtc_video_browser.py index a6bfd55e985..ff1de1e9ac3 100644 --- a/dimos/e2e_tests/test_rtc_video_browser.py +++ b/dimos/e2e_tests/test_rtc_video_browser.py @@ -28,7 +28,7 @@ import numpy as np import pytest -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_from_array from dimos.web.cockpit import Video, cockpit from dimos.web.relay_bridge.relay_bridge_module import RTC_AVAILABLE @@ -49,8 +49,8 @@ # A static picture still yields a steady stream: aiortc encodes every fed frame. _ROWS = np.linspace(0, 255, 240, dtype=np.uint8)[:, None] _COLS = np.linspace(0, 255, 320, dtype=np.uint8)[None, :] -IMAGE = Image( - data=np.stack( +IMAGE = image_from_array( + np.stack( [ np.broadcast_to(_ROWS, (240, 320)), np.broadcast_to(_COLS, (240, 320)), @@ -58,7 +58,7 @@ ], axis=-1, ).copy(), - format=ImageFormat.RGB, + encoding="rgb8", ) diff --git a/dimos/evals/agents/test_no_dimos.py b/dimos/evals/agents/test_no_dimos.py index f2371d0b2a1..6b1f44ea59c 100644 --- a/dimos/evals/agents/test_no_dimos.py +++ b/dimos/evals/agents/test_no_dimos.py @@ -92,7 +92,10 @@ def test_recorded_observations_are_exported_as_plain_files(tmp_path: Path) -> No ts=1.0, ) lidar.append( - pointcloud_from_xyz(points * 2, header=Header(stamp=time_from_seconds(2.0))), ts=2.0 + pointcloud_from_xyz( + points * 2, header=Header(stamp=time_from_seconds(2.0), frame_id="") + ), + ts=2.0, ) images = store.stream("camera", Image) images.append(image_from_array(pixels, encoding="rgb8"), ts=3.0) diff --git a/dimos/evals/conftest.py b/dimos/evals/conftest.py index 8ea5e741760..baeee5c22fb 100644 --- a/dimos/evals/conftest.py +++ b/dimos/evals/conftest.py @@ -14,6 +14,7 @@ from pathlib import Path +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.std_msgs.msg import Header import pytest @@ -29,8 +30,11 @@ def dataset(tmp_path: Path) -> str: stream = store.stream("odom", PoseStamped) for i in range(5): pose = PoseStamped( - header=Header(frame_id="world"), - pose=Pose(position=Point(x=float(i)), orientation=Quaternion(w=1.0)), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=float(i), y=0.0, z=0.0), + orientation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), + ), ) stream.append(pose, ts=1000.0 + i) return str(path) diff --git a/dimos/evals/environments/habitat.py b/dimos/evals/environments/habitat.py index 5216c1d9b8f..8569b9c99fd 100644 --- a/dimos/evals/environments/habitat.py +++ b/dimos/evals/environments/habitat.py @@ -22,10 +22,11 @@ import time from typing import TYPE_CHECKING, Any +from dimos_generated.geometry_msgs.msg import PoseStamped from pydantic import model_validator from dimos.evals.environments.sim import Sim, SimConfig -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.time import to_seconds from dimos.simulation.habitat.server import HabitatProp if TYPE_CHECKING: @@ -163,7 +164,10 @@ def wait_ready(self, recording: Store, *, deadline: float) -> None: try: pose = self.latest_pose(recording) image = recording.streams.color_image.last().data - if time.time() - min(pose.ts, image.ts) < 10.0: + if ( + time.time() - min(to_seconds(pose.header.stamp), to_seconds(image.header.stamp)) + < 10.0 + ): self._spawn = pose return except (LookupError, AttributeError): @@ -176,9 +180,7 @@ def latest_pose(self, recording: Store) -> PoseStamped: if "odometry" not in recording.streams: raise LookupError("No Habitat odometry recorded") odom = recording.streams.odometry.last().data - return PoseStamped( - ts=odom.ts, frame_id=odom.frame_id, position=odom.position, orientation=odom.orientation - ) + return PoseStamped(header=odom.header, pose=odom.pose.pose) def episode_metadata(self) -> dict[str, object]: """Record the launch overrides and initial observed pose. @@ -195,5 +197,11 @@ def episode_metadata(self) -> dict[str, object]: ), # Describes the scan generation method, not publication state. "point_cloud_source": "depth_unprojection", - "initial_observed_position_ros": list(self._spawn.position) if self._spawn else None, + "initial_observed_position_ros": [ + self._spawn.pose.position.x, + self._spawn.pose.position.y, + self._spawn.pose.position.z, + ] + if self._spawn + else None, } diff --git a/dimos/evals/environments/image_file.py b/dimos/evals/environments/image_file.py index a03d33ad0c3..b0d41433e65 100644 --- a/dimos/evals/environments/image_file.py +++ b/dimos/evals/environments/image_file.py @@ -20,8 +20,15 @@ from pathlib import Path from typing import TYPE_CHECKING, Any +from dimos_generated.sensor_msgs.msg import Image +import numpy as np +from PIL import Image as PILImage + from dimos.evals.environments.base import Environment, EnvironmentConfig from dimos.evals.types import RunningEnvironment +from dimos.memory.store.memory import MemoryStore +from dimos.msgs.image import image_from_array +from dimos.msgs.time import header_now, to_seconds if TYPE_CHECKING: from dimos.evals.agents.base import Agent @@ -50,14 +57,13 @@ def preflight(self, agent: Agent) -> None: raise FileNotFoundError(f"image does not exist: {self.config.path}") def start(self, modules: Sequence[str]) -> RunningEnvironment: - from dimos.memory.store.memory import MemoryStore - from dimos.msgs.sensor_msgs.Image import Image - - image = Image.from_file(self.config.path) + with PILImage.open(self.config.path) as source: + pixels = np.asarray(source.convert("RGB")) + image = image_from_array(pixels, encoding="rgb8", header=header_now()) recording = MemoryStore() self._resources.callback(recording.stop) stream = recording.stream("image", Image) - stream.append(image, ts=image.ts) + stream.append(image, ts=to_seconds(image.header.stamp)) return RunningEnvironment( mcp_url="", streams=(stream,), artifacts={"image": self.config.path} ) diff --git a/dimos/evals/environments/lib/benchmark_scenes.py b/dimos/evals/environments/lib/benchmark_scenes.py index c6eb69f9e33..ac95b26c5b3 100644 --- a/dimos/evals/environments/lib/benchmark_scenes.py +++ b/dimos/evals/environments/lib/benchmark_scenes.py @@ -94,8 +94,8 @@ def grade_from_spec(spec: GradeSpec) -> Callable[[Outcome], float]: def grade(outcome: Outcome) -> float: with recording(outcome) as store: try: - start = first_body_transform(store, spec.body).translation.z - end = last_body_transform(store, spec.body).translation.z + start = first_body_transform(store, spec.body).transform.translation.z + end = last_body_transform(store, spec.body).transform.translation.z except LookupError: return 0.0 return min(max((end - start) / spec.by_m, 0.0), 1.0) diff --git a/dimos/evals/environments/lib/recorded_poses.py b/dimos/evals/environments/lib/recorded_poses.py index d316fded888..cac19e49c57 100644 --- a/dimos/evals/environments/lib/recorded_poses.py +++ b/dimos/evals/environments/lib/recorded_poses.py @@ -19,26 +19,27 @@ from typing import TYPE_CHECKING, cast if TYPE_CHECKING: + from dimos_generated.geometry_msgs.msg import TransformStamped + from dimos.memory.store.base import Store - from dimos.msgs.geometry_msgs.Transform import Transform -def last_body_transform(recording: Store, body: str) -> Transform: +def last_body_transform(recording: Store, body: str) -> TransformStamped: """The newest ``world -> body`` transform in the recording.""" return _body_transform(recording, body, newest=True) -def first_body_transform(recording: Store, body: str) -> Transform: +def first_body_transform(recording: Store, body: str) -> TransformStamped: """The oldest ``world -> body`` transform in the recording.""" return _body_transform(recording, body, newest=False) -def _body_transform(recording: Store, body: str, *, newest: bool) -> Transform: +def _body_transform(recording: Store, body: str, *, newest: bool) -> TransformStamped: # Several modules publish on tf, so scan messages until one names the body. if "tf" not in recording.streams: raise LookupError("No tf recorded") for record in recording.streams.tf.order_by("ts", desc=newest): for transform in record.data.transforms: if transform.child_frame_id == body: - return cast("Transform", transform) + return cast("TransformStamped", transform) raise LookupError(f"No tf for body {body!r}; is it in tracked_bodies?") diff --git a/dimos/evals/environments/mujoco_sim.py b/dimos/evals/environments/mujoco_sim.py index 5108ac7ee60..f0cf4af156a 100644 --- a/dimos/evals/environments/mujoco_sim.py +++ b/dimos/evals/environments/mujoco_sim.py @@ -26,11 +26,13 @@ from dimos.evals.environments.lib.recorded_poses import last_body_transform from dimos.evals.environments.sim import Sim, SimConfig +from dimos.msgs.time import to_seconds if TYPE_CHECKING: + from dimos_generated.geometry_msgs.msg import PoseStamped + from dimos.e2e_tests.dimos_cli_call import DimosCliCall from dimos.memory.store.base import Store - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped class MujocoEnvironmentConfig(SimConfig): @@ -81,7 +83,8 @@ def wait_ready(self, recording: Store, *, deadline: float) -> None: while time.monotonic() < deadline: try: ages = [ - time.time() - getattr(recording.streams, name).last().data.ts + time.time() + - to_seconds(getattr(recording.streams, name).last().data.header.stamp) for name in self.config.ready_streams if name in recording.streams ] diff --git a/dimos/evals/environments/test_dimsim.py b/dimos/evals/environments/test_dimsim.py index 3ed0818edf7..7b238f246fb 100644 --- a/dimos/evals/environments/test_dimsim.py +++ b/dimos/evals/environments/test_dimsim.py @@ -14,7 +14,7 @@ from unittest.mock import Mock -from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.std_msgs.msg import Header import pytest @@ -39,7 +39,13 @@ def test_dimsim_launch_setup_and_pose(mocker): with MemoryStore() as store: with pytest.raises(LookupError): env.latest_pose(store) - pose = PoseStamped(header=Header(stamp=time_from_seconds(123), frame_id="world")) + pose = PoseStamped( + header=Header(stamp=time_from_seconds(123), frame_id="world"), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) store.stream("odom", PoseStamped).append(pose) assert to_seconds(env.latest_pose(store).header.stamp) == 123 finally: diff --git a/dimos/evals/environments/test_habitat.py b/dimos/evals/environments/test_habitat.py index 1110eac54be..11274bd0d18 100644 --- a/dimos/evals/environments/test_habitat.py +++ b/dimos/evals/environments/test_habitat.py @@ -16,15 +16,26 @@ from types import SimpleNamespace from unittest.mock import Mock +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseWithCovariance, + Quaternion, + Twist, + TwistWithCovariance, + Vector3, +) +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser from dimos.e2e_tests.dimos_cli_call import DimosCliCall from dimos.evals.environments.habitat import HabitatEnvironment -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.msgs.image import image_from_array +from dimos.msgs.time import time_from_seconds from dimos.simulation.habitat.connection import HabitatConnection, HabitatConnectionConfig @@ -182,7 +193,6 @@ def test_readiness_failure_releases_resources(tmp_path, mocker): def test_readiness_and_pose_normalization(): from dimos.memory.store.memory import MemoryStore - from dimos.msgs.sensor_msgs.Image import Image env = environment() with MemoryStore() as store: @@ -190,14 +200,34 @@ def test_readiness_and_pose_normalization(): env.latest_pose(store) with pytest.raises(TimeoutError): env.wait_ready(store, deadline=time.monotonic()) - odom = Odometry(frame_id="world", pose=Pose(position=Vector3(1, 2, 3))) + odom = Odometry( + header=Header(frame_id="world", stamp=time_from_seconds(time.time())), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=1, y=2, z=3), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), + child_frame_id="", + twist=TwistWithCovariance( + twist=Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ), + covariance=np.zeros(36, dtype=np.float64), + ), + ) store.stream("odometry", Odometry).append(odom) - store.stream("color_image", Image).append(Image(np.zeros((2, 2, 3), dtype=np.uint8))) + store.stream("color_image", Image).append( + image_from_array( + np.zeros((2, 2, 3), dtype=np.uint8), encoding="rgb8", header=odom.header + ) + ) env.wait_ready(store, deadline=time.monotonic() + 1) pose = env.latest_pose(store) - assert pose.ts == odom.ts - assert pose.frame_id == "world" - assert tuple(pose.position) == (1, 2, 3) + assert pose.header.stamp == odom.header.stamp + assert pose.header.frame_id == "world" + assert (pose.pose.position.x, pose.pose.position.y, pose.pose.position.z) == (1, 2, 3) assert env.episode_metadata()["initial_observed_position_ros"] == [1, 2, 3] diff --git a/dimos/evals/environments/test_mujoco_sim.py b/dimos/evals/environments/test_mujoco_sim.py index d8d1f12068e..1661373d464 100644 --- a/dimos/evals/environments/test_mujoco_sim.py +++ b/dimos/evals/environments/test_mujoco_sim.py @@ -16,6 +16,18 @@ from pathlib import Path import time +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.sensor_msgs.msg import Image, JointState +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage import numpy as np import pytest @@ -24,12 +36,8 @@ from dimos.evals.environments.mujoco_sim import MujocoEnvironment from dimos.manipulation.manipulation_module import ManipulationModuleConfig from dimos.memory.store.memory import MemoryStore -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.image import image_from_array +from dimos.msgs.time import time_from_seconds, to_seconds def environment(**kwargs): @@ -38,7 +46,16 @@ def environment(**kwargs): def _tf(ts: float, child: str, z: float) -> TFMessage: return TFMessage( - Transform(translation=Vector3(0.4, 0.08, z), frame_id="world", child_frame_id=child, ts=ts) + transforms=[ + TransformStamped( + header=Header(stamp=time_from_seconds(ts), frame_id="world"), + child_frame_id=child, + transform=Transform( + translation=Vector3(x=0.4, y=0.08, z=z), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ] ) @@ -103,12 +120,12 @@ def test_module_env_reaches_blueprint_parser(monkeypatch): assert sim["tracked_bodies"] == ["apple", "cup"] mounted = ManipulationModuleConfig(**parsed.module_kwargs("manipulationmodule")) - assert mounted.model.base_pose.position.z == pytest.approx(0.912) + assert mounted.model.base_pose.pose.position.z == pytest.approx(0.912) defaults = BlueprintConfigParser(xarm_perception_sim).parse(environ={}) default = ManipulationModuleConfig(**defaults.module_kwargs("manipulationmodule")) - assert default.model.base_pose.position.z == pytest.approx(0.12) + assert default.model.base_pose.pose.position.z == pytest.approx(0.12) assert mounted.model.joint_names == default.model.joint_names - assert mounted.model.base_pose.orientation == default.model.base_pose.orientation + assert mounted.model.base_pose.pose.orientation == default.model.base_pose.pose.orientation def test_latest_pose_needs_odom(): @@ -116,8 +133,16 @@ def test_latest_pose_needs_odom(): with MemoryStore() as store: with pytest.raises(LookupError): env.latest_pose(store) - store.stream("odom", PoseStamped).append(PoseStamped(ts=5, frame_id="world")) - assert env.latest_pose(store).ts == 5 + store.stream("odom", PoseStamped).append( + PoseStamped( + header=Header(stamp=time_from_seconds(5), frame_id="world"), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ) + assert to_seconds(env.latest_pose(store).header.stamp) == 5 def test_raw_arm_blueprint_configures_its_bridge(): @@ -180,10 +205,20 @@ def test_ready_needs_fresh_streams_and_tracked_body_poses(): env.wait_ready(store, deadline=time.monotonic() + 0.3) now = time.time() store.stream("color_image", Image).append( - Image(data=np.zeros((1, 1, 3), dtype=np.uint8), format=ImageFormat.RGB, ts=now) + image_from_array( + np.zeros((1, 1, 3), dtype=np.uint8), + encoding="rgb8", + header=Header(stamp=time_from_seconds(now), frame_id=""), + ) ) store.stream("coordinator_joint_state", JointState).append( - JointState(ts=now, name=["j1"], position=[0.0], velocity=[0.0]) + JointState( + header=Header(stamp=time_from_seconds(now), frame_id=""), + name=["j1"], + position=np.array([0.0]), + velocity=np.array([0.0]), + effort=np.array([], dtype=np.float64), + ) ) store.stream("tf", TFMessage).append(_tf(now, "apple", 0.17)) env.wait_ready(store, deadline=time.monotonic() + 2.0) @@ -199,13 +234,29 @@ def test_settle_waits_for_joints_to_stop(): joints = store.stream("coordinator_joint_state", JointState) env.settle(1.0) assert time.monotonic() - started < 0.5 - joints.append(JointState(ts=1.0, name=["j1"], position=[0.0], velocity=[0.5])) + joints.append( + JointState( + header=Header(stamp=time_from_seconds(1.0), frame_id=""), + name=["j1"], + position=np.array([0.0]), + velocity=np.array([0.5]), + effort=np.array([], dtype=np.float64), + ) + ) env._recording = store started = time.monotonic() env.settle(0.2) assert time.monotonic() - started >= 0.2 - joints.append(JointState(ts=2.0, name=["j1"], position=[0.0], velocity=[0.0])) + joints.append( + JointState( + header=Header(stamp=time_from_seconds(2.0), frame_id=""), + name=["j1"], + position=np.array([0.0]), + velocity=np.array([0.0]), + effort=np.array([], dtype=np.float64), + ) + ) started = time.monotonic() env.settle(5.0) assert time.monotonic() - started < 1.0 diff --git a/dimos/evals/scorers.py b/dimos/evals/scorers.py index 230b2032848..9dc74c09ad9 100644 --- a/dimos/evals/scorers.py +++ b/dimos/evals/scorers.py @@ -34,6 +34,7 @@ from dimos.evals.environments.lib.recorded_poses import first_body_transform, last_body_transform from dimos.evals.types import Outcome, recording +from dimos.msgs.geometry import transform_matrix, vector_array T = TypeVar("T") @@ -160,8 +161,8 @@ def lifted(body: str, *, by_m: float) -> Callable[[Outcome], float]: def grade(outcome: Outcome) -> float: with recording(outcome) as store: try: - start = first_body_transform(store, body).translation.z - end = last_body_transform(store, body).translation.z + start = first_body_transform(store, body).transform.translation.z + end = last_body_transform(store, body).transform.translation.z except LookupError: return 0.0 return min(max((end - start) / by_m, 0.0), 1.0) @@ -178,8 +179,8 @@ def stacked_on( def grade(outcome: Outcome) -> float: with recording(outcome) as store: try: - t = last_body_transform(store, top).translation - b = last_body_transform(store, base).translation + t = last_body_transform(store, top).transform.translation + b = last_body_transform(store, base).transform.translation except LookupError: return 0.0 if not rise_m[0] <= t.z - b.z <= rise_m[1]: @@ -192,8 +193,8 @@ def grade(outcome: Outcome) -> float: def opened_door(outcome: Outcome) -> float: with recording(outcome) as store: try: - frame = last_body_transform(store, "Door_frame").rotation.to_rotation_matrix() - panel = last_body_transform(store, "Door_door").rotation.to_rotation_matrix() + frame = transform_matrix(last_body_transform(store, "Door_frame").transform)[:3, :3] + panel = transform_matrix(last_body_transform(store, "Door_door").transform)[:3, :3] except LookupError: return 0.0 relative = frame.T @ panel @@ -207,10 +208,10 @@ def placed_can(outcome: Outcome) -> float: marker = last_body_transform(store, "VisualCan_main") except LookupError: return 0.0 - delta = (can.translation - marker.translation).to_numpy() + delta = vector_array(can.transform.translation) - vector_array(marker.transform.translation) return float( np.all(np.abs(delta) <= [0.05, 0.075, 0.005]) - and can.rotation.to_rotation_matrix()[2, 2] > 0.98 + and transform_matrix(can.transform)[:3, :3][2, 2] > 0.98 ) @@ -221,12 +222,12 @@ def seated_nut(outcome: Outcome) -> float: peg = last_body_transform(store, "peg1") except LookupError: return 0.0 - delta = (nut.translation - peg.translation).to_numpy() + delta = vector_array(nut.transform.translation) - vector_array(peg.transform.translation) # The seated nut's centre is 2 cm below the exported peg body's origin. return float( np.linalg.norm(delta[:2]) < 0.007 and abs(delta[2] + 0.02) < 0.004 - and abs(nut.rotation.to_rotation_matrix()[2, 2]) > 0.98 + and abs(transform_matrix(nut.transform)[:3, :3][2, 2]) > 0.98 ) @@ -239,19 +240,23 @@ def hung_tool(outcome: Outcome) -> float: hole = last_body_transform(store, "tool_hole1_root") except LookupError: return 0.0 - stand_r = stand.rotation.to_rotation_matrix() - frame_r = frame.rotation.to_rotation_matrix() + stand_r = transform_matrix(stand.transform)[:3, :3] + frame_r = transform_matrix(frame.transform)[:3, :3] # Local offsets from the exported frame tip, stand slot and horizontal hook. - tip_world = frame.translation.to_numpy() + frame_r @ np.array([0.04375, 0, -0.13445]) - tip_in_stand = stand_r.T @ (tip_world - stand.translation.to_numpy()) - hole_in_frame = frame_r.T @ (hole.translation - frame.translation).to_numpy() + tip_world = vector_array(frame.transform.translation) + frame_r @ np.array( + [0.04375, 0, -0.13445] + ) + tip_in_stand = stand_r.T @ (tip_world - vector_array(stand.transform.translation)) + hole_in_frame = frame_r.T @ ( + vector_array(hole.transform.translation) - vector_array(frame.transform.translation) + ) return float( stand_r[2, 2] > 0.98 and np.dot(stand_r[:, 2], frame_r[:, 2]) > 0.98 and np.linalg.norm(tip_in_stand - np.array([0, 0.045, -0.07])) < 0.01 and -0.043 < hole_in_frame[0] < 0.04375 and math.hypot(hole_in_frame[1], hole_in_frame[2] - 0.08625) < 0.007 - and abs(np.dot(hole.rotation.to_rotation_matrix()[:, 2], frame_r[:, 0])) > 0.95 + and abs(np.dot(transform_matrix(hole.transform)[:3, :3][:, 2], frame_r[:, 0])) > 0.95 ) diff --git a/dimos/evals/suites/dimsim_house.py b/dimos/evals/suites/dimsim_house.py index c043d6da94f..d54cdd3fd5c 100644 --- a/dimos/evals/suites/dimsim_house.py +++ b/dimos/evals/suites/dimsim_house.py @@ -40,7 +40,7 @@ if TYPE_CHECKING: from dimos.e2e_tests.dim_sim_client import DimSimClient -BED = Vector3(x=-3.567, y=-1.332) +BED = Vector3(x=-3.567, y=-1.332, z=0.0) _HOUSE_TOUR = [ (3.881, 4.803), diff --git a/dimos/evals/suites/lib/generate.py b/dimos/evals/suites/lib/generate.py index c8c6929a284..8066574d242 100644 --- a/dimos/evals/suites/lib/generate.py +++ b/dimos/evals/suites/lib/generate.py @@ -24,6 +24,7 @@ from collections.abc import Sequence from dimos.memory.cli.dataset import open_dataset +from dimos.msgs.geometry import point_distance Row = dict[str, object] @@ -38,7 +39,7 @@ def displacement_rows(dataset: str, windows: Sequence[tuple[float, float]]) -> l obs = store.streams.odom.range_time(t1, t2).to_list() if len(obs) < 2: continue - d = (obs[-1].data.position - obs[0].data.position).length() + d = point_distance(obs[-1].data.pose.position, obs[0].data.pose.position) rows.append( { "id": f"{dataset}_disp_{t1:g}_{t2:g}", @@ -65,9 +66,9 @@ def path_length_rows(dataset: str, windows: Sequence[tuple[float, float]]) -> li for t1, t2 in windows: path, prev = 0.0, None for obs in store.streams.odom.range_time(t1, t2): - p = obs.data.position + p = obs.data.pose.position if prev is not None: - path += (p - prev).length() + path += point_distance(p, prev) prev = p if prev is None: continue diff --git a/dimos/evals/test_evals.py b/dimos/evals/test_evals.py index 251777c1a86..09cbe0d7c7f 100644 --- a/dimos/evals/test_evals.py +++ b/dimos/evals/test_evals.py @@ -29,6 +29,7 @@ from types import SimpleNamespace from typing import Any +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.std_msgs.msg import Header from langchain_core.language_models.fake_chat_models import FakeListChatModel @@ -85,8 +86,10 @@ def _pose(x: float, y: float) -> PoseStamped: return PoseStamped( - header=Header(frame_id="world"), - pose=Pose(position=Point(x=x, y=y), orientation=Quaternion(w=1.0)), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=x, y=y, z=0.0), orientation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0) + ), ) diff --git a/dimos/evals/test_mem2_wiring.py b/dimos/evals/test_mem2_wiring.py index 39cfaf9ac5a..846c436f6c6 100644 --- a/dimos/evals/test_mem2_wiring.py +++ b/dimos/evals/test_mem2_wiring.py @@ -26,6 +26,7 @@ from pathlib import Path +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.sensor_msgs.msg import Image from dimos_generated.std_msgs.msg import Header @@ -54,8 +55,10 @@ def _pose(x: float, y: float) -> PoseStamped: return PoseStamped( - header=Header(frame_id="world"), - pose=Pose(position=Point(x=x, y=y), orientation=Quaternion(w=1.0)), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=x, y=y, z=0.0), orientation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0) + ), ) diff --git a/dimos/evals/vqa/primitives/test_edge_tam.py b/dimos/evals/vqa/primitives/test_edge_tam.py index 48fcabe3c42..c98d7f2a591 100644 --- a/dimos/evals/vqa/primitives/test_edge_tam.py +++ b/dimos/evals/vqa/primitives/test_edge_tam.py @@ -16,6 +16,7 @@ from collections import Counter +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 from dimos_generated.sensor_msgs.msg import Image, PointCloud2 from dimos_generated.std_msgs.msg import Header @@ -158,17 +159,23 @@ def _frame( image = image_from_array( np.zeros((10, 10, 3), dtype=np.uint8), encoding="rgb8", - header=Header(stamp=time_from_seconds(10)), + header=Header(stamp=time_from_seconds(10), frame_id=""), ) return PointCloudFrame( index=index, image=image, pointcloud=pointcloud, - camera_info=camera_info_from_intrinsics(fx, fy, cx, cy, 10, 10, header=Header()), + camera_info=camera_info_from_intrinsics( + fx, fy, cx, cy, 10, 10, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="") + ), pointcloud_to_camera=TransformStamped( - header=Header(frame_id="optical"), + header=Header(frame_id="optical", stamp=Time(sec=0, nanosec=0)), child_frame_id="lidar", - transform=transform or Transform(rotation=Quaternion(w=1)), + transform=transform + or Transform( + rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + translation=Vector3(x=0.0, y=0.0, z=0.0), + ), ), image_observation_timestamp=10.0, pointcloud_observation_timestamp=9.98, @@ -195,7 +202,7 @@ def test_mask_estimator_batches_and_caches_without_pointcloud() -> None: image = image_from_array( np.zeros((10, 10, 3), dtype=np.uint8), encoding="rgb8", - header=Header(stamp=time_from_seconds(10)), + header=Header(stamp=time_from_seconds(10), frame_id=""), ) left_mask = np.zeros((10, 10), dtype=np.uint8) right_mask = np.zeros((10, 10), dtype=np.uint8) @@ -223,7 +230,7 @@ def test_mask_estimator_batches_and_caches_without_pointcloud() -> None: next_image = image_from_array( np.zeros((10, 10, 3), dtype=np.uint8), encoding="rgb8", - header=Header(stamp=time_from_seconds(11)), + header=Header(stamp=time_from_seconds(11), frame_id=""), ) estimator.estimate(next_image, "left person") @@ -242,7 +249,10 @@ def test_projects_transformed_points_and_rejects_invalid_projections() -> None: ) frame = _frame( pointcloud, - transform=Transform(translation=Vector3(x=1.0)), + transform=Transform( + translation=Vector3(x=1.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), cx=0.0, cy=0.0, ) diff --git a/dimos/evals/vqa/test_contracts.py b/dimos/evals/vqa/test_contracts.py index 79fbd9807b5..399fd4ed5c8 100644 --- a/dimos/evals/vqa/test_contracts.py +++ b/dimos/evals/vqa/test_contracts.py @@ -62,7 +62,9 @@ def _image(pixels: np.ndarray, *, ts: float = 0.0) -> Image: - return image_from_array(pixels, encoding="rgb8", header=Header(stamp=time_from_seconds(ts))) + return image_from_array( + pixels, encoding="rgb8", header=Header(stamp=time_from_seconds(ts), frame_id="") + ) class _TestVlModel(VlModel): diff --git a/dimos/evals/vqa/test_pointcloud_frame.py b/dimos/evals/vqa/test_pointcloud_frame.py index 119af072cff..98fa9b30c00 100644 --- a/dimos/evals/vqa/test_pointcloud_frame.py +++ b/dimos/evals/vqa/test_pointcloud_frame.py @@ -12,8 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dataclasses import asdict from pathlib import Path +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 from dimos_generated.std_msgs.msg import Header @@ -50,7 +53,13 @@ def _cloud(data, *, frame_id="", timestamp=0.0): def _calibration(fx, fy, cx, cy, width, height, *, frame_id=""): result = camera_info_from_intrinsics( - max(fx, 1.0), max(fy, 1.0), cx, cy, width, height, header=Header(frame_id=frame_id) + max(fx, 1.0), + max(fy, 1.0), + cx, + cy, + width, + height, + header=Header(frame_id=frame_id, stamp=Time(sec=0, nanosec=0)), ) result.k[0] = fx result.k[4] = fy @@ -62,8 +71,8 @@ def _transform(*, translation=None, frame_id="", child_frame_id="", ts=0.0): header=Header(frame_id=frame_id, stamp=time_from_seconds(ts)), child_frame_id=child_frame_id, transform=Transform( - translation=translation if translation is not None else Vector3(), - rotation=Quaternion(w=1.0), + translation=translation if translation is not None else Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), ), ) @@ -226,7 +235,7 @@ def test_load_uses_recorded_camera_info_and_tf( def rectify(source: Image, calibration: CameraInfo) -> tuple[Image, CameraInfo]: calibrations.append(calibration) - result = CameraInfo.decode(camera_info.encode()) + result = cdr_decode(cdr_encode(camera_info), CameraInfo) result.header.stamp = source.header.stamp return source, result @@ -238,7 +247,9 @@ def rectify(source: Image, calibration: CameraInfo) -> tuple[Image, CameraInfo]: with loader: frame = loader.load(0) - assert calibrations == [camera_info] + np.testing.assert_equal( + [asdict(value) for value in calibrations], [asdict(value) for value in [camera_info]] + ) assert frame.calibration_source == "recorded" assert frame.camera_info.header.stamp == image.header.stamp assert np.allclose( @@ -286,7 +297,7 @@ def test_load_applies_camera_rectification_to_pointcloud_transform( [0.0, 0.0, 1.0], ] ) - camera_info.r = rectification.ravel().tolist() + camera_info.r = rectification.ravel().astype(np.float64) with SqliteStore(path=dataset) as store: store.stream("color_image", Image).append(image, ts=10.0) store.stream("lidar", PointCloud2).append(cloud, ts=10.0) @@ -352,7 +363,7 @@ def test_load_selects_camera_info_for_each_image_timestamp( def rectify(source: Image, calibration: CameraInfo) -> tuple[Image, CameraInfo]: calibrations.append(calibration) - result = CameraInfo.decode(calibration.encode()) + result = cdr_decode(cdr_encode(calibration), CameraInfo) result.header.stamp = source.header.stamp return source, result @@ -361,6 +372,8 @@ def rectify(source: Image, calibration: CameraInfo) -> tuple[Image, CameraInfo]: early = loader.load(0) late = loader.load(1) - assert calibrations == list(camera_infos) + np.testing.assert_equal( + [asdict(value) for value in calibrations], [asdict(value) for value in list(camera_infos)] + ) assert early.camera_info.header.frame_id == "camera_early" assert late.camera_info.header.frame_id == "camera_late" diff --git a/dimos/experimental/isolated_python/test_generated_example.py b/dimos/experimental/isolated_python/test_generated_example.py index 53c59e8a3eb..35c01af79d9 100644 --- a/dimos/experimental/isolated_python/test_generated_example.py +++ b/dimos/experimental/isolated_python/test_generated_example.py @@ -17,6 +17,7 @@ from pathlib import Path +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.std_msgs.msg import Int32 from dimos.experimental.isolated_python.bootstrap import load_class @@ -33,10 +34,10 @@ def test_external_example_uses_generated_contract_and_cdr(monkeypatch): try: streams = {s.name: s.type for s in ExampleExternal.blueprint().blueprints[0].streams} assert streams["value"] is Int32 and streams["doubled"] is Int32 - runtime._publish_doubled(Int32.decode(Int32(data=-7).encode())) + runtime._publish_doubled(cdr_decode(cdr_encode(Int32(data=-7)), Int32)) assert len(received) == 1 assert type(received[0]) is Int32 - assert Int32.decode(received[0].encode()).data == -21 + assert cdr_decode(cdr_encode(received[0]), Int32).data == -21 assert runtime.get_multiplier() == 3 finally: runtime.stop() diff --git a/dimos/experimental/memory/README.md b/dimos/experimental/memory/README.md index 29ed5ea85e4..31cdc591350 100644 --- a/dimos/experimental/memory/README.md +++ b/dimos/experimental/memory/README.md @@ -1,14 +1,14 @@ # Experimental Native Memory Recorder The Rust recorder is an experimental high-throughput alternative to the Python -Memory2 recorder. It remains compatible with the existing Python readers while -its API and operational behavior are evaluated. Experimental imports may change -without compatibility aliases. +Memory recorder. Python supplies each stream's generated type name and complete +ROS message definition. The native process records custom messages without a +message-specific decoder or rebuild. ## Build and runtime packaging -The recorder is built as a locked Nix package. Nix supplies Rust, CMake, NASM, -SQLite, and the native libraries used by TurboJPEG, so none of those tools or +The recorder is built as a locked Nix package. Nix supplies Rust, Python for build-time message generation, +CMake, and SQLite, so none of those tools or development packages need to be installed on the host. The Python module resolves the package through Nix before each launch, so its @@ -40,8 +40,7 @@ from dimos.experimental.memory.rust_recorder import ( RustRecorder, RustSqliteStoreConfig, ) -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 class SensorRecorder(RustRecorder): @@ -52,19 +51,19 @@ class SensorRecorder(RustRecorder): sensor_recorder = SensorRecorder.blueprint( store=RustSqliteStoreConfig(path="session.db"), encoding_threads=4, - stream_codecs={"lidar": "lz4+lcm"}, + stream_codecs={"lidar": "lz4+cdr"}, ) ``` -SQLite supports LCM-backed messages and the `lcm`, `jpeg`, `lz4+lcm`, and `json` -storage codecs. Images default to JPEG quality 50. Configure depth or other -lossless streams explicitly with `lz4+lcm`. The resulting artifact opens and -replays through the stable Python `SqliteStore` API. +SQLite stores generated CDR bytes by default. `lz4+cdr` adds lossless blob +compression. Raw images retain their encoding and pixels, including depth data. +For lossy compression, publish an explicit generated CompressedImage value. +The resulting artifact opens and replays through Python's SqliteStore. ## MCAP -Select MCAP to write the same Memory2 storage encodings into an indexed, -portable container: +Select MCAP for ROS2-profile recordings with CDR channels and embedded, complete +`ros2msg` schemas: ```python from dimos.experimental.memory.rust_recorder import RustMcapStoreConfig @@ -72,36 +71,30 @@ from dimos.experimental.memory.rust_recorder import RustMcapStoreConfig mcap_recorder = SensorRecorder.blueprint( store=RustMcapStoreConfig(path="session.mcap"), encoding_threads=4, - stream_codecs={"lidar": "lz4+lcm"}, ) ``` -MCAP stores each stream's selected `lcm`, `jpeg`, `lz4+lcm`, or `json` representation -in indexed Zstd chunks. Source time is the MCAP publish time and recorder -reception time is the log time. JPEG channels decode automatically. Supply -trusted codecs explicitly for LCM channels instead of trusting artifact -metadata: +The Python reader now expects standard ROS2-profile CDR channels with embedded +`ros2msg` schemas. Installed generated message packages supply the decoders: ```python -from dimos.memory.codecs.lcm import LcmCodec -from dimos.memory.codecs.lz4 import Lz4Codec from dimos.memory.store.mcap import McapStore -from dimos.msgs.sensor_msgs.Imu import Imu -store = McapStore( - path="session.mcap", - codecs={"imu": Lz4Codec(LcmCodec(Imu))}, -) +store = McapStore(path="session.mcap") ``` -Append mode remains unsupported for MCAP. +MCAP applies Zstd chunk compression. It rejects per-payload wrappers such as +`lz4+cdr`. Standard compressed-image messages use their own declared schema. +See `examples/message-codegen/demo_native_recording.py` for a hardware-free +external-message recording demo on both LCM and Zenoh. -MCAP times must be finite, nonnegative seconds that fit an unsigned 64-bit -nanosecond timestamp. Recording fails on an unrepresentable source or reception -time instead of clamping it. Zero is valid. SQLite can retain negative source -times, such as event-relative timestamps. +Append mode remains unsupported for MCAP. -Both backends preserve source timestamps for common stamped messages. Arbitrary +Both backends preserve source timestamps for supported stamped messages. MCAP +keeps source and reception time as integer nanoseconds; SQLite exposes floating +seconds for its existing query API while message payload stamps remain exact. +Unknown timestamp layouts use reception time. Negative source times work in +SQLite; MCAP rejects them because its time fields are unsigned. Arbitrary pickle payloads, Python `pose_setter_for` hooks, and spatial pose attachment remain Python-recorder features; unsupported combinations fail during startup. @@ -109,18 +102,13 @@ remain Python-recorder features; unsupported combinations fail during startup. ## JSON documents The `json` codec accepts the existing `std_msgs.String` transport and stores its -JSON text as UTF-8 without an LCM envelope. Ordinary String streams retain their -selected LCM codec and reception timestamp. Configure an event source timestamp +JSON text as UTF-8 without a CDR envelope. Ordinary String streams retain their +selected CDR codec and reception timestamp. Configure an event source timestamp explicitly with `stream_timestamp_fields={"events": "ts"}` and select the codec with `stream_codecs={"events": "json"}`. Missing, nonnumeric or nonfinite selected timestamps fail recording. Without a selected field, JSON uses reception time. -Unconfigured JSON options are omitted from the native launch configuration, so -existing non-JSON streams remain compatible with older recorder binaries. JSON -streams require a binary built with JSON codec support. - -`stream_json_schemas={"events": schema}` embeds a JSON Schema in MCAP, whose -channel uses `message_encoding="json"`. SQLite retains the same UTF-8 document -in its existing blob table, the event time in `ts`, and reception time in tags. -Native JSON/String MCAP channels decode and replay as String through the existing -memory APIs without importing an application-specific payload class. +JSON storage is supported by SQLite. The ROS2-profile MCAP backend accepts only +CDR channels with complete message schemas; record a typed custom event message +there instead of selecting the JSON storage codec. JSON source times are converted +to signed nanoseconds with checked range and retained in SQLite. diff --git a/dimos/experimental/memory/rust/Cargo.toml b/dimos/experimental/memory/rust/Cargo.toml index 2a7322024bb..6e21f588bbb 100644 --- a/dimos/experimental/memory/rust/Cargo.toml +++ b/dimos/experimental/memory/rust/Cargo.toml @@ -10,7 +10,7 @@ anyhow = "1" crossbeam-channel = "0.5" crossbeam-utils = "0.8" dimos-module = { path = "../../../../native/rust/dimos-module" } -dimos-generated-messages = { path = "../../../message_codegen" } +dimos-generated-messages = { path = "../../../../packages/dimos-generated/build/message-source/rust" } lz4_flex = "0.14.0" mcap = { version = "0.25.0", default-features = false, features = ["zstd"] } rayon = "1" diff --git a/dimos/experimental/memory/rust/flake.nix b/dimos/experimental/memory/rust/flake.nix index 368ebe32b3c..90c4d714227 100644 --- a/dimos/experimental/memory/rust/flake.nix +++ b/dimos/experimental/memory/rust/flake.nix @@ -17,37 +17,45 @@ "aarch64-darwin" ] (system: let pkgs = nixpkgs.legacyPackages.${system}; - nativeDeps = [pkgs.cmake pkgs.nasm pkgs.pkg-config]; + nativeDeps = [pkgs.cmake pkgs.nasm pkgs.pkg-config pkgs.python3]; systemDeps = [pkgs.sqlite pkgs.sqlite.dev] ++ pkgs.lib.optionals pkgs.stdenv.hostPlatform.isDarwin [pkgs.libiconv]; + messageSource = import ../../../../native/message-source.nix { inherit pkgs; }; + workspaceSource = pkgs.lib.fileset.toSource { + root = ../../../..; + fileset = pkgs.lib.fileset.unions [ + ../../../../Cargo.lock + ../../../../Cargo.toml + ../../../../dimos/experimental/memory/rust + ../../../../native/rust/dimos-module + ../../../../native/rust/dimos-lcm-transport + ../../../../native/rust/dimos-module-macros + ../../../../dimos/mapping/ray_tracing/rust + ../../../../dimos/mapping/ray_tracing/rust/py + ../../../../dimos/navigation/global_planner/mls_planner/rust + ../../../../dimos/navigation/global_planner/mls_planner/rust/py + ../../../../dimos/hardware/sensors/lidar/livox/rust + ../../../../dimos/hardware/sensors/lidar/pointlio/rust + ../../../../dimos/hardware/sensors/lidar/virtual_mid360 + ../../../../dimos/hardware/sensors/camera/v4l2/rust + ../../../../examples/native-modules/rust + ]; + }; dimos-memory-recorder = pkgs.rustPlatform.buildRustPackage { pname = "dimos-memory-recorder"; version = "0.1.0"; - src = pkgs.lib.fileset.toSource { - root = ../../../..; - fileset = pkgs.lib.fileset.unions [ - ../../../../Cargo.lock - ../../../../Cargo.toml - ../../../../dimos/experimental/memory/rust - ../../../../native/rust/dimos-module - ../../../../native/rust/dimos-module-macros - ../../../../dimos/mapping/ray_tracing/rust - ../../../../dimos/mapping/ray_tracing/rust/py - ../../../../dimos/navigation/global_planner/mls_planner/rust - ../../../../dimos/navigation/global_planner/mls_planner/rust/py - ../../../../dimos/hardware/sensors/lidar/livox/rust - ../../../../dimos/hardware/sensors/lidar/pointlio/rust - ../../../../dimos/hardware/sensors/lidar/virtual_mid360 - ../../../../dimos/hardware/sensors/camera/v4l2/rust - ../../../../examples/native-modules/rust - ]; - }; + src = pkgs.runCommand "dimos-memory-recorder-source" {} '' + mkdir -p $out + cp -r ${workspaceSource}/. $out/ + chmod -R u+w $out + mkdir -p $out/packages/dimos-generated/build/message-source + cp -r ${messageSource.source}/rust $out/packages/dimos-generated/build/message-source/rust + ''; cargoLock = { lockFile = ../../../../Cargo.lock; outputHashes = { - "dimos-lcm-0.1.0" = "sha256-GGkx4Mn6NYP6KZecmoRLKGWIih/+y8OgNn12DeXX6n8="; "pointlio-core-0.1.0" = "sha256-FjUcsZVI8nTcmKw8WhdxBuHelBjOCeJmyEsyqRxDvX0="; }; }; diff --git a/dimos/experimental/memory/rust/src/decoding.rs b/dimos/experimental/memory/rust/src/decoding.rs index 79b6b98a5c1..fa0f65e1269 100644 --- a/dimos/experimental/memory/rust/src/decoding.rs +++ b/dimos/experimental/memory/rust/src/decoding.rs @@ -13,51 +13,54 @@ // limitations under the License. use anyhow::{Context, Result}; -use lcm_msgs::foxglove_msgs::CompressedVideo; -use lcm_msgs::geometry_msgs::{ - PointStamped, PoseStamped, PoseWithCovarianceStamped, TwistStamped, TwistWithCovarianceStamped, - WrenchStamped, +use dimos_generated_messages::dimos_msgs::msg::line_segments3_d::LineSegments3D; +use dimos_generated_messages::foxglove_msgs::msg::compressed_video::CompressedVideo; +use dimos_generated_messages::geometry_msgs::msg::{ + point_stamped::PointStamped, pose_stamped::PoseStamped, + pose_with_covariance_stamped::PoseWithCovarianceStamped, transform_stamped::TransformStamped, + twist_stamped::TwistStamped, twist_with_covariance_stamped::TwistWithCovarianceStamped, + wrench_stamped::WrenchStamped, }; -use lcm_msgs::nav_msgs::{OccupancyGrid, Odometry, Path}; -use lcm_msgs::sensor_msgs::{ - CameraInfo, CompressedImage, Image, Imu, JointState, Joy, PointCloud2, +use dimos_generated_messages::nav_msgs::msg::{ + occupancy_grid::OccupancyGrid, odometry::Odometry, path::Path, +}; +use dimos_generated_messages::sensor_msgs::msg::{ + camera_info::CameraInfo, compressed_image::CompressedImage, image::Image, imu::Imu, + joint_state::JointState, joy::Joy, point_cloud2::PointCloud2, +}; +use dimos_generated_messages::tf2_msgs::msg::tf_message::TFMessage; +use dimos_generated_messages::vision_msgs::msg::{ + detection2_d::Detection2D, detection2_d_array::Detection2DArray, detection3_d::Detection3D, + detection3_d_array::Detection3DArray, }; -use lcm_msgs::tf2_msgs::TFMessage; -use lcm_msgs::vision_msgs::{Detection2D, Detection2DArray, Detection3D, Detection3DArray}; - -use crate::{Codec, StreamConfig, IMAGE_PAYLOAD_TYPE}; -/// Transport-decoded payload consumed by storage codecs. -#[derive(Debug)] -pub(crate) enum DecodedPayload { - /// An LCM payload whose storage codecs operate on its canonical wire form. - Lcm(Vec), - /// UTF-8 JSON after removing the existing String transport envelope. - Json(Vec), - /// Images stay typed so JPEG encoding reuses the transport decode. - Image(Image), -} +use crate::{Codec, StreamConfig}; +use dimos_generated_messages::dimos_msgs::msg::{ + contacts::Contacts, region_bounds::RegionBounds, region_line_segments3_d::RegionLineSegments3D, + region_point_cloud2::RegionPointCloud2, +}; +use dimos_generated_messages::std_msgs::msg::string::String as StringMessage; +use dimos_module::cdr; /// One decoded and timestamped observation before storage encoding. #[derive(Debug)] pub(crate) struct DecodedObservation { - pub(crate) ts: f64, - pub(crate) payload: DecodedPayload, + pub(crate) ts: i64, + pub(crate) payload: Vec, } /// Decode a transport packet once and normalize it into storage observations. pub(crate) fn decode( stream: &StreamConfig, data: &[u8], - reception_ts: f64, + reception_ts: i64, ) -> Result> { if stream.codec == Codec::Json { anyhow::ensure!( - stream.payload_type == "dimos.msgs.std_msgs.String.String", + stream.schema_name == "std_msgs/msg/String", "JSON codec requires std_msgs.String" ); - let message = - lcm_msgs::std_msgs::String::decode(data).context("invalid String envelope")?; + let message = cdr::decode::(data).context("invalid CDR String envelope")?; let document: serde_json::Value = serde_json::from_str(&message.data).context("invalid JSON document")?; let ts = if let Some(field) = &stream.timestamp_field { @@ -66,137 +69,134 @@ pub(crate) fn decode( .and_then(serde_json::Value::as_f64) .context("JSON timestamp field must be a number")?; anyhow::ensure!(ts.is_finite(), "JSON timestamp must be finite"); - ts + anyhow::ensure!( + ts >= i64::MIN as f64 / 1e9 && ts < i64::MAX as f64 / 1e9, + "JSON timestamp exceeds signed nanosecond range" + ); + (ts * 1e9).round() as i64 } else { reception_ts }; return Ok(vec![DecodedObservation { ts, - payload: DecodedPayload::Json(message.data.into_bytes()), + payload: message.data.into_bytes(), }]); } if stream.is_tf() { return decode_tf(data, reception_ts); } - if stream.payload_type == IMAGE_PAYLOAD_TYPE { - let image = Image::decode(data).context("invalid LCM Image")?; - let ts = header_timestamp( - image.header.stamp.sec, - image.header.stamp.nsec, - reception_ts, - ); - return Ok(vec![DecodedObservation { - ts, - payload: DecodedPayload::Image(image), - }]); - } - - let ts = source_timestamp(&stream.payload_type, data, reception_ts)?; + let ts = source_timestamp(&stream.schema_name, data, reception_ts)?; Ok(vec![DecodedObservation { ts, - payload: DecodedPayload::Lcm(data.to_vec()), + payload: data.to_vec(), }]) } -fn decode_tf(data: &[u8], reception_ts: f64) -> Result> { - let message = TFMessage::decode(data).context("invalid LCM TFMessage")?; - Ok(message +fn decode_tf(data: &[u8], _reception_ts: i64) -> Result> { + let message = cdr::decode::(data).context("invalid CDR TFMessage")?; + message .transforms .into_iter() - .map(|transform| DecodedObservation { - ts: header_timestamp( - transform.header.stamp.sec, - transform.header.stamp.nsec, - reception_ts, - ), - payload: DecodedPayload::Lcm( - TFMessage { + .map(|transform| { + Ok(DecodedObservation { + ts: header_timestamp(transform.header.stamp.sec, transform.header.stamp.nanosec)?, + payload: cdr::encode(&TFMessage { transforms: vec![transform], - } - .encode(), - ), + })?, + }) }) - .collect()) + .collect() } -/// Read source time while transport-decoding common stamped LCM payloads. +/// Read source time while transport-decoding common stamped CDR payloads. /// -/// Unknown LCM payloads remain opaque and use reception time, matching the -/// Python recorder's fallback for messages without a ``ts`` attribute. -fn source_timestamp(payload_type: &str, data: &[u8], reception_ts: f64) -> Result { +/// Unknown CDR payloads remain opaque and use reception time, matching the +/// Python recorder's fallback for messages without a recognized timestamp layout. +fn source_timestamp(payload_type: &str, data: &[u8], reception_ts: i64) -> Result { macro_rules! stamped { ($message_type:ty) => {{ - let message = <$message_type>::decode(data) - .with_context(|| format!("invalid LCM {payload_type}"))?; - (message.header.stamp.sec, message.header.stamp.nsec) + let message = cdr::decode::<$message_type>(data) + .with_context(|| format!("invalid CDR {payload_type}"))?; + (message.header.stamp.sec, message.header.stamp.nanosec) }}; } let (sec, nsec) = match payload_type { - "dimos.msgs.geometry_msgs.PointStamped.PointStamped" => stamped!(PointStamped), - "dimos.msgs.geometry_msgs.PoseStamped.PoseStamped" => stamped!(PoseStamped), - "dimos.msgs.geometry_msgs.PoseWithCovarianceStamped.PoseWithCovarianceStamped" => { + "geometry_msgs/msg/TransformStamped" => stamped!(TransformStamped), + "dimos_msgs/msg/LineSegments3D" => stamped!(LineSegments3D), + "dimos_msgs/msg/RegionBounds" => stamped!(RegionBounds), + "dimos_msgs/msg/Contacts" => stamped!(Contacts), + "dimos_msgs/msg/RegionPointCloud2" => { + let message = cdr::decode::(data)?; + ( + message.cloud.header.stamp.sec, + message.cloud.header.stamp.nanosec, + ) + } + "dimos_msgs/msg/RegionLineSegments3D" => { + let message = cdr::decode::(data)?; + ( + message.lines.header.stamp.sec, + message.lines.header.stamp.nanosec, + ) + } + "geometry_msgs/msg/PointStamped" => stamped!(PointStamped), + "geometry_msgs/msg/PoseStamped" => stamped!(PoseStamped), + "geometry_msgs/msg/PoseWithCovarianceStamped" => { stamped!(PoseWithCovarianceStamped) } - "dimos.msgs.geometry_msgs.TwistStamped.TwistStamped" => stamped!(TwistStamped), - "dimos.msgs.geometry_msgs.TwistWithCovarianceStamped.TwistWithCovarianceStamped" => { + "geometry_msgs/msg/TwistStamped" => stamped!(TwistStamped), + "geometry_msgs/msg/TwistWithCovarianceStamped" => { stamped!(TwistWithCovarianceStamped) } - "dimos.msgs.geometry_msgs.WrenchStamped.WrenchStamped" => stamped!(WrenchStamped), - "dimos.msgs.nav_msgs.LineSegments3D.LineSegments3D" | "dimos.msgs.nav_msgs.Path.Path" => { + "geometry_msgs/msg/WrenchStamped" => stamped!(WrenchStamped), + "nav_msgs/msg/Path" => { stamped!(Path) } - "dimos.msgs.nav_msgs.OccupancyGrid.OccupancyGrid" => stamped!(OccupancyGrid), - "dimos.msgs.nav_msgs.Odometry.Odometry" => stamped!(Odometry), - "dimos.msgs.sensor_msgs.CameraInfo.CameraInfo" => stamped!(CameraInfo), - "dimos.msgs.sensor_msgs.CompressedImage.CompressedImage" => stamped!(CompressedImage), - "dimos.msgs.sensor_msgs.Imu.Imu" => stamped!(Imu), - "dimos.msgs.sensor_msgs.JointState.JointState" => stamped!(JointState), - "dimos.msgs.sensor_msgs.Joy.Joy" => stamped!(Joy), - "dimos.msgs.sensor_msgs.PointCloud2.PointCloud2" => stamped!(PointCloud2), - "dimos.msgs.vision_msgs.Detection2D.Detection2D" => stamped!(Detection2D), - "dimos.msgs.vision_msgs.Detection2DArray.Detection2DArray" => { + "nav_msgs/msg/OccupancyGrid" => stamped!(OccupancyGrid), + "nav_msgs/msg/Odometry" => stamped!(Odometry), + "sensor_msgs/msg/Image" => stamped!(Image), + "sensor_msgs/msg/CameraInfo" => stamped!(CameraInfo), + "sensor_msgs/msg/CompressedImage" => stamped!(CompressedImage), + "sensor_msgs/msg/Imu" => stamped!(Imu), + "sensor_msgs/msg/JointState" => stamped!(JointState), + "sensor_msgs/msg/Joy" => stamped!(Joy), + "sensor_msgs/msg/PointCloud2" => stamped!(PointCloud2), + "vision_msgs/msg/Detection2D" => stamped!(Detection2D), + "vision_msgs/msg/Detection2DArray" => { stamped!(Detection2DArray) } - "dimos.msgs.vision_msgs.Detection3D.Detection3D" => stamped!(Detection3D), - "dimos.msgs.vision_msgs.Detection3DArray.Detection3DArray" => { + "vision_msgs/msg/Detection3D" => stamped!(Detection3D), + "vision_msgs/msg/Detection3DArray" => { stamped!(Detection3DArray) } - "dimos.msgs.foxglove_msgs.CompressedVideo.CompressedVideo" => { - let message = CompressedVideo::decode(data) - .context("invalid LCM foxglove_msgs.CompressedVideo")?; + "foxglove_msgs/msg/CompressedVideo" => { + let message = cdr::decode::(data) + .context("invalid CDR foxglove_msgs.CompressedVideo")?; (message.timestamp.sec, message.timestamp.nanosec) } - "dimos.msgs.geometry_msgs.Transform.Transform" => { - let message = TFMessage::decode(data).context("invalid LCM TFMessage")?; - let Some(transform) = message.transforms.first() else { - return Ok(reception_ts); - }; - (transform.header.stamp.sec, transform.header.stamp.nsec) - } _ => return Ok(reception_ts), }; - Ok(header_timestamp(sec, nsec, reception_ts)) + header_timestamp(sec, nsec) } -pub(crate) fn header_timestamp(sec: i32, nsec: i32, fallback: f64) -> f64 { - if sec > 0 { - f64::from(sec) + f64::from(nsec) / 1_000_000_000.0 - } else { - fallback - } +pub(crate) fn header_timestamp(sec: i32, nanosec: u32) -> Result { + anyhow::ensure!(nanosec < 1_000_000_000, "invalid ROS nanosecond field"); + Ok(i64::from(sec) * 1_000_000_000 + i64::from(nanosec)) } #[cfg(test)] mod tests { use super::*; - use lcm_msgs::std_msgs::String as LcmString; + use dimos_generated_messages::std_msgs::msg::string::String as StringMessage; fn json_stream(field: Option<&str>) -> StreamConfig { StreamConfig { port: "events".into(), name: "events".into(), - payload_type: "dimos.msgs.std_msgs.String.String".into(), + payload_type: "dimos_generated.std_msgs.msg.String".into(), + schema_name: "std_msgs/msg/String".into(), + schema_definition: "string data\n".into(), codec: Codec::Json, timestamp_field: field.map(str::to_string), json_schema: None, @@ -206,36 +206,101 @@ mod tests { #[test] fn json_codec_preserves_document_and_explicit_source_time() { let text = r#"{"sent":42.25,"label":"拿起积木","counter":2147483647}"#; - let message = LcmString { data: text.into() }; - let mut result = decode(&json_stream(Some("sent")), &message.encode(), 99.0).unwrap(); + let message = StringMessage { data: text.into() }; + let mut result = decode( + &json_stream(Some("sent")), + &cdr::encode(&message).unwrap(), + 99_000_000_000, + ) + .unwrap(); let obs = result.pop().unwrap(); - assert_eq!(obs.ts, 42.25); - let DecodedPayload::Json(data) = obs.payload else { - panic!("not JSON") - }; + assert_eq!(obs.ts, 42_250_000_000); + let data = obs.payload; assert_eq!(data, text.as_bytes()); assert_eq!( - decode(&json_stream(None), &message.encode(), 99.0).unwrap()[0].ts, - 99.0 + decode( + &json_stream(None), + &cdr::encode(&message).unwrap(), + 99_000_000_000 + ) + .unwrap()[0] + .ts, + 99_000_000_000 ); } #[test] fn json_codec_rejects_malformed_or_invalid_explicit_time() { for data in ["not json", r#"{"sent":"42"}"#, "{}", r#"{"sent":NaN}"#] { - let message = LcmString { data: data.into() }; - assert!(decode(&json_stream(Some("sent")), &message.encode(), 99.0).is_err()); + let message = StringMessage { data: data.into() }; + assert!(decode( + &json_stream(Some("sent")), + &cdr::encode(&message).unwrap(), + 99_000_000_000 + ) + .is_err()); } } #[test] fn ordinary_strings_remain_opaque_and_use_reception_time() { let mut stream = json_stream(None); - stream.codec = Codec::Lcm; + stream.codec = Codec::Cdr; for data in ["not json", r#"{"ts":42.25}"#] { - let message = LcmString { data: data.into() }; - let result = decode(&stream, &message.encode(), 99.0).unwrap(); - assert_eq!(result[0].ts, 99.0); + let message = StringMessage { data: data.into() }; + let result = decode(&stream, &cdr::encode(&message).unwrap(), 99_000_000_000).unwrap(); + assert_eq!(result[0].ts, 99_000_000_000); + } + } + #[test] + fn region_envelopes_use_nested_source_time_without_losing_signed_id() { + use dimos_generated_messages::builtin_interfaces::msg::time::Time; + use dimos_generated_messages::std_msgs::msg::header::Header; + let header = Header { + stamp: Time { + sec: 0, + nanosec: 34, + }, + frame_id: "map".into(), + }; + let cloud = RegionPointCloud2 { + region_id: -196603, + cloud: PointCloud2 { + header: header.clone(), + height: 1, + width: 0, + fields: vec![], + is_bigendian: false, + point_step: 12, + row_step: 0, + data: vec![], + is_dense: true, + }, + }; + let lines = RegionLineSegments3D { + region_id: cloud.region_id, + lines: LineSegments3D { + header, + segments: vec![], + }, + }; + for (name, bytes) in [ + ( + "dimos_msgs/msg/RegionPointCloud2", + cdr::encode(&cloud).unwrap(), + ), + ( + "dimos_msgs/msg/RegionLineSegments3D", + cdr::encode(&lines).unwrap(), + ), + ] { + assert_eq!(source_timestamp(name, &bytes, 99_000_000_000).unwrap(), 34); } + assert_eq!( + cdr::decode::(&cdr::encode(&cloud).unwrap()) + .unwrap() + .region_id, + -196603 + ); } } diff --git a/dimos/experimental/memory/rust/src/encoding.rs b/dimos/experimental/memory/rust/src/encoding.rs index 00eacc3823b..f84c16138e1 100644 --- a/dimos/experimental/memory/rust/src/encoding.rs +++ b/dimos/experimental/memory/rust/src/encoding.rs @@ -14,19 +14,15 @@ use std::io::Write; -use anyhow::{anyhow, Context, Result}; -use lcm_msgs::sensor_msgs::Image; +use anyhow::{Context, Result}; use lz4_flex::frame::FrameEncoder; -use turbojpeg::{PixelFormat, Subsamp}; -use crate::decoding::{DecodedObservation, DecodedPayload}; +use crate::decoding::DecodedObservation; use crate::{Codec, StreamConfig}; -pub(crate) const JPEG_QUALITY: i32 = 50; - #[derive(Debug)] pub(crate) struct StoredObservation { - pub(crate) ts: f64, + pub(crate) ts: i64, pub(crate) data: Vec, } @@ -35,13 +31,8 @@ pub(crate) fn encode( observation: DecodedObservation, ) -> Result { let data = match stream.codec { - Codec::Lcm => lcm_encode(observation.payload)?, - Codec::Json => match observation.payload { - DecodedPayload::Json(data) => data, - _ => return Err(anyhow!("JSON storage codec requires a JSON payload")), - }, - Codec::Lz4Lcm => lz4_frame(&lcm_encode(observation.payload)?)?, - Codec::Jpeg => jpeg_encode(observation.payload)?, + Codec::Cdr | Codec::Json => observation.payload, + Codec::Lz4Cdr => lz4_frame(&observation.payload)?, }; Ok(StoredObservation { ts: observation.ts, @@ -49,72 +40,6 @@ pub(crate) fn encode( }) } -fn lcm_encode(payload: DecodedPayload) -> Result> { - Ok(match payload { - DecodedPayload::Lcm(data) => data, - DecodedPayload::Image(image) => image.encode(), - DecodedPayload::Json(_) => return Err(anyhow!("LCM codec cannot store JSON payload")), - }) -} - -fn jpeg_encode(payload: DecodedPayload) -> Result> { - let DecodedPayload::Image(mut image) = payload else { - return Err(anyhow!("JPEG storage codec requires an Image payload")); - }; - if image.encoding == "jpeg" { - return Ok(image.encode()); - } - - let format = pixel_format(&image.encoding)?; - let pixels = normalize_pixels(&image)?; - let source = turbojpeg::Image { - pixels: pixels.as_slice(), - width: image.width as usize, - pitch: image.width as usize * format.size(), - height: image.height as usize, - format, - }; - let jpeg = turbojpeg::compress(source, JPEG_QUALITY, Subsamp::Sub2x1) - .context("TurboJPEG compression failed")?; - image.encoding = "jpeg".to_string(); - image.is_bigendian = 0; - image.step = 0; - image.data = jpeg.to_vec(); - Ok(image.encode()) -} - -fn pixel_format(encoding: &str) -> Result { - match encoding { - "rgb8" => Ok(PixelFormat::RGB), - "bgr8" => Ok(PixelFormat::BGR), - "rgba8" => Ok(PixelFormat::RGBA), - "bgra8" => Ok(PixelFormat::BGRA), - "mono8" | "mono16" | "16UC1" | "16SC1" => Ok(PixelFormat::GRAY), - other => Err(anyhow!( - "JPEG codec does not support image encoding {other:?}" - )), - } -} - -fn normalize_pixels(image: &Image) -> Result> { - match image.encoding.as_str() { - "mono16" | "16UC1" | "16SC1" => { - if !image.data.len().is_multiple_of(2) { - return Err(anyhow!("16-bit image has an odd byte count")); - } - let high_byte = usize::from(image.is_bigendian == 0); - Ok(image - .data - .as_chunks::<2>() - .0 - .iter() - .map(|pixel| pixel[high_byte]) - .collect()) - } - _ => Ok(image.data.clone()), - } -} - pub(crate) fn lz4_frame(data: &[u8]) -> Result> { let mut encoder = FrameEncoder::new(Vec::new()); encoder.write_all(data).context("LZ4 compression failed")?; diff --git a/dimos/experimental/memory/rust/src/lib.rs b/dimos/experimental/memory/rust/src/lib.rs index 74df3f3f9f9..8022e620679 100644 --- a/dimos/experimental/memory/rust/src/lib.rs +++ b/dimos/experimental/memory/rust/src/lib.rs @@ -35,8 +35,9 @@ mod decoding; mod encoding; pub mod store; -const TF_PAYLOAD_TYPE: &str = "dimos.msgs.tf2_msgs.TFMessage.TFMessage"; -const IMAGE_PAYLOAD_TYPE: &str = "dimos.msgs.sensor_msgs.Image.Image"; +const TF_PAYLOAD_TYPE: &str = "tf2_msgs/msg/TFMessage"; +#[cfg(test)] +const IMAGE_PAYLOAD_TYPE: &str = "sensor_msgs/msg/Image"; const QUEUE_CAPACITY: usize = 256; const WRITE_BATCH_SIZE: usize = 128; const FLUSH_INTERVAL: Duration = Duration::from_millis(100); @@ -44,22 +45,10 @@ const FLUSH_INTERVAL: Duration = Duration::from_millis(100); #[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)] #[serde(rename_all = "lowercase")] pub enum Codec { - Lcm, + Cdr, Json, - Jpeg, - #[serde(rename = "lz4+lcm")] - Lz4Lcm, -} - -impl Codec { - pub(crate) fn id(self) -> &'static str { - match self { - Self::Lcm => "lcm", - Self::Json => "json", - Self::Jpeg => "jpeg", - Self::Lz4Lcm => "lz4+lcm", - } - } + #[serde(rename = "lz4+cdr")] + Lz4Cdr, } #[derive(Clone, Debug, Deserialize, Serialize)] @@ -68,6 +57,8 @@ pub struct StreamConfig { pub port: String, pub name: String, pub payload_type: String, + pub schema_name: String, + pub schema_definition: String, pub codec: Codec, #[serde(default)] pub timestamp_field: Option, @@ -77,7 +68,7 @@ pub struct StreamConfig { impl StreamConfig { pub fn is_tf(&self) -> bool { - self.payload_type == TF_PAYLOAD_TYPE + self.schema_name == TF_PAYLOAD_TYPE } } @@ -230,7 +221,7 @@ impl RecorderEngine { struct EncodedBatch { sequence: u64, stream: Arc, - reception_ts: f64, + reception_ts: i64, observations: Result, String>, } @@ -244,7 +235,7 @@ pub struct WriterStats { fn process( stream: &StreamConfig, data: &[u8], - reception_ts: f64, + reception_ts: i64, ) -> Result> { decoding::decode(stream, data, reception_ts)? .into_iter() @@ -252,11 +243,13 @@ fn process( .collect() } -fn wall_time() -> f64 { +fn wall_time() -> i64 { SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap_or_default() - .as_secs_f64() + .as_nanos() + .try_into() + .expect("wall time exceeds signed nanosecond range") } struct RecorderInput { @@ -456,15 +449,80 @@ fn write_ready( mod tests { use std::io::Read; - use lcm_msgs::sensor_msgs::{Image, Imu}; - use lcm_msgs::std_msgs::{Header, Time}; - use lcm_msgs::tf2_msgs::TFMessage; + use dimos_generated_messages::builtin_interfaces::msg::time::Time; + use dimos_generated_messages::geometry_msgs::msg::{ + quaternion::Quaternion, transform::Transform, transform_stamped::TransformStamped, + vector3::Vector3, + }; + use dimos_generated_messages::sensor_msgs::msg::{ + compressed_image::CompressedImage, image::Image, imu::Imu, + }; + use dimos_generated_messages::std_msgs::msg::header::Header; + use dimos_generated_messages::tf2_msgs::msg::tf_message::TFMessage; + use dimos_module::cdr; use lz4_flex::frame::FrameDecoder; use rusqlite::Connection; use tempfile::NamedTempFile; use super::*; + fn schema(name: &str) -> String { + let schemas: std::collections::BTreeMap = + serde_json::from_str(dimos_generated_messages::ROS2MSG_SCHEMAS).unwrap(); + schemas[name].clone() + } + + fn imu() -> Imu { + Imu { + header: Header { + stamp: Time { sec: 0, nanosec: 0 }, + frame_id: String::new(), + }, + orientation: Quaternion { + x: 0.0, + y: 0.0, + z: 0.0, + w: 0.0, + }, + orientation_covariance: [0.0; 9], + angular_velocity: Vector3 { + x: 0.0, + y: 0.0, + z: 0.0, + }, + angular_velocity_covariance: [0.0; 9], + linear_acceleration: Vector3 { + x: 0.0, + y: 0.0, + z: 0.0, + }, + linear_acceleration_covariance: [0.0; 9], + } + } + + fn transform() -> TransformStamped { + TransformStamped { + header: Header { + stamp: Time { sec: 0, nanosec: 0 }, + frame_id: String::new(), + }, + child_frame_id: String::new(), + transform: Transform { + translation: Vector3 { + x: 0.0, + y: 0.0, + z: 0.0, + }, + rotation: Quaternion { + x: 0.0, + y: 0.0, + z: 0.0, + w: 0.0, + }, + }, + } + } + fn stream(name: &str, codec: Codec, is_tf: bool) -> Arc { Arc::new(StreamConfig { port: name.to_string(), @@ -475,8 +533,19 @@ mod tests { "test.Raw".to_string() }, codec, + schema_name: if is_tf { + "tf2_msgs/msg/TFMessage" + } else { + "test_msgs/msg/Raw" + } + .to_string(), timestamp_field: None, json_schema: None, + schema_definition: if is_tf { + schema("tf2_msgs/msg/TFMessage") + } else { + "uint8[] data\n".to_string() + }, }) } @@ -485,6 +554,14 @@ mod tests { port: name.to_string(), name: name.to_string(), payload_type: payload_type.to_string(), + schema_name: payload_type.to_string(), + schema_definition: match payload_type { + "sensor_msgs/msg/Image" => schema("sensor_msgs/msg/Image"), + "sensor_msgs/msg/Imu" => schema("sensor_msgs/msg/Imu"), + "sensor_msgs/msg/CompressedImage" => schema("sensor_msgs/msg/CompressedImage"), + _ => "uint8[] data\n".to_string(), + } + .to_string(), codec, timestamp_field: None, json_schema: None, @@ -502,11 +579,13 @@ mod tests { } #[test] - fn jpeg_codec_preserves_the_lcm_envelope() { + fn raw_images_keep_their_declared_encoding_and_pixels() { let image = Image { header: Header { - seq: 9, - stamp: Time { sec: 12, nsec: 34 }, + stamp: Time { + sec: 12, + nanosec: 34, + }, frame_id: "camera".to_string(), }, height: 2, @@ -517,39 +596,44 @@ mod tests { data: vec![255, 0, 0, 0, 255, 0, 0, 0, 255, 255, 255, 255], }; - let stream = typed_stream("camera", IMAGE_PAYLOAD_TYPE, Codec::Jpeg); - let observations = process(&stream, &image.encode(), 100.0).unwrap(); - let decoded = Image::decode(&observations[0].data).unwrap(); + let stream = typed_stream("camera", IMAGE_PAYLOAD_TYPE, Codec::Cdr); + let observations = + process(&stream, &cdr::encode(&image).unwrap(), 100_000_000_000).unwrap(); + let decoded = cdr::decode::(&observations[0].data).unwrap(); - assert_eq!(observations[0].ts, 12.000_000_034); + assert_eq!(observations[0].ts, 12_000_000_034); assert_eq!(decoded.header, image.header); - assert_eq!(decoded.encoding, "jpeg"); - assert_eq!(decoded.step, 0); - assert_eq!(&decoded.data[..2], &[0xff, 0xd8]); + assert_eq!(decoded, image); } #[test] fn tf_batches_become_individually_timestamped_observations() { - let mut first = lcm_msgs::geometry_msgs::TransformStamped::default(); - first.header.stamp = Time { sec: 10, nsec: 5 }; + let mut first = transform(); + first.header.stamp = Time { + sec: 10, + nanosec: 5, + }; first.child_frame_id = "first".to_string(); - let mut second = lcm_msgs::geometry_msgs::TransformStamped::default(); - second.header.stamp = Time { sec: 20, nsec: 7 }; + let mut second = transform(); + second.header.stamp = Time { + sec: 20, + nanosec: 7, + }; second.child_frame_id = "second".to_string(); let message = TFMessage { transforms: vec![first, second], }; - let stream = stream("tf", Codec::Lcm, true); - let data = message.encode(); + let stream = stream("tf", Codec::Cdr, true); + let data = cdr::encode(&message).unwrap(); - let observations = process(&stream, &data, 100.0).unwrap(); + let observations = process(&stream, &data, 100_000_000_000).unwrap(); assert_eq!(observations.len(), 2); - assert_eq!(observations[0].ts, 10.000_000_005); - assert_eq!(observations[1].ts, 20.000_000_007); + assert_eq!(observations[0].ts, 10_000_000_005); + assert_eq!(observations[1].ts, 20_000_000_007); let decoded: Vec = observations .iter() - .map(|observation| TFMessage::decode(&observation.data).unwrap()) + .map(|observation| cdr::decode::(&observation.data).unwrap()) .collect(); assert_eq!(decoded[0].transforms.len(), 1); assert_eq!(decoded[0].transforms[0].child_frame_id, "first"); @@ -559,42 +643,52 @@ mod tests { #[test] fn stamped_sensor_messages_preserve_their_source_timestamp() { - let mut message = Imu::default(); - message.header.stamp = Time { sec: 42, nsec: 25 }; + let mut message = imu(); + message.header.stamp = Time { + sec: 42, + nanosec: 25, + }; let stream = Arc::new(StreamConfig { port: "imu".to_string(), name: "imu".to_string(), - payload_type: "dimos.msgs.sensor_msgs.Imu.Imu".to_string(), - codec: Codec::Lcm, + payload_type: "dimos_generated.sensor_msgs.msg.Imu".to_string(), + schema_name: "sensor_msgs/msg/Imu".to_string(), + schema_definition: schema("sensor_msgs/msg/Imu").to_string(), + codec: Codec::Cdr, timestamp_field: None, json_schema: None, }); - let observations = decoding::decode(&stream, &message.encode(), 100.0).unwrap(); - assert_eq!(observations[0].ts, 42.000_000_025); + let observations = + decoding::decode(&stream, &cdr::encode(&message).unwrap(), 100_000_000_000).unwrap(); + assert_eq!(observations[0].ts, 42_000_000_025); } #[test] - fn lcm_storage_encoding_consumes_the_decoded_image() { + fn cdr_storage_preserves_image_bytes() { let image = Image { header: Header { - stamp: Time { sec: 42, nsec: 25 }, - ..Header::default() + stamp: Time { + sec: 42, + nanosec: 25, + }, + frame_id: String::new(), }, height: 1, width: 1, encoding: "rgb8".to_string(), step: 3, data: vec![1, 2, 3], - ..Image::default() + is_bigendian: 0, }; - let stream = typed_stream("camera", IMAGE_PAYLOAD_TYPE, Codec::Lcm); + let stream = typed_stream("camera", IMAGE_PAYLOAD_TYPE, Codec::Cdr); - let observations = process(&stream, &image.encode(), 100.0).unwrap(); + let observations = + process(&stream, &cdr::encode(&image).unwrap(), 100_000_000_000).unwrap(); assert_eq!(observations.len(), 1); - assert_eq!(observations[0].ts, 42.000_000_025); - assert_eq!(Image::decode(&observations[0].data).unwrap(), image); + assert_eq!(observations[0].ts, 42_000_000_025); + assert_eq!(cdr::decode::(&observations[0].data).unwrap(), image); } #[test] @@ -605,7 +699,7 @@ mod tests { path: file.path().to_string_lossy().into_owned(), }, encoding_threads: 2, - streams: vec![(*stream("samples", Codec::Lcm, false)).clone()], + streams: vec![(*stream("samples", Codec::Cdr, false)).clone()], }; let (sender, receiver) = bounded(8); let (permit_sender, permit_receiver) = bounded(8); @@ -616,10 +710,10 @@ mod tests { sender .send(EncodedBatch { sequence, - stream: stream("samples", Codec::Lcm, false), - reception_ts: sequence as f64, + stream: stream("samples", Codec::Cdr, false), + reception_ts: sequence as i64 * 1_000_000_000, observations: Ok(vec![StoredObservation { - ts: sequence as f64, + ts: sequence as i64 * 1_000_000_000, data: vec![sequence as u8], }]), }) @@ -668,7 +762,7 @@ mod tests { fn tf_observations_use_empty_jsonb_tags_without_a_reception_index() { let file = NamedTempFile::new().unwrap(); let connection = Connection::open(file.path()).unwrap(); - let tf = stream("tf", Codec::Lcm, true); + let tf = stream("tf", Codec::Cdr, true); drop(connection); let mut recording_store = store::open( &store::RecordingStoreConfig::Sqlite { @@ -681,8 +775,8 @@ mod tests { recording_store .write_batch(&[Observation { stream: Arc::clone(&tf), - source_ts: 1.0, - reception_ts: 100.0, + source_ts: 1_000_000_000, + reception_ts: 100_000_000_000, data: vec![1, 2, 3], }]) .unwrap(); @@ -710,7 +804,7 @@ mod tests { #[test] fn mcap_store_writes_indexed_storage_encoded_messages() { let file = NamedTempFile::new().unwrap(); - let samples = stream("samples", Codec::Lz4Lcm, false); + let samples = stream("samples", Codec::Cdr, false); let mut recording_store = store::open( &store::RecordingStoreConfig::Mcap { path: file.path().to_string_lossy().into_owned(), @@ -719,12 +813,12 @@ mod tests { 1, ) .unwrap(); - let encoded = encoding::lz4_frame(&[1, 2, 3]).unwrap(); + let encoded = vec![0, 1, 0, 0, 3, 0, 0, 0, 1, 2, 3]; recording_store .write_batch(&[Observation { stream: Arc::clone(&samples), - source_ts: 12.5, - reception_ts: 13.0, + source_ts: 1_700_000_000_123_456_789, + reception_ts: 1_700_000_000_987_654_321, data: encoded.clone(), }]) .unwrap(); @@ -747,17 +841,21 @@ mod tests { assert_eq!(messages.len(), 1); let message = &messages[0]; assert_eq!(message.channel.topic, "samples"); - assert_eq!(message.channel.message_encoding, "lz4+lcm"); + assert_eq!(message.channel.message_encoding, "cdr"); + let schema = message.channel.schema.as_ref().unwrap(); + assert_eq!(schema.name, "test_msgs/msg/Raw"); + assert_eq!(schema.encoding, "ros2msg"); + assert_eq!(schema.data.as_ref(), b"uint8[] data\n"); assert_eq!(message.channel.metadata["dimos.payload_type"], "test.Raw"); - assert_eq!(message.log_time, 13_000_000_000); - assert_eq!(message.publish_time, 12_500_000_000); + assert_eq!(message.log_time, 1_700_000_000_987_654_321); + assert_eq!(message.publish_time, 1_700_000_000_123_456_789); assert_eq!(message.data.as_ref(), encoded); } #[test] fn engine_uses_the_configured_worker_count_and_flushes_on_shutdown() { let file = NamedTempFile::new().unwrap(); - let raw = stream("raw", Codec::Lcm, false); + let raw = stream("raw", Codec::Cdr, false); let config = RecorderConfig { store: store::RecordingStoreConfig::Sqlite { path: file.path().to_string_lossy().into_owned(), @@ -780,7 +878,7 @@ mod tests { #[test] fn encoding_failure_fails_the_recording_with_stream_and_sequence() { let file = NamedTempFile::new().unwrap(); - let image = typed_stream("camera", IMAGE_PAYLOAD_TYPE, Codec::Jpeg); + let image = typed_stream("camera", IMAGE_PAYLOAD_TYPE, Codec::Cdr); let config = RecorderConfig { store: store::RecordingStoreConfig::Sqlite { path: file.path().to_string_lossy().into_owned(), @@ -789,14 +887,14 @@ mod tests { streams: vec![(*image).clone()], }; let engine = RecorderEngine::start(config).unwrap(); - engine.handle().record(image, b"not an lcm image"); + engine.handle().record(image, b"not a CDR image"); let error = engine.shutdown().unwrap_err(); let message = format!("{error:#}"); assert!(message.contains("camera")); assert!(message.contains("sequence 0")); - assert!(message.contains("invalid LCM Image")); + assert!(message.contains("invalid CDR sensor_msgs/msg/Image")); } #[test] @@ -812,61 +910,73 @@ mod tests { assert!(validator::Validate::validate(&config).is_err()); } #[test] - fn json_storage_is_portable_and_keeps_both_timestamps() { - let text = r#"{"sent":12.5,"label":"拿起积木"}"#; - let message = lcm_msgs::std_msgs::String { data: text.into() }; - let schema = serde_json::json!({"type":"object","properties":{"sent":{"type":"number"}}}); - let stream = Arc::new(StreamConfig { - port: "events".into(), - name: "events".into(), - payload_type: "dimos.msgs.std_msgs.String.String".into(), - codec: Codec::Json, - timestamp_field: Some("sent".into()), - json_schema: Some(schema.clone()), - }); - let mut encoded = process(&stream, &message.encode(), 13.0).unwrap(); - let stored = encoded.pop().unwrap(); - assert_eq!(stored.data, text.as_bytes()); - let observations = [Observation { - stream: Arc::clone(&stream), - source_ts: stored.ts, - reception_ts: 13.0, - data: stored.data, - }]; - for sqlite in [true, false] { - let file = NamedTempFile::new().unwrap(); - let path = file.path().to_string_lossy().into_owned(); - let config = if sqlite { - store::RecordingStoreConfig::Sqlite { path } - } else { - store::RecordingStoreConfig::Mcap { path } - }; - let mut store = store::open(&config, &[(*stream).clone()], 1).unwrap(); - store.write_batch(&observations).unwrap(); - store.finish().unwrap(); - drop(store); - if sqlite { - let db = Connection::open(file.path()).unwrap(); - let row: (f64, f64, Vec) = db.query_row("SELECT ts,json_extract(tags,'$.reception_ts'),data FROM events JOIN events_blob USING(id)", [], |r| Ok((r.get(0)?,r.get(1)?,r.get(2)?))).unwrap(); - assert_eq!(row, (12.5, 13.0, text.as_bytes().to_vec())); - } else { - let bytes = std::fs::read(file.path()).unwrap(); - let record = mcap::MessageStream::new(&bytes) - .unwrap() - .next() - .unwrap() - .unwrap(); - assert_eq!(record.data.as_ref(), text.as_bytes()); - assert_eq!(record.log_time, 13_000_000_000); - assert_eq!(record.publish_time, 12_500_000_000); - assert_eq!(record.channel.message_encoding, "json"); - let actual_schema = record.channel.schema.as_ref().unwrap(); - assert_eq!(actual_schema.encoding, "jsonschema"); - assert_eq!( - serde_json::from_slice::(&actual_schema.data).unwrap(), - schema - ); + fn stamped_cdr_preserves_zero_negative_and_exact_epoch_nanoseconds_in_both_endians() { + let stream = typed_stream("imu", "sensor_msgs/msg/Imu", Codec::Cdr); + for (sec, nanosec, expected) in [ + (0, 0, 0), + (-1, 500_000_000, -500_000_000), + (1_700_000_000, 123_456_789, 1_700_000_000_123_456_789), + ] { + let mut message = imu(); + message.header.stamp = Time { sec, nanosec }; + for little_endian in [false, true] { + let bytes = cdr::encode_endian(&message, little_endian).unwrap(); + let observations = process(&stream, &bytes, 99).unwrap(); + assert_eq!(observations[0].ts, expected); + assert_eq!(observations[0].data, bytes); } } } + + #[test] + fn compressed_image_keeps_its_standard_schema_and_payload() { + let message = CompressedImage { + header: Header { + stamp: Time { sec: 2, nanosec: 3 }, + frame_id: String::new(), + }, + format: "rgb8; jpeg compressed bgr8".into(), + data: vec![0xff, 0xd8, 0xff, 0xd9], + }; + let stream = typed_stream("camera", "sensor_msgs/msg/CompressedImage", Codec::Cdr); + let observations = process(&stream, &cdr::encode(&message).unwrap(), 99).unwrap(); + assert_eq!(observations[0].ts, 2_000_000_003); + assert_eq!( + cdr::decode::(&observations[0].data).unwrap(), + message + ); + } + + #[test] + fn unknown_custom_type_is_opaque_and_uses_reception_nanoseconds() { + let stream = stream("custom", Codec::Cdr, false); + let bytes = vec![0, 1, 0, 0, 3, 0, 0, 0, 1, 2, 3]; + let observations = process(&stream, &bytes, 1_700_000_000_987_654_321).unwrap(); + assert_eq!(observations[0].ts, 1_700_000_000_987_654_321); + assert_eq!(observations[0].data, bytes); + } + + #[test] + fn invalid_ros_nanoseconds_are_rejected() { + let mut message = imu(); + message.header.stamp.nanosec = 1_000_000_000; + let stream = typed_stream("imu", "sensor_msgs/msg/Imu", Codec::Cdr); + assert!(process(&stream, &cdr::encode(&message).unwrap(), 99).is_err()); + } + + #[test] + fn mcap_rejects_payload_wrappers_before_touching_the_artifact() { + let file = NamedTempFile::new().unwrap(); + std::fs::write(file.path(), b"existing").unwrap(); + let stream = stream("custom", Codec::Lz4Cdr, false); + let result = store::open( + &store::RecordingStoreConfig::Mcap { + path: file.path().to_string_lossy().into_owned(), + }, + &[(*stream).clone()], + 1, + ); + assert!(result.is_err()); + assert_eq!(std::fs::read(file.path()).unwrap(), b"existing"); + } } diff --git a/dimos/experimental/memory/rust/src/store/mcap.rs b/dimos/experimental/memory/rust/src/store/mcap.rs index 23213affe24..d3d026b4344 100644 --- a/dimos/experimental/memory/rust/src/store/mcap.rs +++ b/dimos/experimental/memory/rust/src/store/mcap.rs @@ -15,14 +15,13 @@ use std::collections::{BTreeMap, HashMap}; use std::fs::File; use std::io::BufWriter; -use std::time::Duration; use anyhow::{Context, Result}; use mcap::records::MessageHeader; use mcap::{Compression, WriteOptions, Writer}; use super::{Observation, RecordingStore}; -use crate::StreamConfig; +use crate::{Codec, StreamConfig}; pub struct McapRecordingStore { writer: Writer>, @@ -32,9 +31,19 @@ pub struct McapRecordingStore { impl McapRecordingStore { pub fn open(path: &str, streams: &[StreamConfig], compression_threads: usize) -> Result { + for stream in streams { + anyhow::ensure!( + stream.codec == Codec::Cdr, + "MCAP requires cdr channels; use chunk compression instead of a payload wrapper" + ); + anyhow::ensure!( + !stream.schema_name.is_empty() && !stream.schema_definition.is_empty(), + "MCAP requires a complete message schema" + ); + } let file = File::create(path).with_context(|| format!("failed to create {path}"))?; let options = WriteOptions::new() - .profile("dimos") + .profile("ros2") .library("dimos-memory-recorder") .compression(Some(Compression::Zstd)) .compression_threads(compression_threads.try_into().unwrap_or(u32::MAX)); @@ -53,13 +62,12 @@ impl McapRecordingStore { "publish_time".to_string(), ), ]); - let schema_id = if let Some(schema) = &stream.json_schema { - writer.add_schema(&stream.name, "jsonschema", &serde_json::to_vec(schema)?)? - } else { - 0 - }; - let channel = - writer.add_channel(schema_id, &stream.name, stream.codec.id(), &metadata)?; + let schema = writer.add_schema( + &stream.schema_name, + "ros2msg", + stream.schema_definition.as_bytes(), + )?; + let channel = writer.add_channel(schema, &stream.name, "cdr", &metadata)?; channels.insert(stream.name.clone(), channel); } Ok(Self { @@ -73,18 +81,6 @@ impl McapRecordingStore { impl RecordingStore for McapRecordingStore { fn write_batch(&mut self, observations: &[Observation]) -> Result<()> { for observation in observations { - let log_time = timestamp_ns(observation.reception_ts).with_context(|| { - format!( - "invalid MCAP reception time for stream {:?}", - observation.stream.name - ) - })?; - let publish_time = timestamp_ns(observation.source_ts).with_context(|| { - format!( - "invalid MCAP source time for stream {:?}", - observation.stream.name - ) - })?; let channel_id = self.channels[&observation.stream.name]; let sequence = self .sequences @@ -94,8 +90,10 @@ impl RecordingStore for McapRecordingStore { &MessageHeader { channel_id, sequence: *sequence, - log_time, - publish_time, + log_time: u64::try_from(observation.reception_ts) + .context("MCAP reception time precedes Unix epoch")?, + publish_time: u64::try_from(observation.source_ts) + .context("MCAP source time precedes Unix epoch")?, }, &observation.data, )?; @@ -109,117 +107,3 @@ impl RecordingStore for McapRecordingStore { Ok(()) } } - -fn timestamp_ns(timestamp: f64) -> Result { - Duration::try_from_secs_f64(timestamp) - .context("MCAP time must be finite and nonnegative")? - .as_nanos() - .try_into() - .context("MCAP time exceeds the unsigned 64-bit nanosecond range") -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::Codec; - use std::sync::Arc; - use tempfile::NamedTempFile; - - #[test] - fn timestamp_conversion_rejects_unrepresentable_times_without_panicking() { - for timestamp in [ - -12.5, - f64::NAN, - f64::INFINITY, - f64::NEG_INFINITY, - f64::MAX, - u64::MAX as f64 / 1_000_000_000.0, - ] { - assert!(timestamp_ns(timestamp).is_err(), "accepted {timestamp}"); - } - assert_eq!(timestamp_ns(0.0).unwrap(), 0); - assert_eq!(timestamp_ns(12.5).unwrap(), 12_500_000_000); - } - - #[test] - fn mcap_rejects_negative_source_time_instead_of_writing_zero() { - let file = NamedTempFile::new().unwrap(); - let stream = Arc::new(StreamConfig { - port: "events".into(), - name: "events".into(), - payload_type: "dimos.msgs.std_msgs.String.String".into(), - codec: Codec::Json, - timestamp_field: Some("ts".into()), - json_schema: None, - }); - let mut store = - McapRecordingStore::open(file.path().to_str().unwrap(), &[(*stream).clone()], 1) - .unwrap(); - let observation = Observation { - stream, - source_ts: -12.5, - reception_ts: 100.0, - data: br#"{"ts":-12.5}"#.to_vec(), - }; - assert!(store - .write_batch(std::slice::from_ref(&observation)) - .is_err()); - store.finish().unwrap(); - let bytes = std::fs::read(file.path()).unwrap(); - assert_eq!(mcap::MessageStream::new(&bytes).unwrap().count(), 0); - - let sqlite_file = NamedTempFile::new().unwrap(); - let mut sqlite = super::super::open( - &super::super::RecordingStoreConfig::Sqlite { - path: sqlite_file.path().to_str().unwrap().into(), - }, - std::slice::from_ref(observation.stream.as_ref()), - 1, - ) - .unwrap(); - sqlite.write_batch(&[observation]).unwrap(); - sqlite.finish().unwrap(); - let connection = rusqlite::Connection::open(sqlite_file.path()).unwrap(); - let source_ts: f64 = connection - .query_row("SELECT ts FROM events", [], |row| row.get(0)) - .unwrap(); - assert_eq!(source_ts, -12.5); - } - - #[test] - fn mcap_preserves_zero_source_time_and_rejects_invalid_reception_time() { - let file = NamedTempFile::new().unwrap(); - let stream = Arc::new(StreamConfig { - port: "events".into(), - name: "events".into(), - payload_type: "dimos.msgs.std_msgs.String.String".into(), - codec: Codec::Json, - timestamp_field: Some("ts".into()), - json_schema: None, - }); - let mut store = - McapRecordingStore::open(file.path().to_str().unwrap(), &[(*stream).clone()], 1) - .unwrap(); - let mut observation = Observation { - stream, - source_ts: 0.0, - reception_ts: -1.0, - data: br#"{"ts":0}"#.to_vec(), - }; - assert!(store - .write_batch(std::slice::from_ref(&observation)) - .is_err()); - observation.reception_ts = 12.5; - store.write_batch(&[observation]).unwrap(); - store.finish().unwrap(); - let bytes = std::fs::read(file.path()).unwrap(); - let messages = mcap::MessageStream::new(&bytes) - .unwrap() - .collect::>>() - .unwrap(); - assert_eq!(messages.len(), 1); - assert_eq!(messages[0].publish_time, 0); - assert_eq!(messages[0].log_time, 12_500_000_000); - assert_eq!(messages[0].sequence, 0); - } -} diff --git a/dimos/experimental/memory/rust_cli_recorder.py b/dimos/experimental/memory/rust_cli_recorder.py index 3aad51177d6..e110b847762 100644 --- a/dimos/experimental/memory/rust_cli_recorder.py +++ b/dimos/experimental/memory/rust_cli_recorder.py @@ -102,9 +102,7 @@ def make_plan(transports: dict[tuple[str, type], Any]) -> RustRecordingPlan: for index, ((name, payload_type), transport) in enumerate(transports.items()): if name not in selected: continue - if not all( - hasattr(payload_type, member) for member in ("encode", "decode", "msg_name", "schema") - ): + if not all(hasattr(payload_type, member) for member in ("__msgtype__",)): logger.info( "--record: stream is not a generated CDR message type; skipped", stream=name, diff --git a/dimos/experimental/memory/rust_recorder.py b/dimos/experimental/memory/rust_recorder.py index a70d8e22d35..a0bf80ea748 100644 --- a/dimos/experimental/memory/rust_recorder.py +++ b/dimos/experimental/memory/rust_recorder.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Experimental native recorder for high-throughput LCM-backed memory streams.""" +"""Experimental native recorder for high-throughput CDR memory streams.""" from __future__ import annotations @@ -20,6 +20,9 @@ from pathlib import Path from typing import Annotated, Any, Literal, TypeAlias +from dimos_cdrgen.registry import schema as cdr_schema +from dimos_generated.std_msgs.msg import String +from dimos_generated.tf2_msgs.msg import TFMessage from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from dimos.constants import DIMOS_PROJECT_ROOT @@ -27,17 +30,14 @@ from dimos.core.module import Module from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.stream import In +from dimos.memory.recording_policy import OnExisting from dimos.memory.store.sqlite import SqliteStore -from dimos.memory.type.recording import OnExisting -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.std_msgs.String import String -from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.utils.data import backup_file from dimos.utils.logging_config import setup_logger logger = setup_logger() -_SUPPORTED_NATIVE_CODECS = {"lcm", "jpeg", "lz4+lcm", "json"} +_SUPPORTED_NATIVE_CODECS = {"cdr", "lz4+cdr", "json"} class RustStreamSpec(BaseModel): @@ -47,6 +47,8 @@ class RustStreamSpec(BaseModel): name: str payload_type: str codec: str + schema_name: str = Field(min_length=1) + schema_definition: str = Field(min_length=1) timestamp_field: str | None = None json_schema: dict[str, Any] | None = None @@ -58,6 +60,28 @@ def _json_options(self) -> RustStreamSpec: raise ValueError("JSON options require the json codec") return self + @classmethod + def from_type( + cls, + *, + port: str, + name: str, + payload_type: type[Any], + codec: str, + timestamp_field: str | None = None, + json_schema: dict[str, Any] | None = None, + ) -> RustStreamSpec: + return cls( + port=port, + name=name, + codec=codec, + timestamp_field=timestamp_field, + json_schema=json_schema, + payload_type=f"{payload_type.__module__}.{payload_type.__qualname__}", + schema_name=payload_type.__msgtype__, + schema_definition=cdr_schema(payload_type.__msgtype__), + ) + class RustStoreConfig(BaseModel): """Artifact path shared by native recording-store configurations.""" @@ -178,11 +202,11 @@ class RustRecorder(NativeModule): class CameraRecorder(RustRecorder): color_image: In[Image] - Both stores use the Python Mem2 codec configuration: images default to - JPEG, and ``stream_codecs`` may select ``lcm`` or ``lz4+lcm``. MCAP uses - indexed Zstd chunks around those storage-encoded observations. The native + Both stores accept generated CDR messages with complete schema metadata. + SQLite additionally accepts ``lz4+cdr`` blob compression. MCAP uses + indexed Zstd chunks around unwrapped CDR observations. The native path preserves source timestamps for common stamped message types. Other - LCM messages use their reception timestamp. Spatial pose attachment is not + timestamp layouts use their reception timestamp. Spatial pose attachment is not supported yet. """ @@ -223,11 +247,11 @@ def _stream_specs(self) -> list[RustStreamSpec]: if port is self.tf: if self.config.record_tf: specs.append( - RustStreamSpec( + RustStreamSpec.from_type( port=port_name, name="tf", - payload_type=f"{TFMessage.__module__}.{TFMessage.__qualname__}", - codec="lcm", + payload_type=TFMessage, + codec="cdr", ) ) continue @@ -236,10 +260,10 @@ def _stream_specs(self) -> list[RustStreamSpec]: codec = self.config.stream_codecs.get(stream_name, self._default_codec(port.type)) self._validate_codec(stream_name, port.type, codec) specs.append( - RustStreamSpec( + RustStreamSpec.from_type( port=port_name, name=stream_name, - payload_type=f"{port.type.__module__}.{port.type.__qualname__}", + payload_type=port.type, codec=codec, timestamp_field=self.config.stream_timestamp_fields.get(stream_name), json_schema=self.config.stream_json_schemas.get(stream_name), @@ -261,32 +285,27 @@ def _collect_topics(self) -> dict[str, str]: @staticmethod def _default_codec(payload_type: type[Any]) -> str: - if issubclass(payload_type, Image): - return "jpeg" - if hasattr(payload_type, "lcm_encode") and hasattr(payload_type, "lcm_decode"): - return "lcm" + if all(hasattr(payload_type, member) for member in ("__msgtype__",)): + return "cdr" raise TypeError( - f"RustRecorder only supports LCM-backed messages, got {payload_type.__qualname__}" + f"RustRecorder requires a generated CDR message, got {payload_type.__qualname__}" ) @staticmethod def _validate_codec(stream_name: str, payload_type: type[Any], codec: str) -> None: if codec not in _SUPPORTED_NATIVE_CODECS: raise ValueError( - f"Unsupported native codec {codec!r} for stream {stream_name!r}; " - f"choose one of {sorted(_SUPPORTED_NATIVE_CODECS)}" + f"Unsupported native codec {codec!r} for stream {stream_name!r}; choose one of {sorted(_SUPPORTED_NATIVE_CODECS)}" ) if codec == "json" and payload_type is not String: raise TypeError("JSON codec requires std_msgs.String") - if codec == "jpeg" and not issubclass(payload_type, Image): - raise TypeError(f"JPEG codec requires Image, got {payload_type.__qualname__}") - if not hasattr(payload_type, "lcm_encode") or not hasattr(payload_type, "lcm_decode"): - raise TypeError( - f"Native codec {codec!r} requires an LCM-backed type, " - f"got {payload_type.__qualname__}" - ) + RustRecorder._default_codec(payload_type) def _prepare_store(self, specs: list[RustStreamSpec]) -> None: + if self.config.store.kind == "mcap" and any(spec.codec != "cdr" for spec in specs): + raise ValueError( + "MCAP requires cdr channels; use chunk compression instead of payload wrappers" + ) path = Path(self.config.store.path) if self.config.store.kind == "mcap" and self.config.on_existing is OnExisting.APPEND: raise ValueError("MCAP append is unsupported; choose overwrite, backup, or error") diff --git a/dimos/experimental/memory/test_rust_recorder.py b/dimos/experimental/memory/test_rust_recorder.py index 21524681992..47b8073cd1b 100644 --- a/dimos/experimental/memory/test_rust_recorder.py +++ b/dimos/experimental/memory/test_rust_recorder.py @@ -17,11 +17,17 @@ from pathlib import Path from typing import Any, TypeVar +from dimos_cdrgen.registry import schema as cdr_schema +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header, String import pytest from dimos.core.global_config import GlobalConfig from dimos.core.module import Module from dimos.core.stream import In +from dimos.core.transport_factory import rpc_backend from dimos.experimental.memory.rust_recorder import ( RustMcapStoreConfig, RustRecorder, @@ -29,10 +35,8 @@ RustSqliteStoreConfig, RustStreamSpec, ) +from dimos.memory.recording_policy import OnExisting from dimos.memory.store.sqlite import SqliteStore -from dimos.memory.type.recording import OnExisting -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.Image import Image class SampleRustRecorder(RustRecorder): @@ -60,7 +64,9 @@ def stop(self) -> None: @pytest.fixture -def make_recorder() -> Any: +def make_recorder(mocker) -> Any: + for method in ("start", "serve_module_rpc", "stop"): + mocker.patch.object(rpc_backend(), method) recorders: list[RustRecorder] = [] def make(recorder_type: type[TRecorder], **kwargs: Any) -> TRecorder: @@ -87,7 +93,7 @@ def test_specs_use_native_defaults_remapping_and_configured_workers( record_tf=False, encoding_threads=7, stream_remapping={"odometry": "pose"}, - stream_codecs={"color_image": "lz4+lcm"}, + stream_codecs={"color_image": "lz4+cdr"}, ) connect(recorder, color_image="/camera", odometry="/odom") @@ -104,14 +110,18 @@ def test_specs_use_native_defaults_remapping_and_configured_workers( { "port": "color_image", "name": "color_image", - "payload_type": "dimos.msgs.sensor_msgs.Image.Image", - "codec": "lz4+lcm", + "payload_type": f"{Image.__module__}.{Image.__qualname__}", + "schema_name": Image.__msgtype__, + "schema_definition": cdr_schema(Image.__msgtype__), + "codec": "lz4+cdr", }, { "port": "odometry", "name": "pose", - "payload_type": "dimos.msgs.geometry_msgs.PoseStamped.PoseStamped", - "codec": "lcm", + "payload_type": f"{PoseStamped.__module__}.{PoseStamped.__qualname__}", + "schema_name": PoseStamped.__msgtype__, + "schema_definition": cdr_schema(PoseStamped.__msgtype__), + "codec": "cdr", }, ] assert set(config) == {"encoding_threads", "store", "streams"} @@ -124,7 +134,9 @@ def test_json_options_cross_the_native_boundary_only_when_configured() -> None: RustStreamSpec( port="events", name="events", - payload_type="dimos.msgs.std_msgs.String.String", + payload_type=f"{String.__module__}.{String.__qualname__}", + schema_name=String.__msgtype__, + schema_definition=cdr_schema(String.__msgtype__), codec="json", timestamp_field="sent", json_schema=schema, @@ -136,7 +148,9 @@ def test_json_options_cross_the_native_boundary_only_when_configured() -> None: { "port": "events", "name": "events", - "payload_type": "dimos.msgs.std_msgs.String.String", + "payload_type": f"{String.__module__}.{String.__qualname__}", + "schema_name": String.__msgtype__, + "schema_definition": cdr_schema(String.__msgtype__), "codec": "json", "timestamp_field": "sent", "json_schema": schema, @@ -144,16 +158,25 @@ def test_json_options_cross_the_native_boundary_only_when_configured() -> None: ] -@pytest.mark.parametrize("codec", ["lcm", "jpeg", "lz4+lcm", "json"]) +@pytest.mark.parametrize("codec", ["cdr", "lz4+cdr", "json"]) def test_unconfigured_json_options_are_omitted_from_native_streams(codec: str) -> None: config = RustRecorderConfig( streams=[ - RustStreamSpec(port="samples", name="samples", payload_type="test.Raw", codec=codec) + RustStreamSpec.from_type( + port="samples", name="samples", payload_type=String, codec=codec + ) ] ) assert config.to_config_dict()["streams"] == [ - {"port": "samples", "name": "samples", "payload_type": "test.Raw", "codec": codec} + { + "port": "samples", + "name": "samples", + "payload_type": f"{String.__module__}.{String.__qualname__}", + "schema_name": String.__msgtype__, + "schema_definition": cdr_schema(String.__msgtype__), + "codec": codec, + } ] @@ -169,7 +192,7 @@ def test_store_preparation_creates_a_python_readable_registry( connect(recorder, color_image="/camera", odometry="/odom") specs = recorder._stream_specs() - assert [spec.codec for spec in specs] == ["jpeg", "lcm"] + assert [spec.codec for spec in specs] == ["cdr", "cdr"] recorder._prepare_store(specs) @@ -182,8 +205,26 @@ def test_store_preparation_creates_a_python_readable_registry( def test_append_replaces_only_the_recorded_streams(tmp_path: Path, make_recorder: Any) -> None: path = tmp_path / "recording.db" with SqliteStore(path=str(path)) as store: - store.stream("keep", PoseStamped).append(PoseStamped(), ts=1.0) - store.stream("odometry", PoseStamped).append(PoseStamped(), ts=2.0) + store.stream("keep", PoseStamped).append( + PoseStamped( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + ts=1.0, + ) + store.stream("odometry", PoseStamped).append( + PoseStamped( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + ts=2.0, + ) recorder = make_recorder( SampleRustRecorder, @@ -209,7 +250,7 @@ def test_unsupported_python_payload_fails_before_native_process_starts( ) connect(recorder, values="/values") - with pytest.raises(TypeError, match="only supports LCM-backed messages"): + with pytest.raises(TypeError, match="requires a generated CDR message"): recorder._stream_specs() @@ -239,7 +280,7 @@ def test_invalid_codec_fails_during_preflight(tmp_path: Path, make_recorder: Any recorder._stream_specs() -def test_jpeg_requires_an_image_stream(tmp_path: Path, make_recorder: Any) -> None: +def test_obsolete_jpeg_codec_is_rejected(tmp_path: Path, make_recorder: Any) -> None: recorder = make_recorder( SampleRustRecorder, store=RustSqliteStoreConfig(path=str(tmp_path / "recording.db")), @@ -248,7 +289,7 @@ def test_jpeg_requires_an_image_stream(tmp_path: Path, make_recorder: Any) -> No ) connect(recorder, odometry="/odom") - with pytest.raises(TypeError, match="JPEG codec requires Image"): + with pytest.raises(ValueError, match="Unsupported native codec"): recorder._stream_specs() @@ -278,7 +319,7 @@ def test_mcap_store_uses_python_codec_defaults_and_does_not_precreate_the_artifa specs = recorder._stream_specs() recorder._prepare_store(specs) - assert [spec.codec for spec in specs] == ["jpeg", "lcm"] + assert [spec.codec for spec in specs] == ["cdr", "cdr"] assert not path.exists() recorder.config.streams = specs assert recorder.config.to_config_dict()["store"] == { @@ -293,11 +334,13 @@ def test_mcap_accepts_storage_codecs_and_rejects_append(tmp_path: Path, make_rec SampleRustRecorder, store=RustMcapStoreConfig(path=str(path)), record_tf=False, - stream_codecs={"color_image": "jpeg"}, + stream_codecs={"color_image": "lz4+cdr"}, ) connect(recorder, color_image="/camera") - assert [spec.codec for spec in recorder._stream_specs()] == ["jpeg"] + with pytest.raises(ValueError, match="MCAP requires cdr"): + recorder._prepare_store(recorder._stream_specs()) + assert not path.exists() append_recorder = make_recorder( SampleRustRecorder, @@ -341,6 +384,21 @@ def test_duplicate_remapped_stream_names_fail_before_launch( recorder._stream_specs() +@pytest.mark.parametrize("record_tf", [False, True]) +def test_launch_topics_match_selected_streams_with_connected_tf(tmp_path, make_recorder, record_tf): + recorder = make_recorder( + SampleRustRecorder, + store=RustMcapStoreConfig(path=str(tmp_path / "recording.mcap")), + record_tf=record_tf, + ) + connect(recorder, odometry="/odom", tf="/tf") + recorder.config.streams = recorder._stream_specs() + topics = recorder._collect_topics() + expected = {"odometry": "/odom", "tf": "/tf"} if record_tf else {"odometry": "/odom"} + assert topics == expected + assert set(topics) == {spec.port for spec in recorder.config.streams} + + def test_replay_never_prepares_native_sources_or_process(tmp_path, make_recorder, mocker): fetch = mocker.patch( "dimos.core.native_module.get_project_root", side_effect=AssertionError("fetch") diff --git a/dimos/experimental/memory/test_rust_recorder_e2e.py b/dimos/experimental/memory/test_rust_recorder_e2e.py index 7d3410383cf..bede9bb0b78 100644 --- a/dimos/experimental/memory/test_rust_recorder_e2e.py +++ b/dimos/experimental/memory/test_rust_recorder_e2e.py @@ -25,6 +25,13 @@ from typing import cast import uuid +from dimos_cdrgen.registry import encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import Image, Imu +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage +from mcap.reader import make_reader import numpy as np import pytest @@ -40,16 +47,11 @@ RustRecordingStoreConfig, RustSqliteStoreConfig, ) -from dimos.memory.codecs.lcm import LcmCodec -from dimos.memory.codecs.lz4 import Lz4Codec from dimos.memory.store.mcap import McapStore from dimos.memory.store.sqlite import SqliteStore from dimos.memory.type.observation import Observation -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.sensor_msgs.Imu import Imu -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.image import image_from_array, image_view +from dimos.msgs.time import time_from_seconds from dimos.protocol.pubsub.impl.zenohpubsub import Topic as ZenohTopic from dimos.protocol.service.zenohservice import ZenohConfig, ZenohSessionPool @@ -58,7 +60,7 @@ _RUST_PACKAGE = DIMOS_PROJECT_ROOT / "dimos" / "experimental" / "memory" / "rust" _EXECUTABLE = Path( os.environ.get( - "DIMOS_MEMORY_RECORDER_TEST_EXECUTABLE", + "DIMOS_RECORDER_TEST_EXECUTABLE", str(_RUST_PACKAGE / "result" / "bin" / "dimos-memory-recorder"), ) ) @@ -162,7 +164,7 @@ def test_rust_artifact_is_readable_by_python_memory2( store=store, record_tf=False, encoding_threads=2, - stream_codecs={"imu": "lz4+lcm"}, + stream_codecs={"imu": "lz4+cdr" if store_kind == "sqlite" else "cdr"}, session=ZenohConfig( mode="peer", connect=[], @@ -221,15 +223,18 @@ def test_rust_artifact_is_readable_by_python_memory2( _wait_for_log(process, "memory recorder ready") expected = Imu( - ts=12.5, - frame_id="imu_link", - angular_velocity=Vector3(1.0, 2.0, 3.0), + header=Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id="imu_link"), + angular_velocity=Vector3(x=1.0, y=2.0, z=3.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + orientation_covariance=np.zeros(9, dtype=np.float64), + angular_velocity_covariance=np.zeros(9, dtype=np.float64), + linear_acceleration=Vector3(x=0.0, y=0.0, z=0.0), + linear_acceleration_covariance=np.zeros(9, dtype=np.float64), ) - expected_image = Image( - data=np.full((16, 16, 3), [20, 80, 140], dtype=np.uint8), - format=ImageFormat.RGB, - frame_id="camera", - ts=12.75, + expected_image = image_from_array( + np.full((16, 16, 3), [20, 80, 140], dtype=np.uint8), + encoding="rgb8", + header=Header(frame_id="camera", stamp=time_from_seconds(12.75)), ) publisher.broadcast(None, expected) image_publisher.broadcast(None, expected_image) @@ -249,24 +254,25 @@ def test_rust_artifact_is_readable_by_python_memory2( if store_kind == "sqlite": memory = SqliteStore(path=str(artifact)) else: - memory = McapStore( - path=str(artifact), - codecs={"imu": Lz4Codec(LcmCodec(Imu))}, - ) + with artifact.open("rb") as file: + recorded_imu = [ + row + for _schema, channel, row in make_reader(file).iter_messages() + if channel.topic == "imu" + ] + assert len(recorded_imu) == 1 + assert recorded_imu[0].publish_time == 1700000000123456789 + memory = McapStore(path=str(artifact)) with memory: observation = cast("Observation[Imu]", memory.stream("imu").first()) - assert observation.ts == 12.5 - assert observation.data.lcm_encode() == expected.lcm_encode() + assert observation.ts == 1700000000.123456789 + assert cdr_encode(observation.data) == cdr_encode(expected) image_observation = cast("Observation[Image]", memory.stream("color_image").first()) decoded_image = image_observation.data assert image_observation.ts == 12.75 - assert decoded_image.frame_id == "camera" - assert decoded_image.format is ImageFormat.RGB - assert decoded_image.data.shape == expected_image.data.shape - assert ( - np.mean(np.abs(decoded_image.data.astype(float) - expected_image.data.astype(float))) - < 5 - ) + assert decoded_image.header == expected_image.header + assert decoded_image.encoding == "rgb8" + np.testing.assert_array_equal(image_view(decoded_image), image_view(expected_image)) @pytest.mark.parametrize( @@ -296,7 +302,15 @@ def test_cli_recording_uses_existing_binary_for_both_formats( channel = f"/rust-recorder-{uuid.uuid4().hex[:8]}" publisher: LCMTransport[Imu] = LCMTransport(channel, Imu, url=lcm_url) session = RustRecordingSession(rust_cli_recorder.make_plan({("imu", Imu): publisher})) - expected = Imu(ts=22.5, frame_id="imu_link", angular_velocity=Vector3(1.0, 2.0, 3.0)) + expected = Imu( + header=Header(stamp=time_from_seconds(22.5), frame_id="imu_link"), + angular_velocity=Vector3(x=1, y=2, z=3), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + orientation_covariance=np.zeros(9, dtype=np.float64), + angular_velocity_covariance=np.zeros(9, dtype=np.float64), + linear_acceleration=Vector3(x=0.0, y=0.0, z=0.0), + linear_acceleration_covariance=np.zeros(9, dtype=np.float64), + ) try: publisher.start() session.start() @@ -316,11 +330,11 @@ def test_cli_recording_uses_existing_binary_for_both_formats( if store_kind == "sqlite": memory = SqliteStore(path=str(artifact)) else: - memory = McapStore(path=str(artifact), codecs={"imu": LcmCodec(Imu)}) + memory = McapStore(path=str(artifact)) with memory: observation = cast("Observation[Imu]", memory.stream("imu").first()) assert observation.ts == 22.5 - assert observation.data.lcm_encode() == expected.lcm_encode() + assert cdr_encode(observation.data) == cdr_encode(expected) def test_tf_records_over_zenoh_and_replays_through_python( @@ -378,18 +392,24 @@ def test_tf_records_over_zenoh_and_replays_through_python( stderr=subprocess.PIPE, ) expected = TFMessage( - Transform( - translation=Vector3(1.0, 2.0, 3.0), - frame_id="world", - child_frame_id="base_link", - ts=10.25, - ), - Transform( - translation=Vector3(4.0, 5.0, 6.0), - frame_id="base_link", - child_frame_id="camera", - ts=11.5, - ), + transforms=[ + TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(10.25)), + child_frame_id="base_link", + transform=Transform( + translation=Vector3(x=1, y=2, z=3), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + TransformStamped( + header=Header(frame_id="base_link", stamp=time_from_seconds(11.5)), + child_frame_id="camera", + transform=Transform( + translation=Vector3(x=4, y=5, z=6), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + ] ) try: assert process.stdin is not None diff --git a/dimos/experimental/robot/bosdyn/spot/effectors/high_level.py b/dimos/experimental/robot/bosdyn/spot/effectors/high_level.py index a465b40124c..8ddb8e23678 100644 --- a/dimos/experimental/robot/bosdyn/spot/effectors/high_level.py +++ b/dimos/experimental/robot/bosdyn/spot/effectors/high_level.py @@ -57,6 +57,7 @@ from dimos_generated.sensor_msgs.msg import CameraInfo, Image from dimos_generated.std_msgs.msg import Header from dimos_generated.tf2_msgs.msg import TFMessage +import numpy as np from dimos.agents.annotation import skill from dimos.core.core import rpc @@ -491,8 +492,8 @@ def _publish_odom(self, vision_tform_body: Any, velocity: Any, ts: float) -> Non odometry = Odometry( header=Header(stamp=time_from_seconds(ts), frame_id=self.config.odom_frame_id), child_frame_id=self.config.base_frame_id, - pose=PoseWithCovariance(pose=pose), - twist=TwistWithCovariance(twist=twist), + pose=PoseWithCovariance(pose=pose, covariance=np.zeros(36, dtype=np.float64)), + twist=TwistWithCovariance(twist=twist, covariance=np.zeros(36, dtype=np.float64)), ) self.odometry.publish(odometry) self.tf.publish(TFMessage(transforms=[transform_from_odometry(odometry)])) @@ -561,7 +562,7 @@ async def move_velocity( yaw: Rotational velocity (rad/s). duration: Seconds to move. 0 uses one `cmd_vel_timeout` window. """ - twist = Twist(linear=Vector3(x=x, y=y), angular=Vector3(z=yaw)) + twist = Twist(linear=Vector3(x=x, y=y, z=0.0), angular=Vector3(z=yaw, x=0.0, y=0.0)) if await self.move(twist, duration=duration): return f"Moving with velocity=({x}, {y}, {yaw}) for {duration} seconds" return f"Failed to move with velocity=({x}, {y}, {yaw})" diff --git a/dimos/experimental/robot/bosdyn/spot/effectors/test_high_level.py b/dimos/experimental/robot/bosdyn/spot/effectors/test_high_level.py index fe8905fbc63..e1b19ad184a 100644 --- a/dimos/experimental/robot/bosdyn/spot/effectors/test_high_level.py +++ b/dimos/experimental/robot/bosdyn/spot/effectors/test_high_level.py @@ -16,6 +16,7 @@ from types import SimpleNamespace +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.nav_msgs.msg import Odometry from dimos_generated.tf2_msgs.msg import TFMessage import pytest @@ -44,8 +45,8 @@ def test_odometry_and_tf_preserve_sdk_pose_velocity_and_source_time(module, mock odometry_publish.assert_called_once() tf_publish.assert_called_once() - odometry = Odometry.decode(odometry_publish.call_args.args[0].encode()) - tf = TFMessage.decode(tf_publish.call_args.args[0].encode()) + odometry = cdr_decode(cdr_encode(odometry_publish.call_args.args[0]), Odometry) + tf = cdr_decode(cdr_encode(tf_publish.call_args.args[0]), TFMessage) assert odometry.header.frame_id == "vision" assert odometry.child_frame_id == "body" assert (odometry.header.stamp.sec, odometry.header.stamp.nanosec) == (-1, 500000000) diff --git a/dimos/experimental/robot/bosdyn/spot/recorder.py b/dimos/experimental/robot/bosdyn/spot/recorder.py index eb160dd4c26..e9434057acf 100644 --- a/dimos/experimental/robot/bosdyn/spot/recorder.py +++ b/dimos/experimental/robot/bosdyn/spot/recorder.py @@ -25,18 +25,17 @@ from datetime import datetime from pathlib import Path +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import CameraInfo, Image from pydantic import Field from dimos.core.stream import In from dimos.experimental.robot.bosdyn.spot.config import CAMERA_STREAM_SUFFIXES from dimos.memory.module import Recorder, RecorderConfig from dimos.memory.type.recording import OnExisting -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image # jpeg codec quantises depth it to ~25cm and adds block artifacts (horrible) -LOSSLESS_CODEC = "lz4+lcm" +LOSSLESS_CODEC = "lz4+cdr" class SpotRecorderConfig(RecorderConfig): diff --git a/dimos/experimental/robot/bosdyn/spot/rerun.py b/dimos/experimental/robot/bosdyn/spot/rerun.py index 750bde20dfb..a40b790fb30 100644 --- a/dimos/experimental/robot/bosdyn/spot/rerun.py +++ b/dimos/experimental/robot/bosdyn/spot/rerun.py @@ -18,15 +18,17 @@ from typing import TYPE_CHECKING, Any +from dimos.msgs.time import to_seconds from dimos.protocol.tf.tf import MultiTBuffer +from dimos.visualization.rerun.message_helpers import image_archetype, tf_archetypes if TYPE_CHECKING: from collections.abc import Callable + from dimos_generated.sensor_msgs.msg import CameraInfo, Image + from dimos_generated.tf2_msgs.msg import TFMessage import rerun.blueprint as rrb - from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo - from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.visualization.rerun.bridge import RerunData # Optical tf frame_id (SpotHighLevelConfig defaults) -> stream-name suffix. @@ -88,11 +90,11 @@ def _camera_entity(origin: str) -> str: def _tf_to_rerun(tf_message: TFMessage) -> RerunData: """Log tf as usual, and keep a copy for the camera pose lookups.""" _tf_buffer.receive_tfmessage(tf_message) - return list(tf_message.to_rerun()) + return [(path, archetype) for path, archetype in tf_archetypes(tf_message)] def _camera_info_pinhole(camera_info: CameraInfo, origin: Callable[[str], str]) -> RerunData | None: - suffix = _OPTICAL_FRAME_TO_SUFFIX.get(camera_info.frame_id) + suffix = _OPTICAL_FRAME_TO_SUFFIX.get(camera_info.header.frame_id) if suffix is None: return None import rerun as rr @@ -101,8 +103,8 @@ def _camera_info_pinhole(camera_info: CameraInfo, origin: Callable[[str], str]) ( _camera_entity(origin(suffix)), rr.Pinhole( - focal_length=[camera_info.K[0], camera_info.K[4]], - principal_point=[camera_info.K[2], camera_info.K[5]], + focal_length=[camera_info.k[0], camera_info.k[4]], + principal_point=[camera_info.k[2], camera_info.k[5]], width=camera_info.width, height=camera_info.height, image_plane_distance=_FRUSTUM_PLANE_DISTANCE, @@ -237,11 +239,11 @@ class _ImageBakedIntoAnchor: def __init__(self, origin: str) -> None: self.origin = origin - def __call__(self, image: Any) -> RerunData | None: + def __call__(self, image: Image) -> RerunData | None: anchor = _anchor_frame() if anchor is None: return None - pose = _tf_buffer.get(anchor, image.frame_id, image.ts) + pose = _tf_buffer.get(anchor, image.header.frame_id, to_seconds(image.header.stamp)) if pose is None: return None import rerun as rr @@ -250,10 +252,21 @@ def __call__(self, image: Any) -> RerunData | None: ( self.origin, rr.Transform3D( - translation=[pose.translation.x, pose.translation.y, pose.translation.z], - rotation=pose.rotation.to_rerun(), + translation=[ + pose.transform.translation.x, + pose.transform.translation.y, + pose.transform.translation.z, + ], + quaternion=rr.Quaternion( + xyzw=[ + pose.transform.rotation.x, + pose.transform.rotation.y, + pose.transform.rotation.z, + pose.transform.rotation.w, + ] + ), parent_frame=f"tf#/{anchor}", ), ), - (_camera_entity(self.origin), image.to_rerun()), + (_camera_entity(self.origin), image_archetype(image)), ] diff --git a/dimos/experimental/robot/bosdyn/spot/test_utils.py b/dimos/experimental/robot/bosdyn/spot/test_utils.py index 4acbe37c943..1ae124549fb 100644 --- a/dimos/experimental/robot/bosdyn/spot/test_utils.py +++ b/dimos/experimental/robot/bosdyn/spot/test_utils.py @@ -18,8 +18,9 @@ from types import ModuleType, SimpleNamespace import cv2 +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time -from dimos_generated.geometry_msgs.msg import Transform, TransformStamped, Vector3 +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 from dimos_generated.sensor_msgs.msg import Image from dimos_generated.std_msgs.msg import Header import numpy as np @@ -74,7 +75,9 @@ def test_optical_roll_preserves_mount_and_exact_stamp(turns): edge = TransformStamped( header=Header(frame_id="body", stamp=Time(sec=1700000000, nanosec=123456789)), child_frame_id="optical", - transform=Transform(translation=Vector3(x=1, y=2, z=3)), + transform=Transform( + translation=Vector3(x=1, y=2, z=3), rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), ) rolled = roll_optical_frame(edge, turns) assert rolled.header == edge.header @@ -130,7 +133,7 @@ def test_sdk_raw_camera_boundary_emits_generated_pixels( clock = SimpleNamespace(local_seconds_from_robot_timestamp=lambda stamp: 1700000000.25) value = decode_image(response, "camera-optical", clock) assert type(value) is Image and value.encoding == encoding - decoded = Image.decode(value.encode()) + decoded = cdr_decode(cdr_encode(value), Image) assert decoded.header.frame_id == "camera-optical" assert decoded.header.stamp == Time(sec=1700000000, nanosec=250000000) np.testing.assert_array_equal(image_view(decoded), pixels) @@ -145,7 +148,7 @@ def test_sdk_jpeg_camera_boundary_decodes_generated_image(image_sdk): clock = SimpleNamespace(local_seconds_from_robot_timestamp=lambda stamp: 1700000000.25) value = decode_image(response, "camera-optical", clock) assert type(value) is Image and value.encoding == "bgr8" - np.testing.assert_allclose(image_view(Image.decode(value.encode())), pixels, atol=2) + np.testing.assert_allclose(image_view(cdr_decode(cdr_encode(value), Image)), pixels, atol=2) @pytest.mark.parametrize("turns", [-1, 0, 1, 2, 4]) @@ -154,7 +157,7 @@ def test_generated_camera_rotation_preserves_depth_header_and_intrinsics(turns): pixels = np.arange(6, dtype=np.uint16).reshape(2, 3) * 1000 image = image_from_array(pixels, encoding="16UC1", header=header) info = camera_info_from_intrinsics(10, 20, 1, 0.5, 3, 2, header=header) - before_image, before_info = image.encode(), info.encode() + before_image, before_info = cdr_encode(image), cdr_encode(info) rotated = rotate_image_quarter_turns(image, turns) calibration = rotate_camera_info_quarter_turns(info, turns) np.testing.assert_array_equal(image_view(rotated), np.rot90(pixels, k=turns)) @@ -166,7 +169,7 @@ def test_generated_camera_rotation_preserves_depth_header_and_intrinsics(turns): cx, cy = cy, width - 1 - cx width, height = height, width assert list(calibration.k) == [expected_fx, 0, cx, 0, expected_fy, cy, 0, 0, 1] - assert image.encode() == before_image and info.encode() == before_info + assert cdr_encode(image) == before_image and cdr_encode(info) == before_info def test_sdk_calibration_boundary_emits_plain_generated_value(): diff --git a/dimos/experimental/robot/bosdyn/spot/utils.py b/dimos/experimental/robot/bosdyn/spot/utils.py index 30e5d7c80b3..a90a178bef3 100644 --- a/dimos/experimental/robot/bosdyn/spot/utils.py +++ b/dimos/experimental/robot/bosdyn/spot/utils.py @@ -20,6 +20,7 @@ from pathlib import Path from typing import Any +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Transform, TransformStamped, Vector3 from dimos_generated.sensor_msgs.msg import CameraInfo, Image from dimos_generated.std_msgs.msg import Header @@ -87,7 +88,7 @@ def decode_image(response: Any, frame_id: str, time_converter: Any) -> Image | N def joint_to_transform(joint: JointDescription) -> TransformStamped: return TransformStamped( - header=Header(frame_id=joint.parent_link), + header=Header(frame_id=joint.parent_link, stamp=Time(sec=0, nanosec=0)), child_frame_id=joint.child_link, transform=Transform( translation=Vector3( @@ -140,9 +141,12 @@ def roll_optical_frame(transform: TransformStamped, quarter_turns: int) -> Trans if not quarter_turns: return transform roll = TransformStamped( - header=Header(frame_id=transform.child_frame_id), + header=Header(frame_id=transform.child_frame_id, stamp=Time(sec=0, nanosec=0)), child_frame_id=transform.child_frame_id, - transform=Transform(rotation=quaternion_from_euler(0, 0, quarter_turns * math.pi / 2)), + transform=Transform( + rotation=quaternion_from_euler(0, 0, quarter_turns * math.pi / 2), + translation=Vector3(x=0.0, y=0.0, z=0.0), + ), ) return compose_transforms(transform, roll) diff --git a/dimos/experimental/security_demo/conftest.py b/dimos/experimental/security_demo/conftest.py index 6e7cc79d682..5b0fded37ad 100644 --- a/dimos/experimental/security_demo/conftest.py +++ b/dimos/experimental/security_demo/conftest.py @@ -14,7 +14,9 @@ from __future__ import annotations -from dimos_generated.sensor_msgs.msg import CameraInfo +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.sensor_msgs.msg import CameraInfo, RegionOfInterest +from dimos_generated.std_msgs.msg import Header import numpy as np from PIL import Image as PILImage import pytest @@ -57,7 +59,21 @@ def security_module(mocker): mocker.patch("dimos.experimental.security_demo.security_module.EdgeTAMProcessor") mocker.patch("dimos.experimental.security_demo.security_module.DepthEstimator") - module = SecurityModule(camera_info=CameraInfo()) + module = SecurityModule( + camera_info=CameraInfo( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + height=0, + width=0, + distortion_model="", + d=np.array([], dtype=np.float64), + k=np.zeros(9, dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), + ) + ) # Replace output streams with mocks for test assertions module.detection = mocker.MagicMock() diff --git a/dimos/experimental/security_demo/security_module.py b/dimos/experimental/security_demo/security_module.py index 9ba004e8c70..7853fd4cfa8 100644 --- a/dimos/experimental/security_demo/security_module.py +++ b/dimos/experimental/security_demo/security_module.py @@ -20,7 +20,7 @@ from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT import numpy as np -from dimos_lcm.std_msgs import String, Bool +from dimos_generated.std_msgs.msg import String, Bool from reactivex.disposable import Disposable from dimos.agents.annotation import skill @@ -32,19 +32,22 @@ from dimos.models.segmentation.edge_tam import EdgeTAMProcessor from dimos.perception.detection.detectors.person.yolo import YoloPersonDetector from dimos.perception.detection.type.detection2d.person import Detection2DPerson -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.geometry_msgs.msg import Twist +from dimos_generated.nav_msgs.msg import OccupancyGrid +from dimos_generated.sensor_msgs.msg import CameraInfo +from dimos_generated.sensor_msgs.msg import Image from dimos.navigation.experimental.patrolling.create_patrol_router import create_patrol_router from dimos.navigation.experimental.patrolling.routers.patrol_router import PatrolRouter from dimos.agents.skills.speak_skill_spec import SpeakSkillSpec from dimos.navigation.go2.replanning_a_star.module_spec import ReplanningAStarPlannerSpec from dimos.agents.skills.visual_servoing.visual_servoing_2d import VisualServoing2D from dimos.utils.logging_config import setup_logger +from dimos.msgs.image import image_to_bgr, image_from_array from dimos.navigation.experimental.patrolling.constants import EXTRA_CLEARANCE +from dimos_generated.geometry_msgs.msg import Vector3 + if TYPE_CHECKING: from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox @@ -253,7 +256,9 @@ def _main_loop(self) -> None: case "FOLLOWING": self._follow_step() - self.cmd_vel.publish(Twist.zero()) + self.cmd_vel.publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) self._transition_to("IDLE") def _patrol_step(self) -> None: @@ -298,11 +303,11 @@ def _patrol_step(self) -> None: area=f"{best.bbox_2d_volume():.0f}px", ) - annotated = image.data.copy() + annotated = image_to_bgr(image) best.draw_on(annotated) if isinstance(best, Detection2DPerson): _draw_skeleton(annotated, best) - self.detection.publish(Image.from_numpy(annotated, format=image.format)) + self.detection.publish(image_from_array(annotated, encoding="bgr8", header=image.header)) # Init EdgeTAM with YOLO bbox for continuous tracking box = np.array(list(best.bbox), dtype=np.float32) @@ -327,7 +332,9 @@ def _follow_step(self) -> None: detections = self._tracker.process_image(latest_image) if len(detections) == 0: - self.cmd_vel.publish(Twist.zero()) + self.cmd_vel.publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) self._speak_skill.speak("Lost sight of intruder, resuming patrol", blocking=False) self._router.reset() self._has_active_goal = False @@ -338,7 +345,7 @@ def _follow_step(self) -> None: twist = self._visual_servo.compute_twist(best.bbox, latest_image.width) self.cmd_vel.publish(twist) - overlay = latest_image.data.copy() + overlay = image_to_bgr(latest_image) if hasattr(best, "mask") and best.mask is not None: mask_bool = best.mask > 0 green = np.zeros_like(overlay) @@ -356,7 +363,7 @@ def _follow_step(self) -> None: _draw_skeleton(overlay, person) self.tracking_image.publish( - Image.from_numpy(overlay, format=latest_image.format, ts=latest_image.ts) + image_from_array(overlay, encoding="bgr8", header=latest_image.header) ) time.sleep(1.0 / self.config.follow_frequency) @@ -382,7 +389,7 @@ def _transition_to(self, new_state: State) -> None: self._state = new_state logger.info("state transition", old=old, new=new_state) - self.security_state.publish(String(new_state)) + self.security_state.publish(String(data=new_state)) def _stop_security_patrol_internal(self) -> None: self._stop_event.set() diff --git a/dimos/experimental/security_demo/test_security_module.py b/dimos/experimental/security_demo/test_security_module.py index fa308ff475a..3526c9c2cd7 100644 --- a/dimos/experimental/security_demo/test_security_module.py +++ b/dimos/experimental/security_demo/test_security_module.py @@ -16,7 +16,9 @@ import threading +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Twist, Vector3 +from dimos_generated.std_msgs.msg import Header import pytest from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D @@ -50,7 +52,10 @@ def test_patrol_step_transitions_to_following_on_detection( module._has_active_goal = True module._latest_image = person_image module._latest_pose = PoseStamped( - pose=Pose(position=Point(x=1, y=2), orientation=Quaternion(w=1)) + pose=Pose( + position=Point(x=1, y=2, z=0.0), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) det = make_detection() @@ -79,7 +84,12 @@ def test_patrol_step_requests_goal_when_no_active_goal(security_module): module._has_active_goal = False module._latest_image = None # causes early return after goal logic - goal = PoseStamped(pose=Pose(position=Point(x=5, y=5), orientation=Quaternion(w=1))) + goal = PoseStamped( + pose=Pose( + position=Point(x=5, y=5, z=0.0), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) module._router.next_goal.return_value = goal module._patrol_step() @@ -101,7 +111,7 @@ def test_follow_step_publishes_twist_when_tracking( image=person_image, detections=[det] ) - twist = Twist(linear=Vector3(x=0.3), angular=Vector3(z=0.1)) + twist = Twist(linear=Vector3(x=0.3, y=0.0, z=0.0), angular=Vector3(z=0.1, x=0.0, y=0.0)) module._visual_servo.compute_twist.return_value = twist mocker.patch("dimos.experimental.security_demo.security_module.time.sleep") @@ -126,7 +136,9 @@ def test_follow_step_transitions_to_patrolling_on_person_lost(security_module, p # Zero twist should be published to stop the robot published_twist = module.cmd_vel.publish.call_args[0][0] - assert published_twist == Twist() + assert published_twist == Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) module._speak_skill.speak.assert_called_with( "Lost sight of intruder, resuming patrol", blocking=False @@ -158,7 +170,9 @@ def patrol_side_effect(): # Verify zero twist published on shutdown last_twist = module.cmd_vel.publish.call_args[0][0] - assert last_twist == Twist() + assert last_twist == Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) # Verify state transitions: PATROLLING then IDLE state_values = [call.args[0].data for call in module.security_state.publish.call_args_list] diff --git a/dimos/experimental/world_belief/test_absence.py b/dimos/experimental/world_belief/test_absence.py index 0f7172da35d..ff34eb68a1f 100644 --- a/dimos/experimental/world_belief/test_absence.py +++ b/dimos/experimental/world_belief/test_absence.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 from dimos_generated.std_msgs.msg import Header import numpy as np @@ -32,32 +33,36 @@ ) def test_visibility_inverts_generated_camera_transform(depth, expected): calibration = camera_info_from_intrinsics( - 100, 100, 16, 16, 32, 32, header=Header(frame_id="camera") + 100, 100, 16, 16, 32, 32, header=Header(frame_id="camera", stamp=Time(sec=0, nanosec=0)) ) world_from_camera = TransformStamped( - header=Header(frame_id="world"), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), child_frame_id="camera", - transform=Transform(translation=Vector3(x=10), rotation=Quaternion(w=1)), + transform=Transform( + translation=Vector3(x=10, y=0.0, z=0.0), rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), ) # In world coordinates x=10; after inverse transform it is on the optical axis. - center = Vector3(x=10, z=1) + center = Vector3(x=10, z=1, y=0.0) depth_m = np.full((32, 32), depth, dtype=np.float32) assert classify_visibility(center, calibration, world_from_camera, depth_m) == expected assert ( - classify_visibility(Vector3(x=10, z=-1), calibration, world_from_camera, depth_m) + classify_visibility(Vector3(x=10, z=-1, y=0.0), calibration, world_from_camera, depth_m) == OUT_OF_VIEW ) def test_visibility_requires_all_valid_depth_samples_to_show_absence(): calibration = camera_info_from_intrinsics( - 100, 100, 16, 16, 32, 32, header=Header(frame_id="camera") + 100, 100, 16, 16, 32, 32, header=Header(frame_id="camera", stamp=Time(sec=0, nanosec=0)) ) identity = TransformStamped( - header=Header(frame_id="world"), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), child_frame_id="camera", - transform=Transform(rotation=Quaternion(w=1)), + transform=Transform( + rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), translation=Vector3(x=0.0, y=0.0, z=0.0) + ), ) depth = np.full((32, 32), 3, dtype=np.float32) depth[16, 16] = 1 - assert classify_visibility(Vector3(z=1), calibration, identity, depth) == PRESENT + assert classify_visibility(Vector3(z=1, x=0.0, y=0.0), calibration, identity, depth) == PRESENT diff --git a/dimos/experimental/world_belief/test_generated_recall.py b/dimos/experimental/world_belief/test_generated_recall.py index 9b696b79c4b..712196274ba 100644 --- a/dimos/experimental/world_belief/test_generated_recall.py +++ b/dimos/experimental/world_belief/test_generated_recall.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos_generated.geometry_msgs.msg import Pose, Quaternion +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion from dimos_generated.sensor_msgs.msg import Image from dimos_generated.std_msgs.msg import Header import numpy as np @@ -42,7 +42,9 @@ def test_index_selects_sharp_generated_frame_and_preserves_stamp_and_thumbnail() image = image_from_array( sharp, encoding="rgb8", header=Header(stamp=time_from_seconds(0.2), frame_id="camera") ) - identity = Pose(orientation=Quaternion(w=1.0)) + identity = Pose( + orientation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), position=Point(x=0.0, y=0.0, z=0.0) + ) frames.append( image_from_array(np.full_like(sharp, 150), encoding="rgb8", header=image.header), ts=0.1, diff --git a/dimos/experimental/world_belief/worldbelief_recorder.py b/dimos/experimental/world_belief/worldbelief_recorder.py index cfe344623a8..7f6e1e29786 100644 --- a/dimos/experimental/world_belief/worldbelief_recorder.py +++ b/dimos/experimental/world_belief/worldbelief_recorder.py @@ -18,7 +18,7 @@ from pathlib import Path from typing import Any, Protocol -from dimos_generated.geometry_msgs.msg import Pose, Quaternion +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion from dimos_generated.sensor_msgs.msg import CameraInfo, Image, JointState from dimos.constants import STATE_DIR @@ -89,4 +89,6 @@ def _prepare_streams(self) -> None: @pose_setter_for("coordinator_joint_state") async def _proprio_pose(self, msg: Any) -> Any: """Use an identity pose for proprioceptive joint-state records.""" - return Pose(orientation=Quaternion(w=1.0)) + return Pose( + orientation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), position=Point(x=0.0, y=0.0, z=0.0) + ) diff --git a/dimos/experimental/world_belief/xarm6_blueprint.py b/dimos/experimental/world_belief/xarm6_blueprint.py index 543bb16ef40..d59f00dd6f8 100644 --- a/dimos/experimental/world_belief/xarm6_blueprint.py +++ b/dimos/experimental/world_belief/xarm6_blueprint.py @@ -19,7 +19,9 @@ from functools import partial from typing import TYPE_CHECKING, Any +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.std_msgs.msg import Header from dimos.agents.mcp.mcp_server import McpServer from dimos.constants import STATE_DIR @@ -44,7 +46,9 @@ transform=Transform( translation=Vector3(x=0.06693724, y=-0.0309563, z=0.00691482), rotation=Quaternion(x=0.70513398, y=0.00535696, z=0.70897578, w=-0.01052180), - ) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + child_frame_id="", ) diff --git a/dimos/hardware/drive_trains/flowbase/config.py b/dimos/hardware/drive_trains/flowbase/config.py index a002be9596d..af54e8677db 100644 --- a/dimos/hardware/drive_trains/flowbase/config.py +++ b/dimos/hardware/drive_trains/flowbase/config.py @@ -17,4 +17,6 @@ from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion # Mid-360 mount pose on the FlowBase in the base frame. -FLOWBASE_MID360_MOUNT = Pose(position=Point(x=0.20, y=-0.20, z=0.10), orientation=Quaternion(w=1.0)) +FLOWBASE_MID360_MOUNT = Pose( + position=Point(x=0.20, y=-0.20, z=0.10), orientation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0) +) diff --git a/dimos/hardware/drive_trains/transport/test_adapter.py b/dimos/hardware/drive_trains/transport/test_adapter.py index eb2f7d8905d..18b3204a2e6 100644 --- a/dimos/hardware/drive_trains/transport/test_adapter.py +++ b/dimos/hardware/drive_trains/transport/test_adapter.py @@ -14,7 +14,10 @@ import math +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Twist, Vector3 +from dimos_generated.std_msgs.msg import Header import pytest from dimos.hardware.drive_trains.transport.adapter import TransportTwistAdapter @@ -26,9 +29,10 @@ def test_nested_pose_becomes_planar_odometry_without_sharing_state(): adapter._on_odom( PoseStamped( pose=Pose( - position=Point(x=2.0, y=-1.0), - orientation=Quaternion(z=math.sqrt(0.5), w=math.sqrt(0.5)), - ) + position=Point(x=2.0, y=-1.0, z=0.0), + orientation=Quaternion(z=math.sqrt(0.5), w=math.sqrt(0.5), x=0.0, y=0.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) ) odometry = adapter.read_odometry() @@ -40,9 +44,21 @@ def test_nested_pose_becomes_planar_odometry_without_sharing_state(): @pytest.mark.parametrize( "dof,values,expected", [ - (1, [0.1], Twist(linear=Vector3(x=0.1))), - (2, [0.1, 0.2], Twist(linear=Vector3(x=0.1, y=0.2))), - (3, [0.1, 0.2, 0.3], Twist(linear=Vector3(x=0.1, y=0.2), angular=Vector3(z=0.3))), + ( + 1, + [0.1], + Twist(linear=Vector3(x=0.1, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + ), + ( + 2, + [0.1, 0.2], + Twist(linear=Vector3(x=0.1, y=0.2, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + ), + ( + 3, + [0.1, 0.2, 0.3], + Twist(linear=Vector3(x=0.1, y=0.2, z=0.0), angular=Vector3(z=0.3, x=0.0, y=0.0)), + ), ], ) def test_generated_commands_preserve_enable_and_stop_contract( @@ -56,8 +72,10 @@ def test_generated_commands_preserve_enable_and_stop_contract( adapter.write_enable(True) assert adapter.write_velocities(values) message = transport.publish.call_args.args[0] - assert Twist.decode(message.encode()) == expected + assert cdr_decode(cdr_encode(message), Twist) == expected assert adapter.read_velocities() == values adapter.write_enable(False) - assert transport.publish.call_args.args[0] == Twist() + assert transport.publish.call_args.args[0] == Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) assert adapter.read_velocities() == [0.0] * dof diff --git a/dimos/hardware/sensors/camera/gstreamer/test_generated_frames.py b/dimos/hardware/sensors/camera/gstreamer/test_generated_frames.py index cc413b20089..af1b107373c 100644 --- a/dimos/hardware/sensors/camera/gstreamer/test_generated_frames.py +++ b/dimos/hardware/sensors/camera/gstreamer/test_generated_frames.py @@ -18,6 +18,7 @@ from types import ModuleType, SimpleNamespace from unittest.mock import Mock +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.sensor_msgs.msg import Image import numpy as np import pytest @@ -70,7 +71,7 @@ def test_generated_bgr_frame_owns_pixels_after_sdk_buffer_unmap(camera_module, m try: assert camera._on_new_sample(sink) == 0 assert len(received) == 1 and type(received[0]) is Image - value = Image.decode(received[0].encode()) + value = cdr_decode(cdr_encode(received[0]), Image) assert value.header.frame_id == "camera_test" assert value.encoding == "bgr8" and value.step == 9 assert to_nanoseconds(value.header.stamp) == 1700000000125000000 diff --git a/dimos/hardware/sensors/camera/module.py b/dimos/hardware/sensors/camera/module.py index 2ef9856b843..880974628b5 100644 --- a/dimos/hardware/sensors/camera/module.py +++ b/dimos/hardware/sensors/camera/module.py @@ -13,10 +13,11 @@ # limitations under the License. from collections.abc import Callable -from copy import copy +from copy import deepcopy import time -from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 from dimos_generated.sensor_msgs.msg import CameraInfo, Image from dimos_generated.std_msgs.msg import Header from dimos_generated.tf2_msgs.msg import TFMessage @@ -39,9 +40,12 @@ def default_transform() -> TransformStamped: return TransformStamped( - header=Header(frame_id="base_link"), + header=Header(frame_id="base_link", stamp=Time(sec=0, nanosec=0)), child_frame_id="camera_link", - transform=Transform(rotation=Quaternion(w=1.0)), + transform=Transform( + rotation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), + translation=Vector3(x=0.0, y=0.0, z=0.0), + ), ) @@ -83,19 +87,22 @@ def start(self) -> None: ) def publish_metadata(self) -> None: - camera_info = copy(self.hardware.camera_info) + camera_info = deepcopy(self.hardware.camera_info) camera_info.header.stamp = time_from_nanoseconds(time.time_ns()) self.camera_info.publish(camera_info) if not self.config.transform: return - camera_link = copy(self.config.transform) + camera_link = deepcopy(self.config.transform) camera_link.header.stamp = camera_info.header.stamp camera_optical = TransformStamped( header=Header(frame_id=camera_link.child_frame_id, stamp=camera_info.header.stamp), child_frame_id=camera_info.header.frame_id, - transform=Transform(rotation=Quaternion(x=-0.5, y=0.5, z=-0.5, w=0.5)), + transform=Transform( + rotation=Quaternion(x=-0.5, y=0.5, z=-0.5, w=0.5), + translation=Vector3(x=0.0, y=0.0, z=0.0), + ), ) self.tf.publish(TFMessage(transforms=[camera_link, camera_optical])) diff --git a/dimos/hardware/sensors/camera/realsense/blueprints.py b/dimos/hardware/sensors/camera/realsense/blueprints.py index c7c449b58fd..939cf661545 100644 --- a/dimos/hardware/sensors/camera/realsense/blueprints.py +++ b/dimos/hardware/sensors/camera/realsense/blueprints.py @@ -25,7 +25,8 @@ from typing import Any -from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 from dimos_generated.sensor_msgs.msg import PointCloud2 from dimos_generated.std_msgs.msg import Header @@ -50,9 +51,12 @@ class RealSenseMountTf(StaticTfPublisher): def transforms(self) -> list[TransformStamped]: return [ TransformStamped( - header=Header(frame_id="world"), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), child_frame_id="camera_link", - transform=Transform(rotation=Quaternion(w=1.0)), + transform=Transform( + rotation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), + translation=Vector3(x=0.0, y=0.0, z=0.0), + ), ) ] diff --git a/dimos/hardware/sensors/camera/realsense/camera.py b/dimos/hardware/sensors/camera/realsense/camera.py index 3afd870ba36..2ea3a8b04c0 100644 --- a/dimos/hardware/sensors/camera/realsense/camera.py +++ b/dimos/hardware/sensors/camera/realsense/camera.py @@ -19,8 +19,10 @@ from typing import TYPE_CHECKING, Any +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.dimos_msgs.msg import ImuInfo from dimos_generated.sensor_msgs.msg import CameraInfo, Image, Imu, PointCloud2 +from dimos_generated.std_msgs.msg import Header from dimos_generated.tf2_msgs.msg import TFMessage from pydantic import Field @@ -64,6 +66,8 @@ class RealSenseCameraConfig(NativeModuleConfig, DepthCameraConfig): gyro_random_walk=1.0e-5, accel_noise_density=1.8e-3, accel_random_walk=1.0e-4, + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + frequency=0.0, ) ) pointcloud_fps: float = 5.0 diff --git a/dimos/hardware/sensors/camera/realsense/rust/Cargo.lock b/dimos/hardware/sensors/camera/realsense/rust/Cargo.lock index 2ab58459a48..5e5fb148703 100644 --- a/dimos/hardware/sensors/camera/realsense/rust/Cargo.lock +++ b/dimos/hardware/sensors/camera/realsense/rust/Cargo.lock @@ -557,7 +557,9 @@ dependencies = [ name = "dimos-generated-messages" version = "0.1.0" dependencies = [ + "heck", "re_cdr", + "ros2msg", "serde", "serde-big-array", ] @@ -988,6 +990,12 @@ version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + [[package]] name = "hermit-abi" version = "0.5.2" @@ -1930,6 +1938,16 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + [[package]] name = "proc-macro-crate" version = "3.5.0" @@ -2318,6 +2336,23 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "ros2msg" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fe291be5cb6d4b8e6bf9c9de80c40bd1cb09ecac50deafa2fdcae0027e04988" +dependencies = [ + "heck", + "pest", + "pest_derive", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "syn 2.0.119", + "thiserror 2.0.20", +] + [[package]] name = "rustc-hash" version = "2.1.3" diff --git a/dimos/hardware/sensors/camera/realsense/rust/Cargo.toml b/dimos/hardware/sensors/camera/realsense/rust/Cargo.toml index 4f66e1f7922..0f0041f26bc 100644 --- a/dimos/hardware/sensors/camera/realsense/rust/Cargo.toml +++ b/dimos/hardware/sensors/camera/realsense/rust/Cargo.toml @@ -36,7 +36,7 @@ infrared_right_camera_info = "sensor_msgs/msg/CameraInfo" [dependencies] dimos-module = { path = "../../../../../../native/rust/dimos-module" } -dimos-generated-messages = { path = "../../../../../message_codegen" } +dimos-generated-messages = { path = "../../../../../../packages/dimos-generated/build/message-source/rust" } realsense-rust = "1.3" realsense-sys = "2.56" tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal"] } diff --git a/dimos/hardware/sensors/camera/realsense/rust/src/main.rs b/dimos/hardware/sensors/camera/realsense/rust/src/main.rs index a6d22933cf4..3a91e1035ae 100644 --- a/dimos/hardware/sensors/camera/realsense/rust/src/main.rs +++ b/dimos/hardware/sensors/camera/realsense/rust/src/main.rs @@ -24,15 +24,16 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{mpsc, Arc, Mutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use dimos_generated_messages::builtin_interfaces::msg::Time; -use dimos_generated_messages::dimos_msgs::msg::ImuInfo; +use dimos_generated_messages::builtin_interfaces::msg::time::Time; +use dimos_generated_messages::dimos_msgs::msg::imu_info::ImuInfo; use dimos_generated_messages::geometry_msgs::msg::{ - Quaternion as QuaternionMsg, Vector3 as Vector3Msg, + quaternion::Quaternion as QuaternionMsg, vector3::Vector3 as Vector3Msg, }; use dimos_generated_messages::sensor_msgs::msg::{ - CameraInfo, Image, Imu, PointCloud2, PointField, RegionOfInterest, + camera_info::CameraInfo, image::Image, imu::Imu, point_cloud2::PointCloud2, + point_field::PointField, region_of_interest::RegionOfInterest, }; -use dimos_generated_messages::std_msgs::msg::Header; +use dimos_generated_messages::std_msgs::msg::header::Header; use dimos_module::cdr; use dimos_module::nalgebra::{ Isometry3, Matrix3, Quaternion, Rotation3, Translation3, UnitQuaternion, Vector3, diff --git a/dimos/hardware/sensors/camera/test_module.py b/dimos/hardware/sensors/camera/test_module.py index ec1c5b30519..4dd42050b71 100644 --- a/dimos/hardware/sensors/camera/test_module.py +++ b/dimos/hardware/sensors/camera/test_module.py @@ -14,6 +14,7 @@ """Camera metadata uses the same generated header and frame chain as its images.""" +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.sensor_msgs.msg import CameraInfo from dimos_generated.tf2_msgs.msg import TFMessage import numpy as np @@ -40,14 +41,14 @@ def test_metadata_preserves_calibration_and_does_not_mutate_mount_template(modul mocker.patch( "dimos.hardware.sensors.camera.module.time.time_ns", return_value=1700000000123456789 ) - original = module.config.transform.encode() + original = cdr_encode(module.config.transform) module.publish_metadata() info_publish.assert_called_once() tf_publish.assert_called_once() - info = CameraInfo.decode(info_publish.call_args.args[0].encode()) - transforms = TFMessage.decode(tf_publish.call_args.args[0].encode()).transforms + info = cdr_decode(cdr_encode(info_publish.call_args.args[0]), CameraInfo) + transforms = cdr_decode(cdr_encode(tf_publish.call_args.args[0]), TFMessage).transforms assert (info.header.stamp.sec, info.header.stamp.nanosec) == (1700000000, 123456789) assert len(transforms) == 2 assert [(edge.header.frame_id, edge.child_frame_id) for edge in transforms] == [ @@ -55,5 +56,5 @@ def test_metadata_preserves_calibration_and_does_not_mutate_mount_template(modul ("camera_link", "robot/camera_optical"), ] assert all(edge.header.stamp == info.header.stamp for edge in transforms) - assert module.config.transform.encode() == original + assert cdr_encode(module.config.transform) == original np.testing.assert_allclose(transform_matrix(transforms[0].transform), np.eye(4)) diff --git a/dimos/hardware/sensors/camera/test_webcam.py b/dimos/hardware/sensors/camera/test_webcam.py index 44264d4af3e..94998be5fca 100644 --- a/dimos/hardware/sensors/camera/test_webcam.py +++ b/dimos/hardware/sensors/camera/test_webcam.py @@ -12,7 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, RegionOfInterest from dimos_generated.std_msgs.msg import Header import numpy as np import pytest @@ -30,12 +32,28 @@ def test_webcam_without_intrinsics_reports_a_nominal_pinhole() -> None: def test_webcam_keeps_configured_intrinsics() -> None: - configured = camera_info_from_fov(90.0, 640, 480, header=Header()) + configured = camera_info_from_fov( + 90.0, 640, 480, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="") + ) assert Webcam(camera_info=configured).camera_info is configured def test_webcam_with_size_but_no_focal_length_gets_a_nominal_pinhole() -> None: - info = Webcam(camera_info=CameraInfo(width=1280, height=720)).camera_info + info = Webcam( + camera_info=CameraInfo( + width=1280, + height=720, + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + distortion_model="", + d=np.array([], dtype=np.float64), + k=np.zeros(9, dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), + ) + ).camera_info assert info.k[0] > 0 @@ -53,7 +71,7 @@ def test_capture_converts_bgr_and_crops_before_cdr_encoding(side, start, mocker) camera._capture = capture mocker.patch("dimos.msgs.time.time.time_ns", return_value=1700000000123456789) try: - message = Image.decode(camera.capture_frame().encode()) + message = cdr_decode(cdr_encode(camera.capture_frame()), Image) finally: camera.stop() assert message.encoding == "rgb8" diff --git a/dimos/hardware/sensors/camera/v4l2/module.py b/dimos/hardware/sensors/camera/v4l2/module.py index 65cf3214d1e..4620c3f6b3f 100644 --- a/dimos/hardware/sensors/camera/v4l2/module.py +++ b/dimos/hardware/sensors/camera/v4l2/module.py @@ -24,11 +24,11 @@ from typing import TYPE_CHECKING +from dimos_generated.sensor_msgs.msg import CompressedImage from pydantic import Field from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.stream import Out -from dimos.msgs.sensor_msgs.CompressedImage import CompressedImage class V4L2CameraConfig(NativeModuleConfig): diff --git a/dimos/hardware/sensors/camera/v4l2/rust/Cargo.toml b/dimos/hardware/sensors/camera/v4l2/rust/Cargo.toml index 40540bf7eac..3234d8271c9 100644 --- a/dimos/hardware/sensors/camera/v4l2/rust/Cargo.toml +++ b/dimos/hardware/sensors/camera/v4l2/rust/Cargo.toml @@ -22,7 +22,7 @@ jpeg_out = "sensor_msgs/msg/CompressedImage" [dependencies] dimos-module = { path = "../../../../../../native/rust/dimos-module" } -dimos-generated-messages = { path = "../../../../../../packages/dimos-generated/src/dimos_generated_schemas/package/rust" } +dimos-generated-messages = { path = "../../../../../../packages/dimos-generated/build/message-source/rust" } tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal"] } serde = { version = "1", features = ["derive"] } tracing = "0.1" diff --git a/dimos/hardware/sensors/camera/webcam.py b/dimos/hardware/sensors/camera/webcam.py index a64f6081b9d..f6393d71056 100644 --- a/dimos/hardware/sensors/camera/webcam.py +++ b/dimos/hardware/sensors/camera/webcam.py @@ -18,8 +18,10 @@ import time from typing import Annotated, Any, Literal -from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, RegionOfInterest from dimos_generated.std_msgs.msg import Header +import numpy as np from pydantic import BeforeValidator, Field from reactivex import create from reactivex.observable import Observable @@ -48,7 +50,21 @@ class WebcamConfig(CameraConfig): width: int = 640 height: int = 480 fps: float = 15.0 - camera_info: CameraInfo = Field(default_factory=CameraInfo) + camera_info: CameraInfo = Field( + default_factory=lambda: CameraInfo( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + height=0, + width=0, + distortion_model="", + d=np.array([], dtype=np.float64), + k=np.zeros(9, dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), + ) + ) frame_id_prefix: str | None = None stereo_slice: Literal["left", "right"] | None = None # For stereo cameras @@ -198,7 +214,7 @@ def camera_info(self) -> CameraInfo: width, height, axis="horizontal", - header=Header(frame_id=self._frame("camera_optical")), + header=Header(frame_id=self._frame("camera_optical"), stamp=Time(sec=0, nanosec=0)), ) def emit(self, image: Image) -> None: ... diff --git a/dimos/hardware/sensors/camera/zed/camera.py b/dimos/hardware/sensors/camera/zed/camera.py index b240f7adcad..7f0bd2ee717 100644 --- a/dimos/hardware/sensors/camera/zed/camera.py +++ b/dimos/hardware/sensors/camera/zed/camera.py @@ -18,10 +18,12 @@ import threading import time +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 -from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2, RegionOfInterest from dimos_generated.std_msgs.msg import Header from dimos_generated.tf2_msgs.msg import TFMessage +import numpy as np from pydantic import Field import pyzed.sl as sl import reactivex as rx @@ -49,8 +51,8 @@ def default_base_transform() -> Transform: """Default identity transform for camera mounting.""" return Transform( - translation=Vector3(), - rotation=Quaternion(w=1.0), + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), ) @@ -236,11 +238,14 @@ def _intrinsics_to_camera_info( height=self._stream_height, width=self._stream_width, distortion_model="plumb_bob", - d=D, - k=K, - r=[1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0], - p=P, - header=Header(frame_id=frame_id), + d=np.asarray(D, dtype=np.float64), + k=np.asarray(K, dtype=np.float64), + r=np.array([1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0], dtype=np.float64), + p=np.asarray(P, dtype=np.float64), + header=Header(frame_id=frame_id, stamp=Time(sec=0, nanosec=0)), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), ) def _get_extrinsics(self) -> None: @@ -388,7 +393,9 @@ def _publish_tf(self, ts: float) -> None: TransformStamped( header=Header(frame_id=frame, stamp=stamp), child_frame_id=optical_frame, - transform=Transform(rotation=OPTICAL_ROTATION), + transform=Transform( + rotation=OPTICAL_ROTATION, translation=Vector3(x=0.0, y=0.0, z=0.0) + ), ) ) tracking_tf = self._tracking_transform(ts) diff --git a/dimos/hardware/sensors/camera/zed/test_generated_tf.py b/dimos/hardware/sensors/camera/zed/test_generated_tf.py index 84ea16dd3ae..4d628cdf85a 100644 --- a/dimos/hardware/sensors/camera/zed/test_generated_tf.py +++ b/dimos/hardware/sensors/camera/zed/test_generated_tf.py @@ -19,7 +19,8 @@ from types import ModuleType, SimpleNamespace from unittest.mock import Mock -from dimos_generated.geometry_msgs.msg import Transform, Vector3 +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, Vector3 from dimos_generated.tf2_msgs.msg import TFMessage import numpy as np import pytest @@ -56,7 +57,16 @@ def sdk_pose(translation): ) -@pytest.mark.parametrize("mount", [None, Transform(translation=Vector3(x=0.25))]) +@pytest.mark.parametrize( + "mount", + [ + None, + Transform( + translation=Vector3(x=0.25, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ], +) def test_tracking_uses_generated_frame_chain(camera_module, mount): camera = camera_module.ZEDCamera(base_transform=mount) camera._zed = Mock() @@ -84,7 +94,7 @@ def test_generated_tf_preserves_mount_extrinsics_and_optical_frames(camera_modul monkeypatch.setattr(camera.tf, "publish", received.append) try: camera._publish_tf(1700000000.125) - message = TFMessage.decode(received[0].encode()) + message = cdr_decode(cdr_encode(received[0]), TFMessage) assert len(message.transforms) == 5 frames = [(t.header.frame_id, t.child_frame_id) for t in message.transforms] assert frames == [ diff --git a/dimos/hardware/sensors/lidar/common/point_cloud_utils.hpp b/dimos/hardware/sensors/lidar/common/point_cloud_utils.hpp index f9b95a3e4d3..92dee481c7b 100644 --- a/dimos/hardware/sensors/lidar/common/point_cloud_utils.hpp +++ b/dimos/hardware/sensors/lidar/common/point_cloud_utils.hpp @@ -5,12 +5,16 @@ #pragma once +#include +#include +#include +#include #include #include #include #include #include -#include + namespace dimos { diff --git a/dimos/hardware/sensors/lidar/fastlio2/cpp/main.cpp b/dimos/hardware/sensors/lidar/fastlio2/cpp/main.cpp index b8c89ab074f..076549cbd0a 100644 --- a/dimos/hardware/sensors/lidar/fastlio2/cpp/main.cpp +++ b/dimos/hardware/sensors/lidar/fastlio2/cpp/main.cpp @@ -5,6 +5,8 @@ // FAST-LIO-NON-ROS and publishes sensor-frame point clouds on lidar plus // odometry with covariance on odometry. No inputs, so it overrides handle(). +#include +#include #include #include #include @@ -21,7 +23,7 @@ #include "point_cloud_utils.hpp" #include "dimos/native.hpp" -#include + // FAST-LIO (header-only core, compiled sources linked via CMake) diff --git a/dimos/hardware/sensors/lidar/fastlio2/module.py b/dimos/hardware/sensors/lidar/fastlio2/module.py index 811c19c8c19..793a9603f03 100644 --- a/dimos/hardware/sensors/lidar/fastlio2/module.py +++ b/dimos/hardware/sensors/lidar/fastlio2/module.py @@ -24,9 +24,11 @@ from __future__ import annotations -import time from typing import TYPE_CHECKING +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.tf2_msgs.msg import TFMessage from pydantic import Field from reactivex.disposable import Disposable @@ -46,12 +48,7 @@ SDK_POINT_DATA_PORT, SDK_PUSH_MSG_PORT, ) -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.geometry import transform_from_odometry from dimos.spec import perception @@ -139,17 +136,7 @@ def start(self) -> None: ) def _on_odom_for_tf(self, msg: Odometry) -> None: - self.tf.publish( - TFMessage( - Transform( - frame_id=self.frame_id, - child_frame_id=self.config.sensor_frame_id, - translation=Vector3(msg.pose.position), - rotation=Quaternion(msg.pose.orientation), - ts=msg.ts or time.time(), - ) - ) - ) + self.tf.publish(TFMessage(transforms=[transform_from_odometry(msg)])) @rpc def stop(self) -> None: diff --git a/dimos/hardware/sensors/lidar/fastlio2/recorder.py b/dimos/hardware/sensors/lidar/fastlio2/recorder.py index 175041bd688..de338a5ab49 100644 --- a/dimos/hardware/sensors/lidar/fastlio2/recorder.py +++ b/dimos/hardware/sensors/lidar/fastlio2/recorder.py @@ -24,12 +24,13 @@ from __future__ import annotations +from dimos_generated.geometry_msgs.msg import Pose +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import PointCloud2 + from dimos.core.stream import In from dimos.memory.module import Recorder, RecorderConfig, pose_setter_for from dimos.memory.type.recording import OnExisting -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 class FastLio2RecorderConfig(RecorderConfig): diff --git a/dimos/hardware/sensors/lidar/livox/rust/Cargo.toml b/dimos/hardware/sensors/lidar/livox/rust/Cargo.toml index 4f2a56f4cac..f5fed29b2ec 100644 --- a/dimos/hardware/sensors/lidar/livox/rust/Cargo.toml +++ b/dimos/hardware/sensors/lidar/livox/rust/Cargo.toml @@ -24,7 +24,7 @@ imu = "sensor_msgs/msg/Imu" crc = "3" dimos-module = { path = "../../../../../../native/rust/dimos-module" } etherparse = "0.21.0" -dimos-generated-messages = { path = "../../../../../message_codegen" } +dimos-generated-messages = { path = "../../../../../../packages/dimos-generated/build/message-source/rust" } pcap-parser = "0.17.0" serde = { version = "1", features = ["derive"] } serde_json = "1" diff --git a/dimos/hardware/sensors/lidar/livox/rust/src/module.rs b/dimos/hardware/sensors/lidar/livox/rust/src/module.rs index 097c84621c4..a4df07c7202 100644 --- a/dimos/hardware/sensors/lidar/livox/rust/src/module.rs +++ b/dimos/hardware/sensors/lidar/livox/rust/src/module.rs @@ -18,10 +18,12 @@ use crate::live::{LiveConfig, LiveSource, Ports}; use crate::pcap::PcapSource; use crate::pipeline::{Frame, ImuRecord, PacketSource, Pipeline}; use crate::wire::{DataPacket, DataType}; -use dimos_generated_messages::builtin_interfaces::msg::Time; -use dimos_generated_messages::geometry_msgs::msg::{Quaternion, Vector3}; -use dimos_generated_messages::sensor_msgs::msg::{Imu, PointCloud2, PointField}; -use dimos_generated_messages::std_msgs::msg::Header; +use dimos_generated_messages::builtin_interfaces::msg::time::Time; +use dimos_generated_messages::geometry_msgs::msg::{quaternion::Quaternion, vector3::Vector3}; +use dimos_generated_messages::sensor_msgs::msg::{ + imu::Imu, point_cloud2::PointCloud2, point_field::PointField, +}; +use dimos_generated_messages::std_msgs::msg::header::Header; use dimos_module::cdr; use dimos_module::{native_config, Module, Output}; use serde::{Deserialize, Serialize}; diff --git a/dimos/hardware/sensors/lidar/livox/test_e2e.py b/dimos/hardware/sensors/lidar/livox/test_e2e.py index d13177895b8..b5980135fc5 100644 --- a/dimos/hardware/sensors/lidar/livox/test_e2e.py +++ b/dimos/hardware/sensors/lidar/livox/test_e2e.py @@ -30,6 +30,8 @@ import subprocess import time +from dimos_cdrgen.registry import decode as cdr_decode +from dimos_generated.sensor_msgs.msg import Imu, PointCloud2 import lcm as lcm_module import pytest @@ -42,8 +44,8 @@ SDK_POINT_DATA_PORT, ) from dimos.hardware.sensors.lidar.virtual_mid360.module import VirtualMid360Config -from dimos.msgs.sensor_msgs.Imu import Imu -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_view, pointcloud_xyz +from dimos.msgs.time import to_seconds from dimos.utils.data import get_data _RELEASE = DIMOS_PROJECT_ROOT / "target" / "release" @@ -173,7 +175,7 @@ def on_msg(_ch: str, data: bytes, topic: str = topic) -> None: raw[topic].append(data) msg_type = topics[topic] - lc.subscribe(f"{topic}#sensor_msgs.{msg_type.__name__}", on_msg) + lc.subscribe(f"{topic}#{msg_type.__msgtype__}", on_msg) end = time.monotonic() + seconds while time.monotonic() < end: if all(len(raw[topic]) >= count for topic, count in enough.items()): @@ -196,12 +198,11 @@ def _first_capture_ts(pcap: Path) -> float: def _assert_cloud_shape(cloud: PointCloud2, min_points: int) -> None: - assert cloud.frame_id == "lidar_link" - points, _ = cloud.as_numpy() + assert cloud.header.frame_id == "lidar_link" + points = pointcloud_xyz(cloud) assert len(points) > min_points # Full point format carries per-point deskew offsets within one frame. - offsets = cloud.offset_times_u32() - assert offsets is not None + offsets = pointcloud_view(cloud)["offset_time"] assert offsets.max() < 150_000_000, "offsets exceed one frame interval" @@ -218,8 +219,8 @@ def test_real_capture_replay_publishes_streams(spawn: Spawn) -> None: ) blob = { "topics": { - "lidar": "/e2e_real_lidar#sensor_msgs.PointCloud2", - "imu": "/e2e_real_imu#sensor_msgs.Imu", + "lidar": "/e2e_real_lidar#sensor_msgs/msg/PointCloud2", + "imu": "/e2e_real_imu#sensor_msgs/msg/Imu", }, "config": config.to_config_dict(), "session": {"id": "pointlio-e2e", "links": []}, @@ -231,8 +232,8 @@ def test_real_capture_replay_publishes_streams(spawn: Spawn) -> None: enough={"/e2e_real_lidar": 20, "/e2e_real_imu": 800}, ) - clouds = [PointCloud2.lcm_decode(data) for data in received["/e2e_real_lidar"]] - imus = [Imu.lcm_decode(data) for data in received["/e2e_real_imu"]] + clouds = [cdr_decode(data, PointCloud2) for data in received["/e2e_real_lidar"]] + imus = [cdr_decode(data, Imu) for data in received["/e2e_real_imu"]] # Thresholds leave headroom for receive-buffer drops on the big clouds. assert len(clouds) >= 20, f"expected >=20 clouds, got {len(clouds)}" assert len(imus) >= 800, f"expected >=800 imu samples, got {len(imus)}" @@ -243,14 +244,16 @@ def test_real_capture_replay_publishes_streams(spawn: Spawn) -> None: median = magnitudes[len(magnitudes) // 2] assert 5.0 < median < 15.0, f"median accel magnitude {median}" assert imus[0].orientation_covariance[0] == -1.0 - assert imus[0].frame_id == "imu_link" + assert imus[0].header.frame_id == "imu_link" # Replayed stamps sit on the recording's clock, not the sensor's power-on # clock: within the collected window of the first capture record. The offset # filter strips delivery latency, so a stamp can sit a few ms before its record. capture_start = _first_capture_ts(pcap) for msg in (clouds[0], clouds[-1], imus[0], imus[-1]): - assert -0.1 <= msg.ts - capture_start < 30.0, f"stamp {msg.ts} vs capture {capture_start}" + assert -0.1 <= to_seconds(msg.header.stamp) - capture_start < 30.0, ( + f"stamp {to_seconds(msg.header.stamp)} vs capture {capture_start}" + ) @pytest.mark.native_e2e @@ -279,8 +282,8 @@ def test_live_loopback_handshake_and_stream(spawn: Spawn, synth_pcap: Path) -> N ) driver_blob = { "topics": { - "lidar": "/e2e_live_lidar#sensor_msgs.PointCloud2", - "imu": "/e2e_live_imu#sensor_msgs.Imu", + "lidar": "/e2e_live_lidar#sensor_msgs/msg/PointCloud2", + "imu": "/e2e_live_imu#sensor_msgs/msg/Imu", }, "config": config.to_config_dict(), "session": {"id": "pointlio-e2e", "links": []}, @@ -295,8 +298,8 @@ def test_live_loopback_handshake_and_stream(spawn: Spawn, synth_pcap: Path) -> N enough={"/e2e_live_lidar": 3, "/e2e_live_imu": 120}, ) - clouds = [PointCloud2.lcm_decode(data) for data in received["/e2e_live_lidar"]] - imus = [Imu.lcm_decode(data) for data in received["/e2e_live_imu"]] + clouds = [cdr_decode(data, PointCloud2) for data in received["/e2e_live_lidar"]] + imus = [cdr_decode(data, Imu) for data in received["/e2e_live_imu"]] # Large clouds drop under the default receive buffer, so the IMU packets # carry the rate assertion. assert len(clouds) >= 3, f"expected >=3 clouds, got {len(clouds)}" @@ -305,7 +308,7 @@ def test_live_loopback_handshake_and_stream(spawn: Spawn, synth_pcap: Path) -> N # The synthetic capture is a 5 m ring at z ~1 m, so a decode regression # in scaling or field layout shows up as broken geometry here. - points, _ = clouds[len(clouds) // 2].as_numpy() + points = pointcloud_xyz(clouds[len(clouds) // 2]) radii = sorted(math.hypot(x, y) for x, y, _z in points) assert abs(radii[len(radii) // 2] - 5.0) < 0.1, f"median ring radius {radii[len(radii) // 2]}" z_values = sorted(z for _x, _y, z in points) @@ -315,9 +318,11 @@ def test_live_loopback_handshake_and_stream(spawn: Spawn, synth_pcap: Path) -> N sample = imus[len(imus) // 2] assert abs(_magnitude(sample) - GRAVITY_MS2) < 1e-3 assert sample.orientation_covariance[0] == -1.0 - assert sample.frame_id == "imu_link" + assert sample.header.frame_id == "imu_link" # Live stamps are anchored to the host clock, whatever the device reports. now = time.time() for msg in (clouds[-1], imus[-1]): - assert abs(msg.ts - now) < 15.0, f"stamp {msg.ts} vs host {now}" + assert abs(to_seconds(msg.header.stamp) - now) < 15.0, ( + f"stamp {to_seconds(msg.header.stamp)} vs host {now}" + ) diff --git a/dimos/hardware/sensors/lidar/pointlio/module.py b/dimos/hardware/sensors/lidar/pointlio/module.py index f7c609438dc..fc4f2779da5 100644 --- a/dimos/hardware/sensors/lidar/pointlio/module.py +++ b/dimos/hardware/sensors/lidar/pointlio/module.py @@ -38,12 +38,12 @@ from dimos_generated.nav_msgs.msg import Odometry from dimos_generated.sensor_msgs.msg import Imu, PointCloud2 +from dimos_generated.tf2_msgs.msg import TFMessage from pydantic import Field from dimos.core.core import rpc from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.spec import perception # iVox local-map neighbor stencil. rust/src/module.rs maps the strings to diff --git a/dimos/hardware/sensors/lidar/pointlio/recorder.py b/dimos/hardware/sensors/lidar/pointlio/recorder.py index 8013ec8577a..b0ea2b52d93 100644 --- a/dimos/hardware/sensors/lidar/pointlio/recorder.py +++ b/dimos/hardware/sensors/lidar/pointlio/recorder.py @@ -24,12 +24,13 @@ from __future__ import annotations +from dimos_generated.geometry_msgs.msg import Pose +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import PointCloud2 + from dimos.core.stream import In from dimos.memory.module import Recorder, RecorderConfig, pose_setter_for from dimos.memory.type.recording import OnExisting -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 class PointlioRecorderConfig(RecorderConfig): diff --git a/dimos/hardware/sensors/lidar/pointlio/rust/Cargo.toml b/dimos/hardware/sensors/lidar/pointlio/rust/Cargo.toml index 1982e2093cd..35958a55570 100644 --- a/dimos/hardware/sensors/lidar/pointlio/rust/Cargo.toml +++ b/dimos/hardware/sensors/lidar/pointlio/rust/Cargo.toml @@ -21,19 +21,19 @@ path = "dimos_pointlio::module::PointLioModule" python = "dimos.hardware.sensors.lidar.pointlio.module:PointLio" [package.metadata.dimos.module.pointlio.inputs] -lidar_raw = "sensor_msgs.PointCloud2" -imu_raw = "sensor_msgs.Imu" -tf = "tf2_msgs.TFMessage" +lidar_raw = "sensor_msgs/msg/PointCloud2" +imu_raw = "sensor_msgs/msg/Imu" +tf = "tf2_msgs/msg/TFMessage" [package.metadata.dimos.module.pointlio.outputs] -lidar = "sensor_msgs.PointCloud2" -odometry = "nav_msgs.Odometry" +lidar = "sensor_msgs/msg/PointCloud2" +odometry = "nav_msgs/msg/Odometry" [dependencies] dimos-livox = { path = "../../livox/rust" } pointlio-core = { git = "https://github.com/dimensionalOS/dimos-module-pointlio.git", rev = "22a5403f827d030b248ae473c2808b5d405b3d55" } dimos-module = { path = "../../../../../../native/rust/dimos-module" } -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +dimos-generated-messages = { path = "../../../../../../packages/dimos-generated/build/message-source/rust" } serde = { version = "1", features = ["derive"] } serde_json = "1" tokio = { version = "1", features = ["rt-multi-thread", "macros"] } diff --git a/dimos/hardware/sensors/lidar/pointlio/rust/src/module.rs b/dimos/hardware/sensors/lidar/pointlio/rust/src/module.rs index 41157a2d2ef..d66f099613b 100644 --- a/dimos/hardware/sensors/lidar/pointlio/rust/src/module.rs +++ b/dimos/hardware/sensors/lidar/pointlio/rust/src/module.rs @@ -19,11 +19,20 @@ use dimos_module::{native_config, Input, Module, Output, Tf, Transform}; /// The driver publishes linear acceleration in m/s^2; Point-LIO wants it in g. const GRAVITY_MS2: f64 = 9.80665; -use lcm_msgs::geometry_msgs::{Point, Pose, PoseWithCovariance, Twist, TwistWithCovariance}; -use lcm_msgs::geometry_msgs::{Quaternion as QuatMsg, Vector3}; -use lcm_msgs::nav_msgs::Odometry; -use lcm_msgs::sensor_msgs::{Imu, PointCloud2, PointField}; -use lcm_msgs::std_msgs::{Header, Time}; +use dimos_generated_messages::builtin_interfaces::msg::time::Time; +use dimos_generated_messages::geometry_msgs::msg::{ + point::Point, pose::Pose, pose_with_covariance::PoseWithCovariance, twist::Twist, + twist_with_covariance::TwistWithCovariance, +}; +use dimos_generated_messages::geometry_msgs::msg::{ + quaternion::Quaternion as QuatMsg, vector3::Vector3, +}; +use dimos_generated_messages::nav_msgs::msg::odometry::Odometry; +use dimos_generated_messages::sensor_msgs::msg::{ + imu::Imu, point_cloud2::PointCloud2, point_field::PointField, +}; +use dimos_generated_messages::std_msgs::msg::header::Header; +use dimos_module::cdr; use pointlio_core::{LivoxPoint, PointLio, PointXYZI}; use serde::{Deserialize, Serialize}; use std::time::Duration; @@ -117,16 +126,16 @@ fn core_config_parses(cfg: &Config) -> Result<(), ValidationError> { #[derive(Module)] #[module(name = "pointlio", setup = start)] pub struct PointLioModule { - #[input(decode = PointCloud2::decode)] + #[input(decode = cdr::decode)] lidar_raw: Input, - #[input(decode = Imu::decode)] + #[input(decode = cdr::decode)] imu_raw: Input, - #[output(encode = PointCloud2::encode)] + #[output(encode = cdr::encode)] lidar: Output, - #[output(encode = Odometry::encode)] + #[output(encode = cdr::encode)] odometry: Output, #[tf] @@ -239,20 +248,19 @@ fn due(last: &mut Option, ts: f64, hz: f64) -> bool { } fn stamp_secs(t: &Time) -> f64 { - f64::from(t.sec) + f64::from(t.nsec) * 1e-9 + f64::from(t.sec) + f64::from(t.nanosec) * 1e-9 } fn stamp_ns(t: &Time) -> u64 { - t.sec as u64 * 1_000_000_000 + t.nsec as u64 + t.sec as u64 * 1_000_000_000 + t.nanosec as u64 } fn header(frame_id: &str, ts: f64) -> Header { let sec = ts.floor(); Header { - seq: 0, stamp: Time { sec: sec as i32, - nsec: (((ts - sec) * 1e9).round() as i32).min(999_999_999), + nanosec: (((ts - sec) * 1e9).round() as u32).min(999_999_999), }, frame_id: frame_id.to_string(), } @@ -270,11 +278,15 @@ struct Layout { fn layout(cloud: &PointCloud2) -> Result { let step = cloud.point_step as usize; - let find = |name: &str, datatype: i8| -> Result, String> { - let Some(f) = cloud.fields.iter().find(|f| f.name == name) else { + let find = |name: &str, datatype: u8| -> Result, String> { + let mut matches = cloud.fields.iter().filter(|field| field.name == name); + let Some(f) = matches.next() else { return Ok(None); }; - if f.datatype != datatype as u8 { + if matches.next().is_some() || f.count != 1 { + return Err(format!("field {name} must be unique and scalar")); + } + if f.datatype != datatype { return Err(format!( "field {name} has datatype {}, want {datatype}", f.datatype @@ -291,12 +303,9 @@ fn layout(cloud: &PointCloud2) -> Result { } Ok(Some(offset)) }; - let need = |name: &str, datatype: i8| { + let need = |name: &str, datatype: u8| { find(name, datatype)?.ok_or_else(|| format!("cloud has no {name} field")) }; - if cloud.is_bigendian { - return Err("big-endian clouds are not supported".into()); - } Ok(Layout { x: need("x", PointField::FLOAT32)?, y: need("y", PointField::FLOAT32)?, @@ -311,19 +320,40 @@ fn layout(cloud: &PointCloud2) -> Result { fn livox_points(cloud: &PointCloud2) -> Result, String> { let l = layout(cloud)?; let step = cloud.point_step as usize; - let n = cloud.width as usize * cloud.height as usize; - if step == 0 || cloud.data.len() < n * step { - return Err(format!( - "data has {} bytes for {n} points of {step}", - cloud.data.len() - )); + let width = cloud.width as usize; + let row_step = cloud.row_step as usize; + let row_bytes = width + .checked_mul(step) + .ok_or("point cloud row size overflow")?; + let total = row_step + .checked_mul(cloud.height as usize) + .ok_or("point cloud data size overflow")?; + if step == 0 || row_step < row_bytes || cloud.data.len() != total { + return Err("point cloud dimensions/strides do not match data".into()); } - let f32_at = |p: &[u8], off: usize| f32::from_le_bytes(p[off..off + 4].try_into().unwrap()); - let u32_at = |p: &[u8], off: usize| u32::from_le_bytes(p[off..off + 4].try_into().unwrap()); + if width == 0 || cloud.height == 0 { + return Ok(Vec::new()); + } + let f32_at = |p: &[u8], off: usize| { + let bytes = p[off..off + 4].try_into().unwrap(); + if cloud.is_bigendian { + f32::from_be_bytes(bytes) + } else { + f32::from_le_bytes(bytes) + } + }; + let u32_at = |p: &[u8], off: usize| { + let bytes = p[off..off + 4].try_into().unwrap(); + if cloud.is_bigendian { + u32::from_be_bytes(bytes) + } else { + u32::from_le_bytes(bytes) + } + }; Ok(cloud .data - .chunks_exact(step) - .take(n) + .chunks_exact(row_step) + .flat_map(|row| row[..row_bytes].chunks_exact(step)) .map(|p| LivoxPoint { x: f32_at(p, l.x), y: f32_at(p, l.y), @@ -403,20 +433,20 @@ fn cloud_message(frame_id: &str, ts: f64, cloud: &[PointXYZI]) -> PointCloud2 { .enumerate() .map(|(i, name)| PointField { name: (*name).into(), - offset: (i * 4) as i32, - datatype: PointField::FLOAT32 as u8, + offset: (i * 4) as u32, + datatype: PointField::FLOAT32, count: 1, }) .collect(); - let n = cloud.len() as i32; + let n = cloud.len() as u32; PointCloud2 { header: header(frame_id, ts), height: 1, width: n, fields, is_bigendian: false, - point_step: STEP as i32, - row_step: STEP as i32 * n, + point_step: STEP as u32, + row_step: STEP as u32 * n, data, is_dense: true, } @@ -426,11 +456,11 @@ fn cloud_message(frame_id: &str, ts: f64, cloud: &[PointXYZI]) -> PointCloud2 { mod tests { use super::*; - fn field(name: &str, offset: usize, datatype: i8) -> PointField { + fn field(name: &str, offset: usize, datatype: u8) -> PointField { PointField { name: name.into(), - offset: offset as i32, - datatype: datatype as u8, + offset: offset as u32, + datatype, count: 1, } } @@ -439,16 +469,60 @@ mod tests { PointCloud2 { header: header("lidar_link", 1.5), height: 1, - width: (data.len() / step) as i32, + width: (data.len() / step) as u32, fields, is_bigendian: false, - point_step: step as i32, - row_step: data.len() as i32, + point_step: step as u32, + row_step: data.len() as u32, data, is_dense: true, } } + #[test] + fn padded_rows_and_big_endian_preserve_deskew_offsets() { + for big in [false, true] { + let mut data = vec![0; 48]; + for row in 0..2 { + for (axis, value) in [1f32 + row as f32, 2., 3.].into_iter().enumerate() { + let bytes = if big { + value.to_be_bytes() + } else { + value.to_le_bytes() + }; + data[row * 24 + axis * 4..row * 24 + axis * 4 + 4].copy_from_slice(&bytes); + } + let offset = 42u32 + row as u32; + let bytes = if big { + offset.to_be_bytes() + } else { + offset.to_le_bytes() + }; + data[row * 24 + 12..row * 24 + 16].copy_from_slice(&bytes); + } + let mut message = cloud( + vec![ + field("x", 0, PointField::FLOAT32), + field("y", 4, PointField::FLOAT32), + field("z", 8, PointField::FLOAT32), + field("offset_time", 12, PointField::UINT32), + ], + 16, + data, + ); + message.width = 1; + message.height = 2; + message.row_step = 24; + message.is_bigendian = big; + let points = livox_points(&message).unwrap(); + assert_eq!(points.len(), 2); + assert_eq!((points[0].x, points[1].x), (1., 2.)); + assert_eq!((points[0].offset_time, points[1].offset_time), (42, 43)); + message.row_step = 15; + assert!(livox_points(&message).is_err()); + } + } + #[test] fn full_layout_is_read_by_name() { // Driver `full`: x y z intensity offset_time tag line, 22 B. diff --git a/dimos/hardware/sensors/lidar/test_generated_recorders.py b/dimos/hardware/sensors/lidar/test_generated_recorders.py index 09137c3ff48..0fe4a41ce41 100644 --- a/dimos/hardware/sensors/lidar/test_generated_recorders.py +++ b/dimos/hardware/sensors/lidar/test_generated_recorders.py @@ -15,7 +15,17 @@ # Copyright 2026 Dimensional Inc. # SPDX-License-Identifier: Apache-2.0 -from dimos_generated.geometry_msgs.msg import Point, Pose, PoseWithCovariance, Quaternion +from dimos_cdrgen.registry import encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseWithCovariance, + Quaternion, + Twist, + TwistWithCovariance, + Vector3, +) from dimos_generated.nav_msgs.msg import Odometry from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 from dimos_generated.std_msgs.msg import Header @@ -37,10 +47,25 @@ async def test_generated_recorders_anchor_clouds_to_last_odometry(recorder_type, tmp_path): recorder = recorder_type(db_path=tmp_path / "unused.db") try: - cloud = pointcloud_from_xyz(np.array([[1, 2, 3]], dtype=np.float32), header=Header()) + cloud = pointcloud_from_xyz( + np.array([[1, 2, 3]], dtype=np.float32), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) assert await recorder._lidar_pose(cloud) is None - pose = Pose(position=Point(x=2, y=3), orientation=Quaternion(w=1)) - odometry = Odometry(header=Header(frame_id="world"), pose=PoseWithCovariance(pose=pose)) + pose = Pose( + position=Point(x=2, y=3, z=0.0), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ) + odometry = Odometry( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=PoseWithCovariance(pose=pose, covariance=np.zeros(36, dtype=np.float64)), + child_frame_id="", + twist=TwistWithCovariance( + twist=Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ), + covariance=np.zeros(36, dtype=np.float64), + ), + ) assert await recorder._odom_pose(odometry) == pose assert await recorder._lidar_pose(cloud) == pose assert not (tmp_path / "unused.db").exists() @@ -78,11 +103,14 @@ def test_realsense_mount_is_an_identity_and_generated_cloud_views_are_typed(): np.testing.assert_allclose(transform_matrix(edge.transform), np.eye(4)) finally: mount.stop() - cloud = pointcloud_from_xyz(np.array([[1, 2, 3]], dtype=np.float32), header=Header()) - before = cloud.encode() + cloud = pointcloud_from_xyz( + np.array([[1, 2, 3]], dtype=np.float32), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + before = cdr_encode(cloud) assert _cloud(cloud).positions is not None assert _fine_points(cloud).positions is not None - assert cloud.encode() == before + assert cdr_encode(cloud) == before def test_g1_depth_recorder_declares_generated_lossless_cdr(): diff --git a/dimos/hardware/whole_body/transport/adapter.py b/dimos/hardware/whole_body/transport/adapter.py index 2b496609a1f..ff85afcf920 100644 --- a/dimos/hardware/whole_body/transport/adapter.py +++ b/dimos/hardware/whole_body/transport/adapter.py @@ -26,6 +26,7 @@ from dimos_generated.dimos_msgs.msg import MotorCommandArray from dimos_generated.sensor_msgs.msg import Imu, JointState +import numpy as np from dimos.core.transport_factory import make_transport from dimos.hardware.spec import JointLimits @@ -137,11 +138,11 @@ def write_motor_commands(self, commands: list[MotorCommand]) -> bool: msg = MotorCommandArray( header=header_now(), - q=[c.q for c in commands], - dq=[c.dq for c in commands], - kp=[c.kp for c in commands], - kd=[c.kd for c in commands], - tau=[c.tau for c in commands], + q=np.asarray([c.q for c in commands], dtype=np.float64), + dq=np.asarray([c.dq for c in commands], dtype=np.float64), + kp=np.asarray([c.kp for c in commands], dtype=np.float64), + kd=np.asarray([c.kd for c in commands], dtype=np.float64), + tau=np.asarray([c.tau for c in commands], dtype=np.float64), ) self._motor_command_transport.publish(msg) return True diff --git a/dimos/hardware/whole_body/transport/test_adapter.py b/dimos/hardware/whole_body/transport/test_adapter.py index b779185731b..2aa386ffae9 100644 --- a/dimos/hardware/whole_body/transport/test_adapter.py +++ b/dimos/hardware/whole_body/transport/test_adapter.py @@ -12,8 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Quaternion, Vector3 from dimos_generated.sensor_msgs.msg import Imu, JointState +from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from dimos.hardware.whole_body.spec import IMUState, MotorCommand, MotorState @@ -30,7 +33,13 @@ def test_unconnected_adapter_does_not_send_commands(): @pytest.mark.parametrize("field", ["position", "velocity", "effort"]) def test_short_feedback_preserves_last_complete_snapshot(field): adapter = TransportWholeBodyAdapter(dof=2) - message = JointState(position=[1.0, 2.0], velocity=[3.0, 4.0], effort=[5.0, 6.0]) + message = JointState( + position=np.array([1.0, 2.0], dtype=np.float64), + velocity=np.array([3.0, 4.0], dtype=np.float64), + effort=np.array([5.0, 6.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + ) adapter._on_motor_states(message) setattr(message, field, [0.0]) adapter._on_motor_states(message) @@ -47,6 +56,10 @@ def test_generated_imu_converts_xyzw_to_hardware_wxyz(): orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.4), angular_velocity=Vector3(x=1.0, y=2.0, z=3.0), linear_acceleration=Vector3(x=4.0, y=5.0, z=6.0), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + orientation_covariance=np.zeros(9, dtype=np.float64), + angular_velocity_covariance=np.zeros(9, dtype=np.float64), + linear_acceleration_covariance=np.zeros(9, dtype=np.float64), ) ) assert adapter.read_imu() == IMUState( diff --git a/dimos/hosted/fragment.py b/dimos/hosted/fragment.py index 4eeddb03fda..73d33299da7 100644 --- a/dimos/hosted/fragment.py +++ b/dimos/hosted/fragment.py @@ -37,7 +37,7 @@ def message_type_name(message_type: type) -> str: return str( getattr( message_type, - "msg_name", + "__msgtype__", f"{message_type.__module__}.{message_type.__qualname__}", ) ) @@ -171,7 +171,7 @@ def validate_boundary_streams(run_id: str, payload: PythonFragmentPayload) -> No if not isinstance(transport, TransportSpec): raise ValueError(f"Boundary stream {boundary.name!r} has no pinned transport") - if getattr(boundary.message_type, "lcm_encode", None) is None: + if not isinstance(getattr(boundary.message_type, "__msgtype__", None), str): valid = transport.cls is pZenohTransport and transport.args == (expected_key,) else: expected_base = run_stream_base_topic(run_id, boundary.name) diff --git a/dimos/hosted/fragment_compiler.py b/dimos/hosted/fragment_compiler.py index d7530ad6ecd..fc4156f51bc 100644 --- a/dimos/hosted/fragment_compiler.py +++ b/dimos/hosted/fragment_compiler.py @@ -407,7 +407,7 @@ def _boundary_stream(run_id: str, stream_key: StreamKey) -> BoundaryStream: def _boundary_transport(run_id: str, boundary: BoundaryStream) -> TransportSpec: - if getattr(boundary.message_type, "lcm_encode", None) is None: + if not isinstance(getattr(boundary.message_type, "__msgtype__", None), str): return pZenohTransport.spec(boundary.key_expr) return ZenohTransport.spec( run_stream_base_topic(run_id, boundary.name), diff --git a/dimos/hosted/test_fragment_compiler.py b/dimos/hosted/test_fragment_compiler.py index 46f82075112..9b265c57c6f 100644 --- a/dimos/hosted/test_fragment_compiler.py +++ b/dimos/hosted/test_fragment_compiler.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.std_msgs.msg import Bool, String import pytest from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser @@ -27,8 +28,6 @@ run_stream_key, ) from dimos.hosted.fragment_compiler import compile_fragments, resolve_hosted_assignments -from dimos.msgs.std_msgs.Bool import Bool -from dimos.msgs.std_msgs.String import String class SourceModule(Module): @@ -99,7 +98,8 @@ def test_compiler_splits_modules_and_pins_boundary_stream_to_existing_topic_rule source_payload = fragments["host-a"].load_payload() sink_payload = fragments["host-b"].load_payload() - expected_key = run_stream_key("run-1", "messages", String) + expected_key = "dimos/runs/run-1/streams/messages/std_msgs/msg/String" + assert run_stream_key("run-1", "messages", String) == expected_key assert {atom.name for atom in source_payload.blueprint.active_blueprints} == {"sourcemodule"} assert {atom.name for atom in sink_payload.blueprint.active_blueprints} == {"sinkmodule"} diff --git a/dimos/imitation/collection/episode_monitor.py b/dimos/imitation/collection/episode_monitor.py index 36d43dda068..a9a47b9a1ef 100644 --- a/dimos/imitation/collection/episode_monitor.py +++ b/dimos/imitation/collection/episode_monitor.py @@ -26,6 +26,7 @@ import time from typing import Any, Literal, Protocol, TypeAlias +from dimos_generated.std_msgs.msg import String, UInt32 from pydantic import Field, field_validator from reactivex.abc import DisposableBase from reactivex.disposable import Disposable @@ -38,7 +39,6 @@ EpisodeStatus, RecordingState, ) -from dimos.msgs.std_msgs.String import String from dimos.spec.utils import Spec from dimos.teleop.webxr.controller_types import BUTTON_ALIASES, Buttons from dimos.utils.logging_config import setup_logger @@ -92,7 +92,7 @@ def _validate_button_map(cls, value: dict[EpisodeCommand, str]) -> dict[EpisodeC class EpisodeMonitorModule(Module): config: EpisodeMonitorModuleConfig - teleop_buttons: In[Buttons] + teleop_buttons: In[UInt32] status: Out[String] def __init__(self, **kwargs: Any) -> None: @@ -134,7 +134,7 @@ def stop(self) -> None: # ── port handlers ──────────────────────────────────────────────────────── - def _on_buttons(self, msg: Buttons) -> None: + def _on_buttons(self, msg: UInt32) -> None: """Detect raw rising edges, preserving short taps without held repeats.""" ts = time.time() fired: list[EpisodeCommand] = [] diff --git a/dimos/imitation/collection/recorder.py b/dimos/imitation/collection/recorder.py index 5fd527a165a..5f843d3565d 100644 --- a/dimos/imitation/collection/recorder.py +++ b/dimos/imitation/collection/recorder.py @@ -27,13 +27,13 @@ from typing import Any +from dimos_generated.geometry_msgs.msg import Pose +from dimos_generated.sensor_msgs.msg import Image, JointState +from dimos_generated.std_msgs.msg import String + from dimos.core.stream import In from dimos.imitation.collection.episode import EpisodeStatus from dimos.memory.module import Recorder, RecorderConfig -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.std_msgs.String import String class CollectionRecorderConfig(RecorderConfig): diff --git a/dimos/imitation/collection/test_blueprint.py b/dimos/imitation/collection/test_blueprint.py index cde7556849d..f633290dd42 100644 --- a/dimos/imitation/collection/test_blueprint.py +++ b/dimos/imitation/collection/test_blueprint.py @@ -14,6 +14,7 @@ from __future__ import annotations +from dimos_generated.sensor_msgs.msg import JointState import pytest from dimos.core.coordination.blueprints import Blueprint @@ -26,7 +27,6 @@ EpisodeMonitorModuleConfig, ) from dimos.imitation.collection.recorder import CollectionRecorder -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.robot.unitree.g1.blueprints.basic.unitree_g1_teleop import unitree_g1_teleop from dimos.teleop.webxr.extensions import ArmTeleopModule diff --git a/dimos/imitation/collection/test_episode_monitor.py b/dimos/imitation/collection/test_episode_monitor.py index 8feb91ba9c3..ab9c89ead45 100644 --- a/dimos/imitation/collection/test_episode_monitor.py +++ b/dimos/imitation/collection/test_episode_monitor.py @@ -30,9 +30,9 @@ import pytest import pytest_mock +from dimos.core.transport_factory import rpc_backend from dimos.imitation.collection.episode import EpisodeStatus from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule -from dimos.protocol.rpc.pubsubrpc import LCMRPC from dimos.teleop.webxr.controller_types import BUTTON_ALIASES, Buttons from dimos.teleop.webxr.module import WebXRTeleopModule @@ -49,10 +49,9 @@ def make_monitor( stopped on teardown. """ mocker.patch("dimos.core.module.get_loop", return_value=(mocker.MagicMock(), None)) - mocker.patch.object(LCMRPC, "__init__", return_value=None) - mocker.patch.object(LCMRPC, "serve_module_rpc", return_value=None) - mocker.patch.object(LCMRPC, "start", return_value=None) - mocker.patch.object(LCMRPC, "stop", return_value=None) + backend = rpc_backend() + for method in ("start", "serve_module_rpc", "stop"): + mocker.patch.object(backend, method) built: list[EpisodeMonitorModule] = [] @@ -80,8 +79,8 @@ def _press(monitor: EpisodeMonitorModule, alias: str) -> None: attr = BUTTON_ALIASES[alias] pressed = Buttons() pressed.set_attribute(attr, True) - monitor._on_buttons(pressed) - monitor._on_buttons(Buttons()) + monitor._on_buttons(pressed.to_message()) + monitor._on_buttons(Buttons().to_message()) @pytest.mark.parametrize("duration", [0.03, 0.06]) diff --git a/dimos/imitation/dataprep/core.py b/dimos/imitation/dataprep/core.py index 5e882b63297..d273cee10cb 100644 --- a/dimos/imitation/dataprep/core.py +++ b/dimos/imitation/dataprep/core.py @@ -33,6 +33,8 @@ import math from typing import TYPE_CHECKING, Any +from dimos_generated.sensor_msgs.msg import CompressedImage, Image, JointState +from dimos_generated.std_msgs.msg import Header, String import numpy as np from numpy.typing import NDArray @@ -58,8 +60,7 @@ ) from dimos.memory.store.mcap import McapStore from dimos.msgs.image import image_from_compressed, image_view -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.std_msgs.String import String +from dimos.msgs.time import time_from_seconds if TYPE_CHECKING: from dimos.memory.store.base import Store @@ -85,7 +86,8 @@ def resolve_field(msg: Any, ref: FeatureSpec) -> NDArray[Any]: Single source of truth for obs/action construction across train and live inference. Behavior: - `ref.field is None`: best-effort coerce the whole message - (Image → `.data`, ndarray pass-through, list/tuple → asarray). + (generated Image → pixel view, CompressedImage → decoded pixels, + ndarray pass-through, list/tuple → asarray). - `ref.field` set: `getattr(msg, ref.field)` (or `msg[ref.field]` for dict payloads) then coerce. """ @@ -103,10 +105,14 @@ def resolve_field(msg: Any, ref: FeatureSpec) -> NDArray[Any]: else: value = getattr(msg, ref.field) + if isinstance(value, Image): + return image_view(value) + if isinstance(value, CompressedImage): + return image_view(image_from_compressed(value)) if isinstance(value, np.ndarray): return value if hasattr(value, "data") and isinstance(value.data, np.ndarray): - # e.g. Image → use its underlying ndarray + # Application payloads may expose an underlying ndarray. return value.data return np.asarray(value) @@ -316,6 +322,16 @@ def _source_messages( yield observation.ts, observation.data +def _joint_snapshot(timestamp: float, targets: dict[str, float]) -> JointState: + return JointState( + header=Header(stamp=time_from_seconds(timestamp), frame_id=""), + name=list(targets), + position=np.asarray(list(targets.values()), dtype=np.float64), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + + def _joint_position_updates( messages: Iterator[tuple[float, Any]], start_ts: float, stream: str ) -> Iterator[tuple[float, JointState]]: @@ -326,7 +342,7 @@ def _joint_position_updates( if not seeded and timestamp >= start_ts: yield ( start_ts, - JointState(ts=start_ts, name=list(targets), position=list(targets.values())), + _joint_snapshot(start_ts, targets), ) seeded = True try: @@ -341,10 +357,10 @@ def _joint_position_updates( if timestamp >= start_ts: yield ( timestamp, - JointState(ts=timestamp, name=list(targets), position=list(targets.values())), + _joint_snapshot(timestamp, targets), ) if not seeded: - yield start_ts, JointState(ts=start_ts, name=list(targets), position=list(targets.values())) + yield start_ts, _joint_snapshot(start_ts, targets) def _episode_features( diff --git a/dimos/imitation/dataprep/test_core.py b/dimos/imitation/dataprep/test_core.py index 8b87031a1de..571ade2551d 100644 --- a/dimos/imitation/dataprep/test_core.py +++ b/dimos/imitation/dataprep/test_core.py @@ -28,6 +28,8 @@ import sys from typing import Any +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import String import numpy as np import pytest @@ -54,8 +56,7 @@ SyncConfig, ) from dimos.memory.store.sqlite import SqliteStore -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.std_msgs.String import String +from dimos.msgs.time import header_now @pytest.mark.parametrize( @@ -170,7 +171,13 @@ class Msg: def test_resolve_field_orders_joint_state_by_feature_names() -> None: arr = resolve_field( - JointState(name=["joint_b", "joint_a"], position=[2.0, 1.0]), + JointState( + header=header_now(), + velocity=[], + effort=[], + name=["joint_b", "joint_a"], + position=[2.0, 1.0], + ), FeatureSpec( stream="state", field="position", @@ -189,6 +196,8 @@ def test_resolve_field_orders_joint_state_by_feature_names() -> None: ) def test_resolve_field_orders_all_joint_state_vectors(field: str, expected: list[float]) -> None: state = JointState( + header=header_now(), + position=[], name=["joint_b", "joint_a"], velocity=[2.0, 1.0], effort=[4.0, 3.0], @@ -207,7 +216,13 @@ def test_resolve_field_orders_all_joint_state_vectors(field: str, expected: list def test_resolve_field_rejects_duplicate_joint_names() -> None: with pytest.raises(ValueError, match="duplicate joint names"): resolve_field( - JointState(name=["joint_a", "joint_a"], position=[1.0, 2.0]), + JointState( + header=header_now(), + velocity=[], + effort=[], + name=["joint_a", "joint_a"], + position=[1.0, 2.0], + ), FeatureSpec( stream="state", field="position", @@ -221,7 +236,9 @@ def test_resolve_field_rejects_duplicate_joint_names() -> None: def test_resolve_field_rejects_missing_configured_joint() -> None: with pytest.raises(ValueError, match="missing configured joints.*joint_b"): resolve_field( - JointState(name=["joint_a"], position=[1.0]), + JointState( + header=header_now(), velocity=[], effort=[], name=["joint_a"], position=[1.0] + ), FeatureSpec( stream="state", field="position", @@ -543,8 +560,22 @@ def test_snapshots_can_align_forward_but_targets_remain_causal(): "anchor": _scalar_stream([(10, 0), (11, 0)]), "measured": _scalar_stream([(10.01, 3), (10.99, 4)]), "commands": [ - _Obs(1, JointState(name=["left", "right", "gripper"], position=[1, 2, 0])), - _Obs(10.001, JointState(name=["left"], position=[5])), + _Obs( + 1, + JointState( + header=header_now(), + velocity=[], + effort=[], + name=["left", "right", "gripper"], + position=[1, 2, 0], + ), + ), + _Obs( + 10.001, + JointState( + header=header_now(), velocity=[], effort=[], name=["left"], position=[5] + ), + ), ], } ) @@ -584,9 +615,36 @@ def test_joint_state_message_class_does_not_determine_source_semantics(source_ki { "anchor": _scalar_stream([(0.0, 0), (0.1, 0)]), "joints": [ - _Obs(-0.1, JointState(name=["left", "right"], position=[1, 2])), - _Obs(0.01, JointState(name=["left", "right"], position=[3, 4])), - _Obs(0.1, JointState(name=["left", "right"], position=[5, 6])), + _Obs( + -0.1, + JointState( + header=header_now(), + velocity=[], + effort=[], + name=["left", "right"], + position=[1, 2], + ), + ), + _Obs( + 0.01, + JointState( + header=header_now(), + velocity=[], + effort=[], + name=["left", "right"], + position=[3, 4], + ), + ), + _Obs( + 0.1, + JointState( + header=header_now(), + velocity=[], + effort=[], + name=["left", "right"], + position=[5, 6], + ), + ), ], } ) @@ -620,7 +678,16 @@ def test_shared_update_source_is_read_once_for_multiple_projections(mocker): { "anchor": _scalar_stream([(10, 0)]), "commands": [ - _Obs(1, JointState(name=["left", "right", "gripper"], position=[1, 2, 0])) + _Obs( + 1, + JointState( + header=header_now(), + velocity=[], + effort=[], + name=["left", "right", "gripper"], + position=[1, 2, 0], + ), + ) ], } ) @@ -651,11 +718,52 @@ def test_held_action_uses_causal_recording_history_across_episodes(): { "anchor": _scalar_stream([(10.0, 0.0), (11.0, 0.0), (12.0, 0.0), (20.0, 0.0)]), "commands": [ - _Obs(1.0, JointState(name=["right", "left"], position=[2.0, 1.0])), - _Obs(2.0, JointState(name=["gripper"], position=[0.0])), - _Obs(10.0, JointState(name=["gripper"], position=[1.0])), - _Obs(10.01, JointState(name=["left"], position=[3.0])), - _Obs(12.0, JointState(name=["gripper", "right", "left"], position=[0.0, 5.0, 4.0])), + _Obs( + 1.0, + JointState( + header=header_now(), + velocity=[], + effort=[], + name=["right", "left"], + position=[2.0, 1.0], + ), + ), + _Obs( + 2.0, + JointState( + header=header_now(), + velocity=[], + effort=[], + name=["gripper"], + position=[0.0], + ), + ), + _Obs( + 10.0, + JointState( + header=header_now(), + velocity=[], + effort=[], + name=["gripper"], + position=[1.0], + ), + ), + _Obs( + 10.01, + JointState( + header=header_now(), velocity=[], effort=[], name=["left"], position=[3.0] + ), + ), + _Obs( + 12.0, + JointState( + header=header_now(), + velocity=[], + effort=[], + name=["gripper", "right", "left"], + position=[0.0, 5.0, 4.0], + ), + ), ], } ) @@ -689,8 +797,26 @@ def test_held_action_does_not_initialize_from_future_commands(): { "anchor": _scalar_stream([(10.0, 0.0), (11.0, 0.0)]), "commands": [ - _Obs(1.0, JointState(name=["left", "right"], position=[1.0, 2.0])), - _Obs(10.001, JointState(name=["gripper"], position=[1.0])), + _Obs( + 1.0, + JointState( + header=header_now(), + velocity=[], + effort=[], + name=["left", "right"], + position=[1.0, 2.0], + ), + ), + _Obs( + 10.001, + JointState( + header=header_now(), + velocity=[], + effort=[], + name=["gripper"], + position=[1.0], + ), + ), ], } ) @@ -709,9 +835,22 @@ def test_held_action_does_not_initialize_from_future_commands(): @pytest.mark.parametrize( "message, reason", [ - (JointState(name=["left", "left"], position=[1, 2]), "duplicate"), - (JointState(name=["left"], position=[]), "names but"), - (JointState(name=["left"], position=[float("nan")]), "non-finite"), + ( + JointState( + header=header_now(), velocity=[], effort=[], name=["left", "left"], position=[1, 2] + ), + "duplicate", + ), + ( + JointState(header=header_now(), velocity=[], effort=[], name=["left"], position=[]), + "names but", + ), + ( + JointState( + header=header_now(), velocity=[], effort=[], name=["left"], position=[float("nan")] + ), + "non-finite", + ), ({"position": [1]}, "JointState"), ], ) @@ -721,7 +860,16 @@ def test_held_action_rejects_malformed_updates_even_between_samples(message, rea "anchor": _scalar_stream([(10.0, 0.0)]), "commands": [ _Obs(1, message), - _Obs(2, JointState(name=["left", "right", "gripper"], position=[1, 2, 0])), + _Obs( + 2, + JointState( + header=header_now(), + velocity=[], + effort=[], + name=["left", "right", "gripper"], + position=[1, 2, 0], + ), + ), ], } ) @@ -740,9 +888,22 @@ def test_nearest_validates_selected_samples_not_unused_partial_messages(): { "anchor": _scalar_stream([(10.0, 0.0), (11.0, 0.0)]), "commands": [ - _Obs(10, JointState(name=["joint"], position=[1])), - _Obs(10.5, JointState(name=[], position=[])), - _Obs(11, JointState(name=["joint"], position=[2])), + _Obs( + 10, + JointState( + header=header_now(), velocity=[], effort=[], name=["joint"], position=[1] + ), + ), + _Obs( + 10.5, + JointState(header=header_now(), velocity=[], effort=[], name=[], position=[]), + ), + _Obs( + 11, + JointState( + header=header_now(), velocity=[], effort=[], name=["joint"], position=[2] + ), + ), ], } ) @@ -779,6 +940,9 @@ def test_held_actions_do_not_relax_camera_rate_or_gap_checks(): _Obs( 1, JointState( + header=header_now(), + velocity=[], + effort=[], name=["left", "right", "gripper"], position=[1, 2, 0], ), diff --git a/dimos/imitation/dataprep/test_generated_images.py b/dimos/imitation/dataprep/test_generated_images.py index 7f4e9a4f72e..5ff3ff0e692 100644 --- a/dimos/imitation/dataprep/test_generated_images.py +++ b/dimos/imitation/dataprep/test_generated_images.py @@ -15,27 +15,37 @@ # Copyright 2026 Dimensional Inc. # SPDX-License-Identifier: Apache-2.0 +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import Image from dimos_generated.std_msgs.msg import Header import numpy as np import pytest -from dimos.imitation.dataprep.core import StreamField, is_image_array, resolve_field +from dimos.imitation.dataprep.core import is_image_array, resolve_field +from dimos.imitation.dataprep.schema import FeatureSpec from dimos.msgs.image import compressed_image_from_image, image_from_array @pytest.mark.parametrize("compressed", [False, True]) def test_generated_camera_field_is_a_pixel_array_after_cdr(compressed): pixels = np.full((8, 10, 3), [30, 90, 150], dtype=np.uint8) - image = image_from_array(pixels, encoding="bgr8", header=Header(frame_id="camera")) + image = image_from_array( + pixels, encoding="bgr8", header=Header(frame_id="camera", stamp=Time(sec=0, nanosec=0)) + ) message = compressed_image_from_image(image, quality=100) if compressed else image - message = type(message).decode(message.encode()) - before = message.encode() - array = resolve_field(message, StreamField(stream="camera")) + message = cdr_decode(cdr_encode(message), type(message)) + before = cdr_encode(message) + array = resolve_field( + message, + FeatureSpec( + stream="camera", dtype="video", shape=(8, 10, 3), names=["height", "width", "channels"] + ), + ) assert array.shape == pixels.shape and array.dtype == pixels.dtype assert is_image_array(array) np.testing.assert_allclose(array, pixels, atol=2 if compressed else 0) - assert message.encode() == before + assert cdr_encode(message) == before def test_generated_mono16_camera_respects_row_padding_and_endian(): @@ -43,9 +53,20 @@ def test_generated_mono16_camera_respects_row_padding_and_endian(): padded = np.zeros((2, 6), dtype=np.uint8) padded[:, :4] = pixels.view(np.uint8).reshape(2, 4) message = Image( - width=2, height=2, encoding="mono16", is_bigendian=1, step=6, data=padded.ravel() + width=2, + height=2, + encoding="mono16", + is_bigendian=1, + step=6, + data=np.asarray(padded.ravel(), dtype=np.uint8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + array = resolve_field( + {"camera": message}, + FeatureSpec( + stream="dict", field="camera", dtype="uint16", shape=(2, 2), names=["height", "width"] + ), ) - array = resolve_field({"camera": message}, StreamField(stream="dict", field="camera")) assert is_image_array(array) np.testing.assert_array_equal(array, pixels) assert array.strides == (6, 2) diff --git a/dimos/imitation/dataprep/test_mcap_source.py b/dimos/imitation/dataprep/test_mcap_source.py index fe2b21031b0..1fa8a6b5d8a 100644 --- a/dimos/imitation/dataprep/test_mcap_source.py +++ b/dimos/imitation/dataprep/test_mcap_source.py @@ -21,9 +21,9 @@ from typing import Any import weakref +from dimos_cdrgen.registry import encode as cdr_encode, schema as message_schema from dimos_generated.sensor_msgs.msg import Image, JointState from dimos_generated.std_msgs.msg import Header, String -from dimos_message_build.registry import encode as cdr_encode, schema as message_schema from mcap.writer import Writer as McapWriter import numpy as np import pytest diff --git a/dimos/imitation/dataprep/test_schema.py b/dimos/imitation/dataprep/test_schema.py index 4909da63584..daa80895a8a 100644 --- a/dimos/imitation/dataprep/test_schema.py +++ b/dimos/imitation/dataprep/test_schema.py @@ -16,11 +16,11 @@ import json +from dimos_generated.sensor_msgs.msg import JointState from pydantic import ValidationError import pytest from dimos.imitation.dataprep.schema import FeatureSpec -from dimos.msgs.sensor_msgs.JointState import JointState def test_live_feature_round_trips_as_a_portable_offline_declaration(): diff --git a/dimos/imitation/test_datacollection_e2e.py b/dimos/imitation/test_datacollection_e2e.py index 151599f69aa..94b7d7f96d8 100644 --- a/dimos/imitation/test_datacollection_e2e.py +++ b/dimos/imitation/test_datacollection_e2e.py @@ -22,6 +22,8 @@ from typing import Any import cv2 +from dimos_generated.sensor_msgs.msg import Image, JointState +from dimos_generated.std_msgs.msg import Header, String import h5py import numpy as np import pyarrow.parquet as pq @@ -47,15 +49,13 @@ SyncConfig, ) from dimos.memory.store.sqlite import SqliteStore -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.std_msgs.String import String +from dimos.msgs.image import image_from_array, image_view +from dimos.msgs.time import time_from_seconds from dimos.utils.testing.waiting import wait_until pytestmark = [ pytest.mark.self_hosted, pytest.mark.skipif_aarch64, - pytest.mark.skipif_no_turbojpeg, ] @@ -187,22 +187,20 @@ def publish(name: str, message: Any) -> None: pixel = int(base + frame) * 4 publish( "color_image", - Image( - data=np.full((16, 16, 3), pixel, dtype=np.uint8), - format=ImageFormat.RGB, - frame_id="camera", - ts=ts, + image_from_array( + np.full((16, 16, 3), pixel, dtype=np.uint8), + encoding="rgb8", + header=Header(frame_id="camera", stamp=time_from_seconds(ts)), ), ) publish( "coordinator_joint_state", JointState( - ts=ts, - frame_id="arm", + header=Header(frame_id="arm", stamp=time_from_seconds(ts)), name=["joint_0", "joint_1"], - position=[base + frame, base + 100.0 + frame], - velocity=[0.0, 0.0], - effort=[0.0, 0.0], + position=np.array([base + frame, base + 100.0 + frame], dtype=np.float64), + velocity=np.array([0.0, 0.0], dtype=np.float64), + effort=np.array([0.0, 0.0], dtype=np.float64), ), ) if success: @@ -260,7 +258,7 @@ def recorded_session( (108.0, 110.0, True, "place"), ] recorded_images: dict[float, np.ndarray[Any, Any]] = { - observation.ts: observation.data.data + observation.ts: image_view(observation.data) for observation in store.stream("color_image", Image).to_list() } diff --git a/dimos/manipulation/conftest.py b/dimos/manipulation/conftest.py index 706bbec57fb..52f84b15ba4 100644 --- a/dimos/manipulation/conftest.py +++ b/dimos/manipulation/conftest.py @@ -19,6 +19,7 @@ from typing import Any, Protocol, cast from unittest.mock import DEFAULT, MagicMock +from dimos_generated.builtin_interfaces.msg import Duration from dimos_generated.dimos_msgs.msg import TrajectoryStatus import pytest from pytest_mock import MockerFixture @@ -58,7 +59,12 @@ def invoke(task: str, method: str, args: dict | None = None): return DEFAULT coordinator.task_invoke.return_value = TrajectoryStatus( - header=header_now(), state=TrajectoryState.IDLE + header=header_now(), + state=TrajectoryState.IDLE, + progress=0.0, + time_elapsed=Duration(sec=0, nanosec=0), + time_remaining=Duration(sec=0, nanosec=0), + error="", ) coordinator.task_invoke.side_effect = invoke return coordinator diff --git a/dimos/manipulation/control/servo_control/cartesian_motion_controller.py b/dimos/manipulation/control/servo_control/cartesian_motion_controller.py index 457cd4e97a0..347b3b3702e 100644 --- a/dimos/manipulation/control/servo_control/cartesian_motion_controller.py +++ b/dimos/manipulation/control/servo_control/cartesian_motion_controller.py @@ -35,6 +35,7 @@ from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Twist, Vector3 from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header +import numpy as np from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.core import rpc @@ -548,8 +549,8 @@ def _control_loop(self) -> None: # Publish joint command joint_cmd = JointCommand( - header=Header(stamp=time_from_seconds(current_time)), - positions=list(target_joints), + header=Header(stamp=time_from_seconds(current_time), frame_id=""), + positions=np.asarray(list(target_joints), dtype=np.float64), ) # Always try to publish - the Out stream will handle transport availability diff --git a/dimos/manipulation/control/test_generated_legacy_controllers.py b/dimos/manipulation/control/test_generated_legacy_controllers.py index 75dc43288b7..c559c6a62c0 100644 --- a/dimos/manipulation/control/test_generated_legacy_controllers.py +++ b/dimos/manipulation/control/test_generated_legacy_controllers.py @@ -15,10 +15,14 @@ from types import SimpleNamespace from unittest.mock import MagicMock +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.dimos_msgs.msg import JointCommand from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import numpy as np import pytest from dimos.manipulation.control.servo_control.cartesian_motion_controller import ( @@ -38,9 +42,22 @@ def test_trajectory_loop_publishes_generated_sample_with_source_time( controller._trajectory = JointTrajectory( joint_names=["a", "b"], points=[ - JointTrajectoryPoint(positions=[0.0, 0.0], time_from_start=duration_from_seconds(0)), - JointTrajectoryPoint(positions=[1.0, 2.0], time_from_start=duration_from_seconds(1)), + JointTrajectoryPoint( + positions=np.array([0.0, 0.0], dtype=np.float64), + time_from_start=duration_from_seconds(0), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointTrajectoryPoint( + positions=np.array([1.0, 2.0], dtype=np.float64), + time_from_start=duration_from_seconds(1), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) controller._state = TrajectoryState.EXECUTING controller._start_time = 100.0 @@ -56,8 +73,9 @@ def test_trajectory_loop_publishes_generated_sample_with_source_time( try: controller._execution_loop() controller.joint_position_command.publish.assert_called_once() - message = JointCommand.decode( - controller.joint_position_command.publish.call_args.args[0].encode() + message = cdr_decode( + cdr_encode(controller.joint_position_command.publish.call_args.args[0]), + JointCommand, ) assert list(message.positions) == [0.5, 1.0] assert (message.header.stamp.sec, message.header.stamp.nanosec) == (100, 500000000) @@ -73,9 +91,18 @@ def test_cartesian_generated_command_preserves_joint_count_and_pid_clamp( driver.get_forward_kinematics.return_value = (0, [0.0] * 6) driver.get_inverse_kinematics.return_value = (0, [0.1, 0.2, 0.3]) controller._arm_driver = driver - controller._latest_joint_state = JointState(position=[0.0, 0.0]) + controller._latest_joint_state = JointState( + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) controller._target_pose_ = PoseStamped( - pose=Pose(position=Point(x=0.1), orientation=Quaternion(w=1.0)) + pose=Pose( + position=Point(x=0.1, y=0.0, z=0.0), orientation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) controller._is_tracking = True controller._last_target_time = 100.0 @@ -92,8 +119,9 @@ def test_cartesian_generated_command_preserves_joint_count_and_pid_clamp( try: controller._control_loop() controller.joint_position_command.publish.assert_called_once() - message = JointCommand.decode( - controller.joint_position_command.publish.call_args.args[0].encode() + message = cdr_decode( + cdr_encode(controller.joint_position_command.publish.call_args.args[0]), + JointCommand, ) assert list(message.positions) == [0.1, 0.2] assert (message.header.stamp.sec, message.header.stamp.nanosec) == (100, 500000000) diff --git a/dimos/manipulation/control/trajectory_controller/joint_trajectory_controller.py b/dimos/manipulation/control/trajectory_controller/joint_trajectory_controller.py index 52f6c30b923..11c98bd2b3f 100644 --- a/dimos/manipulation/control/trajectory_controller/joint_trajectory_controller.py +++ b/dimos/manipulation/control/trajectory_controller/joint_trajectory_controller.py @@ -36,6 +36,7 @@ from dimos_generated.dimos_msgs.msg import JointCommand, RobotState, TrajectoryStatus from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.trajectory_msgs.msg import JointTrajectory +import numpy as np from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.core import rpc @@ -325,7 +326,9 @@ def _execution_loop(self) -> None: q_ref, _qd_ref = sample_trajectory(self._trajectory, t) # Create and publish command (outside lock would be better but simpler here) - cmd = JointCommand(positions=q_ref, header=header_now()) + cmd = JointCommand( + positions=np.asarray(q_ref, dtype=np.float64), header=header_now() + ) # Publish - must release lock first for thread safety trajectory_active = True diff --git a/dimos/manipulation/execution_manager.py b/dimos/manipulation/execution_manager.py index d9535409222..50276d668ed 100644 --- a/dimos/manipulation/execution_manager.py +++ b/dimos/manipulation/execution_manager.py @@ -23,6 +23,7 @@ from dimos_generated.dimos_msgs.msg import TrajectoryStatus from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import numpy as np from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.control.coordinator import ControlCoordinator @@ -479,14 +480,23 @@ def _columns(trajectory: JointTrajectory, columns: list[int]) -> JointTrajectory points=[ JointTrajectoryPoint( time_from_start=point.time_from_start, - positions=[point.positions[index] for index in columns], - velocities=[point.velocities[index] for index in columns] - if point.velocities - else [], - accelerations=[point.accelerations[index] for index in columns] - if point.accelerations - else [], - effort=[point.effort[index] for index in columns] if point.effort else [], + positions=np.asarray( + [point.positions[index] for index in columns], dtype=np.float64 + ), + velocities=np.asarray( + [point.velocities[index] for index in columns] if len(point.velocities) else [], + dtype=np.float64, + ), + accelerations=np.asarray( + [point.accelerations[index] for index in columns] + if len(point.accelerations) + else [], + dtype=np.float64, + ), + effort=np.asarray( + [point.effort[index] for index in columns] if len(point.effort) else [], + dtype=np.float64, + ), ) for point in trajectory.points ], diff --git a/dimos/manipulation/grasping/grasp_gen_x/test_module.py b/dimos/manipulation/grasping/grasp_gen_x/test_module.py index edcbdf69d0d..e13f04e36e1 100644 --- a/dimos/manipulation/grasping/grasp_gen_x/test_module.py +++ b/dimos/manipulation/grasping/grasp_gen_x/test_module.py @@ -23,8 +23,9 @@ import sys from typing import Any +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.dimos_msgs.msg import GraspCandidate, GraspCandidateArray -from dimos_generated.geometry_msgs.msg import Point, Pose +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion from dimos_generated.std_msgs.msg import Header import numpy as np import pytest @@ -70,18 +71,29 @@ def test_public_adapter_import_does_not_load_optional_runtime() -> None: def test_messages_round_trip_empty_and_score() -> None: value = GraspCandidateArray( header=Header(stamp=time_from_seconds(3.0), frame_id="camera"), - candidates=[GraspCandidate(pose=Pose(position=Point(x=1, y=2)), score=0.25)], + candidates=[ + GraspCandidate( + pose=Pose( + position=Point(x=1, y=2, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + score=0.25, + ) + ], ) - decoded = GraspCandidateArray.decode(value.encode()) + decoded = cdr_decode(cdr_encode(value), GraspCandidateArray) assert decoded.header.frame_id == "camera" assert decoded.header.stamp == time_from_seconds(3.0) assert decoded.candidates[0].score == pytest.approx(0.25) assert ( - GraspCandidateArray.decode( - GraspCandidateArray( - header=Header(stamp=time_from_seconds(3.0), frame_id="camera"), candidates=[] - ).encode() + cdr_decode( + cdr_encode( + GraspCandidateArray( + header=Header(stamp=time_from_seconds(3.0), frame_id="camera"), candidates=[] + ) + ), + GraspCandidateArray, ).candidates == [] ) @@ -97,7 +109,10 @@ def test_ranked_spec_is_canonical_during_legacy_contract_transition() -> None: "scene_pointcloud", ] assert list(signature.parameters) == ["self", "object_pointcloud"] - assert signature.parameters["object_pointcloud"].annotation.__name__ == "PointCloud2" + assert ( + signature.parameters["object_pointcloud"].annotation.__msgtype__ + == "sensor_msgs/msg/PointCloud2" + ) assert signature.return_annotation is GraspCandidateArray diff --git a/dimos/manipulation/grasping/heuristic_grasp.py b/dimos/manipulation/grasping/heuristic_grasp.py index 74fec3113db..3979b9487d6 100644 --- a/dimos/manipulation/grasping/heuristic_grasp.py +++ b/dimos/manipulation/grasping/heuristic_grasp.py @@ -16,6 +16,7 @@ from __future__ import annotations +from copy import deepcopy import math from dimos_generated.dimos_msgs.msg import GraspCandidate, GraspCandidateArray @@ -59,7 +60,7 @@ def propose_grasps(self, object_pointcloud: PointCloud2) -> GraspCandidateArray: orientation=quaternion_from_euler(-math.pi, 0.0, self._narrow_axis_yaw(xy)), ) return GraspCandidateArray( - header=object_pointcloud.header, + header=deepcopy(object_pointcloud.header), candidates=[GraspCandidate(pose=pose, score=1.0)], ) diff --git a/dimos/manipulation/grasping/test_heuristic_grasp.py b/dimos/manipulation/grasping/test_heuristic_grasp.py index 5fa2c160067..2a7fa74c442 100644 --- a/dimos/manipulation/grasping/test_heuristic_grasp.py +++ b/dimos/manipulation/grasping/test_heuristic_grasp.py @@ -32,7 +32,8 @@ def _cloud( points: np.ndarray, *, frame_id: str = "world", stamp: Time | None = None ) -> PointCloud2: return pointcloud_from_xyz( - points.astype(np.float32), header=Header(frame_id=frame_id, stamp=stamp or Time(sec=1)) + points.astype(np.float32), + header=Header(frame_id=frame_id, stamp=stamp or Time(sec=1, nanosec=0)), ) @@ -63,7 +64,7 @@ def test_heuristic_grasp_proposes_centered_top_down_pose(module: HeuristicGraspM assert len(proposals.candidates) == 1 assert proposals.header.frame_id == "world" - assert proposals.header.stamp == Time(sec=1) + assert proposals.header.stamp == Time(sec=1, nanosec=0) pose = proposals.candidates[0].pose assert pose.position.x == pytest.approx(0.0) assert pose.position.y == pytest.approx(0.0) @@ -111,10 +112,10 @@ def test_heuristic_grasp_canonicalizes_pca_eigenvector_sign( @pytest.mark.parametrize( "points, frame_id, stamp, error", [ - (np.zeros((2, 3)), "world", Time(), "at least three"), - (np.asarray([[0.0, 0.0, math.nan]] * 3), "world", Time(), "finite"), - (np.zeros((3, 3)), "", Time(), "frame_id"), - (np.zeros((3, 3)), "world", Time(nanosec=1000000000), "nanosec"), + (np.zeros((2, 3)), "world", Time(sec=0, nanosec=0), "at least three"), + (np.asarray([[0.0, 0.0, math.nan]] * 3), "world", Time(sec=0, nanosec=0), "finite"), + (np.zeros((3, 3)), "", Time(sec=0, nanosec=0), "frame_id"), + (np.zeros((3, 3)), "world", Time(nanosec=1000000000, sec=0), "nanosec"), ], ) def test_heuristic_grasp_rejects_invalid_pointclouds(module, points, frame_id, stamp, error): @@ -122,7 +123,7 @@ def test_heuristic_grasp_rejects_invalid_pointclouds(module, points, frame_id, s module.propose_grasps(_cloud(points, frame_id=frame_id, stamp=stamp)) -@pytest.mark.parametrize("stamp", [Time(), Time(sec=1700000000, nanosec=123456789)]) +@pytest.mark.parametrize("stamp", [Time(sec=0, nanosec=0), Time(sec=1700000000, nanosec=123456789)]) def test_heuristic_grasp_preserves_exact_header_without_aliasing(module, stamp): cloud = _cloud(np.zeros((3, 3)), stamp=stamp) proposals = module.propose_grasps(cloud) diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index 48b73ef6bc7..2d7ee813035 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -19,6 +19,7 @@ import asyncio from collections import Counter from collections.abc import Mapping, Sequence +from copy import deepcopy from enum import Enum import math import threading @@ -26,6 +27,20 @@ import traceback from typing import Any, Literal, TypeAlias +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.dimos_msgs.msg import GraspCandidateArray +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.sensor_msgs.msg import JointState, PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage import numpy as np from pydantic import Field, model_validator @@ -90,15 +105,9 @@ ) from dimos.manipulation.visualization.factory import create_manipulation_visualization from dimos.manipulation.visualization.operator import ManipulationOperator -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.geometry import transform_from_pose +from dimos.msgs.pointcloud import pointcloud_xyz +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.perception.experimental.object import Object as DetObject from dimos.utils.logging_config import setup_logger @@ -159,7 +168,7 @@ class ManipulationModuleConfig(ModuleConfig): # to a link the model already publishes -- an eye-in-hand camera, say. One # publisher for the whole chain: a second module publishing the mount at its # own rate leaves the two edges of one chain stamped up to a period apart. - static_transforms: list[Transform] = Field(default_factory=list) + static_transforms: list[TransformStamped] = Field(default_factory=list) # Frame the voxel_map port's clouds must already be expressed in. world_frame: str = "world" # Edge length of a voxel_map cell (meters). Must match the mapper's @@ -308,12 +317,14 @@ def _initialize_planning(self) -> None: fz = self.config.floor_z thickness = 0.2 floor_pose = Pose( - Vector3(0.7, 0.0, fz - thickness / 2), - Quaternion(0.0, 0.0, 0.0, 1.0), + position=Point(x=0.7, y=0.0, z=fz - thickness / 2), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), ) floor_obs = Obstacle( name="floor", - pose=floor_pose, + pose=PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), pose=floor_pose + ), obstacle_type=ObstacleType.BOX, dimensions=(0.6, 1.2, thickness), ) @@ -355,7 +366,7 @@ def _on_joint_state(self, msg: JointState) -> None: len(msg.name) != len(msg.position) or len(set(canonical_names)) != len(canonical_names) or not np.isfinite(msg.position).all() - or not np.isfinite(msg.ts) + or not np.isfinite(to_seconds(msg.header.stamp)) ): logger.warning("Skipping malformed model state") return @@ -365,12 +376,16 @@ def _on_joint_state(self, msg: JointState) -> None: return indices = [name_to_idx[name] for name in names] state = JointState( - ts=msg.ts, + header=deepcopy(msg.header), name=list(names), - position=[msg.position[index] for index in indices], - velocity=[msg.velocity[index] for index in indices] - if len(msg.velocity) == len(msg.name) - else [], + position=np.asarray([msg.position[index] for index in indices], dtype=np.float64), + velocity=np.asarray( + [msg.velocity[index] for index in indices] + if len(msg.velocity) == len(msg.name) + else [], + dtype=np.float64, + ), + effort=np.array([], dtype=np.float64), ) self._world_monitor.on_joint_state(state) if self._init_joints is None: @@ -387,7 +402,7 @@ def _tf_publish_loop(self) -> None: try: if self._world_monitor is None: break - transforms: list[Transform] = [] + transforms: list[TransformStamped] = [] config = self.config.model # Use one measurement for every dynamic edge in this publication. state = self._world_monitor.current_model_joint_state() @@ -396,26 +411,25 @@ def _tf_publish_loop(self) -> None: continue ee_pose = self._world_monitor.get_group_ee_pose(group.id, state) if ee_pose is not None and group.tip_link is not None: - ee_tf = Transform.from_pose(group.tip_link, ee_pose) - ee_tf.frame_id = "world" + ee_tf = transform_from_pose(ee_pose, child_frame_id=group.tip_link) transforms.append(ee_tf) for link_name in config.tf_extra_links: if not state.name: break link_pose = self._world_monitor.get_link_pose(link_name, state) if link_pose is not None: - link_tf = Transform.from_pose(link_name, link_pose) - link_tf.frame_id = "world" + link_tf = transform_from_pose(link_pose, child_frame_id=link_name) transforms.append(link_tf) # Fixed mounts are valid at the same measurement time as their parent. - stamp = state.ts if state.name else time.time() + stamp = state.header.stamp if state.name else time_from_seconds(time.time()) for static in self.config.static_transforms: - static.ts = stamp + static = deepcopy(static) + static.header.stamp = deepcopy(stamp) transforms.append(static) if transforms: - self.tf.publish(TFMessage(*transforms)) + self.tf.publish(TFMessage(transforms=transforms)) except Exception as e: logger.warning("TF publish failed", error=str(e)) @@ -544,7 +558,7 @@ def get_current_joints(self) -> list[float] | None: return list(state.position) return None - def get_ee_pose(self, group_id: PlanningGroupID | None = None) -> Pose | None: + def get_ee_pose(self, group_id: PlanningGroupID | None = None) -> PoseStamped | None: """Get a planning group's current tip pose.""" if self._world_monitor: try: @@ -562,7 +576,13 @@ def is_collision_free(self, joints: list[float]) -> bool: joints: Joint configuration to check """ if self._world_monitor: - joint_state = JointState(name=self.config.model.joint_names, position=joints) + joint_state = JointState( + name=self.config.model.joint_names, + position=np.asarray(joints, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) return self._world_monitor.is_state_valid(joint_state) return False @@ -739,7 +759,7 @@ def inverse_kinematics( if seed_state is None: selection = self._world_monitor.planning_groups.select(group_ids) current = self._world_monitor.current_model_joint_state() - if not current.name and not current.position: + if not current.name and not len(current.position): return IKResult(status=IKStatus.NO_SOLUTION, message="No joint state") seed_state = filter_joint_state_to_selected_joints(current, selection.joint_names) except (KeyError, ValueError) as exc: @@ -820,14 +840,20 @@ def generate_plan_to_joint_targets( goal_names.extend(normalized_target.name) goal_positions.extend(normalized_target.position) - goal = JointState(name=goal_names, position=goal_positions) + goal = JointState( + name=goal_names, + position=np.asarray(goal_positions, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) return self._plan_selected_path( group_ids, start, goal, planning_epoch, resolved_speed_scale ) def generate_plan_to_pose_targets( self, - pose_targets: Mapping[PlanningGroupID, Pose], + pose_targets: Mapping[PlanningGroupID, PoseStamped], auxiliary_groups: Sequence[PlanningGroupID] = (), speed_scale: float | None = None, ) -> GeneratedPlan | None: @@ -837,14 +863,7 @@ def generate_plan_to_pose_targets( if not pose_targets: self._fail("At least one pose target is required") return None - stamped_targets = { - group_id: PoseStamped( - frame_id="world", - position=pose.position, - orientation=pose.orientation, - ) - for group_id, pose in pose_targets.items() - } + stamped_targets = dict(pose_targets) auxiliary_ids = tuple(auxiliary_groups) group_ids = tuple(dict.fromkeys((*stamped_targets.keys(), *auxiliary_ids))) planning = self._begin_group_planning(speed_scale) @@ -942,12 +961,28 @@ def move_linear( plan_result = PlanResult(PlanStatus.AMBIGUOUS_GROUP, group.message) return MoveResult(plan_result, None, delta, check_collision) resolved_speed = self.config.linear_speed_scale if speed_scale is None else speed_scale - relative = Transform( - translation=Vector3(*delta), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), + relative = TransformStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + transform=Transform( + translation=Vector3(x=delta[0], y=delta[1], z=delta[2]), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + child_frame_id="", ) plan = self.generate_cartesian_plan( - {group.id: (Transform.identity(), relative)}, + { + group.id: ( + TransformStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + child_frame_id="", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + relative, + ) + }, CartesianPathConfig(), speed_scale=resolved_speed, check_collision=check_collision, @@ -1031,12 +1066,23 @@ def selected(state: JointState) -> JointState: positions = dict(zip(state.name, state.position, strict=True)) return JointState( name=list(group.joint_names), - position=[positions[name] for name in group.joint_names], + position=np.asarray( + [positions[name] for name in group.joint_names], dtype=np.float64 + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) if config.home_joints is not None: presets["home"] = selected( - JointState(name=config.joint_names, position=config.home_joints) + JointState( + name=config.joint_names, + position=np.asarray(config.home_joints, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) ) if self._init_joints is not None: presets["init"] = selected(self._init_joints) @@ -1242,7 +1288,10 @@ def add_obstacle( obstacle = Obstacle( name=name, obstacle_type=obstacle_type, - pose=PoseStamped(position=pose.position, orientation=pose.orientation), + pose=PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose(position=pose.position, orientation=pose.orientation), + ), dimensions=tuple(dimensions) if dimensions else (), mesh_path=mesh_path, ) @@ -1277,7 +1326,10 @@ def update_obstacle( obstacle = Obstacle( name=name, obstacle_type=obstacle_type, - pose=PoseStamped(position=pose.position, orientation=pose.orientation), + pose=PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose(position=pose.position, orientation=pose.orientation), + ), dimensions=tuple(dimensions) if dimensions else (), color=rgba, mesh_path=mesh_path, @@ -1298,17 +1350,17 @@ def _apply_voxel_map(self, cloud: PointCloud2) -> None: if self._world_monitor is None: return frame = self.config.world_frame - if cloud.frame_id != frame: + if cloud.header.frame_id != frame: # The points are metric positions in the planning frame. Registering # them through a guessed transform would invent geometry. logger.warning( "Voxel map is in frame '%s', not the planning frame '%s'; dropped a map.", - cloud.frame_id, + cloud.header.frame_id, frame, ) return - points = cloud.points_f32() + points = pointcloud_xyz(cloud) if not len(points): # An empty map is how a mapper says the space it owns is now clear. self._world_monitor.remove_obstacle(VOXEL_MAP_OBSTACLE_ID) @@ -1320,7 +1372,13 @@ def _apply_voxel_map(self, cloud: PointCloud2) -> None: obstacle = Obstacle( name=VOXEL_MAP_OBSTACLE_ID, obstacle_type=ObstacleType.OCTREE, - pose=PoseStamped(frame_id=frame), + pose=PoseStamped( + header=Header(frame_id=frame, stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), points=tuple(map(tuple, points.tolist())), octree_resolution=self.config.voxel_map_resolution, ) @@ -1345,7 +1403,10 @@ def update_obstacle_pose(self, name: str, pose: Pose) -> bool: return False return self._world_monitor.update_obstacle_pose( name, - PoseStamped(position=pose.position, orientation=pose.orientation), + PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose(position=pose.position, orientation=pose.orientation), + ), ) @rpc diff --git a/dimos/manipulation/manipulation_skills.py b/dimos/manipulation/manipulation_skills.py index 4dcb5aeb8db..2648a20b41d 100644 --- a/dimos/manipulation/manipulation_skills.py +++ b/dimos/manipulation/manipulation_skills.py @@ -16,6 +16,12 @@ from __future__ import annotations +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header +import numpy as np + from dimos.agents.annotation import skill from dimos.agents.capabilities import CAP_MOVEMENT from dimos.agents.skill_result import SkillResult @@ -28,10 +34,7 @@ PlanResult, ) from dimos.manipulation.planning.spec.models import PlanningGroupID -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.geometry import quaternion_euler, quaternion_from_euler def _status(result: PlanResult | ExecutionResult | CommandResult) -> str: @@ -184,20 +187,17 @@ def move_to_pose( if current is None: raise RuntimeError(f"End-effector pose for planning group {group_id!r} is unavailable") if roll is None and pitch is None and yaw is None: - orientation = current.orientation + orientation = current.pose.orientation else: - euler = current.orientation.to_euler() - orientation = Quaternion.from_euler( - Vector3( - euler.x if roll is None else roll, - euler.y if pitch is None else pitch, - euler.z if yaw is None else yaw, - ) + euler = quaternion_euler(current.pose.orientation) + orientation = quaternion_from_euler( + euler[0] if roll is None else roll, + euler[1] if pitch is None else pitch, + euler[2] if yaw is None else yaw, ) target = PoseStamped( - frame_id="world", - position=Vector3(x, y, z), - orientation=orientation, + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose(position=Point(x=x, y=y, z=z), orientation=orientation), ) plan = self.manipulation.plan_to_poses({group_id: target}) if no_plan := self._planning_result(plan): @@ -231,7 +231,15 @@ def move_to_joints( f"got {len(values)}" ) plan = self.manipulation.plan_to_joints( - {group_id: JointState(name=list(group.joint_names), position=values)} + { + group_id: JointState( + name=list(group.joint_names), + position=np.asarray(values, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + } ) if no_plan := self._planning_result(plan): return no_plan diff --git a/dimos/manipulation/pick_and_place_module.py b/dimos/manipulation/pick_and_place_module.py index 76c03e1271b..35730720455 100644 --- a/dimos/manipulation/pick_and_place_module.py +++ b/dimos/manipulation/pick_and_place_module.py @@ -16,8 +16,13 @@ from __future__ import annotations +import math from typing import Any, Literal +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.dimos_msgs.msg import GraspCandidateArray +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Vector3 +from dimos_generated.std_msgs.msg import Header from pydantic import Field from dimos.agents.annotation import skill @@ -35,10 +40,12 @@ from dimos.manipulation.grasping.grasp_gen_spec import GraspGenSpec from dimos.manipulation.manipulation_spec import ExecutionResult, ManipulationSpec, PlanResult from dimos.manipulation.planning.spec.models import PlanningGroupID -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray +from dimos.msgs.geometry import ( + pose_matrix, + quaternion_euler, + quaternion_from_euler, + translate_pose_local, +) from dimos.perception.experimental.object_scene_registration_spec import ObjectSceneRegistrationSpec @@ -79,7 +86,9 @@ class PickAndPlaceModule(Module): def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) self._objects: dict[str, dict[str, Any]] = {} - self._grasp_candidates = GraspCandidateArray() + self._grasp_candidates = GraspCandidateArray( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), candidates=[] + ) self._selected_object_id: str | None = None self._selected_grasp: PoseStamped | None = None self._holding_object = False @@ -108,7 +117,7 @@ def scan_objects(self, prompts: list[str]) -> SkillResult: ] self._objects = {str(obj["object_id"]): obj for obj in objects if "object_id" in obj} return SkillResult( - f"Detected {detections.detections_length} object(s)", + f"Detected {len(detections.detections)} object(s)", metadata={"prompts": prompts, "objects": list(self._objects.values())}, ) @@ -144,6 +153,10 @@ def pick_object( f"Object {object_id} has no point cloud in the latest scan. Use scan_objects again." ) candidates = self._grasp_generator.propose_grasps(pointcloud) + for candidate in candidates.candidates: + if not math.isfinite(candidate.score): + raise ValueError("GraspCandidate.score must be finite") + pose_matrix(candidate.pose) self._grasp_candidates = candidates self._manipulation.show_grasp_proposals(candidates) if candidates.header.frame_id != self.config.planning_frame: @@ -161,10 +174,10 @@ def pick_object( for rank, candidate in enumerate(candidates.candidates[: self.config.max_grasp_attempts]): grasp = self._apply_yaw_policy( PoseStamped( - ts=candidates.header.timestamp, - frame_id=candidates.header.frame_id, - position=candidate.pose.position, - orientation=candidate.pose.orientation, + header=candidates.header, + pose=Pose( + position=candidate.pose.position, orientation=candidate.pose.orientation + ), ), group, ) @@ -230,9 +243,10 @@ def place_at( ) group = self._gripper_group(planning_group) place = PoseStamped( - frame_id=self.config.planning_frame, - position=Vector3(x, y, z), - orientation=self._selected_grasp.orientation, + header=Header(frame_id=self.config.planning_frame, stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=x, y=y, z=z), orientation=self._selected_grasp.pose.orientation + ), ) preplace = self._offset_pose(place, self.config.pregrasp_offset) if blocked := self._move(preplace, group): @@ -248,8 +262,14 @@ def place_at( return SkillResult("Place complete") def _clear_selection(self) -> None: - self._grasp_candidates = GraspCandidateArray() - self._manipulation.show_grasp_proposals(GraspCandidateArray()) + self._grasp_candidates = GraspCandidateArray( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), candidates=[] + ) + self._manipulation.show_grasp_proposals( + GraspCandidateArray( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), candidates=[] + ) + ) self._selected_object_id = None self._selected_grasp = None @@ -298,22 +318,21 @@ def _apply_yaw_policy(self, pose: PoseStamped, group: PlanningGroupID) -> PoseSt current = self._manipulation.get_state().groups[group].end_effector_pose if current is None: return pose - euler = pose.orientation.to_euler() - current_euler = current.orientation.to_euler() + euler = quaternion_euler(pose.pose.orientation) + current_euler = quaternion_euler(current.pose.orientation) return PoseStamped( - ts=pose.ts, - frame_id=pose.frame_id, - position=pose.position, - orientation=Quaternion.from_euler(Vector3(euler.x, euler.y, current_euler.z)), + header=pose.header, + pose=Pose( + position=pose.pose.position, + orientation=quaternion_from_euler(euler[0], euler[1], current_euler[2]), + ), ) @staticmethod def _offset_pose(pose: PoseStamped, offset: float) -> PoseStamped: return PoseStamped( - ts=pose.ts, - frame_id=pose.frame_id, - position=pose.position + pose.orientation.rotate_vector(Vector3(0.0, 0.0, -offset)), - orientation=pose.orientation, + header=pose.header, + pose=translate_pose_local(pose.pose, Vector3(z=-offset, x=0.0, y=0.0)), ) def _servo( @@ -329,9 +348,9 @@ def _servo( when the arm arrived. """ result = self._manipulation.move_linear( - end.position.x - start.position.x, - end.position.y - start.position.y, - end.position.z - start.position.z, + end.pose.position.x - start.pose.position.x, + end.pose.position.y - start.pose.position.y, + end.pose.position.z - start.pose.position.z, planning_group, check_collision=False, ) diff --git a/dimos/manipulation/planning/groups/test_planning_groups.py b/dimos/manipulation/planning/groups/test_planning_groups.py index c93b5f8f331..e82495d51f8 100644 --- a/dimos/manipulation/planning/groups/test_planning_groups.py +++ b/dimos/manipulation/planning/groups/test_planning_groups.py @@ -20,6 +20,7 @@ from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from dimos.manipulation.planning.groups.discovery import ( @@ -164,20 +165,47 @@ def test_normalize_joint_target_accepts_exact_or_unnamed_target() -> None: group = PlanningGroup("left_arm", ("left/j1", "left/j2"), "base", "tool") named = normalize_joint_target( group, - JointState(name=["left/j1", "left/j2"], position=[1.0, 2.0]), + JointState( + name=["left/j1", "left/j2"], + position=np.array([1.0, 2.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + ) + unnamed = normalize_joint_target( + group, + JointState( + position=np.array([3.0, 4.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) - unnamed = normalize_joint_target(group, JointState(position=[3.0, 4.0])) assert list(named.name) == ["left/j1", "left/j2"] assert list(unnamed.name) == ["left/j1", "left/j2"] with pytest.raises(ValueError, match="missing joints"): normalize_joint_target( group, - JointState(name=["j1", "j2"], position=[1.0, 2.0]), + JointState( + name=["j1", "j2"], + position=np.array([1.0, 2.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) def test_state_projection_requires_exact_canonical_names() -> None: - state = JointState(name=["right/j1", "left/j2", "left/j1"], position=[3.0, 2.0, 1.0]) + state = JointState( + name=["right/j1", "left/j2", "left/j1"], + position=np.array([3.0, 2.0, 1.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) projected = filter_joint_state_to_selected_joints(state, ("left/j1", "right/j1")) assert list(projected.name) == ["left/j1", "right/j1"] assert list(projected.position) == [1.0, 3.0] @@ -193,7 +221,9 @@ def test_projection_and_target_normalization_copy_source_header(): source = JointState( header=Header(stamp=Time(sec=-1, nanosec=999999999), frame_id="robot"), name=["arm/b", "arm/a"], - position=[2.0, 1.0], + position=np.array([2.0, 1.0], dtype=np.float64), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) projected = filter_joint_state_to_selected_joints(source, ["arm/a"]) normalized = normalize_joint_target(group, source) diff --git a/dimos/manipulation/planning/groups/utils.py b/dimos/manipulation/planning/groups/utils.py index a92768aff10..d00d6af1625 100644 --- a/dimos/manipulation/planning/groups/utils.py +++ b/dimos/manipulation/planning/groups/utils.py @@ -15,6 +15,7 @@ """Planning-group selection and canonical joint-state projection.""" from collections.abc import Sequence +from copy import deepcopy from dimos_generated.sensor_msgs.msg import JointState import numpy as np @@ -32,9 +33,13 @@ def filter_joint_state_to_selected_joints( if missing: raise ValueError(f"Joint state is missing selected joints: {missing}") return JointState( - header=joint_state.header, + header=deepcopy(joint_state.header), name=list(joint_names), - position=[float(positions_by_name[name]) for name in joint_names], + position=np.asarray( + [float(positions_by_name[name]) for name in joint_names], dtype=np.float64 + ), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) @@ -47,7 +52,11 @@ def normalize_joint_target(group: PlanningGroup, target: JointState) -> JointSta f"expected {len(group.joint_names)}" ) return JointState( - header=target.header, name=list(group.joint_names), position=list(target.position) + header=deepcopy(target.header), + name=list(group.joint_names), + position=np.asarray(list(target.position), dtype=np.float64), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) if len(target.name) != len(target.position): raise ValueError( @@ -62,9 +71,13 @@ def normalize_joint_target(group: PlanningGroup, target: JointState) -> JointSta if extra: raise ValueError(f"Target for '{group.id}' has extra joints: {sorted(extra)}") return JointState( - header=target.header, + header=deepcopy(target.header), name=list(group.joint_names), - position=[float(positions[name]) for name in group.joint_names], + position=np.asarray( + [float(positions[name]) for name in group.joint_names], dtype=np.float64 + ), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) diff --git a/dimos/manipulation/planning/kinematics/drake_optimization_ik.py b/dimos/manipulation/planning/kinematics/drake_optimization_ik.py index 80a4f9ee35b..fd73e715c47 100644 --- a/dimos/manipulation/planning/kinematics/drake_optimization_ik.py +++ b/dimos/manipulation/planning/kinematics/drake_optimization_ik.py @@ -19,8 +19,10 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, cast +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import PoseStamped from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header import numpy as np from dimos.manipulation.planning.groups.models import PlanningGroup @@ -322,7 +324,13 @@ def _solve_single( joint_solution = np.clip(joint_solution, lower_limits, upper_limits) # Compute actual error using FK - solution_state = JointState(name=joint_names, position=joint_solution.tolist()) + solution_state = JointState( + name=joint_names, + position=np.asarray(joint_solution.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) with world.scratch_context() as ctx: world.set_joint_state(ctx, solution_state) actual_matrix = world.get_link_pose(ctx, target_frame_name) @@ -350,7 +358,13 @@ def _create_success_result( ) -> IKResult: return IKResult( status=IKStatus.SUCCESS, - joint_state=JointState(name=joint_names, position=joint_positions.tolist()), + joint_state=JointState( + name=joint_names, + position=np.asarray(joint_positions.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), position_error=position_error, orientation_error=orientation_error, iterations=iterations, diff --git a/dimos/manipulation/planning/kinematics/jacobian_ik.py b/dimos/manipulation/planning/kinematics/jacobian_ik.py index fd37e8b44bb..ca2b47a6e49 100644 --- a/dimos/manipulation/planning/kinematics/jacobian_ik.py +++ b/dimos/manipulation/planning/kinematics/jacobian_ik.py @@ -27,6 +27,8 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.std_msgs.msg import Header import numpy as np from dimos.manipulation.planning.groups.models import PlanningGroup @@ -145,7 +147,13 @@ def solve( else: # Random seed within joint limits random_positions = np.random.uniform(lower_limits, upper_limits) - current_seed = JointState(name=joint_names, position=random_positions.tolist()) + current_seed = JointState( + name=joint_names, + position=np.asarray(random_positions.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) # Solve iterative IK result = self.solve_iterative( @@ -210,7 +218,13 @@ def solve_pose_targets( if request is None: return _create_failure_result(IKStatus.NO_SOLUTION, "Invalid pose target request") - full_seed = JointState(name=request.joint_names, position=request.seed_positions.tolist()) + full_seed = JointState( + name=request.joint_names, + position=np.asarray(request.seed_positions.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) result = self.solve_iterative( world=world, target_pose=request.target_pose, @@ -229,7 +243,13 @@ def solve_pose_targets( full_positions[request.group_indices] = np.asarray( result.joint_state.position, dtype=np.float64 ) - full_state = JointState(name=request.joint_names, position=full_positions.tolist()) + full_state = JointState( + name=request.joint_names, + position=np.asarray(full_positions.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) if not world.check_config_collision_free(full_state): return _create_failure_result(IKStatus.COLLISION, "IK solution is in collision") return result @@ -274,7 +294,13 @@ def solve_iterative( for iteration in range(max_iterations): with world.scratch_context() as ctx: # Set current position (convert to JointState for API) - current_state = JointState(name=joint_names, position=current_joints.tolist()) + current_state = JointState( + name=joint_names, + position=np.asarray(current_joints.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) world.set_joint_state(ctx, current_state) if group is None: @@ -331,7 +357,13 @@ def solve_iterative( # Compute final error with world.scratch_context() as ctx: - final_state = JointState(name=joint_names, position=current_joints.tolist()) + final_state = JointState( + name=joint_names, + position=np.asarray(current_joints.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) world.set_joint_state(ctx, final_state) if group is None: final_pose = pose_matrix(world.get_ee_pose(ctx).pose) @@ -402,7 +434,13 @@ def solve_differential( if max_ratio > 1.0: q_dot = q_dot / max_ratio - return JointState(name=joint_names, velocity=q_dot.tolist()) + return JointState( + name=joint_names, + velocity=np.asarray(q_dot.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + position=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) def solve_differential_position_only( self, @@ -448,7 +486,13 @@ def solve_differential_position_only( # Compute joint velocities q_dot = J_pinv @ vel_array - return JointState(name=joint_names, velocity=q_dot.tolist()) + return JointState( + name=joint_names, + velocity=np.asarray(q_dot.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + position=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) # Result Helpers @@ -464,7 +508,13 @@ def _create_success_result( """Create a successful IK result.""" return IKResult( status=IKStatus.SUCCESS, - joint_state=JointState(name=joint_names, position=joint_positions.tolist()), + joint_state=JointState( + name=joint_names, + position=np.asarray(joint_positions.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), position_error=position_error, orientation_error=orientation_error, iterations=iterations, diff --git a/dimos/manipulation/planning/kinematics/pink_ik.py b/dimos/manipulation/planning/kinematics/pink_ik.py index 667a562cefc..5a921c52642 100644 --- a/dimos/manipulation/planning/kinematics/pink_ik.py +++ b/dimos/manipulation/planning/kinematics/pink_ik.py @@ -20,8 +20,10 @@ from dataclasses import replace from typing import TYPE_CHECKING +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import PoseStamped from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header import numpy as np import pink from pink.exceptions import NoSolutionFound @@ -408,7 +410,13 @@ def _success( ) -> IKResult: return IKResult( status=IKStatus.SUCCESS, - joint_state=JointState(name=joint_names, position=joint_positions.tolist()), + joint_state=JointState( + name=joint_names, + position=np.asarray(joint_positions.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), position_error=position_error, orientation_error=orientation_error, iterations=iterations, @@ -430,7 +438,13 @@ def _selected_result(result: IKResult, selected_names: Sequence[str]) -> IKResul return IKResult( status=result.status, joint_state=JointState( - name=list(selected_names), position=[float(positions[name]) for name in selected_names] + name=list(selected_names), + position=np.asarray( + [float(positions[name]) for name in selected_names], dtype=np.float64 + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), position_error=result.position_error, orientation_error=result.orientation_error, diff --git a/dimos/manipulation/planning/kinematics/test_drake_optimization_ik_selection.py b/dimos/manipulation/planning/kinematics/test_drake_optimization_ik_selection.py index c1a4b6354ed..6017cf18cca 100644 --- a/dimos/manipulation/planning/kinematics/test_drake_optimization_ik_selection.py +++ b/dimos/manipulation/planning/kinematics/test_drake_optimization_ik_selection.py @@ -17,7 +17,8 @@ from contextlib import contextmanager from pathlib import Path -from dimos_generated.geometry_msgs.msg import Pose, PoseStamped +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header import numpy as np @@ -45,7 +46,10 @@ def __init__(self) -> None: ) self.current_state = JointState( name=["arm/base", "arm/shoulder", "arm/elbow", "arm/wrist"], - position=[1.0, 2.0, 3.0, 4.0], + position=np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) self.collision_checked_state: JointState | None = None self.prepared = PreparedRobotModel( @@ -96,7 +100,10 @@ def fake_solve_single(self, **kwargs) -> IKResult: status=IKStatus.SUCCESS, joint_state=JointState( name=["arm/base", "arm/shoulder", "arm/elbow", "arm/wrist"], - position=[10.0, 20.0, 30.0, 40.0], + position=np.array([10.0, 20.0, 30.0, 40.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), position_error=0.0, orientation_error=0.0, @@ -107,8 +114,22 @@ def fake_solve_single(self, **kwargs) -> IKResult: result = DrakeOptimizationIK().solve_pose_targets( world=world, # type: ignore[arg-type] - pose_targets={group: PoseStamped(header=Header(frame_id=""), pose=Pose())}, - seed=JointState(name=["arm/shoulder"], position=[22.0]), + pose_targets={ + group: PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + }, + seed=JointState( + name=["arm/shoulder"], + position=np.array([22.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), check_collision=False, max_attempts=1, ) @@ -116,7 +137,7 @@ def fake_solve_single(self, **kwargs) -> IKResult: assert result.is_success() assert result.joint_state is not None assert result.joint_state.name == ["arm/shoulder", "arm/wrist"] - assert result.joint_state.position == [20.0, 40.0] + np.testing.assert_array_equal(result.joint_state.position, [20.0, 40.0]) assert calls[0]["target_frame_name"] == "group_tip_link" np.testing.assert_allclose(calls[0]["seed"], [1.0, 22.0, 3.0, 4.0]) assert calls[0]["locked_joint_positions"] == {0: 1.0, 2: 3.0} diff --git a/dimos/manipulation/planning/kinematics/test_jacobian_ik_selection.py b/dimos/manipulation/planning/kinematics/test_jacobian_ik_selection.py index e75e7551b61..21b86f1ed7c 100644 --- a/dimos/manipulation/planning/kinematics/test_jacobian_ik_selection.py +++ b/dimos/manipulation/planning/kinematics/test_jacobian_ik_selection.py @@ -19,6 +19,7 @@ from contextlib import nullcontext from pathlib import Path +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header @@ -39,7 +40,7 @@ def _pose(x: float = 0.0) -> PoseStamped: return PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=Point(x=x, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) ), @@ -97,7 +98,11 @@ def scratch_context(self) -> nullcontext[None]: def get_joint_state(self, ctx: object) -> JointState: return JointState( - name=["arm/joint_a", "arm/joint_b", "arm/gripper"], position=[0.0, 0.0, 0.9] + name=["arm/joint_a", "arm/joint_b", "arm/gripper"], + position=np.array([0.0, 0.0, 0.9], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) def set_joint_state(self, ctx: object, joint_state: JointState) -> None: @@ -121,7 +126,11 @@ def test_solve_pose_targets_filters_to_group_and_uses_group_world_methods() -> N world=world, pose_targets={_group(): _pose()}, seed=JointState( - name=["arm/joint_a", "arm/joint_b", "arm/gripper"], position=[0.0, 0.0, 0.9] + name=["arm/joint_a", "arm/joint_b", "arm/gripper"], + position=np.array([0.0, 0.0, 0.9], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), max_attempts=1, ) diff --git a/dimos/manipulation/planning/kinematics/test_pink_ik.py b/dimos/manipulation/planning/kinematics/test_pink_ik.py index b2370c81485..befa21782dd 100644 --- a/dimos/manipulation/planning/kinematics/test_pink_ik.py +++ b/dimos/manipulation/planning/kinematics/test_pink_ik.py @@ -24,6 +24,7 @@ from types import MappingProxyType, ModuleType, SimpleNamespace from typing import Any, cast +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header @@ -93,6 +94,16 @@ """ +def _target_pose(x: float = 0.0, y: float = 0.0, z: float = 0.0) -> PoseStamped: + """Create an independent, identity-oriented target for each test call.""" + return PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=x, y=y, z=z), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + ) + + class _StreamingTestPinkIK(PinkPoseTargetSolver): """Expose private control-side streaming primitives for unit tests.""" @@ -395,10 +406,7 @@ def _install_fake_modules(mocker: MockerFixture, converge: bool = True) -> _Fake def _robot_config() -> RobotModelConfig: return RobotModelConfig( model=RobotModel.from_file(Path("/tmp/fake.urdf")), - base_pose=PoseStamped( - header=Header(frame_id=""), - pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), - ), + base_pose=_target_pose(), joint_names=["joint_a", "joint_b", "joint_c"], base_link="base", planning_groups=[ @@ -550,7 +558,13 @@ def scratch_context(self) -> nullcontext[None]: def get_joint_state(self, ctx: object) -> JointState: self.joint_state_calls += 1 - return JointState(name=["joint_b", "joint_c", "joint_a"], position=[0.0, 0.0, 0.0]) + return JointState( + name=["joint_b", "joint_c", "joint_a"], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) def get_joint_limits(self) -> tuple[np.ndarray, np.ndarray]: return np.array([-1.0, -1.0, -1.0]), np.array([1.0, 1.0, 1.0]) @@ -638,7 +652,13 @@ def test_pink_ik_config_overrides_are_applied(mocker: MockerFixture) -> None: def test_joint_order_mapping_uses_names_not_positions() -> None: mapping = _build_joint_mapping(_FakeModel(), _test_joint_space()) - seed = JointState(name=["joint_b", "joint_c", "joint_a"], position=[20.0, 30.0, 10.0]) + seed = JointState( + name=["joint_b", "joint_c", "joint_a"], + position=np.array([20.0, 30.0, 10.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) assert mapping.idx_q == [1, 0, 2] assert mapping.idx_v == [1, 0, 2] @@ -769,7 +789,13 @@ def apply_increment(**kwargs: Any) -> None: configuration.update(q) mocker.patch.object(ik, "_step_configuration", side_effect=apply_increment) - initial = JointState(name=["joint_a"], position=[0.0]) + initial = JointState( + name=["joint_a"], + position=np.array([0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) alpha = 1.0 - np.exp(-2.0 * np.pi * 5.0 * 0.01) commands = [initial] @@ -777,7 +803,7 @@ def apply_increment(**kwargs: Any) -> None: commands.append( ik.step_frame_targets( robot_model=_robot_config(), - frame_targets={"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())}, + frame_targets={"tool": _target_pose()}, controlled_joints=["joint_a"], command_state=commands[-1], measured_state=initial, @@ -810,7 +836,13 @@ def apply_increment(**kwargs: Any) -> None: configuration.update(q) mocker.patch.object(ik, "_step_configuration", side_effect=apply_increment) - initial = JointState(name=["joint_a"], position=[0.0]) + initial = JointState( + name=["joint_a"], + position=np.array([0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) alpha = 1.0 - np.exp(-2.0 * np.pi * 5.0 * 0.01) commands = [initial] @@ -818,7 +850,7 @@ def apply_increment(**kwargs: Any) -> None: commands.append( ik.step_frame_targets( robot_model=_robot_config(), - frame_targets={"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())}, + frame_targets={"tool": _target_pose()}, controlled_joints=["joint_a"], command_state=commands[-1], measured_state=initial, @@ -852,10 +884,16 @@ def apply_increment(**kwargs: Any) -> None: configuration.update(q) mocker.patch.object(ik, "_step_configuration", side_effect=apply_increment) - initial = JointState(name=["joint_a"], position=[0.0]) + initial = JointState( + name=["joint_a"], + position=np.array([0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) first = ik.step_frame_targets( robot_model=_robot_config(), - frame_targets={"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())}, + frame_targets={"tool": _target_pose()}, controlled_joints=["joint_a"], command_state=initial, measured_state=initial, @@ -867,10 +905,16 @@ def apply_increment(**kwargs: Any) -> None: second = ik.step_frame_targets( robot_model=_robot_config(), - frame_targets={"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())}, + frame_targets={"tool": _target_pose()}, controlled_joints=["joint_a"], command_state=first, - measured_state=JointState(name=["joint_a"], position=[-0.1]), + measured_state=JointState( + name=["joint_a"], + position=np.array([-0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), max_command_tracking_error_rad=0.15, dt=0.01, max_joint_velocity_rad_s=100.0, @@ -893,18 +937,22 @@ def test_step_frame_targets_preserves_controlled_joint_order( result = ik.step_frame_targets( robot_model=_robot_config(), - frame_targets={ - "tool": PoseStamped( - header=Header(frame_id=""), - pose=Pose( - position=Point(x=0.1, y=0.2, z=0.3), - orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), - ), - ) - }, + frame_targets={"tool": _target_pose(x=0.1, y=0.2, z=0.3)}, controlled_joints=["joint_c", "joint_a"], - command_state=JointState(name=["joint_a", "joint_c"], position=[0.0, 0.0]), - measured_state=JointState(name=["joint_a", "joint_c"], position=[0.0, 0.0]), + command_state=JointState( + name=["joint_a", "joint_c"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + measured_state=JointState( + name=["joint_a", "joint_c"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), max_command_tracking_error_rad=_TRACKING_ERROR_RAD, dt=0.02, ) @@ -936,21 +984,24 @@ def test_step_frame_targets_builds_both_frame_tasks_with_tuning( result = ik.step_frame_targets( robot_model=_robot_config(), frame_targets={ - "tool": PoseStamped( - header=Header(frame_id=""), - pose=Pose( - position=Point(x=0.1, y=0.2, z=0.3), - orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), - ), - ), - "base": PoseStamped( - header=Header(frame_id=""), - pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), - ), + "tool": _target_pose(x=0.1, y=0.2, z=0.3), + "base": _target_pose(), }, controlled_joints=["joint_a", "joint_b", "joint_c"], - command_state=JointState(name=["joint_a", "joint_b", "joint_c"], position=[0.0, 0.0, 0.0]), - measured_state=JointState(name=["joint_a", "joint_b", "joint_c"], position=[0.0, 0.0, 0.0]), + command_state=JointState( + name=["joint_a", "joint_b", "joint_c"], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + measured_state=JointState( + name=["joint_a", "joint_b", "joint_c"], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), max_command_tracking_error_rad=_TRACKING_ERROR_RAD, dt=0.02, ) @@ -1016,25 +1067,27 @@ def test_streaming_reuses_task_stack_across_steps( context = _combined_control_context(("tool",), ["joint_a", "joint_b", "joint_c"]) mocker.patch.object(ik, "_get_control_context", return_value=context) create_tasks = mocker.spy(ik, "_create_tasks") - first_target = PoseStamped( - header=Header(frame_id=""), - pose=Pose( - position=Point(x=0.1, y=0.2, z=0.3), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) - ), - ) - second_target = PoseStamped( - header=Header(frame_id=""), - pose=Pose( - position=Point(x=0.3, y=0.2, z=0.1), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) - ), - ) + first_target = _target_pose(x=0.1, y=0.2, z=0.3) + second_target = _target_pose(x=0.3, y=0.2, z=0.1) ik.step_frame_targets( _robot_config(), {"tool": first_target}, ["joint_a", "joint_b", "joint_c"], - JointState(name=["joint_a", "joint_b", "joint_c"], position=[0.0, 0.0, 0.0]), - JointState(name=["joint_a", "joint_b", "joint_c"], position=[0.0, 0.0, 0.0]), + JointState( + name=["joint_a", "joint_b", "joint_c"], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=["joint_a", "joint_b", "joint_c"], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), _TRACKING_ERROR_RAD, ) assert context.tasks is not None @@ -1043,8 +1096,20 @@ def test_streaming_reuses_task_stack_across_steps( _robot_config(), {"tool": second_target}, ["joint_a", "joint_b", "joint_c"], - JointState(name=["joint_a", "joint_b", "joint_c"], position=[0.1, 0.2, 0.3]), - JointState(name=["joint_a", "joint_b", "joint_c"], position=[0.1, 0.2, 0.3]), + JointState( + name=["joint_a", "joint_b", "joint_c"], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=["joint_a", "joint_b", "joint_c"], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), _TRACKING_ERROR_RAD, ) @@ -1063,18 +1128,22 @@ def test_task_hooks_receive_read_only_stack_and_successful_velocity( ik.step_frame_targets( _robot_config(), - { - "tool": PoseStamped( - header=Header(frame_id=""), - pose=Pose( - position=Point(x=0.1, y=0.2, z=0.3), - orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), - ), - ) - }, + {"tool": _target_pose(x=0.1, y=0.2, z=0.3)}, ["joint_a", "joint_b", "joint_c"], - JointState(name=["joint_a", "joint_b", "joint_c"], position=[0.0, 0.0, 0.0]), - JointState(name=["joint_a", "joint_b", "joint_c"], position=[0.0, 0.0, 0.0]), + JointState( + name=["joint_a", "joint_b", "joint_c"], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=["joint_a", "joint_b", "joint_c"], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), _TRACKING_ERROR_RAD, dt=0.05, ) @@ -1097,15 +1166,21 @@ def test_after_solve_hook_is_not_called_when_solver_raises( with pytest.raises(RuntimeError, match="no solution"): ik.step_frame_targets( _robot_config(), - {"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())}, + {"tool": _target_pose()}, ["joint_a", "joint_b", "joint_c"], JointState( name=["joint_a", "joint_b", "joint_c"], - position=[0.0, 0.0, 0.0], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), JointState( name=["joint_a", "joint_b", "joint_c"], - position=[0.0, 0.0, 0.0], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), _TRACKING_ERROR_RAD, ) @@ -1151,15 +1226,21 @@ def step_outside_limits(**kwargs: Any) -> None: result = ik.step_frame_targets( robot_model=_robot_config(), - frame_targets={"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())}, + frame_targets={"tool": _target_pose()}, controlled_joints=["joint_a", "joint_b", "joint_c"], command_state=JointState( name=["joint_a", "joint_b", "joint_c"], - position=[-1.0005, 0.0, 0.0], + position=np.array([-1.0005, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), measured_state=JointState( name=["joint_a", "joint_b", "joint_c"], - position=[-1.0005, 0.0, 0.0], + position=np.array([-1.0005, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), max_command_tracking_error_rad=_TRACKING_ERROR_RAD, ) @@ -1181,10 +1262,22 @@ def test_step_frame_targets_rejects_feedback_beyond_tolerance( with pytest.raises(PinkJointLimitError, match="joint_a.*lower limit"): ik.step_frame_targets( robot_model=_robot_config(), - frame_targets={"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())}, + frame_targets={"tool": _target_pose()}, controlled_joints=["joint_a"], - command_state=JointState(name=["joint_a"], position=[0.0]), - measured_state=JointState(name=["joint_a"], position=[-1.0011]), + command_state=JointState( + name=["joint_a"], + position=np.array([0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + measured_state=JointState( + name=["joint_a"], + position=np.array([-1.0011], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), max_command_tracking_error_rad=_TRACKING_ERROR_RAD, ) @@ -1209,10 +1302,22 @@ def step_unbounded(**kwargs: Any) -> None: result = ik.step_frame_targets( robot_model=_robot_config(), - frame_targets={"tool": PoseStamped(header=Header(frame_id=""), pose=Pose())}, + frame_targets={"tool": _target_pose()}, controlled_joints=["joint_b"], - command_state=JointState(name=["joint_b"], position=[5.0]), - measured_state=JointState(name=["joint_b"], position=[5.0]), + command_state=JointState( + name=["joint_b"], + position=np.array([5.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + measured_state=JointState( + name=["joint_b"], + position=np.array([5.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), max_command_tracking_error_rad=_TRACKING_ERROR_RAD, ) @@ -1256,10 +1361,22 @@ def test_step_frame_targets_rejects_unknown_frame(mocker: MockerFixture, tmp_pat with pytest.raises(ValueError, match="missing_frame"): _StreamingTestPinkIK(PinkIKConfig()).step_frame_targets( robot_model=config, - frame_targets={"missing_frame": PoseStamped(header=Header(frame_id=""), pose=Pose())}, + frame_targets={"missing_frame": _target_pose()}, controlled_joints=config.joint_names, - command_state=JointState(name=config.joint_names, position=[0.0, 0.0, 0.0]), - measured_state=JointState(name=config.joint_names, position=[0.0, 0.0, 0.0]), + command_state=JointState( + name=config.joint_names, + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + measured_state=JointState( + name=config.joint_names, + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), max_command_tracking_error_rad=_TRACKING_ERROR_RAD, ) @@ -1369,7 +1486,13 @@ def test_pose_target_solve_constrains_joints_outside_planning_group(tmp_path: Pa tip_link="tool", ) seed_positions = np.array([0.0, 0.0, 0.0, -0.4, 0.2, 0.0, 1.2, 0.0, 0.0, 0.0]) - seed = JointState(name=joint_names, position=seed_positions.tolist()) + seed = JointState( + name=joint_names, + position=np.asarray(seed_positions.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) prepared = prepare_robot_model(config) class World: @@ -1406,7 +1529,7 @@ def is_collision_free(self, ctx: object) -> bool: cast("Any", World()), { group: PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose(position=target_pose.position, orientation=target_pose.orientation), ) }, @@ -1429,13 +1552,7 @@ def test_solve_rejects_collision_candidate(mocker: MockerFixture) -> None: result = ik.solve( world=cast("Any", _FakeWorld(collision_free=False)), - target_pose=PoseStamped( - header=Header(frame_id=""), - pose=Pose( - position=Point(x=0.1, y=0.0, z=0.0), - orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), - ), - ), + target_pose=_target_pose(x=0.1), check_collision=True, max_attempts=1, ) @@ -1463,7 +1580,10 @@ def fake_solve_targets(**_: object) -> IKResult: status=IKStatus.SUCCESS, joint_state=JointState( name=["joint_a", "joint_b", "joint_c"], - position=[0.1, 0.2, 0.3], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), position_error=0.0, orientation_error=0.0, @@ -1474,13 +1594,7 @@ def fake_solve_targets(**_: object) -> IKResult: result = ik.solve( world=cast("Any", _FakeWorld(collision_free=True)), - target_pose=PoseStamped( - header=Header(frame_id=""), - pose=Pose( - position=Point(x=0.1, y=0.0, z=0.0), - orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), - ), - ), + target_pose=_target_pose(x=0.1), check_collision=True, max_attempts=2, ) @@ -1546,7 +1660,11 @@ def test_solve_pose_targets_uses_group_tip_and_filters_group_joints( return_value=IKResult( status=IKStatus.SUCCESS, joint_state=JointState( - name=["joint_a", "joint_b", "joint_c"], position=[0.1, 0.2, 0.3] + name=["joint_a", "joint_b", "joint_c"], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ), ) @@ -1554,13 +1672,14 @@ def test_solve_pose_targets_uses_group_tip_and_filters_group_joints( result = ik.solve_pose_targets( world=cast("Any", world), - pose_targets={ - world.groups["manipulator"]: PoseStamped( - header=Header(frame_id=""), - pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), - ) - }, - seed=JointState(name=["joint_a", "joint_b", "joint_c"], position=[0.0, 0.0, 0.0]), + pose_targets={world.groups["manipulator"]: _target_pose()}, + seed=JointState( + name=["joint_a", "joint_b", "joint_c"], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), max_attempts=1, ) @@ -1568,7 +1687,7 @@ def test_solve_pose_targets_uses_group_tip_and_filters_group_joints( assert result.status == IKStatus.SUCCESS assert result.joint_state is not None assert result.joint_state.name == ["joint_a", "joint_b"] - assert result.joint_state.position == [0.1, 0.2] + np.testing.assert_array_equal(result.joint_state.position, [0.1, 0.2]) assert world.joint_state_calls == 0 @@ -1578,12 +1697,7 @@ def test_solve_pose_targets_rejects_group_without_tip(mocker: MockerFixture) -> result = ik.solve_pose_targets( world=cast("Any", world), - pose_targets={ - world.groups["no_tip"]: PoseStamped( - header=Header(frame_id=""), - pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), - ) - }, + pose_targets={world.groups["no_tip"]: _target_pose()}, ) assert result.status == IKStatus.UNSUPPORTED @@ -1599,7 +1713,11 @@ def test_solve_pose_targets_partial_seed_reads_world_state(mocker: MockerFixture return_value=IKResult( status=IKStatus.SUCCESS, joint_state=JointState( - name=["joint_a", "joint_b", "joint_c"], position=[0.1, 0.2, 0.3] + name=["joint_a", "joint_b", "joint_c"], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ), ) @@ -1607,13 +1725,14 @@ def test_solve_pose_targets_partial_seed_reads_world_state(mocker: MockerFixture result = ik.solve_pose_targets( world=cast("Any", world), - pose_targets={ - world.groups["manipulator"]: PoseStamped( - header=Header(frame_id=""), - pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), - ) - }, - seed=JointState(name=["joint_a"], position=[0.0]), + pose_targets={world.groups["manipulator"]: _target_pose()}, + seed=JointState( + name=["joint_a"], + position=np.array([0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), max_attempts=1, ) @@ -1631,7 +1750,11 @@ def test_solve_pose_targets_multi_target_uses_multi_frame_solve(mocker: MockerFi return_value=IKResult( status=IKStatus.SUCCESS, joint_state=JointState( - name=["joint_a", "joint_b", "joint_c"], position=[0.1, 0.2, 0.3] + name=["joint_a", "joint_b", "joint_c"], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), position_error=0.0, orientation_error=0.0, @@ -1641,16 +1764,16 @@ def test_solve_pose_targets_multi_target_uses_multi_frame_solve(mocker: MockerFi result = ik.solve_pose_targets( world=cast("Any", world), pose_targets={ - world.groups["manipulator"]: PoseStamped( - header=Header(frame_id=""), - pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), - ), - world.groups["wrist"]: PoseStamped( - header=Header(frame_id=""), - pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), - ), + world.groups["manipulator"]: _target_pose(), + world.groups["wrist"]: _target_pose(), }, - seed=JointState(name=["joint_a", "joint_b", "joint_c"], position=[0.0, 0.0, 0.0]), + seed=JointState( + name=["joint_a", "joint_b", "joint_c"], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), max_attempts=1, ) @@ -1658,7 +1781,7 @@ def test_solve_pose_targets_multi_target_uses_multi_frame_solve(mocker: MockerFi assert len(solve_targets.call_args.kwargs["targets"]) == 2 assert result.joint_state is not None assert result.joint_state.name == ["joint_a", "joint_b", "joint_c"] - assert result.joint_state.position == [0.1, 0.2, 0.3] + np.testing.assert_array_equal(result.joint_state.position, [0.1, 0.2, 0.3]) def test_solve_pose_targets_checks_multi_group_solution_together( @@ -1673,7 +1796,11 @@ def test_solve_pose_targets_checks_multi_group_solution_together( return_value=IKResult( status=IKStatus.SUCCESS, joint_state=JointState( - name=["joint_a", "joint_b", "joint_c"], position=[0.1, 0.2, 0.3] + name=["joint_a", "joint_b", "joint_c"], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ), ) @@ -1681,16 +1808,16 @@ def test_solve_pose_targets_checks_multi_group_solution_together( result = ik.solve_pose_targets( world=cast("Any", world), pose_targets={ - world.groups["manipulator"]: PoseStamped( - header=Header(frame_id=""), - pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), - ), - world.groups["wrist"]: PoseStamped( - header=Header(frame_id=""), - pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), - ), + world.groups["manipulator"]: _target_pose(), + world.groups["wrist"]: _target_pose(), }, - seed=JointState(name=["joint_a", "joint_b", "joint_c"], position=[0.0, 0.0, 0.0]), + seed=JointState( + name=["joint_a", "joint_b", "joint_c"], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), max_attempts=1, ) @@ -1709,17 +1836,29 @@ def test_solve_pose_targets_auxiliary_only_retains_seed_selection_order( world=cast("Any", world), pose_targets={}, auxiliary_groups=[world.groups["no_tip"], world.groups["manipulator"]], - seed=JointState(name=["joint_a", "joint_b", "joint_c"], position=[0.1, 0.2, 0.3]), + seed=JointState( + name=["joint_a", "joint_b", "joint_c"], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) assert result.status == IKStatus.SUCCESS assert result.joint_state is not None assert result.joint_state.name == ["joint_c", "joint_a", "joint_b"] - assert result.joint_state.position == [0.3, 0.1, 0.2] + np.testing.assert_array_equal(result.joint_state.position, [0.3, 0.1, 0.2]) def _solved_joint_state() -> JointState: - return JointState(name=["joint_a", "joint_b", "joint_c"], position=[0.1, 0.2, 0.3]) + return JointState( + name=["joint_a", "joint_b", "joint_c"], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) def _qp_infeasible() -> NoSolutionFound: @@ -1753,13 +1892,7 @@ def fake_solve_targets(**_: object) -> IKResult: result = ik.solve( world=cast("Any", _FakeWorld(collision_free=True)), - target_pose=PoseStamped( - header=Header(frame_id=""), - pose=Pose( - position=Point(x=0.1, y=0.0, z=0.0), - orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), - ), - ), + target_pose=_target_pose(x=0.1), check_collision=True, max_attempts=2, ) @@ -1782,13 +1915,7 @@ def fake_solve_targets(**_: object) -> IKResult: result = ik.solve( world=cast("Any", _FakeWorld(collision_free=True)), - target_pose=PoseStamped( - header=Header(frame_id=""), - pose=Pose( - position=Point(x=0.1, y=0.0, z=0.0), - orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), - ), - ), + target_pose=_target_pose(x=0.1), check_collision=True, max_attempts=3, ) @@ -1807,13 +1934,7 @@ def test_solve_does_not_retry_unexpected_exception(mocker: MockerFixture) -> Non result = ik.solve( world=cast("Any", _FakeWorld(collision_free=True)), - target_pose=PoseStamped( - header=Header(frame_id=""), - pose=Pose( - position=Point(x=0.1, y=0.0, z=0.0), - orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), - ), - ), + target_pose=_target_pose(x=0.1), check_collision=True, max_attempts=3, ) @@ -1832,13 +1953,7 @@ def test_solve_mapping_value_error_fails_without_retrying(mocker: MockerFixture) result = ik.solve( world=cast("Any", _FakeWorld(collision_free=True)), - target_pose=PoseStamped( - header=Header(frame_id=""), - pose=Pose( - position=Point(x=0.1, y=0.0, z=0.0), - orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), - ), - ), + target_pose=_target_pose(x=0.1), check_collision=True, max_attempts=5, ) @@ -1849,7 +1964,13 @@ def test_solve_mapping_value_error_fails_without_retrying(mocker: MockerFixture) def _pose_targets_seed() -> JointState: - return JointState(name=["joint_a", "joint_b", "joint_c"], position=[0.0, 0.0, 0.0]) + return JointState( + name=["joint_a", "joint_b", "joint_c"], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) def test_solve_pose_targets_retries_after_no_solution_found(mocker: MockerFixture) -> None: @@ -1876,12 +1997,7 @@ def fake_solve_targets(**_: object) -> IKResult: result = ik.solve_pose_targets( world=cast("Any", world), - pose_targets={ - world.groups["manipulator"]: PoseStamped( - header=Header(frame_id=""), - pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), - ) - }, + pose_targets={world.groups["manipulator"]: _target_pose()}, seed=_pose_targets_seed(), max_attempts=2, ) @@ -1905,12 +2021,7 @@ def fake_solve_targets(**_: object) -> IKResult: result = ik.solve_pose_targets( world=cast("Any", world), - pose_targets={ - world.groups["manipulator"]: PoseStamped( - header=Header(frame_id=""), - pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), - ) - }, + pose_targets={world.groups["manipulator"]: _target_pose()}, seed=_pose_targets_seed(), max_attempts=3, ) @@ -1932,12 +2043,7 @@ def test_solve_pose_targets_does_not_retry_unexpected_exception( result = ik.solve_pose_targets( world=cast("Any", world), - pose_targets={ - world.groups["manipulator"]: PoseStamped( - header=Header(frame_id=""), - pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), - ) - }, + pose_targets={world.groups["manipulator"]: _target_pose()}, seed=_pose_targets_seed(), max_attempts=3, ) @@ -1959,12 +2065,7 @@ def test_solve_pose_targets_mapping_value_error_fails_without_retrying( result = ik.solve_pose_targets( world=cast("Any", world), - pose_targets={ - world.groups["manipulator"]: PoseStamped( - header=Header(frame_id=""), - pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), - ) - }, + pose_targets={world.groups["manipulator"]: _target_pose()}, seed=_pose_targets_seed(), max_attempts=5, ) diff --git a/dimos/manipulation/planning/monitor/robot_state_monitor.py b/dimos/manipulation/planning/monitor/robot_state_monitor.py index 9c8ba52311b..7130566f016 100644 --- a/dimos/manipulation/planning/monitor/robot_state_monitor.py +++ b/dimos/manipulation/planning/monitor/robot_state_monitor.py @@ -26,12 +26,13 @@ from __future__ import annotations +import copy import time from typing import TYPE_CHECKING +from dimos_generated.sensor_msgs.msg import JointState import numpy as np -from dimos.msgs.sensor_msgs.JointState import JointState from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -84,6 +85,7 @@ def __init__( self._timeout = timeout # Latest state + self._latest_state: JointState | None = None self._latest_positions: NDArray[np.float64] | None = None self._latest_velocities: NDArray[np.float64] | None = None self._last_update_time: float | None = None @@ -142,16 +144,19 @@ def on_joint_state(self, msg: JointState) -> None: self._latest_positions = positions self._latest_velocities = velocities self._last_update_time = current_time + joint_state = JointState( + header=copy.deepcopy(msg.header), + name=self._joint_names, + position=np.asarray(positions.tolist(), dtype=np.float64), + velocity=np.asarray( + velocities.tolist() if velocities is not None else [], dtype=np.float64 + ), + effort=np.array([], dtype=np.float64), + ) + self._latest_state = copy.deepcopy(joint_state) # Sync to world's live context (for visualization) try: - # Create JointState for world sync (API uses JointState) - joint_state = JointState( - ts=msg.ts, - name=self._joint_names, - position=positions.tolist(), - velocity=velocities.tolist() if velocities is not None else [], - ) self._world.sync_from_joint_state(joint_state) except Exception as e: logger.error(f"Failed to sync joint state to live context: {e}") @@ -198,7 +203,7 @@ def _extract_velocities(self, msg: JointState) -> NDArray[np.float64] | None: Uses the same canonical-name lookup as _extract_positions. """ - if not msg.velocity or len(msg.velocity) == 0: + if len(msg.velocity) == 0: return None name_to_idx = {name: i for i, name in enumerate(msg.name)} @@ -215,6 +220,14 @@ def _extract_velocities(self, msg: JointState) -> NDArray[np.float64] | None: return np.array(velocities, dtype=np.float64) + def get_current_joint_state(self) -> JointState | None: + """Copy one atomic sample, preserving its source header and joint ordering.""" + with self._lock: + state = self._latest_state + if state is None: + return None + return copy.deepcopy(state) + def get_current_positions(self) -> NDArray[np.float64] | None: """Get current joint positions (thread-safe). diff --git a/dimos/manipulation/planning/monitor/test_robot_state_monitor.py b/dimos/manipulation/planning/monitor/test_robot_state_monitor.py index f22d99c99d4..beed62f443a 100644 --- a/dimos/manipulation/planning/monitor/test_robot_state_monitor.py +++ b/dimos/manipulation/planning/monitor/test_robot_state_monitor.py @@ -12,11 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dataclasses import asdict import threading +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from dimos.manipulation.planning.monitor.robot_state_monitor import RobotStateMonitor @@ -41,31 +44,54 @@ def test_generated_feedback_reorders_joints_and_preserves_source_header(monitor) source = JointState( header=Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id="robot"), name=["arm/b", "other", "arm/a"], - position=[2.0, 9.0, 1.0], - velocity=[0.2, 0.9, 0.1], + position=np.array([2.0, 9.0, 1.0], dtype=np.float64), + velocity=np.array([0.2, 0.9, 0.1], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) - monitor.on_joint_state(JointState.decode(source.encode())) + monitor.on_joint_state(cdr_decode(cdr_encode(source), JointState)) assert monitor.get_current_positions().tolist() == [1.0, 2.0] assert monitor.get_current_velocities().tolist() == [0.1, 0.2] expected = JointState( - header=source.header, name=["arm/a", "arm/b"], position=[1.0, 2.0], velocity=[0.1, 0.2] + header=source.header, + name=["arm/a", "arm/b"], + position=np.array([1.0, 2.0], dtype=np.float64), + velocity=np.array([0.1, 0.2], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) - world.sync_from_joint_state.assert_called_once_with(expected) - assert received == [expected] + world.sync_from_joint_state.assert_called_once() + np.testing.assert_equal(asdict(world.sync_from_joint_state.call_args.args[0]), asdict(expected)) + assert len(received) == 1 + np.testing.assert_equal(asdict(received[0]), asdict(expected)) snapshot = monitor.get_current_joint_state() - assert snapshot == expected + np.testing.assert_equal(asdict(snapshot), asdict(expected)) snapshot.header.stamp.nanosec = 0 snapshot.position[0] = 99 received[0].position[0] = 42 - assert monitor.get_current_joint_state() == expected + np.testing.assert_equal(asdict(monitor.get_current_joint_state()), asdict(expected)) source.header.stamp.nanosec = 0 assert received[0].header.stamp.nanosec == 123456789 def test_missing_joint_does_not_replace_valid_state(monitor): monitor, world = monitor - monitor.on_joint_state(JointState(name=["arm/a", "arm/b"], position=[1.0, 2.0])) - monitor.on_joint_state(JointState(name=["arm/a"], position=[9.0])) + monitor.on_joint_state( + JointState( + name=["arm/a", "arm/b"], + position=np.array([1.0, 2.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) + monitor.on_joint_state( + JointState( + name=["arm/a"], + position=np.array([9.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) assert monitor.get_current_positions().tolist() == [1.0, 2.0] assert world.sync_from_joint_state.call_count == 1 @@ -73,6 +99,14 @@ def test_missing_joint_does_not_replace_valid_state(monitor): def test_stopped_monitor_does_not_sync(monitor): monitor, world = monitor monitor.stop() - monitor.on_joint_state(JointState(name=["arm/a", "arm/b"], position=[1.0, 2.0])) + monitor.on_joint_state( + JointState( + name=["arm/a", "arm/b"], + position=np.array([1.0, 2.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) assert monitor.get_current_positions() is None world.sync_from_joint_state.assert_not_called() diff --git a/dimos/manipulation/planning/monitor/test_world_monitor.py b/dimos/manipulation/planning/monitor/test_world_monitor.py index 96b2e0ae632..f3893907770 100644 --- a/dimos/manipulation/planning/monitor/test_world_monitor.py +++ b/dimos/manipulation/planning/monitor/test_world_monitor.py @@ -15,9 +15,14 @@ from __future__ import annotations from pathlib import Path -from types import SimpleNamespace from typing import Any +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Vector3 +from dimos_generated.sensor_msgs.msg import Image, JointState +from dimos_generated.std_msgs.msg import Header +from dimos_generated.trajectory_msgs.msg import JointTrajectory +from dimos_generated.vision_msgs.msg import BoundingBox3D, Detection3D import numpy as np import pytest from pytest_mock import MockerFixture @@ -37,13 +42,9 @@ ) from dimos.manipulation.planning.spec.protocols import VisualizationSpec from dimos.manipulation.planning.utils import mesh_utils -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.geometry import quaternion_from_matrix +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_to_open3d +from dimos.msgs.time import header_now from dimos.perception.experimental.object import Object from dimos.robot.assets.model import RobotModel @@ -112,9 +113,19 @@ def set_joint_state(self, ctx, joint_state) -> None: def get_joint_state(self, ctx): if self.config is None: - return JointState() + return JointState( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + position=np.array([], dtype=np.float64), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) return JointState( - name=self.config.joint_names, position=[0.0] * len(self.config.joint_names) + name=self.config.joint_names, + position=np.asarray([0.0] * len(self.config.joint_names), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) def is_collision_free(self, ctx): @@ -134,7 +145,13 @@ def get_ee_pose(self, ctx): def get_group_ee_pose(self, ctx, group_id): self.calls.append(("get_group_ee_pose", ctx, group_id)) - return PoseStamped(position=Vector3(1, 2, 3), orientation=Quaternion([0, 0, 0, 1])) + return PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=1, y=2, z=3), + orientation=Quaternion(x=0, y=0, z=0, w=1), + ), + ) def get_link_pose(self, ctx, link_name): return [] @@ -224,7 +241,13 @@ def _robot_config() -> RobotModelConfig: ) return RobotModelConfig( model=RobotModel.from_file(model_path), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion([0, 0, 0, 1])), + base_pose=PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0, y=0, z=0, w=1), + ), + ), joint_names=["j1", "j2"], base_link="base", planning_groups=[ @@ -292,7 +315,7 @@ def test_world_monitor_syncs_planning_scene_to_visualization() -> None: def test_world_monitor_forwards_raw_trajectory_preview_protocol() -> None: fake_viz = FakeViz() monitor = world_monitor_module.WorldMonitor(world=FakeWorld(), visualization=fake_viz) # type: ignore[arg-type] - trajectory = JointTrajectory(joint_names=["j1"], points=[]) + trajectory = JointTrajectory(header=header_now(), joint_names=["j1"], points=[]) assert isinstance(fake_viz, VisualizationSpec) monitor.cancel_preview_animation() @@ -327,7 +350,13 @@ def test_obstacle_monitor_routes_mutations_through_parent_world_monitor( obstacle_monitor = parent.obstacle_monitor assert obstacle_monitor is not None - pose = PoseStamped(position=Vector3(1, 2, 3), orientation=Quaternion([0, 0, 0, 1])) + pose = PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=1, y=2, z=3), + orientation=Quaternion(x=0, y=0, z=0, w=1), + ), + ) obstacle_monitor.on_collision_object( CollisionObjectMessage( id="source-id", @@ -427,12 +456,20 @@ def test_current_group_joint_state_uses_public_names_in_group_order() -> None: monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] monitor.load_model(_three_joint_reordered_group_config()) monitor.start_state_monitor() - monitor.on_joint_state(JointState(name=["j1", "j2", "j3"], position=[0.1, 0.2, 0.3])) + monitor.on_joint_state( + JointState( + name=["j1", "j2", "j3"], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) state = monitor.current_group_joint_state("manipulator") assert state.name == ["j2", "j1"] - assert state.position == [0.2, 0.1] + np.testing.assert_array_equal(state.position, [0.2, 0.1]) def test_current_model_joint_state_rejects_stale_state(mocker) -> None: @@ -440,13 +477,21 @@ def test_current_model_joint_state_rejects_stale_state(mocker) -> None: monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] monitor.load_model(_three_joint_reordered_group_config()) monitor.start_state_monitor() - monitor.on_joint_state(JointState(name=["j1", "j2", "j3"], position=[1.0, 2.0, 3.0])) + monitor.on_joint_state( + JointState( + name=["j1", "j2", "j3"], + position=np.array([1.0, 2.0, 3.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) mocker.patch.object(monitor, "is_state_stale", return_value=True) state = monitor.current_model_joint_state(max_age=0.5) assert state.name == [] - assert state.position == [] + np.testing.assert_array_equal(state.position, []) def test_current_group_joint_state_rejects_stale_state(mocker) -> None: @@ -454,7 +499,15 @@ def test_current_group_joint_state_rejects_stale_state(mocker) -> None: stale_monitor = world_monitor_module.WorldMonitor(world=stale_world) # type: ignore[arg-type] stale_monitor.load_model(_three_joint_reordered_group_config()) stale_monitor.start_state_monitor() - stale_monitor.on_joint_state(JointState(name=["j1", "j2", "j3"], position=[0.1, 0.2, 0.3])) + stale_monitor.on_joint_state( + JointState( + name=["j1", "j2", "j3"], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) mocker.patch.object(stale_monitor, "is_state_stale", return_value=True) with pytest.raises(ValueError, match="stale"): @@ -466,20 +519,28 @@ def test_group_ee_pose_uses_current_state_when_no_joint_state_is_provided() -> N monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] monitor.load_model(_three_joint_reordered_group_config()) monitor.start_state_monitor() - monitor.on_joint_state(JointState(ts=123.5, name=["j1", "j2", "j3"], position=[0.1, 0.2, 0.3])) + monitor.on_joint_state( + JointState( + name=["j1", "j2", "j3"], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=123, nanosec=500000000), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) pose = monitor.get_group_ee_pose("manipulator") set_calls = [call for call in fake_world.calls if call[0] == "set_joint_state"] assert set_calls[0][2].name == ["j1", "j2", "j3"] - assert set_calls[0][2].position == [0.1, 0.2, 0.3] - assert pose.position.x == 1 - assert pose.ts == 123.5 + np.testing.assert_array_equal(set_calls[0][2].position, [0.1, 0.2, 0.3]) + assert pose.pose.position.x == 1 + assert pose.header.stamp == Time(sec=123, nanosec=500000000) snapshot = monitor.get_current_joint_state() assert snapshot is not None snapshot.position[0] = 99.0 - assert monitor.get_current_joint_state().position == [0.1, 0.2, 0.3] + np.testing.assert_array_equal(monitor.get_current_joint_state().position, [0.1, 0.2, 0.3]) monitor.stop_all_monitors() monitor.start_state_monitor() assert monitor.get_current_joint_state() is None @@ -490,7 +551,15 @@ def test_group_ee_pose_without_joint_state_rejects_stale_state(mocker) -> None: stale_monitor = world_monitor_module.WorldMonitor(world=stale_world) # type: ignore[arg-type] stale_monitor.load_model(_three_joint_reordered_group_config()) stale_monitor.start_state_monitor() - stale_monitor.on_joint_state(JointState(name=["j1", "j2", "j3"], position=[0.1, 0.2, 0.3])) + stale_monitor.on_joint_state( + JointState( + name=["j1", "j2", "j3"], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) mocker.patch.object(stale_monitor, "is_state_stale", return_value=True) with pytest.raises(ValueError, match="stale"): @@ -504,13 +573,19 @@ def test_group_kinematics_with_full_state_does_not_require_current_state() -> No pose = monitor.get_group_ee_pose( "manipulator", - JointState(name=["j1", "j2", "j3"], position=[0.1, 0.2, 0.3]), + JointState( + name=["j1", "j2", "j3"], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) set_calls = [call for call in fake_world.calls if call[0] == "set_joint_state"] assert set_calls[0][2].name == ["j1", "j2", "j3"] - assert set_calls[0][2].position == [0.1, 0.2, 0.3] - assert pose.position.x == 1 + np.testing.assert_array_equal(set_calls[0][2].position, [0.1, 0.2, 0.3]) + assert pose.pose.position.x == 1 def test_group_kinematics_route_full_state_to_backend() -> None: @@ -520,19 +595,31 @@ def test_group_kinematics_route_full_state_to_backend() -> None: pose = monitor.get_group_ee_pose( "manipulator", - JointState(name=["j1", "j2", "j3"], position=[0.9, 0.8, 0.3]), + JointState( + name=["j1", "j2", "j3"], + position=np.array([0.9, 0.8, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) jacobian = monitor.get_group_jacobian( "manipulator", - JointState(name=["j1", "j2", "j3"], position=[0.4, 0.3, 0.3]), + JointState( + name=["j1", "j2", "j3"], + position=np.array([0.4, 0.3, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) set_calls = [call for call in fake_world.calls if call[0] == "set_joint_state"] assert set_calls[0][2].name == ["j1", "j2", "j3"] - assert set_calls[0][2].position == [0.9, 0.8, 0.3] + np.testing.assert_array_equal(set_calls[0][2].position, [0.9, 0.8, 0.3]) assert set_calls[1][2].name == ["j1", "j2", "j3"] - assert set_calls[1][2].position == [0.4, 0.3, 0.3] - assert pose.position.x == 1 + np.testing.assert_array_equal(set_calls[1][2].position, [0.4, 0.3, 0.3]) + assert pose.pose.position.x == 1 assert jacobian.shape == (6, 2) assert ("get_group_ee_pose", "scratch", "manipulator") in fake_world.calls assert ("get_group_jacobian", "scratch", "manipulator") in fake_world.calls @@ -547,7 +634,15 @@ def test_convenience_wrappers_fail_for_no_pose_and_ambiguous_pose_groups() -> No ) ) with pytest.raises(ValueError, match="no unique pose-targetable"): - monitor.get_ee_pose(JointState(name=["j1", "j2"], position=[0.0, 0.0])) + monitor.get_ee_pose( + JointState( + name=["j1", "j2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) fake_world2 = FakeWorld() monitor2 = world_monitor_module.WorldMonitor(world=fake_world2) # type: ignore[arg-type] @@ -564,7 +659,15 @@ def test_convenience_wrappers_fail_for_no_pose_and_ambiguous_pose_groups() -> No ) ) with pytest.raises(ValueError, match="pose-targetable planning groups"): - monitor2.get_jacobian(JointState(name=["j1", "j2"], position=[0.0, 0.0])) + monitor2.get_jacobian( + JointState( + name=["j1", "j2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) def test_world_monitor_obstacle_mutations_cover_failure_and_visualization_errors( @@ -596,7 +699,13 @@ def test_world_monitor_updates_obstacle_pose_with_backend_result(mocker: MockerF world = FakeWorld() viz = FakeViz() monitor = world_monitor_module.WorldMonitor(world=world, visualization=viz) # type: ignore[arg-type] - pose = PoseStamped(position=Vector3(1, 2, 3), orientation=Quaternion([0, 0, 0, 1])) + pose = PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=1, y=2, z=3), + orientation=Quaternion(x=0, y=0, z=0, w=1), + ), + ) update = mocker.patch.object(world, "update_obstacle_pose", side_effect=[True, False]) assert monitor.update_obstacle_pose("obstacle-id", pose) is True @@ -617,7 +726,13 @@ def test_world_monitor_forwards_only_successful_complete_updates( obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, - pose=PoseStamped(), + pose=PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), dimensions=(1.0, 1.0, 1.0), ) update = mocker.patch.object(world, "update_obstacle", side_effect=[True, False]) @@ -638,7 +753,13 @@ def test_obstacle_monitor_routes_complete_and_rejects_incomplete_updates( parent.start_obstacle_monitor() monitor = parent.obstacle_monitor assert monitor is not None - pose = PoseStamped(position=Vector3(1, 2, 3), orientation=Quaternion([0, 0, 0, 1])) + pose = PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=1, y=2, z=3), + orientation=Quaternion(x=0, y=0, z=0, w=1), + ), + ) monitor.on_collision_object( CollisionObjectMessage( @@ -687,7 +808,13 @@ def test_obstacle_monitor_adds_complete_unknown_update_and_stops_on_failed_updat parent.start_obstacle_monitor() monitor = parent.obstacle_monitor assert monitor is not None - pose = PoseStamped(position=Vector3(1, 2, 3), orientation=Quaternion([0, 0, 0, 1])) + pose = PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=1, y=2, z=3), + orientation=Quaternion(x=0, y=0, z=0, w=1), + ), + ) monitor.on_collision_object( CollisionObjectMessage( @@ -726,7 +853,13 @@ def test_world_monitor_clear_updates_world_tracking_and_survives_visualization_e monitor.start_obstacle_monitor() obstacle_monitor = monitor.obstacle_monitor assert obstacle_monitor is not None - pose = PoseStamped(position=Vector3(1, 2, 3), orientation=Quaternion([0, 0, 0, 1])) + pose = PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=1, y=2, z=3), + orientation=Quaternion(x=0, y=0, z=0, w=1), + ), + ) obstacle_monitor.on_collision_object( CollisionObjectMessage( id="source", @@ -736,12 +869,17 @@ def test_world_monitor_clear_updates_world_tracking_and_survives_visualization_e dimensions=(1.0, 2.0, 3.0), ) ) - detection = SimpleNamespace( + detection = Detection3D( id="detection", - bbox=SimpleNamespace( - center=SimpleNamespace(position=Vector3(4, 5, 6), orientation=Quaternion([0, 0, 0, 1])), - size=Vector3(1, 1, 1), + bbox=BoundingBox3D( + center=Pose( + position=Point(x=4, y=5, z=6), + orientation=Quaternion(x=0, y=0, z=0, w=1), + ), + size=Vector3(x=1, y=1, z=1), ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + results=[], ) obstacle_monitor.on_detections([detection]) # type: ignore[arg-type] assert obstacle_monitor.get_obstacle_count() == 2 @@ -803,7 +941,13 @@ def test_world_obstacle_monitor_rejects_invalid_add_and_handles_update_and_callb remove = mocker.patch.object(parent, "remove_obstacle", return_value=True) callback = mocker.Mock(side_effect=RuntimeError("callback failed")) obstacle_monitor.add_obstacle_callback(callback) - pose = PoseStamped(position=Vector3(1, 2, 3), orientation=Quaternion([0, 0, 0, 1])) + pose = PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=1, y=2, z=3), + orientation=Quaternion(x=0, y=0, z=0, w=1), + ), + ) obstacle_monitor.on_collision_object(CollisionObjectMessage(id="bad", operation="add")) assert obstacle_monitor.get_obstacle_count() == 0 @@ -838,17 +982,21 @@ def test_world_obstacle_monitor_detection_add_update_and_stale_cleanup( add = mocker.patch.object(parent, "add_obstacle", return_value="first-id") update = mocker.patch.object(parent, "update_obstacle_pose") remove = mocker.patch.object(parent, "remove_obstacle", return_value=False) - timestamps = iter([1.0, 1.0, 1.0, 1.0, 1.0, 10.0]) - monkeypatch.setattr( - "dimos.manipulation.planning.monitor.world_obstacle_monitor.time.time", - lambda: next(timestamps, 10.0), + clock = mocker.patch( + "dimos.manipulation.planning.monitor.world_obstacle_monitor.time.time", return_value=1.0 + ) + center = Pose( + position=Point(x=1, y=2, z=3), + orientation=Quaternion(x=0, y=0, z=0, w=1), + ) + bbox = BoundingBox3D(center=center, size=Vector3(x=1, y=2, z=3)) + detection = Detection3D( + id="det-1", bbox=bbox, header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), results=[] ) - center = SimpleNamespace(position=Vector3(1, 2, 3), orientation=Quaternion([0, 0, 0, 1])) - bbox = SimpleNamespace(center=center, size=Vector3(1, 2, 3)) - detection = SimpleNamespace(id="det-1", bbox=bbox) obstacle_monitor.on_detections([detection]) # type: ignore[arg-type] - obstacle_monitor.on_detections([detection]) # type: ignore[arg-type] + obstacle_monitor.on_detections([detection]) + clock.return_value = 10.0 obstacle_monitor.on_detections([]) assert add.call_count == 1 @@ -887,19 +1035,23 @@ def test_mesh_obstacle_is_placed_at_the_hull_centroid_without_the_bbox_rotation( add_obstacle = mocker.patch.object(parent, "add_obstacle", return_value="parent-id") points = _tilted_box_cloud() - cloud = PointCloud2.from_numpy(points, frame_id="world") + cloud = pointcloud_from_xyz( + points, header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)) + ) # Mirrors Object.from_detections with use_aabb=False: pose carries the # oriented-box center and rotation. - obb = cloud.pointcloud.get_oriented_bounding_box() + obb = pointcloud_to_open3d(cloud).get_oriented_bounding_box() obj = Object( object_id="tilted-box", name="box", - center=Vector3(obb.center), - size=Vector3(obb.extent), + center=Vector3(x=obb.center[0], y=obb.center[1], z=obb.center[2]), + size=Vector3(x=obb.extent[0], y=obb.extent[1], z=obb.extent[2]), pose=PoseStamped( - frame_id="world", - position=Vector3(obb.center), - orientation=Quaternion.from_rotation_matrix(np.asarray(obb.R)), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=obb.center[0], y=obb.center[1], z=obb.center[2]), + orientation=quaternion_from_matrix(np.asarray(obb.R)), + ), ), pointcloud=cloud, bbox=(0.0, 0.0, 1.0, 1.0), @@ -907,7 +1059,15 @@ def test_mesh_obstacle_is_placed_at_the_hull_centroid_without_the_bbox_rotation( class_id=0, confidence=1.0, ts=0.0, - image=Image(), + image=Image( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + height=0, + width=0, + encoding="", + is_bigendian=0, + step=0, + data=np.array([], dtype=np.uint8), + ), ) monitor = WorldObstacleMonitor(parent=parent, use_mesh_obstacles=True) @@ -919,23 +1079,71 @@ def test_mesh_obstacle_is_placed_at_the_hull_centroid_without_the_bbox_rotation( centroid = points.mean(axis=0) assert obstacle.obstacle_type == ObstacleType.MESH np.testing.assert_allclose( - [obstacle.pose.position.x, obstacle.pose.position.y, obstacle.pose.position.z], + [ + obstacle.pose.pose.position.x, + obstacle.pose.pose.position.y, + obstacle.pose.pose.position.z, + ], centroid, atol=1e-3, ) np.testing.assert_allclose( [ - obstacle.pose.orientation.x, - obstacle.pose.orientation.y, - obstacle.pose.orientation.z, - obstacle.pose.orientation.w, + obstacle.pose.pose.orientation.x, + obstacle.pose.pose.orientation.y, + obstacle.pose.pose.orientation.z, + obstacle.pose.pose.orientation.w, ], [0.0, 0.0, 0.0, 1.0], atol=1e-9, ) - assert obstacle.pose.frame_id == "world" + assert obstacle.pose.header.frame_id == "world" # Teeth: the pose the old code used is a genuinely different placement, so # this cannot pass on the bug. assert np.linalg.norm(np.asarray(obb.center) - centroid) > 3e-3 - assert abs(Quaternion.from_rotation_matrix(np.asarray(obb.R)).w) < 0.999 + assert abs(quaternion_from_matrix(np.asarray(obb.R)).w) < 0.999 + + +@pytest.mark.parametrize("use_current_state", [False, True]) +def test_group_pose_preserves_joint_sample_stamp_and_backend_frame(use_current_state): + world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=world) + monitor.load_model(_three_joint_reordered_group_config()) + state = JointState( + header=Header(frame_id="encoder", stamp=Time(sec=1700000000, nanosec=123456789)), + name=["j1", "j2", "j3"], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + if use_current_state: + monitor.start_state_monitor() + monitor.on_joint_state(state) + pose = monitor.get_group_ee_pose("manipulator", None if use_current_state else state) + assert pose.header.stamp == state.header.stamp + assert pose.header.frame_id == "" + pose.header.stamp.nanosec = 0 + assert state.header.stamp.nanosec == 123456789 + + +def test_link_pose_preserves_joint_sample_stamp(mocker): + world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=world) + monitor.load_model(_three_joint_reordered_group_config()) + matrix = np.eye(4) + matrix[:3, 3] = [1, 2, 3] + mocker.patch.object(world, "get_link_pose", return_value=matrix) + state = JointState( + header=Header(stamp=Time(sec=-1, nanosec=987654321), frame_id=""), + name=["j1", "j2", "j3"], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + pose = monitor.get_link_pose("ee", state) + assert pose is not None + assert pose.header == Header(frame_id="world", stamp=state.header.stamp) + np.testing.assert_array_equal(pose.pose.position, Point(x=1, y=2, z=3)) + pose.header.stamp.nanosec = 0 + assert state.header.stamp.nanosec == 987654321 diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index 4aed33d44cf..be040c575c7 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -17,9 +17,18 @@ from __future__ import annotations from contextlib import contextmanager +from copy import deepcopy import threading from typing import TYPE_CHECKING, Any +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.dimos_msgs.msg import GraspCandidateArray +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header +from dimos_generated.trajectory_msgs.msg import JointTrajectory +import numpy as np + from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.manipulation.planning.groups.identifiers import assert_valid_group_id from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry @@ -33,17 +42,13 @@ ) from dimos.manipulation.planning.spec.protocols import VisualizationSpec, WorldSpec from dimos.manipulation.planning.spec.validation import PreparedRobotModel, prepare_robot_model -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.geometry import quaternion_from_matrix from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: from collections.abc import Generator - import numpy as np + from dimos_generated.vision_msgs.msg import Detection3D from numpy.typing import NDArray from dimos.manipulation.planning.spec.config import RobotModelConfig @@ -53,7 +58,6 @@ Obstacle, PlanningGroupID, ) - from dimos.msgs.vision_msgs.Detection3D import Detection3D from dimos.perception.experimental.object import Object logger = setup_logger() @@ -215,7 +219,7 @@ def start_state_monitor(self) -> None: def record_state(state: JointState) -> None: # RobotStateMonitor invokes callbacks under our shared lock. - self._latest_joint_state = JointState(state) + self._latest_joint_state = deepcopy(state) monitor.add_state_callback(record_state) monitor.start() @@ -275,7 +279,7 @@ def on_collision_object(self, msg: CollisionObjectMessage) -> None: self._obstacle_monitor.on_collision_object(msg) def on_detections(self, detections: list[Detection3D]) -> None: - """Handle perception detections (Detection3D from dimos.msgs.vision_msgs).""" + """Handle perception detections (generated Detection3D).""" if self._obstacle_monitor is not None: self._obstacle_monitor.on_detections(detections) @@ -332,7 +336,7 @@ def get_current_joint_state(self) -> JointState | None: """Get current joint state. Returns None if not yet received.""" with self._lock: if self._state_monitor is not None: - return JointState(self._latest_joint_state) if self._latest_joint_state else None + return deepcopy(self._latest_joint_state) if self._latest_joint_state else None # Fall back to world's live context when no monitor is running. ctx = self._world.get_live_context() return self._world.get_joint_state(ctx) @@ -340,8 +344,20 @@ def get_current_joint_state(self) -> JointState | None: def current_model_joint_state(self, max_age: float = 1.0) -> JointState: """Return the fresh canonical model state.""" if self._state_monitor is not None and self.is_state_stale(max_age): - return JointState() - return self.get_current_joint_state() or JointState() + return JointState( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + position=np.array([], dtype=np.float64), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + return self.get_current_joint_state() or JointState( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + position=np.array([], dtype=np.float64), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) def current_group_joint_state( self, group_id: PlanningGroupID, max_age: float = 1.0 @@ -375,10 +391,14 @@ def _validate_planning_group_config(self, config: RobotModelConfig) -> None: def get_current_velocities(self) -> JointState | None: """Get current joint velocities as JointState. Returns None if not available.""" if self._state_monitor is not None: - velocities = self._state_monitor.get_current_velocities() - if velocities is not None: + state = self._state_monitor.get_current_joint_state() + if state is not None and len(state.velocity): return JointState( - name=self.get_model_config().joint_names, velocity=velocities.tolist() + header=deepcopy(state.header), + name=state.name, + velocity=np.asarray(state.velocity, dtype=np.float64), + position=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) return None @@ -461,8 +481,12 @@ def get_group_ee_pose( self._world.set_joint_state(ctx, joint_state) pose = self._world.get_group_ee_pose(ctx, group_id) - pose.ts = joint_state.ts - return pose + return PoseStamped( + header=Header( + frame_id=pose.header.frame_id, stamp=deepcopy(joint_state.header.stamp) + ), + pose=deepcopy(pose.pose), + ) def get_link_pose( self, link_name: str, joint_state: JointState | None = None @@ -487,12 +511,18 @@ def get_link_pose( pos = mat[:3, 3] rot = mat[:3, :3] - quat = Quaternion.from_rotation_matrix(rot) + quat = quaternion_from_matrix(rot) + header = Header(frame_id="world", stamp=Time(sec=0, nanosec=0)) + if joint_state is not None: + header.stamp = deepcopy(joint_state.header.stamp) return PoseStamped( - ts=joint_state.ts, - frame_id="world", - position=[float(pos[0]), float(pos[1]), float(pos[2])], - orientation=[float(quat.x), float(quat.y), float(quat.z), float(quat.w)], + header=header, + pose=Pose( + position=Point(x=float(pos[0]), y=float(pos[1]), z=float(pos[2])), + orientation=Quaternion( + x=float(quat.x), y=float(quat.y), z=float(quat.z), w=float(quat.w) + ), + ), ) def get_jacobian(self, joint_state: JointState) -> NDArray[np.float64]: diff --git a/dimos/manipulation/planning/planners/roboplan_planner.py b/dimos/manipulation/planning/planners/roboplan_planner.py index ecdd485859b..35e6e727796 100644 --- a/dimos/manipulation/planning/planners/roboplan_planner.py +++ b/dimos/manipulation/planning/planners/roboplan_planner.py @@ -21,6 +21,8 @@ import time from typing import TYPE_CHECKING, Any +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.std_msgs.msg import Header import numpy as np from dimos.msgs.geometry import transform_matrix @@ -135,7 +137,13 @@ def plan_joint_path( with self._world.scratch_context() as ctx: self._world.set_joint_state( ctx, - JointState(name=list(config.joint_names), position=q_start.tolist()), + JointState( + name=list(config.joint_names), + position=np.asarray(q_start.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) return self._plan_group( ctx, @@ -536,12 +544,22 @@ def _path_from_cartesian_trajectory( path.append( JointState( name=expected_names, - position=[ - float(position_by_native[native_by_public[name]]) for name in expected_names - ], - velocity=[ - float(velocity_by_native[native_by_public[name]]) for name in expected_names - ], + position=np.asarray( + [ + float(position_by_native[native_by_public[name]]) + for name in expected_names + ], + dtype=np.float64, + ), + velocity=np.asarray( + [ + float(velocity_by_native[native_by_public[name]]) + for name in expected_names + ], + dtype=np.float64, + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + effort=np.array([], dtype=np.float64), ) ) return path, timestamps @@ -559,7 +577,12 @@ def _combined_path_collision_free( for fraction in np.linspace(0.0, 1.0, steps + 1): sample = JointState( name=start.name, - position=(q_start + fraction * (q_end - q_start)).tolist(), + position=np.asarray( + (q_start + fraction * (q_end - q_start)).tolist(), dtype=np.float64 + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) self._apply_selected_state(ctx, sample) if not self._world.is_collision_free(ctx): @@ -739,7 +762,12 @@ def _path_from_native( path.append( JointState( name=list(ordered_output_names), - position=[float(positions[name]) for name in ordered_output_names], + position=np.asarray( + [float(positions[name]) for name in ordered_output_names], dtype=np.float64 + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) ) return path diff --git a/dimos/manipulation/planning/planners/rrt_planner.py b/dimos/manipulation/planning/planners/rrt_planner.py index 8be423a13eb..7193994bdf6 100644 --- a/dimos/manipulation/planning/planners/rrt_planner.py +++ b/dimos/manipulation/planning/planners/rrt_planner.py @@ -25,7 +25,9 @@ import time from typing import TYPE_CHECKING +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header import numpy as np from dimos.manipulation.planning.groups.models import PlanningGroupSelection @@ -471,8 +473,20 @@ def _extend_tree( new_config = nearest.config + step_size * (diff / dist) # Check validity of edge using context-free method - start_state = JointState(name=joint_names, position=nearest.config.tolist()) - end_state = JointState(name=joint_names, position=new_config.tolist()) + start_state = JointState( + name=joint_names, + position=np.asarray(nearest.config.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + end_state = JointState( + name=joint_names, + position=np.asarray(new_config.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) if world.check_edge_collision_free(start_state, end_state, self._collision_step_size): new_node = TreeNode(config=new_config, parent=nearest) nearest.children.append(new_node) @@ -519,7 +533,16 @@ def _extract_path( full_path_arrays = start_path + list(reversed(goal_path)) # Convert to list of JointState - return [JointState(name=joint_names, position=q.tolist()) for q in full_path_arrays] + return [ + JointState( + name=joint_names, + position=np.asarray(q.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + for q in full_path_arrays + ] def _simplify_path( self, diff --git a/dimos/manipulation/planning/planners/selected_joint_space.py b/dimos/manipulation/planning/planners/selected_joint_space.py index ff792ae0f84..9aa31ce8bf4 100644 --- a/dimos/manipulation/planning/planners/selected_joint_space.py +++ b/dimos/manipulation/planning/planners/selected_joint_space.py @@ -18,7 +18,9 @@ from dataclasses import dataclass +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header import numpy as np from numpy.typing import NDArray @@ -58,7 +60,13 @@ def lift_path(self, path: JointPath) -> JointPath: configurations = [self.joint_space.from_joint_state(state) for state in path] positions = self.joint_space.lifted_positions(configurations) return [ - JointState(name=list(self.selected_joint_names), position=list(values)) + JointState( + name=list(self.selected_joint_names), + position=np.asarray(list(values), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) for values in positions ] @@ -69,7 +77,13 @@ def project_config(self, selected_positions: NDArray[np.float64]) -> JointState: indices = {name: i for i, name in enumerate(self.model_joint_names)} for name, value in zip(self.selected_joint_names, selected_positions, strict=True): positions[indices[name]] = value - return JointState(name=list(self.model_joint_names), position=positions.tolist()) + return JointState( + name=list(self.model_joint_names), + position=np.asarray(positions.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) def config_collision_free( self, world: WorldSpec, selected_positions: NDArray[np.float64] @@ -125,7 +139,13 @@ def normalize_selection_target( raise ValueError( f"{label} target has {len(target.position)} positions, expected {len(names)}" ) - return JointState(name=names, position=list(target.position)) + return JointState( + name=names, + position=np.asarray(list(target.position), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) positions = dict(zip(target.name, target.position, strict=True)) missing = [name for name in names if name not in positions] extra = sorted(set(positions) - set(names)) @@ -133,7 +153,13 @@ def normalize_selection_target( raise ValueError(f"{label} target is missing joints: {missing}") if extra: raise ValueError(f"{label} target has extra joints: {extra}") - return JointState(name=names, position=[float(positions[name]) for name in names]) + return JointState( + name=names, + position=np.asarray([float(positions[name]) for name in names], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) def _ordered_positions(state: JointState, names: list[str], label: str) -> NDArray[np.float64]: diff --git a/dimos/manipulation/planning/planners/test_rrt_planner_selection.py b/dimos/manipulation/planning/planners/test_rrt_planner_selection.py index fe5c85c0d92..f6858cc249f 100644 --- a/dimos/manipulation/planning/planners/test_rrt_planner_selection.py +++ b/dimos/manipulation/planning/planners/test_rrt_planner_selection.py @@ -20,6 +20,7 @@ import math from pathlib import Path +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, Pose, @@ -56,8 +57,10 @@ def _pose() -> PoseStamped: return PoseStamped( - header=Header(frame_id=""), - pose=Pose(position=Point(), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), ) @@ -112,7 +115,13 @@ def scratch_context(self) -> nullcontext[None]: return nullcontext(None) def get_joint_state(self, ctx: object) -> JointState: - return JointState(name=["arm/joint_a", "arm/joint_b", "arm/gripper"], position=self.current) + return JointState( + name=["arm/joint_a", "arm/joint_b", "arm/gripper"], + position=np.asarray(self.current, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) def get_joint_limits(self) -> tuple[np.ndarray, np.ndarray]: return np.array([-1.0, -1.0, -1.0]), np.array([1.0, 1.0, 1.0]) @@ -191,7 +200,13 @@ def get_joint_velocity_limits(self) -> np.ndarray: return np.array([2.0, 1.0, 4.0]) def get_joint_state(self, ctx: object) -> JointState: - return JointState(name=list(self.config.joint_names), position=self.current) + return JointState( + name=list(self.config.joint_names), + position=np.asarray(self.current, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) class _WallPlanarWorld(_PlanarWorld): @@ -204,14 +219,38 @@ def check_config_collision_free(self, joint_state: JointState) -> bool: ("start", "goal", "expected_start", "expected_goal"), [ ( - JointState(position=[0.1, 0.2]), - JointState(position=[0.3, 0.4]), + JointState( + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + position=np.array([0.3, 0.4], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), [0.1, 0.2], [0.3, 0.4], ), ( - JointState(name=["arm/joint_b", "arm/joint_a"], position=[0.2, 0.1]), - JointState(name=["arm/joint_b", "arm/joint_a"], position=[0.4, 0.3]), + JointState( + name=["arm/joint_b", "arm/joint_a"], + position=np.array([0.2, 0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=["arm/joint_b", "arm/joint_a"], + position=np.array([0.4, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), [0.1, 0.2], [0.3, 0.4], ), @@ -228,28 +267,64 @@ def test_plan_selected_joint_path_normalizes_target_forms( assert result.status == PlanningStatus.SUCCESS assert result.path is not None assert result.path[0].name == ["arm/joint_a", "arm/joint_b"] - assert result.path[0].position == expected_start - assert result.path[-1].position == expected_goal + np.testing.assert_array_equal(result.path[0].position, expected_start) + np.testing.assert_array_equal(result.path[-1].position, expected_goal) @pytest.mark.parametrize( ("start", "goal", "status", "message"), [ ( - JointState(name=["arm/joint_a"], position=[0.0]), - JointState(position=[0.0, 0.0]), + JointState( + name=["arm/joint_a"], + position=np.array([0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), PlanningStatus.INVALID_START, "missing", ), ( - JointState(position=[0.0, 0.0]), - JointState(name=["arm/joint_a", "arm/joint_b", "arm/extra"], position=[0.0, 0.0, 0.0]), + JointState( + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=["arm/joint_a", "arm/joint_b", "arm/extra"], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), PlanningStatus.INVALID_GOAL, "extra", ), ( - JointState(name=["arm/joint_a", "joint_b"], position=[0.0, 0.0]), - JointState(position=[0.0, 0.0]), + JointState( + name=["arm/joint_a", "joint_b"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), PlanningStatus.INVALID_START, "missing", ), @@ -275,8 +350,20 @@ def test_plan_selected_joint_path_rejects_noncanonical_names() -> None: result = RRTConnectPlanner().plan_selected_joint_path( _World(), selection, - JointState(name=["joint_a", "gripper"], position=[0.0, 0.0]), - JointState(position=[0.1, 0.2]), + JointState( + name=["joint_a", "gripper"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) assert result.status == PlanningStatus.INVALID_START @@ -290,8 +377,20 @@ def test_plan_selected_joint_path_direct_edge_projects_full_state_with_unselecte result = RRTConnectPlanner().plan_selected_joint_path( world, PlanningGroupSelection.from_groups((group,)), - JointState(position=[0.1, 0.2]), - JointState(position=[0.3, 0.4]), + JointState( + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + position=np.array([0.3, 0.4], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) assert result.status == PlanningStatus.SUCCESS @@ -310,13 +409,29 @@ def test_plan_cartesian_path_is_explicitly_unsupported() -> None: result = RRTConnectPlanner().plan_cartesian_path( _World(), selection, - JointState(position=[0.0, 0.0]), + JointState( + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), { group.id: ( - TransformStamped(header=Header(frame_id="world"), child_frame_id=""), TransformStamped( - header=Header(frame_id="world"), - transform=Transform(translation=Vector3(x=0.1, y=0.0, z=0.0)), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + child_frame_id="", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + TransformStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + transform=Transform( + translation=Vector3(x=0.1, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), child_frame_id="", ), ) @@ -352,8 +467,20 @@ def test_planar_direct_path_lifts_yaw_across_wrap_boundary() -> None: result = RRTConnectPlanner().plan_selected_joint_path( world, PlanningGroupSelection.from_groups((group,)), - JointState(position=[0.0, 0.0, math.pi - 0.1]), - JointState(position=[0.0, 0.0, -math.pi + 0.1]), + JointState( + position=np.array([0.0, 0.0, math.pi - 0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + position=np.array([0.0, 0.0, -math.pi + 0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) assert result.status == PlanningStatus.SUCCESS @@ -377,8 +504,20 @@ def test_planar_rrt_finds_obstacle_detour( ).plan_selected_joint_path( world, PlanningGroupSelection.from_groups((group,)), - JointState(position=[-2.0, 0.0, 0.0]), - JointState(position=[2.0, 0.0, 0.0]), + JointState( + position=np.array([-2.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + position=np.array([2.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), timeout=10.0, max_iterations=2000, ) @@ -418,8 +557,20 @@ def exhaust_connection_budget(*args: object) -> tuple[None, int]: result = planner.plan_selected_joint_path( world, PlanningGroupSelection.from_groups((group,)), - JointState(position=[-2.0, 0.0, 0.0]), - JointState(position=[2.0, 0.0, 0.0]), + JointState( + position=np.array([-2.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + position=np.array([2.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), max_iterations=max_iterations, ) diff --git a/dimos/manipulation/planning/spec/config.py b/dimos/manipulation/planning/spec/config.py index dfa6aa8d9aa..5fcc15d398c 100644 --- a/dimos/manipulation/planning/spec/config.py +++ b/dimos/manipulation/planning/spec/config.py @@ -19,13 +19,15 @@ from pathlib import Path from typing import Any +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.std_msgs.msg import Header from pydantic import Field, field_validator from dimos.core.module import ModuleConfig from dimos.manipulation.grasp_verification import GraspVerificationConfig from dimos.manipulation.planning.groups.identifiers import assert_valid_joint_names from dimos.manipulation.planning.groups.models import PlanningGroupDefinition -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.robot.assets.model import RobotModel @@ -48,7 +50,15 @@ class RobotModelConfig(ModuleConfig): model: RobotModel srdf_path: Path | None = None - base_pose: PoseStamped = Field(default_factory=PoseStamped) + base_pose: PoseStamped = Field( + default_factory=lambda: PoseStamped( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ) joint_names: list[str] base_link: str = "base_link" planning_groups: list[PlanningGroupDefinition] = Field(default_factory=list) diff --git a/dimos/manipulation/planning/spec/test_joint_space.py b/dimos/manipulation/planning/spec/test_joint_space.py index 4b3c5563f81..06adae3739c 100644 --- a/dimos/manipulation/planning/spec/test_joint_space.py +++ b/dimos/manipulation/planning/spec/test_joint_space.py @@ -18,7 +18,9 @@ from pathlib import Path from typing import Any +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest @@ -160,7 +162,15 @@ def test_joint_space_wraps_interpolates_lifts_and_builds_finite_domains( base_link="base", ) ).joint_space - start = space.from_joint_state(JointState(name=["yaw", "x"], position=[math.pi - 0.1, 10.0])) + start = space.from_joint_state( + JointState( + name=["yaw", "x"], + position=np.array([math.pi - 0.1, 10.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) goal = space.normalize_positions([-4.0, -math.pi + 0.1]) assert space.delta(start, goal) == pytest.approx((-14.0, 0.2)) @@ -195,4 +205,12 @@ def test_joint_space_rejects_invalid_input_states(tmp_path, names, positions, me ).joint_space with pytest.raises(ValueError, match=message): - space.from_joint_state(JointState(name=names, position=positions)) + space.from_joint_state( + JointState( + name=names, + position=np.asarray(positions, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) diff --git a/dimos/manipulation/planning/trajectory_generator/joint_trajectory_generator.py b/dimos/manipulation/planning/trajectory_generator/joint_trajectory_generator.py index 81226240591..2b27328ec38 100644 --- a/dimos/manipulation/planning/trajectory_generator/joint_trajectory_generator.py +++ b/dimos/manipulation/planning/trajectory_generator/joint_trajectory_generator.py @@ -33,6 +33,7 @@ import math from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import numpy as np from dimos.msgs.time import duration_from_seconds, header_now, to_seconds from dimos.msgs.trajectory import trajectory_duration @@ -121,8 +122,10 @@ def generate(self, waypoints: list[list[float]]) -> JointTrajectory: all_points.append( JointTrajectoryPoint( time_from_start=duration_from_seconds(0.0), - positions=list(waypoints[0]), - velocities=[0.0] * self.num_joints, + positions=np.asarray(list(waypoints[0]), dtype=np.float64), + velocities=np.asarray([0.0] * self.num_joints, dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) ) @@ -138,7 +141,7 @@ def generate(self, waypoints: list[list[float]]) -> JointTrajectory: all_points.extend(segment_points[1:]) current_time += segment_duration - return JointTrajectory(header=header_now(), points=all_points) + return JointTrajectory(header=header_now(), points=all_points, joint_names=[]) def _generate_segment( self, @@ -208,8 +211,10 @@ def _generate_segment( points.append( JointTrajectoryPoint( time_from_start=duration_from_seconds(t), - positions=positions, - velocities=velocities, + positions=np.asarray(positions, dtype=np.float64), + velocities=np.asarray(velocities, dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) ) diff --git a/dimos/manipulation/planning/trajectory_generator/parametrizer.py b/dimos/manipulation/planning/trajectory_generator/parametrizer.py index 6cc100c94ba..c0693632efc 100644 --- a/dimos/manipulation/planning/trajectory_generator/parametrizer.py +++ b/dimos/manipulation/planning/trajectory_generator/parametrizer.py @@ -21,6 +21,7 @@ from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import numpy as np from dimos.manipulation.planning.groups.models import PlanningGroupSelection from dimos.manipulation.planning.spec.models import GeneratedPlan, PlanningResult @@ -139,8 +140,10 @@ def _timed_trajectory( points.append( JointTrajectoryPoint( time_from_start=duration_from_seconds(float(timestamp)), - positions=list(state.position), - velocities=velocities, + positions=np.asarray(list(state.position), dtype=np.float64), + velocities=np.asarray(velocities, dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) ) return JointTrajectory( diff --git a/dimos/manipulation/planning/trajectory_generator/roboplan_toppra_parametrizer.py b/dimos/manipulation/planning/trajectory_generator/roboplan_toppra_parametrizer.py index 061f9d57496..768e89c6577 100644 --- a/dimos/manipulation/planning/trajectory_generator/roboplan_toppra_parametrizer.py +++ b/dimos/manipulation/planning/trajectory_generator/roboplan_toppra_parametrizer.py @@ -233,8 +233,12 @@ def _canonical_result( points = [ JointTrajectoryPoint( time_from_start=duration_from_seconds(time), - positions=list(position), - velocities=[float(velocity[index]) for index in output_indices], + positions=np.asarray(list(position), dtype=np.float64), + velocities=np.asarray( + [float(velocity[index]) for index in output_indices], dtype=np.float64 + ), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) for time, position, velocity in zip( times, diff --git a/dimos/manipulation/planning/trajectory_generator/test_parametrizer.py b/dimos/manipulation/planning/trajectory_generator/test_parametrizer.py index 7c53e113f86..0b91ade4a7d 100644 --- a/dimos/manipulation/planning/trajectory_generator/test_parametrizer.py +++ b/dimos/manipulation/planning/trajectory_generator/test_parametrizer.py @@ -16,8 +16,11 @@ from unittest.mock import MagicMock +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import numpy as np import pytest from dimos.manipulation.planning.groups.models import ( @@ -65,9 +68,27 @@ def _selection() -> PlanningGroupSelection: def _path() -> list[JointState]: names = ["arm/a", "arm/b"] return [ - JointState(name=names, position=[0.0, 0.0]), - JointState(name=names, position=[0.2, 0.1]), - JointState(name=names, position=[0.4, 0.0]), + JointState( + name=names, + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=names, + position=np.array([0.2, 0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=names, + position=np.array([0.4, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ] @@ -78,13 +99,17 @@ def _output() -> JointTrajectory: points=[ JointTrajectoryPoint( time_from_start=duration_from_seconds(0.0), - positions=[0.0, 0.0], - velocities=[0.0, 0.0], + positions=np.array([0.0, 0.0], dtype=np.float64), + velocities=np.array([0.0, 0.0], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), JointTrajectoryPoint( time_from_start=duration_from_seconds(0.5), - positions=[0.4, 0.0], - velocities=[0.0, 0.0], + positions=np.array([0.4, 0.0], dtype=np.float64), + velocities=np.array([0.0, 0.0], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ], ) @@ -106,11 +131,14 @@ def test_materializes_trajectory_and_preserves_source_path() -> None: speed_scale=0.4, ) - assert [state.position for state in plan.path] == [ - [0.0, 0.0], - [0.2, 0.1], - [0.4, 0.0], - ] + np.testing.assert_array_equal( + [state.position for state in plan.path], + [ + [0.0, 0.0], + [0.2, 0.1], + [0.4, 0.0], + ], + ) assert plan.path is not source_path assert plan.trajectory is parametrizer.output assert plan.planning_time == 0.2 @@ -123,13 +151,17 @@ def test_timed_planner_result_bypasses_backend_path_conversion() -> None: path = [ JointState( name=["arm/a", "arm/b"], - position=[0.0, 0.0], - velocity=[0.0, 0.0], + position=np.array([0.0, 0.0], dtype=np.float64), + velocity=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + effort=np.array([], dtype=np.float64), ), JointState( name=["arm/a", "arm/b"], - position=[0.4, 0.0], - velocity=[0.3, 0.0], + position=np.array([0.4, 0.0], dtype=np.float64), + velocity=np.array([0.3, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + effort=np.array([], dtype=np.float64), ), ] @@ -148,7 +180,7 @@ def test_timed_planner_result_bypasses_backend_path_conversion() -> None: 0.0, 0.75, ] - assert plan.trajectory.points[-1].velocities == [0.3, 0.0] + np.testing.assert_array_equal(plan.trajectory.points[-1].velocities, [0.3, 0.0]) def test_timed_planner_result_requires_velocity_for_each_joint() -> None: @@ -202,22 +234,58 @@ def test_rejects_backend_trajectory_that_changes_path_goal() -> None: [ ( [ - JointState(name=["arm/a", "arm/b"], position=[0.0, 0.0]), - JointState(name=["wrong/a", "wrong/b"], position=[0.4, 0.0]), + JointState( + name=["arm/a", "arm/b"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=["wrong/a", "wrong/b"], + position=np.array([0.4, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ], "joint names", ), ( [ - JointState(name=["arm/a", "arm/b"], position=[0.0, 0.0]), - JointState(name=["arm/a", "arm/b"], position=[0.4]), + JointState( + name=["arm/a", "arm/b"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=["arm/a", "arm/b"], + position=np.array([0.4], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ], "dimension", ), ( [ - JointState(name=["arm/a", "arm/b"], position=[0.0, 0.0]), - JointState(name=["arm/a", "arm/b"], position=[float("nan"), 0.0]), + JointState( + name=["arm/a", "arm/b"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=["arm/a", "arm/b"], + position=np.array([float("nan"), 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ], "non-finite", ), diff --git a/dimos/manipulation/planning/trajectory_generator/test_roboplan_toppra_parametrizer.py b/dimos/manipulation/planning/trajectory_generator/test_roboplan_toppra_parametrizer.py index 2359a76e8c7..ae98eca360e 100644 --- a/dimos/manipulation/planning/trajectory_generator/test_roboplan_toppra_parametrizer.py +++ b/dimos/manipulation/planning/trajectory_generator/test_roboplan_toppra_parametrizer.py @@ -21,6 +21,8 @@ from pathlib import Path from types import SimpleNamespace +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from pytest_mock import MockerFixture @@ -152,11 +154,21 @@ def _selection_and_result( path=[ JointState( name=list(names), - position=[positions_by_name[name][0] for name in names], + position=np.asarray( + [positions_by_name[name][0] for name in names], dtype=np.float64 + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), JointState( name=list(names), - position=[positions_by_name[name][1] for name in names], + position=np.asarray( + [positions_by_name[name][1] for name in names], dtype=np.float64 + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ], ) @@ -311,8 +323,20 @@ def test_roboplan_toppra_parametrizes_unbounded_planar_base(tmp_path: Path) -> N result = PlanningResult( status=PlanningStatus.SUCCESS, path=[ - JointState(name=list(selection.joint_names), position=start), - JointState(name=list(selection.joint_names), position=goal), + JointState( + name=list(selection.joint_names), + position=np.asarray(start, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=list(selection.joint_names), + position=np.asarray(goal, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ], ) diff --git a/dimos/manipulation/planning/trajectory_generator/test_simple_parametrizer.py b/dimos/manipulation/planning/trajectory_generator/test_simple_parametrizer.py index 376bb2dca19..e514790ceea 100644 --- a/dimos/manipulation/planning/trajectory_generator/test_simple_parametrizer.py +++ b/dimos/manipulation/planning/trajectory_generator/test_simple_parametrizer.py @@ -17,9 +17,11 @@ from pathlib import Path from unittest.mock import MagicMock -from dimos_generated.geometry_msgs.msg import Pose, PoseStamped +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from pytest_mock import MockerFixture @@ -65,7 +67,13 @@ def _selection() -> PlanningGroupSelection: def _world(*, velocity: float = 2.0, acceleration: float = 6.0) -> WorldSpec: config = RobotModelConfig( model=RobotModel.from_file(Path("/robot.urdf")), - base_pose=PoseStamped(header=Header(frame_id=""), pose=Pose()), + base_pose=PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), joint_names=["arm/a", "arm/b"], base_link="base", ) @@ -98,9 +106,27 @@ def _result() -> PlanningResult: return PlanningResult( status=PlanningStatus.SUCCESS, path=[ - JointState(name=names, position=[0.0, 0.0]), - JointState(name=names, position=[0.2, 0.1]), - JointState(name=names, position=[0.4, 0.0]), + JointState( + name=names, + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=names, + position=np.array([0.2, 0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=names, + position=np.array([0.4, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ], ) @@ -125,14 +151,17 @@ def test_simple_parametrizer_materializes_segmented_trapezoid_plan() -> None: assert plan.group_ids == ("manipulator",) assert plan.trajectory.joint_names == ["arm/a", "arm/b"] assert len(plan.trajectory.points) == 9 - assert plan.trajectory.points[0].positions == [0.0, 0.0] - assert plan.trajectory.points[4].positions == [0.2, 0.1] - assert plan.trajectory.points[-1].positions == [0.4, 0.0] - assert [state.position for state in result.path] == [ - [0.0, 0.0], - [0.2, 0.1], - [0.4, 0.0], - ] + np.testing.assert_array_equal(plan.trajectory.points[0].positions, [0.0, 0.0]) + np.testing.assert_array_equal(plan.trajectory.points[4].positions, [0.2, 0.1]) + np.testing.assert_array_equal(plan.trajectory.points[-1].positions, [0.4, 0.0]) + np.testing.assert_array_equal( + [state.position for state in result.path], + [ + [0.0, 0.0], + [0.2, 0.1], + [0.4, 0.0], + ], + ) def test_simple_parametrizer_reports_generator_failure(mocker: MockerFixture) -> None: diff --git a/dimos/manipulation/planning/utils/path_utils.py b/dimos/manipulation/planning/utils/path_utils.py index 290c00949aa..df0c678fb7d 100644 --- a/dimos/manipulation/planning/utils/path_utils.py +++ b/dimos/manipulation/planning/utils/path_utils.py @@ -30,7 +30,9 @@ from typing import TYPE_CHECKING +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header import numpy as np if TYPE_CHECKING: @@ -81,7 +83,15 @@ def interpolate_path( for step in range(1, num_steps + 1): alpha = step / num_steps q_interp = q_start + alpha * diff - interpolated.append(JointState(name=joint_names, position=q_interp.tolist())) + interpolated.append( + JointState( + name=joint_names, + position=np.asarray(q_interp.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) return interpolated @@ -127,7 +137,15 @@ def interpolate_segment( for i in range(num_steps + 1): alpha = i / num_steps q_interp = q_start + alpha * diff - segment.append(JointState(name=joint_names, position=q_interp.tolist())) + segment.append( + JointState( + name=joint_names, + position=np.asarray(q_interp.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) return segment @@ -247,7 +265,15 @@ def clip_path_to_limits( for state in path: q = np.array(state.position, dtype=np.float64) q_clipped = np.clip(q, lower_limits, upper_limits) - clipped.append(JointState(name=state.name, position=q_clipped.tolist())) + clipped.append( + JointState( + name=state.name, + position=np.asarray(q_clipped.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) return clipped diff --git a/dimos/manipulation/planning/utils/point_cloud_self_filter.py b/dimos/manipulation/planning/utils/point_cloud_self_filter.py index 2df7dd74a1a..f87e3a1c4a5 100644 --- a/dimos/manipulation/planning/utils/point_cloud_self_filter.py +++ b/dimos/manipulation/planning/utils/point_cloud_self_filter.py @@ -17,8 +17,11 @@ from __future__ import annotations import asyncio +from copy import deepcopy from threading import RLock +from dimos_generated.sensor_msgs.msg import JointState, PointCloud2 +from dimos_generated.tf2_msgs.msg import TFMessage import numpy as np from numpy.typing import NDArray from pydantic import Field @@ -26,10 +29,9 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.types.timestamped import TimestampedBufferCollection +from dimos.msgs.pointcloud import pointcloud_xyz, select_points +from dimos.msgs.time import to_seconds +from dimos.types.timestamped import TimestampedBufferCollection, TimestampedData from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -55,7 +57,9 @@ class PointCloudSelfFilter(Module): def __init__(self, **kwargs: object) -> None: super().__init__(**kwargs) self._filter_lock = RLock() - self._states = TimestampedBufferCollection[JointState](self.config.state_history_s) + self._states = TimestampedBufferCollection[TimestampedData[JointState]]( + self.config.state_history_s + ) self._last_capture: float | None = None @rpc @@ -73,9 +77,10 @@ async def handle_coordinator_joint_state(self, state: JointState) -> None: def add_joint_state(self, state: JointState) -> None: """Buffer full model state; out-of-order arrivals remain timestamped.""" with self._filter_lock: - if np.isfinite(state.ts): - self._states.remove_by_timestamp(state.ts) - self._states.add(JointState(state)) + stamp = to_seconds(state.header.stamp) + if np.isfinite(stamp): + self._states.remove_by_timestamp(stamp) + self._states.add(TimestampedData(deepcopy(state), stamp)) latest = self._states.last() if latest is not None: self._states.prune_old(latest.ts - self.config.state_history_s) @@ -83,21 +88,20 @@ def add_joint_state(self, state: JointState) -> None: def filter_cloud(self, cloud: PointCloud2) -> PointCloud2 | None: """Filter one aligned capture without changing its frame or point fields.""" with self._filter_lock: - if not np.isfinite(cloud.ts) or ( - self._last_capture is not None and cloud.ts < self._last_capture + stamp = to_seconds(cloud.header.stamp) + if not np.isfinite(stamp) or ( + self._last_capture is not None and stamp < self._last_capture ): logger.warning("Dropping cloud: invalid or out-of-order capture timestamp") return None - points = cloud.points_f32() + points = pointcloud_xyz(cloud).astype(np.float32) if not np.isfinite(points).all(): return None keep = self._compute_keep_mask(cloud, points) if keep is None: return None - filtered = PointCloud2(frame_id=cloud.frame_id, ts=cloud.ts, seq=cloud.seq) - for name, values in cloud.pointcloud_tensor.point.items(): - filtered.pointcloud_tensor.point[name] = values[keep] - self._last_capture = cloud.ts + filtered = select_points(cloud, keep) + self._last_capture = stamp return filtered def _on_pointcloud(self, cloud: PointCloud2) -> None: @@ -112,6 +116,6 @@ def _compute_keep_mask( """Return one keep flag per point, or None when the capture cannot be filtered. Called under the filter lock. Implementations must align any required - state and transforms to cloud.ts and preserve the input point order. + state and transforms to the cloud header stamp and preserve the input point order. """ raise NotImplementedError("Choose a concrete point cloud self filter") diff --git a/dimos/manipulation/planning/utils/roboplan_point_cloud_self_filter.py b/dimos/manipulation/planning/utils/roboplan_point_cloud_self_filter.py index d5cfaa0d9c7..311a10469c3 100644 --- a/dimos/manipulation/planning/utils/roboplan_point_cloud_self_filter.py +++ b/dimos/manipulation/planning/utils/roboplan_point_cloud_self_filter.py @@ -19,6 +19,7 @@ import importlib import xml.etree.ElementTree as ET +from dimos_generated.sensor_msgs.msg import PointCloud2 import numpy as np from numpy.typing import NDArray from pydantic import Field @@ -27,7 +28,8 @@ PointCloudSelfFilter, PointCloudSelfFilterConfig, ) -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.geometry import transform_matrix +from dimos.msgs.time import to_seconds from dimos.robot.assets.model import RobotModel from dimos.utils.logging_config import setup_logger from dimos.utils.transform_utils import matrix_to_pose @@ -54,22 +56,23 @@ def _compute_keep_mask( ) -> NDArray[np.bool_] | None: base_from_sensor = self.tfbuffer.get( self._base_link, - cloud.frame_id, - time_point=cloud.ts, + cloud.header.frame_id, + time_point=to_seconds(cloud.header.stamp), time_tolerance=self.config.tf_tolerance_s, forward_tolerance=self.config.tf_forward_tolerance_s, ) - q = self._capture_configuration(cloud.ts) + q = self._capture_configuration(to_seconds(cloud.header.stamp)) if base_from_sensor is None or q is None: logger.warning("Dropping cloud: capture-time robot state or TF unavailable") return None # This Scene is rooted at the URDF base, so its world is the base frame. - transform = base_from_sensor.to_matrix() + transform = transform_matrix(base_from_sensor.transform) base_points = np.asarray(points @ transform[:3, :3].T + transform[:3, 3], dtype=np.float64) return ~np.asarray(self._body_filter.computeMask(q, base_points), dtype=bool) def _capture_configuration(self, stamp: float) -> NDArray[np.float64] | None: - state = self._states.find_closest(stamp, self.config.state_tolerance_s) + sample = self._states.find_closest(stamp, self.config.state_tolerance_s) + state = sample.value if sample is not None else None positions: dict[str, float] = {} if state is not None: if len(state.name) != len(state.position) or len(set(state.name)) != len(state.name): @@ -104,7 +107,7 @@ def _capture_configuration(self, stamp: float) -> NDArray[np.float64] | None: self._neutral_q, joint.child_link, joint.parent_link ) ) - motion = np.linalg.inv(origin) @ parent_from_child.to_matrix() + motion = np.linalg.inv(origin) @ transform_matrix(parent_from_child.transform) pose = matrix_to_pose(motion) if info.num_position_dofs == 4: if not np.allclose(motion[2], [0, 0, 1, 0], atol=1e-6): diff --git a/dimos/manipulation/planning/utils/test_point_cloud_self_filter.py b/dimos/manipulation/planning/utils/test_point_cloud_self_filter.py index fd5fc274f0b..6ba03951f9f 100644 --- a/dimos/manipulation/planning/utils/test_point_cloud_self_filter.py +++ b/dimos/manipulation/planning/utils/test_point_cloud_self_filter.py @@ -16,13 +16,15 @@ from collections.abc import Iterator +from dimos_generated.sensor_msgs.msg import JointState, PointCloud2, PointField +from dimos_generated.std_msgs.msg import Header import numpy as np from numpy.typing import NDArray -from open3d.core import Tensor import pytest from dimos.manipulation.planning.utils.point_cloud_self_filter import PointCloudSelfFilter -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_view, pointcloud_xyz +from dimos.msgs.time import time_from_seconds, to_seconds class _PositiveXFilter(PointCloudSelfFilter): @@ -40,23 +42,54 @@ def point_filter() -> Iterator[_PositiveXFilter]: def test_backend_mask_preserves_fields_and_capture_metadata(point_filter: _PositiveXFilter) -> None: - cloud = PointCloud2.from_numpy( + cloud = pointcloud_from_xyz( np.asarray([[-1, 0, 0], [1, 2, 3]], dtype=np.float32), - frame_id="camera", - timestamp=2.0, + header=Header(frame_id="camera", stamp=time_from_seconds(2.0)), ) - cloud.seq = 7 - cloud.pointcloud_tensor.point["intensity"] = Tensor(np.asarray([[10], [20]], dtype=np.float32)) + records = np.array( + [(-1, 0, 0, 10), (1, 2, 3, 20)], + dtype=[("x", " None: + state = JointState( + header=Header(frame_id="base", stamp=time_from_seconds(2.0)), + name=["joint"], + position=np.array([0.5]), + velocity=np.array([]), + effort=np.array([]), + ) + point_filter.add_joint_state(state) + state.position[0] = 99 + sample = point_filter._states.find_closest(2.0, 0.001) + assert sample is not None + np.testing.assert_array_equal(sample.value.position, [0.5]) + point_filter.add_joint_state(state) + sample = point_filter._states.find_closest(2.0, 0.001) + assert sample is not None + np.testing.assert_array_equal(sample.value.position, [99]) diff --git a/dimos/manipulation/planning/utils/test_roboplan_point_cloud_self_filter.py b/dimos/manipulation/planning/utils/test_roboplan_point_cloud_self_filter.py index 8858067072b..7310b17e022 100644 --- a/dimos/manipulation/planning/utils/test_roboplan_point_cloud_self_filter.py +++ b/dimos/manipulation/planning/utils/test_roboplan_point_cloud_self_filter.py @@ -18,8 +18,10 @@ from pathlib import Path from typing import Any, cast +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import JointState, PointCloud2, PointField +from dimos_generated.std_msgs.msg import Header import numpy as np -from open3d.core import Tensor import pinocchio as pin import pytest import trimesh @@ -27,11 +29,8 @@ from dimos.manipulation.planning.utils.roboplan_point_cloud_self_filter import ( RoboPlanPointCloudSelfFilter, ) -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_view, pointcloud_xyz +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.protocol.tf.tf import MultiTBuffer from dimos.robot.assets.git_cache import DEFAULT_ROBOT_ASSET_CACHE_ROOT, GitAssetCache from dimos.robot.assets.model import RobotModel @@ -87,20 +86,21 @@ def _place_arm( """Move the rigid robot base at capture time in both frames.""" for parent in ("camera", "world"): module.tfbuffer.receive_transform( - Transform( - translation=Vector3(*at), - frame_id=parent, + TransformStamped( + header=Header(frame_id=parent, stamp=time_from_seconds(ts)), child_frame_id="base", - ts=ts, + transform=Transform( + translation=Vector3(x=at[0], y=at[1], z=at[2]), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) ) def _cloud(points: list[list[float]], ts: float = 1.0) -> PointCloud2: - return PointCloud2.from_numpy( + return pointcloud_from_xyz( np.asarray(points, dtype=np.float32).reshape((-1, 3)), - frame_id="camera", - timestamp=ts, + header=Header(frame_id="camera", stamp=time_from_seconds(ts)), ) @@ -114,7 +114,7 @@ def test_points_on_the_robot_are_dropped_and_the_rest_survive( assert result is not None filtered = result - np.testing.assert_allclose(filtered.points_f32(), [[2.0, 0.0, 0.0]], atol=1e-6) + np.testing.assert_allclose(pointcloud_xyz(filtered), [[2.0, 0.0, 0.0]], atol=1e-06) def test_a_cloud_without_capture_time_tf_is_dropped_whole( @@ -129,31 +129,72 @@ def test_a_cloud_without_capture_time_tf_is_dropped_whole( def _joint_robot(kind: str) -> str: return _URDF.replace('name="shoulder" type="fixed"', f'name="shoulder" type="{kind}"').replace( '', - '' - '', + '', ) def test_capture_state_is_matched_by_timestamp_instead_of_latest(make_filter): module = make_filter(urdf_xml=_joint_robot("prismatic"), state_tolerance_s=0.001) _place_arm(module, (0, 0, 0), 1.0) - module.add_joint_state(JointState(ts=2.0, name=["shoulder"], position=[2.0])) - module.add_joint_state(JointState(ts=1.0, name=["shoulder"], position=[0.0])) - module.add_joint_state(JointState(ts=1.0, name=["shoulder"], position=[0.5])) + module.add_joint_state( + JointState( + name=["shoulder"], + header=Header(frame_id="", stamp=time_from_seconds(2.0)), + position=np.asarray([2.0], dtype=np.float64), + velocity=np.asarray([], dtype=np.float64), + effort=np.asarray([], dtype=np.float64), + ) + ) + module.add_joint_state( + JointState( + name=["shoulder"], + header=Header(frame_id="", stamp=time_from_seconds(1.0)), + position=np.asarray([0.0], dtype=np.float64), + velocity=np.asarray([], dtype=np.float64), + effort=np.asarray([], dtype=np.float64), + ) + ) + module.add_joint_state( + JointState( + name=["shoulder"], + header=Header(frame_id="", stamp=time_from_seconds(1.0)), + position=np.asarray([0.5], dtype=np.float64), + velocity=np.asarray([], dtype=np.float64), + effort=np.asarray([], dtype=np.float64), + ) + ) result = module.filter_cloud(_cloud([[0.5, 0, 0], [2, 0, 0]], ts=1.0)) assert result is not None - np.testing.assert_allclose(result.points_f32(), [[2, 0, 0]]) + np.testing.assert_allclose(pointcloud_xyz(result), [[2, 0, 0]]) @pytest.mark.parametrize( "state", [ None, - JointState(ts=2.0, name=["shoulder"], position=[0.5]), - JointState(ts=1.0, name=["wrong"], position=[0.5]), - JointState(ts=1.0, name=["shoulder"], position=[float("nan")]), + JointState( + name=["shoulder"], + header=Header(frame_id="", stamp=time_from_seconds(2.0)), + position=np.asarray([0.5], dtype=np.float64), + velocity=np.asarray([], dtype=np.float64), + effort=np.asarray([], dtype=np.float64), + ), + JointState( + name=["wrong"], + header=Header(frame_id="", stamp=time_from_seconds(1.0)), + position=np.asarray([0.5], dtype=np.float64), + velocity=np.asarray([], dtype=np.float64), + effort=np.asarray([], dtype=np.float64), + ), + JointState( + name=["shoulder"], + header=Header(frame_id="", stamp=time_from_seconds(1.0)), + position=np.asarray([float("nan")], dtype=np.float64), + velocity=np.asarray([], dtype=np.float64), + effort=np.asarray([], dtype=np.float64), + ), ], ) def test_missing_late_or_invalid_state_drops_the_capture(make_filter, state): @@ -170,12 +211,20 @@ def test_continuous_joint_uses_cos_sin_configuration(make_filter): urdf = urdf.replace("", '') module = make_filter(urdf_xml=urdf) _place_arm(module, (0, 0, 0), 1.0) - module.add_joint_state(JointState(ts=1.0, name=["shoulder"], position=[np.pi / 2])) + module.add_joint_state( + JointState( + name=["shoulder"], + header=Header(frame_id="", stamp=time_from_seconds(1.0)), + position=np.asarray([np.pi / 2], dtype=np.float64), + velocity=np.asarray([], dtype=np.float64), + effort=np.asarray([], dtype=np.float64), + ) + ) result = module.filter_cloud(_cloud([[0, 0.5, 0], [0.5, 0, 0]])) assert result is not None - np.testing.assert_allclose(result.points_f32(), [[0.5, 0, 0]]) + np.testing.assert_allclose(pointcloud_xyz(result), [[0.5, 0, 0]]) @pytest.mark.parametrize("kind", ["planar", "floating"]) @@ -187,55 +236,84 @@ def test_multidof_joints_use_capture_time_relative_tf(make_filter, kind): _place_arm(module, (0, 0, 0), 1.0) assert module.filter_cloud(_cloud([[0.5, 0, 0]])) is None module.tfbuffer.receive_transform( - Transform(translation=Vector3(0.5, 0, 0), frame_id="base", child_frame_id="arm", ts=1.0) + TransformStamped( + header=Header(frame_id="base", stamp=time_from_seconds(1.0)), + child_frame_id="arm", + transform=Transform( + translation=Vector3(x=0.5, y=0, z=0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) ) result = module.filter_cloud(_cloud([[0.5, 0, 0], [0, 0, 0]])) assert result is not None - np.testing.assert_allclose(result.points_f32(), [[0, 0, 0]]) + np.testing.assert_allclose(pointcloud_xyz(result), [[0, 0, 0]]) def test_narrowphase_retains_sphere_corner_obstacles_and_ancillary_fields(make_filter): urdf = _URDF.replace('', '') module = make_filter(urdf_xml=urdf, padding_m=0.05) _place_arm(module, (0, 0, 0), 1.0) - cloud = PointCloud2.from_numpy( + cloud = pointcloud_from_xyz( np.array([[0.14, 0, 0], [0.12, 0.12, 0], [0.2, 0, 0]], dtype=np.float32), - frame_id="camera", - timestamp=1.0, - intensities=np.array([1, 2, 3], dtype=np.float32), + header=Header(frame_id="camera", stamp=time_from_seconds(1.0)), + ) + records = np.zeros( + 3, dtype=[("x", " JointState: plant_ctx = self._diagram.GetSubsystemContext(self._plant, ctx) full_positions = self._plant.GetPositions(plant_ctx) positions = [float(full_positions[idx]) for idx in robot_data.joint_indices] - return JointState(name=robot_data.config.joint_names, position=positions) + return JointState( + name=robot_data.config.joint_names, + position=np.asarray(positions, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) # Collision Checking (context-based) @@ -955,7 +962,13 @@ def check_edge_collision_free( for i in range(n_steps): t = i / (n_steps - 1) q = q_start + t * (q_end - q_start) - interp_state = JointState(name=start.name, position=q.tolist()) + interp_state = JointState( + name=start.name, + position=np.asarray(q.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) self.set_joint_state(ctx, interp_state) if not self.is_collision_free(ctx): return False @@ -992,7 +1005,7 @@ def _get_group_ee_pose(self, ctx: Context, group_id: PlanningGroupID) -> PoseSta quat = X_WE.rotation().ToQuaternion() # Drake returns [w, x, y, z] return PoseStamped( - header=Header(frame_id="world"), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=Point(x=float(pos[0]), y=float(pos[1]), z=float(pos[2])), orientation=Quaternion( diff --git a/dimos/manipulation/planning/world/roboplan_world.py b/dimos/manipulation/planning/world/roboplan_world.py index ee8789f4e05..db1d5b80350 100644 --- a/dimos/manipulation/planning/world/roboplan_world.py +++ b/dimos/manipulation/planning/world/roboplan_world.py @@ -29,6 +29,7 @@ from threading import RLock from typing import TYPE_CHECKING, Any +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion from dimos_generated.std_msgs.msg import Header import numpy as np @@ -305,7 +306,13 @@ def get_joint_state(self, ctx: RoboPlanContext) -> JointState: q = ctx.q if not len(q): q = np.zeros(len(model_data.config.joint_names), dtype=np.float64) - return JointState(name=model_data.config.joint_names, position=q.astype(float).tolist()) + return JointState( + name=model_data.config.joint_names, + position=np.asarray(q.astype(float).tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) # Collision Checking @@ -359,7 +366,7 @@ def get_group_ee_pose(self, ctx: RoboPlanContext, group_id: PlanningGroupID) -> mat = self.get_link_pose(ctx, group.tip_link) pose = pose_from_matrix(mat) return PoseStamped( - header=Header(frame_id="world"), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=Point(x=pose.position.x, y=pose.position.y, z=pose.position.z), orientation=Quaternion( diff --git a/dimos/manipulation/planning/world/test_drake_world_planning_groups.py b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py index 1693fa99a38..cfab6d89a44 100644 --- a/dimos/manipulation/planning/world/test_drake_world_planning_groups.py +++ b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py @@ -17,6 +17,7 @@ from dataclasses import replace from pathlib import Path +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header @@ -46,13 +47,17 @@ def _trajectory(names: list[str], first: list[float], second: list[float]) -> Jo points=[ JointTrajectoryPoint( time_from_start=duration_from_seconds(0.0), - positions=first, - velocities=[0.0] * len(names), + positions=np.asarray(first, dtype=np.float64), + velocities=np.asarray([0.0] * len(names), dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), JointTrajectoryPoint( time_from_start=duration_from_seconds(2.0), - positions=second, - velocities=[0.0] * len(names), + positions=np.asarray(second, dtype=np.float64), + velocities=np.asarray([0.0] * len(names), dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ], ) @@ -134,7 +139,7 @@ def _config( return RobotModelConfig( model=RobotModel.from_file(path).with_default_joint_acceleration_limit(2.0), base_pose=PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose(position=Point(x=0, y=0, z=0), orientation=Quaternion(x=0, y=0, z=0, w=1)), ), joint_names=joints or ["joint1", "joint2"], @@ -228,7 +233,7 @@ def test_drake_obstacle_ids_are_world_owned_and_invalid_insertions_are_rejected( name="box", obstacle_type=ObstacleType.BOX, pose=PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose(position=Point(x=2, y=0, z=0), orientation=Quaternion(x=0, y=0, z=0, w=1)), ), dimensions=(0.1, 0.1, 0.1), @@ -249,7 +254,13 @@ def test_drake_obstacle_ids_are_world_owned_and_invalid_insertions_are_rejected( world.finalize() assert world.add_obstacle(unnamed) is None assert world.add_obstacle(obstacle) == "box" - joint_state = JointState(name=["joint1", "joint2"], position=[0.0, 0.0]) + joint_state = JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) assert world.check_config_collision_free(joint_state) original_geometry_id = world._obstacles["box"].geometry_id assert world.add_obstacle(obstacle) is None @@ -257,7 +268,7 @@ def test_drake_obstacle_ids_are_world_owned_and_invalid_insertions_are_rejected( assert world.update_obstacle(replace(obstacle, name="missing")) is False assert world.update_obstacle_pose("missing", obstacle.pose) is False moved_pose = PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) ), @@ -294,7 +305,7 @@ def test_drake_obstacle_replacement_failure_invalidates_world( name="box", obstacle_type=ObstacleType.BOX, pose=PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose(position=Point(x=0, y=0, z=0), orientation=Quaternion(x=0, y=0, z=0, w=1)), ), dimensions=(0.1, 0.1, 0.1), @@ -320,7 +331,16 @@ def test_drake_group_fk_uses_tip_link_and_unique_pose_group(tmp_path: Path) -> N _load(world, _config(urdf, [_arm_group("joint1", "joint2")])) world.finalize() ctx = world.get_live_context() - world.set_joint_state(ctx, JointState(name=["joint1", "joint2"], position=[0.0, 0.0])) + world.set_joint_state( + ctx, + JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + ) group_pose = world.get_group_ee_pose(ctx, "arm") default_pose = world.get_ee_pose(ctx) @@ -342,7 +362,7 @@ def test_drake_applies_config_base_pose_when_urdf_has_world_base_joint( RobotModelConfig( model=RobotModel.from_file(urdf).with_default_joint_acceleration_limit(2.0), base_pose=PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=Point(x=0, y=0.5, z=0), orientation=Quaternion(x=0, y=0, z=0, w=1) ), @@ -379,7 +399,7 @@ def test_drake_planar_base_coordinates_move_original_robot_root(tmp_path: Path) .with_planar_base(planar_base) ), base_pose=PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=Point(x=0, y=0, z=0.5), orientation=Quaternion(x=0, y=0, z=0, w=1) ), @@ -400,7 +420,13 @@ def test_drake_planar_base_coordinates_move_original_robot_root(tmp_path: Path) context = world.get_live_context() world.set_joint_state( context, - JointState(name=joint_names, position=[1.0, 2.0, np.pi / 2.0, 0.0, 0.0]), + JointState( + name=joint_names, + position=np.array([1.0, 2.0, np.pi / 2.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) original_root_pose = world.get_link_pose(context, "base_link") @@ -429,7 +455,16 @@ def test_drake_group_jacobian_shape_and_group_local_order(tmp_path: Path) -> Non ) world.finalize() ctx = world.get_live_context() - world.set_joint_state(ctx, JointState(name=["joint1", "joint2"], position=[0.0, 0.0])) + world.set_joint_state( + ctx, + JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + ) forward_jacobian = world.get_group_jacobian(ctx, "wrist_forward") reverse_jacobian = world.get_group_jacobian(ctx, "wrist_reverse") @@ -486,7 +521,13 @@ def test_drake_animate_trajectory_projects_selected_joints_on_shared_ticks( world._meshcat = object() # type: ignore[assignment] world.set_joint_state( world.get_live_context(), - JointState(name=["joint1", "joint2"], position=[0.1, 0.2]), + JointState( + name=["joint1", "joint2"], + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) updates: list[list[float]] = [] visibility: list[bool] = [] @@ -519,7 +560,13 @@ def test_drake_animate_trajectory_validates_before_visibility_and_cleans_up( world._meshcat = object() # type: ignore[assignment] world.set_joint_state( world.get_live_context(), - JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), + JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) visibility: list[bool] = [] monkeypatch.setattr(world, "_set_preview_visibility", visibility.append) @@ -550,7 +597,13 @@ def test_drake_cancel_preview_hides_model_before_animation_resumes( world._meshcat = object() # type: ignore[assignment] world.set_joint_state( world.get_live_context(), - JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), + JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) visibility: list[bool] = [] snapshots_during_sleep: list[list[bool]] = [] @@ -580,7 +633,13 @@ def test_drake_animate_trajectory_cancellation_stops_stale_frames( world._meshcat = object() # type: ignore[assignment] world.set_joint_state( world.get_live_context(), - JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), + JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) updates: list[list[float]] = [] visibility: list[bool] = [] diff --git a/dimos/manipulation/sdk.py b/dimos/manipulation/sdk.py index ae31e9b4b92..34e82b02710 100644 --- a/dimos/manipulation/sdk.py +++ b/dimos/manipulation/sdk.py @@ -19,6 +19,7 @@ import math from typing import TYPE_CHECKING +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header @@ -138,7 +139,13 @@ def move_joints( ) -> ExecutionResult: """Plan and move to positions in group joint order; angular units are radians.""" values = _vector(positions, len(self.info.joint_names), "positions") - target = JointState(name=list(self.info.joint_names), position=values.tolist()) + target = JointState( + name=list(self.info.joint_names), + position=np.asarray(values.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) return self._move_joint_target("move_joints", target, speed_scale, timeout) def _move_joint_target( @@ -173,7 +180,7 @@ def move_pose( if math.hypot(rotation.x, rotation.y, rotation.z, rotation.w) == 0.0: raise ValueError("orientation must have nonzero norm") target = PoseStamped( - header=Header(frame_id="world"), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), pose=Pose(position=Point(x=xyz[0], y=xyz[1], z=xyz[2]), orientation=rotation), ) plan = self.rpc.plan_to_poses({self.info.id: target}, speed_scale=speed_scale) diff --git a/dimos/manipulation/test_execution_manager.py b/dimos/manipulation/test_execution_manager.py index 42570bb7652..be2713509a1 100644 --- a/dimos/manipulation/test_execution_manager.py +++ b/dimos/manipulation/test_execution_manager.py @@ -16,11 +16,13 @@ from unittest.mock import MagicMock +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Duration, Time from dimos_generated.dimos_msgs.msg import TrajectoryStatus from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import numpy as np import pytest from dimos.control.coordinator import ControlCoordinator @@ -45,16 +47,33 @@ def _plan( ) -> GeneratedPlan: points = [ JointTrajectoryPoint( - positions=[0.0] * len(names), time_from_start=duration_from_seconds(0.0) + positions=np.asarray([0.0] * len(names), dtype=np.float64), + time_from_start=duration_from_seconds(0.0), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), JointTrajectoryPoint( - positions=[1.0] * len(names), time_from_start=duration_from_seconds(1.0) + positions=np.asarray([1.0] * len(names), dtype=np.float64), + time_from_start=duration_from_seconds(1.0), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ] return GeneratedPlan( group_ids=("both_arms",), trajectory=JointTrajectory(header=header_now(), joint_names=list(names), points=points), - path=[JointState(name=list(names), position=point.positions) for point in points], + path=[ + JointState( + name=list(names), + position=np.asarray(point.positions, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + for point in points + ], status=status, ) @@ -69,7 +88,14 @@ def invoke(task: str, method: str, args: dict | None = None): return TrajectoryExecutionResult(TrajectoryExecutionStatus.ACCEPTED) if method == "cancel": return TrajectoryCancellationResult(TrajectoryCancellationStatus.ALREADY_STOPPED) - return TrajectoryStatus(header=header_now(), state=TrajectoryState.IDLE) + return TrajectoryStatus( + header=header_now(), + state=TrajectoryState.IDLE, + progress=0.0, + time_elapsed=Duration(sec=0, nanosec=0), + time_remaining=Duration(sec=0, nanosec=0), + error="", + ) coordinator.task_invoke.side_effect = invoke return coordinator @@ -177,7 +203,12 @@ def _task_invoke(self, task, method, args=None): task, TrajectoryCancellationResult(TrajectoryCancellationStatus.CANCELLED) ) return TrajectoryStatus( - header=header_now(), state=self.states[task], error=self.errors.get(task, "") + header=header_now(), + state=self.states[task], + error=self.errors.get(task, ""), + progress=0.0, + time_elapsed=Duration(sec=0, nanosec=0), + time_remaining=Duration(sec=0, nanosec=0), ) def manager(self) -> PlanExecutionManager: @@ -283,7 +314,7 @@ def test_split_cdr_plan_preserves_header_durations_and_optional_point_fields(): plan.trajectory.points = points try: assert manager.execute(plan, blocking=False).status == ExecutionStatus.ACCEPTED - part = JointTrajectory.decode(robot.dispatched["base_traj"].encode()) + part = cdr_decode(cdr_encode(robot.dispatched["base_traj"]), JointTrajectory) assert part.header == plan.trajectory.header assert part.points[-1].time_from_start == Duration(sec=1, nanosec=9) assert list(part.points[-1].velocities) == [0.3, 0.4, 0.1] diff --git a/dimos/manipulation/test_generated_plan_materialization.py b/dimos/manipulation/test_generated_plan_materialization.py index 3b5becd08cb..0606000e417 100644 --- a/dimos/manipulation/test_generated_plan_materialization.py +++ b/dimos/manipulation/test_generated_plan_materialization.py @@ -17,6 +17,7 @@ from pathlib import Path from unittest.mock import MagicMock +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, Pose, @@ -29,6 +30,7 @@ from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import numpy as np import pytest from dimos.manipulation.manipulation_module import ManipulationState @@ -76,11 +78,14 @@ def generate(self, waypoints: list[list[float]]) -> JointTrajectory: points=[ JointTrajectoryPoint( time_from_start=duration_from_seconds(float(index)), - positions=list(point), - velocities=[0.0] * self.num_joints, + positions=np.asarray(list(point), dtype=np.float64), + velocities=np.asarray([0.0] * self.num_joints, dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) for index, point in enumerate(waypoints) ], + joint_names=[], ) @@ -88,7 +93,11 @@ def _model() -> RobotModelConfig: return RobotModelConfig( model=RobotModel.from_file(Path("/robot.urdf")), base_pose=PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), joint_names=["left/a", "left/b", "right/c"], base_link="base", @@ -151,7 +160,22 @@ def _module(monkeypatch: pytest.MonkeyPatch, module_factory): def _path(names: list[str], first: list[float], second: list[float]) -> list[JointState]: - return [JointState(name=names, position=first), JointState(name=names, position=second)] + return [ + JointState( + name=names, + position=np.asarray(first, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=names, + position=np.asarray(second, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + ] def test_materializes_once_with_reordered_groups_heterogeneous_limits_and_distinct_path( @@ -178,10 +202,28 @@ def test_cartesian_plan_preserves_planner_timestamps_and_velocities(monkeypatch, module = _module(monkeypatch, module_factory) module._state = ManipulationState.IDLE names = ["left/b", "left/a"] - start = JointState(name=names, position=[0.0, 0.0]) + start = JointState( + name=names, + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) path = [ - JointState(name=names, position=[0.0, 0.0], velocity=[0.0, 0.0]), - JointState(name=names, position=[0.2, 0.1], velocity=[0.4, 0.2]), + JointState( + name=names, + position=np.array([0.0, 0.0], dtype=np.float64), + velocity=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=names, + position=np.array([0.2, 0.1], dtype=np.float64), + velocity=np.array([0.4, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + effort=np.array([], dtype=np.float64), + ), ] module._world_monitor.current_model_joint_state.return_value = start module._planner.plan_cartesian_path.return_value = PlanningResult( @@ -193,10 +235,20 @@ def test_cartesian_plan_preserves_planner_timestamps_and_velocities(monkeypatch, plan = module.generate_cartesian_plan( { "left_arm": ( - TransformStamped(header=Header(frame_id="world"), child_frame_id=""), TransformStamped( - header=Header(frame_id="world"), - transform=Transform(translation=Vector3(x=0.01, y=0.0, z=0.0)), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + child_frame_id="", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + TransformStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + transform=Transform( + translation=Vector3(x=0.01, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), child_frame_id="", ), ) @@ -210,7 +262,7 @@ def test_cartesian_plan_preserves_planner_timestamps_and_velocities(monkeypatch, assert plan is not None assert [to_seconds(point.time_from_start) for point in plan.trajectory.points] == [0.0, 0.25] - assert [point.velocities for point in plan.trajectory.points] == [ + assert [point.velocities.tolist() for point in plan.trajectory.points] == [ [0.0, 0.0], [0.4, 0.2], ] diff --git a/dimos/manipulation/test_manipulation_module.py b/dimos/manipulation/test_manipulation_module.py index 15c15fa915b..7068dfdfb82 100644 --- a/dimos/manipulation/test_manipulation_module.py +++ b/dimos/manipulation/test_manipulation_module.py @@ -24,9 +24,12 @@ import importlib.util from unittest.mock import DEFAULT, MagicMock +from dimos_generated.builtin_interfaces.msg import Duration, Time from dimos_generated.dimos_msgs.msg import TrajectoryStatus from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from dimos.control.coordinator import ControlCoordinator @@ -76,9 +79,10 @@ def joint_state_zeros(): "joint6", "joint7", ], - position=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], - velocity=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], - effort=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + position=np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64), + velocity=np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64), + effort=np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) @@ -97,7 +101,12 @@ def invoke(task: str, method: str, args: dict | None = None): return DEFAULT coordinator.task_invoke.return_value = TrajectoryStatus( - header=header_now(), state=TrajectoryState.COMPLETED + header=header_now(), + state=TrajectoryState.COMPLETED, + progress=0.0, + time_elapsed=Duration(sec=0, nanosec=0), + time_remaining=Duration(sec=0, nanosec=0), + error="", ) coordinator.task_invoke.side_effect = invoke mod = ManipulationModule( @@ -150,7 +159,13 @@ def test_plan_to_joints(self, module, joint_state_zeros): """Test planning to a joint configuration.""" module._on_joint_state(joint_state_zeros) - target = JointState(position=[0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]) + target = JointState( + position=np.array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) result = module.plan_to_joints({"manipulator": target}) assert result.succeeded, result.message @@ -166,7 +181,17 @@ def test_plan_to_explicit_joint_target(self, module, joint_state_zeros): """Test planning to an explicit planning-group joint target.""" module._on_joint_state(joint_state_zeros) - result = module.plan_to_joints({"manipulator": JointState(position=[0.05] * 7)}) + result = module.plan_to_joints( + { + "manipulator": JointState( + position=np.asarray([0.05] * 7, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + } + ) assert result.succeeded, result.message assert module._state == ManipulationState.COMPLETED @@ -179,7 +204,9 @@ def test_add_and_remove_obstacle(self, module, joint_state_zeros): """Test adding and removing obstacles.""" module._on_joint_state(joint_state_zeros) - pose = Pose(position=Point(x=0.5, y=0.0, z=0.3), orientation=Quaternion()) + pose = Pose( + position=Point(x=0.5, y=0.0, z=0.3), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) obstacle_id = module.add_obstacle("test_box", pose, "box", [0.1, 0.1, 0.1]) assert obstacle_id != "" @@ -214,7 +241,17 @@ def test_trajectory_uses_canonical_names(self, module, joint_state_zeros): """Test that execution preserves canonical model joint names.""" module._on_joint_state(joint_state_zeros) - result = module.plan_to_joints({"manipulator": JointState(position=[0.05] * 7)}) + result = module.plan_to_joints( + { + "manipulator": JointState( + position=np.asarray([0.05] * 7, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + } + ) assert result.succeeded, result.message assert module._last_plan is not None @@ -232,7 +269,17 @@ def test_execute_with_mock_coordinator(self, module, joint_state_zeros): """Test execute sends trajectory to coordinator.""" module._on_joint_state(joint_state_zeros) - result = module.plan_to_joints({"manipulator": JointState(position=[0.05] * 7)}) + result = module.plan_to_joints( + { + "manipulator": JointState( + position=np.asarray([0.05] * 7, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + } + ) assert result.succeeded, result.message result = module.execute() @@ -251,7 +298,17 @@ def test_execute_rejected_by_coordinator(self, module, joint_state_zeros): """Test handling of coordinator rejection.""" module._on_joint_state(joint_state_zeros) - plan_result = module.plan_to_joints({"manipulator": JointState(position=[0.05] * 7)}) + plan_result = module.plan_to_joints( + { + "manipulator": JointState( + position=np.asarray([0.05] * 7, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + } + ) assert plan_result.succeeded, plan_result.message module._control_coordinator.execute_result = TrajectoryExecutionResult( @@ -271,7 +328,17 @@ def test_state_transitions_during_execution(self, module, joint_state_zeros): module._on_joint_state(joint_state_zeros) # Plan - should go through PLANNING -> COMPLETED - result = module.plan_to_joints({"manipulator": JointState(position=[0.05] * 7)}) + result = module.plan_to_joints( + { + "manipulator": JointState( + position=np.asarray([0.05] * 7, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + } + ) assert result.succeeded, result.message assert module._state == ManipulationState.COMPLETED @@ -280,7 +347,17 @@ def test_state_transitions_during_execution(self, module, joint_state_zeros): assert module.execute().status is ExecutionStatus.NO_PLAN # Plan again - result = module.plan_to_joints({"manipulator": JointState(position=[0.05] * 7)}) + result = module.plan_to_joints( + { + "manipulator": JointState( + position=np.asarray([0.05] * 7, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + } + ) assert result.succeeded, result.message # Execute - should go to EXECUTING then COMPLETED diff --git a/dimos/manipulation/test_manipulation_monitor_preview.py b/dimos/manipulation/test_manipulation_monitor_preview.py index cbe008cf333..cf50c5d6c69 100644 --- a/dimos/manipulation/test_manipulation_monitor_preview.py +++ b/dimos/manipulation/test_manipulation_monitor_preview.py @@ -24,6 +24,7 @@ from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import numpy as np import pytest from dimos.manipulation.manipulation_module import ManipulationModule @@ -45,7 +46,11 @@ def canonical_model_config() -> RobotModelConfig: return RobotModelConfig( model=RobotModel.from_file(Path("/path/to/robot.urdf")), base_pose=PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), joint_names=["left/joint1", "left/joint2", "left/joint3"], base_link="link_base", @@ -64,7 +69,11 @@ def _one_joint_config(name: str = "arm") -> RobotModelConfig: return RobotModelConfig( model=RobotModel.from_file(Path("/path")), base_pose=PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), joint_names=["j0"], base_link="base_link", @@ -86,7 +95,10 @@ def _install_generated_plan( module._world_monitor.planning_groups = PlanningGroupRegistry(config.planning_groups) module._world_monitor.get_current_joint_state.return_value = JointState( name=config.joint_names, - position=[0.0 for _ in config.joint_names], + position=np.asarray([0.0 for _ in config.joint_names], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) module._last_plan = GeneratedPlan( trajectory=JointTrajectory( @@ -95,8 +107,10 @@ def _install_generated_plan( points=[ JointTrajectoryPoint( time_from_start=duration_from_seconds(float(index)), - positions=list(point), - velocities=[0.0 for _ in config.joint_names], + positions=np.asarray(list(point), dtype=np.float64), + velocities=np.asarray([0.0 for _ in config.joint_names], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) for index, point in enumerate(points) ], @@ -106,7 +120,10 @@ def _install_generated_plan( path=[ JointState( name=config.joint_names, - position=list(point), + position=np.asarray(list(point), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) for point in points ], @@ -122,7 +139,13 @@ def _make_module_with_monitor(module_factory) -> ManipulationModule: def _make_joint_state(positions: list[float], name: list[str] | None = None) -> JointState: - return JointState(name=name or [f"j{i}" for i in range(len(positions))], position=positions) + return JointState( + name=name or [f"j{i}" for i in range(len(positions))], + position=np.asarray(positions, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) def _make_path(*points: list[float]) -> list[JointState]: @@ -136,7 +159,11 @@ def _make_trajectory(*points: tuple[float, list[float]]) -> JointTrajectory: joint_names=joint_names, points=[ JointTrajectoryPoint( - time_from_start=duration_from_seconds(time_from_start), positions=positions + time_from_start=duration_from_seconds(time_from_start), + positions=np.asarray(positions, dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) for time_from_start, positions in points ], @@ -197,15 +224,16 @@ def test_routes_canonical_state_in_model_order(self, canonical_model_config, mod msg = JointState( header=Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id="robot"), name=["left/joint3", "unrelated", "left/joint1", "left/joint2"], - position=[0.3, 9.0, 0.1, 0.2], - velocity=[3.0, 9.0, 1.0, 2.0], + position=np.array([0.3, 9.0, 0.1, 0.2], dtype=np.float64), + velocity=np.array([3.0, 9.0, 1.0, 2.0], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) module._on_joint_state(msg) state = module._world_monitor.on_joint_state.call_args.args[0] assert state.name == canonical_model_config.joint_names - assert state.position == [0.1, 0.2, 0.3] - assert state.velocity == [1.0, 2.0, 3.0] + np.testing.assert_array_equal(state.position, [0.1, 0.2, 0.3]) + np.testing.assert_array_equal(state.velocity, [1.0, 2.0, 3.0]) assert state.header == msg.header msg.header.frame_id = "changed" assert state.header.frame_id == "robot" @@ -219,13 +247,16 @@ def test_aliases_coordinator_names_onto_model_joints( msg = JointState( name=["left/joint3", "chassis/vx", "left/joint2"], - position=[0.3, 0.1, 0.2], + position=np.array([0.3, 0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) module._on_joint_state(msg) state = module._world_monitor.on_joint_state.call_args.args[0] assert state.name == canonical_model_config.joint_names - assert state.position == [0.1, 0.2, 0.3] + np.testing.assert_array_equal(state.position, [0.1, 0.2, 0.3]) def test_skips_incomplete_model_state(self, canonical_model_config, module_factory): module = _make_module_with_monitor(module_factory) @@ -233,7 +264,10 @@ def test_skips_incomplete_model_state(self, canonical_model_config, module_facto msg = JointState( name=["left/joint1", "left/joint2"], - position=[0.5, 0.6], + position=np.array([0.5, 0.6], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) module._on_joint_state(msg) @@ -245,19 +279,25 @@ def test_captures_init_joints_once(self, canonical_model_config, module_factory) first_msg = JointState( name=["left/joint1", "left/joint2", "left/joint3"], - position=[0.1, 0.2, 0.3], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) module._on_joint_state(first_msg) assert module._init_joints is not None - assert module._init_joints.position == [0.1, 0.2, 0.3] + np.testing.assert_array_equal(module._init_joints.position, [0.1, 0.2, 0.3]) # Second call should NOT overwrite second_msg = JointState( name=["left/joint1", "left/joint2", "left/joint3"], - position=[0.9, 0.8, 0.7], + position=np.array([0.9, 0.8, 0.7], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) module._on_joint_state(second_msg) - assert module._init_joints.position == [0.1, 0.2, 0.3] + np.testing.assert_array_equal(module._init_joints.position, [0.1, 0.2, 0.3]) def test_no_monitor_returns_early(self, canonical_model_config, module_factory): """When world_monitor is None, _on_joint_state returns without error.""" @@ -268,7 +308,10 @@ def test_no_monitor_returns_early(self, canonical_model_config, module_factory): # Should not raise msg = JointState( name=["left/joint1", "left/joint2", "left/joint3"], - position=[0.1, 0.2, 0.3], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) module._on_joint_state(msg) @@ -285,16 +328,20 @@ def test_visualization_routing_and_stop_all_monitors(self): monitor._viz_thread = MagicMock() monitor._viz_thread.is_alive.return_value = False monitor._world.get_live_context.return_value = object() - monitor._world.get_joint_state.return_value = JointState() + monitor._world.get_joint_state.return_value = JointState( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + position=np.array([], dtype=np.float64), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) assert monitor.get_visualization_url() == "123" monitor.update_visualization_state() monitor.cancel_preview_animation() path = _make_path([1.0], [2.0], [3.0]) plan = GeneratedPlan( - trajectory=JointTrajectory( - header=header_now(), - ), + trajectory=JointTrajectory(header=header_now(), joint_names=[], points=[]), group_ids=("manipulator",), status=PlanningStatus.SUCCESS, path=path, @@ -318,9 +365,7 @@ def test_visualization_none_is_noop(self): monitor.update_visualization_state() monitor.cancel_preview_animation() monitor.animate_trajectory( - JointTrajectory( - header=header_now(), - ), + JointTrajectory(header=header_now(), joint_names=[], points=[]), 1.0, ) monitor.start_visualization_thread() @@ -340,9 +385,7 @@ def test_preview_without_a_plan_returns_rejected(self, module_factory): def test_clear_planned_path_invalidates_before_dismissing_preview(self, module_factory): module = module_factory() plan = GeneratedPlan( - trajectory=JointTrajectory( - header=header_now(), - ), + trajectory=JointTrajectory(header=header_now(), joint_names=[], points=[]), group_ids=("manipulator",), path=[], ) @@ -362,9 +405,7 @@ def test_clear_planned_path_invalidates_before_dismissing_preview(self, module_f def test_clear_planned_path_clears_without_a_world_monitor(self, module_factory): module = module_factory() module._last_plan = GeneratedPlan( - trajectory=JointTrajectory( - header=header_now(), - ), + trajectory=JointTrajectory(header=header_now(), joint_names=[], points=[]), group_ids=("manipulator",), path=[], ) diff --git a/dimos/manipulation/test_manipulation_primitives.py b/dimos/manipulation/test_manipulation_primitives.py index 7c4d15dc7be..79376f64588 100644 --- a/dimos/manipulation/test_manipulation_primitives.py +++ b/dimos/manipulation/test_manipulation_primitives.py @@ -15,13 +15,16 @@ """Behavior tests for the primitive manipulation RPCs and skill adapter.""" from collections.abc import Iterator +from dataclasses import asdict from pathlib import Path from unittest.mock import MagicMock -from dimos_generated.geometry_msgs.msg import TransformStamped +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import numpy as np import pytest from pytest_mock import MockerFixture @@ -73,8 +76,20 @@ def _plan() -> GeneratedPlan: header=header_now(), joint_names=names, points=[ - JointTrajectoryPoint(positions=[0.0], time_from_start=duration_from_seconds(0.0)), - JointTrajectoryPoint(positions=[0.1], time_from_start=duration_from_seconds(1.0)), + JointTrajectoryPoint( + positions=np.array([0.0], dtype=np.float64), + time_from_start=duration_from_seconds(0.0), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointTrajectoryPoint( + positions=np.array([0.1], dtype=np.float64), + time_from_start=duration_from_seconds(1.0), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ], ), ) @@ -105,7 +120,14 @@ def test_move_linear_uses_world_relative_target_and_default_speed( assert result.succeeded targets, config = generate.call_args.args start, relative = targets["tool"] - assert start == TransformStamped(header=Header(frame_id="world"), child_frame_id="") + assert start == TransformStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + child_frame_id="", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) assert relative.transform.translation.x == pytest.approx(0.02) assert relative.transform.translation.y == pytest.approx(0.0) assert relative.transform.translation.z == pytest.approx(-0.01) @@ -117,9 +139,19 @@ def test_move_linear_uses_world_relative_target_and_default_speed( def test_get_state_returns_every_group_with_presets(module_factory) -> None: module = module_factory() _set_groups(module, _model(gripper=True, home=[0.3])) - module._init_joints = JointState(name=["j0"], position=[-0.2]) + module._init_joints = JointState( + name=["j0"], + position=np.array([-0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) module._world_monitor.current_group_joint_state.return_value = JointState( - name=["j0"], position=[0.1] + name=["j0"], + position=np.array([0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) module._world_monitor.get_group_ee_pose.return_value = None module._control_coordinator.task_invoke.return_value = [0.04] @@ -127,10 +159,21 @@ def test_get_state_returns_every_group_with_presets(module_factory) -> None: snapshot = module.get_state() group = snapshot.groups["tool"] - assert group.joints == JointState(name=["j0"], position=[0.1]) + np.testing.assert_equal( + asdict(group.joints), + asdict( + JointState( + name=["j0"], + position=np.array([0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ), + ) assert group.gripper_position == pytest.approx(0.04) - assert group.joint_presets["home"].position == [0.3] - assert group.joint_presets["init"].position == [-0.2] + np.testing.assert_array_equal(group.joint_presets["home"].position, [0.3]) + np.testing.assert_array_equal(group.joint_presets["init"].position, [-0.2]) module._control_coordinator.task_invoke.assert_called_once_with( "arm_gripper", "get_normalized", {} ) @@ -140,7 +183,11 @@ def test_get_state_skips_gripper_telemetry_without_control_hardware(module_facto module = module_factory() _set_groups(module, _model(gripper=False)) module._world_monitor.current_group_joint_state.return_value = JointState( - name=["j0"], position=[0.1] + name=["j0"], + position=np.array([0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) module._world_monitor.get_group_ee_pose.return_value = None @@ -252,5 +299,5 @@ def test_legacy_skill_adapter_delegates_to_primitive_rpcs( assert "COMPLETED" in result.message target = manipulation.plan_to_joints.call_args.args[0]["tool"] assert target.name == ["j0"] - assert target.position == [0.25] + np.testing.assert_array_equal(target.position, [0.25]) manipulation.execute.assert_called_once_with(blocking=True) diff --git a/dimos/manipulation/test_manipulation_spec.py b/dimos/manipulation/test_manipulation_spec.py index a9c383835d9..5d9ca6efd68 100644 --- a/dimos/manipulation/test_manipulation_spec.py +++ b/dimos/manipulation/test_manipulation_spec.py @@ -20,8 +20,11 @@ from types import UnionType from typing import get_args, get_origin, get_type_hints +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import numpy as np from dimos.manipulation.manipulation_module import ManipulationModule import dimos.manipulation.manipulation_spec as manipulation_spec @@ -97,10 +100,24 @@ def test_public_result_repr_does_not_truncate_message() -> None: def test_state_repr_exposes_values_for_repl_use() -> None: state = PlanningGroupState( - JointState(name=["arm/j0"], position=[0.25]), + JointState( + name=["arm/j0"], + position=np.array([0.25], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), None, 0.75, - {"home": JointState(name=["arm/j0"], position=[0.0])}, + { + "home": JointState( + name=["arm/j0"], + position=np.array([0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + }, ) snapshot = ManipulationSnapshot( timestamp=12.5, @@ -129,18 +146,36 @@ def test_plan_result_repr_summarizes_trajectory_without_dumping_points() -> None joint_names=["arm/j0"], points=[ JointTrajectoryPoint( - positions=[0.0], velocities=[0.0], time_from_start=duration_from_seconds(0.0) + positions=np.array([0.0], dtype=np.float64), + velocities=np.array([0.0], dtype=np.float64), + time_from_start=duration_from_seconds(0.0), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), JointTrajectoryPoint( - positions=[123.456], - velocities=[0.0], + positions=np.array([123.456], dtype=np.float64), + velocities=np.array([0.0], dtype=np.float64), time_from_start=duration_from_seconds(1.0), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ], ), path=[ - JointState(name=["arm/j0"], position=[0.0]), - JointState(name=["arm/j0"], position=[123.456]), + JointState( + name=["arm/j0"], + position=np.array([0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=["arm/j0"], + position=np.array([123.456], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ], status=PlanningStatus.SUCCESS, path_length=123.456, diff --git a/dimos/manipulation/test_manipulation_unit.py b/dimos/manipulation/test_manipulation_unit.py index e3c1838f6b6..6e5122758e1 100644 --- a/dimos/manipulation/test_manipulation_unit.py +++ b/dimos/manipulation/test_manipulation_unit.py @@ -16,10 +16,25 @@ from __future__ import annotations +from dataclasses import asdict from pathlib import Path import pickle from unittest.mock import ANY, DEFAULT, MagicMock, call +from dimos_generated.builtin_interfaces.msg import Duration, Time +from dimos_generated.dimos_msgs.msg import TrajectoryStatus +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.sensor_msgs.msg import JointState, PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint import numpy as np import pytest from pytest_mock import MockerFixture @@ -66,16 +81,9 @@ from dimos.manipulation.planning.trajectory_generator.simple_parametrizer import ( SimpleTrapezoidParametrizer, ) -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint -from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState, TrajectoryStatus +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import duration_from_seconds, header_now, time_from_seconds, to_nanoseconds +from dimos.msgs.trajectory import TrajectoryState from dimos.robot.assets.model import LoadedRobotModel, RobotModel @@ -95,7 +103,14 @@ def invoke(task: str, method: str, args: dict | None = None): return DEFAULT coordinator.task_invoke.side_effect = invoke - coordinator.task_invoke.return_value = TrajectoryStatus(state=TrajectoryState.IDLE) + coordinator.task_invoke.return_value = TrajectoryStatus( + header=header_now(), + state=TrajectoryState.IDLE, + progress=0.0, + time_elapsed=Duration(sec=0, nanosec=0), + time_remaining=Duration(sec=0, nanosec=0), + error="", + ) return coordinator @@ -104,7 +119,13 @@ def robot_config(): """Create a robot config for testing.""" return RobotModelConfig( model=RobotModel.from_file(Path("/path/to/robot.urdf")), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + base_pose=PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), joint_names=["joint1", "joint2", "joint3"], base_link="link_base", planning_groups=[ @@ -121,7 +142,13 @@ def robot_config(): def _one_joint_config() -> RobotModelConfig: return RobotModelConfig( model=RobotModel.from_file(Path("/path")), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + base_pose=PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), joint_names=["j0"], base_link="base_link", planning_groups=[ @@ -157,16 +184,22 @@ def _install_generated_plan( module._world_monitor.planning_groups = PlanningGroupRegistry(config.planning_groups) module._world_monitor.current_model_joint_state.return_value = JointState( name=config.joint_names, - position=[0.0 for _ in config.joint_names], + position=np.asarray([0.0 for _ in config.joint_names], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) module._last_plan = GeneratedPlan( trajectory=JointTrajectory( + header=header_now(), joint_names=config.joint_names, points=[ - TrajectoryPoint( - time_from_start=float(index), - positions=list(point), - velocities=[0.0 for _ in config.joint_names], + JointTrajectoryPoint( + time_from_start=duration_from_seconds(float(index)), + positions=np.asarray(list(point), dtype=np.float64), + velocities=np.asarray([0.0 for _ in config.joint_names], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) for index, point in enumerate(points) ], @@ -176,7 +209,10 @@ def _install_generated_plan( path=[ JointState( name=config.joint_names, - position=list(point), + position=np.asarray(list(point), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) for point in points ], @@ -186,12 +222,15 @@ def _install_generated_plan( def _generated_plan_trajectory(joint_names: list[str], *points: list[float]) -> JointTrajectory: return JointTrajectory( + header=header_now(), joint_names=joint_names, points=[ - TrajectoryPoint( - time_from_start=float(index), - positions=list(point), - velocities=[0.0 for _ in joint_names], + JointTrajectoryPoint( + time_from_start=duration_from_seconds(float(index)), + positions=np.asarray(list(point), dtype=np.float64), + velocities=np.asarray([0.0 for _ in joint_names], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) for index, point in enumerate(points) ], @@ -201,9 +240,16 @@ def _generated_plan_trajectory(joint_names: list[str], *points: list[float]) -> def _make_trajectory(*points: tuple[float, list[float]]) -> JointTrajectory: joint_names = [f"j{i}" for i in range(len(points[0][1]))] if points else [] return JointTrajectory( + header=header_now(), joint_names=joint_names, points=[ - TrajectoryPoint(time_from_start=time_from_start, positions=positions) + JointTrajectoryPoint( + time_from_start=duration_from_seconds(time_from_start), + positions=np.asarray(positions, dtype=np.float64), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) for time_from_start, positions in points ], ) @@ -242,10 +288,9 @@ class TestVoxelMap: @staticmethod def _cloud(points: list[list[float]], frame_id: str = "world") -> PointCloud2: - return PointCloud2.from_numpy( + return pointcloud_from_xyz( np.asarray(points, dtype=np.float32).reshape((-1, 3)), - frame_id=frame_id, - timestamp=1.0, + header=Header(stamp=time_from_seconds(1.0), frame_id=frame_id), ) def test_a_map_becomes_one_octree_obstacle(self, module_factory) -> None: @@ -325,7 +370,13 @@ async def test_perception_objects_are_refreshable_and_queryable(self, module_fac module = module_factory() module._world_monitor = MagicMock(spec=WorldMonitor) detected = MagicMock() - pose = PoseStamped(position=Vector3(0.4, 0.1, 0.2)) + pose = PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.4, y=0.1, z=0.2), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) obstacle = Obstacle(name="object-1", pose=pose, obstacle_type=ObstacleType.BOX) module._world_monitor.refresh_obstacles.return_value = [{"object_id": "object-1"}] module._world_monitor.world.get_obstacles.return_value = [obstacle] @@ -343,8 +394,7 @@ def test_complete_update_forwards_new_obstacle_value(self, module_factory) -> No module._world_monitor = MagicMock(spec=WorldMonitor) module._world_monitor.update_obstacle.return_value = True pose = Pose( - position=Vector3(1.0, 2.0, 3.0), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) ) result = module.update_obstacle( @@ -361,7 +411,7 @@ def test_complete_update_forwards_new_obstacle_value(self, module_factory) -> No assert obstacle.obstacle_type == ObstacleType.SPHERE assert obstacle.dimensions == (0.4,) assert obstacle.color == (0.1, 0.2, 0.3, 0.9) - assert obstacle.pose.position.x == pytest.approx(1.0) + assert obstacle.pose.pose.position.x == pytest.approx(1.0) assert module.update_obstacle("default-color", pose, "box", [1.0, 1.0, 1.0]) default_obstacle = module._world_monitor.update_obstacle.call_args.args[0] @@ -372,8 +422,7 @@ def test_pose_update_forwards_only_name_and_pose(self, module_factory) -> None: module._world_monitor = MagicMock(spec=WorldMonitor) module._world_monitor.update_obstacle_pose.return_value = True pose = Pose( - position=Vector3(4.0, 5.0, 6.0), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + position=Point(x=4.0, y=5.0, z=6.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) ) result = module.update_obstacle_pose("moving-shape", pose) @@ -381,9 +430,9 @@ def test_pose_update_forwards_only_name_and_pose(self, module_factory) -> None: assert result is True name, stamped = module._world_monitor.update_obstacle_pose.call_args.args assert name == "moving-shape" - assert stamped.position.x == pytest.approx(4.0) - assert stamped.position.y == pytest.approx(5.0) - assert stamped.position.z == pytest.approx(6.0) + assert stamped.pose.position.x == pytest.approx(4.0) + assert stamped.pose.position.y == pytest.approx(5.0) + assert stamped.pose.position.z == pytest.approx(6.0) def test_complete_update_rejects_unknown_shape_before_world_mutation( self, module_factory @@ -392,7 +441,15 @@ def test_complete_update_rejects_unknown_shape_before_world_mutation( module._world_monitor = MagicMock(spec=WorldMonitor) with pytest.raises(ValueError, match="Unknown obstacle shape"): - module.update_obstacle("shape", Pose(), "capsule", [1.0]) + module.update_obstacle( + "shape", + Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + "capsule", + [1.0], + ) module._world_monitor.update_obstacle.assert_not_called() @@ -403,9 +460,25 @@ def test_complete_update_rejects_incomplete_mesh_and_invalid_color( module._world_monitor = MagicMock(spec=WorldMonitor) with pytest.raises(ValueError, match="mesh_path required"): - module.update_obstacle("mesh", Pose(), "mesh") + module.update_obstacle( + "mesh", + Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + "mesh", + ) with pytest.raises(ValueError, match="four values"): - module.update_obstacle("box", Pose(), "box", [1.0, 1.0, 1.0], color=[1.0]) + module.update_obstacle( + "box", + Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + "box", + [1.0, 1.0, 1.0], + color=[1.0], + ) module._world_monitor.update_obstacle.assert_not_called() @@ -413,8 +486,28 @@ def test_updates_fail_when_world_monitor_is_unavailable(self, module_factory) -> module = module_factory() module._world_monitor = None - assert module.update_obstacle("box", Pose(), "box", [1.0, 1.0, 1.0]) is False - assert module.update_obstacle_pose("box", Pose()) is False + assert ( + module.update_obstacle( + "box", + Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + "box", + [1.0, 1.0, 1.0], + ) + is False + ) + assert ( + module.update_obstacle_pose( + "box", + Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + is False + ) class TestStateMachine: @@ -440,7 +533,9 @@ def test_cancel_hides_active_plan_preview(self, module_factory): module = module_factory() module._state = ManipulationState.EXECUTING module._last_plan = GeneratedPlan( - trajectory=JointTrajectory(), group_ids=("manipulator",), path=[] + trajectory=JointTrajectory(header=header_now(), joint_names=[], points=[]), + group_ids=("manipulator",), + path=[], ) module._world_monitor = MagicMock() @@ -453,7 +548,14 @@ def test_cancel_completed_execution_cancels_coordinator_task(self, module_factor config = _one_joint_config() _install_generated_plan(module, config, [0.0], [0.1]) coordinator = _control_coordinator(cancel_status=TrajectoryCancellationStatus.CANCELLED) - coordinator.task_invoke.return_value = TrajectoryStatus(state=TrajectoryState.ABORTED) + coordinator.task_invoke.return_value = TrajectoryStatus( + header=header_now(), + state=TrajectoryState.ABORTED, + progress=0.0, + time_elapsed=Duration(sec=0, nanosec=0), + time_remaining=Duration(sec=0, nanosec=0), + error="", + ) module._control_coordinator = coordinator module._initialize_execution() @@ -514,7 +616,12 @@ def test_reset_cancels_an_active_trajectory(self, module_factory): cancel_status=TrajectoryCancellationStatus.CANCELLED ) module._control_coordinator.task_invoke.return_value = TrajectoryStatus( - state=TrajectoryState.ABORTED + header=header_now(), + state=TrajectoryState.ABORTED, + progress=0.0, + time_elapsed=Duration(sec=0, nanosec=0), + time_remaining=Duration(sec=0, nanosec=0), + error="", ) module._initialize_execution() module.execute(blocking=False) @@ -676,15 +783,30 @@ def test_inverse_kinematics_uses_current_seed_without_planning( module._world_monitor.world = MagicMock() module._world_monitor.world.get_prepared_model.return_value = _prepared_model(robot_config) module._world_monitor.planning_groups = PlanningGroupRegistry(robot_config.planning_groups) - current = JointState(name=robot_config.joint_names, position=[0.0, 0.0, 0.0]) + current = JointState( + name=robot_config.joint_names, + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) current_model_state = JointState( name=["joint1", "joint2", "joint3"], - position=[0.0, 0.0, 0.0], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) module._world_monitor.current_model_joint_state.return_value = current_model_state expected = IKResult( status=IKStatus.SUCCESS, - joint_state=JointState(name=robot_config.joint_names, position=[0.1, 0.2, 0.3]), + joint_state=JointState( + name=robot_config.joint_names, + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), position_error=0.0001, orientation_error=0.0002, iterations=3, @@ -694,7 +816,11 @@ def test_inverse_kinematics_uses_current_seed_without_planning( module._kinematics.solve_pose_targets.return_value = expected pose = PoseStamped( - frame_id="world", position=Vector3(x=0.45, y=0.0, z=0.25), orientation=Quaternion() + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.45, y=0.0, z=0.25), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) module._state = operation_state result = module.inverse_kinematics({"manipulator": pose}) @@ -706,12 +832,12 @@ def test_inverse_kinematics_uses_current_seed_without_planning( _, kwargs = module._kinematics.solve_pose_targets.call_args assert kwargs["world"] is module._world_monitor.world assert kwargs["seed"].name == current_model_state.name - assert kwargs["seed"].position == current.position + np.testing.assert_array_equal(kwargs["seed"].position, current.position) assert kwargs["check_collision"] is True [(group, target_pose)] = kwargs["pose_targets"].items() assert group.id == "manipulator" - assert target_pose.frame_id == "world" - assert target_pose.position.x == 0.45 + assert target_pose.header.frame_id == "world" + assert target_pose.pose.position.x == 0.45 def test_inverse_kinematics_returns_failure_without_joint_state( self, robot_config, module_factory @@ -722,12 +848,20 @@ def test_inverse_kinematics_returns_failure_without_joint_state( module._world_monitor = MagicMock() module._world_monitor.planning_groups = PlanningGroupRegistry(robot_config.planning_groups) module._world_monitor.current_model_joint_state.return_value = JointState( - name=[], position=[] + name=[], + position=np.array([], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) module._kinematics = MagicMock() pose = PoseStamped( - frame_id="world", position=Vector3(x=0.45, y=0.0, z=0.25), orientation=Quaternion() + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.45, y=0.0, z=0.25), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) result = module.inverse_kinematics({"manipulator": pose}) @@ -747,10 +881,18 @@ def test_inverse_kinematics_accepts_explicit_seed_without_current_state( module._world_monitor.world.get_prepared_model.return_value = _prepared_model(robot_config) module._world_monitor.planning_groups = PlanningGroupRegistry(robot_config.planning_groups) module._world_monitor.current_model_joint_state.return_value = None - explicit_seed = JointState(name=robot_config.joint_names, position=[0.2, 0.1, 0.0]) + explicit_seed = JointState( + name=robot_config.joint_names, + position=np.array([0.2, 0.1, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) pending_plan = GeneratedPlan( group_ids=("manipulator",), - trajectory=JointTrajectory(joint_names=robot_config.joint_names), + trajectory=JointTrajectory( + header=header_now(), joint_names=robot_config.joint_names, points=[] + ), path=[explicit_seed], status=PlanningStatus.SUCCESS, ) @@ -761,7 +903,11 @@ def test_inverse_kinematics_accepts_explicit_seed_without_current_state( module._kinematics.solve_pose_targets.return_value = expected pose = PoseStamped( - frame_id="world", position=Vector3(x=0.45, y=0.0, z=0.25), orientation=Quaternion() + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.45, y=0.0, z=0.25), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) result = module.inverse_kinematics({"manipulator": pose}, seed=explicit_seed) @@ -802,16 +948,25 @@ def test_generate_joint_plan_stores_generated_plan(self, robot_config, module_fa module._world_monitor.planning_groups = registry module._world_monitor.current_model_joint_state.return_value = JointState( name=["joint1", "joint2", "joint3"], - position=[0.0, 0.0, 0.0], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) result_path = [ JointState( name=["joint1", "joint2", "joint3"], - position=[0.0, 0.0, 0.0], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), JointState( name=["joint1", "joint2", "joint3"], - position=[0.1, 0.2, 0.3], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ] module._planner = MagicMock() @@ -828,7 +983,10 @@ def test_generate_joint_plan_stores_generated_plan(self, robot_config, module_fa { "manipulator": JointState( name=robot_config.joint_names, - position=[0.1, 0.2, 0.3], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) } ) @@ -836,8 +994,13 @@ def test_generate_joint_plan_stores_generated_plan(self, robot_config, module_fa assert plan is not None assert module._last_plan is not None assert module._last_plan.group_ids == ("manipulator",) - assert module._last_plan.path == result_path - assert module._last_plan.trajectory.points[-1].positions == [0.1, 0.2, 0.3] + np.testing.assert_equal( + [asdict(value) for value in module._last_plan.path], + [asdict(value) for value in result_path], + ) + np.testing.assert_array_equal( + module._last_plan.trajectory.points[-1].positions, [0.1, 0.2, 0.3] + ) module._planner.plan_selected_joint_path.assert_called_once() _, kwargs = module._planner.plan_selected_joint_path.call_args assert kwargs["selection"].group_ids == ("manipulator",) @@ -851,7 +1014,10 @@ def test_generate_joint_plan_stores_generated_plan(self, robot_config, module_fa { "manipulator": JointState( name=robot_config.joint_names, - position=[0.1, 0.2, 0.3], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) } ) @@ -870,11 +1036,17 @@ def test_generate_pose_plan_uses_group_ik_and_selected_path(self, robot_config, module._world_monitor.planning_groups = registry module._world_monitor.current_model_joint_state.return_value = JointState( name=["joint1", "joint2", "joint3"], - position=[0.0, 0.0, 0.0], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) ik_goal = JointState( name=["joint1", "joint2", "joint3"], - position=[0.1, 0.2, 0.3], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) module._kinematics = MagicMock() module._kinematics.solve_pose_targets.return_value = IKResult( @@ -887,12 +1059,21 @@ def test_generate_pose_plan_uses_group_ik_and_selected_path(self, robot_config, path=[ JointState( name=["joint1", "joint2", "joint3"], - position=[0.0, 0.0, 0.0], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ik_goal, ], ) - pose = Pose(position=Vector3(x=0.45, y=0.0, z=0.25), orientation=Quaternion()) + pose = PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.45, y=0.0, z=0.25), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) plan = module.generate_plan_to_pose_targets({"manipulator": pose}) @@ -902,7 +1083,7 @@ def test_generate_pose_plan_uses_group_ik_and_selected_path(self, robot_config, target_groups = list(ik_kwargs["pose_targets"].keys()) assert [group.id for group in target_groups] == ["manipulator"] target_pose = ik_kwargs["pose_targets"][target_groups[0]] - assert target_pose.position.x == 0.45 + assert target_pose.pose.position.x == 0.45 assert ik_kwargs["seed"].name == [ "joint1", "joint2", @@ -920,7 +1101,10 @@ def test_failed_plan_materialization_clears_generated_plan(self, robot_config, m module._world_monitor.planning_groups = registry module._world_monitor.current_model_joint_state.return_value = JointState( name=["joint1", "joint2", "joint3"], - position=[0.0, 0.0, 0.0], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) module._world_monitor.current_model_joint_state.return_value = None module._last_plan = GeneratedPlan( @@ -934,11 +1118,17 @@ def test_failed_plan_materialization_clears_generated_plan(self, robot_config, m path=[ JointState( name=["joint1", "joint2", "joint3"], - position=[0.0, 0.0, 0.0], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), JointState( name=["joint1", "joint2", "joint3"], - position=[0.2, 0.2, 0.2], + position=np.array([0.2, 0.2, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ], ) @@ -948,13 +1138,24 @@ def test_failed_plan_materialization_clears_generated_plan(self, robot_config, m path=[ JointState( name=["joint1", "joint2", "joint3"], - position=[0.1, 0.2, 0.3], + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) ], ) plan = module.generate_plan_to_joint_targets( - {"manipulator": JointState(position=[0.1, 0.2, 0.3])} + { + "manipulator": JointState( + position=np.array([0.1, 0.2, 0.3], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + } ) assert plan is None @@ -968,7 +1169,13 @@ def test_get_state_exposes_selected_group_init_preset( model = _bimanual_config() module = module_factory() module.config.model = model - module._init_joints = JointState(name=model.joint_names, position=[0.1, -0.1]) + module._init_joints = JointState( + name=model.joint_names, + position=np.array([0.1, -0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) module._world_monitor = MagicMock(spec=WorldMonitor) module._world_monitor.planning_groups = PlanningGroupRegistry(model.planning_groups) module._world_monitor.current_group_joint_state.return_value = None @@ -977,7 +1184,7 @@ def test_get_state_exposes_selected_group_init_preset( preset = module.get_state().groups["left_arm"].joint_presets["init"] assert preset.name == ["left/j1"] - assert preset.position == [0.1] + np.testing.assert_array_equal(preset.position, [0.1]) def test_tf_loop_publishes_every_pose_group_for_bimanual_model( self, module_factory, mocker: MockerFixture @@ -987,11 +1194,29 @@ def test_tf_loop_publishes_every_pose_group_for_bimanual_model( module.config.model = model module._world_monitor = MagicMock(spec=WorldMonitor) module._world_monitor.planning_groups = PlanningGroupRegistry(model.planning_groups) - state = JointState(ts=123.5, name=model.joint_names, position=[0.0, 0.0]) + state = JointState( + header=Header(stamp=Time(sec=123, nanosec=500000000), frame_id=""), + name=model.joint_names, + position=np.array([0.0, 0.0]), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) module._world_monitor.current_model_joint_state.return_value = state module._world_monitor.get_group_ee_pose.side_effect = [ - PoseStamped(position=Vector3(0.4, 0.2, 0.3)), - PoseStamped(position=Vector3(0.4, -0.2, 0.3)), + PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.4, y=0.2, z=0.3), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.4, y=-0.2, z=0.3), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), ] publish = mocker.patch.object(module.tf, "publish") @@ -1019,15 +1244,30 @@ def test_tf_loop_publishes_mount_edges_alongside_the_robot( module = module_factory() module.config.model = model module.config.static_transforms = [ - Transform(frame_id="left/tool", child_frame_id="camera_link") + TransformStamped( + header=Header(frame_id="left/tool", stamp=Time(sec=0, nanosec=0)), + child_frame_id="camera_link", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) ] module._world_monitor = MagicMock(spec=WorldMonitor) module._world_monitor.planning_groups = PlanningGroupRegistry(model.planning_groups) module._world_monitor.current_model_joint_state.return_value = JointState( - ts=123.5, name=model.joint_names, position=[0.0, 0.0] + header=Header(stamp=Time(sec=123, nanosec=500000000), frame_id=""), + name=model.joint_names, + position=np.array([0.0, 0.0]), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) module._world_monitor.get_group_ee_pose.return_value = PoseStamped( - ts=123.5, position=Vector3(0.4, 0.2, 0.3) + header=Header(frame_id="", stamp=Time(sec=123, nanosec=500000000)), + pose=Pose( + position=Point(x=0.4, y=0.2, z=0.3), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) publish = mocker.patch.object(module.tf, "publish") @@ -1040,11 +1280,11 @@ def stop_after_first_iteration(_period: float) -> bool: module._tf_publish_loop() - published = list(publish.call_args.args[0]) + published = list(publish.call_args.args[0].transforms) mount = next(t for t in published if t.child_frame_id == "camera_link") - assert mount.frame_id == "left/tool" - assert mount.ts == 123.5 - assert {tf.ts for tf in published} == {123.5} + assert mount.header.frame_id == "left/tool" + assert mount.header.stamp == Time(sec=123, nanosec=500000000) + assert {to_nanoseconds(tf.header.stamp) for tf in published} == {123500000000} def test_get_ee_pose_fails_safely_without_pose_group(self, robot_config, module_factory): no_pose_config = RobotModelConfig( @@ -1108,14 +1348,27 @@ def test_inverse_kinematics_preserves_backend_failure_detail( module._world_monitor.planning_groups = PlanningGroupRegistry(robot_config.planning_groups) module._world_monitor.current_model_joint_state.return_value = JointState( name=robot_config.joint_names, - position=[0.0, 0.0, 0.0], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) module._kinematics = MagicMock() module._kinematics.solve_pose_targets.return_value = IKResult( status=IKStatus.NO_SOLUTION, message="target is outside the workspace" ) - result = module.inverse_kinematics({"manipulator": PoseStamped(frame_id="world")}) + result = module.inverse_kinematics( + { + "manipulator": PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + } + ) assert result.status == IKStatus.NO_SOLUTION assert result.message == "target is outside the workspace" @@ -1132,7 +1385,10 @@ def test_planner_failure_preserves_backend_detail(self, robot_config, module_fac module._world_monitor.planning_groups = PlanningGroupRegistry(robot_config.planning_groups) module._world_monitor.current_model_joint_state.return_value = JointState( name=robot_config.joint_names, - position=[0.0, 0.0, 0.0], + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) module._planner = MagicMock() module._planner.plan_selected_joint_path.return_value = PlanningResult( @@ -1142,7 +1398,10 @@ def test_planner_failure_preserves_backend_detail(self, robot_config, module_fac { "manipulator": JointState( name=robot_config.joint_names, - position=[1.0, 1.0, 1.0], + position=np.array([1.0, 1.0, 1.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) } ) diff --git a/dimos/manipulation/test_pick_and_place_unit.py b/dimos/manipulation/test_pick_and_place_unit.py index 858d10b857e..73f13d715f3 100644 --- a/dimos/manipulation/test_pick_and_place_unit.py +++ b/dimos/manipulation/test_pick_and_place_unit.py @@ -17,8 +17,13 @@ from typing import Any from unittest.mock import MagicMock +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.dimos_msgs.msg import GraspCandidate, GraspCandidateArray +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.std_msgs.msg import Header import pytest +from dimos.core.transport_factory import rpc_backend from dimos.manipulation.grasp_verification import GripperSettle from dimos.manipulation.manipulation_skills import ManipulationSkills from dimos.manipulation.manipulation_spec import ( @@ -31,13 +36,8 @@ PlanStatus, ) from dimos.manipulation.pick_and_place_module import PickAndPlaceModule -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.manipulation_msgs.GraspCandidate import GraspCandidate -from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray -from dimos.msgs.std_msgs.Header import Header +from dimos.msgs.geometry import quaternion_euler, quaternion_from_euler +from dimos.msgs.time import time_from_seconds PLANNED = PlanResult(PlanStatus.SUCCEEDED) NO_PATH = PlanResult(PlanStatus.FAILED, "unreachable") @@ -50,7 +50,9 @@ def _linear_move(execution: ExecutionResult = COMPLETED) -> MoveResult: @pytest.fixture -def module() -> Iterator[PickAndPlaceModule]: +def module(mocker) -> Iterator[PickAndPlaceModule]: + for method in ("start", "serve_module_rpc", "stop"): + mocker.patch.object(rpc_backend(), method) instance = PickAndPlaceModule(planning_frame="world") instance._scene = MagicMock() instance._grasp_generator = MagicMock() @@ -63,7 +65,11 @@ def module() -> Iterator[PickAndPlaceModule]: "arm/tool": SimpleNamespace( gripper_position=0.5, end_effector_pose=PoseStamped( - frame_id="world", orientation=Quaternion.from_euler(Vector3(0, 0, 0.7)) + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + orientation=quaternion_from_euler(0, 0, 0.7), + position=Point(x=0.0, y=0.0, z=0.0), + ), ), ) } @@ -75,7 +81,7 @@ def module() -> Iterator[PickAndPlaceModule]: instance._objects = {"cup-1": {"object_id": "cup-1", "name": "cup"}} instance._scene.get_object_pointcloud_by_object_id.return_value = MagicMock() instance._grasp_generator.propose_grasps.return_value = GraspCandidateArray( - Header(1.0, "world"), [_candidate(0.1)] + header=Header(stamp=time_from_seconds(1.0), frame_id="world"), candidates=[_candidate(0.1)] ) yield instance instance.stop() @@ -92,8 +98,11 @@ def settle(read: Any, target: float, config: Any, **_: Any) -> GripperSettle: def _candidate(x: float, score: float = 1.0) -> GraspCandidate: return GraspCandidate( - Pose(Vector3(x, 0.0, 0.2), Quaternion.from_euler(Vector3(-3.141592653589793, 0.0, 0.0))), - score, + pose=Pose( + position=Point(x=x, y=0.0, z=0.2), + orientation=quaternion_from_euler(-3.141592653589793, 0.0, 0.0), + ), + score=score, ) @@ -121,7 +130,7 @@ def test_pick_object_uses_first_provider_candidate( grasp_generator: Any = module._grasp_generator first, second = _candidate(0.1, 0.1), _candidate(0.2, 0.9) grasp_generator.propose_grasps.return_value = GraspCandidateArray( - Header(1.0, "world"), [first, second] + header=Header(stamp=time_from_seconds(1.0), frame_id="world"), candidates=[first, second] ) result = module.pick_object("cup-1") @@ -129,14 +138,14 @@ def test_pick_object_uses_first_provider_candidate( assert result.message == "Pick complete" assert module.get_grasp_candidates().candidates == [first, second] assert module._selected_grasp is not None - assert module._selected_grasp.position.x == pytest.approx(0.1) + assert module._selected_grasp.pose.position.x == pytest.approx(0.1) assert result.metadata["rank"] == 0 def test_pick_object_rejects_non_planning_frame(module: PickAndPlaceModule) -> None: grasp_generator: Any = module._grasp_generator grasp_generator.propose_grasps.return_value = GraspCandidateArray( - Header(1.0, "camera"), [_candidate(0.1)] + header=Header(stamp=time_from_seconds(1.0), frame_id="camera"), candidates=[_candidate(0.1)] ) with pytest.raises(RuntimeError, match="frame 'camera'.*planning frame is 'world'"): @@ -149,7 +158,8 @@ def test_pick_falls_through_to_the_next_reachable_candidate( """A learned provider's best-scoring pose is not always kinematically reachable.""" manipulation: Any = module._manipulation module._grasp_generator.propose_grasps.return_value = GraspCandidateArray( - Header(1.0, "world"), [_candidate(0.1, score=0.9), _candidate(0.3, score=0.4)] + header=Header(stamp=time_from_seconds(1.0), frame_id="world"), + candidates=[_candidate(0.1, score=0.9), _candidate(0.3, score=0.4)], ) manipulation.plan_to_poses.side_effect = [NO_PATH, PLANNED, PLANNED, PLANNED] @@ -164,7 +174,8 @@ def test_pick_stops_walking_candidates_on_a_drive_fault(module: PickAndPlaceModu """An execution fault would repeat for every candidate, so it is not a demotion.""" manipulation: Any = module._manipulation module._grasp_generator.propose_grasps.return_value = GraspCandidateArray( - Header(1.0, "world"), [_candidate(0.1), _candidate(0.3)] + header=Header(stamp=time_from_seconds(1.0), frame_id="world"), + candidates=[_candidate(0.1), _candidate(0.3)], ) manipulation.execute.return_value = ExecutionResult(ExecutionStatus.FAULT, "drive fault") @@ -183,7 +194,8 @@ def test_pick_reports_no_reachable_candidate_when_every_attempt_fails( ) -> None: manipulation: Any = module._manipulation module._grasp_generator.propose_grasps.return_value = GraspCandidateArray( - Header(1.0, "world"), [_candidate(0.1), _candidate(0.3)] + header=Header(stamp=time_from_seconds(1.0), frame_id="world"), + candidates=[_candidate(0.1), _candidate(0.3)], ) manipulation.plan_to_poses.return_value = NO_PATH @@ -210,7 +222,7 @@ def test_proposals_reach_the_viewer_as_they_are_generated(module: PickAndPlaceMo def test_pick_object_rejects_empty_candidates(module: PickAndPlaceModule) -> None: module._grasp_generator.propose_grasps.return_value = GraspCandidateArray( - Header(1.0, "world"), [] + header=Header(stamp=time_from_seconds(1.0), frame_id="world"), candidates=[] ) result = module.pick_object("cup-1") @@ -222,14 +234,14 @@ def test_pick_object_rejects_empty_candidates(module: PickAndPlaceModule) -> Non def test_pick_preserves_current_yaw_when_configured(module: PickAndPlaceModule) -> None: module.config.yaw_policy = "preserve_current" module._grasp_generator.propose_grasps.return_value = GraspCandidateArray( - Header(1.0, "world"), - [ + header=Header(stamp=time_from_seconds(1.0), frame_id="world"), + candidates=[ GraspCandidate( - Pose( - Vector3(0.1, 0.0, 0.2), - Quaternion.from_euler(Vector3(-3.141592653589793, 0.0, 0.1)), + pose=Pose( + position=Point(x=0.1, y=0.0, z=0.2), + orientation=quaternion_from_euler(-3.141592653589793, 0.0, 0.1), ), - 1.0, + score=1.0, ) ], ) @@ -238,15 +250,18 @@ def test_pick_preserves_current_yaw_when_configured(module: PickAndPlaceModule) assert result.message == "Pick complete" assert module._selected_grasp is not None - assert module._selected_grasp.orientation.to_euler().z == pytest.approx(0.7) + assert quaternion_euler(module._selected_grasp.pose.orientation)[2] == pytest.approx(0.7) assert module._holding_object def test_place_uses_local_axis_and_clears_held_state(module: PickAndPlaceModule) -> None: manipulation: Any = module._manipulation module._selected_grasp = PoseStamped( - frame_id="world", - orientation=Quaternion.from_euler(Vector3(-3.141592653589793, 0.0, 0.0)), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + orientation=quaternion_from_euler(-3.141592653589793, 0.0, 0.0), + position=Point(x=0.0, y=0.0, z=0.0), + ), ) module._holding_object = True @@ -254,14 +269,19 @@ def test_place_uses_local_axis_and_clears_held_state(module: PickAndPlaceModule) assert result.message == "Place complete" preplace = manipulation.plan_to_poses.call_args_list[0].args[0]["arm/tool"] - assert preplace.position.z == pytest.approx(0.3) + assert preplace.pose.position.z == pytest.approx(0.3) assert not module._holding_object assert module._selected_grasp is None def test_scan_failure_clears_stale_selection(module: PickAndPlaceModule) -> None: scene: Any = module._scene - module._selected_grasp = PoseStamped(frame_id="world") + module._selected_grasp = PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + ) scene.scan_scene.side_effect = RuntimeError("No aligned RGB-D frame") with pytest.raises(RuntimeError, match="No aligned RGB-D frame"): @@ -275,7 +295,12 @@ def test_pick_rejects_when_already_holding(module: PickAndPlaceModule) -> None: manipulation: Any = module._manipulation module._holding_object = True module._selected_object_id = "cup-0" - module._selected_grasp = PoseStamped(frame_id="world") + module._selected_grasp = PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + ) pick = module.pick_object("cup-1") @@ -287,7 +312,12 @@ def test_pick_rejects_when_already_holding(module: PickAndPlaceModule) -> None: def test_failed_pick_clears_previous_selection(module: PickAndPlaceModule) -> None: - module._selected_grasp = PoseStamped(frame_id="world") + module._selected_grasp = PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + ) result = module.pick_object("missing") @@ -413,7 +443,12 @@ def settle(read: Any, target: float, config: Any, **_: Any) -> GripperSettle: def test_place_retains_held_state_when_release_fails( module: PickAndPlaceModule, monkeypatch: pytest.MonkeyPatch ) -> None: - module._selected_grasp = PoseStamped(frame_id="world") + module._selected_grasp = PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + ) module._holding_object = True def settle(read: Any, target: float, config: Any, **_: Any) -> GripperSettle: @@ -445,3 +480,15 @@ def test_motion_skills_declare_movement_capability() -> None: ] assert all(skill.__skill_uses__ == ["movement"] for skill in skills) + + +@pytest.mark.parametrize("score", [float("nan"), float("inf"), -float("inf")]) +def test_pick_rejects_invalid_score_before_commanding_motion(module, score): + module._grasp_generator.propose_grasps.return_value = GraspCandidateArray( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + candidates=[_candidate(0.1, score=score)], + ) + with pytest.raises(ValueError, match="GraspCandidate.score must be finite"): + module.pick_object("cup-1") + module._manipulation.plan_to_poses.assert_not_called() + module._manipulation.set_gripper_position.assert_not_called() diff --git a/dimos/manipulation/test_plan_execution.py b/dimos/manipulation/test_plan_execution.py index 0676e7e494d..21ffa2fe2ff 100644 --- a/dimos/manipulation/test_plan_execution.py +++ b/dimos/manipulation/test_plan_execution.py @@ -16,9 +16,12 @@ from unittest.mock import DEFAULT, MagicMock +from dimos_generated.builtin_interfaces.msg import Duration, Time from dimos_generated.dimos_msgs.msg import TrajectoryStatus from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import numpy as np import pytest from dimos.control.coordinator import ControlCoordinator @@ -45,14 +48,18 @@ def _plan(final_position: float = 1.0, joint_name: str = "arm/j0") -> GeneratedP joint_names=names, points=[ JointTrajectoryPoint( - positions=[0.0], - velocities=[0.0], + positions=np.array([0.0], dtype=np.float64), + velocities=np.array([0.0], dtype=np.float64), time_from_start=duration_from_seconds(0.0), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), JointTrajectoryPoint( - positions=[final_position], - velocities=[0.0], + positions=np.array([final_position], dtype=np.float64), + velocities=np.array([0.0], dtype=np.float64), time_from_start=duration_from_seconds(1.0), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ], ) @@ -60,8 +67,20 @@ def _plan(final_position: float = 1.0, joint_name: str = "arm/j0") -> GeneratedP group_ids=("manipulator",), trajectory=trajectory, path=[ - JointState(name=names, position=[0.0]), - JointState(name=names, position=[final_position]), + JointState( + name=names, + position=np.array([0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=names, + position=np.array([final_position], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ], status=PlanningStatus.SUCCESS, ) @@ -151,7 +170,14 @@ def test_execute_sends_planar_base_columns_to_the_base_task(module_factory) -> N coordinator.task_invoke.side_effect = lambda task, method, args: ( TrajectoryExecutionResult(TrajectoryExecutionStatus.ACCEPTED) if method == "execute" - else TrajectoryStatus(header=header_now(), state=state) + else TrajectoryStatus( + header=header_now(), + state=state, + progress=0.0, + time_elapsed=Duration(sec=0, nanosec=0), + time_remaining=Duration(sec=0, nanosec=0), + error="", + ) ) module = _module_with_coordinator(coordinator, module_factory) planar_base = PlanarBaseDefinition( @@ -167,13 +193,34 @@ def test_execute_sends_planar_base_columns_to_the_base_task(module_factory) -> N module._initialize_execution() names = ["arm/j0", *planar_base.joint_names] points = [ - JointTrajectoryPoint(positions=[0.0] * 4, time_from_start=duration_from_seconds(0.0)), - JointTrajectoryPoint(positions=[1.0] * 4, time_from_start=duration_from_seconds(1.0)), + JointTrajectoryPoint( + positions=np.asarray([0.0] * 4, dtype=np.float64), + time_from_start=duration_from_seconds(0.0), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointTrajectoryPoint( + positions=np.asarray([1.0] * 4, dtype=np.float64), + time_from_start=duration_from_seconds(1.0), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ] module._last_plan = GeneratedPlan( group_ids=("manipulator",), trajectory=JointTrajectory(header=header_now(), joint_names=names, points=points), - path=[JointState(name=names, position=point.positions) for point in points], + path=[ + JointState( + name=names, + position=np.asarray(point.positions, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + for point in points + ], status=PlanningStatus.SUCCESS, ) @@ -252,11 +299,23 @@ def read_status(): assert operator.execute(_plan()) is True coordinator.task_invoke.return_value = TrajectoryStatus( - header=header_now(), state=TrajectoryState.EXECUTING + header=header_now(), + state=TrajectoryState.EXECUTING, + progress=0.0, + time_elapsed=Duration(sec=0, nanosec=0), + time_remaining=Duration(sec=0, nanosec=0), + error="", ) assert read_status() == "EXECUTING" - coordinator.task_invoke.return_value = TrajectoryStatus(header=header_now(), state=terminal) + coordinator.task_invoke.return_value = TrajectoryStatus( + header=header_now(), + state=terminal, + progress=0.0, + time_elapsed=Duration(sec=0, nanosec=0), + time_remaining=Duration(sec=0, nanosec=0), + error="", + ) assert read_status() == operation assert module.get_state().execution_status.name == terminal.name coordinator.task_invoke.reset_mock() diff --git a/dimos/manipulation/test_planning_factory.py b/dimos/manipulation/test_planning_factory.py index 0b2f3376079..67170e89586 100644 --- a/dimos/manipulation/test_planning_factory.py +++ b/dimos/manipulation/test_planning_factory.py @@ -23,6 +23,7 @@ from typing import Any from unittest.mock import ANY +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.std_msgs.msg import Header from pydantic import ValidationError @@ -75,7 +76,11 @@ def robot_config() -> RobotModelConfig: return RobotModelConfig( model=RobotModel.from_file(Path("/path/to/robot.urdf")), base_pose=PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), # type: ignore[call-arg] joint_names=["joint1", "joint2"], planning_groups=[ diff --git a/dimos/manipulation/test_roboplan.py b/dimos/manipulation/test_roboplan.py index 61d2e8d0642..6567ebf3611 100644 --- a/dimos/manipulation/test_roboplan.py +++ b/dimos/manipulation/test_roboplan.py @@ -26,6 +26,7 @@ from typing import Any, ClassVar import xml.etree.ElementTree as ET +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, Pose, @@ -521,7 +522,11 @@ def robot_config(tmp_path: Path) -> RobotModelConfig: return RobotModelConfig( model=RobotModel.from_file(model_path).with_default_joint_acceleration_limit(2.0), base_pose=PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), # type: ignore[call-arg] joint_names=["joint1", "joint2"], base_link="base", @@ -549,7 +554,10 @@ def _make_world( world.sync_from_joint_state( JointState( name=list(robot_config.joint_names), - position=[0.0] * len(robot_config.joint_names), + position=np.asarray([0.0] * len(robot_config.joint_names), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ) planner_module = _import_roboplan_planner(fake_roboplan) @@ -610,13 +618,27 @@ def _selection( def _relative_target(*waypoints: TransformStamped) -> tuple[TransformStamped, ...]: - return (TransformStamped(header=Header(frame_id="world"), child_frame_id=""), *waypoints) + return ( + TransformStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + child_frame_id="", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + *waypoints, + ) def _absolute_target(*waypoints: PoseStamped) -> tuple[PoseStamped, ...]: return ( PoseStamped( - header=Header(frame_id="world"), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), *waypoints, ) @@ -701,17 +723,32 @@ def test_context_cloning_and_joint_state_round_trip( ) -> None: world = _make_world(fake_roboplan, robot_config) - live_state = JointState(name=["joint1", "joint2"], position=[0.1, 0.2]) + live_state = JointState( + name=["joint1", "joint2"], + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) world.sync_from_joint_state(live_state) with world.scratch_context() as scratch: scratch_state = world.get_joint_state(scratch) assert scratch_state.name == ["joint1", "joint2"] - assert scratch_state.position == [0.1, 0.2] - world.set_joint_state(scratch, JointState(name=["joint1", "joint2"], position=[0.3, 0.4])) + np.testing.assert_array_equal(scratch_state.position, [0.1, 0.2]) + world.set_joint_state( + scratch, + JointState( + name=["joint1", "joint2"], + position=np.array([0.3, 0.4], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + ) live_round_trip = world.get_joint_state(world.get_live_context()) - assert live_round_trip.position == [0.1, 0.2] + np.testing.assert_array_equal(live_round_trip.position, [0.1, 0.2]) def test_obstacle_mutation_updates_scene_and_stored_pose( @@ -731,7 +768,11 @@ def test_obstacle_mutation_updates_scene_and_stored_pose( name="box", obstacle_type=ObstacleType.BOX, pose=PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), # type: ignore[call-arg] dimensions=(0.1, 0.2, 0.3), ) @@ -739,8 +780,10 @@ def test_obstacle_mutation_updates_scene_and_stored_pose( assert add_box.call_args.args[2] == "dimos_world" assert "box" in world._scene.geometry updated_pose = PoseStamped( - header=Header(frame_id=""), - pose=Pose(position=Point(x=1, y=0, z=0), orientation=Quaternion()), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=1, y=0, z=0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), ) # type: ignore[call-arg] assert world.update_obstacle_pose( "box", @@ -766,7 +809,11 @@ def test_octree_obstacle_reaches_the_scene_as_occupied_cells( name="voxel-map", obstacle_type=ObstacleType.OCTREE, pose=PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), points=((0.0, 0.0, 0.0), (0.05, 0.0, 0.0)), octree_resolution=0.05, @@ -796,7 +843,11 @@ def test_octree_obstacle_survives_the_deepcopy_and_equality_the_world_does( name="voxel-map", obstacle_type=ObstacleType.OCTREE, pose=PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), points=((0.0, 0.0, 0.0),), octree_resolution=0.05, @@ -820,7 +871,11 @@ def octree(**kwargs: object) -> Obstacle: "name": "voxel-map", "obstacle_type": ObstacleType.OCTREE, "pose": PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), "points": ((0.0, 0.0, 0.0),), "octree_resolution": 0.05, @@ -851,7 +906,11 @@ def test_obstacle_operations_require_finalization( name="box", obstacle_type=ObstacleType.BOX, pose=PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), dimensions=(0.1, 0.2, 0.3), ) @@ -879,7 +938,11 @@ def test_failed_obstacle_add_rolls_back_and_can_be_retried( name="retryable", obstacle_type=ObstacleType.BOX, pose=PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), dimensions=(0.1, 0.2, 0.3), ) @@ -909,7 +972,11 @@ def test_concurrent_remove_waits_for_obstacle_add( name="concurrent", obstacle_type=ObstacleType.BOX, pose=PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), dimensions=(0.1, 0.2, 0.3), ) @@ -959,7 +1026,11 @@ def test_obstacle_ids_are_world_owned_and_invalid_insertions_are_rejected( name="", obstacle_type=ObstacleType.BOX, pose=PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), # type: ignore[call-arg] dimensions=(0.1, 0.1, 0.1), ) @@ -981,7 +1052,11 @@ def test_complete_update_rejects_invalid_obstacle_values( name="shape", obstacle_type=ObstacleType.BOX, pose=PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), # type: ignore[call-arg] dimensions=(0.1, 0.2, 0.3), ) @@ -997,7 +1072,7 @@ def test_complete_update_rejects_invalid_obstacle_values( replace( valid, pose=PoseStamped( - header=Header(frame_id=""), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=Point(x=np.nan, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), @@ -1023,7 +1098,11 @@ def test_complete_replacement_and_defensive_obstacle_snapshots( name="shape", obstacle_type=ObstacleType.BOX, pose=PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), # type: ignore[call-arg] dimensions=(0.1, 0.2, 0.3), color=(1.0, 0.0, 0.0, 1.0), @@ -1036,8 +1115,10 @@ def test_complete_replacement_and_defensive_obstacle_snapshots( name="shape", obstacle_type=ObstacleType.SPHERE, pose=PoseStamped( - header=Header(frame_id=""), - pose=Pose(position=Point(x=1, y=2, z=3), orientation=Quaternion()), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=1, y=2, z=3), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), ), # type: ignore[call-arg] dimensions=(0.4,), color=(0.0, 1.0, 0.0, 0.5), @@ -1067,7 +1148,11 @@ def test_collision_query_blocks_during_obstacle_replacement( name="box", obstacle_type=ObstacleType.BOX, pose=PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), # type: ignore[call-arg] dimensions=(0.1, 0.2, 0.3), ) @@ -1091,7 +1176,13 @@ def blocking_remove(obstacle_id: str) -> None: query_thread = threading.Thread( target=lambda: ( world.check_config_collision_free( - JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), + JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ), query_finished.set(), ) @@ -1114,7 +1205,11 @@ def test_obstacle_replacement_blocks_during_collision_query( name="box", obstacle_type=ObstacleType.BOX, pose=PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), # type: ignore[call-arg] dimensions=(0.1, 0.2, 0.3), ) @@ -1132,7 +1227,13 @@ def blocking_query(q: np.ndarray) -> bool: monkeypatch.setattr(world._scene, "hasCollisions", blocking_query) query_thread = threading.Thread( target=lambda: world.check_config_collision_free( - JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), + JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) ) query_thread.start() @@ -1161,7 +1262,11 @@ def test_native_update_failure_invalidates_world( name="box", obstacle_type=ObstacleType.BOX, pose=PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), # type: ignore[call-arg] dimensions=(0.1, 0.2, 0.3), ) @@ -1178,7 +1283,13 @@ def test_native_update_failure_invalidates_world( world.get_obstacles() with pytest.raises(RuntimeError, match="invalid"): world.check_config_collision_free( - JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), + JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) @@ -1192,7 +1303,11 @@ def test_native_pose_update_failure_invalidates_world( name="box", obstacle_type=ObstacleType.BOX, pose=PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), # type: ignore[call-arg] dimensions=(0.1, 0.2, 0.3), ) @@ -1214,8 +1329,20 @@ def test_collision_config_and_edge_checks( ) -> None: world = _make_world(fake_roboplan, robot_config) - safe = JointState(name=["joint1", "joint2"], position=[0.1, 0.2]) - colliding = JointState(name=["joint1", "joint2"], position=[0.95, 0.2]) + safe = JointState( + name=["joint1", "joint2"], + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + colliding = JointState( + name=["joint1", "joint2"], + position=np.array([0.95, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) assert world.check_config_collision_free(safe) assert not world.check_config_collision_free(colliding) @@ -1227,8 +1354,20 @@ def test_collision_check_uses_scene_queries( ) -> None: world = _make_world(fake_roboplan, robot_config) - safe = JointState(name=["joint1", "joint2"], position=[0.1, 0.2]) - colliding = JointState(name=["joint1", "joint2"], position=[0.95, 0.2]) + safe = JointState( + name=["joint1", "joint2"], + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + colliding = JointState( + name=["joint1", "joint2"], + position=np.array([0.95, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) assert world.check_config_collision_free(safe) assert not world.check_config_collision_free(colliding) @@ -1240,8 +1379,20 @@ def test_generic_rrt_planner_uses_roboplan_world_collision_checks( world = _make_world(fake_roboplan, robot_config) planner = RRTConnectPlanner(step_size=0.5, connect_step_size=0.5, goal_tolerance=10.0) - start = JointState(name=["joint1", "joint2"], position=[0.0, 0.0]) - goal = JointState(name=["joint1", "joint2"], position=[0.2, 0.1]) + start = JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + goal = JointState( + name=["joint1", "joint2"], + position=np.array([0.2, 0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) result = planner.plan_joint_path(world, start, goal, timeout=1.0, max_iterations=3) assert result.status == PlanningStatus.SUCCESS @@ -1258,7 +1409,11 @@ def test_generic_planner_allows_update_between_collision_checks( name="box", obstacle_type=ObstacleType.BOX, pose=PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), # type: ignore[call-arg] dimensions=(0.1, 0.2, 0.3), ) @@ -1279,8 +1434,20 @@ def checking_with_interleaved_update(state: JointState) -> bool: planner = RRTConnectPlanner(step_size=0.5, connect_step_size=0.5, goal_tolerance=10.0) result = planner.plan_joint_path( world, - JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), - JointState(name=["joint1", "joint2"], position=[0.2, 0.1]), + JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=["joint1", "joint2"], + position=np.array([0.2, 0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), timeout=1.0, max_iterations=3, ) @@ -1295,7 +1462,16 @@ def test_fk_jacobian_and_explicit_min_distance_unsupported( ) -> None: world = _make_world(fake_roboplan, robot_config) ctx = world.get_live_context() - world.set_joint_state(ctx, JointState(name=["joint1", "joint2"], position=[0.25, 0.5])) + world.set_joint_state( + ctx, + JointState( + name=["joint1", "joint2"], + position=np.array([0.25, 0.5], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + ) pose = world.get_ee_pose(ctx) assert pose.pose.position.x == pytest.approx(0.75) @@ -1346,7 +1522,13 @@ def fake_jacobian( ctx = world.get_live_context() world.set_joint_state( ctx, - JointState(name=["joint1", "joint2", "joint3"], position=[1.0, 2.0, 3.0]), + JointState( + name=["joint1", "joint2", "joint3"], + position=np.array([1.0, 2.0, 3.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) pose = world.get_group_ee_pose(ctx, "wrist") @@ -1384,7 +1566,16 @@ def test_group_jacobian_validates_projection_shape( ) -> None: world = _make_world(fake_roboplan, robot_config) ctx = world.get_live_context() - world.set_joint_state(ctx, JointState(name=["joint1", "joint2"], position=[0.0, 0.0])) + world.set_joint_state( + ctx, + JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + ) monkeypatch.setattr( FakeScene, @@ -1448,12 +1639,26 @@ def test_group_lookup_rejects_unknown_group_id( def test_native_planner_converts_path(fake_roboplan: None, robot_config: RobotModelConfig) -> None: world = _make_world(fake_roboplan, robot_config) - start = JointState(name=["joint1", "joint2"], position=[0.0, 0.0]) - goal = JointState(name=["joint1", "joint2"], position=[0.4, 0.2]) + start = JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + goal = JointState( + name=["joint1", "joint2"], + position=np.array([0.4, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) result = _planner_for(world).plan_joint_path(world, start, goal, timeout=1.0) assert result.status == PlanningStatus.SUCCESS - assert [state.position for state in result.path] == [[0.0, 0.0], [0.2, 0.1], [0.4, 0.2]] + np.testing.assert_array_equal( + [state.position for state in result.path], [[0.0, 0.0], [0.2, 0.1], [0.4, 0.2]] + ) assert [state.name for state in result.path] == [["joint1", "joint2"]] * 3 @@ -1487,12 +1692,26 @@ def endpoints_only(shortcutter: FakePathShortcutter, path: FakeJointPath) -> Fak result = _planner_for(world).plan_joint_path( world, - JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), - JointState(name=["joint1", "joint2"], position=[0.4, 0.2]), + JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=["joint1", "joint2"], + position=np.array([0.4, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), timeout=1.0, ) - assert [state.position for state in result.path] == [[0.0, 0.0], [0.4, 0.2]] + np.testing.assert_array_equal( + [state.position for state in result.path], [[0.0, 0.0], [0.4, 0.2]] + ) options = FakePathShortcutter.instances[-1].options assert options.group_name == "__dimos_all_configured__" assert options.max_step_size == 0.02 @@ -1517,12 +1736,26 @@ def test_native_planner_can_disable_path_shortcutting( result = _planner_for(world).plan_joint_path( world, - JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), - JointState(name=["joint1", "joint2"], position=[0.4, 0.2]), + JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=["joint1", "joint2"], + position=np.array([0.4, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), timeout=1.0, ) - assert [state.position for state in result.path] == [[0.0, 0.0], [0.2, 0.1], [0.4, 0.2]] + np.testing.assert_array_equal( + [state.position for state in result.path], [[0.0, 0.0], [0.2, 0.1], [0.4, 0.2]] + ) shortcut.assert_not_called() @@ -1541,13 +1774,27 @@ def test_native_planner_uses_raw_path_when_shortcutting_fails( result = _planner_for(world).plan_joint_path( world, - JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), - JointState(name=["joint1", "joint2"], position=[0.4, 0.2]), + JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=["joint1", "joint2"], + position=np.array([0.4, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), timeout=1.0, ) assert result.status == PlanningStatus.SUCCESS - assert [state.position for state in result.path] == [[0.0, 0.0], [0.2, 0.1], [0.4, 0.2]] + np.testing.assert_array_equal( + [state.position for state in result.path], [[0.0, 0.0], [0.2, 0.1], [0.4, 0.2]] + ) def test_native_planner_surfaces_unexpected_shortcutting_error( @@ -1566,8 +1813,20 @@ def test_native_planner_surfaces_unexpected_shortcutting_error( with pytest.raises(TypeError, match="unexpected shortcut integration error"): _planner_for(world).plan_joint_path( world, - JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), - JointState(name=["joint1", "joint2"], position=[0.4, 0.2]), + JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=["joint1", "joint2"], + position=np.array([0.4, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), timeout=1.0, ) @@ -1595,13 +1854,27 @@ def changed_endpoint(shortcutter: FakePathShortcutter, path: FakeJointPath) -> F result = _planner_for(world).plan_joint_path( world, - JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), - JointState(name=["joint1", "joint2"], position=[0.4, 0.2]), + JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=["joint1", "joint2"], + position=np.array([0.4, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), timeout=1.0, ) assert result.status == PlanningStatus.SUCCESS - assert [state.position for state in result.path] == [[0.0, 0.0], [0.2, 0.1], [0.4, 0.2]] + np.testing.assert_array_equal( + [state.position for state in result.path], [[0.0, 0.0], [0.2, 0.1], [0.4, 0.2]] + ) def test_native_planner_uses_raw_path_when_shortcutting_returns_empty_path( @@ -1619,13 +1892,27 @@ def test_native_planner_uses_raw_path_when_shortcutting_returns_empty_path( result = _planner_for(world).plan_joint_path( world, - JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), - JointState(name=["joint1", "joint2"], position=[0.4, 0.2]), + JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=["joint1", "joint2"], + position=np.array([0.4, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), timeout=1.0, ) assert result.status == PlanningStatus.SUCCESS - assert [state.position for state in result.path] == [[0.0, 0.0], [0.2, 0.1], [0.4, 0.2]] + np.testing.assert_array_equal( + [state.position for state in result.path], [[0.0, 0.0], [0.2, 0.1], [0.4, 0.2]] + ) def test_roboplan_planner_is_distinct_and_bound_to_world( @@ -1662,7 +1949,11 @@ def test_native_planning_blocks_obstacle_replacement( name="box", obstacle_type=ObstacleType.BOX, pose=PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(), orientation=Quaternion()) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), # type: ignore[call-arg] dimensions=(0.1, 0.2, 0.3), ) @@ -1682,8 +1973,20 @@ def blocking_plan( return original_plan(self, q_start, q_goal) monkeypatch.setattr(FakeRRT, "plan", blocking_plan) - start = JointState(name=["joint1", "joint2"], position=[0.0, 0.0]) - goal = JointState(name=["joint1", "joint2"], position=[0.2, 0.1]) + start = JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + goal = JointState( + name=["joint1", "joint2"], + position=np.array([0.2, 0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) planning_thread = threading.Thread( target=lambda: _planner_for(world).plan_joint_path(world, start, goal, timeout=1.0) ) @@ -1708,8 +2011,20 @@ def test_native_planner_names_path_from_robot_config_when_start_is_unnamed( ) -> None: world = _make_world(fake_roboplan, robot_config) - start = JointState(name=[], position=[0.0, 0.0]) - goal = JointState(name=["joint1", "joint2"], position=[0.4, 0.2]) + start = JointState( + name=[], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + goal = JointState( + name=["joint1", "joint2"], + position=np.array([0.4, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) result = _planner_for(world).plan_joint_path(world, start, goal, timeout=1.0) assert result.status == PlanningStatus.SUCCESS @@ -1725,14 +2040,28 @@ def test_native_selected_planner_returns_canonical_selected_joint_names( result = _planner_for(world).plan_selected_joint_path( world, selection, - JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), - JointState(name=["joint1", "joint2"], position=[0.4, 0.2]), + JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=["joint1", "joint2"], + position=np.array([0.4, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), timeout=1.0, ) assert result.status == PlanningStatus.SUCCESS assert [state.name for state in result.path] == [["joint1", "joint2"]] * 3 - assert [state.position for state in result.path] == [[0.0, 0.0], [0.2, 0.1], [0.4, 0.2]] + np.testing.assert_array_equal( + [state.position for state in result.path], [[0.0, 0.0], [0.2, 0.1], [0.4, 0.2]] + ) def test_native_selected_planner_uses_explicit_start_after_live_state_advances( @@ -1754,16 +2083,34 @@ def capture_scene_start( return native_plan(planner, q_start, q_goal) mocker.patch.object(FakeRRT, "plan", autospec=True, side_effect=capture_scene_start) - start = JointState(name=list(selection.joint_names), position=[0.1, -0.1]) + start = JointState( + name=list(selection.joint_names), + position=np.array([0.1, -0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) world.sync_from_joint_state( - JointState(name=["joint1", "joint2"], position=[0.3, 0.2]), + JointState( + name=["joint1", "joint2"], + position=np.array([0.3, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) result = _planner_for(world).plan_selected_joint_path( world, selection, start, - JointState(name=list(selection.joint_names), position=[0.4, 0.2]), + JointState( + name=list(selection.joint_names), + position=np.array([0.4, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) assert result.status == PlanningStatus.SUCCESS @@ -1788,8 +2135,20 @@ def test_native_selected_planner_supports_non_overlapping_multi_group_selection( result = _planner_for(world).plan_selected_joint_path( world, selection, - JointState(name=list(selection.joint_names), position=[0.0, 0.0]), - JointState(name=list(selection.joint_names), position=[0.1, 0.1]), + JointState( + name=list(selection.joint_names), + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointState( + name=list(selection.joint_names), + position=np.array([0.1, 0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) assert result.status == PlanningStatus.SUCCESS @@ -1825,11 +2184,17 @@ def test_cartesian_planner_returns_timed_canonical_joint_states_and_options( result = _planner_for(world).plan_cartesian_path( world, selection, - JointState(name=list(selection.joint_names), position=[0.0, 0.0]), + JointState( + name=list(selection.joint_names), + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), { "manipulator": _relative_target( TransformStamped( - header=Header(frame_id="world"), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), transform=Transform( translation=Vector3(x=0.1, y=0.0, z=0.0), rotation=quaternion_from_euler(0.0, 0.0, np.pi / 2.0), @@ -1878,17 +2243,29 @@ def test_cartesian_zero_rotation_preserves_start_orientation( result = _planner_for(world).plan_cartesian_path( world, selection, - JointState(name=list(selection.joint_names), position=[0.0, 0.0]), + JointState( + name=list(selection.joint_names), + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), { "manipulator": _relative_target( TransformStamped( - header=Header(frame_id="world"), - transform=Transform(translation=Vector3(x=0.05, y=0.02, z=0.0)), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + transform=Transform( + translation=Vector3(x=0.05, y=0.02, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), child_frame_id="", ), TransformStamped( - header=Header(frame_id="world"), - transform=Transform(translation=Vector3(x=0.1, y=0.0, z=0.0)), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + transform=Transform( + translation=Vector3(x=0.1, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), child_frame_id="", ), ) @@ -1908,9 +2285,21 @@ def test_cartesian_uses_explicit_start_after_live_state_advances( ) -> None: world = _make_world(fake_roboplan, robot_config) selection = _selection(robot_config, "manipulator") - start = JointState(name=list(selection.joint_names), position=[0.1, 0.0]) + start = JointState( + name=list(selection.joint_names), + position=np.array([0.1, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) world.sync_from_joint_state( - JointState(name=["joint1", "joint2"], position=[0.3, 0.2]), + JointState( + name=["joint1", "joint2"], + position=np.array([0.3, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) result = _planner_for(world).plan_cartesian_path( @@ -1920,8 +2309,11 @@ def test_cartesian_uses_explicit_start_after_live_state_advances( { "manipulator": _relative_target( TransformStamped( - header=Header(frame_id="world"), - transform=Transform(translation=Vector3(x=0.1, y=0.0, z=0.0)), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + transform=Transform( + translation=Vector3(x=0.1, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), child_frame_id="", ) ) @@ -1952,12 +2344,21 @@ def test_cartesian_rejects_official_planner_failure( result = _planner_for(world).plan_cartesian_path( world, selection, - JointState(name=list(selection.joint_names), position=[0.0, 0.0]), + JointState( + name=list(selection.joint_names), + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), { "manipulator": _relative_target( TransformStamped( - header=Header(frame_id="world"), - transform=Transform(translation=Vector3(x=0.1, y=0.0, z=0.0)), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + transform=Transform( + translation=Vector3(x=0.1, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), child_frame_id="", ) ) @@ -2012,12 +2413,21 @@ def collides_only_mid_edge(scene: FakeScene, q: np.ndarray) -> bool: result = _planner_for(world).plan_cartesian_path( world, selection, - JointState(name=list(selection.joint_names), position=[0.0, 0.0]), + JointState( + name=list(selection.joint_names), + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), { "manipulator": _relative_target( TransformStamped( - header=Header(frame_id="world"), - transform=Transform(translation=Vector3(x=0.05, y=0.0, z=0.0)), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + transform=Transform( + translation=Vector3(x=0.05, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), child_frame_id="", ) ) @@ -2042,8 +2452,20 @@ def plan( monkeypatch.setattr(sys.modules["roboplan.rrt"], "RRT", EmptyPathRRT) world = _make_world(fake_roboplan, robot_config) - start = JointState(name=["joint1", "joint2"], position=[0.0, 0.0]) - goal = JointState(name=["joint1", "joint2"], position=[0.4, 0.2]) + start = JointState( + name=["joint1", "joint2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + goal = JointState( + name=["joint1", "joint2"], + position=np.array([0.4, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) result = _planner_for(world).plan_joint_path(world, start, goal, timeout=1.0) assert result.status == PlanningStatus.NO_SOLUTION @@ -2146,8 +2568,10 @@ def test_base_pose_is_written_to_composed_model( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: robot_config.base_pose = PoseStamped( - header=Header(frame_id=""), - pose=Pose(position=Point(x=1, y=0, z=0), orientation=Quaternion()), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=1, y=0, z=0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), ) world = _make_world(fake_roboplan, robot_config) diff --git a/dimos/manipulation/test_roboplan_integration.py b/dimos/manipulation/test_roboplan_integration.py index 175f5f043a8..8422a4eab79 100644 --- a/dimos/manipulation/test_roboplan_integration.py +++ b/dimos/manipulation/test_roboplan_integration.py @@ -18,6 +18,18 @@ from pathlib import Path from typing import Any +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest @@ -40,16 +52,12 @@ RoboPlanTOPPRAParametrizationConfig, ) from dimos.manipulation.planning.utils.kinematics_utils import compute_pose_error -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.geometry import pose_matrix from dimos.robot.assets.model import RobotModel from dimos.robot.manipulators.xarm.config import ( make_dual_xarm6_model_config, make_xarm6_model_config, ) -from dimos.utils.transform_utils import pose_to_matrix pytestmark = pytest.mark.self_hosted @@ -67,7 +75,13 @@ def _sync_zero_state( joint_names: list[str], ) -> None: world.sync_from_joint_state( - JointState(name=joint_names, position=[0.0] * len(joint_names)), + JointState( + name=joint_names, + position=np.asarray([0.0] * len(joint_names), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) @@ -88,7 +102,11 @@ def test_real_roboplan_plans_fixed_orientation_cartesian_path( group_id = "manipulator" selection = world._planning_groups.select((group_id,)) start = JointState( - name=list(selection.joint_names), position=[0.0] * len(selection.joint_names) + name=list(selection.joint_names), + position=np.asarray([0.0] * len(selection.joint_names), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) with world.scratch_context() as ctx: planner._apply_selected_state(ctx, start) @@ -100,8 +118,22 @@ def test_real_roboplan_plans_fixed_orientation_cartesian_path( start, { group_id: ( - Transform.identity(), - Transform(translation=Vector3(0.005, 0.0, 0.0)), + TransformStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + child_frame_id="", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + TransformStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + transform=Transform( + translation=Vector3(x=0.005, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + child_frame_id="", + ), ) }, RoboPlanCartesianPathConfig( @@ -117,9 +149,9 @@ def test_real_roboplan_plans_fixed_orientation_cartesian_path( with world.scratch_context() as ctx: planner._apply_selected_state(ctx, result.path[-1]) final_pose = world.get_group_ee_pose(ctx, group_id) - expected = pose_to_matrix(start_pose) + expected = pose_matrix(start_pose.pose) expected[0, 3] += 0.005 - position_error, orientation_error = compute_pose_error(pose_to_matrix(final_pose), expected) + position_error, orientation_error = compute_pose_error(pose_matrix(final_pose.pose), expected) assert position_error <= 0.005 assert orientation_error <= 0.01 @@ -142,7 +174,11 @@ def test_real_roboplan_synchronizes_different_length_dual_arm_targets( right_group_id = "right_arm" selection = world._planning_groups.select((left_group_id, right_group_id)) start = JointState( - name=list(selection.joint_names), position=[0.0] * len(selection.joint_names) + name=list(selection.joint_names), + position=np.asarray([0.0] * len(selection.joint_names), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) result = planner.plan_cartesian_path( @@ -151,13 +187,48 @@ def test_real_roboplan_synchronizes_different_length_dual_arm_targets( start, { left_group_id: ( - Transform.identity(), - Transform(translation=Vector3(0.0015, 0.001, 0.0)), - Transform(translation=Vector3(0.003, 0.0, 0.0)), + TransformStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + child_frame_id="", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + TransformStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + transform=Transform( + translation=Vector3(x=0.0015, y=0.001, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + child_frame_id="", + ), + TransformStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + transform=Transform( + translation=Vector3(x=0.003, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + child_frame_id="", + ), ), right_group_id: ( - Transform.identity(), - Transform(translation=Vector3(0.005, 0.0, 0.0)), + TransformStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + child_frame_id="", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + TransformStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + transform=Transform( + translation=Vector3(x=0.005, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + child_frame_id="", + ), ), }, RoboPlanCartesianPathConfig(), @@ -186,7 +257,13 @@ def scalar_world(tmp_path, roboplan_types): model=RobotModel.from_file(model_path), joint_names=["slide"], base_link="base", - base_pose=PoseStamped(frame_id="world", position=[1, 0, 0]), + base_pose=PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=1, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), planning_groups=[ PlanningGroupDefinition( name="arm", joint_names=("slide",), base_link="base", tip_link="arm" @@ -196,15 +273,41 @@ def scalar_world(tmp_path, roboplan_types): world = roboplan_types[0]() world.load_model(prepare_robot_model(config)) world.finalize() - world.sync_from_joint_state(JointState(name=["slide"], position=[0.0])) + world.sync_from_joint_state( + JointState( + name=["slide"], + position=np.asarray([0.0], dtype=np.float64), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) return world def test_real_contexts_keep_independent_state_and_refresh_geometry(scalar_world): world = scalar_world with world.scratch_context() as first, world.scratch_context() as second: - world.set_joint_state(first, JointState(name=["slide"], position=[0.0])) - world.set_joint_state(second, JointState(name=["slide"], position=[0.4])) + world.set_joint_state( + first, + JointState( + name=["slide"], + position=np.asarray([0.0], dtype=np.float64), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + ) + world.set_joint_state( + second, + JointState( + name=["slide"], + position=np.asarray([0.4], dtype=np.float64), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + ) np.testing.assert_allclose(world.get_link_pose(first, "arm")[:3, 3], [1, 0, 0]) np.testing.assert_allclose(world.get_link_pose(second, "arm")[:3, 3], [1.4, 0, 0]) with world.parametrization_model() as model: @@ -214,13 +317,26 @@ def test_real_contexts_keep_independent_state_and_refresh_geometry(scalar_world) name="block", obstacle_type=ObstacleType.BOX, dimensions=(0.1, 0.1, 0.1), - pose=PoseStamped(frame_id="world", position=[1.4, 0, 0]), + pose=PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=1.4, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), ) ) assert not world.is_collision_free(second) assert world.is_collision_free(first) assert world.update_obstacle_pose( - "block", PoseStamped(frame_id="world", position=[1.8, 0, 0]) + "block", + PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=1.8, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), ) assert world.is_collision_free(second) assert world.remove_obstacle("block") @@ -238,8 +354,20 @@ def test_real_rrt_and_toppra_use_finite_prepared_limits(scalar_world, roboplan_t assert world.get_prepared_model().joint_space.acceleration_limits == (0.4,) planner = roboplan_types[1](world, RoboPlanPlannerConfig()) - start = JointState(name=["slide"], position=[0.0]) - goal = JointState(name=["slide"], position=[0.5]) + start = JointState( + name=["slide"], + position=np.asarray([0.0], dtype=np.float64), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + goal = JointState( + name=["slide"], + position=np.asarray([0.5], dtype=np.float64), + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) result = planner.plan_joint_path(world, start, goal, timeout=1.0) assert result.status == PlanningStatus.SUCCESS, result.message selection = world._planning_groups.select(("arm",)) diff --git a/dimos/manipulation/test_sdk.py b/dimos/manipulation/test_sdk.py index cc5b1bc51d5..253021f0a95 100644 --- a/dimos/manipulation/test_sdk.py +++ b/dimos/manipulation/test_sdk.py @@ -15,6 +15,7 @@ import atexit from dataclasses import replace +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header @@ -60,9 +61,15 @@ def rpc(mocker): execution_status=ExecutionStatus.IDLE, groups={ "arm": PlanningGroupState( - joints=JointState(name=["j0", "j1"], position=[0.1, 0.2]), + joints=JointState( + name=["j0", "j1"], + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), end_effector_pose=PoseStamped( - header=Header(frame_id="world"), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=Point(x=0.4, y=0.0, z=0.3), orientation=Quaternion(x=0.0, y=1.0, z=0.0, w=0.0), @@ -70,8 +77,20 @@ def rpc(mocker): ), gripper_position=0.5, joint_presets={ - "home": JointState(name=["j0", "j1"], position=[0.0, 0.5]), - "init": JointState(name=["j0", "j1"], position=[0.1, 0.2]), + "home": JointState( + name=["j0", "j1"], + position=np.array([0.0, 0.5], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + "init": JointState( + name=["j0", "j1"], + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), }, ), }, @@ -80,18 +99,14 @@ def rpc(mocker): PlanStatus.SUCCEEDED, plan=GeneratedPlan( ("arm",), - JointTrajectory( - header=header_now(), - ), + JointTrajectory(header=header_now(), joint_names=[], points=[]), ), ) proxy.plan_to_poses.return_value = PlanResult( PlanStatus.SUCCEEDED, plan=GeneratedPlan( ("arm",), - JointTrajectory( - header=header_now(), - ), + JointTrajectory(header=header_now(), joint_names=[], points=[]), ), ) proxy.execute.return_value = ExecutionResult(ExecutionStatus.COMPLETED) @@ -223,7 +238,7 @@ def test_joint_inputs_plan_before_blocking_execution(arm, rpc): target = rpc.plan_to_joints.call_args.args[0]["arm"] assert target.name == ["j0", "j1"] - assert target.position == [0.3, 0.4] + np.testing.assert_array_equal(target.position, [0.3, 0.4]) assert rpc.plan_to_joints.call_args.kwargs == {"speed_scale": 0.2} assert [call[0] for call in rpc.method_calls][-2:] == ["plan_to_joints", "execute"] rpc.execute.assert_called_once_with( @@ -450,7 +465,13 @@ def test_home_uses_server_preset_and_preserves_names(arm, rpc): def test_init_preset_is_fetched_fresh_each_time(arm, rpc): first = rpc.get_state.return_value - new_target = JointState(name=["j1", "j0"], position=[0.7, 0.8]) + new_target = JointState( + name=["j1", "j0"], + position=np.array([0.7, 0.8], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) second = replace( first, groups={"arm": replace(first.groups["arm"], joint_presets={"init": new_target})} ) diff --git a/dimos/manipulation/visualization/operator.py b/dimos/manipulation/visualization/operator.py index 19d832bea25..fdc2bc51956 100644 --- a/dimos/manipulation/visualization/operator.py +++ b/dimos/manipulation/visualization/operator.py @@ -22,8 +22,11 @@ import math from typing import TYPE_CHECKING +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import PoseStamped from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header +import numpy as np from dimos.manipulation.planning.groups.models import PlanningGroup from dimos.manipulation.planning.planners.config import CartesianPathConfig @@ -159,7 +162,12 @@ def plan_to_joints(self, request: JointTargetRequest) -> GeneratedPlan | None: group.id: JointState( name=list(group.joint_names), header=request.target.header, - position=list(request.target.position)[offset : offset + len(group.joint_names)], + position=np.asarray( + list(request.target.position)[offset : offset + len(group.joint_names)], + dtype=np.float64, + ), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) for group, offset in self._group_offsets(groups) } @@ -314,7 +322,13 @@ def _complete_states( if value is None: return None positions.append(value) - return JointState(name=list(config.joint_names), position=positions) + return JointState( + name=list(config.joint_names), + position=np.asarray(positions, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) def _evaluate_complete_target( self, diff --git a/dimos/manipulation/visualization/test_factory.py b/dimos/manipulation/visualization/test_factory.py index d1b770e120e..74db96e88a1 100644 --- a/dimos/manipulation/visualization/test_factory.py +++ b/dimos/manipulation/visualization/test_factory.py @@ -18,8 +18,9 @@ from pathlib import Path from unittest.mock import MagicMock +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.dimos_msgs.msg import GraspCandidateArray -from dimos_generated.geometry_msgs.msg import Pose, PoseStamped +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header from dimos_generated.trajectory_msgs.msg import JointTrajectory @@ -99,7 +100,13 @@ def load_model(self, model: PreparedRobotModel) -> None: def get_model_config(self) -> RobotModelConfig: return RobotModelConfig( model=RobotModel.from_file(Path("fake.urdf")), - base_pose=PoseStamped(header=Header(frame_id=""), pose=Pose()), + base_pose=PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), joint_names=["joint1"], planning_groups=[ PlanningGroupDefinition( @@ -162,7 +169,13 @@ def set_joint_state(self, ctx: object, joint_state: JointState) -> None: return None def get_joint_state(self, ctx: object) -> JointState: - return JointState() + return JointState( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + position=np.array([], dtype=np.float64), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) def is_collision_free(self, ctx: object) -> bool: return True @@ -182,7 +195,13 @@ def check_edge_collision_free( return True def get_ee_pose(self, ctx: object) -> PoseStamped: - return PoseStamped(header=Header(frame_id=""), pose=Pose()) + return PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) def get_link_pose(self, ctx: object, link_name: str) -> NDArray[np.float64]: return np.eye(4, dtype=np.float64) @@ -191,7 +210,13 @@ def get_jacobian(self, ctx: object) -> NDArray[np.float64]: return np.zeros((6, 0), dtype=np.float64) def get_group_ee_pose(self, ctx: object, group_id: str) -> PoseStamped: - return PoseStamped(header=Header(frame_id=""), pose=Pose()) + return PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) def get_group_jacobian(self, ctx: object, group_id: str) -> NDArray[np.float64]: return np.zeros((6, 0), dtype=np.float64) @@ -329,7 +354,13 @@ def test_create_visualization_meshcat_accepts_structural_world() -> None: obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, - pose=PoseStamped(header=Header(frame_id=""), pose=Pose()), + pose=PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), dimensions=(1.0, 1.0, 1.0), ) visualization.initialize(session) @@ -341,7 +372,9 @@ def test_create_visualization_meshcat_accepts_structural_world() -> None: visualization.add_vis_obstacle("box", obstacle) visualization.remove_vis_obstacle("box") visualization.clear_vis_obstacles() - proposals = GraspCandidateArray() + proposals = GraspCandidateArray( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), candidates=[] + ) visualization.show_grasp_proposals(proposals) assert fake_world.visualization_calls == [ ("initialize", session), @@ -396,13 +429,21 @@ def test_drake_meshcat_visualization_lifecycle_is_noop_without_meshcat() -> None obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, - pose=PoseStamped(header=Header(frame_id=""), pose=Pose()), + pose=PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), dimensions=(1.0, 1.0, 1.0), ) world.add_vis_obstacle("box", obstacle) world.remove_vis_obstacle("box") world.clear_vis_obstacles() - world.show_grasp_proposals(GraspCandidateArray()) + world.show_grasp_proposals( + GraspCandidateArray(header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), candidates=[]) + ) world.cancel_preview_animation() world.close() diff --git a/dimos/manipulation/visualization/test_operator.py b/dimos/manipulation/visualization/test_operator.py index 8e9cfbc483c..817223fe5fb 100644 --- a/dimos/manipulation/visualization/test_operator.py +++ b/dimos/manipulation/visualization/test_operator.py @@ -14,14 +14,17 @@ """Focused tests for the single-model visualization operator.""" +from dataclasses import asdict from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock -from dimos_generated.geometry_msgs.msg import Pose, PoseStamped +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header from dimos_generated.trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint +import numpy as np from dimos.manipulation.manipulation_spec import ( CommandResult, @@ -63,13 +66,27 @@ def _plan() -> GeneratedPlan: header=header_now(), joint_names=["left/j1", "left/j2"], points=[ - JointTrajectoryPoint(positions=[0.4, 0.5], time_from_start=duration_from_seconds(1.0)) + JointTrajectoryPoint( + positions=np.array([0.4, 0.5], dtype=np.float64), + time_from_start=duration_from_seconds(1.0), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) ], ) return GeneratedPlan( group_ids=("left_arm",), trajectory=trajectory, - path=[JointState(name=trajectory.joint_names, position=[0.4, 0.5])], + path=[ + JointState( + name=trajectory.joint_names, + position=np.array([0.4, 0.5], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ], status=PlanningStatus.SUCCESS, ) @@ -85,7 +102,13 @@ def _operator() -> tuple[ManipulationOperator, MagicMock, MagicMock]: module.get_init_joints.return_value = None module.inverse_kinematics.return_value = IKResult( status=IKStatus.SUCCESS, - joint_state=JointState(name=["left/j1", "left/j2"], position=[0.4, 0.5]), + joint_state=JointState( + name=["left/j1", "left/j2"], + position=np.array([0.4, 0.5], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) module.generate_plan_to_joint_targets.return_value = _plan() module.generate_plan_to_pose_targets.return_value = _plan() @@ -98,11 +121,18 @@ def _operator() -> tuple[ManipulationOperator, MagicMock, MagicMock]: monitor = MagicMock() monitor.planning_groups = PlanningGroupRegistry(config.planning_groups) monitor.get_current_joint_state.return_value = JointState( - name=config.joint_names, position=[0.0, 0.0, 0.0] + name=config.joint_names, + position=np.array([0.0, 0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) monitor.is_state_valid.return_value = True monitor.get_group_ee_pose.return_value = PoseStamped( - header=Header(frame_id="world"), pose=Pose() + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), ) return ManipulationOperator(module, monitor), module, monitor @@ -118,24 +148,52 @@ def test_joint_evaluation_overlays_selected_target_on_complete_model_state() -> operator, _, monitor = _operator() request = JointTargetRequest( ("left_arm",), - JointState(name=["left/j1", "left/j2"], position=[0.1, 0.2]), + JointState( + name=["left/j1", "left/j2"], + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) result = operator.evaluate_joint_target(request) assert result.success complete = monitor.is_state_valid.call_args.args[0] assert complete.name == ["left/j1", "left/j2", "right/j1"] - assert complete.position == [0.1, 0.2, 0.0] + np.testing.assert_array_equal(complete.position, [0.1, 0.2, 0.0]) def test_joint_evaluation_rejects_noncanonical_unknown_and_overlapping_selection() -> None: operator, _, _ = _operator() noncanonical = JointTargetRequest( - ("left_arm",), JointState(name=["j1", "j2"], position=[0.1, 0.2]) + ("left_arm",), + JointState( + name=["j1", "j2"], + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + ) + unknown = JointTargetRequest( + ("missing",), + JointState( + name=["left/j1"], + position=np.array([0.1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) - unknown = JointTargetRequest(("missing",), JointState(name=["left/j1"], position=[0.1])) duplicate = JointTargetRequest( ("left_arm", "left_arm"), - JointState(name=["left/j1", "left/j2"], position=[0.1, 0.2]), + JointState( + name=["left/j1", "left/j2"], + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) assert all( not operator.evaluate_joint_target(request).success @@ -145,23 +203,65 @@ def test_joint_evaluation_rejects_noncanonical_unknown_and_overlapping_selection def test_pose_evaluation_routes_group_id_and_canonical_seed() -> None: operator, module, _ = _operator() - pose = PoseStamped(header=Header(frame_id="world"), pose=Pose()) - seed = JointState(name=["left/j1", "left/j2"], position=[0.0, 0.0]) + pose = PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + ) + seed = JointState( + name=["left/j1", "left/j2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) result = operator.evaluate_pose_target(PoseTargetRequest({"left_arm": pose}, seed=seed)) assert result.success - module.inverse_kinematics.assert_called_once_with( - pose_targets={"left_arm": pose}, auxiliary_group_ids=(), seed=seed, check_collision=True + module.inverse_kinematics.assert_called_once() + assert module.inverse_kinematics.call_args.args == () + kwargs = module.inverse_kinematics.call_args.kwargs + np.testing.assert_equal( + {**kwargs, "seed": asdict(kwargs["seed"])}, + dict( + pose_targets={"left_arm": pose}, + auxiliary_group_ids=(), + seed=asdict(seed), + check_collision=True, + ), ) def test_pose_evaluation_rejects_non_world_frame_and_ambiguous_local_seed() -> None: operator, _, _ = _operator() bad_frame = PoseTargetRequest( - {"left_arm": PoseStamped(header=Header(frame_id="camera"), pose=Pose())} + { + "left_arm": PoseStamped( + header=Header(frame_id="camera", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + } ) bad_seed = PoseTargetRequest( - {"left_arm": PoseStamped(header=Header(frame_id="world"), pose=Pose())}, - seed=JointState(name=["j1", "j2"], position=[0.0, 0.0]), + { + "left_arm": PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + }, + seed=JointState( + name=["j1", "j2"], + position=np.array([0.0, 0.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) assert not operator.evaluate_pose_target(bad_frame).success assert not operator.evaluate_pose_target(bad_seed).success @@ -170,10 +270,25 @@ def test_pose_evaluation_rejects_non_world_frame_and_ambiguous_local_seed() -> N def test_planning_and_actions_return_exact_generated_plan() -> None: operator, module, _ = _operator() joint_request = JointTargetRequest( - ("left_arm",), JointState(name=["left/j1", "left/j2"], position=[0.1, 0.2]) + ("left_arm",), + JointState( + name=["left/j1", "left/j2"], + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) pose_request = PoseTargetRequest( - {"left_arm": PoseStamped(header=Header(frame_id="world"), pose=Pose())} + { + "left_arm": PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + } ) assert ( operator.plan_to_joints(joint_request) is module.generate_plan_to_joint_targets.return_value @@ -187,7 +302,12 @@ def test_planning_and_actions_return_exact_generated_plan() -> None: def test_cartesian_planning_uses_current_group_pose() -> None: operator, module, monitor = _operator() - target = PoseStamped(header=Header(frame_id="world"), pose=Pose()) + target = PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + ) config = RoboPlanCartesianPathConfig() result = operator.plan_cartesian(CartesianTargetRequest({"left_arm": target}, config)) assert result is module.generate_cartesian_plan.return_value diff --git a/dimos/manipulation/visualization/viser/gui.py b/dimos/manipulation/visualization/viser/gui.py index b977c3f83fc..e451356dc18 100644 --- a/dimos/manipulation/visualization/viser/gui.py +++ b/dimos/manipulation/visualization/viser/gui.py @@ -20,9 +20,11 @@ from threading import RLock from typing import Any, TypeAlias, cast +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header +import numpy as np from dimos.manipulation.planning.groups.models import PlanningGroup from dimos.manipulation.planning.planners.roboplan_config import RoboPlanCartesianPathConfig @@ -206,7 +208,11 @@ def _current_target_for_group(self, group: PlanningGroup) -> dict[PlanningGroupI return {} return { group.id: JointState( - name=list(group.joint_names), position=[values[name] for name in group.joint_names] + name=list(group.joint_names), + position=np.asarray([values[name] for name in group.joint_names], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) } @@ -233,7 +239,16 @@ def evaluate_joint_target_set( names.extend(str(name) for name in target.name) positions.extend(float(value) for value in target.position) return self.operator.evaluate_joint_target( - JointTargetRequest(tuple(group_ids), JointState(name=names, position=positions)) + JointTargetRequest( + tuple(group_ids), + JointState( + name=names, + position=np.asarray(positions, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + ) ) def evaluate_pose_target_set( @@ -265,7 +280,16 @@ def plan_to_selected_joints( names.extend(str(name) for name in target.name) positions.extend(float(value) for value in target.position) plan = self.operator.plan_to_joints( - JointTargetRequest(tuple(group_ids), JointState(name=names, position=positions)) + JointTargetRequest( + tuple(group_ids), + JointState( + name=names, + position=np.asarray(positions, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + ) ) self.state.plan_state.plan = plan return plan is not None @@ -689,7 +713,12 @@ def _initialize_selected_group_targets(self) -> None: continue self.state.group_joint_targets[group_id] = JointState( name=list(group.joint_names), - position=[float(values[str(name)]) for name in group.joint_names], + position=np.asarray( + [float(values[str(name)]) for name in group.joint_names], dtype=np.float64 + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) if group.has_pose_target and group_id not in self.state.pose_targets: pose = self.get_group_ee_pose(group_id) @@ -720,7 +749,17 @@ def _refresh_target_joints_from_groups(self) -> None: if target is not None: names.extend(str(name) for name in target.name) positions.extend(float(value) for value in target.position) - self.state.target_joints = JointState(name=names, position=positions) if names else None + self.state.target_joints = ( + JointState( + name=names, + position=np.asarray(positions, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + if names + else None + ) def _active_pose_targets(self) -> dict[PlanningGroupID, PoseStamped]: return { @@ -789,7 +828,13 @@ def _apply_preset(self, preset: str) -> None: ) return positions = [float(values[str(name)]) for name in group.joint_names] - targets[group_id] = JointState(name=list(group.joint_names), position=positions) + targets[group_id] = JointState( + name=list(group.joint_names), + position=np.asarray(positions, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) slider_values.append((group_id, group.joint_names, positions)) self.state.group_joint_targets.update(targets) with self._joint_controls_lock: @@ -837,7 +882,13 @@ def _target_set_from_sliders(self) -> dict[PlanningGroupID, JointState] | None: self._set_error(f"Missing target slider for {group_id}/{joint_name}") return None positions.append(float(handle.value)) - targets[group_id] = JointState(name=list(group.joint_names), position=positions) + targets[group_id] = JointState( + name=list(group.joint_names), + position=np.asarray(positions, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) return targets def _on_joint_slider_update(self, _group_id: PlanningGroupID, _joint_name: str) -> None: @@ -925,7 +976,11 @@ def _target_ghost_state( if not all(name in values for name in config.joint_names): return None return JointState( - name=list(config.joint_names), position=[values[name] for name in config.joint_names] + name=list(config.joint_names), + position=np.asarray([values[name] for name in config.joint_names], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) def _sync_target_ghost_visibility(self) -> None: @@ -1004,7 +1059,12 @@ def _split_target_joints_by_group(self, target_joints: JointState) -> None: continue self.state.group_joint_targets[group_id] = JointState( name=list(group.joint_names), - position=[positions[str(name)] for name in group.joint_names], + position=np.asarray( + [positions[str(name)] for name in group.joint_names], dtype=np.float64 + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) def _sync_pose_targets_from_group_poses(self) -> None: @@ -1393,7 +1453,7 @@ def _pose_from_transform_target(self, target: TransformControlsHandle) -> PoseSt px, py, pz = (float(value) for value in target.position) qw, qx, qy, qz = (float(value) for value in target.wxyz) return PoseStamped( - header=Header(frame_id="world"), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), pose=Pose( position=Point(x=px, y=py, z=pz), orientation=Quaternion(x=qx, y=qy, z=qz, w=qw) ), diff --git a/dimos/manipulation/visualization/viser/test_gui.py b/dimos/manipulation/visualization/viser/test_gui.py index b89c7e42583..e93bb77dfe9 100644 --- a/dimos/manipulation/visualization/viser/test_gui.py +++ b/dimos/manipulation/visualization/viser/test_gui.py @@ -21,6 +21,9 @@ from threading import Barrier, Lock import time +from dimos_generated.builtin_interfaces.msg import Duration, Time +from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from dimos.msgs.time import duration_from_seconds, header_now @@ -271,14 +274,31 @@ def executable_gui(monkeypatch, mocker): joint_names=["arm/j0"], points=[ JointTrajectoryPoint( - positions=[0.0], time_from_start=duration_from_seconds(0.0) + positions=np.array([0.0], dtype=np.float64), + time_from_start=duration_from_seconds(0.0), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), JointTrajectoryPoint( - positions=[1.0], time_from_start=duration_from_seconds(1.0) + positions=np.array([1.0], dtype=np.float64), + time_from_start=duration_from_seconds(1.0), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ], ), - path=[JointState(name=["arm/j0"], position=[value]) for value in (0.0, 1.0)], + path=[ + JointState( + name=["arm/j0"], + position=np.array([value], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + for value in (0.0, 1.0) + ], status=PlanningStatus.SUCCESS, ), ) @@ -361,7 +381,14 @@ def test_gui_completion_enables_next_plan_without_cancel(executable_gui, module_ def invoke(task, method, args=None): if method == "execute": return TrajectoryExecutionResult(TrajectoryExecutionStatus.ACCEPTED) - return TrajectoryStatus(header=header_now(), state=reported["state"]) + return TrajectoryStatus( + header=header_now(), + state=reported["state"], + progress=0.0, + time_elapsed=Duration(sec=0, nanosec=0), + time_remaining=Duration(sec=0, nanosec=0), + error="", + ) mocker.patch.object(module._control_coordinator, "task_invoke", side_effect=invoke) cancel = mocker.spy(module, "cancel") @@ -442,7 +469,11 @@ def test_planar_joint_controls_use_unbounded_translation_inputs_and_wrapped_yaw( ) panel.state.selected_group_ids = (group.id,) panel.state.group_joint_targets[group.id] = JointState( - name=list(planar.joint_names), position=[6.0, -7.0, 4.0] + name=list(planar.joint_names), + position=np.array([6.0, -7.0, 4.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) controls = FakeJointGui() @@ -614,10 +645,16 @@ def test_gui_preview_enters_previewing_before_worker_runs( gui.state.plan_state.target_sequence_id = gui.state.latest_sequence_id gui.state.plan_state.plan = GeneratedPlan( group_ids=gui.state.selected_group_ids, - trajectory=JointTrajectory( - header=header_now(), - ), - path=[JointState(name=[], position=[])], + trajectory=JointTrajectory(header=header_now(), joint_names=[], points=[]), + path=[ + JointState( + name=[], + position=np.array([], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ], ) assert gui.state.can_execute() is True @@ -653,10 +690,16 @@ def test_gui_selection_change_clears_invalidated_preview( gui.state.plan_state.target_sequence_id = gui.state.latest_sequence_id gui.state.plan_state.plan = GeneratedPlan( group_ids=gui.state.selected_group_ids, - trajectory=JointTrajectory( - header=header_now(), - ), - path=[JointState(name=[], position=[])], + trajectory=JointTrajectory(header=header_now(), joint_names=[], points=[]), + path=[ + JointState( + name=[], + position=np.array([], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ], ) gui._submit_preview() @@ -691,10 +734,16 @@ def test_gui_selection_change_ignores_invalidated_preview_error( gui.state.plan_state.target_sequence_id = gui.state.latest_sequence_id gui.state.plan_state.plan = GeneratedPlan( group_ids=gui.state.selected_group_ids, - trajectory=JointTrajectory( - header=header_now(), - ), - path=[JointState(name=[], position=[])], + trajectory=JointTrajectory(header=header_now(), joint_names=[], points=[]), + path=[ + JointState( + name=[], + position=np.array([], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ], ) gui._submit_preview() diff --git a/dimos/manipulation/visualization/viser/test_scene_grasp_proposals.py b/dimos/manipulation/visualization/viser/test_scene_grasp_proposals.py index f81f7de9131..e5243ad6626 100644 --- a/dimos/manipulation/visualization/viser/test_scene_grasp_proposals.py +++ b/dimos/manipulation/visualization/viser/test_scene_grasp_proposals.py @@ -17,6 +17,7 @@ import inspect from unittest.mock import MagicMock +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point import numpy as np import pytest @@ -118,7 +119,9 @@ def test_an_empty_array_clears_the_markers() -> None: scene = _scene() scene.show_grasp_proposals(_proposals(4)) - scene.show_grasp_proposals(GraspCandidateArray()) + scene.show_grasp_proposals( + GraspCandidateArray(header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), candidates=[]) + ) assert scene._grasp_proposal_handles == [] diff --git a/dimos/manipulation/visualization/viser/test_scene_obstacles.py b/dimos/manipulation/visualization/viser/test_scene_obstacles.py index 02bda7e1916..ffca5a5282b 100644 --- a/dimos/manipulation/visualization/viser/test_scene_obstacles.py +++ b/dimos/manipulation/visualization/viser/test_scene_obstacles.py @@ -16,7 +16,8 @@ from unittest.mock import MagicMock -from dimos_generated.geometry_msgs.msg import Pose +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion from dimos_generated.std_msgs.msg import Header import numpy as np import pytest @@ -38,7 +39,13 @@ def _octree(points: np.ndarray, resolution: float = 0.025) -> Obstacle: return Obstacle( name="mapping/voxel-map", obstacle_type=ObstacleType.OCTREE, - pose=PoseStamped(header=Header(frame_id="world"), pose=Pose()), + pose=PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), points=points, octree_resolution=resolution, ) diff --git a/dimos/manipulation/visualization/viser/test_viser_visualization.py b/dimos/manipulation/visualization/viser/test_viser_visualization.py index e7f57cdd812..5e28843a2e8 100644 --- a/dimos/manipulation/visualization/viser/test_viser_visualization.py +++ b/dimos/manipulation/visualization/viser/test_viser_visualization.py @@ -17,6 +17,9 @@ from pathlib import Path from unittest.mock import MagicMock +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from dimos.msgs.time import duration_from_seconds, header_now @@ -69,11 +72,25 @@ def test_preview_timing_uses_one_model_track() -> None: def test_visualizer_builds_full_model_preview_from_selected_canonical_joints() -> None: visualizer = ViserManipulationVisualizer() visualizer._model_config = _model() - visualizer._current_state = JointState(name=["left/j1", "right/j1"], position=[0.1, 0.2]) + visualizer._current_state = JointState( + name=["left/j1", "right/j1"], + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) trajectory = JointTrajectory( header=header_now(), joint_names=["right/j1"], - points=[JointTrajectoryPoint(positions=[0.8], time_from_start=duration_from_seconds(1.0))], + points=[ + JointTrajectoryPoint( + positions=np.array([0.8], dtype=np.float64), + time_from_start=duration_from_seconds(1.0), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ], ) preview = visualizer._raw_preview_animation(trajectory) assert preview == PreviewAnimation(("left/j1", "right/j1"), (PreviewFrame(1.0, (0.1, 0.8)),)) @@ -82,14 +99,24 @@ def test_visualizer_builds_full_model_preview_from_selected_canonical_joints() - def test_visualizer_rejects_unknown_or_duplicate_trajectory_joints() -> None: visualizer = ViserManipulationVisualizer() visualizer._model_config = _model() - visualizer._current_state = JointState(name=["left/j1", "right/j1"], position=[0.1, 0.2]) + visualizer._current_state = JointState( + name=["left/j1", "right/j1"], + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) for names in (["unknown"], ["left/j1", "left/j1"]): trajectory = JointTrajectory( header=header_now(), joint_names=names, points=[ JointTrajectoryPoint( - positions=[0.0] * len(names), time_from_start=duration_from_seconds(1.0) + positions=np.asarray([0.0] * len(names), dtype=np.float64), + time_from_start=duration_from_seconds(1.0), + velocities=np.array([], dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) ], ) @@ -104,6 +131,12 @@ def test_visualizer_initializes_and_updates_one_scene_model() -> None: visualizer._initialize_scene(PlanningSceneInfo(model=_model())) scene.register_model.assert_called_once() - state = JointState(name=["left/j1", "right/j1"], position=[0.1, 0.2]) + state = JointState( + name=["left/j1", "right/j1"], + position=np.array([0.1, 0.2], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) visualizer.update_state(VisualizationStateFrame(joint_state=state)) scene.update_current_model.assert_called_once_with(state) diff --git a/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py b/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py index 4d12c4409c6..cbcb00f7d57 100644 --- a/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py +++ b/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py @@ -17,8 +17,10 @@ from pathlib import Path from types import SimpleNamespace -from dimos_generated.geometry_msgs.msg import Pose +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from dimos.msgs.time import header_now @@ -91,7 +93,13 @@ def stop(self) -> None: def fake_robot_config(name: str) -> RobotModelConfig: return RobotModelConfig( model=RobotModel.from_file(Path(f"{name}.urdf")), - base_pose=PoseStamped(header=Header(frame_id=""), pose=Pose()), + base_pose=PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), joint_names=[], planning_groups=[ PlanningGroupDefinition( @@ -429,7 +437,13 @@ def test_visualizer_publish_preview_and_close_paths( monkeypatch: pytest.MonkeyPatch, ) -> None: calls: list[tuple[str, str]] = [] - current = JointState(name=["joint1"], position=[0.5]) + current = JointState( + name=["joint1"], + position=np.array([0.5], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) class FakeRuntime: url = "http://localhost:8095" @@ -482,7 +496,7 @@ def close(self) -> None: visualizer.update_state(VisualizationStateFrame(current)) visualizer.cancel_preview_animation() visualizer.animate_trajectory( - JointTrajectory(header=header_now(), joint_names=["arm/joint1"]), duration=1.5 + JointTrajectory(header=header_now(), joint_names=["arm/joint1"], points=[]), duration=1.5 ) visualizer.close() visualizer.update_state(VisualizationStateFrame(current)) @@ -602,7 +616,15 @@ def test_selected_display_mode_survives_primary_recreation_and_joint_updates( scene.register_model(prepared) current = scene._urdfs["current"] - scene.update_current_model(JointState(name=["joint1"], position=[0.75])) + scene.update_current_model( + JointState( + name=["joint1"], + position=np.array([0.75], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) assert current is not old_current assert scene.robot_display_mode == mode diff --git a/dimos/mapping/cli/map.py b/dimos/mapping/cli/map.py index ec0197b1e6b..dc521a90cc7 100644 --- a/dimos/mapping/cli/map.py +++ b/dimos/mapping/cli/map.py @@ -20,14 +20,18 @@ import subprocess from typing import TYPE_CHECKING, Any +from dimos_cdrgen.registry import encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time import typer +from dimos.mapping.cli.streams import select_stream + if TYPE_CHECKING: + from dimos_generated.geometry_msgs.msg import TransformStamped + from dimos_generated.sensor_msgs.msg import Image, PointCloud2 + from dimos.memory.stream import Stream from dimos.memory.type.observation import Observation - from dimos.msgs.geometry_msgs.Transform import Transform - from dimos.msgs.sensor_msgs.Image import Image - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.navigation.go2.loop_closure.pgo import PoseGraph PATH_THICKNESS = 0.01 @@ -109,7 +113,7 @@ def _accumulate( block_count: int, device: str, graph: PoseGraph | None = None, - register: Callable[[Observation[Any]], Transform | None] | None = None, + register: Callable[[Observation[Any]], TransformStamped | None] | None = None, carve_columns: bool = False, progress_cb: Callable[[Observation[Any]], None] | None = None, ) -> PointCloud2 | None: @@ -122,18 +126,24 @@ def _accumulate( Returns the final ``PointCloud2`` (or ``None`` if the input was empty). Disposal of the underlying ``VoxelGrid`` is handled by ``VoxelMapTransformer``. """ + from dimos_generated.geometry_msgs.msg import TransformStamped + from dimos_generated.std_msgs.msg import Header + import numpy as np + from dimos.mapping.voxels.module import VoxelMapTransformer + from dimos.msgs.geometry import transform_from_matrix, transform_matrix + from dimos.msgs.pointcloud import transform_cloud def prepared() -> Iterable[Observation[PointCloud2]]: for obs in obs_iter: if progress_cb is not None: progress_cb(obs) - if len(obs.data) == 0: + if obs.data.width * obs.data.height == 0: continue # sensor->world via `register`, unless the clouds are already # world-registered. graph adds the PGO correction on top # (correction ∘ tf), applied after the registration. - tf: Transform | None = None + tf: TransformStamped | None = None if register is not None: tf = register(obs) if tf is None: @@ -141,9 +151,14 @@ def prepared() -> Iterable[Observation[PointCloud2]]: if graph is not None: if obs.pose_tuple is None: continue - correction = graph.correction_at(obs.ts) - tf = correction if tf is None else correction + tf - yield obs if tf is None else obs.derive(data=obs.data.transform(tf)) + correction = transform_matrix(graph.correction_at(obs.ts).transform) + matrix = np.eye(4) if tf is None else transform_matrix(tf.transform) + tf = TransformStamped( + header=Header(stamp=obs.data.header.stamp, frame_id="world_corrected"), + child_frame_id=obs.data.header.frame_id, + transform=transform_from_matrix(correction @ matrix), + ) + yield obs if tf is None else obs.derive(data=transform_cloud(obs.data, tf)) vmt = VoxelMapTransformer( emit_every=0, # batch mode: emit once on exhaustion @@ -158,12 +173,26 @@ def prepared() -> Iterable[Observation[PointCloud2]]: def _denoise(cloud: PointCloud2 | None) -> PointCloud2 | None: """Statistical outlier removal via o3d; drops sparse floaters, keeps colors.""" - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 - - if cloud is None or len(cloud.pointcloud.points) < 20: + if cloud is None or cloud.width * cloud.height < 20: return cloud - clean, _ = cloud.pointcloud_tensor.remove_statistical_outliers(nb_neighbors=20, std_ratio=2.0) - return PointCloud2(pointcloud=clean, frame_id=cloud.frame_id, ts=cloud.ts) + import numpy as np + import open3d as o3d + + from dimos.msgs.pointcloud import ( + pointcloud_from_xyz, + pointcloud_from_xyz_rgb, + pointcloud_to_open3d, + ) + + if {field.name for field in cloud.fields} - {"x", "y", "z", "rgb"}: + raise ValueError("map denoising supports only XYZ/RGB point fields") + tensor = o3d.t.geometry.PointCloud.from_legacy(pointcloud_to_open3d(cloud)) + clean, _ = tensor.remove_statistical_outliers(nb_neighbors=20, std_ratio=2.0) + points = clean.point.positions.numpy() + if "colors" in clean.point: + colors = (np.clip(clean.point.colors.numpy(), 0.0, 1.0) * 255).astype(np.uint8) + return pointcloud_from_xyz_rgb(points, colors, header=cloud.header) + return pointcloud_from_xyz(points, header=cloud.header) def _log_reconstruction( @@ -180,17 +209,19 @@ def _log_reconstruction( bottom_cutoff: float | None = None, ) -> None: """Log maps, paths, the PGO graph, and markers to the active rerun recording.""" + from dimos_generated.geometry_msgs.msg import Point, Pose import rerun as rr import rerun.blueprint as rrb from dimos.memory.vis.color import Color - from dimos.msgs.geometry_msgs.Transform import Transform + from dimos.msgs.geometry import pose_from_matrix, pose_matrix, transform_matrix + from dimos.visualization.rerun.message_helpers import cloud_archetype rr.send_blueprint(rrb.Blueprint(rrb.Spatial3DView(origin="world"))) if global_map is not None: rr.log( "world/raw_map/pointcloud", - global_map.to_rerun(voxel_size=voxel / 2, bottom_cutoff=bottom_cutoff), + cloud_archetype(global_map, ui_radius=voxel / 2, bottom_cutoff=bottom_cutoff), static=True, ) if path: @@ -202,13 +233,13 @@ def _log_reconstruction( if pgo_map is not None: rr.log( "world/pgo_map/pointcloud", - pgo_map.to_rerun(voxel_size=voxel / 2, bottom_cutoff=bottom_cutoff), + cloud_archetype(pgo_map, ui_radius=voxel / 2, bottom_cutoff=bottom_cutoff), static=True, ) if full_pgo_map is not None: rr.log( "world/full_pgo_map/pointcloud", - full_pgo_map.to_rerun(voxel_size=voxel / 2, bottom_cutoff=bottom_cutoff), + cloud_archetype(full_pgo_map, ui_radius=voxel / 2, bottom_cutoff=bottom_cutoff), static=True, ) if pgo_path: @@ -225,8 +256,16 @@ def _log_reconstruction( if graph is not None and graph.loops: loop_strips = [ [ - (lc.source.translation.x, lc.source.translation.y, lc.source.translation.z), - (lc.target.translation.x, lc.target.translation.y, lc.target.translation.z), + ( + lc.source.transform.translation.x, + lc.source.transform.translation.y, + lc.source.transform.translation.z, + ), + ( + lc.target.transform.translation.x, + lc.target.transform.translation.y, + lc.target.transform.translation.z, + ), ] for lc in graph.loops ] @@ -269,23 +308,23 @@ def _log_reconstruction( pgo_centers: list[tuple[float, float, float]] = [] pgo_quats: list[tuple[float, float, float, float]] = [] for d in marker_dets: - raw_tf = Transform( - translation=d.data.center, - rotation=d.data.orientation, - frame_id="world", - child_frame_id=f"marker_{d.data.marker_id}", - ts=d.ts, + center = d.data.center + raw_pose = Pose( + position=Point(x=center.x, y=center.y, z=center.z), + orientation=d.data.orientation, + ) + corrected = pose_from_matrix( + transform_matrix(graph.correction_at(d.ts).transform) @ pose_matrix(raw_pose) ) - corrected = graph.correct(raw_tf) pgo_centers.append( - (corrected.translation.x, corrected.translation.y, corrected.translation.z) + (corrected.position.x, corrected.position.y, corrected.position.z) ) pgo_quats.append( ( - corrected.rotation.x, - corrected.rotation.y, - corrected.rotation.z, - corrected.rotation.w, + corrected.orientation.x, + corrected.orientation.y, + corrected.orientation.z, + corrected.orientation.w, ) ) _log_markers( @@ -300,9 +339,12 @@ def _log_reconstruction( def main( - dataset: str = typer.Argument(..., help="Dataset .db: bare name (cwd or data/) or path"), - lidar_stream: str = typer.Option( - "lidar", "--lidar", help="Lidar point-cloud stream to reconstruct" + dataset: str = typer.Argument(..., help="Dataset .db or .mcap: bare name or path"), + lidar_stream: str | None = typer.Option( + None, "--lidar", help="PointCloud2 stream; auto-select only a unique candidate" + ), + image_stream: str | None = typer.Option( + None, "--image", help="Image stream for --markers; auto-select only a unique candidate" ), seek: float = typer.Option(0.0, "--seek", help="Skip the first N seconds of the recording"), duration: float | None = typer.Option( @@ -318,9 +360,9 @@ def main( help="Run pose graph optimization and rebuild from spatially-deduped frames", ), pgo_tol: float = typer.Option( - 0.3, + 0.0, "--pgo-tol", - help="Spatial dedup tolerance (meters); applies to both raw and --pgo maps. 0 disables dedup (keep every posed frame)", + help="Spatial dedup tolerance (meters); applies to both raw and --pgo maps. 0 disables dedup (default; keep every registered frame)", ), block_count: int = typer.Option( 2_000_000, "--block-count", help="VoxelBlockGrid capacity (raw and PGO rebuilds)" @@ -328,7 +370,7 @@ def main( export: bool = typer.Option( False, "--export", - help="Export PGO map to ./.pc2.lcm in cwd (implies --pgo)", + help="Export PGO map to ./.pc2.cdr in cwd (implies --pgo)", ), full_pgo: bool = typer.Option( False, @@ -418,27 +460,41 @@ def main( ), ) -> None: """Rebuild a voxel map from a recorded SQLite dataset, write a .rrd, and open it in rerun.""" + from dimos_generated.sensor_msgs.msg import Image, PointCloud2 + from dimos_generated.std_msgs.msg import Header import rerun as rr - from dimos.memory.cli.dataset import open_store, resolve_dataset + from dimos.memory.cli.dataset import open_store, resolve_dataset, stream_payload_types from dimos.memory.transform import QualityWindow, SpeedLimit from dimos.memory.utils.progress import progress - from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo - from dimos.msgs.sensor_msgs.Image import Image - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + from dimos.msgs.camera_info import camera_info_from_yaml + from dimos.msgs.image import image_sharpness from dimos.navigation.go2.loop_closure.pgo import PGO from dimos.perception.fiducial.marker_transformer import DetectMarkers - from dimos.robot.unitree.go2.connection import BASE_TO_OPTICAL, _camera_info_static + from dimos.robot.unitree.go2.camera_calibration import front_camera_calibration + from dimos.robot.unitree.go2.connection import BASE_TO_OPTICAL from dimos.visualization.rerun.init import rerun_init db_path = resolve_dataset(dataset) store = open_store(db_path) + try: + types = stream_payload_types(store) + selected_lidar = select_stream(types, PointCloud2, lidar_stream, "--lidar") + selected_image = ( + select_stream(types, Image, image_stream, "--image") + if markers or image_stream is not None + else None + ) + assert selected_lidar is not None + except Exception: + store.stop() + raise if out is None: out = Path.cwd() / f"{db_path.stem}.rrd" if export or full_pgo: pgo = True - lidar = store.stream(lidar_stream, PointCloud2).from_time(seek or None).to_time(duration) + lidar = store.stream(selected_lidar, PointCloud2).from_time(seek or None).to_time(duration) print(lidar.summary()) @@ -454,7 +510,7 @@ def main( tf_buf = StreamTF.from_store(store) # Streams are homogeneous: read the cloud frame from the first observation. first_obs = next(iter(lidar), None) - cloud_frame: str | None = first_obs.data.frame_id if first_obs is not None else None + cloud_frame: str | None = first_obs.data.header.frame_id if first_obs is not None else None world = frame if world is None and first_obs is not None and cloud_frame is not None: @@ -473,7 +529,7 @@ def main( # Registration: sensor-frame clouds get a per-frame tf lookup lifting them # into the world frame (frames with no tf answer are dropped); clouds # already stamped with the world frame accumulate verbatim (register=None). - register: Callable[[Observation[Any]], Transform | None] | None = None + register: Callable[[Observation[Any]], TransformStamped | None] | None = None if first_obs is not None and cloud_frame is not None and cloud_frame != world: # Fail fast when registration is impossible: probe the first cloud's # timestamp (unbounded tolerance — "possible at all", not "in range"). @@ -490,13 +546,14 @@ def main( print(f"registering clouds {world!r} ← {cloud_frame!r} via tf") buf = tf_buf - def _register(obs: Observation[Any]) -> Transform | None: - return buf.get(world, obs.data.frame_id, time_point=obs.ts, time_tolerance=tf_tolerance) + def _register(obs: Observation[Any]) -> TransformStamped | None: + return buf.get( + world, obs.data.header.frame_id, time_point=obs.ts, time_tolerance=tf_tolerance + ) register = _register elif cloud_frame is not None: print(f"clouds already in world frame {world!r}; accumulating verbatim") - print("warning: trajectory positions come from stored obs.pose (old dataset)") def _position(obs: Observation[Any]) -> tuple[float, float, float] | None: """Trajectory position for dedup/path: registration tf, else the stored pose.""" @@ -504,45 +561,66 @@ def _position(obs: Observation[Any]) -> tuple[float, float, float] | None: tf = register(obs) if tf is None: return None - return (tf.translation.x, tf.translation.y, tf.translation.z) + return ( + tf.transform.translation.x, + tf.transform.translation.y, + tf.transform.translation.z, + ) pose = obs.pose # Reject placeholder poses: zero translation OR uninitialized rotation. # Same condition as pgo_keyframes so dedup and PGO see the same frames. - if pose is not None and not (pose.position.is_zero() or pose.orientation.is_zero()): + if ( + pose is not None + and any((pose.position.x, pose.position.y, pose.position.z)) + and any( + (pose.orientation.x, pose.orientation.y, pose.orientation.z, pose.orientation.w) + ) + ): return (pose.position.x, pose.position.y, pose.position.z) return None - # Spatial dedup: bucket frames by 3D cell using the trajectory position, - # keep the latest per cell. Shared by raw and PGO rebuilds. Doesn't touch - # obs.data so it stays cheap (no pointcloud loading). With pgo_tol<=0 the - # bucketing is disabled and every positioned frame is kept (keyed by index). - seen: dict[Any, tuple[Observation[Any], tuple[float, float, float]]] = {} + # World-frame accumulation needs no trajectory. Dedup and PGO do. + seen: dict[Any, Observation[Any]] = {} + path: list[tuple[float, float, float]] = [] for i, obs in enumerate(lidar): pos = _position(obs) - if pos is None: - continue - if pgo_tol > 0: - # math.floor so negative coords bucket consistently; int() truncates - # toward zero and silently folds -0.5 and 0.5 into the same cell. - key: Any = ( - math.floor(pos[0] / pgo_tol), - math.floor(pos[1] / pgo_tol), - math.floor(pos[2] / pgo_tol), + if register is not None and pos is None: + raise typer.BadParameter( + f"Missing registration TF at timestamp {obs.ts}", param_hint="--frame" ) - else: - key = i - seen[key] = (obs, pos) - - n_kept = len(seen) - pct = 100 * n_kept / total if total else 0 - if pgo_tol > 0: - print(f"dedup: kept [{n_kept}/{total}] frames ({pct:.1f}%) at tol={pgo_tol}m") - else: - print(f"dedup: disabled, kept all [{n_kept}/{total}] positioned frames") + if pgo_tol > 0 and pos is None: + raise typer.BadParameter( + "Spatial dedup requires trajectory positions; use --pgo-tol 0 for plain accumulation", + param_hint="--pgo-tol", + ) + if pgo and ( + obs.pose is None + or not any((obs.pose.position.x, obs.pose.position.y, obs.pose.position.z)) + or not any( + ( + obs.pose.orientation.x, + obs.pose.orientation.y, + obs.pose.orientation.z, + obs.pose.orientation.w, + ) + ) + ): + raise typer.BadParameter( + "PGO requires valid stored trajectory poses; registration TF alone does not populate them", + param_hint="--pgo", + ) + key: Any = i + if pgo_tol > 0: + assert pos is not None + key = tuple(math.floor(value / pgo_tol) for value in pos) + seen[key] = obs - # Dict insertion order = lidar iteration order = chronological. - kept = [obs for obs, _ in seen.values()] - path: list[tuple[float, float, float]] = [pos for _, pos in seen.values()] + kept = list(seen.values()) + path = [pos for obs in kept if (pos := _position(obs)) is not None] + n_kept = len(kept) + if not n_kept: + raise typer.BadParameter("No point-cloud observations in the selected interval") + print(f"dedup: kept [{n_kept}/{total}] frames at tol={pgo_tol}m") pgo_map = None pgo_path: list[tuple[float, float, float]] = [] @@ -553,7 +631,11 @@ def _position(obs: Observation[Any]) -> tuple[float, float, float] | None: graph = lidar.tap(bar).transform(PGO()).last().data pgo_path = [ - (kf.optimized.translation.x, kf.optimized.translation.y, kf.optimized.translation.z) + ( + kf.optimized.transform.translation.x, + kf.optimized.transform.translation.y, + kf.optimized.transform.translation.z, + ) for kf in graph.keyframes ] @@ -596,6 +678,9 @@ def _position(obs: Observation[Any]) -> tuple[float, float, float] | None: progress_cb=bar, ) + if global_map is None: + raise typer.BadParameter("Selected observations contain no points to accumulate") + if denoise: print("denoising maps (statistical outlier removal)...") global_map = _denoise(global_map) @@ -609,7 +694,8 @@ def _position(obs: Observation[Any]) -> tuple[float, float, float] | None: # (verified: matches lidar_base_pose + BASE_TO_OPTICAL to ~1mm). With # --image-pose, swap that stored pose for a different source (e.g. # fastlio_odometry), composing the base→optical mount onto it first. - color_image = store.stream("color_image", Image).from_time(seek or None).to_time(duration) + assert selected_image is not None + color_image = store.stream(selected_image, Image).from_time(seek or None).to_time(duration) n_images = color_image.count() if image_pose is not None: from dimos.mapping.cli.pose_fill import pose_fill @@ -619,7 +705,13 @@ def _position(obs: Observation[Any]) -> tuple[float, float, float] | None: ) print(f"re-posing color_image from {image_pose!r} + camera optical mount") color_image = pose_fill(color_image, src_pose, tolerance=0.1, mount=BASE_TO_OPTICAL) - cam_info = CameraInfo.from_yaml(str(camera_info)) if camera_info else _camera_info_static() + cam_info = ( + camera_info_from_yaml( + camera_info, header=Header(frame_id="camera_optical", stamp=Time(sec=0, nanosec=0)) + ) + if camera_info + else front_camera_calibration() + ) xf = DetectMarkers( camera_info=cam_info, marker_length_m=marker_size, @@ -631,7 +723,7 @@ def _position(obs: Observation[Any]) -> tuple[float, float, float] | None: # than the limits. Defaults match replay_marker.py so positions agree. with progress(n_images, "detecting markers") as bar: pipeline: Stream[Image] = color_image.tap(bar).transform( - QualityWindow(lambda img: img.sharpness, window=marker_quality_window) + QualityWindow(image_sharpness, window=marker_quality_window) ) if marker_max_speed > 0: pipeline = pipeline.transform( @@ -677,14 +769,15 @@ def _position(obs: Observation[Any]) -> tuple[float, float, float] | None: subprocess.Popen(["rerun", str(out)]) if export and pgo_map is not None: - out_path = Path.cwd() / f"{db_path.stem}.pc2.lcm" + out_path = Path.cwd() / f"{db_path.stem}.pc2.cdr" print(f"exporting PGO twopass map to {out_path}...") - out_path.write_bytes(pgo_map.lcm_encode()) + out_path.write_bytes(cdr_encode(pgo_map)) print(f"wrote {out_path}") print() print("load back with:") - print(" from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2") - print(f' pcd = PointCloud2.lcm_decode(open("{out_path.name}", "rb").read())') + print(" from dimos_generated.sensor_msgs.msg import PointCloud2") + print(" from dimos_cdrgen.registry import decode") + print(f' pcd = decode(open("{out_path.name}", "rb").read(), PointCloud2)') if __name__ == "__main__": diff --git a/dimos/mapping/cli/test_pgo_accumulate.py b/dimos/mapping/cli/test_pgo_accumulate.py index 0c25f3e57bf..29f2561f4b0 100644 --- a/dimos/mapping/cli/test_pgo_accumulate.py +++ b/dimos/mapping/cli/test_pgo_accumulate.py @@ -12,23 +12,27 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos_generated.geometry_msgs.msg import Transform, TransformStamped, Vector3 +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 from dimos_generated.std_msgs.msg import Header import numpy as np from dimos.mapping.cli.map import _accumulate -from dimos.mapping.loop_closure.pgo import PoseGraph from dimos.memory.store.memory import MemoryStore from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_xyz from dimos.msgs.time import time_from_nanoseconds +from dimos.navigation.go2.loop_closure.pgo import PoseGraph def test_pgo_accumulate_applies_generated_correction_before_mapping(mocker): graph = mocker.Mock(spec=PoseGraph) graph.correction_at.return_value = TransformStamped( - header=Header(frame_id="world_corrected"), + header=Header(frame_id="world_corrected", stamp=Time(sec=0, nanosec=0)), child_frame_id="world", - transform=Transform(translation=Vector3(x=2.0)), + transform=Transform( + translation=Vector3(x=2.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) mapper = mocker.Mock(side_effect=lambda observations: observations) mocker.patch("dimos.mapping.voxels.module.VoxelMapTransformer", return_value=mapper) diff --git a/dimos/mapping/cli/test_pose_fill.py b/dimos/mapping/cli/test_pose_fill.py index 4db1cab4c57..f4e6588f3c4 100644 --- a/dimos/mapping/cli/test_pose_fill.py +++ b/dimos/mapping/cli/test_pose_fill.py @@ -22,6 +22,7 @@ import math +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, Pose, @@ -30,6 +31,7 @@ TransformStamped, Vector3, ) +from dimos_generated.std_msgs.msg import Header import pytest from dimos.mapping.cli.pose_fill import pose_fill @@ -42,7 +44,13 @@ def test_pose_fill_attaches_nearest_pose() -> None: target = store.stream("image", str) poses = store.stream("odom", Pose) target.append("img0", ts=0.0) - poses.append(Pose(position=Point(x=1.0), orientation=Quaternion(w=1)), ts=0.001) + poses.append( + Pose( + position=Point(x=1.0, y=0.0, z=0.0), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), + ts=0.001, + ) out = pose_fill(target, poses, tolerance=0.05).to_list() @@ -59,9 +67,20 @@ def test_pose_fill_mount_composes_static_child_transform() -> None: poses = store.stream("odom", Pose) target.append("img0", ts=0.0) # Base pose 1m forward in x; mount offsets 1m in y (identity rotations). - poses.append(Pose(position=Point(x=1.0), orientation=Quaternion(w=1)), ts=0.001) + poses.append( + Pose( + position=Point(x=1.0, y=0.0, z=0.0), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), + ts=0.001, + ) mount = TransformStamped( - transform=Transform(translation=Vector3(y=1), rotation=Quaternion(w=1)) + transform=Transform( + translation=Vector3(y=1, x=0.0, z=0.0), + rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + child_frame_id="", ) out = pose_fill(target, poses, tolerance=0.05, mount=mount).to_list() @@ -80,13 +99,18 @@ def test_pose_fill_rotates_mount_offset_in_base_axes() -> None: target.append("camera frame", ts=5.0) poses.append( Pose( - position=Point(x=2), - orientation=Quaternion(z=math.sqrt(0.5), w=math.sqrt(0.5)), + position=Point(x=2, y=0.0, z=0.0), + orientation=Quaternion(z=math.sqrt(0.5), w=math.sqrt(0.5), x=0.0, y=0.0), ), ts=5.001, ) mount = TransformStamped( - transform=Transform(translation=Vector3(x=1), rotation=Quaternion(w=1)) + transform=Transform( + translation=Vector3(x=1, y=0.0, z=0.0), + rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + child_frame_id="", ) result = pose_fill(target, poses, tolerance=0.05, mount=mount).to_list() assert len(result) == 1 @@ -100,5 +124,10 @@ def test_pose_fill_drops_unmatched_frame_without_loading_payload() -> None: target = store.stream("image", str) poses = store.stream("odom", Pose) target.append("unmatched", ts=1.0) - poses.append(Pose(position=Point(x=1), orientation=Quaternion(w=1)), ts=3.0) + poses.append( + Pose( + position=Point(x=1, y=0.0, z=0.0), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), + ts=3.0, + ) assert pose_fill(target, poses, tolerance=0.05).to_list() == [] diff --git a/dimos/mapping/cli/test_stream_selection.py b/dimos/mapping/cli/test_stream_selection.py index b5e01b3494c..10379f47a56 100644 --- a/dimos/mapping/cli/test_stream_selection.py +++ b/dimos/mapping/cli/test_stream_selection.py @@ -14,6 +14,9 @@ """Select recorded streams by compatible type, with explicit ambiguity errors.""" +from click import unstyle +from dimos_cdrgen.registry import encode as cdr_encode, schema as cdr_schema +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 from dimos_generated.sensor_msgs.msg import Image, PointCloud2 from dimos_generated.std_msgs.msg import Header @@ -29,14 +32,17 @@ def recording(path, clouds, images=(), frame="world", transforms=()): - cloud = pointcloud_from_xyz(np.array([[1.0, 2.0, 3.0]]), header=Header(frame_id=frame)) + cloud = pointcloud_from_xyz( + np.array([[1.0, 2.0, 3.0]]), header=Header(frame_id=frame, stamp=Time(sec=0, nanosec=0)) + ) image = Image( - header=Header(frame_id="camera"), + header=Header(frame_id="camera", stamp=Time(sec=0, nanosec=0)), height=1, width=1, encoding="rgb8", step=3, - data=[255, 0, 0], + data=np.array([255, 0, 0], dtype=np.uint8), + is_bigendian=0, ) with CdrMcapWriter(path) as writer: for name, value in ( @@ -46,9 +52,9 @@ def recording(path, clouds, images=(), frame="world", transforms=()): ): writer.write( name, - value.encode(), - schema_name=value.msg_name, - schema=value.schema, + cdr_encode(value), + schema_name=value.__msgtype__, + schema=cdr_schema(value.__msgtype__), log_time_ns=1_000_000_000, publish_time_ns=1_000_000_000, ) @@ -71,7 +77,7 @@ def test_multiple_clouds_require_selection_even_when_one_is_named_lidar(tmp_path result = CliRunner().invoke(map_app, [command, str(source), "--no-gui", "--out", str(output)]) assert result.exit_code == 2, result.output assert "Multiple compatible streams: lidar, second_cloud" in result.output - assert "--lidar" in result.output + assert "--lidar" in unstyle(result.output) assert not output.exists() @@ -123,7 +129,7 @@ def test_multiple_images_require_selection(tmp_path): result = CliRunner().invoke(map_app, ["replay", str(source), "--no-gui", "--out", str(output)]) assert result.exit_code == 2, result.output assert "Multiple compatible streams: left, right" in result.output - assert "--image" in result.output + assert "--image" in unstyle(result.output) assert not output.exists() @@ -178,9 +184,12 @@ def test_world_cloud_trajectory_options_require_pose(tmp_path, option): def test_sensor_cloud_requires_real_registration_tf(tmp_path, with_tf): source, output = tmp_path / "source.mcap", tmp_path / "out.rrd" tf = TransformStamped( - header=Header(frame_id="world"), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), child_frame_id="sensor", - transform=Transform(translation=Vector3(x=2.0), rotation=Quaternion(w=1.0)), + transform=Transform( + translation=Vector3(x=2.0, y=0.0, z=0.0), + rotation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), + ), ) recording(source, ["cloud"], frame="sensor", transforms=[tf] if with_tf else []) result = CliRunner().invoke( diff --git a/dimos/mapping/cli/test_view.py b/dimos/mapping/cli/test_view.py index d7ddb137033..7fbafc3af9e 100644 --- a/dimos/mapping/cli/test_view.py +++ b/dimos/mapping/cli/test_view.py @@ -17,6 +17,8 @@ from pathlib import Path +from dimos_cdrgen.registry import encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.std_msgs.msg import Header import numpy as np import rerun as rr @@ -28,9 +30,10 @@ def test_view_cdr_cloud_writes_offline_recording(tmp_path: Path, capsys): pc2 = tmp_path / "map.pc2.cdr" cloud = pointcloud_from_xyz( - np.random.default_rng(0).random((256, 3)), header=Header(frame_id="world") + np.random.default_rng(0).random((256, 3)), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), ) - pc2.write_bytes(cloud.encode()) + pc2.write_bytes(cdr_encode(cloud)) out = tmp_path / "view.rrd" try: main(pc2, voxel=0.05, bottom_cutoff=None, out=out) diff --git a/dimos/mapping/cli/view.py b/dimos/mapping/cli/view.py index 60c330ad25d..adafecda282 100644 --- a/dimos/mapping/cli/view.py +++ b/dimos/mapping/cli/view.py @@ -18,6 +18,7 @@ from pathlib import Path +from dimos_cdrgen.registry import decode as cdr_decode import typer @@ -40,7 +41,7 @@ def main( from dimos.visualization.rerun.init import rerun_init from dimos.visualization.rerun.message_helpers import cloud_archetype - cloud = PointCloud2.decode(path.read_bytes()) + cloud = cdr_decode(path.read_bytes(), PointCloud2) print(f"{path}: {cloud.width * cloud.height} points in {cloud.header.frame_id!r}") rerun_init("dimos map view") diff --git a/dimos/mapping/dim_slam/rust/Cargo.lock b/dimos/mapping/dim_slam/rust/Cargo.lock index 37dafef1184..e92995fc370 100644 --- a/dimos/mapping/dim_slam/rust/Cargo.lock +++ b/dimos/mapping/dim_slam/rust/Cargo.lock @@ -593,7 +593,9 @@ dependencies = [ name = "dimos-generated-messages" version = "0.1.0" dependencies = [ + "heck", "re_cdr", + "ros2msg", "serde", "serde-big-array", ] @@ -1008,6 +1010,12 @@ version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + [[package]] name = "hermit-abi" version = "0.5.2" @@ -1943,6 +1951,16 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + [[package]] name = "proc-macro-crate" version = "3.5.0" @@ -2309,6 +2327,23 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "ros2msg" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fe291be5cb6d4b8e6bf9c9de80c40bd1cb09ecac50deafa2fdcae0027e04988" +dependencies = [ + "heck", + "pest", + "pest_derive", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "syn 2.0.119", + "thiserror", +] + [[package]] name = "rustc-hash" version = "2.1.3" diff --git a/dimos/mapping/dim_slam/rust/Cargo.toml b/dimos/mapping/dim_slam/rust/Cargo.toml index 2807007898c..43b47abc1ef 100644 --- a/dimos/mapping/dim_slam/rust/Cargo.toml +++ b/dimos/mapping/dim_slam/rust/Cargo.toml @@ -16,7 +16,7 @@ dim_slam = { git = "https://github.com/dimensionalOS/dimSLAM", tag = "v0.4.1" } # and build.rs needs it for the rpath. cu_vslam_rs = "0.2.0" dimos-module = { path = "../../../../native/rust/dimos-module" } -dimos-generated-messages = { path = "../../../message_codegen" } +dimos-generated-messages = { path = "../../../../packages/dimos-generated/build/message-source/rust" } tokio = { version = "1", features = ["rt-multi-thread", "macros"] } tracing = "0.1" # native_config expands to ::serde and ::validator derives, so both must be direct deps. diff --git a/dimos/mapping/dim_slam/rust/flake.nix b/dimos/mapping/dim_slam/rust/flake.nix index fe2288ee502..c570aa9f9ff 100644 --- a/dimos/mapping/dim_slam/rust/flake.nix +++ b/dimos/mapping/dim_slam/rust/flake.nix @@ -29,6 +29,8 @@ sdkPackages = nixpkgs.lib.filterAttrs (name: _: nixpkgs.lib.hasPrefix "sdk-" name) cu-vslam-rs.packages.${system}; variants = map (nixpkgs.lib.removePrefix "sdk-") (builtins.attrNames sdkPackages); + messageSource = import "${dimos-repo}/native/message-source.nix" { inherit pkgs; }; + src = pkgs.runCommand "dim-slam-module-src" {} '' mkdir -p $out/dimos/mapping/dim_slam/rust cp -r ${./src} $out/dimos/mapping/dim_slam/rust/src @@ -40,8 +42,8 @@ cp -r ${dimos-repo}/native/rust/dimos-module $out/native/rust/dimos-module cp -r ${dimos-repo}/native/rust/dimos-module-macros $out/native/rust/dimos-module-macros cp -r ${dimos-repo}/native/rust/dimos-lcm-transport $out/native/rust/dimos-lcm-transport - mkdir -p $out/dimos/message_codegen - cp -r ${dimos-repo}/dimos/message_codegen/. $out/dimos/message_codegen/ + mkdir -p $out/packages/dimos-generated/build/message-source + cp -r ${messageSource.source}/rust $out/packages/dimos-generated/build/message-source/rust ''; generatedCargoNix = crate2nix.tools.${system}.generatedCargoNix { @@ -55,7 +57,6 @@ inherit pkgs; buildRustCrateForPkgs = cratePkgs: cratePkgs.buildRustCrate.override { defaultCrateOverrides = cratePkgs.defaultCrateOverrides // { - dimos-generated-messages = _: { nativeBuildInputs = [cratePkgs.python3]; }; # cu_vslam_rs's build.rs compiles its shim against this SDK. cu_vslam_rs = _: { CUVSLAM_SDK_DIR = sdkPackage; }; # buildRustCrate names DEP_ vars after the crate, cargo after the diff --git a/dimos/mapping/dim_slam/rust/src/main.rs b/dimos/mapping/dim_slam/rust/src/main.rs index aa938075c39..c7441c9fc95 100644 --- a/dimos/mapping/dim_slam/rust/src/main.rs +++ b/dimos/mapping/dim_slam/rust/src/main.rs @@ -17,8 +17,10 @@ mod msg_convert; use dim_slam::{ CameraConfig, DimSlamCore, DimSlamCoreConfig, ImuConfig, InitialStds, SourceConfig, }; -use dimos_generated_messages::nav_msgs::msg::Odometry; -use dimos_generated_messages::sensor_msgs::msg::{CameraInfo, Image, Imu, PointCloud2}; +use dimos_generated_messages::nav_msgs::msg::odometry::Odometry; +use dimos_generated_messages::sensor_msgs::msg::{ + camera_info::CameraInfo, image::Image, imu::Imu, point_cloud2::PointCloud2, +}; use dimos_module::cdr; use dimos_module::{native_config, run_with_transport, Input, Module, Output, Tf}; use msg_convert::{ diff --git a/dimos/mapping/dim_slam/rust/src/msg_convert.rs b/dimos/mapping/dim_slam/rust/src/msg_convert.rs index 2f4b15883fc..402ed67abcb 100644 --- a/dimos/mapping/dim_slam/rust/src/msg_convert.rs +++ b/dimos/mapping/dim_slam/rust/src/msg_convert.rs @@ -16,11 +16,14 @@ use dim_slam::nalgebra::{Isometry3, Matrix6, Quaternion, Translation3, UnitQuaternion, Vector3}; use dim_slam::{CameraModel, ImageFrame, ImuSample, OdometryEstimate, PointCloud, Twist}; -use dimos_generated_messages::builtin_interfaces::msg::Time; +use dimos_generated_messages::builtin_interfaces::msg::time::Time; use dimos_generated_messages::geometry_msgs::msg as geometry_msgs; -use dimos_generated_messages::nav_msgs::msg::Odometry; -use dimos_generated_messages::sensor_msgs::msg::{CameraInfo, Image, Imu, PointCloud2, PointField}; -use dimos_generated_messages::std_msgs::msg::Header; +use dimos_generated_messages::nav_msgs::msg::odometry::Odometry; +use dimos_generated_messages::sensor_msgs::msg::{ + camera_info::CameraInfo, image::Image, imu::Imu, point_cloud2::PointCloud2, + point_field::PointField, +}; +use dimos_generated_messages::std_msgs::msg::header::Header; use dimos_module::{Tf, Transform}; const NS_PER_SEC: i64 = 1_000_000_000; diff --git a/dimos/mapping/experimental/mesh.py b/dimos/mapping/experimental/mesh.py index cefd1f823ad..82a930c43ea 100644 --- a/dimos/mapping/experimental/mesh.py +++ b/dimos/mapping/experimental/mesh.py @@ -21,6 +21,9 @@ import time from typing import TYPE_CHECKING, Any +from dimos_generated.dimos_msgs.msg import RegionPointCloud2, TriangleMesh +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header import numpy as np from reactivex.disposable import Disposable @@ -28,8 +31,9 @@ from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out from dimos.memory.transform import Transformer -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.shape_msgs.TriangleMesh import TriangleMesh +from dimos.msgs.mesh import triangle_mesh_from_arrays, triangle_mesh_to_rerun +from dimos.msgs.pointcloud import pointcloud_xyz +from dimos.msgs.time import time_from_seconds from dimos.visualization.rerun.bridge import RerunEntry if TYPE_CHECKING: @@ -225,8 +229,18 @@ def __call__( if self._mesher is None: self._mesher = OccupancyMesher(self.voxel_size, self.iso, self.device, self.simplify) for obs in upstream: - for key, v, n, f in self._mesher.mesh(obs.data.points_f32()): - yield obs.derive(data=TriangleMesh(v, f, n, key, obs.ts)) + for key, v, n, f in self._mesher.mesh(pointcloud_xyz(obs.data)): + yield obs.derive( + data=triangle_mesh_from_arrays( + v, + f, + n, + key, + header=Header( + stamp=time_from_seconds(obs.ts), frame_id=obs.data.header.frame_id + ), + ) + ) class MeshModuleConfig(ModuleConfig): @@ -311,7 +325,7 @@ class MeshModule(Module): config: MeshModuleConfig - map_regions: In[PointCloud2] + map_regions: In[RegionPointCloud2] mesh: Out[TriangleMesh] def __init__(self, **kwargs: Any) -> None: @@ -342,9 +356,9 @@ def stop(self) -> None: self._worker.join(timeout=5.0) super().stop() - def _on_region(self, msg: PointCloud2) -> None: + def _on_region(self, msg: RegionPointCloud2) -> None: with self._lock: - if self._queue.update(msg.seq, msg.points_f32(), time.time()): + if self._queue.update(msg.region_id, pointcloud_xyz(msg.cloud), time.time()): self._dirty.set() def _run(self) -> None: @@ -368,7 +382,11 @@ def _run(self) -> None: for key, v, n, f in self._mesher.mesh_chunks(points, np.array(chunks, np.int64)): if self._stopping: return - self.mesh.publish(TriangleMesh(v, f, n, key, ts)) + self.mesh.publish( + triangle_mesh_from_arrays( + v, f, n, key, header=Header(stamp=time_from_seconds(ts), frame_id="world") + ) + ) MESH_ENTITY = "world/mesh" @@ -381,7 +399,7 @@ class MeshColours: ``z_range`` pins it. Works as a keyed rerun bridge renderer and in replays. """ - keyed_by_seq = True + keyed_by_region = True def __init__( self, @@ -398,15 +416,21 @@ def __init__( self._range: tuple[float, float] | None = None def __call__(self, msg: TriangleMesh) -> list[RerunEntry]: - out = [RerunEntry(_chunk_path(msg.key), msg.to_rerun(), self.static)] + out = [ + RerunEntry( + _chunk_path(tuple(int(x) for x in msg.key)), + triangle_mesh_to_rerun(msg), + self.static, + ) + ] if len(msg.faces) == 0: - self._heights.pop(msg.key, None) + self._heights.pop(tuple(int(x) for x in msg.key), None) return out - self._heights[msg.key] = msg.vertices[:, 2].copy() + self._heights[tuple(int(x) for x in msg.key)] = msg.vertices.reshape(-1, 3)[:, 2].copy() lo, hi = self.z_range or np.percentile( np.concatenate(list(self._heights.values())), [2, 98] ) - keys = [msg.key] + keys = [tuple(int(x) for x in msg.key)] if ( self._range is None or max(abs(lo - self._range[0]), abs(hi - self._range[1])) > self.recolour_m diff --git a/dimos/mapping/experimental/mesh_rrd.py b/dimos/mapping/experimental/mesh_rrd.py index 977d730da54..176d9e634dd 100644 --- a/dimos/mapping/experimental/mesh_rrd.py +++ b/dimos/mapping/experimental/mesh_rrd.py @@ -19,6 +19,8 @@ from collections.abc import Iterator import itertools +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header import typer from dimos.mapping.experimental.mesh import Mesh, MeshColours @@ -29,8 +31,9 @@ from dimos.memory.transform import FnTransformer from dimos.memory.type.observation import Observation from dimos.memory.utils.progress import progress -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.image import image_to_rgb +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.msgs.time import time_from_seconds from dimos.utils.data import get_data @@ -90,7 +93,7 @@ def log_odom(obs: Observation[PointCloud2]) -> None: ) rr.log( "world/lidar/scan", - rr.Points3D(obs.data.points_f32(), radii=0.01, colors=[255, 255, 255]), + rr.Points3D(pointcloud_xyz(obs.data), radii=0.01, colors=[255, 255, 255]), ) db = dataset if dataset.endswith(".db") else str(get_data(f"{dataset}.db")) @@ -133,11 +136,16 @@ def log_camera(obs: Observation[PointCloud2]) -> None: pending[:] = [img] return rr.set_time("time", timestamp=img.ts) - rr.log("camera/image", img.data.to_rerun()) + rr.log("camera/image", rr.Image(image_to_rgb(img.data))) pending.clear() global_map: FnTransformer[PointCloud2, PointCloud2] = FnTransformer( - lambda o: o.derive(data=PointCloud2.from_numpy(ray.mapper.global_map())) + lambda o: o.derive( + data=pointcloud_from_xyz( + ray.mapper.global_map(), + header=Header(stamp=time_from_seconds(o.ts), frame_id=world_frame), + ) + ) ) with progress(lidar.count(), "meshing") as bar: meshes = ( diff --git a/dimos/mapping/experimental/test_mesh.py b/dimos/mapping/experimental/test_mesh.py index 2180611afc7..c8ccca2eab7 100644 --- a/dimos/mapping/experimental/test_mesh.py +++ b/dimos/mapping/experimental/test_mesh.py @@ -51,12 +51,14 @@ def test_most_recently_changed_chunk_goes_first() -> None: def test_colours_recolour_every_chunk_once_the_range_moves() -> None: + from dimos_generated.dimos_msgs.msg import TriangleMesh + from dimos.mapping.experimental.mesh import MeshColours - from dimos.msgs.shape_msgs.TriangleMesh import TriangleMesh + from dimos.msgs.mesh import triangle_mesh_from_arrays def chunk(key: tuple[int, int, int], z: float) -> TriangleMesh: v = np.array([[0, 0, z], [1, 0, z], [0, 1, z + 1]], np.float32) - return TriangleMesh(v, np.array([[0, 1, 2]], np.uint32), key=key) + return triangle_mesh_from_arrays(v, np.array([[0, 1, 2]], np.uint32), key=key) colours = MeshColours() assert len(colours(chunk((0, 0, 0), 0.0))) == 2 # geometry + its colours diff --git a/dimos/mapping/occupancy/conftest.py b/dimos/mapping/occupancy/conftest.py index 0847c3199cd..1ff9219334c 100644 --- a/dimos/mapping/occupancy/conftest.py +++ b/dimos/mapping/occupancy/conftest.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos_generated.geometry_msgs.msg import Pose, Quaternion +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest @@ -28,10 +30,15 @@ def occupancy() -> OccupancyGrid: info=MapMetaData( width=cells.shape[1], height=cells.shape[0], - resolution=0.05, - origin=Pose(orientation=Quaternion(w=1)), + resolution=float(np.float32(0.05)), + origin=Pose( + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + position=Point(x=0.0, y=0.0, z=0.0), + ), + map_load_time=Time(sec=0, nanosec=0), ), - data=cells.astype(np.int8).ravel(), + data=np.asarray(cells.astype(np.int8).ravel(), dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) diff --git a/dimos/mapping/occupancy/gradient.py b/dimos/mapping/occupancy/gradient.py index 0478058b1c5..81d732aa116 100644 --- a/dimos/mapping/occupancy/gradient.py +++ b/dimos/mapping/occupancy/gradient.py @@ -80,7 +80,9 @@ def gradient( # Create new OccupancyGrid with gradient gradient_grid = OccupancyGrid( - header=occupancy_grid.header, info=occupancy_grid.info, data=gradient_data.ravel() + header=occupancy_grid.header, + info=occupancy_grid.info, + data=np.asarray(gradient_data.ravel(), dtype=np.int8), ) return gradient_grid @@ -130,7 +132,9 @@ def voronoi_gradient( gradient_data = np.zeros_like(occupancy_view(occupancy_grid), dtype=np.int8) gradient_data[unknown_mask] = -1 return OccupancyGrid( - header=occupancy_grid.header, info=occupancy_grid.info, data=gradient_data.ravel() + header=occupancy_grid.header, + info=occupancy_grid.info, + data=np.asarray(gradient_data.ravel(), dtype=np.int8), ) # Label connected obstacle regions (clusters) @@ -195,5 +199,7 @@ def voronoi_gradient( gradient_data[unknown_mask] = -1 return OccupancyGrid( - header=occupancy_grid.header, info=occupancy_grid.info, data=gradient_data.ravel() + header=occupancy_grid.header, + info=occupancy_grid.info, + data=np.asarray(gradient_data.ravel(), dtype=np.int8), ) diff --git a/dimos/mapping/occupancy/operations.py b/dimos/mapping/occupancy/operations.py index 3771bcae436..87253b6c4ae 100644 --- a/dimos/mapping/occupancy/operations.py +++ b/dimos/mapping/occupancy/operations.py @@ -51,7 +51,9 @@ def smooth_occupied( result_grid[unsupported] = 0 return OccupancyGrid( - header=occupancy_grid.header, info=occupancy_grid.info, data=result_grid.ravel() + header=occupancy_grid.header, + info=occupancy_grid.info, + data=np.asarray(result_grid.ravel(), dtype=np.int8), ) @@ -76,4 +78,6 @@ def overlay_occupied(base: OccupancyGrid, overlay: OccupancyGrid) -> OccupancyGr overlay_occupied_mask = occupancy_view(overlay) >= 100 result_grid[overlay_occupied_mask] = 100 - return OccupancyGrid(header=base.header, info=base.info, data=result_grid.ravel()) + return OccupancyGrid( + header=base.header, info=base.info, data=np.asarray(result_grid.ravel(), dtype=np.int8) + ) diff --git a/dimos/mapping/occupancy/path_resampling.py b/dimos/mapping/occupancy/path_resampling.py index 277600e4420..d9237d5e563 100644 --- a/dimos/mapping/occupancy/path_resampling.py +++ b/dimos/mapping/occupancy/path_resampling.py @@ -13,6 +13,7 @@ # limitations under the License. +from copy import deepcopy import math from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion @@ -57,7 +58,7 @@ def _add_orientations_to_path(path: Path, goal_orientation: Quaternion) -> None: path.poses[i] = current_pose # Set last pose orientation - identity_quat = Quaternion(w=1) + identity_quat = Quaternion(w=1, x=0.0, y=0.0, z=0.0) last_pose = path.poses[-1] if goal_orientation != identity_quat: # Use the provided goal orientation if it's not the identity @@ -118,7 +119,9 @@ def simple_resample_path(path: Path, goal_pose: Pose, spacing: float) -> Path: new_y = prev.pose.position.y + dir_y * dist_along new_pose = PoseStamped( header=path.header, - pose=Pose(position=Point(x=new_x, y=new_y), orientation=prev.pose.orientation), + pose=Pose( + position=Point(x=new_x, y=new_y, z=0.0), orientation=prev.pose.orientation + ), ) resampled.append(new_pose) @@ -143,7 +146,7 @@ def simple_resample_path(path: Path, goal_pose: Pose, spacing: float) -> Path: else: resampled.append(last) - ret = Path(header=path.header, poses=resampled) + ret = Path(header=deepcopy(path.header), poses=deepcopy(resampled)) _add_orientations_to_path(ret, goal_pose.orientation) @@ -254,13 +257,13 @@ def smooth_resample_path( new_pose = PoseStamped( header=path.header, pose=Pose( - position=Point(x=float(sampled_x[i]), y=float(sampled_y[i])), - orientation=Quaternion(w=1), + position=Point(x=float(sampled_x[i]), y=float(sampled_y[i]), z=0.0), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), ), ) resampled.append(new_pose) - ret = Path(header=path.header, poses=resampled) + ret = Path(header=deepcopy(path.header), poses=deepcopy(resampled)) _add_orientations_to_path(ret, goal_pose.orientation) diff --git a/dimos/mapping/occupancy/test_generated_visualizations.py b/dimos/mapping/occupancy/test_generated_visualizations.py index a96e03144b8..b96a44bb826 100644 --- a/dimos/mapping/occupancy/test_generated_visualizations.py +++ b/dimos/mapping/occupancy/test_generated_visualizations.py @@ -17,6 +17,7 @@ import math +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid, Path @@ -35,10 +36,19 @@ def test_generated_image_preserves_header_and_colors(palette: Palette) -> None: grid = OccupancyGrid( header=Header(stamp=Time(sec=1, nanosec=999), frame_id="map"), - info=MapMetaData(width=3, height=1, resolution=1), - data=[-1, 0, 100], + info=MapMetaData( + width=3, + height=1, + resolution=1, + map_load_time=Time(sec=0, nanosec=0), + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([-1, 0, 100], dtype=np.int8), ) - image = Image.decode(visualize_occupancy_grid(grid, palette).encode()) + image = cdr_decode(cdr_encode(visualize_occupancy_grid(grid, palette)), Image) assert image.header == grid.header assert image.encoding == "bgr8" assert image_view(image).shape == (1, 3, 3) @@ -52,15 +62,30 @@ def test_generated_image_preserves_header_and_colors(palette: Palette) -> None: def test_path_rendering_is_relative_to_rotated_grid(footprint: bool) -> None: grid = OccupancyGrid( info=MapMetaData( - width=10, height=10, resolution=1, origin=Pose(orientation=Quaternion(w=1)) + width=10, + height=10, + resolution=1, + origin=Pose( + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + position=Point(x=0.0, y=0.0, z=0.0), + ), + map_load_time=Time(sec=0, nanosec=0), ), - data=[0] * 100, + data=np.asarray([0] * 100, dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) path = Path( poses=[ - PoseStamped(pose=Pose(position=Point(x=x, y=3), orientation=Quaternion(w=1))) + PoseStamped( + pose=Pose( + position=Point(x=x, y=3, z=0.0), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) for x in [2, 4, 6] - ] + ], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) expected = ( visualize_path(grid, path, 1, 2) @@ -68,13 +93,15 @@ def test_path_rendering_is_relative_to_rotated_grid(footprint: bool) -> None: else visualize_occupancy_grid(grid, "rainbow", path) ) grid.info.origin = Pose( - position=Point(x=20, y=10), orientation=quaternion_from_euler(0, 0, math.pi / 2) + position=Point(x=20, y=10, z=0.0), orientation=quaternion_from_euler(0, 0, math.pi / 2) ) path.poses = [ PoseStamped( pose=Pose( - position=Point(x=17, y=10 + x), orientation=quaternion_from_euler(0, 0, math.pi / 2) - ) + position=Point(x=17, y=10 + x, z=0.0), + orientation=quaternion_from_euler(0, 0, math.pi / 2), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) for x in [2, 4, 6] ] diff --git a/dimos/mapping/occupancy/test_path_mask.py b/dimos/mapping/occupancy/test_path_mask.py index 57681143baa..18000c982bf 100644 --- a/dimos/mapping/occupancy/test_path_mask.py +++ b/dimos/mapping/occupancy/test_path_mask.py @@ -13,15 +13,15 @@ # limitations under the License. +import cv2 +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion import numpy as np import pytest from dimos.mapping.occupancy.path_mask import make_path_mask from dimos.mapping.occupancy.path_resampling import smooth_resample_path from dimos.mapping.occupancy.visualizations import visualize_occupancy_grid -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_view from dimos.navigation.go2.replanning_a_star.min_cost_astar import min_cost_astar from dimos.utils.data import get_data @@ -34,15 +34,18 @@ ], ) def test_make_path_mask(occupancy_gradient, pose_index, max_length, expected_image) -> None: - start = Vector3(4.0, 2.0, 0) - goal_pose = Pose(6.15, 10.0, 0, 0, 0, 0, 1) - expected = Image.from_file(get_data(expected_image)) + start = Point(x=4, y=2, z=0.0) + goal_pose = Pose( + position=Point(x=6.15, y=10, z=0.0), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ) + expected = cv2.imread(str(get_data(expected_image)), cv2.IMREAD_COLOR) path = min_cost_astar(occupancy_gradient, goal_pose.position, start, use_cpp=False) path = smooth_resample_path(path, goal_pose, 0.1) robot_width = 0.4 path_mask = make_path_mask(occupancy_gradient, path, robot_width, pose_index, max_length) actual = visualize_occupancy_grid(occupancy_gradient, "rainbow") - actual.data[path_mask] = [0, 100, 0] + pixels = image_view(actual).copy() + pixels[path_mask] = [0, 100, 0] - np.testing.assert_array_equal(actual.data, expected.data) + np.testing.assert_array_equal(pixels, expected) diff --git a/dimos/mapping/occupancy/test_path_resampling.py b/dimos/mapping/occupancy/test_path_resampling.py index 52b01424b37..566824d4af0 100644 --- a/dimos/mapping/occupancy/test_path_resampling.py +++ b/dimos/mapping/occupancy/test_path_resampling.py @@ -12,30 +12,51 @@ # See the License for the specific language governing permissions and # limitations under the License. +import cv2 +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from dimos.mapping.occupancy.gradient import gradient from dimos.mapping.occupancy.path_resampling import simple_resample_path, smooth_resample_path from dimos.mapping.occupancy.visualize_path import visualize_path -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_view from dimos.navigation.go2.replanning_a_star.min_cost_astar import min_cost_astar from dimos.utils.data import get_data +def _grid(cells, resolution=0.05): + return OccupancyGrid( + info=MapMetaData( + width=cells.shape[1], + height=cells.shape[0], + resolution=float(np.float32(resolution)), + origin=Pose( + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + position=Point(x=0.0, y=0.0, z=0.0), + ), + map_load_time=Time(sec=0, nanosec=0), + ), + data=np.asarray(cells.astype(np.int8).ravel(), dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + + @pytest.fixture def costmap() -> OccupancyGrid: - return gradient(OccupancyGrid(np.load(get_data("occupancy_simple.npy"))), max_distance=1.5) + return gradient(_grid(np.load(get_data("occupancy_simple.npy"))), max_distance=1.5) @pytest.mark.parametrize("method", ["simple", "smooth"]) def test_resample_path(costmap, method) -> None: - start = Vector3(4.0, 2.0, 0) - goal_pose = Pose(6.15, 10.0, 0, 0, 0, 0, 1) - expected = Image.from_file(get_data(f"resample_path_{method}.png")) + start = Point(x=4, y=2, z=0.0) + goal_pose = Pose( + position=Point(x=6.15, y=10, z=0.0), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ) + expected = cv2.imread(str(get_data(f"resample_path_{method}.png")), cv2.IMREAD_COLOR) path = min_cost_astar(costmap, goal_pose.position, start, use_cpp=False) match method: @@ -47,4 +68,4 @@ def test_resample_path(costmap, method) -> None: raise ValueError(f"Unknown resampling method: {method}") actual = visualize_path(costmap, resampled, 0.2, 0.4) - np.testing.assert_array_equal(actual.data, expected.data) + np.testing.assert_array_equal(image_view(actual), expected) diff --git a/dimos/mapping/odometry_hist.py b/dimos/mapping/odometry_hist.py index eb2594007ae..c9d0f973af3 100644 --- a/dimos/mapping/odometry_hist.py +++ b/dimos/mapping/odometry_hist.py @@ -73,7 +73,7 @@ async def handle_odometry(self, msg: Odometry) -> None: ): header = deepcopy(msg.header) header.frame_id = frame_id - self._poses.append(PoseStamped(header=header, pose=msg.pose.pose)) + self._poses.append(PoseStamped(header=header, pose=deepcopy(msg.pose.pose))) self._unpublished = True if not self._unpublished: @@ -88,4 +88,4 @@ async def handle_odometry(self, msg: Odometry) -> None: self._unpublished = False header = deepcopy(msg.header) header.frame_id = frame_id - self.odom_hist.publish(Path(header=header, poses=list(self._poses))) + self.odom_hist.publish(Path(header=header, poses=deepcopy(list(self._poses)))) diff --git a/dimos/mapping/pointclouds/demo.py b/dimos/mapping/pointclouds/demo.py index 78daaf0db15..da369f178a3 100644 --- a/dimos/mapping/pointclouds/demo.py +++ b/dimos/mapping/pointclouds/demo.py @@ -17,6 +17,7 @@ from typing import TYPE_CHECKING +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.nav_msgs.msg import OccupancyGrid from dimos_generated.std_msgs.msg import Header import numpy as np @@ -49,7 +50,10 @@ def _get_occupancy_grid() -> OccupancyGrid: min_height = 0.15 max_height = 0.6 occupancygrid = simple_occupancy( - pointcloud_from_xyz(np.asarray(_get_sum_map().points), header=Header(frame_id="world")), + pointcloud_from_xyz( + np.asarray(_get_sum_map().points), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + ), resolution=resolution, min_height=min_height, max_height=max_height, diff --git a/dimos/mapping/pointclouds/occupancy.py b/dimos/mapping/pointclouds/occupancy.py index 21a1c8f75c6..bcd3a7318e8 100644 --- a/dimos/mapping/pointclouds/occupancy.py +++ b/dimos/mapping/pointclouds/occupancy.py @@ -14,21 +14,24 @@ from __future__ import annotations +from copy import deepcopy from dataclasses import dataclass +import math from typing import TYPE_CHECKING, Any, Protocol, TypeVar +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid +from dimos_generated.sensor_msgs.msg import PointCloud2 import numpy as np from numpy.typing import NDArray from scipy import ndimage -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid +from dimos.msgs.pointcloud import pointcloud_xyz if TYPE_CHECKING: from collections.abc import Callable - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 - @dataclass(frozen=True) class OccupancyConfig: @@ -37,6 +40,37 @@ class OccupancyConfig: resolution: float = 0.05 frame_id: str | None = None + def __post_init__(self) -> None: + if not math.isfinite(self.resolution) or self.resolution <= 0: + raise ValueError("occupancy resolution must be finite and positive") + + +def _grid_message( + cloud: PointCloud2, + config: OccupancyConfig, + cells: NDArray[np.int8], + origin: Pose | None = None, +) -> OccupancyGrid: + header = deepcopy(cloud.header) + if config.frame_id is not None: + header.frame_id = config.frame_id + return OccupancyGrid( + header=header, + info=MapMetaData( + resolution=config.resolution, + width=cells.shape[1], + height=cells.shape[0], + origin=origin + if origin is not None + else Pose( + orientation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), + position=Point(x=0.0, y=0.0, z=0.0), + ), + map_load_time=Time(sec=0, nanosec=0), + ), + data=np.asarray(cells.ravel(), dtype=np.int8), + ) + ConfigT = TypeVar("ConfigT", bound=OccupancyConfig, covariant=True) @@ -84,17 +118,11 @@ def height_cost_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: OccupancyGrid with costs 0-100 based on terrain slope, -1 for unknown """ cfg = HeightCostConfig(**kwargs) - points, _ = cloud.as_numpy() - points = points.astype(np.float64) # Upcast to avoid float32 rounding - ts = cloud.ts if hasattr(cloud, "ts") and cloud.ts is not None else 0.0 + points = pointcloud_xyz(cloud) + points = points[np.isfinite(points).all(axis=1)] if len(points) == 0: - return OccupancyGrid( - width=1, - height=1, - resolution=cfg.resolution, - frame_id=cfg.frame_id or cloud.frame_id, - ) + return _grid_message(cloud, cfg, np.full((1, 1), -1, dtype=np.int8)) # Find bounds of the point cloud in X-Y plane (use all points) min_x = np.min(points[:, 0]) @@ -114,7 +142,9 @@ def height_cost_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: height = int(np.ceil((max_y - min_y) / cfg.resolution)) # Create origin pose - origin = Pose() + origin = Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) origin.position.x = min_x origin.position.y = min_y origin.position.z = 0.0 @@ -205,13 +235,7 @@ def height_cost_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: else: cost = np.full((height, width), -1, dtype=np.int8) - return OccupancyGrid( - grid=cost, - resolution=cfg.resolution, - origin=origin, - frame_id=cfg.frame_id or cloud.frame_id, - ts=ts, - ) + return _grid_message(cloud, cfg, cost, origin) @dataclass(frozen=True) @@ -236,16 +260,11 @@ def general_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: OccupancyGrid with occupied cells where points were projected """ cfg = GeneralOccupancyConfig(**kwargs) - points, _ = cloud.as_numpy() - points = points.astype(np.float64) # Upcast to avoid float32 rounding + points = pointcloud_xyz(cloud) + points = points[np.isfinite(points).all(axis=1)] if len(points) == 0: - return OccupancyGrid( - width=1, - height=1, - resolution=cfg.resolution, - frame_id=cfg.frame_id or cloud.frame_id, - ) + return _grid_message(cloud, cfg, np.full((1, 1), -1, dtype=np.int8)) # Filter points by height for obstacles obstacle_mask = (points[:, 2] >= cfg.min_height) & (points[:, 2] <= cfg.max_height) @@ -273,7 +292,9 @@ def general_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: height = int(np.ceil((max_y - min_y) / cfg.resolution)) # Create origin pose (bottom-left corner of the grid) - origin = Pose() + origin = Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) origin.position.x = min_x origin.position.y = min_y origin.position.z = 0.0 @@ -329,15 +350,8 @@ def general_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: # Create and return OccupancyGrid # Get timestamp from cloud if available - ts = cloud.ts if hasattr(cloud, "ts") and cloud.ts is not None else 0.0 - return OccupancyGrid( - grid=grid, - resolution=cfg.resolution, - origin=origin, - frame_id=cfg.frame_id or cloud.frame_id, - ts=ts, - ) + return _grid_message(cloud, cfg, grid, origin) @dataclass(frozen=True) @@ -366,16 +380,11 @@ def simple_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: OccupancyGrid with occupied/free cells """ cfg = SimpleOccupancyConfig(**kwargs) - points, _ = cloud.as_numpy() - points = points.astype(np.float64) # Upcast to avoid float32 rounding + points = pointcloud_xyz(cloud) + points = points[np.isfinite(points).all(axis=1)] if len(points) == 0: - return OccupancyGrid( - width=1, - height=1, - resolution=cfg.resolution, - frame_id=cfg.frame_id or cloud.frame_id, - ) + return _grid_message(cloud, cfg, np.full((1, 1), -1, dtype=np.int8)) # Find bounds of the point cloud in X-Y plane min_x = float(np.min(points[:, 0])) - 1.0 @@ -388,7 +397,9 @@ def simple_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: height = int(np.ceil((max_y - min_y) / cfg.resolution)) # Create origin pose (bottom-left corner of the grid) - origin = Pose() + origin = Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) origin.position.x = min_x origin.position.y = min_y origin.position.z = 0.0 @@ -412,15 +423,7 @@ def simple_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: cfg.max_height, ) - ts = cloud.ts if hasattr(cloud, "ts") and cloud.ts is not None else 0.0 - - return OccupancyGrid( - grid=grid, - resolution=cfg.resolution, - origin=origin, - frame_id=cfg.frame_id or cloud.frame_id, - ts=ts, - ) + return _grid_message(cloud, cfg, grid, origin) # Populate algorithm registry diff --git a/dimos/mapping/pointclouds/test_generated_occupancy.py b/dimos/mapping/pointclouds/test_generated_occupancy.py index 4920e8b7db5..2069f38e375 100644 --- a/dimos/mapping/pointclouds/test_generated_occupancy.py +++ b/dimos/mapping/pointclouds/test_generated_occupancy.py @@ -17,6 +17,7 @@ from collections.abc import Callable +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.nav_msgs.msg import OccupancyGrid from dimos_generated.sensor_msgs.msg import PointCloud2 @@ -42,9 +43,9 @@ def test_generated_cloud_to_grid(algorithm: Callable[..., OccupancyGrid]) -> Non x, y = np.meshgrid(xy, xy) points = np.column_stack((x.ravel(), y.ravel(), np.zeros(x.size))) cloud = pointcloud_from_xyz(points, header=header) - decoded = PointCloud2.decode(cloud.encode()) + decoded = cdr_decode(cdr_encode(cloud), PointCloud2) grid = algorithm(decoded, resolution=0.1) - restored = OccupancyGrid.decode(grid.encode()) + restored = cdr_decode(cdr_encode(grid), OccupancyGrid) assert restored.header == header cells = occupancy_view(restored) assert cells.shape == (39, 39) @@ -63,7 +64,9 @@ def test_empty_or_missing_points_remain_unknown( algorithm: Callable[..., OccupancyGrid], points: NDArray[np.float64], ) -> None: - cloud = pointcloud_from_xyz(points, header=Header(stamp=Time(nanosec=123), frame_id="map")) + cloud = pointcloud_from_xyz( + points, header=Header(stamp=Time(nanosec=123, sec=0), frame_id="map") + ) grid = algorithm(cloud) assert grid.header == cloud.header assert occupancy_view(grid).tolist() == [[-1]] @@ -72,14 +75,16 @@ def test_empty_or_missing_points_remain_unknown( @pytest.mark.parametrize("resolution", [0, -1, float("nan"), float("inf")]) def test_invalid_resolution(resolution: float) -> None: - cloud = pointcloud_from_xyz(np.zeros((1, 3)), header=Header()) + cloud = pointcloud_from_xyz( + np.zeros((1, 3)), header=Header(stamp=Time(sec=0, nanosec=0), frame_id="") + ) with pytest.raises(ValueError, match="resolution"): general_occupancy(cloud, resolution=resolution) def test_inflation_preserves_metadata_and_source() -> None: cloud = pointcloud_from_xyz( - np.array([[0, 0, 1.0]]), header=Header(stamp=Time(nanosec=789), frame_id="map") + np.array([[0, 0, 1.0]]), header=Header(stamp=Time(nanosec=789, sec=0), frame_id="map") ) grid = general_occupancy(cloud, resolution=0.1) original = occupancy_view(grid).copy() @@ -90,7 +95,7 @@ def test_inflation_preserves_metadata_and_source() -> None: np.testing.assert_array_equal(occupancy_view(grid), original) inflated.header.frame_id = "changed" assert grid.header.frame_id == "map" - assert simple_inflate(grid, 0) == grid + assert cdr_encode(simple_inflate(grid, 0)) == cdr_encode(grid) for radius in [-1, float("nan"), float("inf")]: with pytest.raises(ValueError, match="radius"): simple_inflate(grid, radius) diff --git a/dimos/mapping/pointclouds/test_occupancy.py b/dimos/mapping/pointclouds/test_occupancy.py index 4f067da2e11..0aab41ba518 100644 --- a/dimos/mapping/pointclouds/test_occupancy.py +++ b/dimos/mapping/pointclouds/test_occupancy.py @@ -16,6 +16,7 @@ from pathlib import Path import cv2 +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.nav_msgs.msg import OccupancyGrid from dimos_generated.std_msgs.msg import Header import numpy as np @@ -58,7 +59,9 @@ def big_office() -> PointCloud: ) def test_occupancy(apartment: PointCloud, occupancy_fn, output_name: str) -> None: expected_image = cv2.imread(str(get_data(output_name)), cv2.IMREAD_GRAYSCALE) - cloud = pointcloud_from_xyz(np.asarray(apartment.points), header=Header(frame_id="map")) + cloud = pointcloud_from_xyz( + np.asarray(apartment.points), header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)) + ) occupancy_grid = occupancy_fn(cloud) @@ -77,7 +80,9 @@ def test_occupancy(apartment: PointCloud, occupancy_fn, output_name: str) -> Non ) def test_occupancy2(big_office, occupancy_fn, output_name): expected_image = image_from_file(get_data(output_name)) - cloud = pointcloud_from_xyz(np.asarray(big_office.points), header=Header()) + cloud = pointcloud_from_xyz( + np.asarray(big_office.points), header=Header(stamp=Time(sec=0, nanosec=0), frame_id="") + ) occupancy_grid = occupancy_fn(cloud) diff --git a/dimos/mapping/ray_tracing/module.py b/dimos/mapping/ray_tracing/module.py index 09185514c4d..238218b96c8 100644 --- a/dimos/mapping/ray_tracing/module.py +++ b/dimos/mapping/ray_tracing/module.py @@ -17,11 +17,12 @@ from typing import TYPE_CHECKING +from dimos_generated.dimos_msgs.msg import RegionBounds, RegionPointCloud2 +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.tf2_msgs.msg import TFMessage + from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.spec import mapping # Max stamp gap between a cloud and the transform used to register it (s). @@ -100,14 +101,14 @@ class RayTracingVoxelMap(NativeModule, mapping.GlobalPointcloud): global_map: Out[PointCloud2] local_map: Out[PointCloud2] local_map_fine: Out[PointCloud2] - region_bounds: Out[PoseStamped] + region_bounds: Out[RegionBounds] # One region of a seeded map as it lands, gated like local_map. seed_bounds # carries its cylinder like region_bounds. - seed_map: Out[PointCloud2] - seed_bounds: Out[PoseStamped] + seed_map: Out[RegionPointCloud2] + seed_bounds: Out[RegionBounds] # The map for viewers, one region cell per message with the cell packed in # the header seq. - map_regions: Out[PointCloud2] + map_regions: Out[RegionPointCloud2] if TYPE_CHECKING: diff --git a/dimos/mapping/ray_tracing/rust/Cargo.toml b/dimos/mapping/ray_tracing/rust/Cargo.toml index be4d6325fea..9f46c73a8cf 100644 --- a/dimos/mapping/ray_tracing/rust/Cargo.toml +++ b/dimos/mapping/ray_tracing/rust/Cargo.toml @@ -21,22 +21,22 @@ python = "dimos.mapping.ray_tracing.module:RayTracingVoxelMap" threads = 2 [package.metadata.dimos.module.ray_tracing.inputs] -lidar = "sensor_msgs.PointCloud2" -loaded_map = "sensor_msgs.PointCloud2" -tf = "tf2_msgs.TFMessage" +lidar = "sensor_msgs/msg/PointCloud2" +loaded_map = "sensor_msgs/msg/PointCloud2" +tf = "tf2_msgs/msg/TFMessage" [package.metadata.dimos.module.ray_tracing.outputs] -global_map = "sensor_msgs.PointCloud2" -local_map = "sensor_msgs.PointCloud2" -local_map_fine = "sensor_msgs.PointCloud2" -seed_map = "sensor_msgs.PointCloud2" -seed_bounds = "geometry_msgs.PoseStamped" -region_bounds = "geometry_msgs.PoseStamped" -map_regions = "sensor_msgs.PointCloud2" +global_map = "sensor_msgs/msg/PointCloud2" +local_map = "sensor_msgs/msg/PointCloud2" +local_map_fine = "sensor_msgs/msg/PointCloud2" +seed_map = "dimos_msgs/msg/RegionPointCloud2" +seed_bounds = "dimos_msgs/msg/RegionBounds" +region_bounds = "dimos_msgs/msg/RegionBounds" +map_regions = "dimos_msgs/msg/RegionPointCloud2" [dependencies] dimos-module = { path = "../../../../native/rust/dimos-module" } -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +dimos-generated-messages = { path = "../../../../packages/dimos-generated/build/message-source/rust" } tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal"] } serde = { version = "1", features = ["derive"] } ahash = "0.8" diff --git a/dimos/mapping/ray_tracing/rust/src/module.rs b/dimos/mapping/ray_tracing/rust/src/module.rs index 9daf0b2f5f8..9865f767613 100644 --- a/dimos/mapping/ray_tracing/rust/src/module.rs +++ b/dimos/mapping/ray_tracing/rust/src/module.rs @@ -12,6 +12,16 @@ // See the License for the specific language governing permissions and // limitations under the License. +use dimos_generated_messages::builtin_interfaces::msg::time::Time; +use dimos_generated_messages::dimos_msgs::msg::{ + region_bounds::RegionBounds, region_point_cloud2::RegionPointCloud2, +}; +use dimos_generated_messages::geometry_msgs::msg::point::Point; +use dimos_generated_messages::sensor_msgs::msg::{ + point_cloud2::PointCloud2, point_field::PointField, +}; +use dimos_generated_messages::std_msgs::msg::header::Header; +use dimos_module::cdr; use std::collections::BTreeMap; use std::time::{Duration, Instant}; @@ -24,9 +34,6 @@ use dimos_module::pointcloud::extract_xyz; use dimos_module::{ debug_throttled, error_throttled, warn_throttled, Input, Module, Output, Tf, Transform, }; -use lcm_msgs::geometry_msgs::{Point, Pose as PoseMsg, PoseStamped, Quaternion}; -use lcm_msgs::sensor_msgs::{PointCloud2, PointField}; -use lcm_msgs::std_msgs::{Header, Time}; use tokio::sync::mpsc; use tokio::sync::mpsc::error::TryRecvError; use tokio::task::JoinHandle; @@ -51,42 +58,42 @@ const SEED_PASS_BUDGET: Duration = Duration::from_millis(20); #[derive(Module)] #[module(name = "ray_tracing", setup = spawn_worker, teardown = stop_worker)] pub struct RayTracingVoxelMap { - #[input(decode = PointCloud2::decode, handler = on_lidar)] + #[input(decode = cdr::decode, handler = on_lidar)] lidar: Input, // An externally loaded map cloud. Only the first one seeds the map. - #[input(decode = PointCloud2::decode, handler = on_loaded_map)] + #[input(decode = cdr::decode, handler = on_loaded_map)] loaded_map: Input, #[tf] tf: Tf, - #[output(encode = PointCloud2::encode)] + #[output(encode = cdr::encode)] global_map: Output, - #[output(encode = PointCloud2::encode)] + #[output(encode = cdr::encode)] local_map: Output, - #[output(encode = PointCloud2::encode)] + #[output(encode = cdr::encode)] local_map_fine: Output, // Cylinder bounds of the local map. Position is the center, orientation holds // radius, z_min, z_max. Stamped like local_map so consumers pair them. - #[output(encode = PoseStamped::encode)] - region_bounds: Output, + #[output(encode = cdr::encode)] + region_bounds: Output, // The map for viewers, one region cell per message with the cell packed in - // the header seq. - #[output(encode = PointCloud2::encode)] - map_regions: Output, + // the region ID. + #[output(encode = cdr::encode)] + map_regions: Output, // One region of a seeded map as it lands, gated like local_map. seed_bounds // carries its cylinder like region_bounds. - #[output(encode = PointCloud2::encode)] - seed_map: Output, + #[output(encode = cdr::encode)] + seed_map: Output, - #[output(encode = PoseStamped::encode)] - seed_bounds: Output, + #[output(encode = cdr::encode)] + seed_bounds: Output, #[config] config: Config, @@ -265,17 +272,17 @@ struct Worker { global_map: Output, local_map: Output, local_map_fine: Output, - region_bounds: Output, - map_regions: Output, - seed_map: Output, - seed_bounds: Output, + region_bounds: Output, + map_regions: Output, + seed_map: Output, + seed_bounds: Output, } impl Worker { async fn run(mut self) { let mut state = State { mapper: Mapper::new(self.config.clone()), - last_frame_stamp: Time::default(), + last_frame_stamp: Time { sec: 0, nanosec: 0 }, seed: SeedState::Idle, viz: RegionSweep::default(), }; @@ -400,7 +407,7 @@ impl Worker { // Bounds pair with local_map by stamp, so publish them on its cadence. if let Some(c) = region { - let bounds_msg = bounds_to_pose(&c, out_frame_id, stamp.clone()); + let bounds_msg = bounds_message(&c, out_frame_id, stamp.clone()); publish_bounds(&self.region_bounds, &bounds_msg).await; } @@ -429,7 +436,7 @@ impl Worker { } /// Publish the map region cells due this viz tick, the cell in each - /// header seq. + /// region ID. async fn publish_map_regions(&self, state: &mut State, stamp: &Time) { let due = tokio::task::block_in_place(|| { map_regions_due( @@ -441,8 +448,10 @@ impl Worker { ) }); for (cell, points) in due { - let mut cloud = points_to_cloud(&points, &self.config.world_frame, stamp.clone()); - cloud.header.seq = pack_cell(cell); + let cloud = RegionPointCloud2 { + region_id: pack_cell(cell), + cloud: points_to_cloud(&points, &self.config.world_frame, stamp.clone()), + }; publish_cloud(&self.map_regions, &cloud).await; } } @@ -508,7 +517,7 @@ impl Worker { "Premap load in progress." ); if let Some(cylinder) = completed { - let seq = load.next_region as i32; + let seq = i32::try_from(load.next_region).expect("seed region ID exceeds int32"); self.publish_seed_region(&state.mapper, &cylinder, seq, &state.last_frame_stamp) .await; } @@ -536,7 +545,7 @@ impl Worker { } /// Publish one seeded region as the map now holds it. Cloud and bounds - /// carry the region number in their header seq, which consumers pair on. + /// carry the region number in their region ID, which consumers pair on. async fn publish_seed_region( &self, mapper: &Mapper, @@ -553,18 +562,20 @@ impl Worker { } /// One seeded region as its bounds and cloud messages, both stamped alike -/// with the region number in their header seq. +/// with the region number in their region ID. fn seed_region_messages( cylinder: &Cylinder, points: &[f32], seq: i32, frame_id: &str, stamp: &Time, -) -> (PoseStamped, PointCloud2) { - let mut bounds = bounds_to_pose(cylinder, frame_id, stamp.clone()); - bounds.header.seq = seq; - let mut cloud = points_to_cloud(points, frame_id, stamp.clone()); - cloud.header.seq = seq; +) -> (RegionBounds, RegionPointCloud2) { + let mut bounds = bounds_message(cylinder, frame_id, stamp.clone()); + bounds.region_id = seq; + let cloud = RegionPointCloud2 { + region_id: seq, + cloud: points_to_cloud(points, frame_id, stamp.clone()), + }; (bounds, cloud) } @@ -628,35 +639,29 @@ fn prepare_seed( const LOADED_MAP_TF_WAIT_TIMEOUT: Duration = Duration::from_secs(2); fn time_secs(t: &Time) -> f64 { - t.sec as f64 + t.nsec as f64 * 1e-9 + t.sec as f64 + t.nanosec as f64 * 1e-9 } -/// A cylinder as the bounds message: position is the center, orientation -/// holds radius, z_min, z_max. -fn bounds_to_pose(c: &Cylinder, frame_id: &str, stamp: Time) -> PoseStamped { - PoseStamped { +/// Explicit cylinder bounds in the source frame; vertical limits are absolute. +fn bounds_message(c: &Cylinder, frame_id: &str, stamp: Time) -> RegionBounds { + RegionBounds { header: Header { - seq: 0, stamp, - frame_id: frame_id.to_string(), + frame_id: frame_id.into(), }, - pose: PoseMsg { - position: Point { - x: c.cx as f64, - y: c.cy as f64, - z: 0.0, - }, - orientation: Quaternion { - x: c.radius as f64, - y: c.z_min as f64, - z: c.z_max as f64, - w: 0.0, - }, + region_id: 0, + center: Point { + x: c.cx as f64, + y: c.cy as f64, + z: 0.0, }, + radius: c.radius as f64, + z_min: c.z_min as f64, + z_max: c.z_max as f64, } } -async fn publish_bounds(out: &Output, msg: &PoseStamped) { +async fn publish_bounds(out: &Output, msg: &RegionBounds) { if let Err(e) = out.publish(msg).await { error_throttled!( Duration::from_secs(1), @@ -693,18 +698,17 @@ fn write_point(data: &mut Vec, n: &mut i32, x: f32, y: f32, z: f32) { fn make_cloud(data: Vec, n: i32, frame_id: &str, stamp: Time) -> PointCloud2 { let make_field = |name: &str, off: i32| PointField { name: name.into(), - offset: off, - datatype: PointField::FLOAT32 as u8, + offset: off as u32, + datatype: PointField::FLOAT32, count: 1, }; PointCloud2 { header: Header { - seq: 0, stamp, frame_id: frame_id.into(), }, height: 1, - width: n, + width: n as u32, fields: vec![ make_field("x", 0), make_field("y", 4), @@ -713,7 +717,7 @@ fn make_cloud(data: Vec, n: i32, frame_id: &str, stamp: Time) -> PointCloud2 ], is_bigendian: false, point_step: 16, - row_step: 16 * n, + row_step: 16 * n as u32, data, is_dense: true, } @@ -729,7 +733,7 @@ fn points_to_cloud(points: &[f32], frame_id: &str, stamp: Time) -> PointCloud2 { make_cloud(data, n, frame_id, stamp) } -async fn publish_cloud(out: &Output, cloud: &PointCloud2) { +async fn publish_cloud(out: &Output, cloud: &T) { if let Err(e) = out.publish(cloud).await { error_throttled!( Duration::from_secs(1), @@ -818,7 +822,7 @@ mod tests { UnitQuaternion::from_axis_angle(&Vector3::z_axis(), std::f64::consts::FRAC_PI_2), ); let t = Transform::new("odom", "map", 0.0, iso); - let cloud = points_to_cloud(&[3.5, 0.0, 0.5], "map", Time::default()); + let cloud = points_to_cloud(&[3.5, 0.0, 0.5], "map", Time { sec: 0, nanosec: 0 }); let part = prepare_seed(&cloud, tf_to_pose(&t), 1.0, (0.0, 0.0, 0.0), 4.0) .expect("loaded map must decode"); @@ -829,7 +833,7 @@ mod tests { assert_eq!(created, 1); let region = mapper.local_points(&part.regions[0].cylinder.bounds()); - let seeded = points_to_cloud(®ion, "odom", Time::default()); + let seeded = points_to_cloud(®ion, "odom", Time { sec: 0, nanosec: 0 }); assert!(cloud_points(&seeded).contains(&voxel_center(10, 3, 0))); } @@ -898,18 +902,18 @@ mod tests { } #[test] - fn seed_region_messages_pair_on_seq_and_stamp() { + fn seed_region_messages_pair_on_region_id_and_stamp() { let part = two_region_seed(); - let stamp = Time { sec: 7, nsec: 9 }; + let stamp = Time { sec: 7, nanosec: 9 }; for (index, region) in part.regions.iter().enumerate() { let seq = index as i32 + 1; let (bounds, cloud) = seed_region_messages(®ion.cylinder, &[1.5, 1.5, 0.5], seq, "odom", &stamp); - assert_eq!((bounds.header.seq, cloud.header.seq), (seq, seq)); + assert_eq!((bounds.region_id, cloud.region_id), (seq, seq)); assert_eq!(bounds.header.stamp, stamp); - assert_eq!(cloud.header.stamp, stamp); - assert_eq!(bounds.pose.position.x, region.cylinder.cx as f64); - assert_eq!(cloud.width, 1); + assert_eq!(cloud.cloud.header.stamp, stamp); + assert_eq!(bounds.center.x, region.cylinder.cx as f64); + assert_eq!(cloud.cloud.width, 1); } } @@ -969,8 +973,9 @@ mod tests { let occupied: Vec = map.voxels.keys().collect(); assert_eq!(occupied, vec![(3, -2, 1)]); - // Encode and decode the occupied voxel center exactly as the port would. - let cloud = points_to_cloud(&[3.5, -1.5, 1.5], "world", Time::default()); + // A mask cloud naming that voxel's center, encoded and decoded exactly + // as the port would. + let cloud = points_to_cloud(&[3.5, -1.5, 1.5], "world", Time { sec: 0, nanosec: 0 }); let Ok(points) = extract_xyz(&cloud) else { panic!("occupied cloud must decode"); }; @@ -997,12 +1002,12 @@ mod tests { let global = points_to_cloud( &emit_points(&map, 1.0, None, &live), "world", - Time::default(), + Time { sec: 0, nanosec: 0 }, ); let local = points_to_cloud( &emit_points(&map, 1.0, Some(&cylinder), &live), "world", - Time::default(), + Time { sec: 0, nanosec: 0 }, ); assert!(cloud_points(&global).contains(&voxel_center(0, 0, 0))); assert!(cloud_points(&local).contains(&voxel_center(0, 0, 0))); @@ -1022,12 +1027,12 @@ mod tests { let global = points_to_cloud( &emit_points(&map, 1.0, None, &live), "world", - Time::default(), + Time { sec: 0, nanosec: 0 }, ); let local = points_to_cloud( &emit_points(&map, 1.0, Some(&cylinder), &live), "world", - Time::default(), + Time { sec: 0, nanosec: 0 }, ); assert!(cloud_points(&global).contains(&voxel_center(5, 0, 0))); assert!(!cloud_points(&local).contains(&voxel_center(5, 0, 0))); @@ -1048,12 +1053,12 @@ mod tests { let global = points_to_cloud( &emit_points(&map, 1.0, None, &live), "world", - Time::default(), + Time { sec: 0, nanosec: 0 }, ); let local = points_to_cloud( &emit_points(&map, 1.0, Some(&cylinder), &live), "world", - Time::default(), + Time { sec: 0, nanosec: 0 }, ); assert!(cloud_points(&global).contains(&voxel_center(0, 0, 5))); assert!(!cloud_points(&local).contains(&voxel_center(0, 0, 5))); @@ -1076,12 +1081,12 @@ mod tests { let global = points_to_cloud( &emit_points(&map, 1.0, None, &live), "world", - Time::default(), + Time { sec: 0, nanosec: 0 }, ); let local = points_to_cloud( &emit_points(&map, 1.0, Some(&cylinder), &live), "world", - Time::default(), + Time { sec: 0, nanosec: 0 }, ); assert!(cloud_points(&global).contains(&voxel_center(1, 0, 0))); assert!(cloud_points(&global).contains(&voxel_center(10, 10, 10))); @@ -1114,7 +1119,7 @@ mod tests { let local = points_to_cloud( &emit_points(&map, 1.0, Some(&cylinder), &live), "world", - Time::default(), + Time { sec: 0, nanosec: 0 }, ); let pts = cloud_points(&local); assert!(pts.contains(&voxel_center(1, 1, 0)), "dense patch kept"); diff --git a/dimos/mapping/ray_tracing/test_transformer.py b/dimos/mapping/ray_tracing/test_transformer.py index 378c1058fc5..61641332a9e 100644 --- a/dimos/mapping/ray_tracing/test_transformer.py +++ b/dimos/mapping/ray_tracing/test_transformer.py @@ -16,22 +16,25 @@ from pathlib import Path +from dimos_generated.builtin_interfaces.msg import Time import numpy as np from numpy.typing import NDArray import pytest pytest.importorskip("dimos_voxel_ray_tracing") +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage + from dimos.mapping.ray_tracing.module import TF_MATCH_TOLERANCE_S from dimos.mapping.ray_tracing.transformer import RayTraceMap, pose_from_tf from dimos.mapping.ray_tracing.voxel_map import VoxelRayMapper from dimos.memory.store.sqlite import SqliteStore from dimos.memory.tf import StreamTF from dimos.memory.type.observation import Observation -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.pointcloud import pointcloud_from_xyz def _default_margin() -> float: @@ -46,7 +49,9 @@ def _obs( id=0, ts=ts, pose=pose, - _data=PointCloud2.from_numpy(points), + _data=pointcloud_from_xyz( + points, header=Header(frame_id="lidar", stamp=Time(sec=0, nanosec=0)) + ), ) @@ -66,7 +71,14 @@ def test_emit_every_n_yields_on_cadence_and_flushes_remainder() -> None: def test_poseless_obs_are_skipped() -> None: points = _cube() - poseless = Observation(id=1, ts=0.0, pose=None, _data=PointCloud2.from_numpy(points)) + poseless = Observation( + id=1, + ts=0.0, + pose=None, + _data=pointcloud_from_xyz( + points, header=Header(frame_id="lidar", stamp=Time(sec=0, nanosec=0)) + ), + ) posed = _obs(points, ts=1.0, pose=(0.0, 0.0, 0.0)) results = list(RayTraceMap()(iter([poseless, posed]))) @@ -119,7 +131,9 @@ def test_registers_sensor_frame_cloud_by_pose() -> None: id=0, ts=1.0, pose=(5.0, 0.0, 2.0, 0.0, s, 0.0, s), - _data=PointCloud2.from_numpy(point), + _data=pointcloud_from_xyz( + point, header=Header(frame_id="lidar", stamp=Time(sec=0, nanosec=0)) + ), ) [emitted] = list(rtm(iter([obs]))) @@ -147,15 +161,19 @@ def test_reused_instance_continues_the_map() -> None: def test_pose_from_tf_attaches_the_pose_at_the_stamp(tmp_path: Path) -> None: ts = 10.0 with SqliteStore(path=str(tmp_path / "tf.db")) as store: - edge = Transform( - frame_id="world", + edge = TransformStamped( + header=Header(frame_id="world", stamp=Time(sec=10, nanosec=0)), child_frame_id="lidar", - translation=Vector3(1.0, 2.0, 3.0), - ts=ts, + transform=Transform( + translation=Vector3(x=1.0, y=2.0, z=3.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) - store.stream("tf", TFMessage).append(TFMessage(edge), ts=ts, pose=None) + store.stream("tf", TFMessage).append(TFMessage(transforms=[edge]), ts=ts, pose=None) tf = StreamTF(store.stream("tf", TFMessage)) - cloud = PointCloud2.from_numpy(_cube(), frame_id="lidar") + cloud = pointcloud_from_xyz( + _cube(), header=Header(frame_id="lidar", stamp=Time(sec=10, nanosec=0)) + ) matched = Observation(id=0, ts=ts, _data=cloud) far = Observation(id=1, ts=ts + 10 * TF_MATCH_TOLERANCE_S, _data=cloud) diff --git a/dimos/mapping/ray_tracing/test_viz.py b/dimos/mapping/ray_tracing/test_viz.py index 6c8d0e0eaa3..a4b506dfd21 100644 --- a/dimos/mapping/ray_tracing/test_viz.py +++ b/dimos/mapping/ray_tracing/test_viz.py @@ -12,10 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.dimos_msgs.msg import RegionPointCloud2 import numpy as np from dimos.mapping.ray_tracing.viz import render_map_region, voxel_map_points -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import header_now HEIGHT_RANGE = (-1.0, 3.0) @@ -30,15 +32,19 @@ def test_the_height_ramp_spans_the_range_and_clamps_outside_it() -> None: def test_map_regions_land_static_on_their_own_cell_entities() -> None: - region = PointCloud2.from_numpy(_at_heights([-1.0, 3.0])) - region.seq = (1 << 16) | (-2 & 0xFFFF) + region = RegionPointCloud2( + region_id=(1 << 16) | (-2 & 0xFFFF), + cloud=pointcloud_from_xyz(_at_heights([-1.0, 3.0]), header=header_now()), + ) (path, arch, static) = render_map_region(region, 0.1, HEIGHT_RANGE)[0] assert path == "world/map_regions/1_-2" and static assert len(arch.positions.as_arrow_array()) == 2 assert arch.class_ids.as_arrow_array().to_pylist() == [0, 255] - emptied = PointCloud2.from_numpy(np.zeros((0, 3), dtype=np.float32)) - emptied.seq = region.seq + emptied = RegionPointCloud2( + region_id=region.region_id, + cloud=pointcloud_from_xyz(np.zeros((0, 3), dtype=np.float32), header=header_now()), + ) (path, arch, static) = render_map_region(emptied, 0.1, HEIGHT_RANGE)[0] assert path == "world/map_regions/1_-2" and static assert len(arch.positions.as_arrow_array()) == 0 diff --git a/dimos/mapping/ray_tracing/transformer.py b/dimos/mapping/ray_tracing/transformer.py index 96bd197fd64..162bef60610 100644 --- a/dimos/mapping/ray_tracing/transformer.py +++ b/dimos/mapping/ray_tracing/transformer.py @@ -16,10 +16,14 @@ from typing import TYPE_CHECKING, Any +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header +import numpy as np + from dimos.mapping.ray_tracing.module import TF_MATCH_TOLERANCE_S from dimos.mapping.ray_tracing.voxel_map import VoxelRayMapper from dimos.memory.transform import FnTransformer, Transformer -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_xyz from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -36,7 +40,10 @@ def pose_from_tf(tf: StreamTF, world_frame: str) -> FnTransformer[PointCloud2, P def attach(obs: Observation[PointCloud2]) -> Observation[PointCloud2]: t = tf.get( - world_frame, obs.data.frame_id, time_point=obs.ts, time_tolerance=TF_MATCH_TOLERANCE_S + world_frame, + obs.data.header.frame_id, + time_point=obs.ts, + time_tolerance=TF_MATCH_TOLERANCE_S, ) return obs.with_pose(t) @@ -76,16 +83,13 @@ def _make_obs( last_obs: Observation[PointCloud2], count: int, ) -> Observation[PointCloud2]: - import open3d as o3d # type: ignore[import-untyped] - import open3d.core as o3c # type: ignore[import-untyped] - tags = {**last_obs.tags, "frame_count": count} cx, cy, radius, z_min, z_max = mapper.take_local_bounds() positions = mapper.local_map((cx, cy, 0.0), radius, z_min, z_max) tags["region_bounds"] = (cx, cy, radius, z_min, z_max) - pcd = o3d.t.geometry.PointCloud() - pcd.point["positions"] = o3c.Tensor.from_numpy(positions) - cloud = PointCloud2(pointcloud=pcd, frame_id="world", ts=last_obs.ts) + cloud = pointcloud_from_xyz( + positions, header=Header(frame_id="world", stamp=last_obs.data.header.stamp) + ) return last_obs.derive(data=cloud, tags=tags) def __call__( @@ -101,7 +105,9 @@ def __call__( logger.warning("RayTraceMap: obs %s has no pose. dropped a cloud", obs.id) continue x, y, z, qx, qy, qz, qw = obs.pose_tuple - mapper.add_frame(obs.data.points_f32(), (x, y, z), (qx, qy, qz, qw)) + mapper.add_frame( + pointcloud_xyz(obs.data).astype(np.float32), (x, y, z), (qx, qy, qz, qw) + ) last_obs = obs count += 1 diff --git a/dimos/mapping/ray_tracing/utils/test_loaded_map.py b/dimos/mapping/ray_tracing/utils/test_loaded_map.py index 8e83b46b104..dea2d30b078 100644 --- a/dimos/mapping/ray_tracing/utils/test_loaded_map.py +++ b/dimos/mapping/ray_tracing/utils/test_loaded_map.py @@ -17,6 +17,10 @@ from collections.abc import Iterator from pathlib import Path +from dimos_generated.geometry_msgs.msg import TransformStamped +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage import numpy as np import pytest from scipy.spatial.transform import Rotation @@ -24,9 +28,9 @@ from dimos.mapping.ray_tracing.utils.loaded_map import first_loaded_map, place_loaded_map from dimos.memory.store.sqlite import SqliteStore from dimos.memory.tf import StreamTF -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.geometry import transform_from_matrix +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.msgs.time import time_from_seconds pytestmark = [pytest.mark.skipif_aarch64, pytest.mark.skipif_macos] @@ -40,17 +44,21 @@ def store(tmp_path: Path) -> Iterator[SqliteStore]: def _cloud(x: float, ts: float) -> PointCloud2: - return PointCloud2.from_numpy( - np.array([[x, 0.0, 0.0]], dtype=np.float32), frame_id="map", timestamp=ts + return pointcloud_from_xyz( + np.array([[x, 0.0, 0.0]], dtype=np.float32), + header=Header(stamp=time_from_seconds(ts), frame_id="map"), ) -def _world_to_map(yaw_deg: float, x: float, y: float, ts: float) -> Transform: +def _world_to_map(yaw_deg: float, x: float, y: float, ts: float) -> TransformStamped: m = np.eye(4) m[:3, :3] = Rotation.from_euler("z", yaw_deg, degrees=True).as_matrix() m[0, 3], m[1, 3] = x, y - tf = Transform.from_matrix(m, frame_id="world", child_frame_id="map") - tf.ts = ts + tf = TransformStamped( + header=Header(stamp=time_from_seconds(ts), frame_id="world"), + child_frame_id="map", + transform=transform_from_matrix(m), + ) return tf @@ -69,17 +77,18 @@ def test_first_loaded_map_is_the_earliest_by_ts(store: SqliteStore) -> None: assert first is not None assert first.ts == T0 - np.testing.assert_allclose(first.data.points_f32(), [[1.0, 0.0, 0.0]]) + np.testing.assert_allclose(pointcloud_xyz(first.data), [[1.0, 0.0, 0.0]]) def test_place_loaded_map_moves_points_into_the_world(store: SqliteStore) -> None: store.stream("tf", TFMessage).append( - TFMessage(_world_to_map(90.0, 1.0, 2.0, T0)), ts=T0, pose=None + TFMessage(transforms=[_world_to_map(90.0, 1.0, 2.0, T0)]), ts=T0, pose=None ) tf = StreamTF(store.stream("tf", TFMessage)) points = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=np.float32) loaded = store.stream("loaded_map", PointCloud2).append( - PointCloud2.from_numpy(points, frame_id="map", timestamp=T0), ts=T0 + pointcloud_from_xyz(points, header=Header(stamp=time_from_seconds(T0), frame_id="map")), + ts=T0, ) placed = place_loaded_map(loaded, tf, "world", T0) diff --git a/dimos/mapping/ray_tracing/viz.py b/dimos/mapping/ray_tracing/viz.py index d74facb6e40..3b642dacb5a 100644 --- a/dimos/mapping/ray_tracing/viz.py +++ b/dimos/mapping/ray_tracing/viz.py @@ -20,13 +20,15 @@ import numpy as np -from dimos.visualization.rerun.bridge import RerunEntry, keyed_by_seq, region_entity +from dimos.msgs.pointcloud import pointcloud_xyz +from dimos.visualization.rerun.bridge import RerunEntry, keyed_by_region, region_entity if TYPE_CHECKING: + from dimos_generated.dimos_msgs.msg import RegionPointCloud2 + from dimos_generated.sensor_msgs.msg import PointCloud2 from numpy.typing import NDArray from rerun._baseclasses import Archetype - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.visualization.rerun.bridge import RerunMulti # rerun is imported inside the renderers: it is heavy and only viewers load it. @@ -56,12 +58,12 @@ def voxel_map_points( def render_voxel_map( msg: PointCloud2, voxel_size: float, height_range: tuple[float, float] ) -> Archetype: - return voxel_map_points(msg.points_f32(), voxel_size, height_range) + return voxel_map_points(pointcloud_xyz(msg).astype(np.float32), voxel_size, height_range) -@keyed_by_seq +@keyed_by_region def render_map_region( - msg: PointCloud2, + msg: RegionPointCloud2, voxel_size: float, height_range: tuple[float, float], z_band: tuple[float, float] | None = None, @@ -70,11 +72,11 @@ def render_map_region( Points outside z_band (world frame) are dropped, to show one storey. """ - pts = msg.points_f32() + pts = pointcloud_xyz(msg.cloud).astype(np.float32) if z_band is not None: pts = pts[(pts[:, 2] >= z_band[0]) & (pts[:, 2] < z_band[1])] cell = voxel_map_points(pts, voxel_size, height_range) - return [RerunEntry(region_entity(MAP_REGIONS_ENTITY, msg.seq), cell, static=True)] + return [RerunEntry(region_entity(MAP_REGIONS_ENTITY, msg.region_id), cell, static=True)] def log_loaded_map(points: NDArray[np.float32]) -> None: diff --git a/dimos/mapping/relocalization/lidar/relocalize.py b/dimos/mapping/relocalization/lidar/relocalize.py index ba885083e76..b270d778cc9 100644 --- a/dimos/mapping/relocalization/lidar/relocalize.py +++ b/dimos/mapping/relocalization/lidar/relocalize.py @@ -37,9 +37,12 @@ from typing import TYPE_CHECKING, NamedTuple +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import TransformStamped +from dimos_generated.std_msgs.msg import Header import numpy as np -from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.geometry import inverse_transform, transform_from_matrix from dimos.protocol.service.spec import BaseConfig from dimos.utils.logging_config import setup_logger @@ -65,7 +68,7 @@ class _Prepared(NamedTuple): class RelocAttempt(NamedTuple): """One decided attempt: the ``world -> map`` fix if accepted, and what it was scored on.""" - fix: Transform | None + fix: TransformStamped | None result: RegistrationResult @@ -259,7 +262,7 @@ def refine(coarse_T: np.ndarray) -> RegistrationResult: def relocalize( self, local_map: PointCloud, world_frame: str, map_frame: str - ) -> Transform | None: + ) -> TransformStamped | None: """The ``world_frame -> map_frame`` transform, or ``None`` when nothing was good enough.""" return self.attempt(local_map, world_frame, map_frame).fix @@ -269,7 +272,9 @@ def attempt(self, local_map: PointCloud, world_frame: str, map_frame: str) -> Re logger.info(f"align: fitness={result.fitness:.3f} rmse={result.inlier_rmse:.3f}") if result.fitness < self.config.fitness_threshold: return RelocAttempt(None, result) - placement = Transform.from_matrix( - np.asarray(result.transformation), frame_id=map_frame, child_frame_id=world_frame + placement = TransformStamped( + header=Header(frame_id=map_frame, stamp=Time(sec=0, nanosec=0)), + child_frame_id=world_frame, + transform=transform_from_matrix(np.asarray(result.transformation)), ) - return RelocAttempt(placement.inverse(), result) + return RelocAttempt(inverse_transform(placement), result) diff --git a/dimos/mapping/relocalization/lidar/replay.py b/dimos/mapping/relocalization/lidar/replay.py index a63b79a9b6d..298b97aeb51 100644 --- a/dimos/mapping/relocalization/lidar/replay.py +++ b/dimos/mapping/relocalization/lidar/replay.py @@ -24,11 +24,11 @@ from pathlib import Path from typing import NamedTuple +from dimos_cdrgen.registry import decode as cdr_decode from dimos_generated.geometry_msgs.msg import TransformStamped as Transform from dimos_generated.sensor_msgs.msg import PointCloud2 from dimos_generated.std_msgs.msg import Header from dimos_generated.tf2_msgs.msg import TFMessage -from dimos_message_build.registry import decode as cdr_decode import numpy as np import typer diff --git a/dimos/mapping/relocalization/lidar/test_replay.py b/dimos/mapping/relocalization/lidar/test_replay.py index c3b7b9c8863..c76aa7547f4 100644 --- a/dimos/mapping/relocalization/lidar/test_replay.py +++ b/dimos/mapping/relocalization/lidar/test_replay.py @@ -16,14 +16,23 @@ from pathlib import Path from types import SimpleNamespace +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from pytest_mock import MockerFixture import rerun as rr from scipy.spatial.transform import Rotation +from dimos.msgs.geometry import transform_from_matrix +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_to_open3d, pointcloud_xyz +from dimos.msgs.time import time_from_seconds, to_seconds + pytest.importorskip("dimos_voxel_ray_tracing") +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.tf2_msgs.msg import TFMessage + from dimos.mapping.ray_tracing.utils.loaded_map import LOADED_MAP_STREAM from dimos.mapping.relocalization.lidar import replay as replay_module from dimos.mapping.relocalization.lidar.relocalize import PRESETS, RelocAttempt @@ -34,21 +43,27 @@ write_loaded_map, ) from dimos.memory.store.sqlite import SqliteStore -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -def _fix(yaw_deg: float, x: float = 0.0, y: float = 0.0) -> Transform: +def _fix(yaw_deg: float, x: float = 0.0, y: float = 0.0) -> TransformStamped: m = np.eye(4) m[:3, :3] = Rotation.from_euler("z", yaw_deg, degrees=True).as_matrix() m[0, 3], m[1, 3] = x, y - return Transform.from_matrix(m, frame_id="odom", child_frame_id="map") + return TransformStamped( + header=Header(stamp=time_from_seconds(0.0), frame_id="odom"), + child_frame_id="map", + transform=transform_from_matrix(m), + ) -def _edge(parent: str, child: str, ts: float, x: float = 0.0) -> Transform: - return Transform(frame_id=parent, child_frame_id=child, translation=Vector3(x, 0.0, 0.0), ts=ts) +def _edge(parent: str, child: str, ts: float, x: float = 0.0) -> TransformStamped: + return TransformStamped( + header=Header(stamp=time_from_seconds(ts), frame_id=parent), + child_frame_id=child, + transform=Transform( + translation=Vector3(x, 0.0, 0.0), rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + ) def test_fix_error_wraps_yaw_and_measures_translation() -> None: @@ -61,16 +76,16 @@ def test_fix_error_wraps_yaw_and_measures_translation() -> None: @pytest.mark.skipif_macos def test_write_loaded_map_writes_once(tmp_path: Path) -> None: pts = np.array([[1, 0, 0], [0, 1, 0]], dtype=np.float32) - premap = PointCloud2.from_numpy(pts, frame_id="map") + premap = pointcloud_from_xyz(pts, header=Header(stamp=time_from_seconds(0.0), frame_id="map")) fix = _fix(90.0, x=1.0) with SqliteStore(path=str(tmp_path / "r.db")) as store: assert write_loaded_map(store, premap, fix, 5.0, "odom") is True stream = store.stream(LOADED_MAP_STREAM, PointCloud2) assert stream.count() == 1 written = stream.first() - assert written.data.frame_id == "odom" + assert written.data.header.frame_id == "odom" assert written.ts == 5.0 - np.testing.assert_allclose(written.data.points_f32(), [[1, 1, 0], [0, 0, 0]], atol=1e-6) + np.testing.assert_allclose(pointcloud_xyz(written.data), [[1, 1, 0], [0, 0, 0]], atol=1e-6) assert write_loaded_map(store, premap, fix, 6.0, "odom") is False assert stream.count() == 1 @@ -87,11 +102,11 @@ def test_recorded_fix_is_the_first_world_to_map_edge(tmp_path: Path) -> None: _edge("odom", "map", 2.0, x=2.0), _edge("odom", "map", 3.0, x=3.0), ): - tf.append(TFMessage(edge), ts=edge.ts, pose=None) + tf.append(TFMessage(transforms=[edge]), ts=to_seconds(edge.header.stamp), pose=None) recorded = recorded_fix(store, "odom", "map") assert recorded is not None - assert recorded.translation.x == 2.0 + assert recorded.transform.translation.x == 2.0 assert recorded_fix(store, "odom", "nowhere") is None @@ -111,7 +126,9 @@ def test_replay_spaces_attempts_and_stops_after_the_fix( set_time = mocker.spy(rr, "set_time") span = np.arange(0.0, 1.0, 0.05) dense = np.array([(x, y, 0.5) for x in span for y in span], dtype=np.float32) - premap = PointCloud2.from_numpy(np.eye(3, dtype=np.float32), frame_id="map") + premap = pointcloud_from_xyz( + np.eye(3, dtype=np.float32), header=Header(stamp=time_from_seconds(0.0), frame_id="map") + ) with SqliteStore(path=str(tmp_path / "r.db")) as store: lidar = store.stream("lidar", PointCloud2) @@ -120,8 +137,13 @@ def test_replay_spaces_attempts_and_stops_after_the_fix( ts = 10.0 + i * 0.5 # The first frame's local map is under min_local_points. points = dense[:3] if i == 0 else dense - tf.append(TFMessage(_edge("odom", "lidar", ts)), ts=ts, pose=None) - lidar.append(PointCloud2.from_numpy(points, frame_id="lidar", timestamp=ts), ts=ts) + tf.append(TFMessage(transforms=[_edge("odom", "lidar", ts)]), ts=ts, pose=None) + lidar.append( + pointcloud_from_xyz( + points, header=Header(stamp=time_from_seconds(ts), frame_id="lidar") + ), + ts=ts, + ) result = replay( store, lidar_stream="lidar", @@ -136,7 +158,7 @@ def test_replay_spaces_attempts_and_stops_after_the_fix( to_time=None, ) - relocalizer_class.assert_called_once_with(premap.pointcloud, PRESETS["mid360"]) + relocalizer_class.assert_called_once_with(pointcloud_to_open3d(premap), PRESETS["mid360"]) assert attempt.call_count == 2 # The sparse frame at 0.0 s is skipped, so the first attempt waits for 0.5 s. assert [a.t_s for a in result.attempts] == [0.5, 1.5] diff --git a/dimos/mapping/relocalization/lidar/tune.py b/dimos/mapping/relocalization/lidar/tune.py index 6237f28b723..e78396a9ed7 100644 --- a/dimos/mapping/relocalization/lidar/tune.py +++ b/dimos/mapping/relocalization/lidar/tune.py @@ -39,6 +39,10 @@ import time from typing import TYPE_CHECKING, Any, NamedTuple +from dimos_cdrgen.registry import decode as cdr_decode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header import numpy as np import typer @@ -49,7 +53,8 @@ RelocalizeConfig, ) from dimos.mapping.relocalization.lidar.replay import main as replay_main -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_to_open3d, pointcloud_xyz +from dimos.msgs.time import time_from_seconds from dimos.utils.data import get_data if TYPE_CHECKING: @@ -190,7 +195,7 @@ class Dataset(NamedTuple): # so probes come from before that. "go2-sf-area1": Dataset( recording="recording_go2_mid360_2026-05-29_4-45pm-PST_corrected.db", - premap="recording_go2_mid360_2026-05-29_4-45pm-PST_corrected.pc2.lcm", + premap="recording_go2_mid360_2026-05-29_4-45pm-PST_corrected.pc2.cdr", # Runs the full stretch before the premap's own scans: 0-250 s is # ground the loop revisits after 400 s and the premap therefore # holds, while 250-400 s is a one-time excursion it never saw. @@ -283,7 +288,7 @@ def fixtures(name: str) -> tuple[PointCloud2, Any]: from dimos.memory.store.sqlite import SqliteStore ds = DATASETS[name] - premap = PointCloud2.lcm_decode(get_data(ds.premap).read_bytes()) + premap = cdr_decode(get_data(ds.premap).read_bytes(), PointCloud2) store = SqliteStore(path=str(get_data(ds.recording)), must_exist=True) store.start() return premap, store @@ -293,13 +298,15 @@ def fixtures(name: str) -> tuple[PointCloud2, Any]: def premap_index(name: str, voxel: float = 0.2) -> Any: """The premap thinned for nearest-neighbour queries, built once.""" premap, _ = fixtures(name) - return premap.pointcloud.voxel_down_sample(voxel) + return pointcloud_to_open3d(premap).voxel_down_sample(voxel) def premap_distance(name: str, points: np.ndarray) -> np.ndarray: """Per-point distance from ``points`` to the nearest premap point.""" - cloud = PointCloud2.from_numpy(points, frame_id="world", timestamp=0.0) - return np.asarray(cloud.pointcloud.compute_point_cloud_distance(premap_index(name))) + cloud = pointcloud_from_xyz( + points, header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)) + ) + return np.asarray(pointcloud_to_open3d(cloud).compute_point_cloud_distance(premap_index(name))) def accumulate( @@ -337,7 +344,7 @@ def accumulate( if obs.pose_tuple is None: continue x, y, z = obs.pose_tuple[:3] - mapper.add_frame_world(obs.data.points_f32(), (x, y, z)) + mapper.add_frame_world(pointcloud_xyz(obs.data).astype(np.float32), (x, y, z)) used += 1 last_ts = obs.ts if used == n_frames: @@ -345,7 +352,9 @@ def accumulate( if used < n_frames: raise ValueError(f"only {used} of {n_frames} scans had a pose from {start:.0f}s") - return PointCloud2.from_numpy(mapper.global_map(), frame_id="world", timestamp=last_ts) + return pointcloud_from_xyz( + mapper.global_map(), header=Header(frame_id="world", stamp=time_from_seconds(last_ts)) + ) @lru_cache(maxsize=512) @@ -386,7 +395,7 @@ def coverage(name: str, cloud: PointCloud2) -> float: whether the aligner answered it: a probe standing where the premap has no points cannot be placed by any parameter setting. """ - distances = premap_distance(name, np.asarray(cloud.pointcloud.points)) + distances = premap_distance(name, pointcloud_xyz(cloud)) return float((distances < COVERAGE_TOL_M).mean()) @@ -467,7 +476,7 @@ def identity_error(T: np.ndarray) -> tuple[float, float, float]: def applied(T: np.ndarray, cloud: PointCloud2) -> np.ndarray: """``cloud``'s points where ``T`` puts them.""" - points: np.ndarray = np.asarray(cloud.pointcloud.points) @ T[:3, :3].T + T[:3, 3] + points: np.ndarray = pointcloud_xyz(cloud) @ T[:3, :3].T + T[:3, 3] return points @@ -495,7 +504,7 @@ def run_probes( premap, _ = fixtures(name) # The premap's preprocessing depends on the config but not the probe, so # it is paid once per trial rather than once per attempt. - relocalizer = LidarRelocalizer(premap.pointcloud, config) + relocalizer = LidarRelocalizer(pointcloud_to_open3d(premap), config) rate = scan_rate(name) probes: list[Probe] = [] @@ -512,7 +521,7 @@ def run_probes( frames = min(max(int(elapsed * rate), min_frames), max_frames) cloud = local_map(name, frames, start, voxel) t0, c0 = time.monotonic(), time.process_time() - result = relocalizer.align(cloud.pointcloud) + result = relocalizer.align(pointcloud_to_open3d(cloud)) dt = time.monotonic() - t0 elapsed += dt wall += dt @@ -525,13 +534,13 @@ def run_probes( break T = np.asarray(result.transformation) - truth_points = np.asarray(cloud.pointcloud.points) + truth_points = pointcloud_xyz(cloud) moved = applied(T, cloud) translation, rotation, tilt = identity_error(T) probes.append( Probe( start=start, - n_points=len(cloud), + n_points=cloud.width * cloud.height, coverage=place_coverage(name, start, voxel), translation_m=translation, rotation_deg=rotation, @@ -614,6 +623,7 @@ def view(name: str, probes: list[Probe], out: str | None) -> None: import rerun as rr from dimos.visualization.rerun.init import rerun_init + from dimos.visualization.rerun.message_helpers import cloud_archetype premap, _ = fixtures(name) rerun_init("dimos relocalize eval") @@ -621,22 +631,20 @@ def view(name: str, probes: list[Probe], out: str | None) -> None: rr.save(out) else: rr.spawn() - rr.log("world/premap/pointcloud", premap.to_rerun(voxel_size=0.05), static=True) + rr.log("world/premap/pointcloud", cloud_archetype(premap, ui_radius=0.05), static=True) for probe in probes: - aligned = PointCloud2.from_numpy( - applied(probe.transform, probe.cloud), - frame_id=probe.cloud.frame_id, - timestamp=probe.cloud.ts, + aligned = pointcloud_from_xyz( + applied(probe.transform, probe.cloud), header=probe.cloud.header ) entity = f"world/probe_{probe.start:04.0f}s" rr.log( f"{entity}/truth", - probe.cloud.to_rerun(voxel_size=0.05, colors=[76, 220, 41]), + rr.Points3D(pointcloud_xyz(probe.cloud), radii=0.05, colors=[76, 220, 41]), static=True, ) rr.log( f"{entity}/aligned", - aligned.to_rerun(voxel_size=0.05, colors=[231, 76, 60]), + rr.Points3D(pointcloud_xyz(aligned), radii=0.05, colors=[231, 76, 60]), static=True, ) if out is not None: @@ -681,7 +689,7 @@ def _register( DatasetOpt = typer.Option(DEFAULT_DATASET, "--dataset", "-d", help="Registered dataset name") RecordingOpt = typer.Option(None, "--recording", help="Recording, resolved through get_data") -PremapOpt = typer.Option(None, "--premap", help="Premap .pc2.lcm, resolved through get_data") +PremapOpt = typer.Option(None, "--premap", help="Premap .pc2.cdr, resolved through get_data") LidarOpt = typer.Option(None, "--lidar", help="Lidar stream in the recording") FromOpt = typer.Option(None, "--from", help="Earliest second probes may come from") ToOpt = typer.Option(None, "--to", help="Latest second probes may come from") @@ -727,7 +735,7 @@ def run( console = Console() table = _table( - f"{name} premap {len(pre):,} pts preset={preset} " + f"{name} premap {pre.width * pre.height:,} pts preset={preset} " f"cutoff={config.fitness_threshold} " f"frames {config.min_frames}-{config.max_frames}", *PROBE_COLUMNS, diff --git a/dimos/mapping/relocalization/module.py b/dimos/mapping/relocalization/module.py index 81b7b81fe18..d171765ef16 100644 --- a/dimos/mapping/relocalization/module.py +++ b/dimos/mapping/relocalization/module.py @@ -39,34 +39,41 @@ from __future__ import annotations +from copy import deepcopy import math +import time from typing import Any +from dimos_cdrgen.registry import decode as cdr_decode +from dimos_generated.geometry_msgs.msg import TransformStamped +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.tf2_msgs.msg import TFMessage import reactivex as rx from reactivex import Observable, Subject, operators as ops from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import Out -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.geometry import yaw +from dimos.msgs.time import time_from_seconds from dimos.utils.data import resolve_named_path from dimos.utils.logging_config import setup_logger logger = setup_logger() -MAP_SUFFIX = ".pc2.lcm" +MAP_SUFFIX = ".pc2.cdr" -def yaw_deg(tf: Transform) -> float: - return math.degrees(tf.rotation.euler.z) +def yaw_deg(tf: TransformStamped) -> float: + return math.degrees(yaw(tf.transform.rotation)) -def fix_stream(fixes: Observable[Transform], interval: float) -> Observable[Transform]: +def fix_stream( + fixes: Observable[TransformStamped], interval: float +) -> Observable[TransformStamped]: """Every accepted fix as it lands, then again every ``interval`` s (once only if <= 0).""" - def now_and_again(fix: Transform) -> Observable[Transform]: + def now_and_again(fix: TransformStamped) -> Observable[TransformStamped]: again = rx.interval(interval).pipe(ops.map(lambda _: fix)) if interval > 0 else rx.empty() return rx.concat(rx.of(fix), again) @@ -75,7 +82,7 @@ def now_and_again(fix: Transform) -> Observable[Transform]: class Config(ModuleConfig): # Premap stem or path, e.g. `--map-file=go2_hongkong_office_twopass_map`; - # `.pc2.lcm` is appended if absent. Without one the module runs but never + # `.pc2.cdr` is appended if absent. Without one the module runs but never # attempts a fix. map_file: str | None = None # What the live fixed frame is called, whatever the odometry is in @@ -102,16 +109,14 @@ class RelocalizationModule(Module): def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) - self.fixes: Subject[Transform] = Subject() + self.fixes: Subject[TransformStamped] = Subject() self.premap: PointCloud2 | None = None @rpc def start(self) -> None: super().start() self.register_disposable( - fix_stream(self.fixes, self.config.tf_interval).subscribe( - lambda tf: self.tf.publish(TFMessage(tf.now())) - ) + fix_stream(self.fixes, self.config.tf_interval).subscribe(self._publish_tf) ) if not self.config.map_file: logger.info("Relocalization module disabled (no map_file configured)") @@ -119,11 +124,16 @@ def start(self) -> None: self._load_premap(self.config.map_file) logger.info(f"Relocalization module started: map_file={self.config.map_file!r}") + def _publish_tf(self, tf: TransformStamped) -> None: + published = deepcopy(tf) + published.header.stamp = time_from_seconds(time.time()) + self.tf.publish(TFMessage(transforms=[published])) + def _load_premap(self, map_file: str) -> None: path = resolve_named_path(map_file, MAP_SUFFIX) - premap = PointCloud2.lcm_decode(path.read_bytes()) + premap = cdr_decode(path.read_bytes(), PointCloud2) - premap.frame_id = self.config.map_frame + premap.header.frame_id = self.config.map_frame self.premap = premap self.register_disposable( @@ -141,19 +151,19 @@ def keep_relocalizing(self) -> bool: """Whether to keep attempting. Implementations gate their input on this.""" return not (self._placed and self.config.relocalize_once) - def submit(self, tf: Transform, source: str = "") -> None: + def submit(self, tf: TransformStamped, source: str = "") -> None: """Publish a ``world_frame -> map`` fix the implementation already decided to believe.""" world, map_frame = self.config.world_frame, self.config.map_frame - assert (tf.frame_id, tf.child_frame_id) == (world, map_frame), ( + assert (tf.header.frame_id, tf.child_frame_id) == (world, map_frame), ( f"relocalize {source}: expected {world!r} -> {map_frame!r}, " - f"got {tf.frame_id!r} -> {tf.child_frame_id!r}" + f"got {tf.header.frame_id!r} -> {tf.child_frame_id!r}" ) logger.info( "relocalize fix", source=source, world=world, map_frame=map_frame, - translation=str(tf.translation), + translation=str(tf.transform.translation), yaw_deg=round(yaw_deg(tf), 1), ) self.fixes.on_next(tf) diff --git a/dimos/mapping/relocalization/test_module.py b/dimos/mapping/relocalization/test_module.py index 3c0384c432a..ea50f81f6b5 100644 --- a/dimos/mapping/relocalization/test_module.py +++ b/dimos/mapping/relocalization/test_module.py @@ -15,6 +15,8 @@ from types import SimpleNamespace from typing import get_type_hints +from dimos_cdrgen.registry import encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import TransformStamped from dimos_generated.std_msgs.msg import Header from dimos_generated.vision_msgs.msg import Detection3DArray @@ -23,6 +25,7 @@ from scipy.spatial.transform import Rotation from dimos.core.stream import In +from dimos.core.transport_factory import rpc_backend from dimos.mapping.relocalization.lidar.module import LidarWindowRelocalization from dimos.mapping.relocalization.module import RelocalizationModule from dimos.msgs.geometry import inverse_transform, transform_from_matrix, transform_matrix @@ -30,8 +33,11 @@ @pytest.fixture -def module(): +def module(mocker): """Build a real module (and dispose it); `cls, **config` picks the class.""" + backend = rpc_backend() + for method in ("start", "serve_module_rpc", "stop"): + mocker.patch.object(backend, method) built = [] def build(cls=RelocalizationModule, **config): @@ -56,7 +62,7 @@ def test_submit_publishes_and_checks_frames(module): m = module() got = fixes(m) tf = TransformStamped( - header=Header(frame_id="world"), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), child_frame_id="map", transform=transform_from_matrix(np.eye(4)), ) @@ -68,7 +74,7 @@ def test_submit_publishes_and_checks_frames(module): with pytest.raises(AssertionError): m.submit( TransformStamped( - header=Header(frame_id="map"), + header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), child_frame_id="world", transform=transform_from_matrix(np.eye(4)), ) @@ -78,7 +84,7 @@ def test_submit_publishes_and_checks_frames(module): def test_relocalize_once_stops_after_the_first_fix(module): """What the flag does, either way. Which one is the default is a policy call.""" tf = TransformStamped( - header=Header(frame_id="world"), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), child_frame_id="map", transform=transform_from_matrix(np.eye(4)), ) @@ -103,7 +109,7 @@ def test_the_fix_goes_out_the_moment_it_is_accepted(interval): fixes, sent = Subject(), [] disposable = fix_stream(fixes, interval=interval).subscribe(sent.append) tf = TransformStamped( - header=Header(frame_id="world"), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), child_frame_id="map", transform=transform_from_matrix(np.eye(4)), ) @@ -116,7 +122,12 @@ def test_premap_defines_the_map_frame_and_waits_for_a_fix(module, tmp_path): """Loading is the base's: every strategy reads a premap and publishes it, once placed.""" path = tmp_path / "somewhere.pc2.cdr" path.write_bytes( - pointcloud_from_xyz(np.zeros((5, 3), dtype=np.float32), header=Header()).encode() + cdr_encode( + pointcloud_from_xyz( + np.zeros((5, 3), dtype=np.float32), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + ) ) m = module() published, disposables = [], [] @@ -132,7 +143,7 @@ def test_premap_defines_the_map_frame_and_waits_for_a_fix(module, tmp_path): assert published == [] # ... which stays silent until a fix lands m.submit( TransformStamped( - header=Header(frame_id="world"), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), child_frame_id="map", transform=transform_from_matrix(np.eye(4)), ) @@ -194,7 +205,12 @@ def test_the_match_runs_on_a_window_of_the_last_scans(): window(scans, cfg, interval=0.001).subscribe(matched.append) for i in range(5): - scans.on_next(pointcloud_from_xyz(np.full((4, 3), i, dtype=np.float32), header=Header())) + scans.on_next( + pointcloud_from_xyz( + np.full((4, 3), i, dtype=np.float32), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + ) time.sleep(0.01) # clear the throttle, so every scan gets its attempt # The first scan is below min_frames; from then on the window is full. @@ -209,7 +225,7 @@ def test_from_matrix_inverse_matches_linalg_inv(): T[:3, 3] = [1.5, -2.0, 0.3] tf = inverse_transform( TransformStamped( - header=Header(frame_id="map"), + header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), child_frame_id="world", transform=transform_from_matrix(T), ) diff --git a/dimos/mapping/test_generated_costmapper.py b/dimos/mapping/test_generated_costmapper.py index 1b918edf3c9..bae0adcdf20 100644 --- a/dimos/mapping/test_generated_costmapper.py +++ b/dimos/mapping/test_generated_costmapper.py @@ -14,6 +14,7 @@ import math +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid @@ -33,7 +34,7 @@ def test_cost_mapper_generates_cdr_grid_with_exact_header(request): header = Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id="map") cloud = pointcloud_from_xyz(np.array([[0, 0, 0], [1, 1, 1]]), header=header) grid = mapper._calculate_costmap(cloud) - decoded = OccupancyGrid.decode(grid.encode()) + decoded = cdr_decode(cdr_encode(grid), OccupancyGrid) assert decoded.header == header assert np.any(occupancy_view(decoded) == 100) assert np.any(occupancy_view(decoded) == 0) @@ -48,10 +49,13 @@ def test_initial_safe_radius_respects_rotated_origin_and_replaces_readonly_data( height=3, resolution=1, origin=Pose( - position=Point(x=1, y=-1), orientation=quaternion_from_euler(0, 0, math.pi / 2) + position=Point(x=1, y=-1, z=0.0), + orientation=quaternion_from_euler(0, 0, math.pi / 2), ), + map_load_time=Time(sec=0, nanosec=0), ), - data=[100] * 9, + data=np.asarray([100] * 9, dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) mapper._apply_initial_safe_radius(grid) expected = np.full((3, 3), 100, dtype=np.int8) diff --git a/dimos/mapping/test_generated_odometry_hist.py b/dimos/mapping/test_generated_odometry_hist.py index 61cdae3fe40..b6a04aaf9f0 100644 --- a/dimos/mapping/test_generated_odometry_hist.py +++ b/dimos/mapping/test_generated_odometry_hist.py @@ -12,19 +12,45 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time -from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseWithCovariance, + Quaternion, + Twist, + TwistWithCovariance, + Vector3, +) from dimos_generated.nav_msgs.msg import Odometry, Path from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from dimos.mapping.odometry_hist import OdometryHist, path_at_true_height def odometry(x: float, nanosec: int, sec: int = 1700000000) -> Odometry: - message = Odometry(header=Header(stamp=Time(sec=sec, nanosec=nanosec), frame_id="odom")) - message.pose.pose = Pose(position=Point(x=x, y=2, z=3), orientation=Quaternion(w=1)) - return Odometry.decode(message.encode()) + message = Odometry( + header=Header(stamp=Time(sec=sec, nanosec=nanosec), frame_id="odom"), + child_frame_id="", + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), + twist=TwistWithCovariance( + twist=Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + covariance=np.zeros(36, dtype=np.float64), + ), + ) + message.pose.pose = Pose( + position=Point(x=x, y=2, z=3), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ) + return cdr_decode(cdr_encode(message), Odometry) @pytest.mark.asyncio @@ -32,7 +58,7 @@ async def test_history_retains_exact_stamp_and_one_nanosecond_interval(request): module = OdometryHist(min_publish_interval_seconds=1e-9) request.addfinalizer(module.stop) paths = [] - module.odom_hist.subscribe(lambda path: paths.append(Path.decode(path.encode()))) + module.odom_hist.subscribe(lambda path: paths.append(cdr_decode(cdr_encode(path), Path))) source = odometry(1, 123456789) await module.handle_odometry(source) source.pose.pose.position.x = 100 @@ -60,4 +86,4 @@ async def test_history_throttles_bounds_and_restarts_at_zero(request): assert paths[-1].header.frame_id == "world" assert all(p.header.frame_id == "world" for p in paths[-1].poses) await module.handle_odometry(odometry(4, 0, sec=0)) - assert len(paths) == 3 and paths[-1].header.stamp == Time() + assert len(paths) == 3 and paths[-1].header.stamp == Time(sec=0, nanosec=0) diff --git a/dimos/mapping/test_generated_voxels.py b/dimos/mapping/test_generated_voxels.py index 6e3dc1e7722..2f964ed0db2 100644 --- a/dimos/mapping/test_generated_voxels.py +++ b/dimos/mapping/test_generated_voxels.py @@ -34,14 +34,21 @@ def test_generated_cloud_union_and_column_carving(backend, carve): ) try: for points in [[[0.1, 0.1, 0.1], [0.1, 0.1, 1.1], [1.1, 1.1, 0.1]], [[0.1, 0.1, 0.6]]]: - grid.add_frame(pointcloud_from_xyz(np.array(points), header=Header(frame_id="map"))) + grid.add_frame( + pointcloud_from_xyz( + np.array(points), header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)) + ) + ) actual = {tuple(point) for point in grid.points()} expected = {(0.25, 0.25, 0.75), (1.25, 1.25, 0.25)} if not carve: expected |= {(0.25, 0.25, 0.25), (0.25, 0.25, 1.25)} assert actual == expected grid.add_frame( - pointcloud_from_xyz(np.array([[np.nan, 0, 0], [0, np.inf, 0]]), header=Header()) + pointcloud_from_xyz( + np.array([[np.nan, 0, 0], [0, np.inf, 0]]), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) ) assert {tuple(point) for point in grid.points()} == expected finally: @@ -52,7 +59,7 @@ def test_voxel_output_retains_exact_latest_stamp_and_invalidates_cache(): grid = VoxelGrid(voxel_size=0.5, device="CPU:0", frame_id="map", show_startup_log=False) try: empty = grid.get_global_pointcloud2() - assert empty.width == 0 and empty.header.stamp == Time() + assert empty.width == 0 and empty.header.stamp == Time(sec=0, nanosec=0) header = Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id="map") source = pointcloud_from_xyz(np.array([[0.1, 0.2, 0.3]]), header=header) grid.add_frame(source) @@ -67,8 +74,12 @@ def test_voxel_output_retains_exact_latest_stamp_and_invalidates_cache(): assert second is not first assert second.header == header and second.width == 2 assert first.width == 1 - grid.add_frame(pointcloud_from_xyz(np.empty((0, 3)), header=Header(frame_id="map"))) - assert grid.get_global_pointcloud2().header.stamp == Time() + grid.add_frame( + pointcloud_from_xyz( + np.empty((0, 3)), header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)) + ) + ) + assert grid.get_global_pointcloud2().header.stamp == Time(sec=0, nanosec=0) assert grid.get_global_pointcloud2().width == 2 finally: grid.dispose() diff --git a/dimos/mapping/test_voxels.py b/dimos/mapping/test_voxels.py index 26963134604..5bce0939a84 100644 --- a/dimos/mapping/test_voxels.py +++ b/dimos/mapping/test_voxels.py @@ -16,6 +16,8 @@ from collections.abc import Generator +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import PointCloud2 from dimos_generated.std_msgs.msg import Header import numpy as np @@ -38,15 +40,19 @@ def grid() -> Generator[VoxelGrid, None, None]: def lidar_frame() -> PointCloud2: axis = np.arange(20) * 0.05 + 0.025 points = np.stack(np.meshgrid(axis, axis, axis, indexing="ij"), axis=-1).reshape(-1, 3) - message = pointcloud_from_xyz(points, header=Header(frame_id="world")) - return PointCloud2.decode(message.encode()) + message = pointcloud_from_xyz( + points, header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)) + ) + return cdr_decode(cdr_encode(message), PointCloud2) def test_ingest_a_few(grid: VoxelGrid) -> None: for offset in (0.0, 1.0, 2.0): points = np.array([[offset + 0.025, 0.025, 0.025], [offset + 0.075, 0.025, 0.025]]) - frame = pointcloud_from_xyz(points, header=Header(frame_id="world")) - grid.add_frame(PointCloud2.decode(frame.encode())) + frame = pointcloud_from_xyz( + points, header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)) + ) + grid.add_frame(cdr_decode(cdr_encode(frame), PointCloud2)) assert grid.get_global_pointcloud2().width == 6 np.testing.assert_allclose( np.sort(pointcloud_xyz(grid.get_global_pointcloud2())[:, 0]), @@ -69,7 +75,7 @@ def test_roundtrip(lidar_frame: PointCloud2, voxel_size: float, expected_points: np.sort(pointcloud_xyz(lidar_frame), axis=0), atol=1e-6, ) - grid.add_frame(PointCloud2.decode(first.encode())) + grid.add_frame(cdr_decode(cdr_encode(first), PointCloud2)) assert grid.get_global_pointcloud2().width == expected_points np.testing.assert_array_equal( pointcloud_xyz(grid.get_global_pointcloud2()), pointcloud_xyz(first) diff --git a/dimos/mapping/voxels/grid.py b/dimos/mapping/voxels/grid.py index a3dc189c0d5..82af0947365 100644 --- a/dimos/mapping/voxels/grid.py +++ b/dimos/mapping/voxels/grid.py @@ -15,12 +15,16 @@ from __future__ import annotations from collections.abc import Callable -import time +from copy import deepcopy from typing import TYPE_CHECKING +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header + from dimos.mapping.voxels.impl.o3d import O3dVoxels from dimos.mapping.voxels.impl.packed import PackedVoxels -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_from_xyz from dimos.utils.decorators.decorators import simple_mcache from dimos.utils.logging_config import setup_logger @@ -55,23 +59,27 @@ def __init__( show_startup_log: bool = True, stamped: bool = False, ) -> None: - import open3d.core as o3c # type: ignore[import-untyped] - self._voxel_size = voxel_size self._frame_id = frame_id - # ponytail: stamps are CPU only; on CUDA they'd ride as hashmap values - use_cuda = not stamped and device.startswith("CUDA") and o3c.cuda.is_available() - self._impl: O3dVoxels | PackedVoxels = ( - O3dVoxels(voxel_size, block_count, carve_columns, o3c.Device(device)) - if use_cuda - else PackedVoxels(voxel_size, carve_columns, stamped) - ) + use_cuda = False + self._impl: O3dVoxels | PackedVoxels + if not stamped and device.startswith("CUDA"): + import open3d.core as o3c # type: ignore[import-untyped] + + use_cuda = o3c.cuda.is_available() + self._impl = ( + O3dVoxels(voxel_size, block_count, carve_columns, o3c.Device(device)) + if use_cuda + else PackedVoxels(voxel_size, carve_columns, stamped) + ) + else: + self._impl = PackedVoxels(voxel_size, carve_columns, stamped) if show_startup_log: logger.info(f"VoxelGrid using device: {device if use_cuda else 'CPU:0 (packed-numpy)'}") - self._latest_frame_ts: float = 0.0 + self._latest_header = Header(frame_id=frame_id, stamp=Time(sec=0, nanosec=0)) self._disposed = False def _check_disposed(self) -> None: @@ -80,10 +88,9 @@ def _check_disposed(self) -> None: def add_frame(self, frame: PointCloud2) -> None: self._check_disposed() - if frame.ts is not None: - self._latest_frame_ts = frame.ts - self._impl.add_frame(frame) + self._latest_header = deepcopy(frame.header) + self._latest_header.frame_id = self._frame_id self.get_global_pointcloud.invalidate_cache(self) self.get_global_pointcloud2.invalidate_cache(self) @@ -99,11 +106,7 @@ def reproject(self, place: Callable[[PointCloud2], PointCloud2]) -> None: @simple_mcache def get_global_pointcloud2(self) -> PointCloud2: self._check_disposed() - return PointCloud2( - ensure_legacy_pcd(self.get_global_pointcloud()), - frame_id=self._frame_id, - ts=self._latest_frame_ts if self._latest_frame_ts else time.time(), - ) + return pointcloud_from_xyz(self._impl.points(), header=self._latest_header) @simple_mcache def get_global_pointcloud(self) -> o3d.t.geometry.PointCloud: @@ -131,18 +134,3 @@ def dispose(self) -> None: self.get_global_pointcloud.invalidate_cache(self) # type: ignore[attr-defined] self.get_global_pointcloud2.invalidate_cache(self) # type: ignore[attr-defined] self._impl.dispose() - - -def ensure_legacy_pcd( - pcd_any: o3d.t.geometry.PointCloud | o3d.geometry.PointCloud, -) -> o3d.geometry.PointCloud: - import open3d as o3d # type: ignore[import-untyped] - - if isinstance(pcd_any, o3d.geometry.PointCloud): - return pcd_any - - assert isinstance(pcd_any, o3d.t.geometry.PointCloud), ( - "Input must be a legacy PointCloud or a tensor PointCloud" - ) - - return pcd_any.to_legacy() diff --git a/dimos/mapping/voxels/impl/packed.py b/dimos/mapping/voxels/impl/packed.py index 83f9d5f2e25..0a082c3e9cd 100644 --- a/dimos/mapping/voxels/impl/packed.py +++ b/dimos/mapping/voxels/impl/packed.py @@ -16,10 +16,13 @@ from collections.abc import Callable +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header import numpy as np from dimos.mapping.voxels.keys import FIELD_BITS as _BITS, FIELD_MASK as _MASK, KEY_OFFSET as _BIAS -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.msgs.time import time_from_seconds, to_seconds class PackedVoxels: @@ -56,7 +59,8 @@ def _centers(self, k: np.ndarray) -> np.ndarray: return (vox + np.float32(0.5)) * np.float32(self._voxel_size) def add_frame(self, frame: PointCloud2) -> None: - pts = frame.points_f32() + pts = pointcloud_xyz(frame).astype(np.float32) + pts = pts[np.isfinite(pts).all(axis=1)] if not len(pts): return new = np.unique(self._pack(pts)) @@ -84,7 +88,7 @@ def add_frame(self, frame: PointCloud2) -> None: pos, new = pos[fresh], new[fresh] self._keys = np.insert(keys, pos, new) if self.stamped: - self._ts = np.insert(self._ts, pos, frame.ts) + self._ts = np.insert(self._ts, pos, to_seconds(frame.header.stamp)) self._raw = np.insert(self._raw, pos, new) def reproject(self, place: Callable[[PointCloud2], PointCloud2]) -> None: @@ -94,7 +98,15 @@ def reproject(self, place: Callable[[PointCloud2], PointCloud2]) -> None: if not len(self._keys): return keys = self._pack( - place(PointCloud2.from_numpy(self._centers(self._raw), stamps=self._ts)).points_f32() + pointcloud_xyz( + place( + pointcloud_from_xyz( + self._centers(self._raw), + header=Header(stamp=time_from_seconds(0.0), frame_id=""), + stamps=self._ts, + ) + ) + ) ) order = np.lexsort((-self._ts, keys)) # by key, latest stamp first keys, ts, raw = keys[order], self._ts[order], self._raw[order] diff --git a/dimos/mapping/voxels/test_grid.py b/dimos/mapping/voxels/test_grid.py index 365661d437f..8ad8731f056 100644 --- a/dimos/mapping/voxels/test_grid.py +++ b/dimos/mapping/voxels/test_grid.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.std_msgs.msg import Header import numpy as np @@ -35,7 +36,7 @@ def test_cpu_accumulation_preserves_latest_exact_stamp_and_output_frame() -> Non grid.add_frame(first) grid.add_frame(second) result = grid.get_global_pointcloud2() - decoded = type(result).decode(result.encode()) + decoded = cdr_decode(cdr_encode(result), type(result)) assert decoded.header.frame_id == "map" assert decoded.header.stamp == second.header.stamp assert first.header.frame_id == second.header.frame_id == "lidar" diff --git a/dimos/mapping/voxels/test_stamped.py b/dimos/mapping/voxels/test_stamped.py index 8fc126fef7d..065d282e7ae 100644 --- a/dimos/mapping/voxels/test_stamped.py +++ b/dimos/mapping/voxels/test_stamped.py @@ -12,11 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from dimos.mapping.voxels.grid import VoxelGrid -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_stamps, pointcloud_xyz +from dimos.msgs.time import time_from_seconds VOXEL = 0.05 @@ -25,11 +27,11 @@ def _wall(x, ts): """A 1 m x 1 m wall of voxel centers in the plane at `x`.""" ys, zs = np.meshgrid(np.arange(0.025, 1, VOXEL), np.arange(0.025, 1, VOXEL)) pts = np.stack([np.full(ys.size, x), ys.ravel(), zs.ravel()], axis=1) - return PointCloud2.from_numpy(pts, timestamp=ts) + return pointcloud_from_xyz(pts, header=Header(stamp=time_from_seconds(ts), frame_id="")) def _voxels(grid): - return {tuple(v) for v in np.floor(grid.get_global_pointcloud2().points_f32() / VOXEL)} + return {tuple(v) for v in np.floor(pointcloud_xyz(grid.get_global_pointcloud2()) / VOXEL)} def _drifted_walls(): @@ -41,8 +43,10 @@ def _drifted_walls(): def _undrift(cloud): - shift = np.where(cloud.stamps_f64()[:, None] == 1.0, [0.5, 0.0, 0.0], 0.0) - return PointCloud2.from_numpy(cloud.points_f32() + shift) + shift = np.where(pointcloud_stamps(cloud)[:, None] == 1.0, [0.5, 0.0, 0.0], 0.0) + return pointcloud_from_xyz( + pointcloud_xyz(cloud) + shift, header=Header(stamp=time_from_seconds(0.0), frame_id="") + ) def test_stamped_holds_the_same_voxels(): diff --git a/dimos/memory/cli/render.py b/dimos/memory/cli/render.py index e60664b6456..9673f0ed64d 100644 --- a/dimos/memory/cli/render.py +++ b/dimos/memory/cli/render.py @@ -176,7 +176,7 @@ def render_store( if out is None: src = getattr(store.config, "path", None) or "store" - out = str(Path(src).with_suffix(".rrd")) + out = str(Path(src, poses=[]).with_suffix(".rrd")) base = root.strip("/") if root else "" diff --git a/dimos/memory/cli/test_render.py b/dimos/memory/cli/test_render.py index 1b3f4540ecc..c439336df30 100644 --- a/dimos/memory/cli/test_render.py +++ b/dimos/memory/cli/test_render.py @@ -15,6 +15,7 @@ from pathlib import Path from types import SimpleNamespace +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 from dimos_generated.sensor_msgs.msg import CameraInfo, Image from dimos_generated.std_msgs.msg import Header @@ -40,13 +41,15 @@ def _info(frame_id: str) -> CameraInfo: cy=240.0, width=640, height=480, - header=Header(frame_id=frame_id), + header=Header(frame_id=frame_id, stamp=Time(sec=0, nanosec=0)), ) def _image(frame_id: str) -> Image: return image_from_array( - np.zeros((4, 4, 3), dtype=np.uint8), encoding="bgr8", header=Header(frame_id=frame_id) + np.zeros((4, 4, 3), dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id=frame_id, stamp=Time(sec=0, nanosec=0)), ) @@ -111,7 +114,10 @@ def _tf(child: str, ts: float) -> TFMessage: TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(ts)), child_frame_id=child, - transform=Transform(translation=Vector3(x=1), rotation=Quaternion(w=1)), + transform=Transform( + translation=Vector3(x=1, y=0.0, z=0.0), + rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), ) ] ) diff --git a/dimos/memory/codecs/README.md b/dimos/memory/codecs/README.md index 3864b6bcfb7..921608a3f41 100644 --- a/dimos/memory/codecs/README.md +++ b/dimos/memory/codecs/README.md @@ -1,62 +1,35 @@ -# codecs +# Storage codecs -Encode/decode payloads for persistent storage. Codecs convert typed Python objects to `bytes` and back, used by backends that store observation data as blobs. +Codecs convert Python values to bytes for blob-backed stores such as SQLite. +Generated messages use the same CDR bytes as typed transports. Raw images retain +their pixels and encoding; lossy compression is an explicit conversion to a +`sensor_msgs/msg/CompressedImage` value before storage. -## Protocol +| Codec | Values | Stored representation | +|-------|--------|-----------------------| +| `CdrCodec` | Generated messages | Encapsulated CDR, with no private wrapper | +| `Lz4Codec` | Values accepted by its inner codec | LZ4 frame around the inner bytes | +| `PickleCodec` | Python objects | Python pickle; separate from typed message storage | ```python -class Codec(Protocol[T]): - def encode(self, value: T) -> bytes: ... - def decode(self, data: bytes) -> T: ... -``` - -## Built-in codecs - -| Codec | Type | Notes | -|-------|------|-------| -| `PickleCodec` | Any Python object | Fallback. Uses `HIGHEST_PROTOCOL`. | -| `JpegCodec` | `Image` | Lossy compression via TurboJPEG. ~10-20x smaller. Preserves `frame_id` in header. | -| `LcmCodec` | `DimosMsg` subclasses | Uses `lcm_encode()`/`lcm_decode()`. Zero-copy for LCM message types. | - -## Auto-selection - -`codec_for(payload_type)` picks the right codec: +from dimos_generated.sensor_msgs.msg import Image, CompressedImage +from dimos.memory.codecs.base import codec_for -```python -from dimos.memory.codecs import codec_for - -codec_for(Image) # → JpegCodec(quality=50) -codec_for(Image, sample=depth_image) # → Lz4Codec(LcmCodec(Image)) (DEPTH, DEPTH16, GRAY16) -codec_for(SomeLcmMsg) # → LcmCodec(SomeLcmMsg) (if has lcm_encode/lcm_decode) -codec_for(dict) # → PickleCodec() (fallback) -codec_for(None) # → PickleCodec() +codec_for(Image) # CdrCodec(Image), lossless +codec_for(CompressedImage) # CdrCodec(CompressedImage) +codec_for(dict) # PickleCodec() ``` -`Image` covers depth as well as colour, and JPEG would destroy depth values, so pass a -first payload as `sample` when one is at hand. `--record` opens each stream on its first -message and does this. - -## Writing a new codec - -1. Create `dimos/memory/codecs/mycodec.py`: +SQLite persists the codec ID and Python payload type so reopening a stream uses +the same codec. Use `codec="lz4+cdr"` for lossless blob compression. Old `lcm`, +`lz4+lcm`, and `jpeg` storage codec IDs are rejected; no legacy decoder is retained. -```python -class MyCodec: - def encode(self, value: MyType) -> bytes: - ... - - def decode(self, data: bytes) -> MyType: - ... -``` - -2. Add a branch in `codec_for()` in `base.py` to auto-select it for the relevant type. - -3. Add a test case to `test_codecs.py` — the grid fixture makes this easy: - -```python -@pytest.fixture(params=[..., ("mycodec", MyCodec(), sample_value)]) -def codec_case(request): - ... -``` +MCAP uses `cdr` channels with embedded `ros2msg` definitions and chunk compression. +The reader resolves installed message types by qualified schema name. Unknown +schemas remain inspectable as raw bytes; payload-module metadata never triggers +an arbitrary import. Compressed images remain CompressedImage values when read. +Multiple channels on one topic must agree on schema and timestamp policy. -No base class needed — `Codec` is a protocol. Just implement `encode` and `decode`. +A custom codec only needs `encode(value) -> bytes` and `decode(bytes) -> value`. +Implement the `Codec` protocol and supply the instance explicitly to a stream. +Adding a generated message does not require a new storage codec. diff --git a/dimos/memory/codecs/base.py b/dimos/memory/codecs/base.py index a33667e36f5..a480220fe10 100644 --- a/dimos/memory/codecs/base.py +++ b/dimos/memory/codecs/base.py @@ -17,6 +17,12 @@ import importlib from typing import Any, Protocol, TypeVar, runtime_checkable +from dimos_cdrgen.registry import message_types + +from dimos.memory.codecs.cdr import CdrCodec +from dimos.memory.codecs.lz4 import Lz4Codec +from dimos.memory.codecs.pickle import PickleCodec + T = TypeVar("T") @@ -29,39 +35,14 @@ def decode(self, data: bytes) -> T: ... def codec_for(payload_type: type[Any] | None = None, sample: Any = None) -> Codec[Any]: - """Auto-select codec based on payload type. - - ``sample`` is a first payload, when one is at hand. ``Image`` covers depth as - well as colour, and JPEG would destroy depth values, so a depth or 16-bit - sample selects lossless ``lz4+lcm`` instead. - """ - from dimos.memory.codecs.pickle import PickleCodec - - if payload_type is not None: - from dimos.msgs.sensor_msgs.Image import Image, ImageFormat - - if issubclass(payload_type, Image): - if isinstance(sample, Image) and sample.format in ( - ImageFormat.DEPTH, - ImageFormat.DEPTH16, - ImageFormat.GRAY16, - ): - from dimos.memory.codecs.lcm import LcmCodec - from dimos.memory.codecs.lz4 import Lz4Codec - - return Lz4Codec(LcmCodec(payload_type)) - from dimos.memory.codecs.jpeg import JpegCodec - - return JpegCodec() - if hasattr(payload_type, "lcm_encode") and hasattr(payload_type, "lcm_decode"): - from dimos.memory.codecs.lcm import LcmCodec - - return LcmCodec(payload_type) + """Auto-select codec based on payload type.""" + if payload_type is not None and all(hasattr(payload_type, name) for name in ("__msgtype__",)): + return CdrCodec(payload_type) return PickleCodec() def codec_id(codec: Codec[Any]) -> str: - """Derive a string ID from a codec instance, e.g. ``'lz4+lcm'``. + """Derive a string ID from a codec instance, e.g. ``'lz4+cdr'``. Walks the ``_inner`` chain for wrapper codecs, joining with ``+``. Uses the naming convention ``FooCodec`` → ``'foo'``. @@ -76,7 +57,7 @@ def codec_id(codec: Codec[Any]) -> str: def codec_from_id(codec_id_str: str, payload_module: str) -> Codec[Any]: - """Reconstruct a codec chain from its string ID (e.g. ``'lz4+lcm'``). + """Reconstruct a codec chain from its string ID (e.g. ``'lz4+cdr'``). Builds inside-out: the rightmost segment is the innermost (base) codec. """ @@ -96,6 +77,11 @@ def _class_to_id(codec: Any) -> str: def resolve_payload_type(payload_module: str) -> type[Any]: + if payload_module.startswith("rosbags.usertypes."): + for cls in message_types().values(): + if payload_module == f"{cls.__module__}.{cls.__qualname__}": + return cls + raise ValueError(f"Unknown native message payload type: {payload_module}") parts = payload_module.rsplit(".", 1) if len(parts) != 2: raise ValueError(f"Cannot resolve payload type from {payload_module!r}") @@ -106,19 +92,11 @@ def resolve_payload_type(payload_module: str) -> type[Any]: def _make_one(name: str, payload_module: str, inner: Codec[Any] | None = None) -> Codec[Any]: """Instantiate a single codec by its short name.""" if name == "lz4": - from dimos.memory.codecs.lz4 import Lz4Codec - if inner is None: raise ValueError("lz4 is a wrapper codec — must have an inner codec") return Lz4Codec(inner) - if name == "jpeg": - from dimos.memory.codecs.jpeg import JpegCodec - - return JpegCodec() - if name == "lcm": - from dimos.memory.codecs.lcm import LcmCodec - - return LcmCodec(resolve_payload_type(payload_module)) + if name == "cdr": + return CdrCodec(resolve_payload_type(payload_module)) if name == "json": from dimos.memory.codecs.json import JsonCodec @@ -127,7 +105,13 @@ def _make_one(name: str, payload_module: str, inner: Codec[Any] | None = None) - raise TypeError("JSON codec requires std_msgs.String") return JsonCodec() if name == "pickle": - from dimos.memory.codecs.pickle import PickleCodec - return PickleCodec() + if name in {"lcm", "jpeg"}: + raise ValueError( + f"Legacy recording codec {name!r} is unsupported after the CDR cutover. " + "Create a new CDR recording with the installed generated message package; " + "historical recordings must be exported with their original compatible checkout. " + "No legacy payload fallback or in-place conversion is provided. " + "See docs/development/messages.md." + ) raise ValueError(f"Unknown codec: {name!r}") diff --git a/dimos/memory/codecs/cdr.py b/dimos/memory/codecs/cdr.py index 51d127d983a..5bb14db57ff 100644 --- a/dimos/memory/codecs/cdr.py +++ b/dimos/memory/codecs/cdr.py @@ -14,7 +14,9 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode if TYPE_CHECKING: from dimos.msgs.protocol import DimosMsg @@ -24,9 +26,7 @@ class CdrCodec: """Store a generated message as its transport-neutral CDR bytes.""" def __init__(self, msg_type: type[DimosMsg]) -> None: - if not all( - hasattr(msg_type, member) for member in ("encode", "decode", "msg_name", "schema") - ): + if not all(hasattr(msg_type, member) for member in ("__msgtype__",)): raise TypeError(f"{msg_type!r} is not a generated CDR message type") self._msg_type = msg_type @@ -37,8 +37,8 @@ def payload_type(self) -> type[DimosMsg]: def encode(self, value: DimosMsg) -> bytes: if not isinstance(value, self._msg_type): - raise TypeError(f"Expected {self._msg_type.msg_name}, got {type(value).__name__}") - return value.encode() + raise TypeError(f"Expected {self._msg_type.__msgtype__}, got {type(value).__name__}") + return cdr_encode(value) def decode(self, data: bytes) -> DimosMsg: - return self._msg_type.decode(data) + return cast("DimosMsg", cdr_decode(data, self._msg_type.__msgtype__)) diff --git a/dimos/memory/codecs/json.py b/dimos/memory/codecs/json.py index 96166d68c81..537e4c536f4 100644 --- a/dimos/memory/codecs/json.py +++ b/dimos/memory/codecs/json.py @@ -16,7 +16,7 @@ import json -from dimos.msgs.std_msgs.String import String +from dimos_generated.std_msgs.msg import String def _invalid_constant(value: str) -> None: @@ -31,6 +31,6 @@ def encode(self, value: String) -> bytes: return value.data.encode("utf-8") def decode(self, data: bytes) -> String: - value = String(data.decode("utf-8")) + value = String(data=data.decode("utf-8")) json.loads(value.data, parse_constant=_invalid_constant) return value diff --git a/dimos/memory/codecs/test_codecs.py b/dimos/memory/codecs/test_codecs.py index a5f247b6511..fde098a5d82 100644 --- a/dimos/memory/codecs/test_codecs.py +++ b/dimos/memory/codecs/test_codecs.py @@ -12,229 +12,136 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Grid tests for Codec implementations. - -Runs roundtrip encode→decode tests across every codec, verifying data preservation. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +"""Storage uses generated CDR without altering message fields or image pixels.""" +from dimos_cdrgen.registry import encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import CompressedImage, Image +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest -from dimos.memory.codecs.base import Codec, codec_for, codec_id -from dimos.memory.codecs.jpeg import JpegCodec -from dimos.memory.codecs.lcm import LcmCodec +from dimos.memory.codecs.base import codec_for, codec_from_id, codec_id +from dimos.memory.codecs.cdr import CdrCodec +from dimos.memory.codecs.lz4 import Lz4Codec from dimos.memory.codecs.pickle import PickleCodec -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat - -if TYPE_CHECKING: - from collections.abc import Callable - - from dimos.msgs.protocol import DimosMsg - - -@dataclass -class Case: - name: str - codec: Codec[Any] - values: list[Any] - eq: Callable[[Any, Any], bool] | None = None # custom equality: (original, decoded) -> bool - - -def _lcm_values() -> list[DimosMsg]: - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - from dimos.msgs.geometry_msgs.Quaternion import Quaternion - from dimos.msgs.geometry_msgs.Vector3 import Vector3 - - return [ - PoseStamped( - ts=1.0, - frame_id="map", - position=Vector3(1.0, 2.0, 3.0), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), - ), - PoseStamped(ts=0.5, frame_id="odom"), - ] - - -def _pickle_case() -> Case: - from dimos.memory.codecs.pickle import PickleCodec - - return Case( - name="pickle", - codec=PickleCodec(), - values=[42, "hello", b"raw bytes", {"key": "value"}], +from dimos.memory.store.sqlite import SqliteStore +from dimos.memory.type.observation import Observation +from dimos.msgs.geometry import transform_from_pose +from dimos.msgs.image import image_from_array, image_to_jpeg, image_view + + +@pytest.fixture(params=["pose", "image", "depth", "compressed"]) +def message(request): + header = Header(frame_id="sensor", stamp=Time(sec=1700000000, nanosec=123456789)) + if request.param == "pose": + return PoseStamped( + header=header, + pose=Pose( + position=Point(x=1.25, y=-2.5, z=3.75), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + if request.param == "depth": + return image_from_array( + np.array([[0, 1, 65535]], dtype=np.uint16), encoding="16UC1", header=header + ) + image = image_from_array( + np.full((8, 8, 3), [20, 80, 140], dtype=np.uint8), encoding="rgb8", header=header ) - - -def _lcm_case() -> Case: - from dimos.memory.codecs.lcm import LcmCodec - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - - return Case( - name="lcm", - codec=LcmCodec(PoseStamped), - values=_lcm_values(), - ) - - -def _lz4_pickle_case() -> Case: - from dimos.memory.codecs.lz4 import Lz4Codec - from dimos.memory.codecs.pickle import PickleCodec - - return Case( - name="lz4+pickle", - codec=Lz4Codec(PickleCodec()), - values=[42, "hello", b"raw bytes", {"key": "value"}, list(range(1000))], - ) - - -def _lz4_lcm_case() -> Case: - from dimos.memory.codecs.lcm import LcmCodec - from dimos.memory.codecs.lz4 import Lz4Codec - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - - return Case( - name="lz4+lcm", - codec=Lz4Codec(LcmCodec(PoseStamped)), - values=_lcm_values(), - ) - - -def _jpeg_eq(original: Any, decoded: Any) -> bool: - """JPEG is lossy and normalizes to RGB — check shape, frame_id, RGB tag, and color closeness. - - Compares against ``original.to_rgb()`` because the codec normalizes everything to RGB on - the wire (so a BGR-tagged input comes back RGB-tagged with channels swapped accordingly). - """ - import numpy as np - - if decoded.data.shape != original.data.shape: - return False - if decoded.frame_id != original.frame_id: - return False - if decoded.format != ImageFormat.RGB: - return False - expected = original.to_rgb().data - return bool(np.mean(np.abs(decoded.data.astype(float) - expected.astype(float))) < 5) - - -def _turbojpeg_available() -> bool: - try: - from turbojpeg import TurboJPEG - - TurboJPEG() # fail fast if native lib is missing - except (ImportError, RuntimeError): - return False - return True - - -def _jpeg_case() -> Case | None: - if not _turbojpeg_available(): - return None - - import numpy as np - - # smooth gradients survive lossy jpeg within the eq tolerance - frames = [] - for shift in (0, 90, 180): - arr = np.zeros((48, 64, 3), np.uint8) - arr[..., 0] = np.linspace(0, 255, 64, dtype=np.uint8) - arr[..., 1] = np.linspace(0, 255, 48, dtype=np.uint8)[:, None] - arr[..., 2] = shift - frames.append(Image(data=arr, format=ImageFormat.RGB, frame_id="cam", ts=1.0)) - - return Case( - name="jpeg", - codec=JpegCodec(quality=95), - values=frames, - eq=_jpeg_eq, + if request.param == "compressed": + return CompressedImage( + header=header, + format="rgb8; jpeg compressed bgr8", + data=np.frombuffer(image_to_jpeg(image), dtype=np.uint8), + ) + return image + + +@pytest.mark.parametrize("wrapped", [False, True]) +def test_codec_roundtrip_and_reconstruction_preserve_generated_value(message, wrapped): + codec = Lz4Codec(CdrCodec(type(message))) if wrapped else codec_for(type(message)) + name = codec_id(codec) + assert name == ("lz4+cdr" if wrapped else "cdr") + restored = codec_from_id(name, f"{type(message).__module__}.{type(message).__qualname__}") + encoded = codec.encode(message) + decoded = restored.decode(encoded) + assert cdr_encode(decoded) == cdr_encode(message) + assert decoded.header.stamp.nanosec == 123456789 + if not wrapped: + assert encoded == cdr_encode(message) + if isinstance(message, Image): + np.testing.assert_array_equal(image_view(decoded), image_view(message)) + + +@pytest.mark.parametrize("wrapped", [False, True]) +def test_sqlite_reopen_uses_cdr_and_preserves_exact_message(message, wrapped, tmp_path): + path = tmp_path / "messages.db" + with SqliteStore(path=str(path)) as store: + stream = store.stream("samples", type(message), codec="lz4+cdr" if wrapped else "cdr") + stream.append(message, ts=12.5) + with SqliteStore(path=str(path), must_exist=True) as store: + observation = store.stream("samples").first() + assert observation.ts == 12.5 + assert cdr_encode(observation.data) == cdr_encode(message) + assert cdr_encode(observation.data) == cdr_encode(message) + + +@pytest.mark.parametrize("identifier", ["lcm", "lz4+lcm", "jpeg"]) +def test_obsolete_storage_codecs_are_rejected(identifier): + with pytest.raises(ValueError, match="Legacy recording codec.*unsupported after the CDR"): + codec_from_id(identifier, "builtins.dict") + + +def test_cdr_requires_a_generated_message_type(): + with pytest.raises(TypeError, match="not a generated CDR"): + codec_from_id("cdr", "builtins.dict") + + +def test_cdr_malformed_payload_raises(): + with pytest.raises((ValueError, RuntimeError)): + CdrCodec(PoseStamped).decode(b"bad") + + +@pytest.mark.parametrize("value", [42, "hello", b"raw bytes", {"key": "value"}]) +def test_python_objects_keep_the_explicit_python_storage_path(value): + codec = codec_for(type(value)) + assert isinstance(codec, PickleCodec) + assert codec.decode(codec.encode(value)) == value + + +@pytest.mark.parametrize("stamped_transform", [False, True]) +def test_observation_pose_metadata_accepts_nested_generated_values(stamped_transform): + pose = PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=1, y=2, z=3), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), ) - - -_case_factories = { - "pickle": _pickle_case, - "lcm": _lcm_case, - "lz4+pickle": _lz4_pickle_case, - "lz4+lcm": _lz4_lcm_case, - "jpeg": _jpeg_case, -} - -case_params: list[Any] = ["pickle", "lcm", "lz4+pickle", "lz4+lcm"] -if _turbojpeg_available(): - case_params.append("jpeg") - - -@pytest.fixture -def case(request: pytest.FixtureRequest) -> Case: - resolved = _case_factories[request.param]() - if resolved is None: - pytest.skip(f"no usable data for the {request.param} case") - return resolved - - -@pytest.mark.parametrize("case", case_params, indirect=True) -class TestCodecRoundtrip: - """Every codec must perfectly roundtrip its values.""" - - def test_roundtrip_preserves_value(self, case: Case) -> None: - eq = case.eq or (lambda a, b: a == b) - for value in case.values: - encoded = case.codec.encode(value) - assert isinstance(encoded, bytes) - decoded = case.codec.decode(encoded) - assert eq(value, decoded), f"Roundtrip failed for {value!r}: got {decoded!r}" - - def test_encode_returns_nonempty_bytes(self, case: Case) -> None: - for value in case.values: - encoded = case.codec.encode(value) - assert len(encoded) > 0, f"Empty encoding for {value!r}" - - def test_different_values_produce_different_bytes(self, case: Case) -> None: - encodings = [case.codec.encode(v) for v in case.values] - assert len(set(encodings)) > 1, "All values encoded to identical bytes" - - -class TestCodecFor: - """codec_for() auto-selects the right codec.""" - - def test_none_returns_pickle(self) -> None: - assert isinstance(codec_for(None), PickleCodec) - - def test_unknown_type_returns_pickle(self) -> None: - assert isinstance(codec_for(dict), PickleCodec) - - def test_lcm_type_returns_lcm(self) -> None: - assert isinstance(codec_for(PoseStamped), LcmCodec) - - def test_image_type_returns_jpeg(self) -> None: - pytest.importorskip("turbojpeg") - from dimos.memory.codecs.jpeg import JpegCodec - - assert isinstance(codec_for(Image), JpegCodec) - - def test_colour_image_sample_returns_jpeg(self) -> None: - pytest.importorskip("turbojpeg") - rgb = Image(data=np.zeros((48, 64, 3), dtype=np.uint8), format=ImageFormat.RGB) - - assert isinstance(codec_for(Image, sample=rgb), JpegCodec) - - @pytest.mark.parametrize("fmt", [ImageFormat.DEPTH, ImageFormat.DEPTH16, ImageFormat.GRAY16]) - def test_depth_image_sample_returns_lossless(self, fmt: ImageFormat) -> None: - """JPEG is 8-bit colour: uint16 depth comes back as wrong metres, float32 does not encode.""" - metres = np.linspace(0.5, 6.0, 48 * 64, dtype=np.float32).reshape(48, 64) - data = metres if fmt is ImageFormat.DEPTH else (metres * 1000).astype(np.uint16) - depth = Image(data=data, format=fmt, frame_id="depth_optical", ts=1.0) - - codec = codec_for(Image, sample=depth) - decoded = codec.decode(codec.encode(depth)) - - assert codec_id(codec) == "lz4+lcm" - assert decoded.format == fmt - assert decoded.frame_id == "depth_optical" - np.testing.assert_array_equal(decoded.data, depth.data) + value = transform_from_pose(pose, child_frame_id="base") if stamped_transform else pose + observation = Observation(id=0, ts=-0.5, pose=value, _data="payload") + assert observation.pose_tuple == (1, 2, 3, 0, 0, 0, 1) + assert observation.pose == pose.pose + assert observation.pose_stamped.pose == pose.pose + assert observation.pose_stamped.header.stamp == Time(sec=-1, nanosec=500000000) + + +def test_cdr_rejects_a_message_of_the_wrong_type(): + with pytest.raises(TypeError, match="Expected geometry_msgs/msg/PoseStamped"): + CdrCodec(PoseStamped).encode( + Image( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + height=0, + width=0, + encoding="", + is_bigendian=0, + step=0, + data=np.array([], dtype=np.uint8), + ) + ) + + +def test_unknown_storage_codec_is_not_misdiagnosed_as_historical_recording(): + with pytest.raises(ValueError, match="Unknown codec: 'misspelled'"): + codec_from_id("misspelled", "builtins.dict") diff --git a/dimos/memory/demo_data.py b/dimos/memory/demo_data.py index d1abd2a84f6..551e6067410 100644 --- a/dimos/memory/demo_data.py +++ b/dimos/memory/demo_data.py @@ -52,8 +52,8 @@ def write_demo_recording(path: Path, *, samples: int = 720) -> SqliteStore: header = Header(stamp=stamp, frame_id="world") angle = index / 80.0 pose = Pose( - position=Point(x=math.cos(angle), y=math.sin(angle)), - orientation=Quaternion(w=1.0), + position=Point(x=math.cos(angle), y=math.sin(angle), z=0.0), + orientation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), ) pixels = np.full((48, 64, 3), 40 + index % 180, dtype=np.uint8) pixels[:, 16:32, 1] = 230 diff --git a/dimos/memory/module.py b/dimos/memory/module.py index a5867f0c53c..f30a3e9ea5a 100644 --- a/dimos/memory/module.py +++ b/dimos/memory/module.py @@ -23,6 +23,7 @@ import time from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast +from dimos_generated.tf2_msgs.msg import TFMessage from pydantic import Field, field_validator from reactivex import operators as ops from reactivex.disposable import Disposable @@ -35,11 +36,13 @@ from dimos.memory.store.sqlite import SqliteStore from dimos.memory.stream import Stream from dimos.memory.type.recording import OnExisting -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.geometry import pose_from_transform +from dimos.msgs.time import message_header, to_seconds from dimos.utils.data import backup_file from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: + from dimos_generated.geometry_msgs.msg import Pose from reactivex.abc import DisposableBase from dimos.core.stream import Out @@ -47,7 +50,6 @@ SemanticSearch as SemanticSearch, SemanticSearchConfig as SemanticSearchConfig, ) - from dimos.msgs.geometry_msgs.Pose import Pose logger = setup_logger() @@ -222,7 +224,7 @@ class RecorderConfig(MemoryModuleConfig): # read the active remappings from inside the module (AFAIK), so this config # arg does the per-stream rename directly. stream_remapping: dict[str, str] = Field(default_factory=dict) - # ex: {"depth_image_1": "lz4+lcm"} for lossless depth recording + # ex: {"depth_image_1": "lz4+cdr"} for lossless depth recording stream_codecs: dict[str, str] = Field(default_factory=dict) # Port names that inherently have no pose to anchor (command streams, etc.). poseless_streams: list[str] = Field(default_factory=list) @@ -334,7 +336,7 @@ def _data_ports(self) -> dict[str, In[Any]]: def _port_to_stream(self, name: str, input_topic: In[Any], stream: Stream[Any]) -> None: """Append each message from *input_topic* to *stream*, attaching world pose via tf. - Stamped messages use their own ``.frame_id`` and ``.ts``; unstamped + Stamped messages use their own ``header.frame_id`` and ``header.stamp``; unstamped messages (or ones whose frame isn't in the tf graph, e.g. a payload already in world coords) fall back to ``config.default_frame_id`` — so every observation gets a robot-pose anchor when tf is publishing. @@ -375,7 +377,8 @@ def _prepare_streams(self) -> None: def _resolve_ts(self, name: str, msg: Any) -> float: """Timestamp to record *msg* at. Override to re-base onto another clock.""" - return getattr(msg, "ts", None) or time.time() + header = message_header(msg) + return to_seconds(header.stamp) if header is not None else time.time() async def _resolve_pose(self, name: str, msg: Any, ts: float) -> Pose | None: """Pose to anchor *msg* with. @@ -391,11 +394,12 @@ async def _resolve_pose(self, name: str, msg: Any, ts: float) -> Pose | None: return cast("Pose | None", await setter(msg)) if self._tf is None: return None - frame_id = getattr(msg, "frame_id", None) or self.config.default_frame_id + header = message_header(msg) + frame_id = (header.frame_id if header is not None else "") or self.config.default_frame_id transform = self._tf.get( self.config.root_frame, frame_id, time_point=ts, time_tolerance=self.config.tf_tolerance ) - return transform.to_pose() if transform is not None else None + return pose_from_transform(transform).pose if transform is not None else None def _collect_pose_setters(self) -> dict[str, PoseSetter]: """Map stream name -> bound ``@pose_setter_for`` method.""" @@ -416,7 +420,11 @@ def _record_tf(self) -> None: def on_tf(msg: TFMessage) -> None: try: for transform in msg.transforms: - tf_stream.append(TFMessage(transform), ts=transform.ts, pose=None) + tf_stream.append( + TFMessage(transforms=[transform]), + ts=to_seconds(transform.header.stamp), + pose=None, + ) except sqlite3.ProgrammingError: # A late LCM callback raced teardown and hit the closed store. pass diff --git a/dimos/memory/semantic_search.py b/dimos/memory/semantic_search.py index 520b17d3c34..d7b32498864 100644 --- a/dimos/memory/semantic_search.py +++ b/dimos/memory/semantic_search.py @@ -16,6 +16,9 @@ from typing import Any, cast +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import Image + from dimos.agents.annotation import skill from dimos.core.core import rpc from dimos.memory.embed import EmbedImages @@ -24,8 +27,7 @@ from dimos.memory.transform import QualityWindow, peaks from dimos.memory.type.observation import EmbeddedObservation, Observation from dimos.models.embedding.base import EmbeddingModel -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_brightness, image_sharpness class SemanticSearchConfig(MemoryModuleConfig): @@ -55,8 +57,8 @@ def start(self) -> None: # fmt: off self.store.streams.color_image \ .live() \ - .filter(lambda obs: obs.data.brightness > 0.1) \ - .transform(QualityWindow(lambda img: img.sharpness, window=0.5)) \ + .filter(lambda obs: image_brightness(obs.data) > 0.1) \ + .transform(QualityWindow(image_sharpness, window=0.5)) \ .transform(EmbedImages(self.model, batch_size=2)) \ .save(self.embeddings) \ .drain_thread() diff --git a/dimos/memory/store/mcap.py b/dimos/memory/store/mcap.py index 1cf84c123fb..0283a1ce79c 100644 --- a/dimos/memory/store/mcap.py +++ b/dimos/memory/store/mcap.py @@ -14,15 +14,15 @@ """Read-only memory store backed by an mcap file. -Generic and robot-independent. JPEG channels decode automatically because their -payload type is fixed. Other formats use a caller-supplied ``codecs`` map (wire +Generic and robot-independent. CDR channels resolve installed message types by +their ROS schema name. Other formats use a caller-supplied ``codecs`` map (wire topic -> codec), while ``streams`` may map friendly stream names to topics. See ``dimos.robot.unitree.go2.dds.store.Go2McapStore`` for the Go2 DDS wiring. Channels published by dimos itself need neither: a dimos topic carries its -message type in its own name, so it is decoded and named from the wire with no -registry at all (``_dimos_wire``). An injected codec still wins where both -apply. +message type in its own name, which supplies a friendly stream name +(``_dimos_wire``). Decoding still requires a declared CDR schema or an injected +codec; a type-like topic never implies the byte encoding. Read-only: no append, blobs, vectors, or embeddings. Payloads decode lazily on ``obs.data``; ts and counts are cheap (counts come from the mcap index). @@ -31,20 +31,22 @@ from __future__ import annotations from collections.abc import Iterator, Mapping -from dataclasses import dataclass, replace +from dataclasses import replace from functools import partial -from typing import Any, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable from dimos.memory.backend import Backend from dimos.memory.codecs.base import codec_for -from dimos.memory.codecs.jpeg import JpegCodec -from dimos.memory.codecs.json import JsonCodec +from dimos.memory.codecs.cdr import CdrCodec from dimos.memory.notifier.subject import SubjectNotifier from dimos.memory.observationstore.base import ObservationStore, ObservationStoreConfig from dimos.memory.store.base import Store, StoreConfig from dimos.memory.type.filter import StreamQuery from dimos.memory.type.observation import Observation -from dimos.protocol.pubsub.impl.rospubsub_conversion import get_dimos_type +from dimos.msgs.helpers import resolve_msg_type + +if TYPE_CHECKING: + from dimos.msgs.protocol import DimosMsg @runtime_checkable @@ -69,31 +71,21 @@ def decode(self, data: bytes) -> bytes: _BYTES_CODEC = _BytesCodec() -@dataclass(frozen=True) -class _DimosCodec: - """Decodes a dimos wire channel through its own message type.""" - - payload_type: type - - def decode(self, data: bytes) -> Any: - return self.payload_type.lcm_decode(data) # type: ignore[attr-defined] - - -def _dimos_wire(topic: str) -> tuple[str, type] | None: +def _dimos_wire(topic: str) -> tuple[str, type[DimosMsg]] | None: """Split a dimos wire topic into its port name and message type. A dimos topic embeds the type in its last segment, - ``dimos//`` (built by ``zenohpubsub.Topic.key_expr``), so + ``dimos///msg/`` (built by ``zenohpubsub.Topic.key_expr``), so every dimos channel describes itself and no per-topic registry is needed. None for anything that is not one, which leaves DDS and app channels to the injected codecs. """ - prefix, _, rest = topic.partition("/") - port, _, msg_name = rest.rpartition("/") - if prefix != "dimos" or not port: + parts = topic.split("/") + if len(parts) < 5 or parts[0] != "dimos" or parts[-2] != "msg": return None - msg_type = get_dimos_type(msg_name) - return (port, msg_type) if msg_type is not None else None + port = "/".join(parts[1:-3]) + msg_type = resolve_msg_type("/".join(parts[-3:])) + return (port, msg_type) if port and msg_type is not None else None def _slug(topic: str) -> str: @@ -102,7 +94,7 @@ def _slug(topic: str) -> str: ``rt/`` is the ROS2-over-DDS topic prefix; ``removeprefix`` only strips it where present (e.g. app-level ``control_log`` is left alone). A dimos topic names itself by its port, dropping the type segment: the dot in - ``nav_msgs.Path`` is not attribute-addressable on ``store.streams``. + ``nav_msgs/msg/Path`` is not attribute-addressable on ``store.streams``. """ wire = _dimos_wire(topic) if wire is not None: @@ -226,6 +218,7 @@ def __init__( # _BYTES_CODEC — reachable but undecoded. _raw maps their stream name to the # source schema so summary() can flag them [raw bytes: ]. self._raw: dict[str, str | None] = {} # raw stream name -> source schema + topic_signatures: dict[str, tuple[Any, ...]] = {} if summary is not None and summary.statistics is not None: for cid, ch in summary.channels.items(): count = summary.statistics.channel_message_counts.get(cid, 0) @@ -238,27 +231,39 @@ def __init__( # `streams=`, not something to guess at. wire, had = _dimos_wire(ch.topic), _dimos_wire(taken) if wire is not None and had is not None and wire[0] == had[0]: - name = f"{name}_{wire[1].__name__}" + type_name = wire[1].__msgtype__.rsplit("/", 1)[-1] + name = f"{name}_{type_name}" if self._stream_topic.get(name, ch.topic) != ch.topic: raise ValueError( f"stream {name!r} would name both {taken!r} and {ch.topic!r};" " pass streams= to name them apart" ) - if ch.topic not in self._codecs and ch.message_encoding == "jpeg": - self._codecs[ch.topic] = JpegCodec() - if ch.topic not in self._codecs and ch.message_encoding == "json": - kind = JsonCodec.payload_type - if ( - ch.metadata.get("dimos.payload_type") - == f"{kind.__module__}.{kind.__qualname__}" - ): - self._codecs[ch.topic] = JsonCodec() + schema = summary.schemas.get(ch.schema_id) + signature = ( + ch.message_encoding, + (schema.name, schema.encoding, schema.data) if schema else None, + ch.metadata.get("dimos.observation_time"), + ) + if ch.topic in topic_signatures and topic_signatures[ch.topic] != signature: + raise ValueError( + f"Topic {ch.topic!r} has conflicting channel schemas or timing policies" + ) + topic_signatures[ch.topic] = signature + if ( + ch.topic not in self._codecs + and ch.message_encoding == "cdr" + and schema is not None + and schema.encoding == "ros2msg" + ): + message_type = resolve_msg_type(schema.name) + if message_type is not None: + self._codecs[ch.topic] = CdrCodec(message_type) self._stream_topic[name] = ch.topic - self._available[name] = count + self._available[name] = self._available.get(name, 0) + count self._observation_uses_publish_time[name] = ( ch.metadata.get("dimos.observation_time") == "publish_time" ) - if ch.topic not in self._codecs and _dimos_wire(ch.topic) is None: + if ch.topic not in self._codecs: sch = summary.schemas.get(ch.schema_id) self._raw[name] = sch.name if sch else None @@ -282,10 +287,9 @@ def _create_backend( if name not in self._available: raise KeyError(f"No stream {name!r}. Available: {sorted(self._available)}") topic = self._stream_topic[name] - # Injected codecs win; a dimos channel decodes itself; the rest stay bytes. - wire = _dimos_wire(topic) - fallback: StreamCodec = _DimosCodec(wire[1]) if wire is not None else _BYTES_CODEC - codec = self._codecs.get(topic) or fallback + # Injected codecs win; known CDR schemas were resolved from channel metadata. + # A type-like topic alone must never select a decoder for arbitrary bytes. + codec = self._codecs.get(topic) or _BYTES_CODEC ptype = codec.payload_type obs = McapObservationStore( name=name, diff --git a/dimos/memory/store/test_mcap.py b/dimos/memory/store/test_mcap.py index 788682ed21c..6bf674b51c9 100644 --- a/dimos/memory/store/test_mcap.py +++ b/dimos/memory/store/test_mcap.py @@ -12,8 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_cdrgen.registry import encode as cdr_encode, schema as cdr_schema from dimos_generated.builtin_interfaces.msg import Time -from dimos_generated.geometry_msgs.msg import Vector3 +from dimos_generated.geometry_msgs.msg import Quaternion, Vector3 from dimos_generated.sensor_msgs.msg import CompressedImage, Imu from dimos_generated.std_msgs.msg import Header from mcap.writer import Writer @@ -33,13 +34,26 @@ def test_cdr_channel_decodes_and_orders_by_declared_source_time(tmp_path, explic expected = Imu( header=Header(stamp=Time(sec=12, nanosec=500000000), frame_id="imu_link"), angular_velocity=Vector3(x=1, y=2, z=3), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + orientation_covariance=np.zeros(9, dtype=np.float64), + angular_velocity_covariance=np.zeros(9, dtype=np.float64), + linear_acceleration=Vector3(x=0.0, y=0.0, z=0.0), + linear_acceleration_covariance=np.zeros(9, dtype=np.float64), + ) + earlier = Imu( + header=Header(stamp=Time(sec=11, nanosec=500000000), frame_id="earlier"), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + orientation_covariance=np.zeros(9, dtype=np.float64), + angular_velocity=Vector3(x=0.0, y=0.0, z=0.0), + angular_velocity_covariance=np.zeros(9, dtype=np.float64), + linear_acceleration=Vector3(x=0.0, y=0.0, z=0.0), + linear_acceleration_covariance=np.zeros(9, dtype=np.float64), ) - earlier = Imu(header=Header(stamp=Time(sec=11, nanosec=500000000), frame_id="earlier")) with path.open("wb") as output: writer = Writer(output) writer.start(profile="ros2", library="test") schema = writer.register_schema( - name=Imu.msg_name, encoding="ros2msg", data=Imu.schema.encode() + name=Imu.__msgtype__, encoding="ros2msg", data=cdr_schema(Imu.__msgtype__).encode() ) channel = writer.register_channel( topic="imu", @@ -52,15 +66,20 @@ def test_cdr_channel_decodes_and_orders_by_declared_source_time(tmp_path, explic channel_id=channel, log_time=13000000000 + index * 1000000000, publish_time=stamp, - data=message.encode(), + data=cdr_encode(message), ) writer.finish() with McapStore(path=str(path), codecs={"imu": CdrCodec(Imu)} if explicit else None) as store: assert store.list_streams() == ["imu"] observations = list(store.stream("imu").order_by("ts")) assert [observation.ts for observation in observations] == [11.5, 12.5] - assert [observation.data for observation in observations] == [earlier, expected] - assert store.stream("imu").order_by("ts", desc=True).first().data == expected + assert [cdr_encode(observation.data) for observation in observations] == [ + cdr_encode(earlier), + cdr_encode(expected), + ] + assert cdr_encode(store.stream("imu").order_by("ts", desc=True).first().data) == cdr_encode( + expected + ) def test_standard_compressed_image_uses_cdr_and_generic_dataset_dispatch(tmp_path): @@ -71,21 +90,23 @@ def test_standard_compressed_image_uses_cdr_and_generic_dataset_dispatch(tmp_pat header=Header(frame_id="camera", stamp=Time(sec=12, nanosec=123456789)), ) expected = CompressedImage( - header=image.header, format="rgb8; jpeg compressed bgr8", data=image_to_jpeg(image) + header=image.header, + format="rgb8; jpeg compressed bgr8", + data=np.frombuffer(image_to_jpeg(image), dtype=np.uint8), ) with CdrMcapWriter(path) as writer: writer.write( "color_image", - expected.encode(), - schema_name=expected.msg_name, - schema=expected.schema, + cdr_encode(expected), + schema_name=expected.__msgtype__, + schema=cdr_schema(expected.__msgtype__), log_time_ns=13000000000, publish_time_ns=12123456789, ) with open_store(path) as store: observation = store.stream("color_image").first() assert observation.ts == 13.0 - assert observation.data == expected + assert cdr_encode(observation.data) == cdr_encode(expected) assert observation.data.header.stamp.nanosec == 123456789 assert observation.data.format == "rgb8; jpeg compressed bgr8" @@ -121,22 +142,42 @@ def test_duplicate_topic_channels_count_together_or_reject_conflicting_schemas( tmp_path, second_type ): path = tmp_path / "channels.mcap" + header = Header(Time(0, 0), "") + imu = Imu( + header, + Quaternion(0.0, 0.0, 0.0, 1.0), + np.zeros(9), + Vector3(0.0, 0.0, 0.0), + np.zeros(9), + Vector3(0.0, 0.0, 0.0), + np.zeros(9), + ) + second = ( + imu if second_type is Imu else CompressedImage(header, "jpeg", np.array([], dtype=np.uint8)) + ) with path.open("wb") as output: writer = Writer(output) writer.start(profile="ros2") - for cls in [Imu, second_type]: + for message in [imu, second]: schema = writer.register_schema( - name=cls.msg_name, encoding="ros2msg", data=cls.schema.encode() + name=message.__msgtype__, + encoding="ros2msg", + data=cdr_schema(message.__msgtype__).encode(), ) channel = writer.register_channel( topic="sensor", message_encoding="cdr", schema_id=schema ) - writer.add_message(channel_id=channel, log_time=1, publish_time=1, data=cls().encode()) + writer.add_message( + channel_id=channel, log_time=1, publish_time=1, data=cdr_encode(message) + ) writer.finish() if second_type is Imu: with McapStore(path=str(path)) as store: assert store.stream("sensor").count() == 2 - assert [observation.data for observation in store.stream("sensor")] == [Imu(), Imu()] + assert [cdr_encode(observation.data) for observation in store.stream("sensor")] == [ + cdr_encode(imu), + cdr_encode(imu), + ] else: with pytest.raises(ValueError, match="conflicting channel schemas"): McapStore(path=str(path)) diff --git a/dimos/memory/store/test_mcap_wire.py b/dimos/memory/store/test_mcap_wire.py index e116afccc79..6958a8cfbf8 100644 --- a/dimos/memory/store/test_mcap_wire.py +++ b/dimos/memory/store/test_mcap_wire.py @@ -14,18 +14,22 @@ """The self-describing dimos wire channel: naming and decode, no file needed.""" +from dimos_cdrgen.registry import encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header +from mcap.writer import Writer import pytest -from dimos.memory.store.mcap import McapStore, _dimos_wire, _DimosCodec, _slug -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.memory.codecs.cdr import CdrCodec +from dimos.memory.store.mcap import McapStore, _dimos_wire, _slug def test_a_dimos_topic_names_its_own_type(): - assert _dimos_wire("dimos/planner_path/nav_msgs.Path") == ("planner_path", Path) - assert _dimos_wire("dimos/local_map/sensor_msgs.PointCloud2") == ( + assert _dimos_wire("dimos/planner_path/nav_msgs/msg/Path") == ("planner_path", Path) + assert _dimos_wire("dimos/local_map/sensor_msgs/msg/PointCloud2") == ( "local_map", PointCloud2, ) @@ -37,30 +41,38 @@ def test_everything_else_stays_the_injected_codecs_business(): assert _dimos_wire("rt/utlidar/cloud") is None assert _dimos_wire("control_log") is None assert _dimos_wire("dimos/mystery/not_a_package.Nope") is None - assert _dimos_wire("dimos/nav_msgs.Path") is None + assert _dimos_wire("dimos/nav_msgs/msg/Path") is None def test_the_stream_name_is_the_port_not_the_wire_name(): - # the type segment is dropped: "dimos_planner_path_nav_msgs.Path" carries a + # the type segment is dropped: "dimos_planner_path_nav_msgs/msg/Path" carries a # dot, so it is not addressable as store.streams. - assert _slug("dimos/planner_path/nav_msgs.Path") == "planner_path" - assert "." not in _slug("dimos/tf/tf2_msgs.TFMessage") + assert _slug("dimos/planner_path/nav_msgs/msg/Path") == "planner_path" + assert "." not in _slug("dimos/tf/tf2_msgs/msg/TFMessage") # and the rt/ and app cases are untouched assert _slug("rt/utlidar/cloud") == "utlidar_cloud" assert _slug("control_log") == "control_log" def test_the_codec_decodes_the_wire_bytes(): + header = Header(stamp=Time(sec=1, nanosec=0), frame_id="odom") path = Path( - ts=1.0, - frame_id="odom", - poses=[PoseStamped(ts=1.0, frame_id="odom", position=Vector3(0.5, 0.0, 0.0))], + header=header, + poses=[ + PoseStamped( + header=header, + pose=Pose( + position=Point(x=0.5, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ], ) - _, msg_type = _dimos_wire("dimos/path/nav_msgs.Path") - got = _DimosCodec(msg_type).decode(path.lcm_encode()) + _, msg_type = _dimos_wire("dimos/path/nav_msgs/msg/Path") + got = CdrCodec(msg_type).decode(cdr_encode(path)) assert isinstance(got, Path) - assert got.frame_id == "odom" - assert abs(got.poses[0].position.x - 0.5) < 1e-9 + assert got.header.frame_id == "odom" + assert abs(got.poses[0].pose.position.x - 0.5) < 1e-9 def _write(path, topics): @@ -70,22 +82,37 @@ def _write(path, topics): writer = mcap_writer.Writer(output) writer.start(profile="dimos", library="test") for topic in topics: - writer.register_channel(topic=topic, message_encoding="lcm", schema_id=0) + writer.register_channel(topic=topic, message_encoding="opaque", schema_id=0) writer.finish() def test_two_types_on_one_port_stay_two_streams(tmp_path): path = tmp_path / "shared.mcap" - _write(path, ["dimos/shared/sensor_msgs.Imu", "dimos/shared/geometry_msgs.Vector3"]) + _write(path, ["dimos/shared/sensor_msgs/msg/Imu", "dimos/shared/geometry_msgs/msg/Vector3"]) with McapStore(path=str(path)) as store: assert store.list_streams() == ["shared", "shared_Vector3"] def test_a_real_collision_is_refused_not_overwritten(tmp_path): path = tmp_path / "slash.mcap" - _write(path, ["dimos/slash/name/nav_msgs.Path", "dimos/slash_name/nav_msgs.Path"]) + _write(path, ["dimos/slash/name/nav_msgs/msg/Path", "dimos/slash_name/nav_msgs/msg/Path"]) with pytest.raises(ValueError, match="slash_name"): McapStore(path=str(path)) # and the way out is to name them yourself - with McapStore(path=str(path), streams={"a": "dimos/slash/name/nav_msgs.Path"}) as store: + with McapStore(path=str(path), streams={"a": "dimos/slash/name/nav_msgs/msg/Path"}) as store: assert "a" in store.list_streams() + + +def test_type_like_topic_does_not_imply_cdr_encoding(tmp_path): + path = tmp_path / "opaque.mcap" + with path.open("wb") as output: + writer = Writer(output) + writer.start() + channel = writer.register_channel( + topic="dimos/path/nav_msgs/msg/Path", message_encoding="opaque", schema_id=0 + ) + writer.add_message(channel_id=channel, log_time=1, publish_time=1, data=b"not CDR") + writer.finish() + with McapStore(path=str(path)) as store: + assert store.stream("path").first().data == b"not CDR" + assert "raw bytes" in store.summary() diff --git a/dimos/memory/tap.py b/dimos/memory/tap.py index 2ce45b57933..30486de97ef 100644 --- a/dimos/memory/tap.py +++ b/dimos/memory/tap.py @@ -33,10 +33,13 @@ import time from typing import TYPE_CHECKING, Any +from dimos_generated.std_msgs.msg import Header + from dimos.constants import RECORDINGS_DIR from dimos.core.global_config import global_config from dimos.memory.codecs.base import codec_for from dimos.memory.store.sqlite import SqliteStore +from dimos.msgs.time import message_header, to_seconds from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -76,8 +79,8 @@ class TransportRecorder: work, so a slow append never stalls a transport's handler thread. When the queue is full the message is dropped and counted. - A stream is opened on its first message, so its codec can depend on what the - stream carries: depth images are stored losslessly, colour images as JPEG. + A stream is opened on its first message. Generated messages, including image + pixels and source headers, are stored losslessly as CDR. """ def __init__(self, store: SqliteStore, topics: str = "*", queue_size: int = 1000) -> None: @@ -115,14 +118,15 @@ def tap( """Subscribe *transport* and record into stream *name*; returns the unsubscribe.""" if not matching(self._topics, [name]): return None - if not hasattr(stream_type, "lcm_encode"): + if not all(hasattr(stream_type, member) for member in ("__msgtype__",)): logger.info("--record: %s (%s) is not a dimos message type, skipped", name, stream_type) return None def on_msg(msg: Any) -> None: try: - ts = getattr(msg, "ts", None) or time.time() - self._queue.put_nowait((name, stream_type, msg, ts)) + header = message_header(msg) + timestamp = to_seconds(header.stamp) if isinstance(header, Header) else time.time() + self._queue.put_nowait((name, stream_type, msg, timestamp)) except queue.Full: self.dropped += 1 if self.dropped % 1000 == 1: diff --git a/dimos/memory/test_replay_module.py b/dimos/memory/test_replay_module.py index 356a335672a..ab67a0a4f7c 100644 --- a/dimos/memory/test_replay_module.py +++ b/dimos/memory/test_replay_module.py @@ -16,7 +16,7 @@ import time from typing import Any -from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.std_msgs.msg import Header import pytest @@ -34,8 +34,26 @@ def recording(tmp_path: Path) -> str: odom = store.stream("odom", PoseStamped) goal = store.stream("goal", PoseStamped) for ts in (1.0, 1.1, 1.2): - odom.append(PoseStamped(header=Header(stamp=time_from_seconds(ts))), ts=ts) - goal.append(PoseStamped(header=Header(stamp=time_from_seconds(1.05))), ts=1.05) + odom.append( + PoseStamped( + header=Header(stamp=time_from_seconds(ts), frame_id=""), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + ts=ts, + ) + goal.append( + PoseStamped( + header=Header(stamp=time_from_seconds(1.05), frame_id=""), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + ts=1.05, + ) store.stop() return str(path) @@ -94,10 +112,24 @@ def test_stream_named_like_a_module_attribute_is_skipped( store = SqliteStore(path=str(path)) store.start() store.stream("start", PoseStamped).append( - PoseStamped(header=Header(stamp=time_from_seconds(1.0))), ts=1.0 + PoseStamped( + header=Header(stamp=time_from_seconds(1.0), frame_id=""), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + ts=1.0, ) store.stream("odom", PoseStamped).append( - PoseStamped(header=Header(stamp=time_from_seconds(1.0))), ts=1.0 + PoseStamped( + header=Header(stamp=time_from_seconds(1.0), frame_id=""), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + ts=1.0, ) store.stop() assert list(replay_module(str(path)).__annotations__) == ["odom"] @@ -114,9 +146,25 @@ def test_single_timestamp_stream_is_republished_and_not_the_anchor(tmp_path: Pat store.start() info = store.stream("camera_info", PoseStamped) for _ in range(3): - info.append(PoseStamped(header=Header(stamp=time_from_seconds(1.0))), ts=1.0) + info.append( + PoseStamped( + header=Header(stamp=time_from_seconds(1.0), frame_id=""), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + ts=1.0, + ) store.stream("odom", PoseStamped).append( - PoseStamped(header=Header(stamp=time_from_seconds(40.0))), ts=40.0 + PoseStamped( + header=Header(stamp=time_from_seconds(40.0), frame_id=""), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + ts=40.0, ) store.stop() module = replay_module(str(path))(dataset=str(path)) diff --git a/dimos/memory/test_tap.py b/dimos/memory/test_tap.py index de73151d778..ddbb0cb3a0a 100644 --- a/dimos/memory/test_tap.py +++ b/dimos/memory/test_tap.py @@ -16,6 +16,11 @@ from pathlib import Path from typing import Any +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.dimos_msgs.msg import RegionPointCloud2 +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Twist, Vector3 +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from pytest_mock import MockerFixture @@ -25,8 +30,9 @@ from dimos.memory import tap from dimos.memory.store.sqlite import SqliteStore from dimos.memory.tap import TransportRecorder -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_from_array, image_view +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import time_from_seconds class _Transport(Transport[Any]): @@ -59,12 +65,44 @@ def test_taps_matching_dimos_streams(tmp_path: Path) -> None: unsub = rec.tap("odom", PoseStamped, odom) assert rec.tap("goal", PoseStamped, goal) is None # filtered out assert rec.tap("lidar", dict, raw) is None # not a dimos message type - odom.publish(PoseStamped(ts=1.0)) - odom.publish(PoseStamped(ts=2.0)) - goal.publish(PoseStamped(ts=3.0)) + odom.publish( + PoseStamped( + header=Header(stamp=time_from_seconds(1.0), frame_id=""), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ) + odom.publish( + PoseStamped( + header=Header(stamp=time_from_seconds(2.0), frame_id=""), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ) + goal.publish( + PoseStamped( + header=Header(stamp=time_from_seconds(3.0), frame_id=""), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ) assert unsub is not None unsub() - odom.publish(PoseStamped(ts=4.0)) + odom.publish( + PoseStamped( + header=Header(stamp=time_from_seconds(4.0), frame_id=""), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ) rec.close() store.stop() @@ -75,8 +113,7 @@ def test_taps_matching_dimos_streams(tmp_path: Path) -> None: store.stop() -def test_depth_is_recorded_losslessly_and_colour_as_jpeg(tmp_path: Path) -> None: - pytest.importorskip("turbojpeg") +def test_depth_and_colour_are_recorded_losslessly(tmp_path: Path) -> None: path = tmp_path / "memory.db" store = SqliteStore(path=str(path)) store.start() @@ -84,9 +121,21 @@ def test_depth_is_recorded_losslessly_and_colour_as_jpeg(tmp_path: Path) -> None mm = np.arange(48 * 64, dtype=np.uint16).reshape(48, 64) * 7 metres = mm.astype(np.float32) / 1000 frames = { - "depth16": Image(data=mm, format=ImageFormat.DEPTH16, frame_id="depth_optical", ts=1.0), - "depth": Image(data=metres, format=ImageFormat.DEPTH, frame_id="depth_optical", ts=1.0), - "color": Image(data=np.zeros((48, 64, 3), dtype=np.uint8), format=ImageFormat.RGB, ts=1.0), + "depth16": image_from_array( + mm, + encoding="16UC1", + header=Header(frame_id="depth_optical", stamp=time_from_seconds(1.0)), + ), + "depth": image_from_array( + metres, + encoding="32FC1", + header=Header(frame_id="depth_optical", stamp=time_from_seconds(1.0)), + ), + "color": image_from_array( + np.zeros((48, 64, 3), dtype=np.uint8), + encoding="rgb8", + header=Header(frame_id="color", stamp=time_from_seconds(1.0)), + ), } for name, frame in frames.items(): transport = _Transport() @@ -98,11 +147,12 @@ def test_depth_is_recorded_losslessly_and_colour_as_jpeg(tmp_path: Path) -> None store = SqliteStore(path=str(path), must_exist=True) store.start() codecs = {name: store._registry.get(name)["codec_id"] for name in frames} # type: ignore[index] - assert codecs == {"depth16": "lz4+lcm", "depth": "lz4+lcm", "color": "jpeg"} - for name in ("depth16", "depth"): + assert codecs == {"depth16": "cdr", "depth": "cdr", "color": "cdr"} + for name in frames: decoded = store.stream(name, Image).first().data - assert decoded.format == frames[name].format - np.testing.assert_array_equal(decoded.data, frames[name].data) + assert decoded.encoding == frames[name].encoding + assert decoded.header == frames[name].header + np.testing.assert_array_equal(image_view(decoded), image_view(frames[name])) store.stop() @@ -116,9 +166,33 @@ def test_full_queue_drops_and_counts(tmp_path: Path, monkeypatch: pytest.MonkeyP rec._writer.join() odom = _Transport() rec.tap("odom", PoseStamped, odom) - odom.publish(PoseStamped(ts=1.0)) - odom.publish(PoseStamped(ts=2.0)) - odom.publish(PoseStamped(ts=3.0)) + odom.publish( + PoseStamped( + header=Header(stamp=time_from_seconds(1.0), frame_id=""), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ) + odom.publish( + PoseStamped( + header=Header(stamp=time_from_seconds(2.0), frame_id=""), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ) + odom.publish( + PoseStamped( + header=Header(stamp=time_from_seconds(3.0), frame_id=""), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ) assert rec.dropped == 2 assert len(warnings) == 1 store.stop() @@ -132,7 +206,15 @@ def test_append_failure_does_not_kill_writer(tmp_path: Path) -> None: odom = _Transport() rec.tap("odom", PoseStamped, odom) odom.publish("not a PoseStamped") # append raises TypeError inside the writer - odom.publish(PoseStamped(ts=2.0)) + odom.publish( + PoseStamped( + header=Header(stamp=time_from_seconds(2.0), frame_id=""), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ) rec.close() # returns: the writer is still draining store.stop() @@ -178,3 +260,73 @@ def test_recording_delegates_to_rust_session( make_plan.assert_called_once_with(transports) create_session.assert_called_once_with("plan") session.stop.assert_called_once_with() + + +def test_zero_source_stamp_and_unstamped_arrival_time_are_distinct(tmp_path, monkeypatch): + monkeypatch.setattr(tap.time, "time", lambda: 123.5) + path = tmp_path / "time.db" + pose_transport, twist_transport = _Transport(), _Transport() + with SqliteStore(path=str(path)) as store: + recorder = TransportRecorder(store) + subscriptions = [] + try: + subscriptions.append(recorder.tap("pose", PoseStamped, pose_transport)) + subscriptions.append(recorder.tap("twist", Twist, twist_transport)) + pose_transport.publish( + PoseStamped( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ) + twist_transport.publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) + finally: + for unsubscribe in subscriptions: + assert unsubscribe is not None + unsubscribe() + recorder.close() + with SqliteStore(path=str(path), must_exist=True) as store: + pose = store.stream("pose").first() + twist = store.stream("twist").first() + assert pose.ts == 0.0 + assert pose.data == PoseStamped( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + assert twist.ts == 123.5 + assert twist.data == Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) + + +def test_region_recorder_preserves_nested_source_time_and_empty_replacement(tmp_path, monkeypatch): + monkeypatch.setattr(tap.time, "time", lambda: 123.5) + path = tmp_path / "regions.db" + transport = _Transport() + header = Header(stamp=Time(sec=0, nanosec=34), frame_id="map") + with SqliteStore(path=str(path)) as store: + recorder = TransportRecorder(store) + unsubscribe = recorder.tap("map_regions", RegionPointCloud2, transport) + try: + for xyz in [np.array([[1.0, 2.0, 3.0]]), np.empty((0, 3))]: + transport.publish( + RegionPointCloud2( + region_id=-196603, cloud=pointcloud_from_xyz(xyz, header=header) + ) + ) + finally: + assert unsubscribe is not None + unsubscribe() + recorder.close() + with SqliteStore(path=str(path), must_exist=True) as store: + observations = list(store.stream("map_regions")) + assert [obs.ts for obs in observations] == [34e-9, 34e-9] + assert [obs.data.region_id for obs in observations] == [-196603, -196603] + assert [obs.data.cloud.width for obs in observations] == [1, 0] diff --git a/dimos/memory/tool_e2e.py b/dimos/memory/tool_e2e.py index 42a8f5e80f8..7b90ebf5fe6 100644 --- a/dimos/memory/tool_e2e.py +++ b/dimos/memory/tool_e2e.py @@ -19,7 +19,7 @@ import bisect from typing import TYPE_CHECKING, Any -from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.sensor_msgs.msg import Image, PointCloud2 from dimos_generated.std_msgs.msg import Header import numpy as np @@ -48,7 +48,13 @@ def source(tmp_path_factory: pytest.TempPathFactory) -> Iterator[SqliteStore]: for index in range(1024): ns = 1_700_000_000_123_456_789 + index * 25_000_000 header = Header(stamp=time_from_nanoseconds(ns), frame_id="world") - pose = PoseStamped(header=header, pose=Pose(position=Point(x=index / 100))) + pose = PoseStamped( + header=header, + pose=Pose( + position=Point(x=index / 100, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) pixels = np.full((8, 8, 3), 64 + index % 128, dtype=np.uint8) pixels[:, index % 8] = (255, 0, 127) video.append(image_from_array(pixels, encoding="rgb8", header=header), ts=ns / 1e9) diff --git a/dimos/memory/type/observation.py b/dimos/memory/type/observation.py index ee3808f94ab..4a8ee2c2d0b 100644 --- a/dimos/memory/type/observation.py +++ b/dimos/memory/type/observation.py @@ -171,7 +171,7 @@ def pose_stamped(self) -> PoseStamped | None: return None x, y, z, qx, qy, qz, qw = self.pose_tuple return PoseStamped( - header=Header(stamp=time_from_seconds(self.ts)), + header=Header(stamp=time_from_seconds(self.ts), frame_id=""), pose=Pose( position=Point(x=x, y=y, z=z), orientation=Quaternion(x=qx, y=qy, z=qz, w=qw) ), diff --git a/dimos/memory/type/test_generated_query.py b/dimos/memory/type/test_generated_query.py index 50a57d8a547..b7397ed16de 100644 --- a/dimos/memory/type/test_generated_query.py +++ b/dimos/memory/type/test_generated_query.py @@ -15,16 +15,21 @@ # Copyright 2026 Dimensional Inc. # SPDX-License-Identifier: Apache-2.0 +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, Pose, PoseStamped, PoseWithCovariance, + Quaternion, Transform, TransformStamped, + Twist, + TwistWithCovariance, Vector3, ) from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest @@ -34,14 +39,43 @@ @pytest.mark.parametrize( "center", [ - Point(x=1), - Vector3(x=1), + Point(x=1, y=0.0, z=0.0), + Vector3(x=1, y=0.0, z=0.0), (1, 0, 0), np.array([1, 0, 0]), - Pose(position=Point(x=1)), - PoseStamped(pose=Pose(position=Point(x=1))), - TransformStamped(transform=Transform(translation=Vector3(x=1))), - Odometry(pose=PoseWithCovariance(pose=Pose(position=Point(x=1)))), + Pose(position=Point(x=1, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), + PoseStamped( + pose=Pose( + position=Point(x=1, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ), + TransformStamped( + transform=Transform( + translation=Vector3(x=1, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + child_frame_id="", + ), + Odometry( + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=1, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + child_frame_id="", + twist=TwistWithCovariance( + twist=Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ), + covariance=np.zeros(36, dtype=np.float64), + ), + ), ], ) def test_near_accepts_generated_nested_positions_and_retains_radius_boundary(center): diff --git a/dimos/memory/vis/space/rerun.py b/dimos/memory/vis/space/rerun.py index 643b6f363c7..3f1e02abfd6 100644 --- a/dimos/memory/vis/space/rerun.py +++ b/dimos/memory/vis/space/rerun.py @@ -19,6 +19,7 @@ import math from typing import TYPE_CHECKING, Any +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 from dimos_generated.nav_msgs.msg import OccupancyGrid from dimos_generated.sensor_msgs.msg import Image, PointCloud2 @@ -53,10 +54,10 @@ def _rgba(el: Any) -> tuple[int, int, int, int]: # base_link → camera_optical extrinsics (applied at render time for image observations) _BASE_TO_OPTICAL = TransformStamped( - header=Header(frame_id="base_link"), + header=Header(frame_id="base_link", stamp=Time(sec=0, nanosec=0)), child_frame_id="camera_optical", transform=Transform( - translation=Vector3(x=0.3), + translation=Vector3(x=0.3, y=0.0, z=0.0), rotation=Quaternion(x=-0.5, y=0.5, z=-0.5, w=0.5), ), ) diff --git a/dimos/memory/vis/space/test_generated_geometry.py b/dimos/memory/vis/space/test_generated_geometry.py index ac0a9f36027..ebc18df5bc9 100644 --- a/dimos/memory/vis/space/test_generated_geometry.py +++ b/dimos/memory/vis/space/test_generated_geometry.py @@ -17,7 +17,14 @@ import math import xml.etree.ElementTree as ET -from dimos_generated.geometry_msgs.msg import Point as GeoPoint, Pose as GeoPose, PoseStamped +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + Point as GeoPoint, + Pose as GeoPose, + PoseStamped, + Quaternion, +) from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid, Path from dimos_generated.sensor_msgs.msg import PointCloud2 from dimos_generated.std_msgs.msg import Header @@ -38,9 +45,22 @@ def test_rerun_logs_generated_pose_point_and_path(mocker): stamped = PoseStamped( pose=GeoPose( position=GeoPoint(x=1, y=2, z=3), orientation=quaternion_from_euler(0, 0, math.pi / 2) - ) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + path = Path( + poses=[ + stamped, + PoseStamped( + pose=GeoPose( + position=GeoPoint(x=4, y=5, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ), + ], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) - path = Path(poses=[stamped, PoseStamped(pose=GeoPose(position=GeoPoint(x=4, y=5)))]) space = Space() space.add(Point(stamped)).add(Pose(stamped)).add(Arrow(stamped)) space.add(Polyline(path)).add(Camera(stamped)) @@ -73,12 +93,15 @@ def test_both_renderers_apply_rotated_occupancy_origin(mocker): height=1, resolution=1.0, origin=GeoPose( - position=GeoPoint(x=2, y=3), orientation=quaternion_from_euler(0, 0, math.pi / 2) + position=GeoPoint(x=2, y=3, z=0.0), + orientation=quaternion_from_euler(0, 0, math.pi / 2), ), + map_load_time=Time(sec=0, nanosec=0), ), - data=[0, 100], + data=np.array([0, 100], dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) - space = Space().base_map(OccupancyGrid.decode(grid.encode())) + space = Space().base_map(cdr_decode(cdr_encode(grid), OccupancyGrid)) svg = ET.fromstring(space.to_svg()) image = svg.find(".//{http://www.w3.org/2000/svg}image") assert image is not None @@ -105,9 +128,9 @@ def test_both_renderers_apply_rotated_occupancy_origin(mocker): def test_generated_cloud_renders_in_svg_and_rerun(mocker): cloud = pointcloud_from_xyz( np.array([[0.0, 0.0, 0.0], [0.2, 0.3, 1.0], [1.0, 1.0, 0.0]]), - header=Header(frame_id="map"), + header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), ) - space = Space().add(PointCloud2.decode(cloud.encode())) + space = Space().add(cdr_decode(cdr_encode(cloud), PointCloud2)) svg = ET.fromstring(space.to_svg()) assert svg.find(".//{http://www.w3.org/2000/svg}image") is not None mocker.patch("dimos.visualization.rerun.init.rerun_init") @@ -127,15 +150,20 @@ def test_generated_cloud_renders_in_svg_and_rerun(mocker): def test_rerun_camera_uses_generated_calibration_and_image_helpers(mocker): - calibration = camera_info_from_intrinsics(20, 30, 2, 2, 4, 4, header=Header(frame_id="optical")) + calibration = camera_info_from_intrinsics( + 20, 30, 2, 2, 4, 4, header=Header(frame_id="optical", stamp=Time(sec=0, nanosec=0)) + ) image = image_from_array( np.zeros((4, 4, 3), dtype=np.uint8), encoding="rgb8", header=calibration.header ) pose = PoseStamped( - pose=GeoPose(position=GeoPoint(x=2), orientation=quaternion_from_euler(0, 0, 0)) + pose=GeoPose( + position=GeoPoint(x=2, y=0.0, z=0.0), orientation=quaternion_from_euler(0, 0, 0) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) space = Space().add(Camera(pose=pose, image=image, camera_info=calibration)) - original = (calibration.encode(), image.encode()) + original = (cdr_encode(calibration), cdr_encode(image)) mocker.patch("dimos.visualization.rerun.init.rerun_init") log = mocker.spy(rr, "log") with rr.RecordingStream("generated-camera-test") as recording: @@ -152,4 +180,4 @@ def test_rerun_camera_uses_generated_calibration_and_image_helpers(mocker): call.args[0] == "scene/cameras/0/image" and isinstance(call.args[1], rr.EncodedImage) for call in log.call_args_list ) - assert (calibration.encode(), image.encode()) == original + assert (cdr_encode(calibration), cdr_encode(image)) == original diff --git a/dimos/memory/vis/space/test_space.py b/dimos/memory/vis/space/test_space.py index cc289827a93..db62447ced4 100644 --- a/dimos/memory/vis/space/test_space.py +++ b/dimos/memory/vis/space/test_space.py @@ -14,6 +14,7 @@ """Tests for Space builder and element types.""" +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point as GeoPoint, Pose as GeoPose, @@ -22,7 +23,8 @@ Vector3, ) from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid, Path as Path -from dimos_generated.vision_msgs.msg import Detection3D +from dimos_generated.std_msgs.msg import Header +from dimos_generated.vision_msgs.msg import BoundingBox3D, Detection3D import numpy as np import pytest @@ -38,7 +40,11 @@ class TestElementTypes: def test_pose_wraps_posestamped(self): ps = PoseStamped( - pose=GeoPose(position=GeoPoint(x=3.2, y=1.5, z=0.0), orientation=Quaternion()) + pose=GeoPose( + position=GeoPoint(x=3.2, y=1.5, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) p = Pose(ps, color="red", label="fridge") assert p.msg is ps @@ -49,7 +55,8 @@ def test_arrow_wraps_posestamped(self): ps = PoseStamped( pose=GeoPose( position=GeoPoint(x=1, y=2, z=0), orientation=Quaternion(x=0, y=0, z=0.1, w=1) - ) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) a = Arrow(ps, color="orange", length=0.8) assert a.msg is ps @@ -62,13 +69,20 @@ def test_point_wraps_geopoint(self): assert p.msg.x == pytest.approx(7.1) def test_point_wraps_posestamped(self): - ps = PoseStamped(pose=GeoPose(position=GeoPoint(x=3, y=1, z=0), orientation=Quaternion())) + ps = PoseStamped( + pose=GeoPose( + position=GeoPoint(x=3, y=1, z=0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) p = Point(ps, radius=0.5) assert p.msg.pose.position.x == pytest.approx(3.0) def test_box3d_from_center_size(self): b = Box3D( - center=GeoPose(position=GeoPoint(x=5, y=3, z=0), orientation=Quaternion()), + center=GeoPose( + position=GeoPoint(x=5, y=3, z=0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), size=Vector3(x=2, y=1, z=0.5), label="table", ) @@ -77,7 +91,12 @@ def test_box3d_from_center_size(self): assert b.label == "table" def test_camera_with_image(self): - ps = PoseStamped(pose=GeoPose(position=GeoPoint(x=1, y=2, z=0), orientation=Quaternion())) + ps = PoseStamped( + pose=GeoPose( + position=GeoPoint(x=1, y=2, z=0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) img = image_from_array(np.zeros((480, 640, 3), dtype=np.uint8), encoding="rgb8") c = Camera(pose=ps, image=img, color="purple") assert c.pose is ps @@ -95,7 +114,12 @@ class TestSpaceExplicitElements: def test_add_pose(self): s = Space() - ps = PoseStamped(pose=GeoPose(position=GeoPoint(x=3, y=1, z=0), orientation=Quaternion())) + ps = PoseStamped( + pose=GeoPose( + position=GeoPoint(x=3, y=1, z=0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) pose = Pose(ps, color="red") s.add(pose) assert len(s) == 1 @@ -103,7 +127,12 @@ def test_add_pose(self): def test_add_multiple_types(self): s = Space() - ps = PoseStamped(pose=GeoPose(position=GeoPoint(x=3, y=1, z=0), orientation=Quaternion())) + ps = PoseStamped( + pose=GeoPose( + position=GeoPoint(x=3, y=1, z=0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) s.add(Pose(ps, color="red")) s.add(Arrow(ps, color="orange")) s.add(Point(GeoPoint(x=1, y=2, z=0), label="x")) @@ -111,7 +140,12 @@ def test_add_multiple_types(self): assert len(s) == 4 def test_chaining(self): - ps = PoseStamped(pose=GeoPose(position=GeoPoint(x=1, y=1, z=0), orientation=Quaternion())) + ps = PoseStamped( + pose=GeoPose( + position=GeoPoint(x=1, y=1, z=0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) s = Space().add(Pose(ps)).add(Arrow(ps)).add(Text((0, 0, 0), "hi")) assert len(s) == 3 @@ -122,7 +156,11 @@ class TestSpaceAutoWrap: def test_posestamped_becomes_pose(self): s = Space() ps = PoseStamped( - pose=GeoPose(position=GeoPoint(x=3.2, y=1.5, z=0), orientation=Quaternion()) + pose=GeoPose( + position=GeoPoint(x=3.2, y=1.5, z=0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) s.add(ps, color="blue", label="auto") assert len(s) == 1 @@ -146,10 +184,15 @@ def test_path_becomes_polyline(self): p = Path( poses=[ PoseStamped( - pose=GeoPose(position=GeoPoint(x=i, y=0, z=0), orientation=Quaternion()) + pose=GeoPose( + position=GeoPoint(x=i, y=0, z=0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) for i in range(3) - ] + ], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) s.add(p, color="blue", width=0.1) el = s.elements[0] @@ -160,12 +203,36 @@ def test_path_becomes_polyline(self): def test_occupancy_grid_passthrough(self): s = Space() - grid = OccupancyGrid() + grid = OccupancyGrid( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + info=MapMetaData( + map_load_time=Time(sec=0, nanosec=0), + resolution=0.0, + width=0, + height=0, + origin=GeoPose( + position=GeoPoint(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([], dtype=np.int8), + ) s.add(grid) assert s.elements[0] is grid def test_detection3d_becomes_box3d(self): - det = Detection3D() + det = Detection3D( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + results=[], + bbox=BoundingBox3D( + center=GeoPose( + position=GeoPoint(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + size=Vector3(x=0.0, y=0.0, z=0.0), + ), + id="", + ) det.bbox.center.position.x = 5.0 det.bbox.center.position.y = 3.0 det.bbox.size.x = 2.0 @@ -213,7 +280,13 @@ def test_posestamped_observation_stored_as_observation(self): id=3, ts=3.0, pose=(1, 2, 0, 0, 0, 0, 1), - _data=PoseStamped(pose=GeoPose(position=GeoPoint(x=5, y=2))), + _data=PoseStamped( + pose=GeoPose( + position=GeoPoint(x=5, y=2, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ), ) s = Space() @@ -242,14 +315,33 @@ class TestSpaceConvenience: """Space convenience methods: base_map.""" def test_base_map(self): - grid = OccupancyGrid() + grid = OccupancyGrid( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + info=MapMetaData( + map_load_time=Time(sec=0, nanosec=0), + resolution=0.0, + width=0, + height=0, + origin=GeoPose( + position=GeoPoint(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([], dtype=np.int8), + ) s = Space().base_map(grid) assert len(s) == 1 assert isinstance(s.elements[0], OccupancyGrid) def test_add_list_of_msgs(self): poses = [ - PoseStamped(pose=GeoPose(position=GeoPoint(x=i, y=0, z=0), orientation=Quaternion())) + PoseStamped( + pose=GeoPose( + position=GeoPoint(x=i, y=0, z=0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) for i in range(3) ] s = Space() @@ -266,7 +358,12 @@ def test_repr_empty(self): def test_repr_with_elements(self): s = Space() - ps = PoseStamped(pose=GeoPose(position=GeoPoint(x=0, y=0, z=0), orientation=Quaternion())) + ps = PoseStamped( + pose=GeoPose( + position=GeoPoint(x=0, y=0, z=0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) s.add(Pose(ps)) s.add(Pose(ps)) s.add(Arrow(ps)) @@ -294,7 +391,11 @@ def test_pose_renders_polygon(self): s.add( Pose( PoseStamped( - pose=GeoPose(position=GeoPoint(x=1, y=2, z=0), orientation=Quaternion()) + pose=GeoPose( + position=GeoPoint(x=1, y=2, z=0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ), color="blue", ) @@ -311,7 +412,8 @@ def test_arrow_renders_polygon(self): pose=GeoPose( position=GeoPoint(x=0, y=0, z=0), orientation=Quaternion(x=0, y=0, z=0.38, w=0.92), - ) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ), color="orange", ) @@ -328,11 +430,14 @@ def test_polyline_renders(self): poses=[ PoseStamped( pose=GeoPose( - position=GeoPoint(x=i, y=i * 0.5, z=0), orientation=Quaternion() - ) + position=GeoPoint(x=i, y=i * 0.5, z=0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) for i in range(5) - ] + ], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ), color="blue", ) @@ -344,7 +449,10 @@ def test_box3d_renders_rect(self): s = Space() s.add( Box3D( - center=GeoPose(position=GeoPoint(x=5, y=3, z=0), orientation=Quaternion()), + center=GeoPose( + position=GeoPoint(x=5, y=3, z=0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), size=Vector3(x=2, y=1, z=0), label="table", ) @@ -365,7 +473,11 @@ def test_camera_without_info_renders_dot(self): s.add( Camera( pose=PoseStamped( - pose=GeoPose(position=GeoPoint(x=1, y=2, z=0), orientation=Quaternion()) + pose=GeoPose( + position=GeoPoint(x=1, y=2, z=0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ), color="purple", ) @@ -376,8 +488,18 @@ def test_camera_without_info_renders_dot(self): def test_occupancy_grid_renders_image(self): grid = OccupancyGrid( - info=MapMetaData(width=10, height=10, resolution=0.1), - data=[0] * 100, + info=MapMetaData( + width=10, + height=10, + resolution=0.1, + map_load_time=Time(sec=0, nanosec=0), + origin=GeoPose( + position=GeoPoint(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.asarray([0] * 100, dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) s = Space().base_map(grid) svg = s.to_svg() @@ -386,7 +508,12 @@ def test_occupancy_grid_renders_image(self): def test_mixed_space(self): s = Space() - ps = PoseStamped(pose=GeoPose(position=GeoPoint(x=3, y=1, z=0), orientation=Quaternion())) + ps = PoseStamped( + pose=GeoPose( + position=GeoPoint(x=3, y=1, z=0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) s.add(Pose(ps, color="red", label="robot")) s.add(Arrow(ps, color="orange")) s.add(Point(GeoPoint(x=5, y=5, z=0), color="green", label="goal")) diff --git a/dimos/message_codegen/test_builtin_catalog.py b/dimos/message_codegen/test_builtin_catalog.py index 23c940c4fd1..5c17805a86f 100644 --- a/dimos/message_codegen/test_builtin_catalog.py +++ b/dimos/message_codegen/test_builtin_catalog.py @@ -36,7 +36,7 @@ def test_unlisted_sources_do_not_expand_the_public_catalog(tmp_path): manifest = json.loads( (source / "dimos_generated_schemas/package/message-package.json").read_text() ) - assert len(manifest["owned"]) == 78 + assert len(manifest["owned"]) == 80 assert "extra_msgs/msg/NotExported" not in manifest["owned"] assert "builtin_interfaces/msg/Time" in manifest["owned"] diff --git a/dimos/message_codegen/test_builtin_install.py b/dimos/message_codegen/test_builtin_install.py index 0b709e6b0f1..97c532478c0 100644 --- a/dimos/message_codegen/test_builtin_install.py +++ b/dimos/message_codegen/test_builtin_install.py @@ -38,7 +38,7 @@ assert type(r.header) is Header and r == m assert type(pickle.loads(pickle.dumps(m)).header) is Header assert 'MSG: std_msgs/Header' in schema(m.__msgtype__) -assert len(message_types()) == 78 +assert len(message_types()) == 80 assert 'visualization_msgs/msg/Marker' not in message_types() root = Path(package_root()) assert (root / 'schema-sources.json').is_file() diff --git a/dimos/message_codegen/test_builtin_preparation.py b/dimos/message_codegen/test_builtin_preparation.py index 82d816bc161..d4c4c7a068b 100644 --- a/dimos/message_codegen/test_builtin_preparation.py +++ b/dimos/message_codegen/test_builtin_preparation.py @@ -48,12 +48,12 @@ def no_network(*args, **kwargs): monkeypatch.setattr("dimos_cdrgen.source_inputs.requests.get", no_network) inputs = prepare(cache, offline=True, project=project) - assert len(list(inputs.rglob("*.msg"))) == 78 + assert len(list(inputs.rglob("*.msg"))) == 80 output = generate(project) resources = output / "dimos_generated_schemas/package" manifest = json.loads((resources / "message-package.json").read_text()) - assert len(manifest["owned"]) == 78 - assert len(list(project.glob("custom/dimos_msgs/msg/*.msg"))) == 19 + assert len(manifest["owned"]) == 80 + assert len(list(project.glob("custom/dimos_msgs/msg/*.msg"))) == 21 lock = json.loads((project / "schema-sources.json").read_text()) assert all(len(item["sha256"]) == 64 for item in lock["sources"]) assert (resources / "dimos_msgs-LICENSE").is_file() diff --git a/dimos/models/embedding/siglip2.py b/dimos/models/embedding/siglip2.py index 963707fffe5..d690cc57145 100644 --- a/dimos/models/embedding/siglip2.py +++ b/dimos/models/embedding/siglip2.py @@ -17,6 +17,7 @@ from functools import cached_property from typing import overload +from dimos_generated.sensor_msgs.msg import Image from PIL import Image as PILImage import torch import torch.nn.functional as functional @@ -24,7 +25,8 @@ from dimos.models.base import HuggingFaceModel from dimos.models.embedding.base import Embedding, EmbeddingModel, HuggingFaceEmbeddingModelConfig -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_to_rgb +from dimos.msgs.time import to_seconds TEXT_MAX_LENGTH = 64 @@ -61,7 +63,7 @@ def embed(self, image: Image, /) -> Embedding: ... def embed(self, *images: Image) -> list[Embedding]: ... def embed(self, *images: Image) -> Embedding | list[Embedding]: """Embed one or more images into the shared image-text space.""" - pil_images = [PILImage.fromarray(img.to_rgb().data) for img in images] + pil_images = [PILImage.fromarray(image_to_rgb(img)) for img in images] with torch.inference_mode(): inputs = self._processor(images=pil_images, return_tensors="pt") @@ -71,7 +73,8 @@ def embed(self, *images: Image) -> Embedding | list[Embedding]: image_features = functional.normalize(image_features, dim=-1) embeddings = [ - Embedding(vector=feat, timestamp=images[i].ts) for i, feat in enumerate(image_features) + Embedding(vector=feat, timestamp=to_seconds(images[i].header.stamp)) + for i, feat in enumerate(image_features) ] return embeddings[0] if len(images) == 1 else embeddings diff --git a/dimos/models/vl/test_generated_images.py b/dimos/models/vl/test_generated_images.py index 01e560fb76a..2795ed5fbf2 100644 --- a/dimos/models/vl/test_generated_images.py +++ b/dimos/models/vl/test_generated_images.py @@ -16,6 +16,7 @@ from unittest.mock import MagicMock import cv2 +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import Image from dimos_generated.std_msgs.msg import Header @@ -49,7 +50,7 @@ def test_generated_vl_resize_and_detection(): assert (resized.width, resized.height) == (30, 20) assert resized.header == image.header detections = model.query_detections(image, "target") - wire = Detection2DArray.decode(detections.to_ros_detection2d_array().encode()) + wire = cdr_decode(cdr_encode(detections.to_ros_detection2d_array()), Detection2DArray) assert len(wire.detections) == 1 assert wire.header == image.header assert wire.detections[0].results[0].hypothesis.class_id == "-1" diff --git a/dimos/msgs/camera_info.py b/dimos/msgs/camera_info.py index 7977a275b22..bb7c6ecc055 100644 --- a/dimos/msgs/camera_info.py +++ b/dimos/msgs/camera_info.py @@ -21,7 +21,7 @@ from typing import Any, Literal from dimos_generated.builtin_interfaces.msg import Time -from dimos_generated.sensor_msgs.msg import CameraInfo +from dimos_generated.sensor_msgs.msg import CameraInfo, RegionOfInterest from dimos_generated.std_msgs.msg import Header import numpy as np from numpy.typing import NDArray @@ -48,10 +48,13 @@ def camera_info_from_intrinsics( width=width, height=height, distortion_model="plumb_bob", - d=[0.0] * 5, - k=[fx, 0, cx, 0, fy, cy, 0, 0, 1], - r=[1, 0, 0, 0, 1, 0, 0, 0, 1], - p=[fx, 0, cx, 0, 0, fy, cy, 0, 0, 0, 1, 0], + d=np.asarray([0.0] * 5, dtype=np.float64), + k=np.array([fx, 0, cx, 0, fy, cy, 0, 0, 1], dtype=np.float64), + r=np.array([1, 0, 0, 0, 1, 0, 0, 0, 1], dtype=np.float64), + p=np.array([fx, 0, cx, 0, 0, fy, cy, 0, 0, 0, 1, 0], dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), ) @@ -109,10 +112,13 @@ def camera_info_from_yaml(path: str | Path, *, header: Header) -> CameraInfo: width=width, height=height, distortion_model=data["distortion_model"], - k=_matrix(data, "camera_matrix", 3, 3), - d=_matrix(data, "distortion_coefficients", 1, None), - r=_matrix(data, "rectification_matrix", 3, 3), - p=_matrix(data, "projection_matrix", 3, 4), + k=np.asarray(_matrix(data, "camera_matrix", 3, 3), dtype=np.float64), + d=np.asarray(_matrix(data, "distortion_coefficients", 1, None), dtype=np.float64), + r=np.asarray(_matrix(data, "rectification_matrix", 3, 3), dtype=np.float64), + p=np.asarray(_matrix(data, "projection_matrix", 3, 4), dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), ) @@ -155,7 +161,9 @@ def __getattr__(self, name: str) -> CameraInfo: for key in (name, snake): path = self._calibration_dir / f"{key}.yaml" if path.is_file(): - value = camera_info_from_yaml(path, header=Header(frame_id=self._frame_id)) + value = camera_info_from_yaml( + path, header=Header(frame_id=self._frame_id, stamp=Time(sec=0, nanosec=0)) + ) self._cache[name] = value self._cache[snake] = value return value diff --git a/dimos/msgs/contacts.py b/dimos/msgs/contacts.py new file mode 100644 index 00000000000..5dc3c81a269 --- /dev/null +++ b/dimos/msgs/contacts.py @@ -0,0 +1,20 @@ +# 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. + +"""Semantic labels used by the simulator's generated contact messages.""" + +from typing import Literal + +Part = Literal["trunk", "lidar", "leg", "foot"] +Kind = Literal["floor", "wall", "ceiling", "clutter"] diff --git a/dimos/msgs/control.py b/dimos/msgs/control.py new file mode 100644 index 00000000000..b9d6ed2664f --- /dev/null +++ b/dimos/msgs/control.py @@ -0,0 +1,65 @@ +# 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. + +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +"""Construction and boundary validation for generated control interface values.""" + +from collections.abc import Sequence +import math +import time + +from dimos_generated.dimos_msgs.msg import ControlValues +import numpy as np +from numpy.typing import ArrayLike + + +def validate_control_values(frame: ControlValues) -> None: + """Reject mismatched arrays or non-finite values at a control boundary.""" + if len(frame.interface_names) != len(frame.values): + raise ValueError( + "interface_names and values must be the same length; " + f"got {len(frame.interface_names)} names and {len(frame.values)} values" + ) + for name, value in zip(frame.interface_names, frame.values, strict=True): + if not math.isfinite(value): + raise ValueError(f"non-finite value for interface {name!r}: {value!r}") + + +def control_values( + source: str, + source_ts: float | None = None, + epoch: int = 0, + sequence: int = 0, + interface_names: Sequence[str] = (), + values: ArrayLike = (), +) -> ControlValues: + """Copy a sparse frame, stamping only an absent publisher timestamp.""" + frame = ControlValues( + source=source, + source_ts=time.time() if source_ts is None else source_ts, + epoch=epoch, + sequence=sequence, + interface_names=list(interface_names), + values=np.array(values, dtype=np.float64, copy=True), + ) + validate_control_values(frame) + return frame + + +def control_values_dict(frame: ControlValues) -> dict[str, float]: + """Pair interface names with values; the contract checks uniqueness.""" + validate_control_values(frame) + return dict(zip(frame.interface_names, frame.values, strict=True)) diff --git a/dimos/msgs/control_msgs/test_ControlValues.py b/dimos/msgs/control_msgs/test_ControlValues.py deleted file mode 100644 index 0fcdb4a5520..00000000000 --- a/dimos/msgs/control_msgs/test_ControlValues.py +++ /dev/null @@ -1,182 +0,0 @@ -# Copyright 2025-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. - -"""Wire contract for control_msgs.ControlValues. - -The byte-identity test is the load-bearing one: everything downstream of this -PR assumes the wrapper is a drop-in for the generated type on the wire. -""" - -import time - -from dimos_lcm.control_msgs import ControlValues as LCMControlValues -import pytest - -from dimos.core.transport import LCMTransport, ZenohTransport -from dimos.msgs.control_msgs.ControlValues import ControlValues -from dimos.msgs.helpers import resolve_msg_type -from dimos.msgs.protocol import DimosMsg - -FIXED_FRAME = { - "source": "xarm", - "source_ts": 1758547200.25, - "epoch": 7, - "sequence": 1234, - "interface_names": ["xarm/j1/position", "xarm/j2/position", "xarm/gripper/position"], - "values": [0.5, -1.25, 0.085], -} - - -def test_round_trip() -> None: - """Every field survives encode/decode unchanged.""" - original = ControlValues(**FIXED_FRAME) # type: ignore[arg-type] - - decoded = ControlValues.lcm_decode(original.lcm_encode()) - - assert decoded.source == original.source - assert decoded.source_ts == original.source_ts - assert decoded.epoch == original.epoch - assert decoded.sequence == original.sequence - assert decoded.interface_names == original.interface_names - assert decoded.values == original.values - - -def test_empty_frame_round_trip() -> None: - """The coordinator heartbeat: an epoch and a sequence, commanding nothing.""" - heartbeat = ControlValues("coordinator", source_ts=99.5, epoch=3, sequence=42) - assert heartbeat.interface_names == [] - assert heartbeat.values == [] - - decoded = ControlValues.lcm_decode(heartbeat.lcm_encode()) - - assert decoded.source == "coordinator" - assert decoded.source_ts == 99.5 - assert decoded.epoch == 3 - assert decoded.sequence == 42 - assert decoded.interface_names == [] - assert decoded.values == [] - assert decoded.as_dict() == {} - - -def test_mismatched_lengths_raise() -> None: - """Parallel arrays that are not parallel are rejected at construction.""" - with pytest.raises(ValueError, match="same length"): - ControlValues("arm", interface_names=["a", "b"], values=[1.0]) - - with pytest.raises(ValueError, match="same length"): - ControlValues("arm", interface_names=[], values=[1.0]) - - -@pytest.mark.parametrize("bad", [float("nan"), float("inf"), float("-inf")]) -def test_non_finite_raises(bad: float) -> None: - """NaN and both infinities are rejected, and the message names the interface.""" - with pytest.raises(ValueError, match="arm/j2/effort"): - ControlValues( - "arm", interface_names=["arm/j1/position", "arm/j2/effort"], values=[0.0, bad] - ) - - -def test_decode_rejects_mismatched_wire_frame() -> None: - """A frame whose arrays disagree on the wire does not become an object.""" - malformed = LCMControlValues( - source="arm", - source_ts=1.0, - epoch=1, - sequence=1, - interface_names_length=2, - values_length=1, - interface_names=["a", "b"], - values=[1.0], - ) - - with pytest.raises(ValueError, match="same length"): - ControlValues.lcm_decode(malformed.lcm_encode()) - - -def test_bytes_identical_to_generated_type() -> None: - """The wrapper is a drop-in for the generated type on the wire.""" - generated = LCMControlValues( - source=FIXED_FRAME["source"], - source_ts=FIXED_FRAME["source_ts"], - epoch=FIXED_FRAME["epoch"], - sequence=FIXED_FRAME["sequence"], - interface_names_length=len(FIXED_FRAME["interface_names"]), - values_length=len(FIXED_FRAME["values"]), - interface_names=list(FIXED_FRAME["interface_names"]), - values=list(FIXED_FRAME["values"]), - ) - - assert ControlValues(**FIXED_FRAME).lcm_encode() == generated.lcm_encode() # type: ignore[arg-type] - - -def test_as_dict() -> None: - """as_dict pairs the two arrays positionally.""" - frame = ControlValues(**FIXED_FRAME) # type: ignore[arg-type] - - assert frame.as_dict() == { - "xarm/j1/position": 0.5, - "xarm/j2/position": -1.25, - "xarm/gripper/position": 0.085, - } - - -def test_source_ts_defaults_to_now() -> None: - """An unstamped frame is stamped at construction, not left at zero.""" - before = time.time() - frame = ControlValues("arm") - after = time.time() - - assert before <= frame.source_ts <= after - - -def test_source_ts_zero_is_kept_not_replaced() -> None: - """0.0 is a real stamp: only None means "stamp me now".""" - assert ControlValues("arm", source_ts=0.0).source_ts == 0.0 - assert ( - ControlValues.lcm_decode(ControlValues("arm", source_ts=0.0).lcm_encode()).source_ts == 0.0 - ) - - -def test_repr_names_every_field() -> None: - """A logged frame shows the epoch and sequence, which is what debugging needs.""" - text = repr( - ControlValues( - "arm", source_ts=1.5, epoch=7, sequence=3, interface_names=["a"], values=[2.0] - ) - ) - - assert text.startswith("ControlValues(") - for fragment in ("source='arm'", "source_ts=1.5", "epoch=7", "sequence=3", "['a']", "[2.0]"): - assert fragment in text, fragment - - -def test_satisfies_dimos_msg_protocol() -> None: - """The wrapper is usable anywhere dimos.msgs.protocol.DimosMsg is required.""" - assert isinstance(ControlValues("arm"), DimosMsg) - assert ControlValues.msg_name == "control_msgs.ControlValues" - - -def test_resolve_msg_type_returns_the_wrapper() -> None: - """Name resolution prefers the wrapper over the generated type.""" - assert resolve_msg_type("control_msgs.ControlValues") is ControlValues - - -def test_transports_accept_the_type() -> None: - """Both transports the control streams will use carry the type (no I/O).""" - lcm = LCMTransport("/control_command", ControlValues) - assert lcm.topic.lcm_type is ControlValues - - spec = ZenohTransport.spec("/control_state", ControlValues) - assert spec.cls is ZenohTransport - assert ControlValues in spec.args diff --git a/dimos/msgs/foxglove_msgs/test_CompressedVideo.py b/dimos/msgs/foxglove_msgs/test_CompressedVideo.py index fa475e0c5b9..616f1bb3fb0 100644 --- a/dimos/msgs/foxglove_msgs/test_CompressedVideo.py +++ b/dimos/msgs/foxglove_msgs/test_CompressedVideo.py @@ -12,7 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.foxglove_msgs.msg import CompressedVideo +import numpy as np import pytest import rerun as rr @@ -25,12 +28,12 @@ def test_cdr_encode_decode() -> None: original = CompressedVideo( - data=PACKET, + data=np.frombuffer(PACKET, dtype=np.uint8), format="h264", frame_id="front_camera", timestamp=time_from_nanoseconds(1234567890123456789), ) - decoded = CompressedVideo.decode(original.encode()) + decoded = cdr_decode(cdr_encode(original), CompressedVideo) assert bytes(decoded.data) == PACKET assert decoded.format == "h264" @@ -45,11 +48,25 @@ def test_resolves_generated_canonical_type() -> None: def test_to_rerun_video_stream() -> None: - stream = video_archetype(CompressedVideo(data=PACKET, format="h264")) + stream = video_archetype( + CompressedVideo( + data=np.frombuffer(PACKET, dtype=np.uint8), + format="h264", + timestamp=Time(sec=0, nanosec=0), + frame_id="", + ) + ) assert isinstance(stream, rr.VideoStream) assert bytes(stream.sample.as_arrow_array().to_pylist()[0]) == PACKET # type: ignore[union-attr] def test_to_rerun_unknown_codec() -> None: with pytest.raises(ValueError, match="mjpeg"): - video_archetype(CompressedVideo(data=PACKET, format="mjpeg")) + video_archetype( + CompressedVideo( + data=np.frombuffer(PACKET, dtype=np.uint8), + format="mjpeg", + timestamp=Time(sec=0, nanosec=0), + frame_id="", + ) + ) diff --git a/dimos/msgs/geometry.py b/dimos/msgs/geometry.py index a5f007eb114..298db4fce71 100644 --- a/dimos/msgs/geometry.py +++ b/dimos/msgs/geometry.py @@ -15,6 +15,7 @@ """Geometry operations on generated ROS value types.""" from collections.abc import Sequence +from copy import deepcopy import math from dimos_generated.geometry_msgs.msg import ( @@ -46,11 +47,11 @@ def transform_from_odometry(message: Odometry) -> TransformStamped: """Copy the pose and declared frames into TF, preserving the exact source stamp.""" position = message.pose.pose.position return TransformStamped( - header=message.header, + header=deepcopy(message.header), child_frame_id=message.child_frame_id, transform=Transform( translation=Vector3(x=position.x, y=position.y, z=position.z), - rotation=message.pose.pose.orientation, + rotation=deepcopy(message.pose.pose.orientation), ), ) @@ -108,7 +109,7 @@ def compose_transforms(first: TransformStamped, second: TransformStamped) -> Tra rotation = _rotation(first.transform.rotation) translation = _translation(first.transform) + rotation.apply(_translation(second.transform)) return TransformStamped( - header=first.header, + header=deepcopy(first.header), child_frame_id=second.child_frame_id, transform=_transform(translation, rotation * _rotation(second.transform.rotation)), ) @@ -118,7 +119,7 @@ def inverse_transform(message: TransformStamped) -> TransformStamped: """Return B←A from A←B, swapping frame names and preserving the exact stamp.""" rotation = _rotation(message.transform.rotation).inv() return TransformStamped( - header=Header(stamp=message.header.stamp, frame_id=message.child_frame_id), + header=Header(stamp=deepcopy(message.header.stamp), frame_id=message.child_frame_id), child_frame_id=message.header.frame_id, transform=_transform(-rotation.apply(_translation(message.transform)), rotation), ) @@ -128,10 +129,10 @@ def pose_from_transform(message: TransformStamped) -> PoseStamped: """Copy the child frame's pose in the parent frame into a generated PoseStamped.""" translation = message.transform.translation return PoseStamped( - header=message.header, + header=deepcopy(message.header), pose=Pose( position=Point(x=translation.x, y=translation.y, z=translation.z), - orientation=message.transform.rotation, + orientation=deepcopy(message.transform.rotation), ), ) @@ -142,11 +143,11 @@ def transform_from_pose(message: PoseStamped, *, child_frame_id: str) -> Transfo raise TypeError("Expected a generated PoseStamped") position = message.pose.position return TransformStamped( - header=message.header, + header=deepcopy(message.header), child_frame_id=child_frame_id, transform=Transform( translation=Vector3(x=position.x, y=position.y, z=position.z), - rotation=message.pose.orientation, + rotation=deepcopy(message.pose.orientation), ), ) @@ -193,7 +194,7 @@ def transform_from_matrix(matrix: NDArray[np.float64]) -> Transform: pose = pose_from_matrix(matrix) return Transform( translation=Vector3(x=pose.position.x, y=pose.position.y, z=pose.position.z), - rotation=pose.orientation, + rotation=deepcopy(pose.orientation), ) @@ -214,7 +215,7 @@ def translate_pose_local(pose: Pose, offset: Vector3) -> Pose: position=Point( x=pose.position.x + delta[0], y=pose.position.y + delta[1], z=pose.position.z + delta[2] ), - orientation=pose.orientation, + orientation=deepcopy(pose.orientation), ) @@ -271,7 +272,7 @@ def normalized_vector(value: Vector3) -> Vector3: """Return a unit vector; zero length remains zero.""" array = vector_array(value) length = float(np.linalg.norm(array)) - return vector_from_array(array / length) if length else Vector3() + return vector_from_array(array / length) if length else Vector3(x=0.0, y=0.0, z=0.0) def quaternion_array(value: Quaternion) -> NDArray[np.float64]: @@ -284,7 +285,7 @@ def quaternion_from_array(values: Sequence[float] | NDArray[np.float64]) -> Quat array = np.asarray(values, dtype=np.float64) if array.shape != (4,): raise ValueError("Quaternion requires exactly 4 components") - return Quaternion(x=array[0], y=array[1], z=array[2], w=array[3]) + return Quaternion(x=float(array[0]), y=float(array[1]), z=float(array[2]), w=float(array[3])) def quaternion_product(first: Quaternion, second: Quaternion) -> Quaternion: diff --git a/dimos/msgs/geometry_msgs/test_PointStamped.py b/dimos/msgs/geometry_msgs/test_PointStamped.py index e4a27f14b88..3bb0c13fbea 100644 --- a/dimos/msgs/geometry_msgs/test_PointStamped.py +++ b/dimos/msgs/geometry_msgs/test_PointStamped.py @@ -14,6 +14,9 @@ """Generated point/stamp roundtrips and explicit pose construction.""" +from copy import deepcopy + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import Point, PointStamped, Pose, PoseStamped, Quaternion from dimos_generated.std_msgs.msg import Header from rosbags.typesys import Stores, get_typestore @@ -23,7 +26,7 @@ def test_point_has_standard_cdr_schema() -> None: point = Point(x=1, y=2, z=3) - decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(point.encode(), Point.msg_name) + decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(cdr_encode(point), Point.__msgtype__) assert (decoded.x, decoded.y, decoded.z) == (1, 2, 3) @@ -32,7 +35,7 @@ def test_cdr_encode_decode() -> None: point=Point(x=1.5, y=-2.5, z=3.5), header=Header(stamp=time_from_nanoseconds(1700000000123456789), frame_id="/world/grid"), ) - dest = PointStamped.decode(source.encode()) + dest = cdr_decode(cdr_encode(source), PointStamped) assert isinstance(dest, PointStamped) assert dest is not source assert dest.point.x == source.point.x @@ -47,7 +50,8 @@ def test_explicit_pose_stamped_conversion() -> None: point=Point(x=1, y=2, z=3), header=Header(stamp=time_from_seconds(500), frame_id="/map") ) pose = PoseStamped( - header=point.header, pose=Pose(position=point.point, orientation=Quaternion(w=1)) + header=deepcopy(point.header), + pose=Pose(position=deepcopy(point.point), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0)), ) assert isinstance(pose, PoseStamped) assert pose.pose.position.x == 1 diff --git a/dimos/msgs/geometry_msgs/test_Pose.py b/dimos/msgs/geometry_msgs/test_Pose.py index 9dd6bb12a88..fc8c52e4766 100644 --- a/dimos/msgs/geometry_msgs/test_Pose.py +++ b/dimos/msgs/geometry_msgs/test_Pose.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from copy import deepcopy import pickle +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion, Vector3 import numpy as np import pytest @@ -32,7 +34,9 @@ def test_pose_default_init() -> None: """Test that default initialization creates a pose at origin with identity orientation.""" - pose = Pose() + pose = Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) assert pose.position.x == 0.0 assert pose.position.y == 0.0 assert pose.position.z == 0.0 @@ -47,7 +51,9 @@ def test_pose_default_init() -> None: def test_pose_pose_init() -> None: """Test initialization with position coordinates only (identity orientation).""" - pose_data = Pose(position=Point(x=1.0, y=2.0, z=3.0)) + pose_data = Pose( + position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) pose = pose_data assert pose.position.x == 1.0 assert pose.position.y == 2.0 @@ -63,7 +69,9 @@ def test_pose_pose_init() -> None: def test_pose_position_init() -> None: """Test initialization with position coordinates only (identity orientation).""" - pose = Pose(position=Point(x=1.0, y=2.0, z=3.0)) + pose = Pose( + position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) assert pose.position.x == 1.0 assert pose.position.y == 2.0 assert pose.position.z == 3.0 @@ -96,7 +104,10 @@ def test_pose_full_init() -> None: def test_pose_vector_position_init() -> None: """Test initialization with Vector3 position (identity orientation).""" position = Vector3(x=4.0, y=5.0, z=6.0) - pose = Pose(position=point_from_array(vector_array(position))) + pose = Pose( + position=point_from_array(vector_array(position)), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ) assert pose.position.x == 4.0 assert pose.position.y == 5.0 assert pose.position.z == 6.0 @@ -175,7 +186,7 @@ def test_pose_copy_init() -> None: original = Pose( position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) ) - copy = Pose.decode(original.encode()) + copy = cdr_decode(cdr_encode(original), Pose) assert copy.position.x == 1.0 assert copy.position.y == 2.0 assert copy.position.z == 3.0 @@ -189,7 +200,9 @@ def test_pose_copy_init() -> None: def test_pose_mutated_value_copy() -> None: """Test initialization from a generated Pose.""" - source_pose = Pose() + source_pose = Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) source_pose.position.x = 1.0 source_pose.position.y = 2.0 source_pose.position.z = 3.0 @@ -197,7 +210,7 @@ def test_pose_mutated_value_copy() -> None: source_pose.orientation.y = 0.2 source_pose.orientation.z = 0.3 source_pose.orientation.w = 0.9 - pose = Pose.decode(source_pose.encode()) + pose = cdr_decode(cdr_encode(source_pose), Pose) assert pose.position.x == 1.0 assert pose.position.y == 2.0 assert pose.position.z == 3.0 @@ -223,7 +236,9 @@ def test_pose_properties() -> None: def test_pose_euler_properties_identity() -> None: """Test pose Euler angle properties with identity orientation.""" - pose = Pose(position=Point(x=1.0, y=2.0, z=3.0)) + pose = Pose( + position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) assert np.isclose(quaternion_euler(pose.orientation)[0], 0.0, atol=1e-10) assert np.isclose(quaternion_euler(pose.orientation)[1], 0.0, atol=1e-10) assert np.isclose(quaternion_euler(pose.orientation)[2], 0.0, atol=1e-10) @@ -237,7 +252,7 @@ def test_pose_independent_cdr_fields() -> None: position=Point(x=1.234, y=2.567, z=3.891), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9), ) - decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(source.encode(), Pose.msg_name) + decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(cdr_encode(source), Pose.__msgtype__) assert (decoded.position.x, decoded.position.y, decoded.position.z) == (1.234, 2.567, 3.891) assert ( decoded.orientation.x, @@ -306,9 +321,9 @@ def test_pose_with_mixed_types() -> None: assert pose1.orientation.w == pose2.orientation.w -def test_pose_nested_value_copy() -> None: - source = Pose(position=Point(x=1, y=2, z=3)) - copied = Pose(position=source.position, orientation=source.orientation) +def test_pose_explicit_nested_value_copy() -> None: + source = Pose(position=Point(x=1, y=2, z=3), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)) + copied = deepcopy(source) copied.position.x = 10 assert source.position == Point(x=1, y=2, z=3) assert copied.position == Point(x=10, y=2, z=3) @@ -361,7 +376,9 @@ def test_pose_euler_roundtrip() -> None: def test_pose_zero_position() -> None: """Test pose with zero position vector.""" - pose = Pose(position=Point(x=0.0, y=0.0, z=0.0)) + pose = Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) assert pose.position.x == 0.0 assert pose.position.y == 0.0 assert pose.position.z == 0.0 @@ -372,15 +389,24 @@ def test_pose_zero_position() -> None: def test_pose_unit_vectors() -> None: """Test pose with unit vector positions.""" - pose_x = Pose(position=point_from_array(vector_array(Vector3(x=1)))) + pose_x = Pose( + position=point_from_array(vector_array(Vector3(x=1, y=0.0, z=0.0))), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ) assert pose_x.position.x == 1.0 assert pose_x.position.y == 0.0 assert pose_x.position.z == 0.0 - pose_y = Pose(position=point_from_array(vector_array(Vector3(y=1)))) + pose_y = Pose( + position=point_from_array(vector_array(Vector3(y=1, x=0.0, z=0.0))), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ) assert pose_y.position.x == 0.0 assert pose_y.position.y == 1.0 assert pose_y.position.z == 0.0 - pose_z = Pose(position=point_from_array(vector_array(Vector3(z=1)))) + pose_z = Pose( + position=point_from_array(vector_array(Vector3(z=1, x=0.0, y=0.0))), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ) assert pose_z.position.x == 0.0 assert pose_z.position.y == 0.0 assert pose_z.position.z == 1.0 @@ -404,7 +430,10 @@ def test_pose_negative_coordinates() -> None: def test_pose_large_coordinates() -> None: """Test pose with large coordinate values.""" large_value = 1000.0 - pose = Pose(position=Point(x=large_value, y=large_value, z=large_value)) + pose = Pose( + position=Point(x=large_value, y=large_value, z=large_value), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ) assert pose.position.x == large_value assert pose.position.y == large_value assert pose.position.z == large_value @@ -420,7 +449,7 @@ def test_pose_large_coordinates() -> None: ) def test_pose_parametrized_positions(x, y, z) -> None: """Parametrized test for various position values.""" - pose = Pose(position=Point(x=x, y=y, z=z)) + pose = Pose(position=Point(x=x, y=y, z=z), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)) assert pose.position.x == x assert pose.position.y == y assert pose.position.z == z @@ -459,8 +488,8 @@ def encodepass() -> None: pose_source = Pose( position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) ) - binary_msg = pose_source.encode() - pose_dest = Pose.decode(binary_msg) + binary_msg = cdr_encode(pose_source) + pose_dest = cdr_decode(binary_msg, Pose) assert isinstance(pose_dest, Pose) assert pose_dest is not pose_source assert pose_dest == pose_source @@ -492,8 +521,12 @@ def encodepass() -> None: def test_pose_addition_translation_only() -> None: """Test pose addition with translation only (identity rotations).""" - pose1 = Pose(position=Point(x=1.0, y=2.0, z=3.0)) - pose2 = Pose(position=Point(x=4.0, y=5.0, z=6.0)) + pose1 = Pose( + position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) + pose2 = Pose( + position=Point(x=4.0, y=5.0, z=6.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) result = compose_poses(pose1, pose2) assert result.position.x == 5.0 assert result.position.y == 7.0 @@ -511,7 +544,9 @@ def test_pose_addition_with_rotation() -> None: position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=np.sin(angle / 2), w=np.cos(angle / 2)), ) - pose2 = Pose(position=Point(x=1.0, y=0.0, z=0.0)) + pose2 = Pose( + position=Point(x=1.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) result = compose_poses(pose1, pose2) assert np.isclose(result.position.x, 0.0, atol=1e-10) assert np.isclose(result.position.y, 1.0, atol=1e-10) @@ -549,7 +584,9 @@ def test_pose_addition_full_transform() -> None: position=Point(x=2.0, y=1.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=np.sin(robot_yaw / 2), w=np.cos(robot_yaw / 2)), ) - object_in_robot = Pose(position=Point(x=3.0, y=-1.0, z=0.0)) + object_in_robot = Pose( + position=Point(x=3.0, y=-1.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) object_in_world = compose_poses(robot_pose, object_in_robot) assert np.isclose(object_in_world.position.x, 3.0, atol=1e-10) assert np.isclose(object_in_world.position.y, 4.0, atol=1e-10) @@ -559,9 +596,15 @@ def test_pose_addition_full_transform() -> None: def test_pose_addition_chain() -> None: """Test chaining multiple pose additions.""" - pose1 = Pose(position=Point(x=1.0, y=0.0, z=0.0)) - pose2 = Pose(position=Point(x=0.0, y=1.0, z=0.0)) - pose3 = Pose(position=Point(x=0.0, y=0.0, z=1.0)) + pose1 = Pose( + position=Point(x=1.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) + pose2 = Pose( + position=Point(x=0.0, y=1.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) + pose3 = Pose( + position=Point(x=0.0, y=0.0, z=1.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) result = compose_poses(compose_poses(pose1, pose2), pose3) assert result.position.x == 1.0 assert result.position.y == 1.0 @@ -570,7 +613,9 @@ def test_pose_addition_chain() -> None: def test_pose_addition_with_convertible() -> None: """Test pose addition with convertible types.""" - pose1 = Pose(position=Point(x=1.0, y=2.0, z=3.0)) + pose1 = Pose( + position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) pose_tuple = ([4.0, 5.0, 6.0], [0.0, 0.0, 0.0, 1.0]) result1 = compose_poses( pose1, @@ -600,7 +645,9 @@ def test_pose_identity_addition() -> None: pose = Pose( position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) ) - identity = Pose() + identity = Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) result = compose_poses(pose, identity) assert result.position.x == pose.position.x assert result.position.y == pose.position.y @@ -618,7 +665,9 @@ def test_pose_addition_3d_rotation() -> None: position=Point(x=1.0, y=0.0, z=0.0), orientation=Quaternion(x=np.sin(roll / 2), y=0.0, z=0.0, w=np.cos(roll / 2)), ) - pose2 = Pose(position=Point(x=0.0, y=1.0, z=1.0)) + pose2 = Pose( + position=Point(x=0.0, y=1.0, z=1.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) result = compose_poses(pose1, pose2) cos45 = np.cos(roll) sin45 = np.sin(roll) diff --git a/dimos/msgs/geometry_msgs/test_PoseStamped.py b/dimos/msgs/geometry_msgs/test_PoseStamped.py index d43b2981eaa..5e6a55744d8 100644 --- a/dimos/msgs/geometry_msgs/test_PoseStamped.py +++ b/dimos/msgs/geometry_msgs/test_PoseStamped.py @@ -14,6 +14,7 @@ import pickle +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.std_msgs.msg import Header from rosbags.typesys import Stores, get_typestore @@ -23,17 +24,17 @@ def test_cdr_encode_decode() -> None: source = PoseStamped( - header=Header(stamp=time_from_nanoseconds(1234567890123456789)), + header=Header(stamp=time_from_nanoseconds(1234567890123456789), frame_id=""), pose=Pose( position=Point(x=1, y=2, z=3), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) ), ) - binary = source.encode() - decoded = PoseStamped.decode(binary) + binary = cdr_encode(source) + decoded = cdr_decode(binary, PoseStamped) assert isinstance(decoded, PoseStamped) assert decoded is not source - assert decoded.encode() == binary - independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(binary, source.msg_name) + assert cdr_encode(decoded) == binary + independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(binary, source.__msgtype__) assert (independent.header.stamp.sec, independent.header.stamp.nanosec) == ( 1234567890, 123456789, @@ -53,7 +54,7 @@ def test_cdr_encode_decode() -> None: def test_pickle_encode_decode() -> None: source = PoseStamped( - header=Header(stamp=time_from_nanoseconds(1234567890123456789)), + header=Header(stamp=time_from_nanoseconds(1234567890123456789), frame_id=""), pose=Pose( position=Point(x=1, y=2, z=3), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) ), @@ -61,4 +62,4 @@ def test_pickle_encode_decode() -> None: decoded = pickle.loads(pickle.dumps(source)) assert isinstance(decoded, PoseStamped) assert decoded is not source - assert decoded.encode() == source.encode() + assert cdr_encode(decoded) == cdr_encode(source) diff --git a/dimos/msgs/geometry_msgs/test_PoseWithCovariance.py b/dimos/msgs/geometry_msgs/test_PoseWithCovariance.py index fa52d25c80a..056687052ec 100644 --- a/dimos/msgs/geometry_msgs/test_PoseWithCovariance.py +++ b/dimos/msgs/geometry_msgs/test_PoseWithCovariance.py @@ -17,14 +17,23 @@ Legacy polymorphic positional constructors and presentation methods are retired. """ +from dataclasses import asdict + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import Point, Pose, PoseWithCovariance, Quaternion import numpy as np import pytest +from rosbags.serde import SerdeError from rosbags.typesys import Stores, get_typestore def test_default_fields_and_covariance() -> None: - source = PoseWithCovariance() + source = PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + covariance=np.zeros(36, dtype=np.float64), + ) value = source.pose np.testing.assert_array_equal( [ @@ -52,11 +61,11 @@ def test_explicit_construction_and_independent_cdr(as_list: bool, covariance: np pose=Pose( position=Point(x=1, y=2, z=3), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) ), - covariance=covariance.tolist() if as_list else covariance, + covariance=np.asarray(covariance.tolist() if as_list else covariance, dtype=np.float64), ) - decoded = PoseWithCovariance.decode(source.encode()) + decoded = cdr_decode(cdr_encode(source), PoseWithCovariance) independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr( - source.encode(), PoseWithCovariance.msg_name + cdr_encode(source), PoseWithCovariance.__msgtype__ ) for result in (source, decoded, independent): value = result.pose @@ -85,22 +94,27 @@ def test_copy_equality_and_independent_storage() -> None: pose=Pose( position=Point(x=1, y=2, z=3), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) ), - covariance=np.arange(36, dtype=float), + covariance=np.asarray(np.arange(36, dtype=float), dtype=np.float64), ) - copied = PoseWithCovariance.decode(original.encode()) - assert copied == original + copied = cdr_decode(cdr_encode(original), PoseWithCovariance) + np.testing.assert_equal(asdict(copied), asdict(original)) assert copied is not original assert copied.pose is not original.pose assert copied.covariance is not original.covariance original.covariance[0] = 999 assert copied.covariance[0] == 0 - assert copied != original + assert not np.array_equal(copied.covariance, original.covariance) assert copied != "not a message" assert copied is not None def test_matrix_assignment() -> None: - source = PoseWithCovariance(covariance=np.arange(36, dtype=float)) + source = PoseWithCovariance( + covariance=np.asarray(np.arange(36, dtype=float), dtype=np.float64), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + ) matrix = np.asarray(source.covariance).reshape(6, 6) assert matrix[0, 0] == 0 assert matrix[5, 5] == 35 @@ -110,13 +124,26 @@ def test_matrix_assignment() -> None: @pytest.mark.parametrize("size", [0, 35, 37]) -def test_fixed_covariance_rejects_invalid_length(size: int) -> None: - with pytest.raises((ValueError, TypeError, RuntimeError)): - PoseWithCovariance(covariance=np.zeros(size)) +def test_cdr_rejects_invalid_covariance_length(size: int) -> None: + source = PoseWithCovariance( + covariance=np.asarray(np.zeros(size), dtype=np.float64), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + with pytest.raises(SerdeError, match="Unexpected array length"): + cdr_encode(source) @pytest.mark.parametrize("xyz", [(0, 0, 0), (1, 2, 3), (-1, -2, -3), (100, -100, 0)]) def test_parameterized_values(xyz: tuple[float, float, float]) -> None: - source = PoseWithCovariance(pose=Pose(position=Point(x=xyz[0], y=xyz[1], z=xyz[2]))) + source = PoseWithCovariance( + pose=Pose( + position=Point(x=xyz[0], y=xyz[1], z=xyz[2]), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ) position = source.pose.position assert (position.x, position.y, position.z) == xyz diff --git a/dimos/msgs/geometry_msgs/test_PoseWithCovarianceStamped.py b/dimos/msgs/geometry_msgs/test_PoseWithCovarianceStamped.py index e2c1a33697d..cf56932887f 100644 --- a/dimos/msgs/geometry_msgs/test_PoseWithCovarianceStamped.py +++ b/dimos/msgs/geometry_msgs/test_PoseWithCovarianceStamped.py @@ -21,6 +21,8 @@ import pickle import time +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, Pose, @@ -37,7 +39,16 @@ def test_generated_defaults() -> None: - source = PoseWithCovarianceStamped() + source = PoseWithCovarianceStamped( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), + ) assert source.header.frame_id == "" assert to_nanoseconds(source.header.stamp) == 0 assert isinstance(source.pose, PoseWithCovariance) @@ -70,12 +81,12 @@ def test_fields_covariance_and_independent_cdr(stamp: int, covariance: np.ndarra pose=Pose( position=Point(x=1, y=2, z=3), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.9) ), - covariance=covariance, + covariance=np.asarray(covariance, dtype=np.float64), ), ) - decoded = PoseWithCovarianceStamped.decode(source.encode()) + decoded = cdr_decode(cdr_encode(source), PoseWithCovarianceStamped) independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr( - source.encode(), PoseWithCovarianceStamped.msg_name + cdr_encode(source), PoseWithCovarianceStamped.__msgtype__ ) for result in (decoded, independent): assert result.header.frame_id == "camera_link" @@ -98,14 +109,23 @@ def test_fields_covariance_and_independent_cdr(stamp: int, covariance: np.ndarra assert matrix.shape == (6, 6) assert np.trace(matrix) == np.trace(covariance.reshape(6, 6)) copied = pickle.loads(pickle.dumps(source)) - assert copied.encode() == source.encode() + assert cdr_encode(copied) == cdr_encode(source) copied.pose.covariance[0] = 999 assert source.pose.covariance[0] == covariance[0] def test_explicit_current_header() -> None: before = time.time_ns() - source = PoseWithCovarianceStamped(header=header_now("test")) + source = PoseWithCovarianceStamped( + header=header_now("test"), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), + ) assert before <= to_nanoseconds(source.header.stamp) <= time.time_ns() assert source.header.frame_id == "test" diff --git a/dimos/msgs/geometry_msgs/test_Quaternion.py b/dimos/msgs/geometry_msgs/test_Quaternion.py index 06ba7e1ee2e..fcf4d78771f 100644 --- a/dimos/msgs/geometry_msgs/test_Quaternion.py +++ b/dimos/msgs/geometry_msgs/test_Quaternion.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import Quaternion, Vector3 import numpy as np import pytest @@ -31,7 +32,7 @@ def test_quaternion_default_init() -> None: """Test that default initialization creates an identity quaternion (w=1, x=y=z=0).""" - q = Quaternion() + q = Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) assert q.x == 0.0 assert q.y == 0.0 assert q.z == 0.0 @@ -61,7 +62,7 @@ def test_quaternion_component_init() -> None: assert q4.y == 2.0 assert q4.z == 3.0 assert q4.w == 4.0 - assert isinstance(q4.x, float) + assert isinstance(cdr_decode(cdr_encode(q4), Quaternion).x, float) def test_quaternion_sequence_init() -> None: @@ -110,7 +111,7 @@ def test_quaternion_numpy_init() -> None: def test_quaternion_copy_init() -> None: """Test initialization from another Quaternion (copy constructor).""" original = Quaternion(x=0.1, y=0.2, z=0.3, w=0.4) - copy = Quaternion.decode(original.encode()) + copy = cdr_decode(cdr_encode(original), Quaternion) assert copy.x == 0.1 assert copy.y == 0.2 assert copy.z == 0.3 @@ -121,12 +122,12 @@ def test_quaternion_copy_init() -> None: def test_quaternion_mutated_value_copy() -> None: """Test initialization from generated Quaternion.""" - source_quaternion = Quaternion() + source_quaternion = Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) source_quaternion.x = 0.1 source_quaternion.y = 0.2 source_quaternion.z = 0.3 source_quaternion.w = 0.4 - q = Quaternion.decode(source_quaternion.encode()) + q = cdr_decode(cdr_encode(source_quaternion), Quaternion) assert q.x == 0.1 assert q.y == 0.2 assert q.z == 0.3 @@ -154,7 +155,7 @@ def test_quaternion_indexing() -> None: def test_quaternion_euler() -> None: """Test quaternion to Euler angles conversion.""" - q_identity = Quaternion() + q_identity = Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) angles = vector_from_array(quaternion_euler(q_identity)) assert np.isclose(angles.x, 0.0, atol=1e-10) assert np.isclose(angles.y, 0.0, atol=1e-10) @@ -174,8 +175,8 @@ def test_quaternion_euler() -> None: def test_cdr_encode_decode() -> None: """Test encoding and decoding of Quaternion to/from binary CDR format.""" q_source = Quaternion(x=1.0, y=2.0, z=3.0, w=4.0) - binary_msg = q_source.encode() - q_dest = Quaternion.decode(binary_msg) + binary_msg = cdr_encode(q_source) + q_dest = cdr_decode(binary_msg, Quaternion) assert isinstance(q_dest, Quaternion) assert q_dest is not q_source assert q_dest == q_source diff --git a/dimos/msgs/geometry_msgs/test_Transform.py b/dimos/msgs/geometry_msgs/test_Transform.py index eb397dbe02c..f830ebb46ba 100644 --- a/dimos/msgs/geometry_msgs/test_Transform.py +++ b/dimos/msgs/geometry_msgs/test_Transform.py @@ -14,6 +14,7 @@ import math +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, @@ -42,7 +43,9 @@ def test_transform_initialization() -> None: - tf = Transform() + tf = Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) assert tf.translation.x == 0.0 assert tf.translation.y == 0.0 assert tf.translation.z == 0.0 @@ -55,33 +58,47 @@ def test_transform_initialization() -> None: tf2 = Transform(translation=trans, rotation=rot) assert tf2.translation == trans assert tf2.rotation == rot - tf5 = Transform(translation=Vector3(x=7.0, y=8.0, z=9.0)) + tf5 = Transform( + translation=Vector3(x=7.0, y=8.0, z=9.0), rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) assert tf5.translation.x == 7.0 assert tf5.translation.y == 8.0 assert tf5.translation.z == 9.0 assert tf5.rotation.w == 1.0 - tf6 = Transform(rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)) + tf6 = Transform( + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), translation=Vector3(x=0.0, y=0.0, z=0.0) + ) assert np.allclose(vector_array(tf6.translation), 0) assert tf6.rotation.w == 1.0 - tf7 = Transform(translation=Vector3(x=1, y=2, z=3), rotation=Quaternion()) + tf7 = Transform( + translation=Vector3(x=1, y=2, z=3), rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) assert tf7.translation == Vector3(x=1, y=2, z=3) - assert tf7.rotation == Quaternion() - tf8 = Transform(translation=Vector3(x=4, y=5, z=6)) + assert tf7.rotation == Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + tf8 = Transform( + translation=Vector3(x=4, y=5, z=6), rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) assert tf8.translation == Vector3(x=4, y=5, z=6) assert tf8.rotation.w == 1.0 - tf9 = Transform(rotation=Quaternion(x=0, y=0, z=1, w=0)) + tf9 = Transform( + rotation=Quaternion(x=0, y=0, z=1, w=0), translation=Vector3(x=0.0, y=0.0, z=0.0) + ) assert np.allclose(vector_array(tf9.translation), 0) assert tf9.rotation == Quaternion(x=0, y=0, z=1, w=0) def test_transform_identity() -> None: - tf = Transform() + tf = Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) assert np.allclose(vector_array(tf.translation), 0) assert tf.rotation.x == 0.0 assert tf.rotation.y == 0.0 assert tf.rotation.z == 0.0 assert tf.rotation.w == 1.0 - assert tf == Transform() + assert tf == Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) def test_transform_equality() -> None: @@ -97,9 +114,12 @@ def test_transform_equality() -> None: def test_transform_independent_cdr_fields() -> None: source = Transform( - translation=Vector3(x=1.5, y=-2.0, z=3.14), rotation=Quaternion(z=0.707107, w=0.707107) + translation=Vector3(x=1.5, y=-2.0, z=3.14), + rotation=Quaternion(z=0.707107, w=0.707107, x=0.0, y=0.0), + ) + decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr( + cdr_encode(source), Transform.__msgtype__ ) - decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(source.encode(), Transform.msg_name) assert [decoded.translation.x, decoded.translation.y, decoded.translation.z] == [ 1.5, -2.0, @@ -114,7 +134,9 @@ def test_transform_independent_cdr_fields() -> None: def test_pose_add_transform() -> None: - initial_pose = Pose(position=Point(x=1.0, y=0.0, z=0.0)) + initial_pose = Pose( + position=Point(x=1.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) angle = np.pi / 2 transform = Transform( translation=Vector3(x=2.0, y=1.0, z=0.0), @@ -180,8 +202,8 @@ def test_encode_decode() -> None: translation=Vector3(x=2.0, y=1.0, z=0.0), rotation=Quaternion(x=0.0, y=0.0, z=np.sin(angle / 2), w=np.cos(angle / 2)), ) - data = transform.encode() - decoded_transform = Transform.decode(data) + data = cdr_encode(transform) + decoded_transform = cdr_decode(data, Transform) assert decoded_transform == transform @@ -218,27 +240,34 @@ def test_transform_addition() -> None: def test_transform_frame_tracking() -> None: first = TransformStamped( - header=Header(frame_id="world"), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), child_frame_id="robot", - transform=Transform(translation=Vector3(x=1)), + transform=Transform( + translation=Vector3(x=1, y=0.0, z=0.0), rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), ) second = TransformStamped( - header=Header(frame_id="robot"), + header=Header(frame_id="robot", stamp=Time(sec=0, nanosec=0)), child_frame_id="sensor", - transform=Transform(translation=Vector3(x=2)), + transform=Transform( + translation=Vector3(x=2, y=0.0, z=0.0), rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), ) result = compose_transforms(first, second) assert result.header.frame_id == "world" assert result.child_frame_id == "sensor" - assert result.transform.translation == Vector3(x=3) + assert result.transform.translation == Vector3(x=3, y=0.0, z=0.0) with pytest.raises(TypeError): compose_transform_values(first.transform, "not a transform") def test_transform_from_pose() -> None: - pose = Pose(position=Point(x=1, y=2, z=3), orientation=Quaternion(z=0.707, w=0.707)) + pose = Pose( + position=Point(x=1, y=2, z=3), orientation=Quaternion(z=0.707, w=0.707, x=0.0, y=0.0) + ) result = transform_from_pose( - PoseStamped(header=Header(frame_id="world"), pose=pose), child_frame_id="base_link" + PoseStamped(header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), pose=pose), + child_frame_id="base_link", ) np.testing.assert_array_equal( vector_array(result.transform.translation), vector_array(pose.position) @@ -254,7 +283,8 @@ def test_transform_from_ros() -> None: PoseStamped( header=Header(stamp=stamp, frame_id="base_link"), pose=Pose( - position=Point(x=1, y=-1), orientation=quaternion_from_euler(0, 0, math.pi / 6) + position=Point(x=1, y=-1, z=0.0), + orientation=quaternion_from_euler(0, 0, math.pi / 6), ), ), child_frame_id="arm_base_link", @@ -263,7 +293,8 @@ def test_transform_from_ros() -> None: PoseStamped( header=Header(stamp=stamp, frame_id="arm_base_link"), pose=Pose( - position=Point(x=1, y=1), orientation=quaternion_from_euler(0, 0, math.pi / 6) + position=Point(x=1, y=1, z=0.0), + orientation=quaternion_from_euler(0, 0, math.pi / 6), ), ), child_frame_id="end", @@ -277,7 +308,9 @@ def test_transform_from_ros() -> None: def test_transform_from_pose_stamped() -> None: pose = PoseStamped( header=Header(stamp=Time(sec=123, nanosec=456789123), frame_id="map"), - pose=Pose(position=Point(x=4, y=5, z=6), orientation=Quaternion(y=0.707, w=0.707)), + pose=Pose( + position=Point(x=4, y=5, z=6), orientation=Quaternion(y=0.707, w=0.707, x=0.0, z=0.0) + ), ) result = transform_from_pose(pose, child_frame_id="robot_base") np.testing.assert_array_equal( @@ -290,10 +323,15 @@ def test_transform_from_pose_stamped() -> None: @pytest.mark.parametrize("values", [[1.0, 2.0, 3.0], (7.0, 8.0, 9.0), [10.0, 11.0, 12.0]]) def test_transform_from_coordinate_arrays(values) -> None: - pose = PoseStamped(pose=Pose(position=point_from_array(values))) + pose = PoseStamped( + pose=Pose( + position=point_from_array(values), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) result = transform_from_pose(pose, child_frame_id="base_link") np.testing.assert_array_equal(vector_array(result.transform.translation), values) - assert result.transform.rotation == Quaternion() + assert result.transform.rotation == Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) @pytest.mark.parametrize("value", ["not a pose", 42, None]) diff --git a/dimos/msgs/geometry_msgs/test_Twist.py b/dimos/msgs/geometry_msgs/test_Twist.py index bff834b4c57..0dc7b6c95bf 100644 --- a/dimos/msgs/geometry_msgs/test_Twist.py +++ b/dimos/msgs/geometry_msgs/test_Twist.py @@ -14,6 +14,7 @@ from copy import deepcopy +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import Quaternion, Twist, Vector3 import numpy as np from rosbags.typesys import Stores, get_typestore @@ -22,7 +23,7 @@ def test_twist_initialization() -> None: - tw = Twist() + tw = Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) assert tw.linear.x == 0.0 assert tw.linear.y == 0.0 assert tw.linear.z == 0.0 @@ -40,10 +41,10 @@ def test_twist_initialization() -> None: assert tw3 == tw2 tw3.linear.x = 10.0 assert tw2.linear.x == 1.0 - source = Twist() + source = Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) source.linear = Vector3(x=4.0, y=5.0, z=6.0) source.angular = Vector3(x=0.4, y=0.5, z=0.6) - tw4 = Twist.decode(source.encode()) + tw4 = cdr_decode(cdr_encode(source), Twist) assert tw4.linear.x == 4.0 assert tw4.linear.y == 5.0 assert tw4.linear.z == 6.0 @@ -62,14 +63,15 @@ def test_twist_initialization() -> None: tw7 = Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.1, y=0.2, z=0.3)) assert tw7.linear == Vector3(x=1, y=2, z=3) assert tw7.angular == Vector3(x=0.1, y=0.2, z=0.3) - tw8 = Twist(linear=Vector3(x=4, y=5, z=6)) + tw8 = Twist(linear=Vector3(x=4, y=5, z=6), angular=Vector3(x=0.0, y=0.0, z=0.0)) assert tw8.linear == Vector3(x=4, y=5, z=6) assert np.allclose(vector_array(tw8.angular), 0) - tw9 = Twist(angular=Vector3(x=0.4, y=0.5, z=0.6)) + tw9 = Twist(angular=Vector3(x=0.4, y=0.5, z=0.6), linear=Vector3(x=0.0, y=0.0, z=0.0)) assert np.allclose(vector_array(tw9.linear), 0) assert tw9.angular == Vector3(x=0.4, y=0.5, z=0.6) tw10 = Twist( - angular=vector_from_array(quaternion_euler(Quaternion(x=0, y=0, z=0.707107, w=0.707107))) + angular=vector_from_array(quaternion_euler(Quaternion(x=0, y=0, z=0.707107, w=0.707107))), + linear=Vector3(x=0.0, y=0.0, z=0.0), ) assert np.allclose(vector_array(tw10.linear), 0) euler = vector_from_array(quaternion_euler(Quaternion(x=0, y=0, z=0.707107, w=0.707107))) @@ -85,11 +87,11 @@ def test_twist_initialization() -> None: def test_twist_zero() -> None: - tw = Twist() + tw = Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) assert np.allclose(vector_array(tw.linear), 0) assert np.allclose(vector_array(tw.angular), 0) assert np.allclose(np.concatenate((vector_array(tw.linear), vector_array(tw.angular))), 0) - assert tw == Twist() + assert tw == Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) def test_twist_equality() -> None: @@ -106,29 +108,29 @@ def test_twist_equality() -> None: def test_twist_independent_cdr_fields() -> None: tw = Twist(linear=Vector3(x=1.5, y=-2.0, z=3.14), angular=Vector3(x=0.1, y=-0.2, z=0.3)) decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr( - tw.encode(), "geometry_msgs/msg/Twist" + cdr_encode(tw), "geometry_msgs/msg/Twist" ) assert [decoded.linear.x, decoded.linear.y, decoded.linear.z] == [1.5, -2.0, 3.14] assert [decoded.angular.x, decoded.angular.y, decoded.angular.z] == [0.1, -0.2, 0.3] def test_twist_is_zero() -> None: - tw1 = Twist() + tw1 = Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) assert np.allclose(np.concatenate((vector_array(tw1.linear), vector_array(tw1.angular))), 0) - tw2 = Twist(linear=Vector3(x=0.1, y=0, z=0)) + tw2 = Twist(linear=Vector3(x=0.1, y=0, z=0), angular=Vector3(x=0.0, y=0.0, z=0.0)) assert not np.allclose(np.concatenate((vector_array(tw2.linear), vector_array(tw2.angular))), 0) - tw3 = Twist(angular=Vector3(x=0, y=0, z=0.1)) + tw3 = Twist(angular=Vector3(x=0, y=0, z=0.1), linear=Vector3(x=0.0, y=0.0, z=0.0)) assert not np.allclose(np.concatenate((vector_array(tw3.linear), vector_array(tw3.angular))), 0) tw4 = Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.1, y=0.2, z=0.3)) assert not np.allclose(np.concatenate((vector_array(tw4.linear), vector_array(tw4.angular))), 0) def test_twist_explicit_nonzero_predicate() -> None: - tw1 = Twist() + tw1 = Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) assert not (np.any(vector_array(tw1.linear)) or np.any(vector_array(tw1.angular))) - tw2 = Twist(linear=Vector3(x=1, y=0, z=0)) + tw2 = Twist(linear=Vector3(x=1, y=0, z=0), angular=Vector3(x=0.0, y=0.0, z=0.0)) assert np.any(vector_array(tw2.linear)) or np.any(vector_array(tw2.angular)) - tw3 = Twist(angular=Vector3(x=0, y=0, z=0.1)) + tw3 = Twist(angular=Vector3(x=0, y=0, z=0.1), linear=Vector3(x=0.0, y=0.0, z=0.0)) assert np.any(vector_array(tw3.linear)) or np.any(vector_array(tw3.angular)) tw4 = Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.1, y=0.2, z=0.3)) assert np.any(vector_array(tw4.linear)) or np.any(vector_array(tw4.angular)) @@ -136,9 +138,9 @@ def test_twist_explicit_nonzero_predicate() -> None: def test_twist_cdr_encoding() -> None: tw = Twist(linear=Vector3(x=1.5, y=2.5, z=3.5), angular=Vector3(x=0.1, y=0.2, z=0.3)) - encoded = tw.encode() + encoded = cdr_encode(tw) assert isinstance(encoded, bytes) - decoded = Twist.decode(encoded) + decoded = cdr_decode(encoded, Twist) assert decoded.linear == tw.linear assert decoded.angular == tw.angular assert isinstance(decoded.linear, Vector3) diff --git a/dimos/msgs/geometry_msgs/test_TwistStamped.py b/dimos/msgs/geometry_msgs/test_TwistStamped.py index 590f31291d8..90db22dfb31 100644 --- a/dimos/msgs/geometry_msgs/test_TwistStamped.py +++ b/dimos/msgs/geometry_msgs/test_TwistStamped.py @@ -14,6 +14,7 @@ import pickle +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import Twist, TwistStamped, Vector3 from dimos_generated.std_msgs.msg import Header from rosbags.typesys import Stores, get_typestore @@ -23,15 +24,15 @@ def test_cdr_encode_decode() -> None: source = TwistStamped( - header=Header(stamp=time_from_nanoseconds(1234567890123456789)), + header=Header(stamp=time_from_nanoseconds(1234567890123456789), frame_id=""), twist=Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.1, y=0.2, z=0.3)), ) - binary = source.encode() - decoded = TwistStamped.decode(binary) + binary = cdr_encode(source) + decoded = cdr_decode(binary, TwistStamped) assert isinstance(decoded, TwistStamped) assert decoded is not source - assert decoded.encode() == binary - independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(binary, source.msg_name) + assert cdr_encode(decoded) == binary + independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(binary, source.__msgtype__) assert (independent.header.stamp.sec, independent.header.stamp.nanosec) == ( 1234567890, 123456789, @@ -50,10 +51,10 @@ def test_cdr_encode_decode() -> None: def test_pickle_encode_decode() -> None: source = TwistStamped( - header=Header(stamp=time_from_nanoseconds(1234567890123456789)), + header=Header(stamp=time_from_nanoseconds(1234567890123456789), frame_id=""), twist=Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.1, y=0.2, z=0.3)), ) decoded = pickle.loads(pickle.dumps(source)) assert isinstance(decoded, TwistStamped) assert decoded is not source - assert decoded.encode() == source.encode() + assert cdr_encode(decoded) == cdr_encode(source) diff --git a/dimos/msgs/geometry_msgs/test_TwistWithCovariance.py b/dimos/msgs/geometry_msgs/test_TwistWithCovariance.py index 2814a6667c3..7ac29bb9276 100644 --- a/dimos/msgs/geometry_msgs/test_TwistWithCovariance.py +++ b/dimos/msgs/geometry_msgs/test_TwistWithCovariance.py @@ -17,14 +17,21 @@ Legacy polymorphic positional constructors and presentation methods are retired. """ +from dataclasses import asdict + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import Twist, TwistWithCovariance, Vector3 import numpy as np import pytest +from rosbags.serde import SerdeError from rosbags.typesys import Stores, get_typestore def test_default_fields_and_covariance() -> None: - source = TwistWithCovariance() + source = TwistWithCovariance( + twist=Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + covariance=np.zeros(36, dtype=np.float64), + ) value = source.twist np.testing.assert_array_equal( [ @@ -49,11 +56,11 @@ def test_default_fields_and_covariance() -> None: def test_explicit_construction_and_independent_cdr(as_list: bool, covariance: np.ndarray) -> None: source = TwistWithCovariance( twist=Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.1, y=0.2, z=0.3)), - covariance=covariance.tolist() if as_list else covariance, + covariance=np.asarray(covariance.tolist() if as_list else covariance, dtype=np.float64), ) - decoded = TwistWithCovariance.decode(source.encode()) + decoded = cdr_decode(cdr_encode(source), TwistWithCovariance) independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr( - source.encode(), TwistWithCovariance.msg_name + cdr_encode(source), TwistWithCovariance.__msgtype__ ) for result in (source, decoded, independent): value = result.twist @@ -79,22 +86,25 @@ def test_explicit_construction_and_independent_cdr(as_list: bool, covariance: np def test_copy_equality_and_independent_storage() -> None: original = TwistWithCovariance( twist=Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.1, y=0.2, z=0.3)), - covariance=np.arange(36, dtype=float), + covariance=np.asarray(np.arange(36, dtype=float), dtype=np.float64), ) - copied = TwistWithCovariance.decode(original.encode()) - assert copied == original + copied = cdr_decode(cdr_encode(original), TwistWithCovariance) + np.testing.assert_equal(asdict(copied), asdict(original)) assert copied is not original assert copied.twist is not original.twist assert copied.covariance is not original.covariance original.covariance[0] = 999 assert copied.covariance[0] == 0 - assert copied != original + assert not np.array_equal(copied.covariance, original.covariance) assert copied != "not a message" assert copied is not None def test_matrix_assignment() -> None: - source = TwistWithCovariance(covariance=np.arange(36, dtype=float)) + source = TwistWithCovariance( + covariance=np.asarray(np.arange(36, dtype=float), dtype=np.float64), + twist=Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + ) matrix = np.asarray(source.covariance).reshape(6, 6) assert matrix[0, 0] == 0 assert matrix[5, 5] == 35 @@ -104,9 +114,13 @@ def test_matrix_assignment() -> None: @pytest.mark.parametrize("size", [0, 35, 37]) -def test_fixed_covariance_rejects_invalid_length(size: int) -> None: - with pytest.raises((ValueError, TypeError, RuntimeError)): - TwistWithCovariance(covariance=np.zeros(size)) +def test_cdr_rejects_invalid_covariance_length(size: int) -> None: + source = TwistWithCovariance( + covariance=np.asarray(np.zeros(size), dtype=np.float64), + twist=Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + ) + with pytest.raises(SerdeError, match="Unexpected array length"): + cdr_encode(source) @pytest.mark.parametrize( @@ -125,7 +139,8 @@ def test_parameterized_values( twist=Twist( linear=Vector3(x=xyz[0], y=xyz[1], z=xyz[2]), angular=Vector3(x=angular[0], y=angular[1], z=angular[2]), - ) + ), + covariance=np.zeros(36, dtype=np.float64), ) assert (source.twist.linear.x, source.twist.linear.y, source.twist.linear.z) == xyz assert (source.twist.angular.x, source.twist.angular.y, source.twist.angular.z) == angular diff --git a/dimos/msgs/geometry_msgs/test_TwistWithCovarianceStamped.py b/dimos/msgs/geometry_msgs/test_TwistWithCovarianceStamped.py index da7c1382ba0..c76698d0492 100644 --- a/dimos/msgs/geometry_msgs/test_TwistWithCovarianceStamped.py +++ b/dimos/msgs/geometry_msgs/test_TwistWithCovarianceStamped.py @@ -21,6 +21,8 @@ import pickle import time +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Twist, TwistWithCovariance, @@ -36,7 +38,13 @@ def test_generated_defaults() -> None: - source = TwistWithCovarianceStamped() + source = TwistWithCovarianceStamped( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + twist=TwistWithCovariance( + twist=Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + covariance=np.zeros(36, dtype=np.float64), + ), + ) assert source.header.frame_id == "" assert to_nanoseconds(source.header.stamp) == 0 assert isinstance(source.twist, TwistWithCovariance) @@ -66,12 +74,12 @@ def test_fields_covariance_and_independent_cdr(stamp: int, covariance: np.ndarra header=Header(stamp=time_from_nanoseconds(stamp), frame_id="camera_link"), twist=TwistWithCovariance( twist=Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.1, y=0.2, z=0.3)), - covariance=covariance, + covariance=np.asarray(covariance, dtype=np.float64), ), ) - decoded = TwistWithCovarianceStamped.decode(source.encode()) + decoded = cdr_decode(cdr_encode(source), TwistWithCovarianceStamped) independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr( - source.encode(), TwistWithCovarianceStamped.msg_name + cdr_encode(source), TwistWithCovarianceStamped.__msgtype__ ) for result in (decoded, independent): assert result.header.frame_id == "camera_link" @@ -93,14 +101,20 @@ def test_fields_covariance_and_independent_cdr(stamp: int, covariance: np.ndarra assert matrix.shape == (6, 6) assert np.trace(matrix) == np.trace(covariance.reshape(6, 6)) copied = pickle.loads(pickle.dumps(source)) - assert copied.encode() == source.encode() + assert cdr_encode(copied) == cdr_encode(source) copied.twist.covariance[0] = 999 assert source.twist.covariance[0] == covariance[0] def test_explicit_current_header() -> None: before = time.time_ns() - source = TwistWithCovarianceStamped(header=header_now("test")) + source = TwistWithCovarianceStamped( + header=header_now("test"), + twist=TwistWithCovariance( + twist=Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + covariance=np.zeros(36, dtype=np.float64), + ), + ) assert before <= to_nanoseconds(source.header.stamp) <= time.time_ns() assert source.header.frame_id == "test" diff --git a/dimos/msgs/geometry_msgs/test_Vector3.py b/dimos/msgs/geometry_msgs/test_Vector3.py index 05941ddea08..2f7f85001c1 100644 --- a/dimos/msgs/geometry_msgs/test_Vector3.py +++ b/dimos/msgs/geometry_msgs/test_Vector3.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import Vector3 import numpy as np import pytest @@ -27,7 +28,7 @@ def test_vector_default_init() -> None: """Test that default initialization of Vector() has x,y,z components all zero.""" - v = Vector3() + v = Vector3(x=0.0, y=0.0, z=0.0) assert v.x == 0.0 assert v.y == 0.0 assert v.z == 0.0 @@ -38,7 +39,7 @@ def test_vector_default_init() -> None: def test_vector_specific_init() -> None: """Test initialization with specific values and different input types.""" - v1 = Vector3(x=1.0, y=2.0) + v1 = Vector3(x=1.0, y=2.0, z=0.0) assert v1.x == 1.0 assert v1.y == 2.0 assert v1.z == 0.0 @@ -59,7 +60,7 @@ def test_vector_specific_init() -> None: assert v5.y == 13.0 assert v5.z == 14.0 original = vector_from_array([15.0, 16.0, 17.0]) - v6 = Vector3.decode(original.encode()) + v6 = cdr_decode(cdr_encode(original), Vector3) assert v6.x == 15.0 assert v6.y == 16.0 assert v6.z == 17.0 @@ -119,7 +120,7 @@ def test_vector_dot_product() -> None: def test_vector_length() -> None: """Test vector length calculation.""" - v1 = Vector3(x=3.0, y=4.0) + v1 = Vector3(x=3.0, y=4.0, z=0.0) assert np.linalg.norm(vector_array(v1)) == 5.0 v2 = Vector3(x=2.0, y=3.0, z=6.0) assert np.linalg.norm(vector_array(v2)) == pytest.approx(7.0, 0.001) @@ -153,12 +154,12 @@ def test_vector_normalize() -> None: def test_vector_to_2d() -> None: """Test conversion to 2D vector.""" v = Vector3(x=2.0, y=3.0, z=6.0) - v_2d = Vector3(x=v.x, y=v.y) + v_2d = Vector3(x=v.x, y=v.y, z=0.0) assert v_2d.x == 2.0 assert v_2d.y == 3.0 assert v_2d.z == 0.0 - v2 = Vector3(x=4.0, y=5.0) - v2_2d = Vector3(x=v2.x, y=v2.y) + v2 = Vector3(x=4.0, y=5.0, z=0.0) + v2_2d = Vector3(x=v2.x, y=v2.y, z=0.0) assert v2_2d.x == 4.0 assert v2_2d.y == 5.0 assert v2_2d.z == 0.0 @@ -189,8 +190,8 @@ def test_vector_cross_product() -> None: assert c.x == -3.0 assert c.y == 6.0 assert c.z == -3.0 - v_2d1 = Vector3(x=1.0, y=2.0) - v_2d2 = Vector3(x=3.0, y=4.0) + v_2d1 = Vector3(x=1.0, y=2.0, z=0.0) + v_2d2 = Vector3(x=3.0, y=4.0, z=0.0) cross_2d = vector_from_array(np.cross(vector_array(v_2d1), vector_array(v_2d2))) assert cross_2d.x == 0.0 assert cross_2d.y == 0.0 @@ -199,7 +200,7 @@ def test_vector_cross_product() -> None: def test_vector_zeros() -> None: """Test Vector3.zeros class method.""" - v_zeros = Vector3() + v_zeros = Vector3(x=0.0, y=0.0, z=0.0) assert v_zeros.x == 0.0 assert v_zeros.y == 0.0 assert v_zeros.z == 0.0 @@ -231,18 +232,18 @@ def test_vector_equality() -> None: v3 = Vector3(x=4, y=5, z=6) assert v1 == v2 assert v1 != v3 - assert v1 != Vector3(x=1, y=2) + assert v1 != Vector3(x=1, y=2, z=0.0) assert v1 != Vector3(x=1.1, y=2, z=3) assert v1 != [1, 2, 3] def test_vector_is_zero() -> None: """Test is_zero method for vectors.""" - v0 = Vector3() + v0 = Vector3(x=0.0, y=0.0, z=0.0) assert np.allclose(vector_array(v0), 0) v1 = Vector3(x=0.0, y=0.0, z=0.0) assert np.allclose(vector_array(v1), 0) - v2 = Vector3(x=0.0, y=0.0) + v2 = Vector3(x=0.0, y=0.0, z=0.0) assert np.allclose(vector_array(v2), 0) v3 = Vector3(x=1.0, y=0.0, z=0.0) assert not np.allclose(vector_array(v3), 0) @@ -258,7 +259,7 @@ def test_vector_is_zero() -> None: def test_vector_bool_conversion(): """Test boolean conversion of vectors.""" - v0 = Vector3() + v0 = Vector3(x=0.0, y=0.0, z=0.0) assert not not np.allclose(vector_array(v0), 0) v1 = Vector3(x=0.0, y=0.0, z=0.0) assert not not np.allclose(vector_array(v1), 0) @@ -292,13 +293,13 @@ def test_vector_add() -> None: assert v_add_op.x == 5.0 assert v_add_op.y == 7.0 assert v_add_op.z == 9.0 - v_zero = Vector3() + v_zero = Vector3(x=0.0, y=0.0, z=0.0) assert vector_from_array(vector_array(v1) + vector_array(v_zero)) == v1 def test_vector_add_dim_mismatch() -> None: """Test vector addition with different input dimensions (now all vectors are 3D).""" - v1 = Vector3(x=1.0, y=2.0) + v1 = Vector3(x=1.0, y=2.0, z=0.0) v2 = Vector3(x=4.0, y=5.0, z=6.0) v_add_op = vector_from_array(vector_array(v1) + vector_array(v2)) assert v_add_op.x == 5.0 @@ -312,11 +313,11 @@ def test_yaw_pitch_roll_accessors() -> None: assert v.x == 1.0 assert v.y == 2.0 assert v.z == 3.0 - v_2d = Vector3(x=4.0, y=5.0) + v_2d = Vector3(x=4.0, y=5.0, z=0.0) assert v_2d.x == 4.0 assert v_2d.y == 5.0 assert v_2d.z == 0.0 - v_empty = Vector3() + v_empty = Vector3(x=0.0, y=0.0, z=0.0) assert v_empty.x == 0.0 assert v_empty.y == 0.0 assert v_empty.z == 0.0 @@ -353,8 +354,8 @@ def test_vector_to_quaternion() -> None: def test_cdr_encode_decode() -> None: v_source = Vector3(x=1.0, y=2.0, z=3.0) - binary_msg = v_source.encode() - v_dest = Vector3.decode(binary_msg) + binary_msg = cdr_encode(v_source) + v_dest = cdr_decode(binary_msg, Vector3) assert isinstance(v_dest, Vector3) assert v_dest is not v_source assert v_dest == v_source diff --git a/dimos/msgs/geometry_msgs/test_WrenchStamped.py b/dimos/msgs/geometry_msgs/test_WrenchStamped.py index 2c3b932f12f..cca6c4b43f8 100644 --- a/dimos/msgs/geometry_msgs/test_WrenchStamped.py +++ b/dimos/msgs/geometry_msgs/test_WrenchStamped.py @@ -17,6 +17,8 @@ import pickle import time +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Vector3, Wrench, WrenchStamped from dimos_generated.std_msgs.msg import Header import numpy as np @@ -28,7 +30,12 @@ def test_wrench_empty_is_zero() -> None: - np.testing.assert_array_equal(wrench_array(Wrench()), np.zeros(6)) + np.testing.assert_array_equal( + wrench_array( + Wrench(force=Vector3(x=0.0, y=0.0, z=0.0), torque=Vector3(x=0.0, y=0.0, z=0.0)) + ), + np.zeros(6), + ) def test_wrench_from_array_and_vectors() -> None: @@ -39,15 +46,19 @@ def test_wrench_from_array_and_vectors() -> None: def test_wrench_keywords() -> None: - assert Wrench(force=Vector3(x=1, y=2, z=3)).torque == Vector3() - assert Wrench(torque=Vector3(x=4, y=5, z=6)).force == Vector3() + assert Wrench( + force=Vector3(x=1, y=2, z=3), torque=Vector3(x=0.0, y=0.0, z=0.0) + ).torque == Vector3(x=0.0, y=0.0, z=0.0) + assert Wrench( + torque=Vector3(x=4, y=5, z=6), force=Vector3(x=0.0, y=0.0, z=0.0) + ).force == Vector3(x=0.0, y=0.0, z=0.0) with pytest.raises(TypeError): - Wrench(bogus=1) + Wrench(bogus=1, force=Vector3(x=0.0, y=0.0, z=0.0), torque=Vector3(x=0.0, y=0.0, z=0.0)) def test_wrench_copy_storage() -> None: source = wrench_from_array([1, 2, 3, 4, 5, 6]) - copied = Wrench.decode(source.encode()) + copied = cdr_decode(cdr_encode(source), Wrench) assert copied == source copied.force.x = 10 assert source.force.x == 1 @@ -89,9 +100,22 @@ def test_stamped_nested_wrench_and_explicit_header() -> None: def test_explicit_current_time_and_zero_default() -> None: before = time.time_ns() - source = WrenchStamped(header=header_now()) + source = WrenchStamped( + header=header_now(), + wrench=Wrench(force=Vector3(x=0.0, y=0.0, z=0.0), torque=Vector3(x=0.0, y=0.0, z=0.0)), + ) assert before <= to_nanoseconds(source.header.stamp) <= time.time_ns() - assert to_nanoseconds(WrenchStamped().header.stamp) == 0 + assert ( + to_nanoseconds( + WrenchStamped( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + wrench=Wrench( + force=Vector3(x=0.0, y=0.0, z=0.0), torque=Vector3(x=0.0, y=0.0, z=0.0) + ), + ).header.stamp + ) + == 0 + ) @pytest.mark.parametrize("stamp", [0, 5.25]) @@ -100,11 +124,11 @@ def test_cdr_and_independent_decoding(stamp: float) -> None: header=Header(stamp=time_from_seconds(stamp), frame_id="ft_sensor"), wrench=wrench_from_array([1, 2, 3, 0.1, 0.2, 0.3]), ) - decoded = WrenchStamped.decode(source.encode()) + decoded = cdr_decode(cdr_encode(source), WrenchStamped) assert decoded is not source assert decoded == source independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr( - source.encode(), WrenchStamped.msg_name + cdr_encode(source), WrenchStamped.__msgtype__ ) assert independent.header.frame_id == "ft_sensor" assert independent.header.stamp.sec * 1000000000 + independent.header.stamp.nanosec == int( diff --git a/dimos/msgs/geometry_msgs/tool_publish.py b/dimos/msgs/geometry_msgs/tool_publish.py index a428513e86d..592e54b1d4f 100644 --- a/dimos/msgs/geometry_msgs/tool_publish.py +++ b/dimos/msgs/geometry_msgs/tool_publish.py @@ -14,6 +14,7 @@ import time +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import Vector3 import lcm @@ -22,7 +23,7 @@ def test_runpublish() -> None: for i in range(10): msg = Vector3(x=-5 + i, y=-5 + i, z=i) lc = lcm.LCM() - lc.publish("thing1_vector3#geometry_msgs/msg/Vector3", msg.encode()) + lc.publish("thing1_vector3#geometry_msgs/msg/Vector3", cdr_encode(msg)) time.sleep(0.1) print(f"Published: {msg}") @@ -31,7 +32,7 @@ def test_receive() -> None: lc = lcm.LCM() def receive(channel: str, msg: bytes) -> None: - print(Vector3.decode(msg)) + print(cdr_decode(msg, Vector3)) lc.subscribe("thing1_vector3#geometry_msgs/msg/Vector3", receive) diff --git a/dimos/msgs/helpers.py b/dimos/msgs/helpers.py index 278092a47fa..6e46bea8f03 100644 --- a/dimos/msgs/helpers.py +++ b/dimos/msgs/helpers.py @@ -17,7 +17,7 @@ from functools import lru_cache from typing import TYPE_CHECKING, cast -from dimos.message_codegen.registry import message_types +from dimos_cdrgen.registry import message_types if TYPE_CHECKING: from dimos.msgs.protocol import DimosMsg diff --git a/dimos/msgs/image.py b/dimos/msgs/image.py index 8b3167ef29e..ea773b13e3c 100644 --- a/dimos/msgs/image.py +++ b/dimos/msgs/image.py @@ -19,6 +19,7 @@ from pathlib import Path from typing import Any +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import CompressedImage, Image from dimos_generated.std_msgs.msg import Header import numpy as np @@ -52,15 +53,15 @@ def image_from_array(pixels: NDArray[Any], *, encoding: str, header: Header | No channels != 1 and (pixels.ndim != 3 or pixels.shape[2] != channels) ): raise ValueError(f"array shape does not match {encoding}") - contiguous = np.ascontiguousarray(pixels) + contiguous = np.array(pixels, copy=True, order="C") return Image( - header=header if header is not None else Header(), + header=header if header is not None else Header(stamp=Time(sec=0, nanosec=0), frame_id=""), width=pixels.shape[1], height=pixels.shape[0], encoding=encoding, is_bigendian=int(not pixels.dtype.isnative if np.little_endian else pixels.dtype.isnative), step=pixels.shape[1] * channels * pixels.dtype.itemsize, - data=contiguous.view(np.uint8).ravel(), + data=np.asarray(contiguous.view(np.uint8).ravel(), dtype=np.uint8), ) @@ -79,7 +80,9 @@ def image_view(msg: Image) -> NDArray[Any]: if channels == 1 else (msg.step, channels * dtype.itemsize, dtype.itemsize) ) - return np.ndarray(shape, dtype=dtype, buffer=msg.data.view(), strides=strides) + view = np.ndarray(shape, dtype=dtype, buffer=msg.data, strides=strides) + view.setflags(write=False) + return view def image_to_jpeg(msg: Image, quality: int = 75) -> bytes: @@ -163,7 +166,7 @@ def compressed_image_from_image( return CompressedImage( header=message.header, format=f"{message.encoding}; {format} compressed {output_encoding}", - data=data, + data=np.asarray(np.frombuffer(data, dtype=np.uint8).copy(), dtype=np.uint8), ) diff --git a/dimos/msgs/manipulation_msgs/test_grasp_candidate.py b/dimos/msgs/manipulation_msgs/test_grasp_candidate.py index eb676b1441e..8313f74b512 100644 --- a/dimos/msgs/manipulation_msgs/test_grasp_candidate.py +++ b/dimos/msgs/manipulation_msgs/test_grasp_candidate.py @@ -12,17 +12,23 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.dimos_msgs.msg import GraspCandidate, GraspCandidateArray -from dimos_generated.geometry_msgs.msg import Point, Pose +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion from dimos_generated.std_msgs.msg import Header from dimos.msgs.time import time_from_seconds def test_grasp_candidate_round_trip_preserves_pose_and_score() -> None: - candidate = GraspCandidate(pose=Pose(position=Point(x=0.4, y=-0.2, z=0.3)), score=0.75) + candidate = GraspCandidate( + pose=Pose( + position=Point(x=0.4, y=-0.2, z=0.3), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + score=0.75, + ) - decoded = GraspCandidate.decode(candidate.encode()) + decoded = cdr_decode(cdr_encode(candidate), GraspCandidate) assert decoded.pose.position.x == 0.4 assert decoded.pose.position.y == -0.2 @@ -32,14 +38,26 @@ def test_grasp_candidate_round_trip_preserves_pose_and_score() -> None: def test_grasp_candidate_array_round_trip_preserves_header_and_order() -> None: candidates = [ - GraspCandidate(pose=Pose(position=Point(x=0.1, z=0.2)), score=0.9), - GraspCandidate(pose=Pose(position=Point(x=0.2, z=0.2)), score=0.7), + GraspCandidate( + pose=Pose( + position=Point(x=0.1, z=0.2, y=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + score=0.9, + ), + GraspCandidate( + pose=Pose( + position=Point(x=0.2, z=0.2, y=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + score=0.7, + ), ] proposals = GraspCandidateArray( header=Header(stamp=time_from_seconds(123.0), frame_id="world"), candidates=candidates ) - decoded = GraspCandidateArray.decode(proposals.encode()) + decoded = cdr_decode(cdr_encode(proposals), GraspCandidateArray) assert decoded.header.stamp == time_from_seconds(123.0) assert decoded.header.frame_id == "world" diff --git a/dimos/msgs/mesh.py b/dimos/msgs/mesh.py new file mode 100644 index 00000000000..88ca71f0d3a --- /dev/null +++ b/dimos/msgs/mesh.py @@ -0,0 +1,62 @@ +# 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. + +"""Array construction and rendering for generated keyed mesh chunks.""" + +from typing import Any + +from dimos_generated.dimos_msgs.msg import TriangleMesh +from dimos_generated.std_msgs.msg import Header +import numpy as np +from numpy.typing import NDArray + +from dimos.msgs.time import header_now + + +def triangle_mesh_from_arrays( + vertices: NDArray[Any], + faces: NDArray[Any], + normals: NDArray[Any] | None = None, + key: tuple[int, int, int] = (0, 0, 0), + *, + header: Header | None = None, +) -> TriangleMesh: + vertices = np.array(vertices, dtype=np.float32, copy=True).reshape(-1, 3) + normals = ( + np.zeros_like(vertices) + if normals is None + else np.array(normals, dtype=np.float32, copy=True).reshape(-1, 3) + ) + faces = np.array(faces, dtype=np.uint32, copy=True).reshape(-1, 3) + if normals.shape != vertices.shape: + raise ValueError("mesh normals must match vertices") + return TriangleMesh( + header=header_now() if header is None else header, + key=np.array(key, dtype=np.int32), + vertices=vertices.ravel(), + normals=normals.ravel(), + faces=faces.ravel(), + ) + + +def triangle_mesh_to_rerun(message: TriangleMesh) -> Any: + import rerun as rr + + if len(message.faces) == 0: + return rr.Clear(recursive=False) + return rr.Mesh3D( + vertex_positions=message.vertices.reshape(-1, 3), + vertex_normals=message.normals.reshape(-1, 3), + triangle_indices=message.faces.reshape(-1, 3), + ) diff --git a/dimos/msgs/nav_msgs/test_LineSegments3D.py b/dimos/msgs/nav_msgs/test_LineSegments3D.py index 4b88685eb57..bc3f9ba5e19 100644 --- a/dimos/msgs/nav_msgs/test_LineSegments3D.py +++ b/dimos/msgs/nav_msgs/test_LineSegments3D.py @@ -12,75 +12,73 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos_lcm.geometry_msgs.Point import Point -from dimos_lcm.geometry_msgs.Pose import Pose -from dimos_lcm.geometry_msgs.PoseStamped import PoseStamped -from dimos_lcm.geometry_msgs.Quaternion import Quaternion -from dimos_lcm.nav_msgs.Path import Path -from dimos_lcm.std_msgs.Header import Header -from dimos_lcm.std_msgs.Time import Time +"""Explicit line messages replace the old Path/orientation weight encoding.""" + +from dimos_cdrgen.registry import ( + decode as cdr_decode, + encode as cdr_encode, + schema as cdr_schema, +) +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.dimos_msgs.msg import LineSegment3D, LineSegments3D +from dimos_generated.geometry_msgs.msg import Point +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest +from rosbags.typesys import Stores, get_types_from_msg, get_typestore -from dimos.msgs.nav_msgs.LineSegments3D import LineSegments3D - - -def encode_edges(n_segments: int, frame_ids: list[str] | None = None, seq: int = 0) -> bytes: - """A Path of consecutive pose pairs with orientation.w carrying the weight.""" - poses = [ - PoseStamped( - header=Header(frame_id=frame_ids[i % len(frame_ids)] if frame_ids else "odom"), - pose=Pose( - position=Point(float(i), float(i) * 0.5, -1.0), - orientation=Quaternion(w=float(i // 2) * 0.1), - ), - ) - for i in range(2 * n_segments) - ] - header = Header(seq=seq, stamp=Time(12, 500_000_000), frame_id="odom") - return Path(poses_length=len(poses), header=header, poses=poses).lcm_encode() - - -def expected_segments(n_segments: int) -> np.ndarray: - i = np.arange(2 * n_segments, dtype=np.float64) - return np.stack([i, i * 0.5, np.full_like(i, -1.0)], axis=1).reshape(-1, 2, 3) - - -def test_decode_matches_the_wire_layout() -> None: - msg = LineSegments3D.lcm_decode(encode_edges(50)) - assert msg.frame_id == "odom" - assert msg.ts == 12.5 - np.testing.assert_array_equal(msg.segments, expected_segments(50)) - np.testing.assert_allclose(msg.weights, np.arange(50) * 0.1) +from dimos.msgs.time import time_from_nanoseconds, to_nanoseconds - empty = LineSegments3D.lcm_decode(encode_edges(0)) - assert len(empty) == 0 - assert empty.segments.shape == (0, 2, 3) - assert empty.weights.shape == (0,) - -def test_seq_round_trips_including_the_empty_message() -> None: - seq = (3 << 16) | 5 - assert LineSegments3D.lcm_decode(encode_edges(4, seq=seq)).seq == seq - assert LineSegments3D.lcm_decode(encode_edges(0, seq=seq)).seq == seq - assert LineSegments3D.lcm_decode(encode_edges(4)).seq == 0 - - -def test_mixed_frame_id_lengths_are_rejected() -> None: - raw = encode_edges(20, frame_ids=["odom", "map", "base_link"]) - with pytest.raises(ValueError, match="frame_id length"): - LineSegments3D.lcm_decode(raw) +@pytest.mark.parametrize("count", [0, 1, 50]) +def test_cdr_segment_layout_and_independent_schema(count: int) -> None: + source = LineSegments3D( + header=Header(stamp=time_from_nanoseconds(12500000000), frame_id="odom"), + segments=[ + LineSegment3D( + start=Point(x=2 * i, y=i, z=-1), + end=Point(x=2 * i + 1, y=(2 * i + 1) * 0.5, z=-1), + weight=i * 0.1, + ) + for i in range(count) + ], + ) + decoded = cdr_decode(cdr_encode(source), LineSegments3D) + assert decoded.header.frame_id == "odom" + assert to_nanoseconds(decoded.header.stamp) == 12500000000 + assert len(decoded.segments) == count + coordinates = np.array( + [[[s.start.x, s.start.y, s.start.z], [s.end.x, s.end.y, s.end.z]] for s in decoded.segments] + ).reshape(-1, 2, 3) + indices = np.arange(2 * count, dtype=np.float64) + expected = np.stack([indices, indices * 0.5, np.full_like(indices, -1)], axis=1).reshape( + -1, 2, 3 + ) + np.testing.assert_array_equal(coordinates, expected) + np.testing.assert_allclose([s.weight for s in decoded.segments], np.arange(count) * 0.1) + store = get_typestore(Stores.ROS2_JAZZY) + store.register( + get_types_from_msg(cdr_schema(LineSegments3D.__msgtype__), LineSegments3D.__msgtype__) + ) + independent = store.deserialize_cdr(cdr_encode(source), LineSegments3D.__msgtype__) + assert len(independent.segments) == count + assert independent.header.frame_id == "odom" + for i, segment in enumerate(independent.segments): + assert segment.start.x == 2 * i + assert segment.end.y == (2 * i + 1) * 0.5 + assert segment.weight == pytest.approx(i * 0.1) -def test_encode_round_trips() -> None: - seq = (-3 << 16) | 5 - msg = LineSegments3D( - ts=12.5, frame_id="world", segments=expected_segments(3), weights=[0.1, 0.2, 0.3], seq=seq +@pytest.mark.parametrize("frame", ["odom", "map", "base_link"]) +def test_explicit_segment_schema_accepts_variable_frame_lengths(frame: str) -> None: + source = LineSegments3D( + header=Header(frame_id=frame, stamp=Time(sec=0, nanosec=0)), + segments=[ + LineSegment3D( + weight=0.7, start=Point(x=0.0, y=0.0, z=0.0), end=Point(x=0.0, y=0.0, z=0.0) + ) + ], ) - back = LineSegments3D.lcm_decode(msg.lcm_encode()) - assert back.frame_id == "world" - assert back.ts == 12.5 - assert back.seq == seq - np.testing.assert_array_equal(back.segments, msg.segments) - np.testing.assert_allclose(back.weights, msg.weights) - assert len(LineSegments3D.lcm_decode(LineSegments3D(ts=1.0).lcm_encode())) == 0 + decoded = cdr_decode(cdr_encode(source), LineSegments3D) + assert decoded.header.frame_id == frame + assert decoded.segments[0].weight == pytest.approx(0.7) diff --git a/dimos/msgs/nav_msgs/test_OccupancyGrid.py b/dimos/msgs/nav_msgs/test_OccupancyGrid.py index 050c6420900..1dd4e6ef075 100644 --- a/dimos/msgs/nav_msgs/test_OccupancyGrid.py +++ b/dimos/msgs/nav_msgs/test_OccupancyGrid.py @@ -15,8 +15,10 @@ """Generated occupancy grids with external array and coordinate operations.""" -from dimos_generated.geometry_msgs.msg import Point, Pose -from dimos_generated.nav_msgs.msg import OccupancyGrid +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid from dimos_generated.std_msgs.msg import Header import numpy as np import pytest @@ -37,7 +39,20 @@ def test_empty_grid() -> None: """Test creating an empty grid.""" - grid = OccupancyGrid() + grid = OccupancyGrid( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + info=MapMetaData( + map_load_time=Time(sec=0, nanosec=0), + resolution=0.0, + width=0, + height=0, + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([], dtype=np.int8), + ) assert grid.info.width == 0 assert grid.info.height == 0 assert occupancy_view(grid).shape == (0, 0) @@ -48,7 +63,9 @@ def test_empty_grid() -> None: def test_grid_with_dimensions() -> None: """Test creating a grid with specified dimensions.""" grid = occupancy_from_array( - np.full((10, 10), -1, dtype=np.int8), resolution=0.1, header=Header(frame_id="map") + np.full((10, 10), -1, dtype=np.int8), + resolution=0.1, + header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), ) assert grid.info.width == 10 assert grid.info.height == 10 @@ -65,9 +82,14 @@ def test_grid_from_numpy_array() -> None: data = np.zeros((20, 30), dtype=np.int8) data[5:10, 10:20] = 100 data[15:18, 5:8] = -1 - origin = Pose(position=Point(x=1, y=2)) + origin = Pose( + position=Point(x=1, y=2, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) grid = occupancy_from_array( - data, resolution=0.05, origin=origin, header=Header(frame_id="odom") + data, + resolution=0.05, + origin=origin, + header=Header(frame_id="odom", stamp=Time(sec=0, nanosec=0)), ) assert grid.info.width == 30 assert grid.info.height == 20 @@ -87,11 +109,16 @@ def test_grid_from_numpy_array() -> None: def test_world_grid_coordinate_conversion() -> None: """Test converting between world and grid coordinates.""" data = np.zeros((20, 30), dtype=np.int8) - origin = Pose(position=Point(x=1, y=2)) + origin = Pose( + position=Point(x=1, y=2, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) grid = occupancy_from_array( - data, resolution=0.05, origin=origin, header=Header(frame_id="odom") + data, + resolution=0.05, + origin=origin, + header=Header(frame_id="odom", stamp=Time(sec=0, nanosec=0)), ) - grid_pos = world_to_grid(grid, Point(x=2.5, y=3.0)) + grid_pos = world_to_grid(grid, Point(x=2.5, y=3.0, z=0.0)) assert grid_pos[0] == pytest.approx(30) assert grid_pos[1] == pytest.approx(20) world_pos = grid_to_world(grid, (10, 5)) @@ -104,16 +131,21 @@ def test_cdr_encode_decode() -> None: data = np.zeros((20, 30), dtype=np.int8) data[5:10, 10:20] = 100 data[15:18, 5:8] = -1 - origin = Pose(position=Point(x=1, y=2)) + origin = Pose( + position=Point(x=1, y=2, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) grid = occupancy_from_array( - data, resolution=0.05, origin=origin, header=Header(frame_id="odom") + data, + resolution=0.05, + origin=origin, + header=Header(frame_id="odom", stamp=Time(sec=0, nanosec=0)), ) - grid_pos = world_to_grid(grid, Point(x=1.5, y=2.25)) + grid_pos = world_to_grid(grid, Point(x=1.5, y=2.25, z=0.0)) grid.data[round(grid_pos[1]) * grid.info.width + round(grid_pos[0])] = 50 - lcm_data = grid.encode() + lcm_data = cdr_encode(grid) assert isinstance(lcm_data, bytes) assert len(lcm_data) > 0 - decoded = OccupancyGrid.decode(lcm_data) + decoded = cdr_decode(lcm_data, OccupancyGrid) assert np.array_equal(occupancy_view(grid), occupancy_view(decoded)) assert grid.info.width == decoded.info.width assert grid.info.height == decoded.info.height @@ -127,11 +159,14 @@ def test_cdr_encode_decode() -> None: def test_cdr_decode_origin_is_dimos_pose() -> None: """Generated origin orientation works with the external yaw helper.""" quat = quaternion_from_euler(0, 0, 0.5) - origin = Pose(position=Point(x=1, y=2), orientation=quat) + origin = Pose(position=Point(x=1, y=2, z=0.0), orientation=quat) grid = occupancy_from_array( - np.zeros((2, 3), dtype=np.int8), resolution=0.05, origin=origin, header=Header(frame_id="") + np.zeros((2, 3), dtype=np.int8), + resolution=0.05, + origin=origin, + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), ) - decoded = OccupancyGrid.decode(grid.encode()) + decoded = cdr_decode(cdr_encode(grid), OccupancyGrid) assert isinstance(decoded.info.origin, Pose) assert yaw(decoded.info.origin.orientation) == pytest.approx(0.5) @@ -139,14 +174,34 @@ def test_cdr_decode_origin_is_dimos_pose() -> None: def test_cdr_decode_empty_grid() -> None: """An empty grid (mapper warming up) must survive the wire; the decoder used to rebuild it as a 1-D array the constructor rejects.""" - decoded = OccupancyGrid.decode(OccupancyGrid().encode()) + decoded = cdr_decode( + cdr_encode( + OccupancyGrid( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + info=MapMetaData( + map_load_time=Time(sec=0, nanosec=0), + resolution=0.0, + width=0, + height=0, + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([], dtype=np.int8), + ) + ), + OccupancyGrid, + ) assert occupancy_view(decoded).size == 0 assert isinstance(decoded.info.origin, Pose) def test_physical_extent() -> None: grid = occupancy_from_array( - np.full((10, 10), -1, dtype=np.int8), header=Header(frame_id="map"), resolution=0.1 + np.full((10, 10), -1, dtype=np.int8), + header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), + resolution=0.1, ) assert occupancy_extent(grid) == pytest.approx((1, 1)) @@ -154,7 +209,9 @@ def test_physical_extent() -> None: def test_grid_property_sync() -> None: """Test that the grid property works correctly.""" grid = occupancy_from_array( - np.full((5, 5), -1, dtype=np.int8), resolution=0.1, header=Header(frame_id="map") + np.full((5, 5), -1, dtype=np.int8), + resolution=0.1, + header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), ) grid.data[2 * grid.info.width + 3] = 100 assert occupancy_view(grid)[2, 3] == 100 @@ -165,14 +222,18 @@ def test_grid_property_sync() -> None: def test_invalid_grid_dimensions() -> None: """Test handling of invalid grid dimensions.""" with pytest.raises(ValueError, match="2D integer array"): - occupancy_from_array(np.zeros(10), resolution=0.1, header=Header(frame_id="")) + occupancy_from_array( + np.zeros(10), resolution=0.1, header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)) + ) def test_from_pointcloud() -> None: """Test creating OccupancyGrid from PointCloud2.""" x, y = np.meshgrid(np.arange(10) * 0.05, np.arange(10) * 0.05) points = np.column_stack((x.ravel(), y.ravel(), np.full(x.size, 0.5))) - pointcloud = pointcloud_from_xyz(points, header=Header(frame_id="world")) + pointcloud = pointcloud_from_xyz( + points, header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)) + ) occupancygrid = general_occupancy(pointcloud, resolution=0.05, min_height=0.1, max_height=2.0) occupancygrid = simple_inflate(occupancygrid, 0.1) assert occupancygrid.info.width > 0 @@ -187,7 +248,9 @@ def test_filter_above() -> None: data = np.array( [[-1, 0, 20, 50], [10, 30, 60, 80], [40, 70, 90, 100], [-1, 15, 25, -1]], dtype=np.int8 ) - grid = occupancy_from_array(data, resolution=0.1, header=Header(frame_id="")) + grid = occupancy_from_array( + data, resolution=0.1, header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)) + ) filtered = occupancy_from_array( np.where(occupancy_view(grid) > 50, occupancy_view(grid), -1).astype(np.int8), header=grid.header, @@ -219,7 +282,9 @@ def test_filter_below() -> None: data = np.array( [[-1, 0, 20, 50], [10, 30, 60, 80], [40, 70, 90, 100], [-1, 15, 25, -1]], dtype=np.int8 ) - grid = occupancy_from_array(data, resolution=0.1, header=Header(frame_id="")) + grid = occupancy_from_array( + data, resolution=0.1, header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)) + ) filtered = occupancy_from_array( np.where(occupancy_view(grid) < 50, occupancy_view(grid), -1).astype(np.int8), header=grid.header, @@ -253,7 +318,9 @@ def test_max() -> None: data = np.array( [[-1, 0, 20, 50], [10, 30, 60, 80], [40, 70, 90, 100], [-1, 15, 25, -1]], dtype=np.int8 ) - grid = occupancy_from_array(data, resolution=0.1, header=Header(frame_id="")) + grid = occupancy_from_array( + data, resolution=0.1, header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)) + ) maxed = occupancy_from_array( np.where(occupancy_view(grid) >= 0, 100, -1).astype(np.int8), header=grid.header, diff --git a/dimos/msgs/nav_msgs/test_Odometry.py b/dimos/msgs/nav_msgs/test_Odometry.py index 73af9d94e05..c47e6a45667 100644 --- a/dimos/msgs/nav_msgs/test_Odometry.py +++ b/dimos/msgs/nav_msgs/test_Odometry.py @@ -12,8 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dataclasses import asdict import time +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, Pose, @@ -33,7 +36,21 @@ def test_odometry_default_init() -> None: - odom = Odometry() + odom = Odometry( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + child_frame_id="", + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), + twist=TwistWithCovariance( + twist=Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + covariance=np.zeros(36, dtype=np.float64), + ), + ) assert to_seconds(odom.header.stamp) == 0 assert odom.header.frame_id == "" assert odom.child_frame_id == "" @@ -50,7 +67,19 @@ def test_odometry_default_init() -> None: def test_odometry_with_frames() -> None: ts = 1234567890.123456 odom = Odometry( - child_frame_id="base_link", header=Header(stamp=time_from_seconds(ts), frame_id="odom") + child_frame_id="base_link", + header=Header(stamp=time_from_seconds(ts), frame_id="odom"), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), + twist=TwistWithCovariance( + twist=Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + covariance=np.zeros(36, dtype=np.float64), + ), ) assert to_seconds(odom.header.stamp) == ts assert odom.header.frame_id == "odom" @@ -64,8 +93,8 @@ def test_odometry_with_pose_and_twist() -> None: twist = Twist(linear=Vector3(x=0.5, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.1)) odom = Odometry( child_frame_id="base_link", - pose=PoseWithCovariance(pose=pose), - twist=TwistWithCovariance(twist=twist), + pose=PoseWithCovariance(pose=pose, covariance=np.zeros(36, dtype=np.float64)), + twist=TwistWithCovariance(twist=twist, covariance=np.zeros(36, dtype=np.float64)), header=Header(stamp=time_from_seconds(1000.0), frame_id="odom"), ) assert odom.pose.pose.position.x == 1.0 @@ -76,12 +105,16 @@ def test_odometry_with_pose_and_twist() -> None: def test_odometry_with_covariances() -> None: - pose = Pose(position=Point(x=1.0, y=2.0, z=3.0)) + pose = Pose( + position=Point(x=1.0, y=2.0, z=3.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) pose_cov = np.arange(36, dtype=float) - pose_with_cov = PoseWithCovariance(pose=pose, covariance=pose_cov) + pose_with_cov = PoseWithCovariance(pose=pose, covariance=np.asarray(pose_cov, dtype=np.float64)) twist = Twist(linear=Vector3(x=0.5, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.1)) twist_cov = np.arange(36, 72, dtype=float) - twist_with_cov = TwistWithCovariance(twist=twist, covariance=twist_cov) + twist_with_cov = TwistWithCovariance( + twist=twist, covariance=np.asarray(twist_cov, dtype=np.float64) + ) odom = Odometry( child_frame_id="base_link", pose=pose_with_cov, @@ -101,8 +134,8 @@ def test_odometry_properties() -> None: twist = Twist(linear=Vector3(x=0.5, y=0.6, z=0.7), angular=Vector3(x=0.1, y=0.2, z=0.3)) odom = Odometry( child_frame_id="base_link", - pose=PoseWithCovariance(pose=pose), - twist=TwistWithCovariance(twist=twist), + pose=PoseWithCovariance(pose=pose, covariance=np.zeros(36, dtype=np.float64)), + twist=TwistWithCovariance(twist=twist, covariance=np.zeros(36, dtype=np.float64)), header=Header(stamp=time_from_seconds(1000.0), frame_id="odom"), ) assert odom.pose.pose.position.x == 1.0 @@ -127,10 +160,21 @@ def test_independent_ros_decoding() -> None: source = Odometry( header=Header(stamp=time_from_seconds(1000), frame_id="odom"), child_frame_id="base_link", - pose=PoseWithCovariance(pose=Pose(position=Point(x=1.234, y=2.567, z=3.891))), - twist=TwistWithCovariance(twist=Twist(linear=Vector3(x=0.5), angular=Vector3(z=0.1))), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=1.234, y=2.567, z=3.891), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), + twist=TwistWithCovariance( + twist=Twist(linear=Vector3(x=0.5, y=0.0, z=0.0), angular=Vector3(z=0.1, x=0.0, y=0.0)), + covariance=np.zeros(36, dtype=np.float64), + ), + ) + decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr( + cdr_encode(source), Odometry.__msgtype__ ) - decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(source.encode(), Odometry.msg_name) assert decoded.header.frame_id == "odom" assert decoded.child_frame_id == "base_link" assert ( @@ -145,15 +189,31 @@ def test_independent_ros_decoding() -> None: def test_odometry_equality() -> None: kwargs = dict( child_frame_id="base_link", - pose=PoseWithCovariance(pose=Pose(position=Point(x=1.0, y=2.0, z=3.0))), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=1.0, y=2.0, z=3.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), twist=TwistWithCovariance( - twist=Twist(linear=Vector3(x=0.5, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.1)) + twist=Twist(linear=Vector3(x=0.5, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.1)), + covariance=np.zeros(36, dtype=np.float64), ), header=Header(stamp=time_from_seconds(1000.0), frame_id="odom"), ) assert Odometry(**kwargs) == Odometry(**kwargs) assert Odometry(**kwargs) != Odometry( - **{**kwargs, "pose": PoseWithCovariance(pose=Pose(position=Point(x=1.1, y=2.0, z=3.0)))} + **{ + **kwargs, + "pose": PoseWithCovariance( + pose=Pose( + position=Point(x=1.1, y=2.0, z=3.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), + } ) assert Odometry(**kwargs) != "not an odometry" @@ -167,26 +227,83 @@ def test_odometry_cdr_roundtrip() -> None: twist_cov = np.arange(36, 72, dtype=float) source = Odometry( child_frame_id="base_link", - pose=PoseWithCovariance(pose=pose, covariance=pose_cov), - twist=TwistWithCovariance(twist=twist, covariance=twist_cov), + pose=PoseWithCovariance(pose=pose, covariance=np.asarray(pose_cov, dtype=np.float64)), + twist=TwistWithCovariance(twist=twist, covariance=np.asarray(twist_cov, dtype=np.float64)), header=Header(stamp=time_from_seconds(1234567890.123456), frame_id="odom"), ) - decoded = Odometry.decode(source.encode()) + decoded = cdr_decode(cdr_encode(source), Odometry) assert abs(to_seconds(decoded.header.stamp) - to_seconds(source.header.stamp)) < 1e-06 assert decoded.header.frame_id == source.header.frame_id assert decoded.child_frame_id == source.child_frame_id - assert decoded.pose == source.pose - assert decoded.twist == source.twist + np.testing.assert_equal(asdict(decoded.pose), asdict(source.pose)) + np.testing.assert_equal(asdict(decoded.twist), asdict(source.twist)) def test_odometry_zero_timestamp() -> None: - assert to_seconds(Odometry().header.stamp) == 0 + assert ( + to_seconds( + Odometry( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + child_frame_id="", + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), + twist=TwistWithCovariance( + twist=Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ), + covariance=np.zeros(36, dtype=np.float64), + ), + ).header.stamp + ) + == 0 + ) before = time.time() - assert before <= to_seconds(Odometry(header=header_now()).header.stamp) <= time.time() + assert ( + before + <= to_seconds( + Odometry( + header=header_now(), + child_frame_id="", + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), + twist=TwistWithCovariance( + twist=Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ), + covariance=np.zeros(36, dtype=np.float64), + ), + ).header.stamp + ) + <= time.time() + ) def test_odometry_with_just_pose() -> None: - odom = Odometry(pose=PoseWithCovariance(pose=Pose(position=Point(x=1.0, y=2.0, z=3.0)))) + odom = Odometry( + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=1.0, y=2.0, z=3.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + child_frame_id="", + twist=TwistWithCovariance( + twist=Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + covariance=np.zeros(36, dtype=np.float64), + ), + ) assert odom.pose.pose.position.x == 1.0 assert np.all(np.asarray(odom.pose.covariance) == 0.0) assert np.all(np.asarray(odom.twist.covariance) == 0.0) @@ -195,8 +312,18 @@ def test_odometry_with_just_pose() -> None: def test_odometry_with_just_twist() -> None: odom = Odometry( twist=TwistWithCovariance( - twist=Twist(linear=Vector3(x=0.5, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.1)) - ) + twist=Twist(linear=Vector3(x=0.5, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.1)), + covariance=np.zeros(36, dtype=np.float64), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + child_frame_id="", + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), ) assert odom.twist.twist.linear.x == 0.5 assert odom.twist.twist.angular.z == 0.1 @@ -210,10 +337,12 @@ def test_odometry_typical_robot_scenario() -> None: pose=Pose( position=Point(x=10.0, y=5.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=np.sin(0.1), w=np.cos(0.1)), - ) + ), + covariance=np.zeros(36, dtype=np.float64), ), twist=TwistWithCovariance( - twist=Twist(linear=Vector3(x=0.5, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.05)) + twist=Twist(linear=Vector3(x=0.5, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.05)), + covariance=np.zeros(36, dtype=np.float64), ), header=Header(stamp=time_from_seconds(1000.0), frame_id="odom"), ) diff --git a/dimos/msgs/nav_msgs/test_Path.py b/dimos/msgs/nav_msgs/test_Path.py index bd68698c5fc..f1251e5f69a 100644 --- a/dimos/msgs/nav_msgs/test_Path.py +++ b/dimos/msgs/nav_msgs/test_Path.py @@ -13,6 +13,8 @@ # limitations under the License. +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.nav_msgs.msg import Path from dimos_generated.std_msgs.msg import Header @@ -24,14 +26,14 @@ def create_test_pose(x: float, y: float, z: float, frame_id: str = "map") -> PoseStamped: """Helper to create a test PoseStamped.""" return PoseStamped( - header=Header(frame_id=frame_id), - pose=Pose(position=Point(x=x, y=y, z=z), orientation=Quaternion(w=1)), + header=Header(frame_id=frame_id, stamp=Time(sec=0, nanosec=0)), + pose=Pose(position=Point(x=x, y=y, z=z), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0)), ) def test_init_empty() -> None: """Test creating an empty path.""" - path = Path(header=Header(frame_id="map")) + path = Path(header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), poses=[]) assert path.header.frame_id == "map" assert len(path.poses) == 0 assert not list(path.poses) # Should be falsy when empty @@ -41,7 +43,7 @@ def test_init_empty() -> None: def test_init_with_poses() -> None: """Test creating a path with initial poses.""" poses = [create_test_pose(i, i, 0) for i in range(3)] - path = Path(header=Header(frame_id="map"), poses=poses) + path = Path(header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), poses=poses) assert len(path.poses) == 3 assert bool(list(path.poses)) # Should be truthy when has poses assert path.poses == poses @@ -50,46 +52,46 @@ def test_init_with_poses() -> None: def test_head() -> None: """Test getting the first pose.""" poses = [create_test_pose(i, i, 0) for i in range(3)] - path = Path(poses=poses) + path = Path(poses=poses, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) assert (path.poses[0] if path.poses else None) == poses[0] # Test empty path - empty_path = Path() + empty_path = Path(header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), poses=[]) assert (empty_path.poses[0] if empty_path.poses else None) is None def test_last() -> None: """Test getting the last pose.""" poses = [create_test_pose(i, i, 0) for i in range(3)] - path = Path(poses=poses) + path = Path(poses=poses, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) assert (path.poses[-1] if path.poses else None) == poses[-1] # Test empty path - empty_path = Path() + empty_path = Path(header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), poses=[]) assert (empty_path.poses[-1] if empty_path.poses else None) is None def test_tail() -> None: """Test getting all poses except the first.""" poses = [create_test_pose(i, i, 0) for i in range(3)] - path = Path(poses=poses) + path = Path(poses=poses, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) tail = Path(header=path.header, poses=list(path.poses)[1:]) assert len(tail.poses) == 2 assert tail.poses == poses[1:] assert tail.header.frame_id == path.header.frame_id # Test single element path - single_path = Path(poses=[poses[0]]) + single_path = Path(poses=[poses[0]], header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) assert len(list(single_path.poses)[1:]) == 0 # Test empty path - empty_path = Path() + empty_path = Path(header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), poses=[]) assert len(list(empty_path.poses)[1:]) == 0 def test_push_immutable() -> None: """Test immutable push operation.""" - path = Path(header=Header(frame_id="map")) + path = Path(header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), poses=[]) pose1 = create_test_pose(1, 1, 0) pose2 = create_test_pose(2, 2, 0) @@ -108,7 +110,7 @@ def test_push_immutable() -> None: def test_push_mutable() -> None: """Test mutable push operation.""" - path = Path(header=Header(frame_id="map")) + path = Path(header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), poses=[]) pose1 = create_test_pose(1, 1, 0) pose2 = create_test_pose(2, 2, 0) @@ -125,7 +127,7 @@ def test_push_mutable() -> None: def test_indexing() -> None: """Test indexing and slicing.""" poses = [create_test_pose(i, i, 0) for i in range(5)] - path = Path(poses=poses) + path = Path(poses=poses, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) # Single index assert next(iter(path.poses)) == poses[0] @@ -140,7 +142,7 @@ def test_indexing() -> None: def test_iteration() -> None: """Test iterating over poses.""" poses = [create_test_pose(i, i, 0) for i in range(3)] - path = Path(poses=poses) + path = Path(poses=poses, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) collected = [] for pose in path.poses: @@ -151,7 +153,7 @@ def test_iteration() -> None: def test_slice_method() -> None: """Test slice method.""" poses = [create_test_pose(i, i, 0) for i in range(5)] - path = Path(header=Header(frame_id="map"), poses=poses) + path = Path(header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), poses=poses) sliced = Path(header=path.header, poses=list(path.poses)[1:4]) assert len(sliced.poses) == 3 @@ -168,8 +170,8 @@ def test_extend_immutable() -> None: poses1 = [create_test_pose(i, i, 0) for i in range(2)] poses2 = [create_test_pose(i + 2, i + 2, 0) for i in range(2)] - path1 = Path(header=Header(frame_id="map"), poses=poses1) - path2 = Path(header=Header(frame_id="odom"), poses=poses2) + path1 = Path(header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), poses=poses1) + path2 = Path(header=Header(frame_id="odom", stamp=Time(sec=0, nanosec=0)), poses=poses2) extended = Path(header=path1.header, poses=[*path1.poses, *path2.poses]) assert len(path1.poses) == 2 # Original unchanged @@ -184,9 +186,9 @@ def test_extend_mutable() -> None: poses2 = [create_test_pose(i + 2, i + 2, 0) for i in range(2)] path1 = Path( - header=Header(frame_id="map"), poses=poses1.copy() + header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), poses=poses1.copy() ) # Use copy to avoid modifying original - path2 = Path(header=Header(frame_id="odom"), poses=poses2) + path2 = Path(header=Header(frame_id="odom", stamp=Time(sec=0, nanosec=0)), poses=poses2) path1.poses.extend(path2.poses) assert len(path1.poses) == 4 @@ -200,7 +202,7 @@ def test_extend_mutable() -> None: def test_reverse() -> None: """Test reverse operation.""" poses = [create_test_pose(i, i, 0) for i in range(3)] - path = Path(poses=poses) + path = Path(poses=poses, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) reversed_path = Path(header=path.header, poses=list(reversed(list(path.poses)))) assert len(path.poses) == 3 # Original unchanged @@ -210,7 +212,7 @@ def test_reverse() -> None: def test_clear() -> None: """Test clear operation.""" poses = [create_test_pose(i, i, 0) for i in range(3)] - path = Path(poses=poses) + path = Path(poses=poses, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) path.poses.clear() assert len(path.poses) == 0 @@ -235,11 +237,13 @@ def test_cdr_roundtrip_preserves_per_pose_frames_and_stamps() -> None: header=Header(stamp=time_from_nanoseconds(1234567890500000000), frame_id="world"), poses=poses, ) - decoded = Path.decode(source.encode()) + decoded = cdr_decode(cdr_encode(source), Path) assert decoded is not source assert decoded.header == source.header assert list(decoded.poses) == poses - independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(source.encode(), Path.msg_name) + independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr( + cdr_encode(source), Path.__msgtype__ + ) assert independent.header.frame_id == "world" assert independent.header.stamp.sec == 1234567890 assert independent.header.stamp.nanosec == 500000000 @@ -261,7 +265,7 @@ def test_cdr_roundtrip_preserves_per_pose_frames_and_stamps() -> None: def test_cdr_empty_path() -> None: - source = Path(header=Header(frame_id="base_link")) - decoded = Path.decode(source.encode()) + source = Path(header=Header(frame_id="base_link", stamp=Time(sec=0, nanosec=0)), poses=[]) + decoded = cdr_decode(cdr_encode(source), Path) assert decoded.header.frame_id == "base_link" assert len(decoded.poses) == 0 diff --git a/dimos/msgs/occupancy.py b/dimos/msgs/occupancy.py index 773c34fe137..9d2d93a5384 100644 --- a/dimos/msgs/occupancy.py +++ b/dimos/msgs/occupancy.py @@ -31,7 +31,9 @@ def occupancy_view(msg: OccupancyGrid) -> NDArray[np.int8]: """Borrow a read-only row-major occupancy grid, retaining its owner.""" if len(msg.data) != msg.info.height * msg.info.width: raise ValueError("occupancy data length does not match dimensions") - return msg.data.view().reshape(msg.info.height, msg.info.width) + view = msg.data.view().reshape(msg.info.height, msg.info.width) + view.setflags(write=False) + return view def block_max_reduce(cells: NDArray[np.int8], factor: int) -> NDArray[np.int8]: @@ -126,7 +128,12 @@ def occupancy_from_array( resolution=resolution, height=cells.shape[0], width=cells.shape[1], - origin=origin if origin is not None else Pose(orientation=Quaternion(w=1)), + origin=origin + if origin is not None + else Pose( + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + position=Point(x=0.0, y=0.0, z=0.0), + ), ), - data=cells.astype(np.int8).reshape(-1), + data=np.asarray(cells.astype(np.int8).reshape(-1), dtype=np.int8), ) diff --git a/dimos/msgs/pointcloud.py b/dimos/msgs/pointcloud.py index aeb81a58a56..9d65119d578 100644 --- a/dimos/msgs/pointcloud.py +++ b/dimos/msgs/pointcloud.py @@ -17,9 +17,11 @@ """NumPy conversions for generated PointCloud2 messages.""" +from copy import deepcopy import math from typing import Any +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import TransformStamped from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2, PointField from dimos_generated.std_msgs.msg import Header @@ -63,12 +65,14 @@ def pointcloud_view(message: PointCloud2) -> NDArray[Any]: dtype = np.dtype( {"names": names, "formats": formats, "offsets": offsets, "itemsize": message.point_step} ) - return np.ndarray( + view = np.ndarray( (message.height, message.width), dtype=dtype, - buffer=message.data.view(), + buffer=message.data, strides=(message.row_step, message.point_step), ) + view.setflags(write=False) + return view def pointcloud_xyz(message: PointCloud2) -> NDArray[np.float64]: @@ -80,7 +84,9 @@ def pointcloud_xyz(message: PointCloud2) -> NDArray[np.float64]: return np.stack([points[name].ravel() for name in ("x", "y", "z")], axis=-1).astype(np.float64) -def pointcloud_from_xyz(points: NDArray[Any], *, header: Header) -> PointCloud2: +def pointcloud_from_xyz( + points: NDArray[Any], *, header: Header, stamps: NDArray[Any] | None = None +) -> PointCloud2: """Copy Nx3 or HxWx3 coordinates into a tightly packed little-endian XYZ cloud. XYZ fields are float32. NaNs/infinities remain missing-point markers and set @@ -96,8 +102,8 @@ def pointcloud_from_xyz(points: NDArray[Any], *, header: Header) -> PointCloud2: raise ValueError("XYZ coordinates exceed float32 range") packed = np.ascontiguousarray(values, dtype=" PointCloud2: is_bigendian=False, point_step=12, row_step=width * 12, - data=packed.view(np.uint8).reshape(-1), + data=np.asarray(packed.view(np.uint8).reshape(-1), dtype=np.uint8), is_dense=bool(finite.all()), ) + if stamps is not None: + times = np.asarray(stamps, dtype=" NDArray[np.float64] | None: + """Read per-point capture times; absence falls back to the header at the caller.""" + fields = [field for field in message.fields if field.name == "stamps"] + if not fields: + return None + if len(fields) != 1 or fields[0].datatype != PointField.FLOAT64 or fields[0].count != 1: + raise ValueError("Per-point stamps must be a scalar FLOAT64 field") + return np.asarray(pointcloud_view(message)["stamps"].reshape(-1), dtype=np.float64) + def select_points(message: PointCloud2, keep: NDArray[np.bool_]) -> PointCloud2: """Copy selected point records in row order, retaining every field and point byte. @@ -124,20 +156,20 @@ def select_points(message: PointCloud2, keep: NDArray[np.bool_]) -> PointCloud2: records: NDArray[np.uint8] = np.ndarray( (message.height, message.width, message.point_step), dtype=np.uint8, - buffer=message.data.view(), + buffer=message.data, strides=(message.row_step, message.point_step, 1), ) selected = records.reshape(message.height * message.width, message.point_step)[keep] width = len(selected) return PointCloud2( - header=message.header, + header=deepcopy(message.header), height=1, width=width, - fields=message.fields, + fields=deepcopy(message.fields), is_bigendian=message.is_bigendian, point_step=message.point_step, row_step=width * message.point_step, - data=selected.tobytes(), + data=np.asarray(np.frombuffer(selected.tobytes(), dtype=np.uint8).copy(), dtype=np.uint8), is_dense=message.is_dense, ) @@ -169,9 +201,9 @@ def concatenate_clouds(first: PointCloud2, second: PointCloud2) -> PointCloud2: pointcloud_view(first) pointcloud_view(second) if not first.width * first.height: - return PointCloud2.decode(second.encode()) + return cdr_decode(cdr_encode(second), PointCloud2) if not second.width * second.height: - return PointCloud2.decode(first.encode()) + return cdr_decode(cdr_encode(first), PointCloud2) if first.header.frame_id != second.header.frame_id: raise ValueError("cannot concatenate point clouds in different frames") if ( @@ -187,11 +219,11 @@ def concatenate_clouds(first: PointCloud2, second: PointCloud2) -> PointCloud2: header=Header(frame_id=first.header.frame_id, stamp=stamp), height=1, width=a.width + b.width, - fields=first.fields, + fields=deepcopy(first.fields), is_bigendian=first.is_bigendian, point_step=first.point_step, row_step=a.row_step + b.row_step, - data=bytes(a.data) + bytes(b.data), + data=np.asarray(np.concatenate((a.data, b.data)), dtype=np.uint8), is_dense=first.is_dense and second.is_dense, ) @@ -225,9 +257,9 @@ def transform_cloud(message: PointCloud2, transform: TransformStamped) -> PointC ) for index, name in enumerate(("x", "y", "z")): writable[name] = transformed[:, index].reshape(message.height, message.width) - output = PointCloud2.decode(message.encode()) - output.data = bytes(payload) - output.header = Header(stamp=message.header.stamp, frame_id=transform.header.frame_id) + output = cdr_decode(cdr_encode(message), PointCloud2) + output.data = np.frombuffer(payload, dtype=np.uint8).copy() + output.header = Header(stamp=deepcopy(message.header.stamp), frame_id=transform.header.frame_id) return output @@ -246,7 +278,7 @@ def pointcloud_from_xyz_rgb( cloud.fields.append(PointField(name="rgb", offset=12, datatype=PointField.UINT32, count=1)) cloud.point_step = 16 cloud.row_step = cloud.width * cloud.point_step - cloud.data = records.tobytes() + cloud.data = np.frombuffer(records.tobytes(), dtype=np.uint8).copy() return cloud diff --git a/dimos/msgs/protocol.py b/dimos/msgs/protocol.py index 11e218e9009..ecac648d0d1 100644 --- a/dimos/msgs/protocol.py +++ b/dimos/msgs/protocol.py @@ -14,21 +14,9 @@ from typing import ClassVar, Protocol, runtime_checkable -from typing_extensions import Self - @runtime_checkable class DimosMsg(Protocol): - """Transport-neutral contract implemented by generated CDR messages.""" - - msg_name: ClassVar[str] - schema: ClassVar[str] - - @classmethod - def decode(cls, data: bytes) -> Self: - """Decode bytes into a message instance.""" - ... + """Native rosbags value identity; serialization belongs to the typestore.""" - def encode(self, little_endian: bool = True) -> bytes: - """Encode this message instance into bytes.""" - ... + __msgtype__: ClassVar[str] diff --git a/dimos/msgs/sensor_msgs/test_CameraInfo.py b/dimos/msgs/sensor_msgs/test_CameraInfo.py index 6967b9c0fd3..da91a779c83 100644 --- a/dimos/msgs/sensor_msgs/test_CameraInfo.py +++ b/dimos/msgs/sensor_msgs/test_CameraInfo.py @@ -13,7 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos_generated.sensor_msgs.msg import CameraInfo +from dataclasses import asdict + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.sensor_msgs.msg import CameraInfo, RegionOfInterest from dimos_generated.std_msgs.msg import Header import numpy as np from rosbags.typesys import Stores, get_typestore @@ -32,38 +36,45 @@ def test_encode_decode() -> None: height=480, width=640, distortion_model="plumb_bob", - d=[-0.1, 0.05, 0.001, -0.002, 0.0], # 5 distortion coefficients - k=[ - 500.0, - 0.0, - 320.0, # fx, 0, cx - 0.0, - 500.0, - 240.0, # 0, fy, cy - 0.0, - 0.0, - 1.0, - ], # 0, 0, 1 - r=[1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0], - p=[ - 500.0, - 0.0, - 320.0, - 0.0, # fx, 0, cx, Tx - 0.0, - 500.0, - 240.0, - 0.0, # 0, fy, cy, Ty - 0.0, - 0.0, - 1.0, - 0.0, - ], # 0, 0, 1, 0 + d=np.array([-0.1, 0.05, 0.001, -0.002, 0.0], dtype=np.float64), # 5 distortion coefficients + k=np.array( + [ + 500.0, + 0.0, + 320.0, # fx, 0, cx + 0.0, + 500.0, + 240.0, # 0, fy, cy + 0.0, + 0.0, + 1.0, + ], + dtype=np.float64, + ), # 0, 0, 1 + r=np.array([1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0], dtype=np.float64), + p=np.array( + [ + 500.0, + 0.0, + 320.0, + 0.0, # fx, 0, cx, Tx + 0.0, + 500.0, + 240.0, + 0.0, # 0, fy, cy, Ty + 0.0, + 0.0, + 1.0, + 0.0, + ], + dtype=np.float64, + ), # 0, 0, 1, 0 binning_x=2, binning_y=2, header=Header( frame_id="camera_optical_frame", stamp=time_from_nanoseconds(1234567890123456789) ), + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), ) # Set ROI @@ -74,9 +85,11 @@ def test_encode_decode() -> None: original.roi.do_rectify = True # Encode and decode - binary_msg = original.encode() - decoded = CameraInfo.decode(binary_msg) - independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(binary_msg, CameraInfo.msg_name) + binary_msg = cdr_encode(original) + decoded = cdr_decode(binary_msg, CameraInfo) + independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr( + binary_msg, CameraInfo.__msgtype__ + ) assert independent.header.stamp.nanosec == 123456789 assert independent.roi.x_offset == 100 np.testing.assert_array_equal(independent.k, original.k) @@ -153,7 +166,19 @@ def test_numpy_matrix_operations() -> None: """Test numpy matrix getter/setter operations.""" print("\nTesting numpy matrix operations...") - camera_info = CameraInfo() + camera_info = CameraInfo( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + height=0, + width=0, + distortion_model="", + d=np.array([], dtype=np.float64), + k=np.zeros(9, dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), + ) # Test K matrix K = np.array([[525.0, 0.0, 319.5], [0.0, 525.0, 239.5], [0.0, 0.0, 1.0]]) @@ -194,28 +219,46 @@ def test_equality() -> None: height=480, width=640, distortion_model="plumb_bob", - d=[-0.1, 0.05, 0.0, 0.0, 0.0], - header=Header(frame_id="camera1"), + d=np.array([-0.1, 0.05, 0.0, 0.0, 0.0], dtype=np.float64), + header=Header(frame_id="camera1", stamp=Time(sec=0, nanosec=0)), + k=np.zeros(9, dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), ) info2 = CameraInfo( height=480, width=640, distortion_model="plumb_bob", - d=[-0.1, 0.05, 0.0, 0.0, 0.0], - header=Header(frame_id="camera1"), + d=np.array([-0.1, 0.05, 0.0, 0.0, 0.0], dtype=np.float64), + header=Header(frame_id="camera1", stamp=Time(sec=0, nanosec=0)), + k=np.zeros(9, dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), ) info3 = CameraInfo( height=720, width=1280, # Different resolution distortion_model="plumb_bob", - d=[-0.1, 0.05, 0.0, 0.0, 0.0], - header=Header(frame_id="camera1"), + d=np.array([-0.1, 0.05, 0.0, 0.0, 0.0], dtype=np.float64), + header=Header(frame_id="camera1", stamp=Time(sec=0, nanosec=0)), + k=np.zeros(9, dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), ) - assert info1 == info2, "Identical CameraInfo objects should be equal" - assert info1 != info3, "Different CameraInfo objects should not be equal" + np.testing.assert_equal(asdict(info1), asdict(info2)) + assert (info1.height, info1.width) != (info3.height, info3.width) assert info1 != "not_camera_info", "CameraInfo should not equal non-CameraInfo object" print("✓ Equality comparison works correctly") @@ -236,7 +279,9 @@ def test_camera_info_from_yaml() -> None: ) # Load CameraInfo from YAML - camera_info = camera_info_from_yaml(yaml_path, header=Header(frame_id="camera_optical")) + camera_info = camera_info_from_yaml( + yaml_path, header=Header(frame_id="camera_optical", stamp=Time(sec=0, nanosec=0)) + ) # Verify loaded values assert camera_info.width == 640 diff --git a/dimos/msgs/sensor_msgs/test_CompressedImage.py b/dimos/msgs/sensor_msgs/test_CompressedImage.py index ecae6d26b06..cc7f4313d5f 100644 --- a/dimos/msgs/sensor_msgs/test_CompressedImage.py +++ b/dimos/msgs/sensor_msgs/test_CompressedImage.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import CompressedImage from dimos_generated.std_msgs.msg import Header import numpy as np @@ -81,16 +83,22 @@ def test_png_decode_alpha_is_bgra() -> None: arr[..., 3] = 128 ok, buf = cv2.imencode(".png", arr) assert ok - img = image_from_compressed(CompressedImage(data=buf.tobytes(), format="png")) + img = image_from_compressed( + CompressedImage( + data=np.frombuffer(buf.tobytes(), dtype=np.uint8), + format="png", + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + ) assert img.encoding == "bgra8" assert np.array_equal(image_view(img), arr) def test_lcm_roundtrip(rgb_image) -> None: ci = compressed_image_from_image(rgb_image) - wire = ci.encode() - out = CompressedImage.decode(wire) - assert out.data == ci.data + wire = cdr_encode(ci) + out = cdr_decode(wire, CompressedImage) + np.testing.assert_array_equal(out.data, ci.data) assert "; jpeg compressed " in out.format assert out.header.frame_id == "cam" assert abs(to_nanoseconds(out.header.stamp) - to_nanoseconds(ci.header.stamp)) < 1e-06 @@ -98,7 +106,9 @@ def test_lcm_roundtrip(rgb_image) -> None: def test_jpeg_rejects_depth_formats() -> None: depth = image_from_array( - np.zeros((10, 10), dtype=np.uint16), encoding="mono16", header=Header() + np.zeros((10, 10), dtype=np.uint16), + encoding="mono16", + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) with pytest.raises(ValueError, match="JPEG requires 8-bit"): compressed_image_from_image(depth) @@ -106,7 +116,9 @@ def test_jpeg_rejects_depth_formats() -> None: def test_png_rejects_float_depth() -> None: depth = image_from_array( - np.zeros((10, 10), dtype=np.float32), encoding="32FC1", header=Header() + np.zeros((10, 10), dtype=np.float32), + encoding="32FC1", + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) with pytest.raises(ValueError, match="PNG cannot encode"): compressed_image_from_image(depth, format="png") @@ -180,7 +192,9 @@ def test_compressed_uint16_depth_keeps_units_and_rerun_depth(codec: str) -> None import rerun as rr pixels = np.arange(80, dtype=np.uint16).reshape(10, 8) - source = image_from_array(pixels, encoding="16UC1", header=Header(frame_id="depth")) + source = image_from_array( + pixels, encoding="16UC1", header=Header(frame_id="depth", stamp=Time(sec=0, nanosec=0)) + ) compressed = compressed_image_from_image(source, format=codec) decoded = image_from_compressed(compressed) assert decoded.encoding == "16UC1" diff --git a/dimos/msgs/sensor_msgs/test_Joy.py b/dimos/msgs/sensor_msgs/test_Joy.py index d49c81bee90..a27697f34bd 100644 --- a/dimos/msgs/sensor_msgs/test_Joy.py +++ b/dimos/msgs/sensor_msgs/test_Joy.py @@ -14,8 +14,13 @@ # limitations under the License. +from dataclasses import asdict + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import Joy from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from rosbags.typesys import Stores, get_typestore @@ -26,50 +31,62 @@ def test_cdr_encode_decode() -> None: """Test CDR encode/decode preserves Joy data.""" print("Testing Joy CDR encode/decode...") original = Joy( - axes=[0.5, -0.25, 1.0, -1.0, 0.0, 0.75], - buttons=[1, 0, 0, 1, 1, 0, 0, 0, 1, 0, 0, 0], + axes=np.array([0.5, -0.25, 1.0, -1.0, 0.0, 0.75], dtype=np.float32), + buttons=np.array([1, 0, 0, 1, 1, 0, 0, 0, 1, 0, 0, 0], dtype=np.int32), header=Header(stamp=time_from_seconds(1234567890.1234567), frame_id="gamepad"), ) - encoded = original.encode() + encoded = cdr_encode(original) assert isinstance(encoded, bytes) assert len(encoded) > 0 - decoded = Joy.decode(encoded) + decoded = cdr_decode(encoded, Joy) assert abs(to_nanoseconds(decoded.header.stamp) - to_nanoseconds(original.header.stamp)) < 1e-09 assert decoded.header.frame_id == original.header.frame_id - assert decoded.axes == original.axes - assert decoded.buttons == original.buttons + np.testing.assert_array_equal(decoded.axes, original.axes) + np.testing.assert_array_equal(decoded.buttons, original.buttons) print("✓ Joy CDR encode/decode test passed") def test_initialization_methods() -> None: """Test various initialization methods for Joy.""" print("Testing Joy initialization methods...") - joy1 = Joy() - assert joy1.axes == [] - assert joy1.buttons == [] + joy1 = Joy( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + axes=np.array([], dtype=np.float32), + buttons=np.array([], dtype=np.int32), + ) + np.testing.assert_array_equal(joy1.axes, []) + np.testing.assert_array_equal(joy1.buttons, []) assert joy1.header.frame_id == "" assert to_nanoseconds(joy1.header.stamp) == 0 joy2 = Joy( - axes=[0.1, 0.2, 0.3], - buttons=[1, 0, 1], + axes=np.array([0.1, 0.2, 0.3], dtype=np.float32), + buttons=np.array([1, 0, 1], dtype=np.int32), header=Header(stamp=time_from_seconds(1234567890.0), frame_id="xbox_controller"), ) assert to_nanoseconds(joy2.header.stamp) == 1234567890000000000 assert joy2.header.frame_id == "xbox_controller" assert list(joy2.axes) == pytest.approx([0.1, 0.2, 0.3], abs=1e-6) - assert joy2.buttons == [1, 0, 1] - joy3 = Joy(axes=[0.5, -0.5], buttons=[1, 1, 0]) - assert joy3.axes == [0.5, -0.5] - assert joy3.buttons == [1, 1, 0] - joy4 = Joy(axes=[0.7, 0.8], buttons=[0, 1], header=Header(frame_id="ps4_controller")) + np.testing.assert_array_equal(joy2.buttons, [1, 0, 1]) + joy3 = Joy( + axes=np.array([0.5, -0.5], dtype=np.float32), + buttons=np.array([1, 1, 0], dtype=np.int32), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + np.testing.assert_array_equal(joy3.axes, [0.5, -0.5]) + np.testing.assert_array_equal(joy3.buttons, [1, 1, 0]) + joy4 = Joy( + axes=np.array([0.7, 0.8], dtype=np.float32), + buttons=np.array([0, 1], dtype=np.int32), + header=Header(frame_id="ps4_controller", stamp=Time(sec=0, nanosec=0)), + ) assert list(joy4.axes) == pytest.approx([0.7, 0.8], abs=1e-6) - assert joy4.buttons == [0, 1] + np.testing.assert_array_equal(joy4.buttons, [0, 1]) assert joy4.header.frame_id == "ps4_controller" - joy5 = Joy.decode(joy2.encode()) + joy5 = cdr_decode(cdr_encode(joy2), Joy) assert to_nanoseconds(joy5.header.stamp) == to_nanoseconds(joy2.header.stamp) assert joy5.header.frame_id == joy2.header.frame_id - assert joy5.axes == joy2.axes - assert joy5.buttons == joy2.buttons + np.testing.assert_array_equal(joy5.axes, joy2.axes) + np.testing.assert_array_equal(joy5.buttons, joy2.buttons) assert joy5 is not joy2 print("✓ Joy initialization methods test passed") @@ -78,28 +95,28 @@ def test_equality() -> None: """Test Joy equality comparison.""" print("Testing Joy equality...") joy1 = Joy( - axes=[0.5, -0.5], - buttons=[1, 0, 1], + axes=np.array([0.5, -0.5], dtype=np.float32), + buttons=np.array([1, 0, 1], dtype=np.int32), header=Header(stamp=time_from_seconds(1000.0), frame_id="controller1"), ) joy2 = Joy( - axes=[0.5, -0.5], - buttons=[1, 0, 1], + axes=np.array([0.5, -0.5], dtype=np.float32), + buttons=np.array([1, 0, 1], dtype=np.int32), header=Header(stamp=time_from_seconds(1000.0), frame_id="controller1"), ) joy3 = Joy( - axes=[0.5, -0.5], - buttons=[1, 0, 1], + axes=np.array([0.5, -0.5], dtype=np.float32), + buttons=np.array([1, 0, 1], dtype=np.int32), header=Header(stamp=time_from_seconds(1000.0), frame_id="controller2"), ) joy4 = Joy( - axes=[0.6, -0.5], - buttons=[1, 0, 1], + axes=np.array([0.6, -0.5], dtype=np.float32), + buttons=np.array([1, 0, 1], dtype=np.int32), header=Header(stamp=time_from_seconds(1000.0), frame_id="controller1"), ) - assert joy1 == joy2 - assert joy1 != joy3 - assert joy1 != joy4 + np.testing.assert_equal(asdict(joy1), asdict(joy2)) + assert joy1.header != joy3.header + assert not np.array_equal(joy1.axes, joy4.axes) assert joy1 != "not a joy" assert joy1 != 42 print("✓ Joy equality test passed") @@ -108,10 +125,10 @@ def test_equality() -> None: def test_independent_schema_decoding() -> None: joy = Joy( header=Header(stamp=time_from_nanoseconds(1234567890123456789), frame_id="test_controller"), - axes=[0.1, -0.2, 0.3, 0.4], - buttons=[1, 0, 1, 0, 0, 1], + axes=np.array([0.1, -0.2, 0.3, 0.4], dtype=np.float32), + buttons=np.array([1, 0, 1, 0, 0, 1], dtype=np.int32), ) - decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(joy.encode(), Joy.msg_name) + decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(cdr_encode(joy), Joy.__msgtype__) assert decoded.header.frame_id == "test_controller" assert decoded.header.stamp.sec == 1234567890 assert decoded.header.stamp.nanosec == 123456789 @@ -122,25 +139,37 @@ def test_independent_schema_decoding() -> None: def test_edge_cases() -> None: """Test Joy with edge cases.""" print("Testing Joy edge cases...") - joy1 = Joy(axes=[], buttons=[]) - assert joy1.axes == [] - assert joy1.buttons == [] - encoded = joy1.encode() - decoded = Joy.decode(encoded) - assert decoded.axes == [] - assert decoded.buttons == [] + joy1 = Joy( + axes=np.array([], dtype=np.float32), + buttons=np.array([], dtype=np.int32), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + np.testing.assert_array_equal(joy1.axes, []) + np.testing.assert_array_equal(joy1.buttons, []) + encoded = cdr_encode(joy1) + decoded = cdr_decode(encoded, Joy) + np.testing.assert_array_equal(decoded.axes, []) + np.testing.assert_array_equal(decoded.buttons, []) many_axes = [float(i) / 100.0 for i in range(20)] many_buttons = [i % 2 for i in range(32)] - joy2 = Joy(axes=many_axes, buttons=many_buttons) + joy2 = Joy( + axes=np.asarray(many_axes, dtype=np.float32), + buttons=np.asarray(many_buttons, dtype=np.int32), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) assert len(joy2.axes) == 20 assert len(joy2.buttons) == 32 - encoded = joy2.encode() - decoded = Joy.decode(encoded) + encoded = cdr_encode(joy2) + decoded = cdr_decode(encoded, Joy) assert len(decoded.axes) == len(many_axes) for i, (a, b) in enumerate(zip(decoded.axes, many_axes, strict=False)): assert abs(a - b) < 1e-06, f"Axis {i}: {a} != {b}" - assert decoded.buttons == many_buttons + np.testing.assert_array_equal(decoded.buttons, many_buttons) extreme_axes = [-1.0, 1.0, 0.0, -0.999999, 0.999999] - joy3 = Joy(axes=extreme_axes) + joy3 = Joy( + axes=np.asarray(extreme_axes, dtype=np.float32), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + buttons=np.array([], dtype=np.int32), + ) assert list(joy3.axes) == pytest.approx(extreme_axes, abs=1e-6) print("✓ Joy edge cases test passed") diff --git a/dimos/msgs/sensor_msgs/test_PointCloud2.py b/dimos/msgs/sensor_msgs/test_PointCloud2.py index fe6394ec6e9..45fa1c5672c 100644 --- a/dimos/msgs/sensor_msgs/test_PointCloud2.py +++ b/dimos/msgs/sensor_msgs/test_PointCloud2.py @@ -14,26 +14,95 @@ # limitations under the License. +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.sensor_msgs.msg import PointCloud2, PointField +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import ( + cloud_bounds_intersect, + pointcloud_from_xyz, + pointcloud_from_xyz_rgb, + pointcloud_to_open3d, + pointcloud_view, + pointcloud_xyz, +) +from dimos.msgs.time import time_from_seconds, to_nanoseconds from dimos.robot.unitree.type.lidar import pointcloud2_from_webrtc_lidar -from dimos.utils.testing.replay import SensorReplay +from dimos.visualization.rerun.message_helpers import cloud_archetype + + +def _cloud( + points, + *, + frame_id="", + timestamp=0.0, + intensities=None, + offset_times=None, + tags=None, + lines=None, +): + extras = [ + (name, values, dtype, datatype) + for name, values, dtype, datatype in [ + ("intensity", intensities, " None: - """Test LCM encode/decode preserves pointcloud data.""" - replay = SensorReplay("office_lidar", autocast=pointcloud2_from_webrtc_lidar) - lidar_msg: PointCloud2 = replay.load_one("lidar_data_021") - binary_msg = lidar_msg.lcm_encode() - decoded = PointCloud2.lcm_decode(binary_msg) +def test_cdr_encode_decode() -> None: + points = np.arange(300, dtype=np.float32).reshape(100, 3) / 10 + lidar_msg = pointcloud2_from_webrtc_lidar({"data": {"stamp": 12.5, "data": {"points": points}}}) + + binary_msg = cdr_encode(lidar_msg) + decoded = cdr_decode(binary_msg, PointCloud2) # 1. Check number of points - original_points, _ = lidar_msg.as_numpy() - decoded_points, _ = decoded.as_numpy() + original_points = pointcloud_xyz(lidar_msg) + decoded_points = pointcloud_xyz(decoded) assert len(original_points) == len(decoded_points), ( f"Point count mismatch: {len(original_points)} vs {len(decoded_points)}" @@ -50,68 +119,72 @@ def test_lcm_encode_decode() -> None: ) # 3. Check frame_id is preserved - assert lidar_msg.frame_id == decoded.frame_id, ( - f"Frame ID mismatch: '{lidar_msg.frame_id}' vs '{decoded.frame_id}'" + assert lidar_msg.header.frame_id == decoded.header.frame_id, ( + f"Frame ID mismatch: '{lidar_msg.header.frame_id}' vs '{decoded.header.frame_id}'" ) # 4. Check timestamp is preserved (within reasonable tolerance for float precision) - if lidar_msg.ts is not None and decoded.ts is not None: - assert abs(lidar_msg.ts - decoded.ts) < 1e-6, ( - f"Timestamp mismatch: {lidar_msg.ts} vs {decoded.ts}" + if ( + to_nanoseconds(lidar_msg.header.stamp) is not None + and to_nanoseconds(decoded.header.stamp) is not None + ): + assert ( + abs(to_nanoseconds(lidar_msg.header.stamp) - to_nanoseconds(decoded.header.stamp)) + < 1e-6 + ), ( + f"Timestamp mismatch: {to_nanoseconds(lidar_msg.header.stamp)} vs {to_nanoseconds(decoded.header.stamp)}" ) # 5. Check pointcloud properties - assert len(lidar_msg.pointcloud.points) == len(decoded.pointcloud.points), ( - "Open3D pointcloud size mismatch" - ) + assert len(pointcloud_to_open3d(lidar_msg).points) == len( + pointcloud_to_open3d(decoded).points + ), "Open3D pointcloud size mismatch" -def test_lcm_intensity_round_trip() -> None: +def test_cdr_intensity_round_trip() -> None: """Test that intensity values survive an lcm_encode → lcm_decode round trip.""" points = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], dtype=np.float32) intensities = np.array([0.25, 1.1, 0.0], dtype=np.float32) - original = PointCloud2.from_numpy( - points, frame_id="map", timestamp=42.0, intensities=intensities - ) + original = _cloud(points, frame_id="map", timestamp=42.0, intensities=intensities) # Verify getter before encoding - got = original.intensities_f32() + got = _field(original, "intensity") assert got is not None, "intensities_f32() returned None on source cloud" np.testing.assert_allclose(got, intensities, atol=1e-6) # Round-trip through LCM - binary = original.lcm_encode() - decoded = PointCloud2.lcm_decode(binary) + binary = cdr_encode(original) + decoded = cdr_decode(binary, PointCloud2) # Positions preserved - decoded_pts, _ = decoded.as_numpy() + decoded_pts = pointcloud_xyz(decoded) np.testing.assert_allclose(decoded_pts.astype(np.float32), points, atol=1e-6) # Intensities preserved - decoded_intensities = decoded.intensities_f32() + decoded_intensities = _field(decoded, "intensity") assert decoded_intensities is not None, "intensities lost after lcm_decode" np.testing.assert_allclose(decoded_intensities, intensities, atol=1e-6) -def test_lcm_no_intensity_round_trip() -> None: +def test_cdr_no_intensity_round_trip() -> None: """Clouds without intensity should round-trip without creating spurious intensities.""" points = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32) - original = PointCloud2.from_numpy(points, frame_id="map", timestamp=1.0) + original = _cloud(points, frame_id="map", timestamp=1.0) - assert original.intensities_f32() is None + assert _field(original, "intensity") is None - binary = original.lcm_encode() - decoded = PointCloud2.lcm_decode(binary) + binary = cdr_encode(original) + decoded = cdr_decode(binary, PointCloud2) # No intensities should appear (all-zero wire data is ignored) - assert decoded.intensities_f32() is None, "Spurious intensities created from zero wire data" + assert _field(decoded, "intensity") is None, "Spurious intensities created from zero wire data" - decoded_pts, _ = decoded.as_numpy() + decoded_pts = pointcloud_xyz(decoded) np.testing.assert_allclose(decoded_pts.astype(np.float32), points, atol=1e-6) -def test_lcm_per_point_timing_round_trip() -> None: +def test_cdr_per_point_timing_round_trip() -> None: """offset_time/tag/line survive an lcm_encode → lcm_decode round trip.""" points = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], dtype=np.float32) intensities = np.array([10.0, 20.0, 30.0], dtype=np.float32) @@ -120,7 +193,7 @@ def test_lcm_per_point_timing_round_trip() -> None: tags = np.array([0, 16, 32], dtype=np.uint8) lines = np.array([0, 1, 3], dtype=np.uint8) - original = PointCloud2.from_numpy( + original = _cloud( points, frame_id="mid360_link", timestamp=100.5, @@ -130,28 +203,28 @@ def test_lcm_per_point_timing_round_trip() -> None: lines=lines, ) - got_offsets = original.offset_times_u32() + got_offsets = _field(original, "offset_time") assert got_offsets is not None np.testing.assert_array_equal(got_offsets, offset_times) - decoded = PointCloud2.lcm_decode(original.lcm_encode()) + decoded = cdr_decode(cdr_encode(original), PointCloud2) - decoded_pts, _ = decoded.as_numpy() + decoded_pts = pointcloud_xyz(decoded) np.testing.assert_allclose(decoded_pts.astype(np.float32), points, atol=1e-6) - decoded_intensities = decoded.intensities_f32() + decoded_intensities = _field(decoded, "intensity") assert decoded_intensities is not None np.testing.assert_allclose(decoded_intensities, intensities, atol=1e-6) - decoded_offsets = decoded.offset_times_u32() + decoded_offsets = _field(decoded, "offset_time") assert decoded_offsets is not None, "offset_time lost after lcm_decode" assert decoded_offsets.dtype == np.uint32 np.testing.assert_array_equal(decoded_offsets, offset_times) - decoded_tags = decoded.tags_u8() + decoded_tags = _field(decoded, "tag") assert decoded_tags is not None, "tag lost after lcm_decode" np.testing.assert_array_equal(decoded_tags, tags) - decoded_lines = decoded.lines_u8() + decoded_lines = _field(decoded, "line") assert decoded_lines is not None, "line lost after lcm_decode" np.testing.assert_array_equal(decoded_lines, lines) @@ -159,80 +232,72 @@ def test_lcm_per_point_timing_round_trip() -> None: def test_bounding_box_intersects() -> None: """Test bounding_box_intersects method with various scenarios.""" # Test 1: Overlapping boxes - pc1 = PointCloud2.from_numpy(np.array([[0, 0, 0], [2, 2, 2]])) - pc2 = PointCloud2.from_numpy(np.array([[1, 1, 1], [3, 3, 3]])) - assert pc1.bounding_box_intersects(pc2) - assert pc2.bounding_box_intersects(pc1) # Should be symmetric + pc1 = _cloud(np.array([[0, 0, 0], [2, 2, 2]])) + pc2 = _cloud(np.array([[1, 1, 1], [3, 3, 3]])) + assert cloud_bounds_intersect(pc1, pc2) + assert cloud_bounds_intersect(pc2, pc1) # Should be symmetric # Test 2: Non-overlapping boxes - pc3 = PointCloud2.from_numpy(np.array([[0, 0, 0], [1, 1, 1]])) - pc4 = PointCloud2.from_numpy(np.array([[2, 2, 2], [3, 3, 3]])) - assert not pc3.bounding_box_intersects(pc4) - assert not pc4.bounding_box_intersects(pc3) + pc3 = _cloud(np.array([[0, 0, 0], [1, 1, 1]])) + pc4 = _cloud(np.array([[2, 2, 2], [3, 3, 3]])) + assert not cloud_bounds_intersect(pc3, pc4) + assert not cloud_bounds_intersect(pc4, pc3) # Test 3: Touching boxes (edge case - should be True) - pc5 = PointCloud2.from_numpy(np.array([[0, 0, 0], [1, 1, 1]])) - pc6 = PointCloud2.from_numpy(np.array([[1, 1, 1], [2, 2, 2]])) - assert pc5.bounding_box_intersects(pc6) - assert pc6.bounding_box_intersects(pc5) + pc5 = _cloud(np.array([[0, 0, 0], [1, 1, 1]])) + pc6 = _cloud(np.array([[1, 1, 1], [2, 2, 2]])) + assert cloud_bounds_intersect(pc5, pc6) + assert cloud_bounds_intersect(pc6, pc5) # Test 4: One box completely inside another - pc7 = PointCloud2.from_numpy(np.array([[0, 0, 0], [3, 3, 3]])) - pc8 = PointCloud2.from_numpy(np.array([[1, 1, 1], [2, 2, 2]])) - assert pc7.bounding_box_intersects(pc8) - assert pc8.bounding_box_intersects(pc7) + pc7 = _cloud(np.array([[0, 0, 0], [3, 3, 3]])) + pc8 = _cloud(np.array([[1, 1, 1], [2, 2, 2]])) + assert cloud_bounds_intersect(pc7, pc8) + assert cloud_bounds_intersect(pc8, pc7) # Test 5: Boxes overlapping only in 2 dimensions (not all 3) - pc9 = PointCloud2.from_numpy(np.array([[0, 0, 0], [2, 2, 1]])) - pc10 = PointCloud2.from_numpy(np.array([[1, 1, 2], [3, 3, 3]])) - assert not pc9.bounding_box_intersects(pc10) - assert not pc10.bounding_box_intersects(pc9) + pc9 = _cloud(np.array([[0, 0, 0], [2, 2, 1]])) + pc10 = _cloud(np.array([[1, 1, 2], [3, 3, 3]])) + assert not cloud_bounds_intersect(pc9, pc10) + assert not cloud_bounds_intersect(pc10, pc9) # Test 6: Real-world detection scenario with floating point coordinates detection1_points = np.array( [[-3.5, -0.3, 0.1], [-3.3, -0.2, 0.1], [-3.5, -0.3, 0.3], [-3.3, -0.2, 0.3]] ) - pc_det1 = PointCloud2.from_numpy(detection1_points) + pc_det1 = _cloud(detection1_points) detection2_points = np.array( [[-3.4, -0.25, 0.15], [-3.2, -0.15, 0.15], [-3.4, -0.25, 0.35], [-3.2, -0.15, 0.35]] ) - pc_det2 = PointCloud2.from_numpy(detection2_points) + pc_det2 = _cloud(detection2_points) - assert pc_det1.bounding_box_intersects(pc_det2) + assert cloud_bounds_intersect(pc_det1, pc_det2) # Test 7: Single point clouds - pc_single1 = PointCloud2.from_numpy(np.array([[1.0, 1.0, 1.0]])) - pc_single2 = PointCloud2.from_numpy(np.array([[1.0, 1.0, 1.0]])) - pc_single3 = PointCloud2.from_numpy(np.array([[2.0, 2.0, 2.0]])) + pc_single1 = _cloud(np.array([[1.0, 1.0, 1.0]])) + pc_single2 = _cloud(np.array([[1.0, 1.0, 1.0]])) + pc_single3 = _cloud(np.array([[2.0, 2.0, 2.0]])) # Same point should intersect - assert pc_single1.bounding_box_intersects(pc_single2) + assert cloud_bounds_intersect(pc_single1, pc_single2) # Different points should not intersect - assert not pc_single1.bounding_box_intersects(pc_single3) + assert not cloud_bounds_intersect(pc_single1, pc_single3) # Test 8: Empty point clouds - pc_empty1 = PointCloud2.from_numpy(np.array([]).reshape(0, 3)) - pc_empty2 = PointCloud2.from_numpy(np.array([]).reshape(0, 3)) - PointCloud2.from_numpy(np.array([[1.0, 1.0, 1.0]])) + pc_empty1 = _cloud(np.array([]).reshape(0, 3)) + pc_empty2 = _cloud(np.array([]).reshape(0, 3)) + _cloud(np.array([[1.0, 1.0, 1.0]])) - # Empty clouds should handle gracefully (Open3D returns inf bounds) - # This might raise an exception or return False - we should handle gracefully - try: - result = pc_empty1.bounding_box_intersects(pc_empty2) - # If no exception, verify behavior is consistent - assert isinstance(result, bool) - except Exception: - # If it raises an exception, that's also acceptable for empty clouds - pass + assert not cloud_bounds_intersect(pc_empty1, pc_empty2) def test_to_rerun_points_mode_is_screen_space() -> None: """ "points" must be flat screen-space dots, not the world-space spheres branch.""" - cloud = PointCloud2.from_numpy(np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]])) + cloud = _cloud(np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]])) - points = cloud.to_rerun(mode="points", voxel_size=0.05, ui_radius=1.5) - spheres = cloud.to_rerun(mode="spheres", voxel_size=0.05) + points = cloud_archetype(cloud, mode="points", voxel_size=0.05, ui_radius=1.5) + spheres = cloud_archetype(cloud, mode="spheres", voxel_size=0.05) # Negative radii are UI points in rerun; positive ones are world-space. assert points.radii.as_arrow_array().to_pylist() == pytest.approx([-1.5]) @@ -241,20 +306,20 @@ def test_to_rerun_points_mode_is_screen_space() -> None: def test_to_rerun_keeps_the_clouds_own_rgb() -> None: """An RGBD cloud renders in its own colors; rgb=False falls back to the height ramp.""" - import open3d as o3d - - pcd = o3d.geometry.PointCloud() - pcd.points = o3d.utility.Vector3dVector(np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]])) - pcd.colors = o3d.utility.Vector3dVector(np.array([[1.0, 0.0, 0.0], [0.0, 0.0, 1.0]])) - cloud = PointCloud2(pointcloud=pcd) + cloud = pointcloud_from_xyz_rgb( + np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]), + np.array([[255, 0, 0], [0, 0, 255]], dtype=np.uint8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) - colored = cloud.to_rerun(mode="points") + colored = cloud_archetype(cloud, mode="points") assert colored.colors is not None assert colored.class_ids is None - ramp = cloud.to_rerun(mode="points", rgb=False) - assert ramp.colors is None - assert ramp.class_ids is not None + ramp = cloud_archetype(cloud, mode="points", rgb=False) + assert ramp.colors is not None + assert ramp.class_ids is None + assert ramp.colors.as_arrow_array().to_pylist() != colored.colors.as_arrow_array().to_pylist() def test_seq_round_trips_through_lcm_with_and_without_intensity_and_empty() -> None: diff --git a/dimos/msgs/sensor_msgs/test_image.py b/dimos/msgs/sensor_msgs/test_image.py index be1cca8a22b..9c258b86a88 100644 --- a/dimos/msgs/sensor_msgs/test_image.py +++ b/dimos/msgs/sensor_msgs/test_image.py @@ -14,8 +14,10 @@ """File, color, CDR and reactive selection checks on generated images.""" +from dataclasses import asdict from pathlib import Path +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.sensor_msgs.msg import Image from dimos_generated.std_msgs.msg import Header import numpy as np @@ -42,7 +44,9 @@ def img(tmp_path: Path) -> Image: pixels[..., 1] = 37 path = tmp_path / "camera.png" PILImage.fromarray(pixels).save(path) - return image_from_file(path, header=Header(stamp=time_from_nanoseconds(1234567890123456789))) + return image_from_file( + path, header=Header(stamp=time_from_nanoseconds(1234567890123456789), frame_id="") + ) def test_file_load(img: Image) -> None: @@ -59,16 +63,16 @@ def test_file_load(img: Image) -> None: def test_cdr_encode_decode(img: Image) -> None: - decoded = Image.decode(img.encode()) + decoded = cdr_decode(cdr_encode(img), Image) assert decoded is not img - assert decoded == img + np.testing.assert_equal(asdict(decoded), asdict(img)) def test_rgb_bgr_conversion(img: Image) -> None: bgr = image_from_array(image_to_bgr(img), encoding="bgr8", header=img.header) assert bgr != img restored = image_from_array(image_to_rgb(bgr), encoding="rgb8", header=img.header) - assert restored == img + np.testing.assert_equal(asdict(restored), asdict(img)) def test_opencv_conversion(img: Image) -> None: diff --git a/dimos/msgs/sim_msgs/Contacts.py b/dimos/msgs/sim_msgs/Contacts.py deleted file mode 100644 index 08df5d2efdf..00000000000 --- a/dimos/msgs/sim_msgs/Contacts.py +++ /dev/null @@ -1,54 +0,0 @@ -# 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. - -"""Which scene geometry a simulated robot touches, by robot part and geometry kind.""" - -from __future__ import annotations - -from dataclasses import dataclass -import json -import time -from typing import Literal - -from dimos_lcm.std_msgs import String as LCMString - -from dimos.types.timestamped import Timestamped - -Part = Literal["trunk", "lidar", "leg", "foot"] -Kind = Literal["floor", "wall", "ceiling", "clutter"] - - -@dataclass(frozen=True, order=True) -class Contact: - part: Part - kind: Kind - - -class Contacts(Timestamped): - """The set of contacts at one instant, published whenever it changes.""" - - msg_name = "sim_msgs.Contacts" - - def __init__(self, contacts: list[Contact] | None = None, ts: float | None = None) -> None: - self.contacts = contacts or [] - self.ts = time.time() if ts is None else ts - - def lcm_encode(self) -> bytes: - payload = {"ts": self.ts, "contacts": [[c.part, c.kind] for c in self.contacts]} - return LCMString(data=json.dumps(payload)).lcm_encode() # type: ignore[no-any-return] - - @classmethod - def lcm_decode(cls, data: bytes) -> Contacts: - payload = json.loads(LCMString.lcm_decode(data).data) - return cls([Contact(part, kind) for part, kind in payload["contacts"]], payload["ts"]) diff --git a/dimos/msgs/sim_msgs/test_Contacts.py b/dimos/msgs/sim_msgs/test_Contacts.py index 840c9c5f2f8..326ade33159 100644 --- a/dimos/msgs/sim_msgs/test_Contacts.py +++ b/dimos/msgs/sim_msgs/test_Contacts.py @@ -12,19 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. -from __future__ import annotations +from dimos_cdrgen.registry import decode, encode +from dimos_generated.dimos_msgs.msg import Contact, Contacts +from dimos_generated.std_msgs.msg import Header +import pytest -from dimos.msgs.sim_msgs.Contacts import Contact, Contacts +from dimos.msgs.time import time_from_seconds, to_seconds -def test_round_trip_keeps_contacts_and_stamp() -> None: - msg = Contacts([Contact("foot", "floor"), Contact("trunk", "wall")], ts=12.5) - back = Contacts.lcm_decode(msg.lcm_encode()) - assert back.contacts == msg.contacts - assert back.ts == 12.5 - - -def test_empty_set_round_trips() -> None: - back = Contacts.lcm_decode(Contacts(ts=1.0).lcm_encode()) - assert back.contacts == [] - assert back.ts == 1.0 +@pytest.mark.parametrize( + "contacts", [[], [Contact(part="foot", kind="floor"), Contact(part="trunk", kind="wall")]] +) +def test_cdr_retains_semantic_contacts_and_source_stamp(contacts): + value = Contacts(header=Header(stamp=time_from_seconds(12.5), frame_id=""), contacts=contacts) + decoded = decode(encode(value), Contacts) + assert decoded.contacts == contacts + assert to_seconds(decoded.header.stamp) == 12.5 diff --git a/dimos/msgs/std_msgs/test_header.py b/dimos/msgs/std_msgs/test_header.py index a1849aae5fd..2e6f4354445 100644 --- a/dimos/msgs/std_msgs/test_header.py +++ b/dimos/msgs/std_msgs/test_header.py @@ -17,6 +17,8 @@ from datetime import datetime, timezone import time +from dimos_cdrgen.registry import encode as cdr_encode, schema as cdr_schema +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.std_msgs.msg import Header from rosbags.typesys import Stores, get_typestore @@ -33,7 +35,7 @@ def test_header_initialization_methods() -> None: after = time.time_ns() assert before <= current.stamp.sec * 1000000000 + current.stamp.nanosec <= after assert current.frame_id == "base_link" - empty = Header() + empty = Header(stamp=Time(sec=0, nanosec=0), frame_id="") assert empty.stamp.sec == empty.stamp.nanosec == 0 assert empty.frame_id == "" dt = datetime(2025, 1, 18, 12, 30, 45, 500000, tzinfo=timezone.utc) @@ -55,8 +57,10 @@ def test_header_datetime_conversion() -> None: def test_header_independent_cdr_fields() -> None: header = Header(stamp=time_from_seconds(100.5), frame_id="map") - decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(header.encode(), Header.msg_name) + decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr( + cdr_encode(header), Header.__msgtype__ + ) assert decoded.stamp.sec == 100 assert decoded.stamp.nanosec == 500000000 assert decoded.frame_id == "map" - assert "seq" not in Header.schema + assert "seq" not in cdr_schema(Header.__msgtype__) diff --git a/dimos/msgs/test_camera_info.py b/dimos/msgs/test_camera_info.py index 641f71e3220..81dcbd8955a 100644 --- a/dimos/msgs/test_camera_info.py +++ b/dimos/msgs/test_camera_info.py @@ -17,6 +17,7 @@ import math from pathlib import Path +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import CameraInfo from dimos_generated.std_msgs.msg import Header @@ -39,7 +40,7 @@ def test_real_go2_calibration_survives_cdr_with_exact_header(): header = Header(frame_id="camera_optical", stamp=Time(sec=1700000000, nanosec=123456789)) source = yaml.safe_load(_CALIBRATION.read_text()) message = camera_info_from_yaml(_CALIBRATION, header=header) - decoded = CameraInfo.decode(message.encode()) + decoded = cdr_decode(cdr_encode(message), CameraInfo) assert decoded.header == header assert (decoded.width, decoded.height) == (1280, 720) assert decoded.distortion_model == "equidistant" @@ -60,7 +61,9 @@ def test_real_go2_calibration_survives_cdr_with_exact_header(): @pytest.mark.parametrize("axis,focal", [("vertical", 240), ("horizontal", 320)]) def test_fov_converts_selected_axis_and_builds_projection(axis, focal): - message = camera_info_from_fov(90, 640, 480, header=Header(frame_id="optical"), axis=axis) + message = camera_info_from_fov( + 90, 640, 480, header=Header(frame_id="optical", stamp=Time(sec=0, nanosec=0)), axis=axis + ) np.testing.assert_allclose(message.k, [focal, 0, 320, 0, focal, 240, 0, 0, 1]) np.testing.assert_allclose(message.p, [focal, 0, 320, 0, 0, focal, 240, 0, 0, 0, 1, 0]) np.testing.assert_array_equal(message.r, np.eye(3).reshape(-1)) @@ -70,19 +73,31 @@ def test_fov_converts_selected_axis_and_builds_projection(axis, focal): @pytest.mark.parametrize("fov", [0, -1, 180, 181, math.nan, math.inf]) def test_fov_rejects_nonphysical_angles(fov): with pytest.raises(ValueError, match="Field of view"): - camera_info_from_fov(fov, 640, 480, header=Header()) + camera_info_from_fov( + fov, 640, 480, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="") + ) @pytest.mark.parametrize("width,height", [(0, 480), (640, -1), (True, 480), (1.5, 480)]) def test_intrinsics_reject_invalid_dimensions(width, height): with pytest.raises(ValueError, match="positive integers"): - camera_info_from_intrinsics(300, 300, 320, 240, width, height, header=Header()) + camera_info_from_intrinsics( + 300, + 300, + 320, + 240, + width, + height, + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) @pytest.mark.parametrize("fx", [0, -1, math.nan, math.inf]) def test_intrinsics_reject_invalid_focal_length(fx): with pytest.raises(ValueError, match="positive focal"): - camera_info_from_intrinsics(fx, 300, 320, 240, 640, 480, header=Header()) + camera_info_from_intrinsics( + fx, 300, 320, 240, 640, 480, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="") + ) @pytest.mark.parametrize("mutation", ["rows", "length", "nonfinite"]) @@ -98,7 +113,7 @@ def test_yaml_rejects_invalid_camera_matrix(tmp_path, mutation): path = tmp_path / "invalid.yaml" path.write_text(yaml.safe_dump(data)) with pytest.raises(ValueError, match="camera_matrix"): - camera_info_from_yaml(path, header=Header()) + camera_info_from_yaml(path, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) def test_stamp_copy_preserves_calibration_and_does_not_mutate_prior_publish(): diff --git a/dimos/msgs/test_control.py b/dimos/msgs/test_control.py new file mode 100644 index 00000000000..3d85ec348d5 --- /dev/null +++ b/dimos/msgs/test_control.py @@ -0,0 +1,114 @@ +# Copyright 2025-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. + +"""Generated control values retain the sparse frame contract over CDR.""" + +import time + +from dimos_cdrgen.registry import decode, encode +from dimos_generated.dimos_msgs.msg import ControlValues +import numpy as np +import pytest + +from dimos.msgs.control import control_values, control_values_dict, validate_control_values +from dimos.msgs.helpers import resolve_msg_type +from dimos.msgs.protocol import DimosMsg + + +def test_cdr_round_trip_preserves_every_field() -> None: + frame = control_values( + "xarm", + source_ts=1758547200.25, + epoch=7, + sequence=1234, + interface_names=["xarm/j1/position", "xarm/gripper/position"], + values=[-1.25, 0.085], + ) + decoded = decode(encode(frame), ControlValues) + assert (decoded.source, decoded.source_ts, decoded.epoch, decoded.sequence) == ( + "xarm", + 1758547200.25, + 7, + 1234, + ) + assert decoded.interface_names == ["xarm/j1/position", "xarm/gripper/position"] + np.testing.assert_array_equal(decoded.values, [-1.25, 0.085]) + assert control_values_dict(decoded) == { + "xarm/j1/position": -1.25, + "xarm/gripper/position": 0.085, + } + + +def test_empty_heartbeat_round_trip() -> None: + decoded = decode( + encode(control_values("coordinator", source_ts=99.5, epoch=3, sequence=42)), ControlValues + ) + assert (decoded.source, decoded.source_ts, decoded.epoch, decoded.sequence) == ( + "coordinator", + 99.5, + 3, + 42, + ) + assert decoded.interface_names == [] + np.testing.assert_array_equal(decoded.values, []) + assert control_values_dict(decoded) == {} + + +@pytest.mark.parametrize("names,values", [(["a", "b"], [1.0]), ([], [1.0])]) +def test_mismatched_lengths_raise(names, values) -> None: + with pytest.raises(ValueError, match="same length"): + control_values("arm", interface_names=names, values=values) + + +@pytest.mark.parametrize("bad", [float("nan"), float("inf"), float("-inf")]) +def test_non_finite_values_name_the_interface(bad: float) -> None: + with pytest.raises(ValueError, match="arm/j2/effort"): + control_values( + "arm", interface_names=["arm/j1/position", "arm/j2/effort"], values=[0.0, bad] + ) + + +def test_decode_boundary_rejects_ragged_generated_frame() -> None: + frame = ControlValues( + source="arm", + source_ts=1.0, + epoch=1, + sequence=1, + interface_names=["a", "b"], + values=np.array([1.0]), + ) + decoded = decode(encode(frame), ControlValues) + with pytest.raises(ValueError, match="same length"): + validate_control_values(decoded) + + +def test_constructor_copies_arrays() -> None: + names, values = ["a"], np.array([2.0]) + frame = control_values("arm", interface_names=names, values=values) + names[0], values[0] = "b", 99 + assert frame.interface_names == ["a"] + np.testing.assert_array_equal(frame.values, [2.0]) + + +def test_only_absent_timestamp_uses_current_time() -> None: + before = time.time() + frame = control_values("arm") + assert before <= frame.source_ts <= time.time() + decoded = decode(encode(control_values("arm", source_ts=0.0)), ControlValues) + assert decoded.source_ts == 0.0 + + +def test_generated_type_is_the_resolved_wire_type() -> None: + assert isinstance(control_values("arm"), DimosMsg) + assert resolve_msg_type("dimos_msgs/msg/ControlValues") is ControlValues diff --git a/dimos/msgs/test_generated_projection.py b/dimos/msgs/test_generated_projection.py index 3c200468093..a6d7d0eaa13 100644 --- a/dimos/msgs/test_generated_projection.py +++ b/dimos/msgs/test_generated_projection.py @@ -14,6 +14,7 @@ """Stored image projection uses generated transforms and padded depth pixels.""" +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, Quaternion, @@ -21,7 +22,7 @@ TransformStamped, Vector3, ) -from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, RegionOfInterest from dimos_generated.std_msgs.msg import Header import numpy as np @@ -32,12 +33,26 @@ def test_stored_pose_projection_inverts_world_camera_chain() -> None: obs = Observation( - id=0, ts=0.0, data_type=Image, _data=Image(), pose_tuple=(2.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0) + id=0, + ts=0.0, + data_type=Image, + _data=Image( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + height=0, + width=0, + encoding="", + is_bigendian=0, + step=0, + data=np.array([], dtype=np.uint8), + ), + pose_tuple=(2.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0), ) extrinsic = TransformStamped( - header=Header(frame_id="base_link"), + header=Header(frame_id="base_link", stamp=Time(sec=0, nanosec=0)), child_frame_id="camera_optical", - transform=Transform(translation=Vector3(x=1), rotation=Quaternion(w=1)), + transform=Transform( + translation=Vector3(x=1, y=0.0, z=0.0), rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), ) result = _world_to_optical(obs, "world", None, extrinsic, "camera_optical", 5.0) assert result is not None @@ -48,14 +63,40 @@ def test_stored_pose_projection_inverts_world_camera_chain() -> None: def test_depth_occlusion_reads_big_endian_rows_and_millimeters() -> None: obs = Observation( - id=0, ts=0.0, data_type=Image, _data=Image(), pose_tuple=(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0) + id=0, + ts=0.0, + data_type=Image, + _data=Image( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + height=0, + width=0, + encoding="", + is_bigendian=0, + step=0, + data=np.array([], dtype=np.uint8), + ), + pose_tuple=(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0), ) extrinsic = TransformStamped( - header=Header(frame_id="base_link"), + header=Header(frame_id="base_link", stamp=Time(sec=0, nanosec=0)), child_frame_id="camera_optical", - transform=Transform(rotation=Quaternion(w=1)), + transform=Transform( + rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), translation=Vector3(x=0.0, y=0.0, z=0.0) + ), + ) + calibration = CameraInfo( + width=2, + height=2, + k=np.array([1, 0, 0, 0, 1, 0, 0, 0, 1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + distortion_model="", + d=np.array([], dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), ) - calibration = CameraInfo(width=2, height=2, k=[1, 0, 0, 0, 1, 0, 0, 0, 1]) # Pixel (0, 1) = 1500mm; padding cannot be interpreted as pixels. depth = Image( width=2, @@ -63,7 +104,8 @@ def test_depth_occlusion_reads_big_endian_rows_and_millimeters() -> None: step=6, encoding="16UC1", is_bigendian=1, - data=[0, 0, 0, 0, 99, 99, 5, 220, 0, 0, 99, 99], + data=np.array([0, 0, 0, 0, 99, 99, 5, 220, 0, 0, 99, 99], dtype=np.uint8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) visible = sees( Point(x=0, y=1, z=1), calibration, base_to_optical=extrinsic, depth=lambda _: depth @@ -82,7 +124,17 @@ def test_timestamp_adapter_matches_generated_cloud_without_mutating_header() -> from dimos.msgs.time import time_from_nanoseconds, to_seconds from dimos.types.timestamped import TimestampedData, align_timestamped - source = PointCloud2(header=Header(stamp=time_from_nanoseconds(1700000000123456789))) + source = PointCloud2( + header=Header(stamp=time_from_nanoseconds(1700000000123456789), frame_id=""), + height=0, + width=0, + fields=[], + is_bigendian=False, + point_step=0, + row_step=0, + data=np.array([], dtype=np.uint8), + is_dense=False, + ) primary = Subject() secondary = Subject() received = [] diff --git a/dimos/msgs/test_generated_views.py b/dimos/msgs/test_generated_views.py index 54f600a24ee..7f1258324c1 100644 --- a/dimos/msgs/test_generated_views.py +++ b/dimos/msgs/test_generated_views.py @@ -19,6 +19,8 @@ import math import cv2 +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid, Path from dimos_generated.sensor_msgs.msg import Image @@ -43,7 +45,15 @@ def test_padded_image_view_retains_owner_and_is_readonly() -> None: - msg = Image(width=2, height=2, step=4, encoding="mono8", data=[1, 2, 99, 99, 3, 4, 99, 99]) + msg = Image( + width=2, + height=2, + step=4, + encoding="mono8", + data=np.array([1, 2, 99, 99, 3, 4, 99, 99], dtype=np.uint8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + is_bigendian=0, + ) pixels = image_view(msg) np.testing.assert_array_equal(pixels, [[1, 2], [3, 4]]) assert pixels.strides == (4, 1) @@ -59,7 +69,13 @@ def test_padded_image_view_retains_owner_and_is_readonly() -> None: def test_big_endian_depth_view() -> None: msg = Image( - width=2, height=1, step=4, encoding="16UC1", is_bigendian=1, data=[0x01, 0x02, 0x03, 0x04] + width=2, + height=1, + step=4, + encoding="16UC1", + is_bigendian=1, + data=np.array([0x01, 0x02, 0x03, 0x04], dtype=np.uint8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) np.testing.assert_array_equal(image_view(msg), [[258, 772]]) @@ -85,11 +101,29 @@ def test_image_array_requires_matching_encoding(pixels: np.ndarray, encoding: st @pytest.mark.parametrize("step,data", [(1, [1]), (2, [1]), (2, [1, 2, 3])]) def test_invalid_image_layout(step: int, data: list[int]) -> None: with pytest.raises(ValueError, match="dimensions"): - image_view(Image(width=2, height=1, step=step, data=data, encoding="mono8")) + image_view( + Image( + width=2, + height=1, + step=step, + data=np.asarray(data, dtype=np.uint8), + encoding="mono8", + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + is_bigendian=0, + ) + ) def test_jpeg_preserves_rgb_color() -> None: - msg = Image(width=16, height=16, step=48, encoding="rgb8", data=[255, 0, 0] * 256) + msg = Image( + width=16, + height=16, + step=48, + encoding="rgb8", + data=np.asarray([255, 0, 0] * 256, dtype=np.uint8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + is_bigendian=0, + ) jpeg = image_to_jpeg(msg, quality=95) bgr = cv2.imdecode(np.frombuffer(jpeg, np.uint8), cv2.IMREAD_COLOR) assert bgr.shape == (16, 16, 3) @@ -97,7 +131,20 @@ def test_jpeg_preserves_rgb_color() -> None: def test_occupancy_view_and_obstacle_preserving_reduction() -> None: - msg = OccupancyGrid(info=MapMetaData(width=4, height=2), data=[-1, -1, 0, 100, -1, -1, 10, 20]) + msg = OccupancyGrid( + info=MapMetaData( + width=4, + height=2, + map_load_time=Time(sec=0, nanosec=0), + resolution=0.0, + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([-1, -1, 0, 100, -1, -1, 10, 20], dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) cells = occupancy_view(msg) assert not cells.flags.writeable np.testing.assert_array_equal(block_max_reduce(cells, 2), [[-1, 100]]) @@ -108,17 +155,20 @@ def test_occupancy_view_and_obstacle_preserving_reduction() -> None: def test_web_pose_path_and_goal_use_nested_ros_fields() -> None: pose = PoseStamped( - header=Header(stamp=time_from_nanoseconds(1500000000)), + header=Header(stamp=time_from_nanoseconds(1500000000), frame_id=""), pose=Pose( - position=Point(x=2, y=3), orientation=Quaternion(z=math.sqrt(0.5), w=math.sqrt(0.5)) + position=Point(x=2, y=3, z=0.0), + orientation=Quaternion(z=math.sqrt(0.5), w=math.sqrt(0.5), x=0.0, y=0.0), ), ) assert yaw(pose.pose.orientation) == pytest.approx(math.pi / 2) value = json.loads(encode_pose(pose)) assert value == {"x": 2, "y": 3, "z": 0, "yaw": pytest.approx(math.pi / 2), "ts": 1.5} - assert json.loads(encode_path(Path(poses=[pose]))) == [[2, 3]] + assert json.loads( + encode_path(Path(poses=[pose], header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""))) + ) == [[2, 3]] goal = decode_point({"x": 4, "y": 5}) - assert goal.point == Point(x=4, y=5) and goal.header.frame_id == "world" + assert goal.point == Point(x=4, y=5, z=0.0) and goal.header.frame_id == "world" @pytest.mark.parametrize("encoding", ["rgb8", "bgr8", "rgba8", "bgra8", "mono8"]) @@ -137,7 +187,17 @@ def test_sharpness_rejects_depth_and_empty_images(): with pytest.raises(ValueError, match="8-bit visual"): image_sharpness(depth) with pytest.raises(ValueError, match="nonempty"): - image_sharpness(Image(encoding="mono8")) + image_sharpness( + Image( + encoding="mono8", + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + height=0, + width=0, + is_bigendian=0, + step=0, + data=np.array([], dtype=np.uint8), + ) + ) @pytest.mark.parametrize( @@ -151,7 +211,7 @@ def test_sharpness_rejects_depth_and_empty_images(): ) def test_rgb_model_input_is_color_correct_and_independent(encoding: str, pixel: list[int]) -> None: source = image_from_array(np.array([[pixel]], dtype=np.uint8), encoding=encoding) - restored = Image.decode(source.encode()) + restored = cdr_decode(cdr_encode(source), Image) rgb = image_to_rgb(restored) np.testing.assert_array_equal(rgb, [[[255, 0, 17]]]) rgb[0, 0] = 0 @@ -165,7 +225,8 @@ def test_rgb_model_input_respects_big_endian_padded_grayscale() -> None: step=6, encoding="mono16", is_bigendian=1, - data=[0x12, 0x34, 0xFF, 0xFF, 99, 99, 0x01, 0x00, 0, 0, 99, 99], + data=np.array([0x12, 0x34, 0xFF, 0xFF, 99, 99, 0x01, 0x00, 0, 0, 99, 99], dtype=np.uint8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) np.testing.assert_array_equal( image_to_rgb(msg), [[[18, 18, 18], [255, 255, 255]], [[1, 1, 1], [0, 0, 0]]] @@ -184,7 +245,7 @@ def test_mosaic_uses_generated_pixels_and_preserves_rgb_colors() -> None: red = image_from_array(np.array([[[255, 0, 0]]], dtype=np.uint8), encoding="rgb8") blue = image_from_array(np.array([[[255, 0, 0]]], dtype=np.uint8), encoding="bgr8") observation = mosaic([red, blue], cols=2, cell_height=1) - result = Image.decode(observation.data.encode()) + result = cdr_decode(cdr_encode(observation.data), Image) assert result.encoding == "bgr8" np.testing.assert_array_equal(image_to_rgb(result), [[[255, 0, 0], [0, 0, 255]]]) assert observation.tags == {"mosaic": True} @@ -198,7 +259,8 @@ def test_brightness_ignores_padding_and_normalizes_big_endian_u16() -> None: step=6, encoding="mono16", is_bigendian=1, - data=[0, 0, 255, 255, 255, 255], + data=np.array([0, 0, 255, 255, 255, 255], dtype=np.uint8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) assert image_brightness(message) == pytest.approx(0.5) depth = image_from_array(np.array([[1.0]], dtype=np.float32), encoding="32FC1") @@ -215,7 +277,7 @@ def test_image_file_decodes_rgb_and_preserves_source_header(tmp_path, mode): header = Header(stamp=time_from_nanoseconds(1234567890), frame_id="camera") message = image_from_file(path, header=header) assert message.encoding == "rgb8" - assert message.header.encode() == header.encode() + assert cdr_encode(message.header) == cdr_encode(header) np.testing.assert_array_equal(image_view(message), [[[7, 7, 7], [19, 19, 19]]]) @@ -226,7 +288,7 @@ def test_image_file_decodes_rgb_and_preserves_source_header(tmp_path, mode): def test_remaining_raw_encodings_keep_signed_float_and_endian_values(encoding, dtype, shape): pixels = (np.arange(np.prod(shape)).reshape(shape) - 5).astype(dtype) source = image_from_array(pixels, encoding=encoding) - decoded = Image.decode(source.encode()) + decoded = cdr_decode(cdr_encode(source), Image) assert decoded.encoding == encoding assert decoded.is_bigendian == 1 np.testing.assert_array_equal(image_view(decoded), pixels) diff --git a/dimos/msgs/test_geometry.py b/dimos/msgs/test_geometry.py index 38780853d5f..08dfff82c9b 100644 --- a/dimos/msgs/test_geometry.py +++ b/dimos/msgs/test_geometry.py @@ -17,14 +17,18 @@ import math +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, Pose, PoseStamped, + PoseWithCovariance, Quaternion, Transform, TransformStamped, + Twist, + TwistWithCovariance, Vector3, ) from dimos_generated.nav_msgs.msg import Odometry @@ -50,11 +54,25 @@ @pytest.mark.parametrize("seconds,nanoseconds", [(0, 0), (-1, 500000000), (1700000000, 123456789)]) def test_odometry_transform_preserves_frames_stamp_and_pose(seconds, nanoseconds): - message = Odometry(child_frame_id="robot/lidar") + message = Odometry( + child_frame_id="robot/lidar", + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), + twist=TwistWithCovariance( + twist=Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + covariance=np.zeros(36, dtype=np.float64), + ), + ) message.header.frame_id = "robot/odom" message.header.stamp = Time(sec=seconds, nanosec=nanoseconds) message.pose.pose.position = Point(x=1, y=2, z=3) - message.pose.pose.orientation = Quaternion(z=0.6, w=0.8) + message.pose.pose.orientation = Quaternion(z=0.6, w=0.8, x=0.0, y=0.0) transformed = transform_from_odometry(message) @@ -84,7 +102,9 @@ def _edge(parent, child, xyz, rotation, stamp=Time(sec=1700000000, nanosec=12345 def test_transform_composition_rotates_translation_and_preserves_exact_source_stamp(): first = _edge("world", "arm", (1, -1, 0), quaternion_from_euler(0, 0, math.pi / 2)) - second = _edge("arm", "tool", (2, 3, 4), Quaternion(w=1), Time(sec=7)) + second = _edge( + "arm", "tool", (2, 3, 4), Quaternion(w=1, x=0.0, y=0.0, z=0.0), Time(sec=7, nanosec=0) + ) result = compose_transforms(first, second) @@ -115,26 +135,31 @@ def test_inverse_composes_to_identity_and_swaps_frames_without_losing_nanosecond def test_transform_composition_rejects_disconnected_frames(): - first = _edge("a", "b", (0, 0, 0), Quaternion(w=1)) - second = _edge("c", "d", (0, 0, 0), Quaternion(w=1)) + first = _edge("a", "b", (0, 0, 0), Quaternion(w=1, x=0.0, y=0.0, z=0.0)) + second = _edge("c", "d", (0, 0, 0), Quaternion(w=1, x=0.0, y=0.0, z=0.0)) with pytest.raises(ValueError, match="Cannot compose frames"): compose_transforms(first, second) @pytest.mark.parametrize( - "rotation", [Quaternion(w=0), Quaternion(x=math.nan, w=1), Quaternion(w=math.inf)] + "rotation", + [ + Quaternion(w=0, x=0.0, y=0.0, z=0.0), + Quaternion(x=math.nan, w=1, y=0.0, z=0.0), + Quaternion(w=math.inf, x=0.0, y=0.0, z=0.0), + ], ) def test_transform_math_rejects_invalid_rotations(rotation): edge = _edge("a", "b", (0, 0, 0), rotation) with pytest.raises(ValueError): inverse_transform(edge) with pytest.raises(ValueError): - compose_transforms(edge, _edge("b", "c", (0, 0, 0), Quaternion(w=1))) + compose_transforms(edge, _edge("b", "c", (0, 0, 0), Quaternion(w=1, x=0.0, y=0.0, z=0.0))) @pytest.mark.parametrize("value", [math.nan, math.inf, -math.inf]) def test_transform_math_rejects_nonfinite_translations_and_angles(value): - edge = _edge("a", "b", (value, 0, 0), Quaternion(w=1)) + edge = _edge("a", "b", (value, 0, 0), Quaternion(w=1, x=0.0, y=0.0, z=0.0)) with pytest.raises(ValueError, match="Translation components"): inverse_transform(edge) with pytest.raises(ValueError, match="Euler angles"): @@ -142,7 +167,7 @@ def test_transform_math_rejects_nonfinite_translations_and_angles(value): def test_pose_conversion_copies_nested_fields_and_preserves_header(): - edge = _edge("map", "robot", (1, 2, 3), Quaternion(z=0.6, w=0.8)) + edge = _edge("map", "robot", (1, 2, 3), Quaternion(z=0.6, w=0.8, x=0.0, y=0.0)) pose = pose_from_transform(edge) assert pose.header == edge.header assert pose.pose.position == Point(x=1, y=2, z=3) @@ -156,10 +181,12 @@ def test_pose_conversion_copies_nested_fields_and_preserves_header(): def test_pose_transform_preserves_header_and_does_not_alias_input(): source = PoseStamped( header=Header(frame_id="odom", stamp=Time(sec=-1, nanosec=987654321)), - pose=Pose(position=Point(x=1, y=2, z=3), orientation=Quaternion(z=0.6, w=0.8)), + pose=Pose( + position=Point(x=1, y=2, z=3), orientation=Quaternion(z=0.6, w=0.8, x=0.0, y=0.0) + ), ) result = transform_from_pose(source, child_frame_id="robot/base") - decoded = TransformStamped.decode(result.encode()) + decoded = cdr_decode(cdr_encode(result), TransformStamped) assert decoded.header == source.header assert decoded.child_frame_id == "robot/base" assert decoded.transform.translation == Vector3(x=1, y=2, z=3) @@ -188,7 +215,7 @@ def test_pose_matrix_roundtrip_preserves_rigid_motion_without_aliasing(): ) matrix = pose_matrix(pose) np.testing.assert_allclose(matrix @ [1, 0, 0, 1], [2, 0, 3, 1], atol=1e-12) - decoded = Pose.decode(pose_from_matrix(matrix).encode()) + decoded = cdr_decode(cdr_encode(pose_from_matrix(matrix)), Pose) np.testing.assert_allclose(pose_matrix(decoded), matrix, atol=1e-12) matrix[0, 3] = 99 decoded.position.y = 42 @@ -214,9 +241,15 @@ def test_pose_from_matrix_rejects_nonrigid_or_invalid_matrices(matrix): @pytest.mark.parametrize( "pose", [ - Pose(position=Point(x=math.inf), orientation=Quaternion(w=1)), - Pose(orientation=Quaternion(w=0)), - Pose(orientation=Quaternion(w=math.nan)), + Pose( + position=Point(x=math.inf, y=0.0, z=0.0), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), + Pose(orientation=Quaternion(w=0, x=0.0, y=0.0, z=0.0), position=Point(x=0.0, y=0.0, z=0.0)), + Pose( + orientation=Quaternion(w=math.nan, x=0.0, y=0.0, z=0.0), + position=Point(x=0.0, y=0.0, z=0.0), + ), ], ) def test_pose_matrix_rejects_invalid_geometry(pose): @@ -229,7 +262,7 @@ def test_local_offset_uses_pose_axes_and_copies_orientation(): position=Point(x=1, y=2, z=3), orientation=quaternion_from_euler(0, 0, math.pi / 2), ) - translated = translate_pose_local(pose, Vector3(x=2, z=-1)) + translated = translate_pose_local(pose, Vector3(x=2, z=-1, y=0.0)) position = translated.position assert (position.x, position.y, position.z) == pytest.approx((1, 4, 2)) assert quaternion_euler(translated.orientation) == pytest.approx((0, 0, math.pi / 2)) @@ -241,4 +274,10 @@ def test_local_offset_uses_pose_axes_and_copies_orientation(): @pytest.mark.parametrize("value", [math.nan, math.inf, -math.inf]) def test_local_offset_rejects_nonfinite_displacements(value): with pytest.raises(ValueError, match="Offset components"): - translate_pose_local(Pose(orientation=Quaternion(w=1)), Vector3(x=value)) + translate_pose_local( + Pose( + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + position=Point(x=0.0, y=0.0, z=0.0), + ), + Vector3(x=value, y=0.0, z=0.0), + ) diff --git a/dimos/msgs/test_mesh.py b/dimos/msgs/test_mesh.py new file mode 100644 index 00000000000..cd64c0bcde9 --- /dev/null +++ b/dimos/msgs/test_mesh.py @@ -0,0 +1,37 @@ +# 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. + +"""Keyed mesh chunks retain geometry and source metadata over CDR.""" + +from dimos_cdrgen.registry import decode, encode +import numpy as np + +from dimos.msgs.mesh import triangle_mesh_from_arrays +from dimos.msgs.time import header_now, to_nanoseconds + + +def test_mesh_cdr_preserves_arrays_key_and_source_stamp() -> None: + header = header_now("map") + vertices = np.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]], dtype=np.float32) + message = triangle_mesh_from_arrays( + vertices, np.array([[0, 1, 2]], dtype=np.uint32), key=(-1, 2, 3), header=header + ) + vertices[:] = 99 + restored = decode(encode(message), message.__msgtype__) + assert restored.vertices.tolist() == [0, 1, 2, 3, 4, 5, 6, 7, 8] + assert restored.normals.tolist() == [0] * 9 + assert restored.faces.tolist() == [0, 1, 2] + assert restored.key.tolist() == [-1, 2, 3] + assert restored.header.frame_id == "map" + assert to_nanoseconds(restored.header.stamp) == to_nanoseconds(header.stamp) diff --git a/dimos/msgs/test_msg_construction.py b/dimos/msgs/test_msg_construction.py index 3b49b4661fd..6a265928fcf 100644 --- a/dimos/msgs/test_msg_construction.py +++ b/dimos/msgs/test_msg_construction.py @@ -20,6 +20,7 @@ from copy import deepcopy +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, @@ -29,6 +30,7 @@ PoseWithCovariance, PoseWithCovarianceStamped, Quaternion, + Transform, TransformStamped, Twist, TwistStamped, @@ -41,6 +43,7 @@ from dimos_generated.std_msgs.msg import Header import numpy as np import pytest +from rosbags.serde.errors import SerdeError from rosbags.typesys import Stores, get_typestore from dimos.msgs.geometry import ( @@ -61,14 +64,14 @@ def test_quaternion_explicit_array_conversion(values) -> None: q = quaternion_from_array(values) np.testing.assert_array_equal(quaternion_array(q), values) assert all(type(value) is float for value in (q.x, q.y, q.z, q.w)) - assert Quaternion.decode(q.encode()) == q + assert cdr_decode(cdr_encode(q), Quaternion) == q def test_quaternion_defaults_and_partial_fields() -> None: - assert Quaternion() == Quaternion(x=0, y=0, z=0, w=1) - assert Quaternion(z=1) == Quaternion(x=0, y=0, z=1, w=1) + assert Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) == Quaternion(x=0, y=0, z=0, w=1) + assert Quaternion(z=1, x=0.0, y=0.0, w=1.0) == Quaternion(x=0, y=0, z=1, w=1) # Explicit zero w is retained, even though it is not a spatial rotation. - assert Quaternion.decode(Quaternion(w=0).encode()).w == 0 + assert cdr_decode(cdr_encode(Quaternion(w=0, x=0.0, y=0.0, z=0.0)), Quaternion).w == 0 @pytest.mark.parametrize("values", [[1, 2, 3], [1, 2, 3, 4, 5], (1, 2), np.array([1.0, 2.0, 3.0])]) @@ -92,7 +95,9 @@ def test_pose_explicit_numeric_fields(position, orientation) -> None: def test_pose_explicit_dictionary_and_pair_unpacking() -> None: data = {"position": [1, 2, 3], "orientation": [0, 0, 0, 1]} pair = (data["position"], data["orientation"]) - expected = Pose(position=Point(x=1, y=2, z=3)) + expected = Pose( + position=Point(x=1, y=2, z=3), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) assert ( Pose( position=point_from_array(data["position"]), @@ -104,7 +109,9 @@ def test_pose_explicit_dictionary_and_pair_unpacking() -> None: Pose(position=point_from_array(pair[0]), orientation=quaternion_from_array(pair[1])) == expected ) - assert Pose(position=Point(x=5)).position == Point(x=5, y=0, z=0) + assert Pose( + position=Point(x=5, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ).position == Point(x=5, y=0, z=0) @pytest.mark.parametrize("values", [[1, 2, 3], (1, 2, 3), np.array([1.0, 2.0, 3.0])]) @@ -112,98 +119,165 @@ def test_twist_explicit_arrays(values) -> None: twist = Twist(linear=vector_from_array(values), angular=vector_from_array([4, 5, 6])) np.testing.assert_array_equal(vector_array(twist.linear), [1, 2, 3]) np.testing.assert_array_equal(vector_array(twist.angular), [4, 5, 6]) - assert Twist(linear=twist.linear).angular == Vector3() - assert Twist(angular=twist.angular).linear == Vector3() + assert Twist(linear=twist.linear, angular=Vector3(x=0.0, y=0.0, z=0.0)).angular == Vector3( + x=0.0, y=0.0, z=0.0 + ) + assert Twist(angular=twist.angular, linear=Vector3(x=0.0, y=0.0, z=0.0)).linear == Vector3( + x=0.0, y=0.0, z=0.0 + ) def test_twist_euler_conversion_is_explicit() -> None: twist = Twist( - linear=Vector3(x=1, y=2, z=3), angular=vector_from_array(quaternion_euler(Quaternion())) + linear=Vector3(x=1, y=2, z=3), + angular=vector_from_array(quaternion_euler(Quaternion(x=0.0, y=0.0, z=0.0, w=1.0))), ) assert twist.linear == Vector3(x=1, y=2, z=3) - assert twist.angular == Vector3() + assert twist.angular == Vector3(x=0.0, y=0.0, z=0.0) @pytest.mark.parametrize( "message", [ Quaternion(x=0.1, y=0.2, z=0.3, w=0.4), - Pose(position=Point(x=1, y=2, z=3)), + Pose(position=Point(x=1, y=2, z=3), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)), Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=4, y=5, z=6)), - PoseWithCovariance(pose=Pose(position=Point(x=1, y=2, z=3)), covariance=list(range(36))), - TwistWithCovariance(twist=Twist(linear=Vector3(x=1, y=2, z=3)), covariance=list(range(36))), - JointState(header=Header(stamp=Time(sec=5), frame_id="f"), name=["a"], position=[1.0]), - Joy(header=Header(stamp=Time(sec=5), frame_id="f"), axes=[1.0], buttons=[1]), + PoseWithCovariance( + pose=Pose( + position=Point(x=1, y=2, z=3), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + covariance=np.asarray(list(range(36)), dtype=np.float64), + ), + TwistWithCovariance( + twist=Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.0, y=0.0, z=0.0)), + covariance=np.asarray(list(range(36)), dtype=np.float64), + ), + JointState( + header=Header(stamp=Time(sec=5, nanosec=0), frame_id="f"), + name=["a"], + position=np.array([1.0], dtype=np.float64), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + Joy( + header=Header(stamp=Time(sec=5, nanosec=0), frame_id="f"), + axes=np.array([1.0], dtype=np.float32), + buttons=np.array([1], dtype=np.int32), + ), ], ) def test_copy_and_cdr_roundtrip_preserve_fields(message) -> None: copied = deepcopy(message) - assert copied == message + assert cdr_encode(copied) == cdr_encode(message) assert copied is not message - assert type(message).decode(message.encode()) == message + assert cdr_encode(cdr_decode(cdr_encode(message), type(message))) == cdr_encode(message) independent = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr( - message.encode(), message.msg_name + cdr_encode(message), message.__msgtype__ ) - assert independent.__msgtype__ == message.msg_name + assert independent.__msgtype__ == message.__msgtype__ def test_copy_mutation_is_independent() -> None: - pose = Pose(position=Point(x=1, y=2, z=3)) + pose = Pose(position=Point(x=1, y=2, z=3), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)) pose_copy = deepcopy(pose) pose_copy.position.x = 99 assert pose.position.x == 1 - twist = Twist(linear=Vector3(x=1, y=2, z=3)) + twist = Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.0, y=0.0, z=0.0)) twist_copy = deepcopy(twist) twist_copy.linear.x = 99 assert twist.linear.x == 1 - covariance = PoseWithCovariance(covariance=list(range(36))) + covariance = PoseWithCovariance( + covariance=np.asarray(list(range(36)), dtype=np.float64), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + ) covariance_copy = deepcopy(covariance) covariance_copy.covariance[0] = 99 assert covariance.covariance[0] == 0 - joints = JointState(name=["a"], position=[1.0]) + joints = JointState( + name=["a"], + position=np.array([1.0], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) joints_copy = deepcopy(joints) joints_copy.name[0] = "b" assert joints.name == ["a"] - joy = Joy(axes=[1.0], buttons=[1]) + joy = Joy( + axes=np.array([1.0], dtype=np.float32), + buttons=np.array([1], dtype=np.int32), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) joy_copy = deepcopy(joy) joy_copy.axes[0] = 2 assert joy.axes == [1.0] @pytest.mark.parametrize("message_type", [PoseWithCovariance, TwistWithCovariance]) -def test_covariance_defaults_layout_and_length(message_type) -> None: - assert list(message_type().covariance) == [0.0] * 36 +def test_covariance_layout_and_length_at_native_codec_boundary(message_type) -> None: + payload = ( + {"pose": Pose(position=Point(0.0, 0.0, 0.0), orientation=Quaternion(0.0, 0.0, 0.0, 1.0))} + if message_type is PoseWithCovariance + else {"twist": Twist(linear=Vector3(0.0, 0.0, 0.0), angular=Vector3(0.0, 0.0, 0.0))} + ) covariance = np.arange(36.0).reshape(6, 6) - message = message_type(covariance=covariance.flatten().tolist()) - np.testing.assert_array_equal(np.asarray(message.covariance).reshape(6, 6), covariance) + message = message_type(**payload, covariance=covariance.flatten()) + restored = cdr_decode(cdr_encode(message), message_type) + np.testing.assert_array_equal(restored.covariance.reshape(6, 6), covariance) for invalid in ([1, 2, 3], [0.0] * 35, [0.0] * 37): - with pytest.raises(RuntimeError, match="Unable to cast"): - message_type(covariance=invalid) - - -def test_sensor_explicit_fields_and_defaults() -> None: - joints = JointState( - header=Header(stamp=Time(sec=5), frame_id="f"), - name=["a"], - position=[1.0], - velocity=[2.0], - effort=[3.0], - ) - assert (joints.name, joints.position, joints.velocity, joints.effort) == ( - ["a"], - [1.0], - [2.0], - [3.0], - ) - assert joints.header == Header(stamp=Time(sec=5), frame_id="f") - assert JointState().name == [] - assert JointState().velocity == [] - assert JointState().effort == [] - pair = ([1.0, 2.0], [1, 0]) - joy = Joy(axes=pair[0], buttons=pair[1]) - assert (joy.axes, joy.buttons) == pair - assert (Joy().axes, Joy().buttons) == ([], []) - assert Joy().header.frame_id == "" + with pytest.raises(SerdeError, match="array length"): + cdr_encode(message_type(**payload, covariance=np.array(invalid, dtype=np.float64))) + + +def test_sensor_explicit_fields_and_empty_arrays() -> None: + header = Header(stamp=Time(sec=5, nanosec=0), frame_id="f") + joints = JointState(header, ["a"], np.array([1.0]), np.array([2.0]), np.array([3.0])) + restored = cdr_decode(cdr_encode(joints), JointState) + assert restored.header == header and restored.name == ["a"] + for actual, expected in [ + (restored.position, [1.0]), + (restored.velocity, [2.0]), + (restored.effort, [3.0]), + ]: + np.testing.assert_array_equal(actual, expected) + joy = Joy(header, np.array([1.0, 2.0], dtype=np.float32), np.array([1, 0], dtype=np.int32)) + restored_joy = cdr_decode(cdr_encode(joy), Joy) + np.testing.assert_array_equal(restored_joy.axes, [1.0, 2.0]) + np.testing.assert_array_equal(restored_joy.buttons, [1, 0]) + empty = Joy(header, np.array([], dtype=np.float32), np.array([], dtype=np.int32)) + restored_empty = cdr_decode(cdr_encode(empty), Joy) + assert restored_empty.axes.size == restored_empty.buttons.size == 0 + assert restored_empty.header == header + + +def stamped_message(message_type, header): + pose = Pose(Point(0.0, 0.0, 0.0), Quaternion(0.0, 0.0, 0.0, 1.0)) + twist = Twist(Vector3(0.0, 0.0, 0.0), Vector3(0.0, 0.0, 0.0)) + pose_covariance = PoseWithCovariance(pose, np.zeros(36, dtype=np.float64)) + twist_covariance = TwistWithCovariance(twist, np.zeros(36, dtype=np.float64)) + fields = { + PoseStamped: {"pose": pose}, + TwistStamped: {"twist": twist}, + PoseWithCovarianceStamped: {"pose": pose_covariance}, + TwistWithCovarianceStamped: {"twist": twist_covariance}, + JointState: { + "name": [], + "position": np.array([], dtype=np.float64), + "velocity": np.array([], dtype=np.float64), + "effort": np.array([], dtype=np.float64), + }, + Joy: {"axes": np.array([], dtype=np.float32), "buttons": np.array([], dtype=np.int32)}, + Path: {"poses": []}, + PointStamped: {"point": pose.position}, + TransformStamped: { + "child_frame_id": "", + "transform": Transform(Vector3(0.0, 0.0, 0.0), pose.orientation), + }, + Odometry: {"child_frame_id": "", "pose": pose_covariance, "twist": twist_covariance}, + } + return message_type(header=header, **fields[message_type]) STAMPED_TYPES = [ @@ -223,35 +297,55 @@ def test_sensor_explicit_fields_and_defaults() -> None: @pytest.mark.parametrize("message_type", STAMPED_TYPES) @pytest.mark.parametrize( "stamp", - [Time(), Time(sec=5), Time(sec=-1, nanosec=999999999), Time(sec=1700000000, nanosec=123456789)], + [ + Time(sec=0, nanosec=0), + Time(sec=5, nanosec=0), + Time(sec=-1, nanosec=999999999), + Time(sec=1700000000, nanosec=123456789), + ], ) def test_exact_stamps_survive_cdr(message_type, stamp) -> None: - message = message_type(header=Header(stamp=stamp, frame_id="odom")) - decoded = message_type.decode(message.encode()) + message = stamped_message(message_type, Header(stamp=stamp, frame_id="odom")) + decoded = cdr_decode(cdr_encode(message), message_type) assert decoded.header == message.header assert to_nanoseconds(decoded.header.stamp) == stamp.sec * 1000000000 + stamp.nanosec @pytest.mark.parametrize("message_type", STAMPED_TYPES) -def test_omitted_stamp_is_zero_without_wall_clock(message_type) -> None: - message = message_type() - assert message.header.stamp == Time() - assert message_type.decode(message.encode()).header.stamp == Time() +def test_native_constructor_requires_fields_and_explicit_zero_never_samples_wall_clock( + message_type, +) -> None: + with pytest.raises(TypeError): + message_type() + message = stamped_message(message_type, Header(stamp=Time(0, 0), frame_id="")) + assert message.header.stamp == Time(sec=0, nanosec=0) + assert cdr_decode(cdr_encode(message), message_type).header.stamp == Time(sec=0, nanosec=0) def test_stamped_payloads_are_nested_values() -> None: - pose = Pose(position=Point(x=1, y=2, z=3), orientation=Quaternion(z=0.1, w=1)) - stamped = PoseStamped(header=Header(stamp=Time(sec=5), frame_id="odom"), pose=pose) + pose = Pose(position=Point(x=1, y=2, z=3), orientation=Quaternion(z=0.1, w=1, x=0.0, y=0.0)) + stamped = PoseStamped(header=Header(stamp=Time(sec=5, nanosec=0), frame_id="odom"), pose=pose) assert not isinstance(stamped, Pose) assert stamped.pose == pose assert deepcopy(stamped.pose) == pose twist = Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=4, y=5, z=6)) - assert TwistStamped(twist=twist).twist == twist - pwc = PoseWithCovarianceStamped(pose=PoseWithCovariance(pose=pose, covariance=list(range(36)))) + assert ( + TwistStamped(twist=twist, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")).twist + == twist + ) + pwc = PoseWithCovarianceStamped( + pose=PoseWithCovariance( + pose=pose, covariance=np.asarray(list(range(36)), dtype=np.float64) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) assert pwc.pose.pose == pose assert list(pwc.pose.covariance) == list(range(36)) twc = TwistWithCovarianceStamped( - twist=TwistWithCovariance(twist=twist, covariance=list(range(36))) + twist=TwistWithCovariance( + twist=twist, covariance=np.asarray(list(range(36)), dtype=np.float64) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) assert twc.twist.twist == twist assert list(twc.twist.covariance) == list(range(36)) @@ -274,10 +368,10 @@ def test_unknown_fields_rejected(message_type) -> None: (Pose, {"orientation": [0, 0, 0, 1]}), (Pose, {"position": None}), (Twist, {"linear": [1, 2, 3]}), - (Twist, {"angular": Quaternion()}), + (Twist, {"angular": Quaternion(x=0.0, y=0.0, z=0.0, w=1.0)}), (PoseStamped, {"ts": 5}), (PoseStamped, {"frame_id": "odom"}), - (PoseStamped, {"position": Point(x=1)}), + (PoseStamped, {"position": Point(x=1, y=0.0, z=0.0)}), ], ) def test_retired_polymorphic_fields_rejected(message_type, kwargs) -> None: diff --git a/dimos/msgs/test_occupancy.py b/dimos/msgs/test_occupancy.py index 8773aa1a0f9..dabeecee18a 100644 --- a/dimos/msgs/test_occupancy.py +++ b/dimos/msgs/test_occupancy.py @@ -14,7 +14,9 @@ import math +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid from dimos_generated.std_msgs.msg import Header import numpy as np @@ -27,9 +29,20 @@ def test_occupancy_cells_keep_ros_row_order_and_do_not_alias_message(): message = OccupancyGrid( - info=MapMetaData(width=3, height=2, resolution=0.1), data=[-1, 0, 100, 25, 50, 75] + info=MapMetaData( + width=3, + height=2, + resolution=0.1, + map_load_time=Time(sec=0, nanosec=0), + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([-1, 0, 100, 25, 50, 75], dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) - decoded = OccupancyGrid.decode(message.encode()) + decoded = cdr_decode(cdr_encode(message), OccupancyGrid) values = occupancy_view(decoded) np.testing.assert_array_equal(values, [[-1, 0, 100], [25, 50, 75]]) with pytest.raises(ValueError, match="read-only"): @@ -40,20 +53,61 @@ def test_occupancy_cells_keep_ros_row_order_and_do_not_alias_message(): def test_occupancy_rejects_dimensions_that_do_not_match_payload(): - message = OccupancyGrid(info=MapMetaData(width=2, height=2, resolution=1.0), data=[0]) + message = OccupancyGrid( + info=MapMetaData( + width=2, + height=2, + resolution=1.0, + map_load_time=Time(sec=0, nanosec=0), + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([0], dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) with pytest.raises(ValueError, match="data length"): occupancy_view(message) @pytest.mark.parametrize("resolution", [0.0, -0.1, math.nan, math.inf]) def test_occupancy_rejects_invalid_resolution_for_nonempty_grid(resolution): - message = OccupancyGrid(info=MapMetaData(width=1, height=1, resolution=resolution), data=[0]) + message = OccupancyGrid( + info=MapMetaData( + width=1, + height=1, + resolution=resolution, + map_load_time=Time(sec=0, nanosec=0), + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([0], dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) with pytest.raises(ValueError, match="resolution"): occupancy_extent(message) def test_empty_default_occupancy_grid_has_no_cells(): - assert occupancy_view(OccupancyGrid()).shape == (0, 0) + assert occupancy_view( + OccupancyGrid( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + info=MapMetaData( + map_load_time=Time(sec=0, nanosec=0), + resolution=0.0, + width=0, + height=0, + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([], dtype=np.int8), + ) + ).shape == (0, 0) @pytest.mark.parametrize("suffix", [".npy", ".png"]) @@ -65,7 +119,10 @@ def test_static_occupancy_loader_preserves_cells_header_and_identity(tmp_path, s else: Image.fromarray(cells.astype(np.uint8)).save(path) header = Header(frame_id="map", stamp=Time(sec=123, nanosec=456)) - output = OccupancyGrid.decode(occupancy_from_file(path, header=header, resolution=0.2).encode()) + output = cdr_decode( + cdr_encode(occupancy_from_file(path, header=header, resolution=0.2)), + OccupancyGrid, + ) np.testing.assert_array_equal(occupancy_view(output), cells) assert output.header == header assert output.info.map_load_time == header.stamp @@ -81,7 +138,7 @@ def test_static_occupancy_loader_rejects_invalid_cells(tmp_path, cells): path = tmp_path / "grid.npy" np.save(path, cells) with pytest.raises(ValueError): - occupancy_from_file(path, header=Header()) + occupancy_from_file(path, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) def test_gradient() -> None: @@ -91,9 +148,18 @@ def test_gradient() -> None: data[4:6, 4:6] = 100 # 2x2 obstacle in center grid = OccupancyGrid( - header=Header(frame_id="world"), - info=MapMetaData(width=10, height=10, resolution=0.1), - data=data.ravel(), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + info=MapMetaData( + width=10, + height=10, + resolution=0.1, + map_load_time=Time(sec=0, nanosec=0), + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.asarray(data.ravel(), dtype=np.int8), ) # 0.1m per cell # Convert to gradient @@ -130,7 +196,9 @@ def test_gradient() -> None: data_with_unknown[8:10, 8:10] = -1 # Add unknown area (far from obstacle) grid_with_unknown = OccupancyGrid( - header=grid.header, info=grid.info, data=data_with_unknown.ravel() + header=grid.header, + info=grid.info, + data=np.asarray(data_with_unknown.ravel(), dtype=np.int8), ) gradient_with_unknown = gradient(grid_with_unknown, max_distance=1.0) # 1m max distance diff --git a/dimos/msgs/test_pointcloud.py b/dimos/msgs/test_pointcloud.py index e38565a6a2a..a349c0800f7 100644 --- a/dimos/msgs/test_pointcloud.py +++ b/dimos/msgs/test_pointcloud.py @@ -18,6 +18,7 @@ import gc import struct +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 from dimos_generated.sensor_msgs.msg import PointCloud2, PointField @@ -33,6 +34,7 @@ pointcloud_from_xyz, pointcloud_from_xyz_rgb, pointcloud_rgb, + pointcloud_stamps, pointcloud_view, pointcloud_xyz, select_points, @@ -65,13 +67,15 @@ def padded_cloud(request): point_step=24, row_step=64, is_bigendian=request.param, - data=bytes(payload), + data=np.frombuffer(bytes(payload), dtype=np.uint8), fields=[ PointField(name="z", offset=0, datatype=8, count=1), PointField(name="x", offset=8, datatype=7, count=1), PointField(name="tags", offset=12, datatype=4, count=2), PointField(name="y", offset=16, datatype=7, count=1), ], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + is_dense=False, ) @@ -93,8 +97,11 @@ def test_view_retains_storage_after_message_is_deleted(): width=1, point_step=4, row_step=4, - fields=[PointField(name="x", datatype=7, count=1)], - data=struct.pack(" None: first = TransformStamped( header=Header(stamp=time_from_seconds(100), frame_id="world"), - transform=Transform(translation=Vector3(x=1, y=2, z=3), rotation=Quaternion(w=1)), + transform=Transform( + translation=Vector3(x=1, y=2, z=3), rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), + child_frame_id="", ) second = TransformStamped( header=Header(stamp=time_from_seconds(101), frame_id="map"), transform=Transform( - translation=Vector3(x=4, y=5, z=6), rotation=Quaternion(z=0.707, w=0.707) + translation=Vector3(x=4, y=5, z=6), rotation=Quaternion(z=0.707, w=0.707, x=0.0, y=0.0) ), + child_frame_id="", ) message = TFMessage(transforms=[first, second]) assert len(message.transforms) == 2 - assert [item.encode() for item in message.transforms] == [first.encode(), second.encode()] + assert [cdr_encode(item) for item in message.transforms] == [ + cdr_encode(first), + cdr_encode(second), + ] def test_tfmessage_empty() -> None: - message = TFMessage() + message = TFMessage(transforms=[]) assert len(message.transforms) == 0 assert list(message.transforms) == [] def test_tfmessage_add_transform() -> None: - message = TFMessage() + message = TFMessage(transforms=[]) transform = TransformStamped( header=Header(stamp=time_from_seconds(200), frame_id="base"), - transform=Transform(translation=Vector3(x=1, y=2, z=3)), + transform=Transform( + translation=Vector3(x=1, y=2, z=3), rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + child_frame_id="", ) message.transforms.append(transform) assert len(message.transforms) == 1 - assert message.transforms[0].encode() == transform.encode() + assert cdr_encode(message.transforms[0]) == cdr_encode(transform) def test_tfmessage_tree() -> None: - message = TFMessage() + message = TFMessage(transforms=[]) edges = [ ("world", "robot", (1, 2, 0.5)), ("robot", "camera", (0.1, 0, 1)), @@ -66,7 +77,10 @@ def test_tfmessage_tree() -> None: TransformStamped( header=Header(stamp=time_from_seconds(200), frame_id=parent), child_frame_id=child, - transform=Transform(translation=Vector3(x=x, y=y, z=z)), + transform=Transform( + translation=Vector3(x=x, y=y, z=z), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) ) assert len(message.transforms) == 4 @@ -76,7 +90,7 @@ def test_tfmessage_tree() -> None: t.child_frame_id, (t.transform.translation.x, t.transform.translation.y, t.transform.translation.z), ) - for t in TFMessage.decode(message.encode()).transforms + for t in cdr_decode(cdr_encode(message), TFMessage).transforms ] == edges @@ -84,17 +98,21 @@ def test_tfmessage_cdr_encode_decode() -> None: first = TransformStamped( header=Header(stamp=time_from_seconds(123.456), frame_id="world"), child_frame_id="robot", - transform=Transform(translation=Vector3(x=1, y=2, z=3), rotation=Quaternion(w=1)), + transform=Transform( + translation=Vector3(x=1, y=2, z=3), rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), ) second = TransformStamped( header=Header(stamp=time_from_seconds(124.567), frame_id="robot"), child_frame_id="target", transform=Transform( - translation=Vector3(x=4, y=5, z=6), rotation=Quaternion(z=0.707, w=0.707) + translation=Vector3(x=4, y=5, z=6), rotation=Quaternion(z=0.707, w=0.707, x=0.0, y=0.0) ), ) message = TFMessage(transforms=[first, second]) - decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr(message.encode(), TFMessage.msg_name) + decoded = get_typestore(Stores.ROS2_JAZZY).deserialize_cdr( + cdr_encode(message), TFMessage.__msgtype__ + ) assert len(decoded.transforms) == 2 first_result, second_result = decoded.transforms assert first_result.header.frame_id == "world" diff --git a/dimos/msgs/tf2_msgs/test_TFMessage_lcmpub.py b/dimos/msgs/tf2_msgs/test_TFMessage_lcmpub.py index c20cc55b291..7a2eb014ae7 100644 --- a/dimos/msgs/tf2_msgs/test_TFMessage_lcmpub.py +++ b/dimos/msgs/tf2_msgs/test_TFMessage_lcmpub.py @@ -49,7 +49,7 @@ def test_publish_transforms(lcm: LCM) -> None: world_to_base = TransformStamped( transform=Transform( translation=Vector3(x=4.0, y=3.0, z=0.0), - rotation=Quaternion(z=0.382683, w=0.923880), # 45 degrees around Z + rotation=Quaternion(z=0.382683, w=0.923880, x=0.0, y=0.0), # 45 degrees around Z ), header=Header(frame_id="world", stamp=time_from_seconds(current_time)), child_frame_id="base_link", @@ -59,7 +59,7 @@ def test_publish_transforms(lcm: LCM) -> None: base_to_arm = TransformStamped( transform=Transform( translation=Vector3(x=0.2, y=0.0, z=1.5), - rotation=Quaternion(y=0.258819, w=0.965926), # 30 degrees around Y + rotation=Quaternion(y=0.258819, w=0.965926, x=0.0, z=0.0), # 30 degrees around Y ), header=Header(frame_id="base_link", stamp=time_from_seconds(current_time)), child_frame_id="arm_link", @@ -69,7 +69,7 @@ def test_publish_transforms(lcm: LCM) -> None: arm_to_gripper = TransformStamped( transform=Transform( translation=Vector3(x=0.5, y=0.0, z=0.0), - rotation=Quaternion(w=1.0), # No rotation + rotation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), # No rotation ), header=Header(frame_id="arm_link", stamp=time_from_seconds(current_time)), child_frame_id="gripper_link", diff --git a/dimos/msgs/time.py b/dimos/msgs/time.py index 020fcd2f2ab..82f4ccf537c 100644 --- a/dimos/msgs/time.py +++ b/dimos/msgs/time.py @@ -16,8 +16,10 @@ import math import time +from typing import Any from dimos_generated.builtin_interfaces.msg import Duration, Time +from dimos_generated.dimos_msgs.msg import RegionLineSegments3D, RegionPointCloud2 from dimos_generated.std_msgs.msg import Header NANOSECONDS_PER_SECOND = 1_000_000_000 @@ -68,3 +70,13 @@ def duration_from_seconds(value: float) -> Duration: def header_now(frame_id: str = "") -> Header: """Construct a wall-clock header; generated defaults themselves remain zero.""" return Header(stamp=time_from_nanoseconds(time.time_ns()), frame_id=frame_id) + + +def message_header(message: Any) -> Header | None: + """Return source metadata, including the payload owned by a region envelope.""" + if isinstance(message, RegionPointCloud2): + return message.cloud.header + if isinstance(message, RegionLineSegments3D): + return message.lines.header + header = getattr(message, "header", None) + return header if isinstance(header, Header) else None diff --git a/dimos/msgs/vision_msgs/test_detection3d_rerun.py b/dimos/msgs/vision_msgs/test_detection3d_rerun.py index df75cfe4936..aaf8ba539ed 100644 --- a/dimos/msgs/vision_msgs/test_detection3d_rerun.py +++ b/dimos/msgs/vision_msgs/test_detection3d_rerun.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion, Vector3 +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseWithCovariance, Quaternion, Vector3 from dimos_generated.std_msgs.msg import Header from dimos_generated.vision_msgs.msg import ( BoundingBox3D, @@ -21,6 +22,7 @@ ObjectHypothesis, ObjectHypothesisWithPose, ) +import numpy as np import pytest import rerun as rr @@ -35,7 +37,18 @@ def _detection3d( marker_id: str = "7", class_id: str = "DICT_APRILTAG_36h11:7", ) -> Detection3D: - det = Detection3D() + det = Detection3D( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + results=[], + bbox=BoundingBox3D( + center=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + size=Vector3(x=0.0, y=0.0, z=0.0), + ), + id="", + ) det.header = Header(stamp=time_from_seconds(ts), frame_id=frame_id) det.id = marker_id det.results = [ @@ -43,15 +56,22 @@ def _detection3d( hypothesis=ObjectHypothesis( class_id=class_id, score=1.0, - ) + ), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), ) ] det.bbox = BoundingBox3D( center=Pose( position=Point(x=1.0, y=2.0, z=3.0), - orientation=Quaternion(z=0.70710678, w=0.70710678), + orientation=Quaternion(z=0.70710678, w=0.70710678, x=0.0, y=0.0), ), - size=Vector3(x=0.2, y=0.4), + size=Vector3(x=0.2, y=0.4, z=0.0), ) return det diff --git a/dimos/navigation/experimental/dannav/geometry/path_distancer.py b/dimos/navigation/experimental/dannav/geometry/path_distancer.py index cfd0975d854..d976707deed 100644 --- a/dimos/navigation/experimental/dannav/geometry/path_distancer.py +++ b/dimos/navigation/experimental/dannav/geometry/path_distancer.py @@ -16,11 +16,10 @@ from dataclasses import dataclass +from dimos_generated.nav_msgs.msg import Path import numpy as np from numpy.typing import NDArray -from dimos.msgs.nav_msgs.Path import Path - @dataclass(frozen=True) class PolylineProjection: @@ -49,7 +48,7 @@ class PathDistancer: """Arc-length projection onto a fixed ``Path`` polyline.""" def __init__(self, path: Path) -> None: - self._path = np.array([[p.position.x, p.position.y] for p in path.poses]) + self._path = np.array([[p.pose.position.x, p.pose.position.y] for p in path.poses]) def project(self, pos: NDArray[np.float64]) -> PolylineProjection: return project_to_polyline(float(pos[0]), float(pos[1]), self._path) diff --git a/dimos/navigation/experimental/dannav/holonomic_tc/holonomic_path_controller.py b/dimos/navigation/experimental/dannav/holonomic_tc/holonomic_path_controller.py index 9531f9fc871..d28c71122de 100644 --- a/dimos/navigation/experimental/dannav/holonomic_tc/holonomic_path_controller.py +++ b/dimos/navigation/experimental/dannav/holonomic_tc/holonomic_path_controller.py @@ -23,11 +23,11 @@ import math +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist, Vector3 +from dimos_generated.nav_msgs.msg import Path import numpy as np -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import quaternion_euler from dimos.navigation.experimental.dannav.geometry.path_distancer import project_to_polyline from dimos.navigation.experimental.dannav.geometry.path_speed_profile import profile_speed from dimos.navigation.experimental.dannav.holonomic_tc.holonomic_tracking_controller import ( @@ -70,7 +70,7 @@ def set_path(self, path: Path | None) -> None: self._path = None if path is None or len(path.poses) < 2: return - xy = np.array([[p.position.x, p.position.y] for p in path.poses]) + xy = np.array([[p.pose.position.x, p.pose.position.y] for p in path.poses]) xy = xy[np.concatenate([[True], np.hypot(*np.diff(xy, axis=0).T) > 1e-6])] if len(xy) < 2: return @@ -95,14 +95,14 @@ def step(self, odom: PoseStamped) -> Twist | None: """The next command; a zero twist on arrival, which also drops the path.""" if self._path is None: return None - x, y = float(odom.position.x), float(odom.position.y) - yaw = float(odom.orientation.euler[2]) + x, y = float(odom.pose.position.x), float(odom.pose.position.y) + yaw = float(quaternion_euler(odom.pose.orientation)[2]) foot = project_to_polyline(x, y, self._xy) s = foot.s_along_path_m end = self._xy[-1] if math.hypot(end[0] - x, end[1] - y) < self._goal_tolerance: self._path = None - return Twist() + return Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) v = float(np.interp(s, self._s, self._speed)) heading = float(np.interp(s, self._s, self._heading)) reference = PlanarState( diff --git a/dimos/navigation/experimental/dannav/holonomic_tc/holonomic_tracking_controller.py b/dimos/navigation/experimental/dannav/holonomic_tc/holonomic_tracking_controller.py index 5c2f57ace79..0c9d53e8e82 100644 --- a/dimos/navigation/experimental/dannav/holonomic_tc/holonomic_tracking_controller.py +++ b/dimos/navigation/experimental/dannav/holonomic_tc/holonomic_tracking_controller.py @@ -25,8 +25,8 @@ from dataclasses import dataclass import math -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos_generated.geometry_msgs.msg import Twist, Vector3 + from dimos.utils.trigonometry import angle_diff diff --git a/dimos/navigation/experimental/dannav/holonomic_tc/module.py b/dimos/navigation/experimental/dannav/holonomic_tc/module.py index e22e21f7da2..e0ce2c5b4bb 100644 --- a/dimos/navigation/experimental/dannav/holonomic_tc/module.py +++ b/dimos/navigation/experimental/dannav/holonomic_tc/module.py @@ -24,7 +24,9 @@ from threading import Event, RLock, Thread from typing import Any -from dimos_lcm.std_msgs import Bool +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist, Vector3 +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Bool from pydantic import Field from reactivex.disposable import Disposable @@ -32,9 +34,6 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Path import Path from dimos.navigation.base import NavigationState from dimos.navigation.experimental.dannav.holonomic_tc.holonomic_path_controller import ( HolonomicPathController, @@ -107,7 +106,9 @@ def stop(self) -> None: self._stop_event.set() if self._thread is not None: self._thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) - self.nav_cmd_vel.publish(Twist()) + self.nav_cmd_vel.publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) super().stop() def _on_odom(self, msg: PoseStamped) -> None: @@ -125,7 +126,9 @@ def _follow(self, path: Path | None) -> None: with self._lock: self._controller.set_path(path) if not self._controller.active: - self.nav_cmd_vel.publish(Twist()) + self.nav_cmd_vel.publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) def _control_loop(self) -> None: period = 1.0 / self.config.control_frequency @@ -142,7 +145,7 @@ def _control_loop(self) -> None: if cmd is not None: self.nav_cmd_vel.publish(cmd) if not self._controller.active: - self.goal_reached.publish(Bool(True)) + self.goal_reached.publish(Bool(data=True)) logger.info("Goal reached") @rpc diff --git a/dimos/navigation/experimental/dannav/holonomic_tc/test_dan_holonomic_tc.py b/dimos/navigation/experimental/dannav/holonomic_tc/test_dan_holonomic_tc.py index bdf0485cf44..31b600bb71d 100644 --- a/dimos/navigation/experimental/dannav/holonomic_tc/test_dan_holonomic_tc.py +++ b/dimos/navigation/experimental/dannav/holonomic_tc/test_dan_holonomic_tc.py @@ -23,6 +23,8 @@ import time from typing import Any, Literal +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Twist from dimos_generated.nav_msgs.msg import Path from dimos_generated.std_msgs.msg import Bool, Header @@ -44,7 +46,7 @@ def __init__(self) -> None: def broadcast(self, _selfstream: Any, value: Any) -> None: for callback in list(self._subscribers): - callback(type(value).decode(value.encode())) + callback(cdr_decode(cdr_encode(value), type(value))) def subscribe( self, callback: Callable[[Any], Any], _selfstream: Stream[Any] | None = None @@ -70,13 +72,13 @@ class _Captured: def _yaw_quaternion(yaw_rad: float) -> Quaternion: - return Quaternion(z=math.sin(yaw_rad / 2.0), w=math.cos(yaw_rad / 2.0)) + return Quaternion(z=math.sin(yaw_rad / 2.0), w=math.cos(yaw_rad / 2.0), x=0.0, y=0.0) def _odom(x: float, y: float, yaw_rad: float, *, ts: float = 1.0) -> PoseStamped: return PoseStamped( header=Header(stamp=time_from_seconds(ts), frame_id="map"), - pose=Pose(position=Point(x=x, y=y), orientation=_yaw_quaternion(yaw_rad)), + pose=Pose(position=Point(x=x, y=y, z=0.0), orientation=_yaw_quaternion(yaw_rad)), ) @@ -90,7 +92,7 @@ def _path_from_points(points: list[tuple[float, float]]) -> Path: prev_point = points[index - 1] yaw = math.atan2(point[1] - prev_point[1], point[0] - prev_point[0]) poses.append(_odom(point[0], point[1], yaw)) - return Path(header=Header(frame_id="map"), poses=poses) + return Path(header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), poses=poses) def _is_zero_twist(cmd: Twist) -> bool: @@ -125,7 +127,9 @@ def feed_path(self, path: Path) -> None: self.module.path.transport.broadcast(None, path) def feed_empty_path(self) -> None: - self.module.path.transport.broadcast(None, Path(header=Header(frame_id="map"), poses=[])) + self.module.path.transport.broadcast( + None, Path(header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), poses=[]) + ) def feed_stop(self, value: bool = True) -> None: self.module.stop_movement.transport.broadcast(None, Bool(data=value)) diff --git a/dimos/navigation/experimental/dannav/local_planner/module.py b/dimos/navigation/experimental/dannav/local_planner/module.py index e2a611fa6ff..ace527e95f9 100644 --- a/dimos/navigation/experimental/dannav/local_planner/module.py +++ b/dimos/navigation/experimental/dannav/local_planner/module.py @@ -36,6 +36,8 @@ import math from typing import Any +from dimos_generated.geometry_msgs.msg import PointStamped, PoseStamped +from dimos_generated.nav_msgs.msg import Path import numpy as np from numpy.typing import NDArray from reactivex.disposable import Disposable @@ -44,10 +46,6 @@ from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out from dimos.mapping.occupancy.path_resampling import smooth_resample_path -from dimos.msgs.geometry_msgs.PointStamped import PointStamped -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.Path import Path from dimos.navigation.experimental.dannav.geometry.path_distancer import PathDistancer from dimos.utils.logging_config import setup_logger @@ -98,7 +96,7 @@ def __init__(self, config: DanLocalPlannerConfig) -> None: self._anchor_progress_m: float = 0.0 def on_odom(self, odom: PoseStamped) -> None: - self._robot_xy = np.array([odom.position.x, odom.position.y], dtype=float) + self._robot_xy = np.array([odom.pose.position.x, odom.pose.position.y], dtype=float) def on_goal(self, goal: PointStamped) -> None: """Arm or disarm the fresh-click commit. @@ -108,11 +106,11 @@ def on_goal(self, goal: PointStamped) -> None: drops the committed path). ``DanLocalPlanner`` does not touch ``stop_movement``. """ - if math.isnan(goal.x) or math.isnan(goal.y): + if math.isnan(goal.point.x) or math.isnan(goal.point.y): self._armed_goal = None self._drop_committed() return - self._armed_goal = np.array([goal.x, goal.y], dtype=float) + self._armed_goal = np.array([goal.point.x, goal.point.y], dtype=float) def on_planner_path(self, path: Path) -> Path | None: """Return the path to forward, or ``None`` to suppress it. @@ -171,7 +169,7 @@ def _smooth(self, path: Path) -> Path: if self._cfg.resample_spacing_m <= 0.0 or len(path.poses) == 0: return path last = path.poses[-1] - goal_pose = Pose(last.position, last.orientation) + goal_pose = last.pose return smooth_resample_path( path, goal_pose, self._cfg.resample_spacing_m, self._cfg.smoothing_window ) @@ -201,7 +199,10 @@ def _progress_on_committed(self) -> float: def _ends_near(path: Path, goal_xy: NDArray[np.float64], tolerance_m: float) -> bool: """True when ``path``'s last pose is within ``tolerance_m`` of ``goal_xy``.""" end = path.poses[-1] - return float(math.hypot(end.x - goal_xy[0], end.y - goal_xy[1])) <= tolerance_m + return ( + float(math.hypot(end.pose.position.x - goal_xy[0], end.pose.position.y - goal_xy[1])) + <= tolerance_m + ) class DanLocalPlanner(Module): diff --git a/dimos/navigation/experimental/dannav/local_planner/test_dan_local_planner.py b/dimos/navigation/experimental/dannav/local_planner/test_dan_local_planner.py index b5f65b0f71c..fb9fbe33a2e 100644 --- a/dimos/navigation/experimental/dannav/local_planner/test_dan_local_planner.py +++ b/dimos/navigation/experimental/dannav/local_planner/test_dan_local_planner.py @@ -19,9 +19,13 @@ import math from typing import Any -from dimos.msgs.geometry_msgs.PointStamped import PointStamped -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.Path import Path +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, PointStamped, Pose, PoseStamped, Quaternion +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Header + +from dimos.msgs.time import time_from_seconds from dimos.navigation.experimental.dannav.local_planner.module import ( DanLocalPlannerConfig, _ReplanGate, @@ -29,26 +33,28 @@ def _path_from_points(points: list[tuple[float, float]]) -> Path: - poses = [ - PoseStamped( - ts=1.0, frame_id="world", position=[x, y, 0.0], orientation=[0.0, 0.0, 0.0, 1.0] - ) - for x, y in points - ] - return Path(frame_id="world", poses=poses) + value = Path( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + poses=[_odom(x, y) for x, y in points], + ) + return cdr_decode(cdr_encode(value), Path) def _odom(x: float, y: float, *, ts: float = 1.0) -> PoseStamped: - return PoseStamped( - ts=ts, - frame_id="world", - position=[x, y, 0.0], - orientation=[0.0, 0.0, 0.0, 1.0], + value = PoseStamped( + header=Header(stamp=time_from_seconds(ts), frame_id="world"), + pose=Pose( + position=Point(x=x, y=y, z=0.0), orientation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0) + ), ) + return cdr_decode(cdr_encode(value), PoseStamped) def _point(x: float, y: float) -> PointStamped: - return PointStamped(x=x, y=y, z=0.0, frame_id="world") + value = PointStamped( + point=Point(x=x, y=y, z=0.0), header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)) + ) + return cdr_decode(cdr_encode(value), PointStamped) def _gate(**config: Any) -> _ReplanGate: @@ -110,7 +116,7 @@ def test_empty_path_published_and_resets_gate() -> None: # An empty path (nothing safe ahead) forwards immediately as a stop and # drops the committed path. - empty = Path(frame_id="world", poses=[]) + empty = Path(header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), poses=[]) assert gate.on_planner_path(empty) is empty assert gate._committed is None @@ -145,3 +151,21 @@ def test_lock_replan_zero_commits_every_non_empty_path() -> None: _path_from_points([(0.0, 0.0), (1.0, 1.0)]), ): assert gate.on_planner_path(path) is not None + + +def test_smoothing_preserves_cdr_header_and_endpoints() -> None: + gate = _gate(resample_spacing_m=0.1, smoothing_window=3) + path = _path_from_points([(0.0, 0.0), (0.5, 0.5), (1.0, 0.0)]) + path.header.stamp.sec = 1700000000 + path.header.stamp.nanosec = 123456789 + output = gate.on_planner_path(path) + assert output is not None + decoded = cdr_decode(cdr_encode(output), Path) + assert decoded.header.frame_id == "world" + assert decoded.header.stamp.sec == 1700000000 + assert decoded.header.stamp.nanosec == 123456789 + assert len(decoded.poses) > len(path.poses) + assert decoded.poses[0].pose.position.x == 0.0 + assert decoded.poses[-1].pose.position.x == 1.0 + assert decoded.poses[-1].pose.position.y == 0.0 + assert path.poses[1].pose.position.y == 0.5 diff --git a/dimos/navigation/experimental/frontier_exploration/test_generated_frontiers.py b/dimos/navigation/experimental/frontier_exploration/test_generated_frontiers.py new file mode 100644 index 00000000000..88635ce5745 --- /dev/null +++ b/dimos/navigation/experimental/frontier_exploration/test_generated_frontiers.py @@ -0,0 +1,118 @@ +# 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. + +import math + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest + +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.occupancy import grid_to_world +from dimos.navigation.experimental.frontier_exploration.wavefront_frontier_goal_selector import ( + WavefrontFrontierExplorer, +) + + +@pytest.mark.parametrize("yaw", [0, math.pi / 2]) +def test_generated_frontier_centroid_respects_map_origin(request, yaw): + explorer = WavefrontFrontierExplorer(min_frontier_perimeter=0.1) + request.addfinalizer(explorer.stop) + cells = np.full((20, 20), -1, dtype=np.int8) + cells[5:15, 5:15] = 0 + grid = OccupancyGrid( + header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), + info=MapMetaData( + width=20, + height=20, + resolution=0.5, + origin=Pose( + position=Point(x=2, y=3, z=1), orientation=quaternion_from_euler(0, 0, yaw) + ), + map_load_time=Time(sec=0, nanosec=0), + ), + data=np.asarray(cells.ravel(), dtype=np.int8), + ) + received = cdr_decode(cdr_encode(grid), OccupancyGrid) + centroids = explorer.detect_frontiers(grid_to_world(received, (10, 10)), received) + assert len(centroids) == 1 + expected = grid_to_world(received, (9.5, 9.5)) + assert (centroids[0].x, centroids[0].y, centroids[0].z) == pytest.approx( + (expected.x, expected.y, expected.z) + ) + + +def test_stopping_exploration_retains_current_pose_header(request): + explorer = WavefrontFrontierExplorer() + request.addfinalizer(explorer.stop) + source = PoseStamped( + header=Header(frame_id="odom", stamp=Time(sec=1700000000, nanosec=123456789)), + pose=Pose( + position=Point(x=2, y=3, z=0.0), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), + ) + explorer.latest_odometry = source + explorer.exploration_active = True + goals = [] + explorer.goal_request.subscribe( + lambda value: goals.append(cdr_decode(cdr_encode(value), PoseStamped)) + ) + assert explorer.stop_exploration() + assert goals == [source] + source.pose.position.x = 100 + assert goals[0].pose.position.x == 2 + + +def test_exploration_loop_publishes_generated_goal_with_map_header(request): + explorer = WavefrontFrontierExplorer(goal_timeout=0.001) + request.addfinalizer(explorer.stop) + cells = np.full((20, 20), -1, dtype=np.int8) + cells[:, :10] = 0 + header = Header(frame_id="odom", stamp=Time(sec=1700000000, nanosec=123456789)) + explorer.latest_costmap = OccupancyGrid( + header=header, + info=MapMetaData( + width=20, + height=20, + resolution=0.5, + origin=Pose( + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + position=Point(x=0.0, y=0.0, z=0.0), + ), + map_load_time=Time(sec=0, nanosec=0), + ), + data=np.asarray(cells.ravel(), dtype=np.int8), + ) + explorer.latest_odometry = PoseStamped( + header=header, + pose=Pose( + position=Point(x=2, y=5, z=0.0), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), + ) + goals = [] + + def receive(value): + goals.append(cdr_decode(cdr_encode(value), PoseStamped)) + explorer.exploration_active = False + + explorer.goal_request.subscribe(receive) + explorer.exploration_active = True + explorer._run_exploration_loop() + assert len(goals) == 1 + assert goals[0].header == header + assert goals[0].pose.position.x == 5 diff --git a/dimos/navigation/experimental/frontier_exploration/wavefront_frontier_goal_selector.py b/dimos/navigation/experimental/frontier_exploration/wavefront_frontier_goal_selector.py index f42128b3c32..cecc1fa7867 100644 --- a/dimos/navigation/experimental/frontier_exploration/wavefront_frontier_goal_selector.py +++ b/dimos/navigation/experimental/frontier_exploration/wavefront_frontier_goal_selector.py @@ -809,7 +809,7 @@ def _run_exploration_loop(self) -> None: # Publish goal to navigator goal_msg = PoseStamped( header=self.latest_costmap.header, - pose=Pose(position=goal, orientation=Quaternion(w=1)), + pose=Pose(position=goal, orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0)), ) self.goal_request.publish(goal_msg) diff --git a/dimos/navigation/experimental/patrolling/module.py b/dimos/navigation/experimental/patrolling/module.py index 8e1d470686e..dc6b9ef1590 100644 --- a/dimos/navigation/experimental/patrolling/module.py +++ b/dimos/navigation/experimental/patrolling/module.py @@ -16,7 +16,9 @@ import asyncio from collections.abc import AsyncGenerator -from dimos_lcm.std_msgs import Bool +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.nav_msgs.msg import OccupancyGrid +from dimos_generated.std_msgs.msg import Bool from dimos.agents.annotation import skill from dimos.agents.capabilities import CAP_MOVEMENT @@ -24,8 +26,6 @@ from dimos.core.global_config import GlobalConfig, global_config from dimos.core.module import Module from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid from dimos.navigation.experimental.patrolling.constants import EXTRA_CLEARANCE from dimos.navigation.experimental.patrolling.create_patrol_router import create_patrol_router from dimos.navigation.experimental.patrolling.routers.patrol_router import PatrolRouter diff --git a/dimos/navigation/experimental/patrolling/routers/base_patrol_router.py b/dimos/navigation/experimental/patrolling/routers/base_patrol_router.py index 46975826bd2..3d654c2e8cd 100644 --- a/dimos/navigation/experimental/patrolling/routers/base_patrol_router.py +++ b/dimos/navigation/experimental/patrolling/routers/base_patrol_router.py @@ -16,11 +16,11 @@ from threading import RLock import time +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.nav_msgs.msg import OccupancyGrid import numpy as np from numpy.typing import NDArray -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid from dimos.navigation.experimental.patrolling.routers.visitation_history import VisitationHistory @@ -62,7 +62,7 @@ def handle_odom(self, msg: PoseStamped) -> None: self._pose = msg if self._occupancy_grid is None: return - self._visitation.handle_odom(msg.position.x, msg.position.y) + self._visitation.handle_odom(msg.pose.position.x, msg.pose.position.y) def get_saturation(self) -> float: with self._lock: diff --git a/dimos/navigation/experimental/patrolling/routers/coverage_patrol_router.py b/dimos/navigation/experimental/patrolling/routers/coverage_patrol_router.py index d887fe975b9..77e74aa0599 100644 --- a/dimos/navigation/experimental/patrolling/routers/coverage_patrol_router.py +++ b/dimos/navigation/experimental/patrolling/routers/coverage_patrol_router.py @@ -12,14 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.nav_msgs.msg import OccupancyGrid, Path import numpy as np from numpy.typing import NDArray from scipy.ndimage import binary_erosion from dimos.mapping.occupancy.gradient import gradient, voronoi_gradient -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.occupancy import grid_to_world, occupancy_view, world_to_grid from dimos.navigation.experimental.patrolling.routers.base_patrol_router import BasePatrolRouter from dimos.navigation.experimental.patrolling.utilities import point_to_pose_stamped from dimos.navigation.go2.replanning_a_star.min_cost_astar import min_cost_astar @@ -48,7 +48,7 @@ def handle_occupancy_grid(self, msg: OccupancyGrid) -> None: # Precompute the safe mask (cells with enough clearance from obstacles). clearance_cells = self._visitation.clearance_radius_cells - free_mask = msg.grid == 0 + free_mask = occupancy_view(msg) == 0 structure = np.ones((2 * clearance_cells + 1, 2 * clearance_cells + 1), dtype=bool) self._safe_mask = binary_erosion(free_mask, structure=structure).astype(bool) @@ -57,7 +57,7 @@ def handle_occupancy_grid(self, msg: OccupancyGrid) -> None: # open areas. Low voronoi cost = on the skeleton (equidistant from # walls) = high sampling weight. voronoi = voronoi_gradient(msg, max_distance=1.5) - voronoi_cost = voronoi.grid.astype(np.float64) + voronoi_cost = occupancy_view(voronoi).astype(np.float64) # Invert: skeleton cells (cost 0) become weight 100, walls (100) become 0. # Clamp negatives (unknown = -1) to 0. weights = np.clip(100.0 - voronoi_cost, 0.0, 100.0) @@ -83,7 +83,7 @@ def next_goal(self) -> PoseStamped | None: if pose is None: return None - start = (pose.position.x, pose.position.y) + start = pose.pose.position # Get candidate points from unvisited safe cells. unvisited_safe = safe_mask & ~visited @@ -112,8 +112,8 @@ def next_goal(self) -> PoseStamped | None: best_point = None for row, col in chosen: - world = occupancy_grid.grid_to_world((col, row, 0)) - candidate = (world.x, world.y) + world = grid_to_world(occupancy_grid, (col, row)) + candidate = world path = min_cost_astar(costmap, candidate, start, unknown_penalty=1.0, use_cpp=True) if path is None: @@ -127,7 +127,7 @@ def next_goal(self) -> PoseStamped | None: if best_point is None: return None - return point_to_pose_stamped(best_point) + return point_to_pose_stamped(best_point, occupancy_grid.header) def _count_new_coverage( self, @@ -142,11 +142,11 @@ def _count_new_coverage( # Sample every few poses to avoid redundant work on dense paths. step = max(1, r) - poses = path.poses[::step] + poses = list(path.poses)[::step] for pose in poses: - grid = occupancy_grid.world_to_grid((pose.position.x, pose.position.y)) - col, row = int(grid.x), int(grid.y) + grid = world_to_grid(occupancy_grid, pose.pose.position) + col, row = int(np.floor(round(grid[0], 9))), int(np.floor(round(grid[1], 9))) r_min = max(0, row - r) r_max = min(h, row + r + 1) c_min = max(0, col - r) diff --git a/dimos/navigation/experimental/patrolling/routers/frontier_patrol_router.py b/dimos/navigation/experimental/patrolling/routers/frontier_patrol_router.py index d9c55fce2f8..0c487b19cbc 100644 --- a/dimos/navigation/experimental/patrolling/routers/frontier_patrol_router.py +++ b/dimos/navigation/experimental/patrolling/routers/frontier_patrol_router.py @@ -14,12 +14,13 @@ from typing import Any +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.nav_msgs.msg import OccupancyGrid import numpy as np from numpy.typing import NDArray from scipy.ndimage import binary_erosion, label -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid +from dimos.msgs.occupancy import grid_to_world, occupancy_view, world_to_grid from dimos.navigation.experimental.patrolling.routers.base_patrol_router import BasePatrolRouter from dimos.navigation.experimental.patrolling.utilities import point_to_pose_stamped @@ -49,7 +50,7 @@ def handle_occupancy_grid(self, msg: OccupancyGrid) -> None: return clearance_cells = self._visitation.clearance_radius_cells - free_mask = msg.grid == 0 + free_mask = occupancy_view(msg) == 0 structure = np.ones((2 * clearance_cells + 1, 2 * clearance_cells + 1), dtype=bool) self._safe_mask = binary_erosion(free_mask, structure=structure).astype(bool) @@ -66,8 +67,8 @@ def next_goal(self) -> PoseStamped | None: return None # Robot position in grid coordinates. - grid_pos = occupancy_grid.world_to_grid((pose.position.x, pose.position.y)) - robot_col, robot_row = grid_pos.x, grid_pos.y + grid_pos = world_to_grid(occupancy_grid, pose.pose.position) + robot_col, robot_row = grid_pos[0], grid_pos[1] # Unvisited safe cells. unvisited_safe = safe_mask & ~visited @@ -121,5 +122,5 @@ def next_goal(self) -> PoseStamped | None: dists_sq = cluster_dr * cluster_dr + cluster_dc * cluster_dc goal_row, goal_col = cluster_indices[np.argmax(dists_sq)] - world = occupancy_grid.grid_to_world((int(goal_col), int(goal_row), 0)) - return point_to_pose_stamped((world.x, world.y)) + world = grid_to_world(occupancy_grid, (int(goal_col), int(goal_row))) + return point_to_pose_stamped(world, occupancy_grid.header) diff --git a/dimos/navigation/experimental/patrolling/routers/random_patrol_router.py b/dimos/navigation/experimental/patrolling/routers/random_patrol_router.py index acd1072a0a0..368975e87fe 100644 --- a/dimos/navigation/experimental/patrolling/routers/random_patrol_router.py +++ b/dimos/navigation/experimental/patrolling/routers/random_patrol_router.py @@ -12,12 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.geometry_msgs.msg import Point, PoseStamped +from dimos_generated.nav_msgs.msg import OccupancyGrid import numpy as np from numpy.typing import NDArray from scipy.ndimage import binary_erosion -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid +from dimos.msgs.occupancy import grid_to_world, occupancy_view from dimos.navigation.experimental.patrolling.routers.base_patrol_router import BasePatrolRouter from dimos.navigation.experimental.patrolling.utilities import point_to_pose_stamped @@ -34,17 +35,17 @@ def next_goal(self) -> PoseStamped | None: ) if point is None: return None - return point_to_pose_stamped(point) + return point_to_pose_stamped(point, occupancy_grid.header) def _random_empty_spot( occupancy_grid: OccupancyGrid, clearance_m: float, visited: NDArray[np.bool_] | None = None, -) -> tuple[float, float] | None: - clearance_cells = int(np.ceil(clearance_m / occupancy_grid.resolution)) +) -> Point | None: + clearance_cells = int(np.ceil(clearance_m / occupancy_grid.info.resolution)) - free_mask = occupancy_grid.grid == 0 + free_mask = occupancy_view(occupancy_grid) == 0 if not np.any(free_mask): return None @@ -64,5 +65,5 @@ def _random_empty_spot( idx = safe_indices[np.random.randint(len(safe_indices))] row, col = idx - world = occupancy_grid.grid_to_world((col, row, 0)) - return (world.x, world.y) + world = grid_to_world(occupancy_grid, (col, row)) + return world diff --git a/dimos/navigation/experimental/patrolling/routers/visitation_history.py b/dimos/navigation/experimental/patrolling/routers/visitation_history.py index e3d9a65e015..5d093806245 100644 --- a/dimos/navigation/experimental/patrolling/routers/visitation_history.py +++ b/dimos/navigation/experimental/patrolling/routers/visitation_history.py @@ -108,7 +108,7 @@ def _stamp(self, x: float, y: float) -> None: if grid is None or visited is None: return r = self._clearance_radius_cells - grid_pos = world_to_grid(grid, Point(x=x, y=y)) + grid_pos = world_to_grid(grid, Point(x=x, y=y, z=0.0)) col, row = int(np.floor(round(grid_pos[0], 9))), int(np.floor(round(grid_pos[1], 9))) if row + r < 0 or row - r >= grid.info.height or col + r < 0 or col - r >= grid.info.width: return diff --git a/dimos/navigation/experimental/patrolling/test_create_patrol_router.py b/dimos/navigation/experimental/patrolling/test_create_patrol_router.py index 420bc8a1d63..668a0a8b959 100644 --- a/dimos/navigation/experimental/patrolling/test_create_patrol_router.py +++ b/dimos/navigation/experimental/patrolling/test_create_patrol_router.py @@ -16,6 +16,10 @@ import os import cv2 +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point +from dimos_generated.nav_msgs.msg import OccupancyGrid +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest @@ -24,8 +28,9 @@ from dimos.mapping.occupancy.visualizations import visualize_occupancy_grid from dimos.mapping.pointclouds.occupancy import height_cost_occupancy from dimos.mapping.pointclouds.util import read_pointcloud -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.image import image_view +from dimos.msgs.occupancy import world_to_grid +from dimos.msgs.pointcloud import pointcloud_from_xyz from dimos.navigation.experimental.patrolling.create_patrol_router import create_patrol_router from dimos.navigation.experimental.patrolling.utilities import point_to_pose_stamped from dimos.navigation.go2.replanning_a_star.min_cost_astar import min_cost_astar @@ -35,7 +40,9 @@ @pytest.fixture def big_office() -> OccupancyGrid: data = read_pointcloud(get_data("big_office.ply")) - cloud = PointCloud2.from_numpy(np.asarray(data.points), frame_id="") + cloud = pointcloud_from_xyz( + np.asarray(data.points), header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)) + ) return height_cost_occupancy(cloud) @@ -44,13 +51,13 @@ def big_office() -> OccupancyGrid: "router_name, saturation", [("random", 0.20), ("coverage", 0.30), ("frontier", 0.20)] ) def test_patrolling_coverage(router_name, saturation, big_office) -> None: - start = (-1.03, -13.48) + start = Point(x=-1.03, y=-13.48, z=0.0) robot_width = 0.4 multiplier = 1.5 big_office_gradient = gradient(big_office, max_distance=1.5) router = create_patrol_router(router_name, robot_width * multiplier) router.handle_occupancy_grid(big_office) - router.handle_odom(point_to_pose_stamped(start)) + router.handle_odom(point_to_pose_stamped(start, big_office.header)) all_poses: list = [] for _ in range(15): @@ -58,15 +65,15 @@ def test_patrolling_coverage(router_name, saturation, big_office) -> None: if goal is None: continue path = min_cost_astar( - big_office_gradient, goal.position, start, unknown_penalty=1.0, use_cpp=True + big_office_gradient, goal.pose.position, start, unknown_penalty=1.0, use_cpp=True ) if path is None: continue - path = smooth_resample_path(path, goal, 0.1) + path = smooth_resample_path(path, goal.pose, 0.1) for pose in path.poses: router.handle_odom(pose) all_poses.append(pose) - start = (path.poses[-1].position.x, path.poses[-1].position.y) + start = path.poses[-1].pose.position assert router.get_saturation() > saturation @@ -76,14 +83,15 @@ def test_patrolling_coverage(router_name, saturation, big_office) -> None: def _save_coverage_image(router_name, router, all_poses, big_office, big_office_gradient) -> None: image = visualize_occupancy_grid(big_office_gradient, "rainbow") - h, w = image.data.shape[:2] + pixels = image_view(image) + h, w = pixels.shape[:2] visit_counts = np.zeros((h, w), dtype=np.float32) - radius = int(np.ceil(router._clearance_radius_m / big_office.resolution)) + radius = int(np.ceil(router._clearance_radius_m / big_office.info.resolution)) stamp = np.zeros((h, w), dtype=np.uint8) for pose in all_poses: - grid = big_office.world_to_grid((pose.position.x, pose.position.y)) - gx, gy = int(grid.x), int(grid.y) + grid = world_to_grid(big_office, pose.pose.position) + gx, gy = int(grid[0]), int(grid[1]) if 0 <= gy < h and 0 <= gx < w: stamp[:] = 0 cv2.circle(stamp, (gx, gy), radius, 1, -1) @@ -93,9 +101,8 @@ def _save_coverage_image(router_name, router, all_poses, big_office, big_office_ mask = visit_counts > 0 blend = 1.0 - (1.0 - alpha) ** visit_counts - overlay = image.data.astype(np.float32) * 0.24 + overlay = pixels.astype(np.float32) * 0.24 for c in range(3): overlay[:, :, c][mask] = overlay[:, :, c][mask] * (1.0 - blend[mask]) + 255.0 * blend[mask] - image.data = overlay.astype(np.uint8) - image.save(f"patrolling_coverage_{router_name}.png") + cv2.imwrite(f"patrolling_coverage_{router_name}.png", overlay.astype(np.uint8)) diff --git a/dimos/navigation/experimental/patrolling/test_generated_patrol.py b/dimos/navigation/experimental/patrolling/test_generated_patrol.py new file mode 100644 index 00000000000..588af6ebf9f --- /dev/null +++ b/dimos/navigation/experimental/patrolling/test_generated_patrol.py @@ -0,0 +1,68 @@ +# 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. + + +import math + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest + +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.occupancy import grid_to_world, occupancy_view, world_to_grid +from dimos.navigation.experimental.patrolling.create_patrol_router import create_patrol_router +from dimos.navigation.experimental.patrolling.utilities import point_to_pose_stamped + + +@pytest.mark.parametrize("router_name", ["random", "coverage", "frontier"]) +@pytest.mark.parametrize("yaw", [0, math.pi / 2]) +def test_generated_patrol_goals(router_name, yaw): + cells = np.zeros((24, 24), dtype=np.int8) + cells[[0, -1], :] = 100 + cells[:, [0, -1]] = 100 + grid = OccupancyGrid( + header=Header(frame_id="map", stamp=Time(sec=1700000000, nanosec=123456789)), + info=MapMetaData( + width=24, + height=24, + resolution=0.5, + origin=Pose( + position=Point(x=2, y=3, z=1), orientation=quaternion_from_euler(0, 0, yaw) + ), + map_load_time=Time(sec=0, nanosec=0), + ), + data=np.asarray(cells.ravel(), dtype=np.int8), + ) + router = create_patrol_router(router_name, 0.5) + assert router.next_goal() is None + received = cdr_decode(cdr_encode(grid), OccupancyGrid) + router.handle_occupancy_grid(received) + router.handle_odom(point_to_pose_stamped(grid_to_world(received, (12, 12)), received.header)) + initial = router.get_saturation() + goal = router.next_goal() + assert goal is not None + decoded = cdr_decode(cdr_encode(goal), PoseStamped) + assert decoded.header == grid.header + assert decoded.pose.position.z == pytest.approx(1) + x, y = world_to_grid(received, decoded.pose.position) + assert occupancy_view(received)[round(y), round(x)] == 0 + router.handle_odom(decoded) + assert router.get_saturation() > initial + router.reset() + assert router.next_goal() is None + assert router.get_saturation() == 0 diff --git a/dimos/navigation/experimental/patrolling/utilities.py b/dimos/navigation/experimental/patrolling/utilities.py index df9805bc85e..cdd8fa2ebae 100644 --- a/dimos/navigation/experimental/patrolling/utilities.py +++ b/dimos/navigation/experimental/patrolling/utilities.py @@ -17,7 +17,9 @@ def point_to_pose_stamped(point: Point, header: Header) -> PoseStamped: - return PoseStamped(header=header, pose=Pose(position=point, orientation=Quaternion(w=1))) + return PoseStamped( + header=header, pose=Pose(position=point, orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0)) + ) def pose_stamped_to_point(pose: PoseStamped) -> tuple[float, float]: diff --git a/dimos/navigation/global_planner/mls_planner/mls_planner_native.py b/dimos/navigation/global_planner/mls_planner/mls_planner_native.py index 46ef21800de..aea5a48a869 100644 --- a/dimos/navigation/global_planner/mls_planner/mls_planner_native.py +++ b/dimos/navigation/global_planner/mls_planner/mls_planner_native.py @@ -16,14 +16,14 @@ from __future__ import annotations +from dimos_generated.dimos_msgs.msg import RegionBounds, RegionLineSegments3D, RegionPointCloud2 +from dimos_generated.geometry_msgs.msg import PointStamped +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.tf2_msgs.msg import TFMessage + from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PointStamped import PointStamped -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.LineSegments3D import LineSegments3D -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.navigation import spec @@ -72,15 +72,15 @@ class MLSPlannerNative(NativeModule, spec.GlobalPlanner): global_map: In[PointCloud2] local_map: In[PointCloud2] - region_bounds: In[PoseStamped] + region_bounds: In[RegionBounds] # A seeded map's regions as RayTracingVoxelMap lands them, applied through # the region pipeline between live updates. Live updates keep priority. - seed_map: In[PointCloud2] - seed_bounds: In[PoseStamped] + seed_map: In[RegionPointCloud2] + seed_bounds: In[RegionBounds] goal: In[PointStamped] tf: In[TFMessage] path: Out[Path] - surface_map: Out[PointCloud2] + surface_map: Out[RegionPointCloud2] nodes: Out[PointCloud2] - node_edges: Out[LineSegments3D] + node_edges: Out[RegionLineSegments3D] diff --git a/dimos/navigation/global_planner/mls_planner/rust/Cargo.toml b/dimos/navigation/global_planner/mls_planner/rust/Cargo.toml index 0d81d18749c..c57949ef5cf 100644 --- a/dimos/navigation/global_planner/mls_planner/rust/Cargo.toml +++ b/dimos/navigation/global_planner/mls_planner/rust/Cargo.toml @@ -22,23 +22,23 @@ python = "dimos.navigation.global_planner.mls_planner.mls_planner_native:MLSPlan threads = 2 [package.metadata.dimos.module.mls_planner.inputs] -global_map = "sensor_msgs.PointCloud2" -local_map = "sensor_msgs.PointCloud2" -region_bounds = "geometry_msgs.PoseStamped" -seed_map = "sensor_msgs.PointCloud2" -seed_bounds = "geometry_msgs.PoseStamped" -goal = "geometry_msgs.PointStamped" -tf = "tf2_msgs.TFMessage" +global_map = "sensor_msgs/msg/PointCloud2" +local_map = "sensor_msgs/msg/PointCloud2" +region_bounds = "dimos_msgs/msg/RegionBounds" +seed_map = "dimos_msgs/msg/RegionPointCloud2" +seed_bounds = "dimos_msgs/msg/RegionBounds" +goal = "geometry_msgs/msg/PointStamped" +tf = "tf2_msgs/msg/TFMessage" [package.metadata.dimos.module.mls_planner.outputs] -surface_map = "sensor_msgs.PointCloud2" -nodes = "sensor_msgs.PointCloud2" -node_edges = "nav_msgs.LineSegments3D" -path = "nav_msgs.Path" +surface_map = "dimos_msgs/msg/RegionPointCloud2" +nodes = "sensor_msgs/msg/PointCloud2" +node_edges = "dimos_msgs/msg/RegionLineSegments3D" +path = "nav_msgs/msg/Path" [dependencies] dimos-module = { path = "../../../../../native/rust/dimos-module" } -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +dimos-generated-messages = { path = "../../../../../packages/dimos-generated/build/message-source/rust" } tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal", "time"] } serde = { version = "1", features = ["derive"] } ahash = "0.8" diff --git a/dimos/navigation/global_planner/mls_planner/rust/src/module.rs b/dimos/navigation/global_planner/mls_planner/rust/src/module.rs index d8a3d497981..9c17b643241 100644 --- a/dimos/navigation/global_planner/mls_planner/rust/src/module.rs +++ b/dimos/navigation/global_planner/mls_planner/rust/src/module.rs @@ -12,20 +12,31 @@ // See the License for the specific language governing permissions and // limitations under the License. +use dimos_generated_messages::builtin_interfaces::msg::time::Time; +use dimos_generated_messages::dimos_msgs::msg::{ + line_segment3_d::LineSegment3D, line_segments3_d::LineSegments3D, region_bounds::RegionBounds, + region_line_segments3_d::RegionLineSegments3D, region_point_cloud2::RegionPointCloud2, +}; +use dimos_generated_messages::geometry_msgs::msg::{ + point::Point, point_stamped::PointStamped, pose::Pose, pose_stamped::PoseStamped, + quaternion::Quaternion, +}; +use dimos_generated_messages::nav_msgs::msg::path::Path; +use dimos_generated_messages::sensor_msgs::msg::{ + point_cloud2::PointCloud2, point_field::PointField, +}; +use dimos_generated_messages::std_msgs::msg::header::Header; +use dimos_module::cdr; use std::collections::{HashMap, VecDeque}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; -use crate::mls_planner::{Config, Planner, RegionBounds}; +use crate::mls_planner::{Config, Planner, RegionBounds as PlannerRegionBounds}; use crate::region_viz::{pack_cell, Cell, RegionContent, RegionViz}; use crate::voxel::{surface_point_xyz, VoxelKey}; use dimos_module::time::now; use dimos_module::{debug_throttled, error_throttled, warn_throttled, Input, Module, Output, Tf}; -use lcm_msgs::geometry_msgs::{Point, PointStamped, Pose, PoseStamped, Quaternion}; -use lcm_msgs::nav_msgs::Path; -use lcm_msgs::sensor_msgs::{PointCloud2, PointField}; -use lcm_msgs::std_msgs::{Header, Time}; use tokio::sync::Notify; use tracing::{debug, info, warn}; @@ -41,7 +52,7 @@ type Shared = Arc>>; enum MapUpdate { Region { cloud: PointCloud2, - bounds: PoseStamped, + bounds: RegionBounds, received: Instant, }, Global { @@ -52,8 +63,8 @@ enum MapUpdate { /// One region of a seeded map, as the ray tracer hands them on. They queue, /// since every one must land. struct SeedRegion { - cloud: PointCloud2, - bounds: PoseStamped, + cloud: RegionPointCloud2, + bounds: RegionBounds, } /// How long half a seed region waits for its counterpart before it is dropped. @@ -143,31 +154,31 @@ impl SeedProgress { } /// Seed clouds and bounds waiting for their counterpart, keyed by the region -/// number in their header seq, since every region of a seed shares one stamp. +/// number in their region ID, since every region of a seed shares one stamp. #[derive(Default)] struct SeedPairs { - clouds: HashMap, - bounds: HashMap, + clouds: HashMap, + bounds: HashMap, } impl SeedPairs { - fn pair_cloud(&mut self, msg: PointCloud2) -> Option { + fn pair_cloud(&mut self, msg: RegionPointCloud2) -> Option { self.expire_before(Instant::now()); - match self.bounds.remove(&msg.header.seq) { + match self.bounds.remove(&msg.region_id) { Some((bounds, _)) => Some(SeedRegion { cloud: msg, bounds }), None => { - self.clouds.insert(msg.header.seq, (msg, Instant::now())); + self.clouds.insert(msg.region_id, (msg, Instant::now())); None } } } - fn pair_bounds(&mut self, msg: PoseStamped) -> Option { + fn pair_bounds(&mut self, msg: RegionBounds) -> Option { self.expire_before(Instant::now()); - match self.clouds.remove(&msg.header.seq) { + match self.clouds.remove(&msg.region_id) { Some((cloud, _)) => Some(SeedRegion { cloud, bounds: msg }), None => { - self.bounds.insert(msg.header.seq, (msg, Instant::now())); + self.bounds.insert(msg.region_id, (msg, Instant::now())); None } } @@ -193,40 +204,40 @@ impl SeedPairs { #[derive(Module)] #[module(name = "mls_planner", setup = spawn_worker, teardown = stop_worker)] pub struct MlsPlanner { - #[input(decode = PointCloud2::decode, handler = on_global_map)] + #[input(decode = cdr::decode, handler = on_global_map)] global_map: Input, - #[input(decode = PointCloud2::decode, handler = on_local_map)] + #[input(decode = cdr::decode, handler = on_local_map)] local_map: Input, - #[input(decode = PoseStamped::decode, handler = on_region_bounds)] - region_bounds: Input, + #[input(decode = cdr::decode, handler = on_region_bounds)] + region_bounds: Input, // A seeded map's regions as the ray tracer lands them, applied through // the region pipeline between live updates. Live updates keep priority. - #[input(decode = PointCloud2::decode, handler = on_seed_map)] - seed_map: Input, + #[input(decode = cdr::decode, handler = on_seed_map)] + seed_map: Input, - #[input(decode = PoseStamped::decode, handler = on_seed_bounds)] - seed_bounds: Input, + #[input(decode = cdr::decode, handler = on_seed_bounds)] + seed_bounds: Input, - #[input(decode = PointStamped::decode, handler = on_goal)] + #[input(decode = cdr::decode, handler = on_goal)] goal: Input, #[tf] tf: Tf, - #[output(encode = PointCloud2::encode)] - surface_map: Output, + #[output(encode = cdr::encode)] + surface_map: Output, - #[output(encode = PointCloud2::encode)] + #[output(encode = cdr::encode)] nodes: Output, // The wire payload is a Path. dimos names the channel LineSegments3D. - #[output(encode = Path::encode, msg = "LineSegments3D")] - node_edges: Output, + #[output(encode = cdr::encode)] + node_edges: Output, - #[output(encode = Path::encode)] + #[output(encode = cdr::encode)] path: Output, #[config] @@ -234,7 +245,7 @@ pub struct MlsPlanner { // Held on the handle loop until stamps match, then handed off paired. pending_local: Option, - pending_bounds: Option, + pending_bounds: Option, pending_seeds: SeedPairs, // Written by the handle loop, read by the worker, so the loop never blocks @@ -281,18 +292,18 @@ impl MlsPlanner { self.try_pair(); } - async fn on_region_bounds(&mut self, msg: PoseStamped) { + async fn on_region_bounds(&mut self, msg: RegionBounds) { self.pending_bounds = Some(msg); self.try_pair(); } - async fn on_seed_map(&mut self, msg: PointCloud2) { + async fn on_seed_map(&mut self, msg: RegionPointCloud2) { if let Some(region) = self.pending_seeds.pair_cloud(msg) { self.queue_seed(region); } } - async fn on_seed_bounds(&mut self, msg: PoseStamped) { + async fn on_seed_bounds(&mut self, msg: RegionBounds) { if let Some(region) = self.pending_seeds.pair_bounds(msg) { self.queue_seed(region); } @@ -334,7 +345,7 @@ impl MlsPlanner { } /// True when bounds and a local cloud are both present with matching stamps. -fn stamps_paired(bounds: Option<&PoseStamped>, cloud: Option<&PointCloud2>) -> bool { +fn stamps_paired(bounds: Option<&RegionBounds>, cloud: Option<&PointCloud2>) -> bool { match (bounds, cloud) { (Some(b), Some(c)) => same_stamp(&b.header.stamp, &c.header.stamp), _ => false, @@ -358,9 +369,9 @@ struct Worker { wake: Arc, tf: Tf, config: Config, - surface_map: Output, + surface_map: Output, nodes: Output, - node_edges: Output, + node_edges: Output, path: Output, } @@ -548,7 +559,13 @@ impl Worker { viz: &mut RegionViz, ) -> Option { let region = region_bounds(&seed.bounds); - self.apply_region(planner, &seed.cloud, ®ion, viz, "seed region processed") + self.apply_region( + planner, + &seed.cloud.cloud, + ®ion, + viz, + "seed region processed", + ) } /// Replace the voxels in a region and repair the graph around them, @@ -558,7 +575,7 @@ impl Worker { &self, planner: &mut Planner, cloud: &PointCloud2, - region: &RegionBounds, + region: &PlannerRegionBounds, viz: &mut RegionViz, label: &'static str, ) -> Option { @@ -641,24 +658,24 @@ fn replan_due(goal_changed: bool, live_update_applied: bool) -> bool { /// The region a bounds message describes: position is the center, orientation /// carries radius, z_min and z_max. -fn region_bounds(msg: &PoseStamped) -> RegionBounds { - RegionBounds { - origin_x: msg.pose.position.x as f32, - origin_y: msg.pose.position.y as f32, - radius: msg.pose.orientation.x as f32, - z_min: msg.pose.orientation.y as f32, - z_max: msg.pose.orientation.z as f32, +fn region_bounds(msg: &RegionBounds) -> PlannerRegionBounds { + PlannerRegionBounds { + origin_x: msg.center.x as f32, + origin_y: msg.center.y as f32, + radius: msg.radius as f32, + z_min: msg.z_min as f32, + z_max: msg.z_max as f32, } } -/// One cell's surface and edge messages, its cell in the header seq. +/// One cell's surface and edge messages, its cell in the region ID. fn region_messages( cell: Cell, content: RegionContent, voxel_size: f32, frame: &str, stamp: Time, -) -> (PointCloud2, Path) { +) -> (RegionPointCloud2, RegionLineSegments3D) { let surface_points: Vec = content .surface .into_iter() @@ -667,10 +684,14 @@ fn region_messages( (x, y, z, clearance) }) .collect(); - let mut surface = build_pc2_xyzi(&surface_points, frame, stamp.clone()); - surface.header.seq = pack_cell(cell); - let mut edges = build_segments_path(content.segments, voxel_size, frame, stamp); - edges.header.seq = pack_cell(cell); + let surface = RegionPointCloud2 { + region_id: pack_cell(cell), + cloud: build_pc2_xyzi(&surface_points, frame, stamp.clone()), + }; + let edges = RegionLineSegments3D { + region_id: pack_cell(cell), + lines: build_segments_path(content.segments, voxel_size, frame, stamp), + }; (surface, edges) } @@ -680,10 +701,10 @@ fn is_at_goal(start: Xyz, goal: Xyz, tol: f32) -> bool { } fn same_stamp(a: &Time, b: &Time) -> bool { - a.sec == b.sec && a.nsec == b.nsec + a.sec == b.sec && a.nanosec == b.nanosec } -async fn publish_cloud(out: &Output, cloud: &PointCloud2) { +async fn publish_cloud(out: &Output, cloud: &T) { if let Err(e) = out.publish(cloud).await { error_throttled!( Duration::from_secs(1), @@ -694,7 +715,7 @@ async fn publish_cloud(out: &Output, cloud: &PointCloud2) { } } -async fn publish_path(out: &Output, msg: &Path) { +async fn publish_path(out: &Output, msg: &T) { if let Err(e) = out.publish(msg).await { error_throttled!( Duration::from_secs(1), @@ -707,7 +728,6 @@ async fn publish_path(out: &Output, msg: &Path) { fn header(frame_id: &str, stamp: Time) -> Header { Header { - seq: 0, stamp, frame_id: frame_id.into(), } @@ -761,17 +781,30 @@ fn build_segments_path( voxel_size: f32, frame_id: &str, stamp: Time, -) -> Path { - let mut poses: Vec = Vec::with_capacity(segments.len() * 2); - for (a, b, cost) in segments { - let pa = surface_point_xyz(a.0, a.1, a.2, voxel_size); - let pb = surface_point_xyz(b.0, b.1, b.2, voxel_size); - poses.push(pose_stamped(pa, cost as f64, frame_id, stamp.clone())); - poses.push(pose_stamped(pb, cost as f64, frame_id, stamp.clone())); - } - Path { +) -> LineSegments3D { + let segments = segments + .into_iter() + .map(|(a, b, weight)| { + let a = surface_point_xyz(a.0, a.1, a.2, voxel_size); + let b = surface_point_xyz(b.0, b.1, b.2, voxel_size); + LineSegment3D { + start: Point { + x: a.0 as f64, + y: a.1 as f64, + z: a.2 as f64, + }, + end: Point { + x: b.0 as f64, + y: b.1 as f64, + z: b.2 as f64, + }, + weight: weight as f64, + } + }) + .collect(); + LineSegments3D { header: header(frame_id, stamp), - poses, + segments, } } @@ -787,14 +820,14 @@ fn build_pc2_xyzi(points: &[Xyzi], frame_id: &str, stamp: Time) -> PointCloud2 { } let make_field = |name: &str, off: i32| PointField { name: name.into(), - offset: off, - datatype: PointField::FLOAT32 as u8, + offset: off as u32, + datatype: PointField::FLOAT32, count: 1, }; PointCloud2 { header: header(frame_id, stamp), height: 1, - width: n, + width: n as u32, fields: vec![ make_field("x", 0), make_field("y", 4), @@ -803,7 +836,7 @@ fn build_pc2_xyzi(points: &[Xyzi], frame_id: &str, stamp: Time) -> PointCloud2 { ], is_bigendian: false, point_step: 16, - row_step: 16 * n, + row_step: 16 * n as u32, data, is_dense: true, } @@ -819,18 +852,18 @@ fn build_pc2_xyz(points: &[(f32, f32, f32)], frame_id: &str, stamp: Time) -> Poi } let make_field = |name: &str, off: i32| PointField { name: name.into(), - offset: off, - datatype: PointField::FLOAT32 as u8, + offset: off as u32, + datatype: PointField::FLOAT32, count: 1, }; PointCloud2 { header: header(frame_id, stamp), height: 1, - width: n, + width: n as u32, fields: vec![make_field("x", 0), make_field("y", 4), make_field("z", 8)], is_bigendian: false, point_step: 12, - row_step: 12 * n, + row_step: 12 * n as u32, data, is_dense: true, } @@ -848,7 +881,7 @@ fn extract_xyz(msg: &PointCloud2) -> Result, ExtractError> let mut y_off: Option = None; let mut z_off: Option = None; for f in &msg.fields { - if f.datatype != PointField::FLOAT32 as u8 { + if f.datatype != PointField::FLOAT32 { continue; } match f.name.as_str() { @@ -915,32 +948,54 @@ mod tests { fn stamped(stamp: Time) -> Header { Header { stamp, - ..Default::default() + frame_id: String::new(), } } - fn bounds_at(stamp: Time) -> PoseStamped { - PoseStamped { + fn bounds_at(stamp: Time) -> RegionBounds { + RegionBounds { header: stamped(stamp), - ..Default::default() + region_id: 0, + center: Point { + x: 0.0, + y: 0.0, + z: 0.0, + }, + radius: 0.0, + z_min: 0.0, + z_max: 0.0, + } + } + + fn region_cloud_at(stamp: Time) -> RegionPointCloud2 { + RegionPointCloud2 { + region_id: 0, + cloud: cloud_at(stamp), } } fn cloud_at(stamp: Time) -> PointCloud2 { PointCloud2 { header: stamped(stamp), - ..Default::default() + height: 1, + width: 0, + fields: Vec::new(), + is_bigendian: false, + point_step: 0, + row_step: 0, + data: Vec::new(), + is_dense: true, } } #[test] fn stamps_paired_only_when_both_present_and_stamps_match() { - let s = Time { sec: 2, nsec: 3 }; + let s = Time { sec: 2, nanosec: 3 }; let b = bounds_at(s.clone()); let c = cloud_at(s); assert!(stamps_paired(Some(&b), Some(&c))); - let other = cloud_at(Time { sec: 2, nsec: 4 }); + let other = cloud_at(Time { sec: 2, nanosec: 4 }); assert!(!stamps_paired(Some(&b), Some(&other))); assert!(!stamps_paired(Some(&b), None)); @@ -950,62 +1005,61 @@ mod tests { #[test] fn seed_pairs_match_on_seq_whichever_side_lands_first() { - let stamp = Time { sec: 2, nsec: 3 }; + let stamp = Time { sec: 2, nanosec: 3 }; let mut pairs = SeedPairs::default(); let mut b1 = bounds_at(stamp.clone()); - b1.header.seq = 1; + b1.region_id = 1; let mut b2 = bounds_at(stamp.clone()); - b2.header.seq = 2; - let mut c1 = cloud_at(stamp.clone()); - c1.header.seq = 1; - let mut c2 = cloud_at(stamp); - c2.header.seq = 2; + b2.region_id = 2; + let mut c1 = region_cloud_at(stamp.clone()); + c1.region_id = 1; + let mut c2 = region_cloud_at(stamp); + c2.region_id = 2; assert!(pairs.pair_bounds(b1).is_none()); assert!(pairs.pair_bounds(b2).is_none()); let first = pairs .pair_cloud(c1) .expect("region 1 pairs with its own bounds"); - assert_eq!(first.bounds.header.seq, 1); + assert_eq!(first.bounds.region_id, 1); let second = pairs .pair_cloud(c2) .expect("region 2 pairs with its own bounds"); - assert_eq!(second.bounds.header.seq, 2); + assert_eq!(second.bounds.region_id, 2); assert!(pairs.clouds.is_empty() && pairs.bounds.is_empty()); } #[test] fn seed_pairs_match_cloud_first_and_interleaved() { - let stamp = Time { sec: 2, nsec: 3 }; + let stamp = Time { sec: 2, nanosec: 3 }; let mut pairs = SeedPairs::default(); - let mut c1 = cloud_at(stamp.clone()); - c1.header.seq = 1; + let mut c1 = region_cloud_at(stamp.clone()); + c1.region_id = 1; let mut b1 = bounds_at(stamp.clone()); - b1.header.seq = 1; + b1.region_id = 1; assert!(pairs.pair_cloud(c1).is_none()); assert_eq!( pairs .pair_bounds(b1) .expect("cloud first pairs") .cloud - .header - .seq, + .region_id, 1 ); // Interleaved: cloud 2, bounds 3, bounds 2, cloud 3. - let mut c2 = cloud_at(stamp.clone()); - c2.header.seq = 2; + let mut c2 = region_cloud_at(stamp.clone()); + c2.region_id = 2; let mut b3 = bounds_at(stamp.clone()); - b3.header.seq = 3; + b3.region_id = 3; let mut b2 = bounds_at(stamp.clone()); - b2.header.seq = 2; - let mut c3 = cloud_at(stamp); - c3.header.seq = 3; + b2.region_id = 2; + let mut c3 = region_cloud_at(stamp); + c3.region_id = 3; assert!(pairs.pair_cloud(c2).is_none()); assert!(pairs.pair_bounds(b3).is_none()); - assert_eq!(pairs.pair_bounds(b2).expect("region 2").cloud.header.seq, 2); - assert_eq!(pairs.pair_cloud(c3).expect("region 3").bounds.header.seq, 3); + assert_eq!(pairs.pair_bounds(b2).expect("region 2").cloud.region_id, 2); + assert_eq!(pairs.pair_cloud(c3).expect("region 3").bounds.region_id, 3); assert!(pairs.clouds.is_empty() && pairs.bounds.is_empty()); } @@ -1016,31 +1070,37 @@ mod tests { surface: vec![((1, 2, 0), 0.5)], segments: vec![((1, 2, 0), (2, 2, 0), 1.5)], }; - let (surface, edges) = region_messages(cell, content, 0.1, "odom", Time::default()); - assert_eq!(surface.header.seq, pack_cell(cell)); - assert_eq!(edges.header.seq, pack_cell(cell)); - assert_eq!((surface.width, edges.poses.len()), (1, 2)); - let (surface, edges) = - region_messages(cell, RegionContent::default(), 0.1, "odom", Time::default()); - assert_eq!(surface.header.seq, pack_cell(cell)); - assert_eq!(edges.header.seq, pack_cell(cell)); - assert_eq!((surface.width, edges.poses.len()), (0, 0)); + region_messages(cell, content, 0.1, "odom", Time { sec: 0, nanosec: 0 }); + assert_eq!(surface.region_id, pack_cell(cell)); + assert_eq!(edges.region_id, pack_cell(cell)); + assert_eq!((surface.cloud.width, edges.lines.segments.len()), (1, 1)); + + let (surface, edges) = region_messages( + cell, + RegionContent::default(), + 0.1, + "odom", + Time { sec: 0, nanosec: 0 }, + ); + assert_eq!(surface.region_id, pack_cell(cell)); + assert_eq!(edges.region_id, pack_cell(cell)); + assert_eq!((surface.cloud.width, edges.lines.segments.len()), (0, 0)); } #[test] fn an_unpaired_seed_half_expires_after_the_timeout() { - let stamp = Time { sec: 2, nsec: 3 }; + let stamp = Time { sec: 2, nanosec: 3 }; let mut pairs = SeedPairs::default(); let mut b1 = bounds_at(stamp.clone()); - b1.header.seq = 1; + b1.region_id = 1; assert!(pairs.pair_bounds(b1).is_none()); pairs.expire_before(Instant::now() + SEED_PAIR_TIMEOUT / 2); assert_eq!(pairs.bounds.len(), 1); pairs.expire_before(Instant::now() + SEED_PAIR_TIMEOUT * 2); assert!(pairs.bounds.is_empty()); - let mut c1 = cloud_at(stamp); - c1.header.seq = 1; + let mut c1 = region_cloud_at(stamp); + c1.region_id = 1; assert!( pairs.pair_cloud(c1).is_none(), "the expired bounds no longer pair" diff --git a/dimos/navigation/global_planner/mls_planner/test_start_relay.py b/dimos/navigation/global_planner/mls_planner/test_start_relay.py index 3756e770952..ffb1760aeb8 100644 --- a/dimos/navigation/global_planner/mls_planner/test_start_relay.py +++ b/dimos/navigation/global_planner/mls_planner/test_start_relay.py @@ -14,17 +14,32 @@ from typing import Any -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage +import pytest + +from dimos.core.transport_factory import rpc_backend from dimos.navigation.global_planner.mls_planner.start_relay import StartRelay from dimos.protocol.tf.tf import MultiTBuffer MOUNT_Z = 0.163 +@pytest.fixture(autouse=True) +def isolate_rpc_boundary(mocker): + for method in ("start", "serve_module_rpc", "stop"): + mocker.patch.object(rpc_backend(), method) + + class FakeTF(MultiTBuffer): """In-memory tf with the dispose() hook and call counter the module tests need.""" @@ -40,7 +55,7 @@ def get( time_tolerance: float | None = None, *, forward_tolerance: float = 0.0, - ) -> Transform | None: + ) -> TransformStamped | None: self.gets += 1 return super().get( parent_frame, @@ -54,23 +69,25 @@ def dispose(self) -> None: pass -def _mount() -> Transform: - return Transform( - translation=Vector3(0.0, 0.0, MOUNT_Z), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="base_link", +def _mount() -> TransformStamped: + return TransformStamped( + header=Header(frame_id="base_link", stamp=Time(sec=1, nanosec=0)), child_frame_id="mid360_link", - ts=1.0, + transform=Transform( + translation=Vector3(z=MOUNT_Z, x=0.0, y=0.0), + rotation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), + ), ) -def _odom_edge() -> Transform: - return Transform( - translation=Vector3(1.0, 2.0, 3.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="odom", +def _odom_edge() -> TransformStamped: + return TransformStamped( + header=Header(frame_id="odom", stamp=Time(sec=2, nanosec=0)), child_frame_id="mid360_link", - ts=2.0, + transform=Transform( + translation=Vector3(x=1.0, y=2.0, z=3.0), + rotation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), + ), ) @@ -78,7 +95,9 @@ def _relay(tf: FakeTF, **config: Any) -> tuple[StartRelay, list[PoseStamped]]: module = StartRelay(**config) module._tf = tf captured: list[PoseStamped] = [] - module.start_pose.subscribe(captured.append) + module.start_pose.subscribe( + lambda pose: captured.append(cdr_decode(cdr_encode(pose), PoseStamped)) + ) return module, captured @@ -88,12 +107,12 @@ def test_start_pose_is_the_tf_base_pose() -> None: tf.receive_transform(_odom_edge()) module, captured = _relay(tf) try: - module._on_tf(TFMessage()) + module._on_tf(TFMessage(transforms=[])) # The base sits MOUNT_Z below the sensor along the mount leg. assert len(captured) == 1 - assert abs(captured[0].position.x - 1.0) < 1e-9 - assert abs(captured[0].position.y - 2.0) < 1e-9 - assert abs(captured[0].position.z - (3.0 - MOUNT_Z)) < 1e-9 + assert abs(captured[0].pose.position.x - 1.0) < 1e-9 + assert abs(captured[0].pose.position.y - 2.0) < 1e-9 + assert abs(captured[0].pose.position.z - (3.0 - MOUNT_Z)) < 1e-9 finally: module.stop() @@ -103,7 +122,7 @@ def test_nothing_published_while_the_chain_is_incomplete() -> None: tf.receive_transform(_mount()) module, captured = _relay(tf) try: - module._on_tf(TFMessage()) + module._on_tf(TFMessage(transforms=[])) assert captured == [] finally: module.stop() @@ -113,8 +132,8 @@ def test_lookup_retries_are_throttled_during_an_outage() -> None: tf = FakeTF() module, captured = _relay(tf) try: - module._on_tf(TFMessage()) - module._on_tf(TFMessage()) + module._on_tf(TFMessage(transforms=[])) + module._on_tf(TFMessage(transforms=[])) assert captured == [] assert tf.gets == 1 finally: diff --git a/dimos/navigation/global_planner/mls_planner/test_transformer.py b/dimos/navigation/global_planner/mls_planner/test_transformer.py index 875e62ae31e..39df630f615 100644 --- a/dimos/navigation/global_planner/mls_planner/test_transformer.py +++ b/dimos/navigation/global_planner/mls_planner/test_transformer.py @@ -14,14 +14,20 @@ from __future__ import annotations +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode import numpy as np from numpy.typing import NDArray import pytest pytest.importorskip("dimos_mls_planner") +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header + from dimos.memory.type.observation import Observation -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_from_xyz from dimos.navigation.global_planner.mls_planner.transformer import MLSPlan @@ -35,7 +41,15 @@ def _obs( ts=0.0, pose=pose, tags={"region_bounds": region_bounds}, - _data=PointCloud2.from_numpy(points), + _data=cdr_decode( + cdr_encode( + pointcloud_from_xyz( + points, + header=Header(frame_id="world", stamp=Time(sec=1700000000, nanosec=123456789)), + ) + ), + PointCloud2, + ), ) @@ -55,6 +69,11 @@ def test_flat_floor_yields_populated_path_and_planned_true() -> None: [out] = list(MLSPlan(goal=(2.0, 2.0, 0.0), voxel_size=0.2, robot_height=1.0)(iter([obs]))) + decoded = cdr_decode(cdr_encode(out.data), Path) + assert decoded.header.frame_id == "world" + assert decoded.header.stamp.sec == 1700000000 + assert decoded.header.stamp.nanosec == 123456789 + assert all(p.header.stamp.nanosec == 123456789 for p in decoded.poses) assert out.tags["planned"] is True assert len(out.data.poses) >= 2 assert out.tags["voxels"] > 0 @@ -70,7 +89,15 @@ def test_poseless_obs_is_skipped() -> None: ts=0.0, pose=None, tags={"region_bounds": (0.0, 0.0, 5.0, -1.0, 2.0)}, - _data=PointCloud2.from_numpy(points), + _data=cdr_decode( + cdr_encode( + pointcloud_from_xyz( + points, + header=Header(frame_id="world", stamp=Time(sec=1700000000, nanosec=123456789)), + ) + ), + PointCloud2, + ), ) posed = _obs(points, pose=(-2.0, -2.0, 1.0), region_bounds=(0.0, 0.0, 5.0, -1.0, 2.0)) @@ -105,4 +132,4 @@ def test_no_route_yields_empty_path_with_planned_false() -> None: [out] = list(MLSPlan(goal=(100.0, 100.0, 100.0))(iter([obs]))) assert out.tags["planned"] is False - assert out.data.poses == [] + assert len(out.data.poses) == 0 diff --git a/dimos/navigation/global_planner/mls_planner/test_viz.py b/dimos/navigation/global_planner/mls_planner/test_viz.py index 01457cf2c7d..29d260cc72c 100644 --- a/dimos/navigation/global_planner/mls_planner/test_viz.py +++ b/dimos/navigation/global_planner/mls_planner/test_viz.py @@ -12,13 +12,95 @@ # See the License for the specific language governing permissions and # limitations under the License. -from functools import partial +"""Generated planner messages retain geometry through the Rerun adapters.""" +from functools import partial +import struct + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.dimos_msgs.msg import ( + LineSegment3D, + LineSegments3D, + RegionLineSegments3D, + RegionPointCloud2, +) +from dimos_generated.geometry_msgs.msg import Point +from dimos_generated.sensor_msgs.msg import PointCloud2, PointField +from dimos_generated.std_msgs.msg import Header import numpy as np -from dimos.msgs.nav_msgs.LineSegments3D import LineSegments3D -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_from_xyz from dimos.navigation.global_planner.mls_planner import viz +from dimos.navigation.global_planner.mls_planner.viz import ( + render_node_edges, + render_nodes, + render_surface_map, +) + + +def test_surface_clearance_filter_after_cdr() -> None: + cloud = PointCloud2( + height=1, + width=2, + point_step=16, + row_step=32, + fields=[ + PointField(name=name, offset=i * 4, datatype=PointField.FLOAT32, count=1) + for i, name in enumerate(("x", "y", "z", "intensity")) + ], + data=np.frombuffer(struct.pack("<8f", 1, 2, 3, 0.1, 4, 5, 6, 0.8), dtype=np.uint8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + is_bigendian=False, + is_dense=False, + ) + decoded = cdr_decode(cdr_encode(cloud), PointCloud2) + before = cdr_encode(decoded) + rendered = render_surface_map(decoded, wall_clearance_m=0.2) + assert rendered.positions.as_arrow_array().to_pylist() == [[4.0, 5.0, 6.0]] + assert cdr_encode(decoded) == before + + +def test_nodes_lift_and_surface_without_intensity() -> None: + cloud = pointcloud_from_xyz( + np.array([[1.0, 2.0, 3.0]]), header=Header(stamp=Time(sec=0, nanosec=0), frame_id="") + ) + before = cdr_encode(cloud) + np.testing.assert_allclose( + render_nodes(cloud).positions.as_arrow_array().to_pylist(), [[1, 2, 3.05]] + ) + assert render_surface_map(cloud).positions.as_arrow_array().to_pylist() == [[1, 2, 3]] + assert cdr_encode(cloud) == before + + +def test_weighted_edges_and_empty_messages() -> None: + message = LineSegments3D( + segments=[ + LineSegment3D(start=Point(x=0.0, y=0.0, z=0.0), end=Point(x=1, y=0.0, z=0.0), weight=1), + LineSegment3D(start=Point(x=1, y=0.0, z=0.0), end=Point(x=2, y=0.0, z=0.0), weight=100), + ], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + decoded = cdr_decode(cdr_encode(message), LineSegments3D) + rendered = render_node_edges(decoded) + np.testing.assert_allclose( + rendered.strips.as_arrow_array().to_pylist(), + [[[0, 0, 0.05], [1, 0, 0.05]], [[1, 0, 0.05], [2, 0, 0.05]]], + ) + assert rendered.colors.as_arrow_array().to_pylist() == [0x00FF3CDC, 0xFF003CDC] + assert decoded.segments[0].start.z == 0 + assert ( + render_node_edges( + LineSegments3D(header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), segments=[]) + ) + .strips.as_arrow_array() + .to_pylist() + == [] + ) + empty = pointcloud_from_xyz( + np.empty((0, 3)), header=Header(stamp=Time(sec=0, nanosec=0), frame_id="") + ) + assert render_nodes(empty).positions.as_arrow_array().to_pylist() == [] def _rgb(packed: int) -> tuple[int, int, int]: @@ -62,19 +144,34 @@ def test_graph_nodes_are_lifted_off_the_surface() -> None: def test_region_renders_are_static_and_an_empty_cell_still_lands() -> None: - cell = PointCloud2.from_numpy(np.array([[1.0, 1.0, 0.0]], dtype=np.float32)) - cell.seq = 1 << 16 + cell = RegionPointCloud2( + region_id=1 << 16, + cloud=pointcloud_from_xyz( + np.array([[1.0, 1.0, 0.0]], dtype=np.float32), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id="map"), + ), + ) (path, arch, static) = viz.render_surface_region(cell, 0.1, 0.1, 1.0)[0] assert path == "world/surface_map/1_0" and static assert len(arch.positions.as_arrow_array()) == 1 - emptied = PointCloud2.from_numpy(np.zeros((0, 3), dtype=np.float32)) - emptied.seq = 1 << 16 + emptied = RegionPointCloud2( + region_id=1 << 16, + cloud=pointcloud_from_xyz( + np.zeros((0, 3), dtype=np.float32), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id="map"), + ), + ) (path, arch, static) = viz.render_surface_region(emptied, 0.1, 0.1, 1.0)[0] assert path == "world/surface_map/1_0" and static assert len(arch.positions.as_arrow_array()) == 0 - edges = LineSegments3D(segments=np.zeros((0, 2, 3)), seq=(2 << 16) | 3) + edges = RegionLineSegments3D( + region_id=(2 << 16) | 3, + lines=LineSegments3D( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id="map"), segments=[] + ), + ) (path, arch, static) = viz.render_edge_region(edges)[0] assert path == "world/node_edges/2_3" and static assert arch.strips.as_arrow_array().to_pylist() == [] diff --git a/dimos/navigation/global_planner/mls_planner/transformer.py b/dimos/navigation/global_planner/mls_planner/transformer.py index 9937c260006..92b6d4e04d3 100644 --- a/dimos/navigation/global_planner/mls_planner/transformer.py +++ b/dimos/navigation/global_planner/mls_planner/transformer.py @@ -17,17 +17,20 @@ import time from typing import TYPE_CHECKING, Any +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header +import numpy as np + from dimos.memory.transform import Transformer -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_xyz from dimos.navigation.global_planner.mls_planner.mls_planner import MLSPlanner from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: from collections.abc import Iterator - import numpy as np from numpy.typing import NDArray from dimos.memory.type.observation import Observation @@ -51,19 +54,20 @@ def __init__( self.robot_height = robot_height self._planner_kwargs = planner_kwargs - def _path_from_waypoints(self, waypoints: NDArray[np.float32] | None, ts: float) -> Path: + def _path_from_waypoints(self, waypoints: NDArray[np.float32] | None, header: Header) -> Path: poses: list[PoseStamped] = [] if waypoints is not None: for x, y, z in waypoints: poses.append( PoseStamped( - ts=ts, - frame_id="world", - position=(float(x), float(y), float(z)), - orientation=(0.0, 0.0, 0.0, 1.0), + header=header, + pose=Pose( + position=Point(x=float(x), y=float(y), z=float(z)), + orientation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), + ), ) ) - return Path(ts=ts, frame_id="world", poses=poses) + return Path(header=header, poses=poses) def __call__( self, @@ -83,11 +87,18 @@ def __call__( ox, oy, radius, z_min, z_max = obs.tags["region_bounds"] t_update = time.perf_counter() - planner.update_region(obs.data.points_f32(), (ox, oy), radius, z_min, z_max, float(z)) + planner.update_region( + pointcloud_xyz(obs.data).astype(np.float32), + (ox, oy), + radius, + z_min, + z_max, + float(z), + ) t_plan = time.perf_counter() waypoints = planner.plan(start, self.goal) t_done = time.perf_counter() - path = self._path_from_waypoints(waypoints, obs.ts) + path = self._path_from_waypoints(waypoints, obs.data.header) timings = { "update_ms": (t_plan - t_update) * 1000, diff --git a/dimos/navigation/global_planner/mls_planner/utils/plan_rrd.py b/dimos/navigation/global_planner/mls_planner/utils/plan_rrd.py index a771b546a13..fef5893e71d 100644 --- a/dimos/navigation/global_planner/mls_planner/utils/plan_rrd.py +++ b/dimos/navigation/global_planner/mls_planner/utils/plan_rrd.py @@ -26,6 +26,9 @@ from time import perf_counter from typing import TYPE_CHECKING, Any, NamedTuple +from dimos_generated.geometry_msgs.msg import TransformStamped +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.tf2_msgs.msg import TFMessage import numpy as np from numpy.typing import NDArray import typer @@ -43,10 +46,7 @@ from dimos.memory.transform import FnTransformer from dimos.memory.type.observation import Observation from dimos.memory.vis.utils import DEFAULT_RENDER_VOXEL, default_render_voxel -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2, register_colormap_annotation -from dimos.msgs.tf2_msgs.TFMessage import TfFrameTree, TFMessage -from dimos.navigation.global_planner.mls_planner.mls_planner import MLSPlanner +from dimos.msgs.pointcloud import pointcloud_xyz from dimos.navigation.global_planner.mls_planner.viz import graph_edges, graph_nodes, surface_points from dimos.navigation.global_planner.viz import ( PATH_COLOR, @@ -62,10 +62,19 @@ ROBOT_WIDTH, ) from dimos.utils.data import resolve_named_path +from dimos.visualization.rerun.message_helpers import register_colormap_annotation, tf_archetypes if TYPE_CHECKING: import rerun.blueprint as rrb + from dimos.navigation.global_planner.mls_planner.mls_planner import MLSPlanner + +# Logging recorded TF/odometry does not require the optional planner extension. +try: + import dimos_mls_planner as _planner_module +except ImportError: + _planner_module = None + from dimos.memory.stream import Stream TIMELINE = "ts" @@ -139,6 +148,45 @@ def _parse_configs( return out +def _pose_from_tf(tf: StreamTF, world_frame: str) -> FnTransformer[PointCloud2, PointCloud2]: + """Attach the tf pose at the cloud stamp. A failed lookup clears any + recorded pose so the cloud is dropped downstream.""" + + def attach(obs: Observation[PointCloud2]) -> Observation[PointCloud2]: + t = tf.get( + world_frame, + obs.data.header.frame_id, + time_point=obs.ts, + time_tolerance=TF_MATCH_TOLERANCE_S, + ) + return obs.with_pose(t) + + return FnTransformer(attach) + + +def _log_edges(edges: NDArray[np.float32], entity: str) -> None: + import rerun as rr + + if edges.size == 0: + rr.log(entity, rr.LineStrips3D([])) + return + segments = [ + [(float(r[0]), float(r[1]), float(r[2])), (float(r[3]), float(r[4]), float(r[5]))] + for r in edges + ] + rr.log(entity, rr.LineStrips3D(segments, colors=[[255, 255, 255]], radii=0.02)) + + +def _log_path_wp(waypoints: NDArray[np.float32] | None, entity: str, color: list[int]) -> None: + import rerun as rr + + if waypoints is None or len(waypoints) == 0: + rr.log(entity, rr.LineStrips3D([])) + return + points = [(float(p[0]), float(p[1]), float(p[2])) for p in waypoints] + rr.log(entity, rr.LineStrips3D([points], colors=[color], radii=0.05)) + + def _tf_over(store: SqliteStore, window: Stream[Any]) -> Stream[TFMessage] | None: """The recorded tf stream clipped to another stream's span. @@ -161,14 +209,13 @@ class _TfSync: def __init__(self, tf: Stream[TFMessage] | None) -> None: self._pending = iter(tf) if tf is not None else iter(()) self._next = next(self._pending, None) - self._tree = TfFrameTree() def up_to(self, ts: float) -> None: import rerun as rr while self._next is not None and self._next.ts <= ts: rr.set_time(TIMELINE, timestamp=self._next.ts) - for path, archetype in self._next.data.to_rerun(self._tree): + for path, archetype in tf_archetypes(self._next.data): rr.log(path, archetype) self._next = next(self._pending, None) @@ -177,7 +224,7 @@ def _log_odometry( pose: tuple[float, ...], ts: float, trail: list[tuple[float, float, float]], - base: Transform | None, + base: TransformStamped | None, ) -> None: """Trace the sensor moving throughout the scene.""" import rerun as rr @@ -194,9 +241,18 @@ def _log_odometry( rr.log( "world/robot_body", rr.Transform3D( - translation=[base.translation.x, base.translation.y, base.translation.z], + translation=[ + base.transform.translation.x, + base.transform.translation.y, + base.transform.translation.z, + ], quaternion=rr.Quaternion( - xyzw=[base.rotation.x, base.rotation.y, base.rotation.z, base.rotation.w] + xyzw=[ + base.transform.rotation.x, + base.transform.rotation.y, + base.transform.rotation.z, + base.transform.rotation.w, + ] ), ), ) @@ -405,9 +461,14 @@ def _build_planners( step_height: float, step_penalty_weight: float, ) -> list[tuple[str, tuple[int, int, int], MLSPlanner]]: + if _planner_module is None: + raise ImportError( + "dimos_mls_planner is not built. Run: uv run maturin develop --release --uv " + "-m dimos/navigation/global_planner/mls_planner/rust/py/Cargo.toml" + ) planners: list[tuple[str, tuple[int, int, int], MLSPlanner]] = [] for i, (clr, buf, wgt) in enumerate(configs): - planner = MLSPlanner( + planner = _planner_module.MLSPlanner( voxel_size=voxel_size, robot_height=robot_height, max_overhead_m=max_overhead, @@ -442,11 +503,13 @@ def _process_frame( bounds = ray_obs.tags["region_bounds"] ox, oy, radius, z_min, z_max = bounds - pts = ray_obs.data.points_f32() + pts = pointcloud_xyz(ray_obs.data).astype(np.float32) rr.set_time(TIMELINE, timestamp=ray_obs.ts) ref_timing: dict[str, float] = {} - surface = nodes = edges = np.empty((0,), dtype=np.float32) + surface: NDArray[np.float32] = np.empty((0,), dtype=np.float32) + nodes: NDArray[np.float32] = np.empty((0,), dtype=np.float32) + edges: NDArray[np.float32] = np.empty((0,), dtype=np.float32) for j, (label, color, planner) in enumerate(planners): t0 = perf_counter() planner.update_region(pts, (ox, oy), radius, z_min, z_max, sensor_z) @@ -702,11 +765,11 @@ def main( if base is None: continue start = ( - float(base.translation.x), - float(base.translation.y), - float(base.translation.z) - start_z_offset, + float(base.transform.translation.x), + float(base.transform.translation.y), + float(base.transform.translation.z) - start_z_offset, ) - sensor_z = float(base.translation.z) + sensor_z = float(base.transform.translation.z) ref_timing = _process_frame( ray_obs, planners, diff --git a/dimos/navigation/global_planner/mls_planner/utils/test_plan_rrd.py b/dimos/navigation/global_planner/mls_planner/utils/test_plan_rrd.py index de962eb5af0..1bc0714464b 100644 --- a/dimos/navigation/global_planner/mls_planner/utils/test_plan_rrd.py +++ b/dimos/navigation/global_planner/mls_planner/utils/test_plan_rrd.py @@ -12,6 +12,70 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Recorded generated TF messages feed the offline planner viewer.""" + +from unittest.mock import patch + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage + +from dimos.memory.type.observation import Observation +from dimos.navigation.global_planner.mls_planner.utils.plan_rrd import _log_odometry, _TfSync +from dimos.visualization.rerun.message_helpers import register_colormap_annotation + + +def test_tf_sync_generated_edges_and_odometry() -> None: + edge = TransformStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + child_frame_id="base_link", + transform=Transform( + translation=Vector3(x=2, y=0.0, z=0.0), rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), + ) + message = cdr_decode(cdr_encode(TFMessage(transforms=[edge])), TFMessage) + sync = _TfSync([Observation(id=0, ts=1.0, _data=message)]) + with patch("rerun.log") as log: + sync.up_to(0.5) + log.assert_not_called() + sync.up_to(1.0) + entity, archetype = log.call_args.args + assert entity == "tf_links/base_link" + assert archetype.parent_frame.as_arrow_array().to_pylist() == ["tf#/world"] + assert archetype.child_frame.as_arrow_array().to_pylist() == ["tf#/base_link"] + assert archetype.translation.as_arrow_array().to_pylist() == [[2.0, 0.0, 0.0]] + sync.up_to(2.0) + assert log.call_count == 1 + _log_odometry((2, 0, 0, 0, 0, 0, 1), 1.0, [], edge) + assert log.call_args.args[0] == "world/robot_body" + assert log.call_args.args[1].translation.as_arrow_array().to_pylist() == [[2.0, 0.0, 0.0]] + + +def test_colormap_annotation_has_all_indices() -> None: + with patch("rerun.log") as log: + register_colormap_annotation() + assert log.call_args.args[0] == "/" + assert log.call_args.kwargs == {"static": True} + context = log.call_args.args[1].context.as_arrow_array().to_pylist() + assert len(context[0]) == 256 + + +# 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. + import numpy as np import pytest from pytest_mock import MockerFixture diff --git a/dimos/navigation/global_planner/mls_planner/viz.py b/dimos/navigation/global_planner/mls_planner/viz.py index 5174a621fbf..3035d478236 100644 --- a/dimos/navigation/global_planner/mls_planner/viz.py +++ b/dimos/navigation/global_planner/mls_planner/viz.py @@ -20,11 +20,19 @@ from functools import partial from typing import TYPE_CHECKING +from dimos_generated.dimos_msgs.msg import ( + LineSegment3D, + LineSegments3D, + RegionLineSegments3D, + RegionPointCloud2, +) +from dimos_generated.geometry_msgs.msg import Point +from dimos_generated.sensor_msgs.msg import PointCloud2 import numpy as np -from dimos.msgs.nav_msgs.LineSegments3D import LineSegments3D -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.visualization.rerun.bridge import RerunEntry, keyed_by_seq, region_entity +from dimos.msgs.pointcloud import pointcloud_view, pointcloud_xyz +from dimos.msgs.time import header_now +from dimos.visualization.rerun.bridge import RerunEntry, keyed_by_region, region_entity if TYPE_CHECKING: from numpy.typing import NDArray @@ -72,9 +80,9 @@ def surface_points( ) -@keyed_by_seq +@keyed_by_region def render_surface_region( - msg: PointCloud2, + msg: RegionPointCloud2, voxel_size: float, wall_clearance_m: float, clearance_clamp_m: float, @@ -85,19 +93,21 @@ def render_surface_region( Clearance rides the intensity channel, flat blue without it. Points outside z_band (world frame) are dropped, to show one storey of a multi-storey map. """ - seq = msg.seq - pts = msg.points_f32() - clearance = msg.intensities_f32() + pts = pointcloud_xyz(msg.cloud).astype(np.float32) + records = pointcloud_view(msg.cloud) + clearance = records["intensity"].ravel() if "intensity" in (records.dtype.names or ()) else None + if z_band is not None: + keep = (pts[:, 2] >= z_band[0]) & (pts[:, 2] <= z_band[1]) + pts = pts[keep] + if clearance is not None: + clearance = clearance[keep] if clearance is None or len(clearance) != len(pts): - if z_band is not None: - msg = msg.filter_by_height(min_height=z_band[0], max_height=z_band[1]) - cell = msg.to_rerun(voxel_size=voxel_size, colors=[40, 75, 130]) + import rerun as rr + + cell: Archetype = rr.Points3D(pts, radii=voxel_size * 0.5, colors=[40, 75, 130]) else: - if z_band is not None: - keep = (pts[:, 2] >= z_band[0]) & (pts[:, 2] <= z_band[1]) - pts, clearance = pts[keep], clearance[keep] cell = surface_points(pts, clearance, voxel_size, wall_clearance_m, clearance_clamp_m) - return [RerunEntry(region_entity(SURFACE_MAP_ENTITY, seq), cell, static=True)] + return [RerunEntry(region_entity(SURFACE_MAP_ENTITY, msg.region_id), cell, static=True)] def graph_nodes(pts: NDArray[np.float32]) -> Archetype: @@ -111,30 +121,75 @@ def graph_nodes(pts: NDArray[np.float32]) -> Archetype: def render_nodes(msg: PointCloud2) -> Archetype: - return graph_nodes(msg.points_f32()) + return graph_nodes(pointcloud_xyz(msg).astype(np.float32)) def graph_edges(edges: NDArray[np.float32]) -> Archetype: """Edges as ``[x0, y0, z0, x1, y1, z1, cost]`` rows, colored green to red by cost.""" segments = LineSegments3D( - segments=edges[:, :6].reshape(-1, 2, 3) if len(edges) else None, - weights=edges[:, 6] if len(edges) else None, + header=header_now(), + segments=[ + LineSegment3D( + start=Point(x=float(row[0]), y=float(row[1]), z=float(row[2])), + end=Point(x=float(row[3]), y=float(row[4]), z=float(row[5])), + weight=float(row[6]), + ) + for row in edges + ], ) return render_node_edges(segments) def render_node_edges(msg: LineSegments3D) -> Archetype: - return msg.to_rerun(z_offset=GRAPH_Z_LIFT, radii=0.01) - + import rerun as rr -@keyed_by_seq -def render_edge_region(msg: LineSegments3D) -> RerunMulti: - """One cell of the edge corridors on its own static entity.""" + if not msg.segments: + return rr.LineStrips3D([]) + strips = np.array( + [ + [ + [segment.start.x, segment.start.y, segment.start.z], + [segment.end.x, segment.end.y, segment.end.z], + ] + for segment in msg.segments + ], + dtype=np.float32, + ) + strips[:, :, 2] += GRAPH_Z_LIFT + weights = np.array([segment.weight for segment in msg.segments], dtype=np.float64) + log_weights = np.log10(np.maximum(weights, 1e-6)) + low, high = float(log_weights.min()), float(log_weights.max()) + norm = (log_weights - low) / (high - low) if high > low else np.zeros_like(log_weights) + red = (255 * norm).astype(np.uint8) + green = (255 * (1 - norm)).astype(np.uint8) + colors = np.column_stack([red, green, np.full_like(red, 60), np.full_like(red, 220)]) + return rr.LineStrips3D(strips, colors=colors, radii=0.01) + + +@keyed_by_region +def render_edge_region(msg: RegionLineSegments3D) -> RerunMulti: return [ - RerunEntry(region_entity(NODE_EDGES_ENTITY, msg.seq), render_node_edges(msg), static=True) + RerunEntry( + region_entity(NODE_EDGES_ENTITY, msg.region_id), + render_node_edges(msg.lines), + static=True, + ) ] +def render_surface_map( + msg: PointCloud2, + voxel_size: float = 0.1, + wall_clearance_m: float = 0.0, + clearance_clamp_m: float = 1.0, +) -> Archetype: + entry = render_surface_region( + RegionPointCloud2(region_id=0, cloud=msg), voxel_size, wall_clearance_m, clearance_clamp_m + )[0] + assert isinstance(entry, RerunEntry) + return entry.archetype + + def planner_visual_override( viz_publish_hz: float, voxel_size: float, diff --git a/dimos/navigation/global_planner/test_viz.py b/dimos/navigation/global_planner/test_viz.py index 1deb9ba4dcb..2761540e7ea 100644 --- a/dimos/navigation/global_planner/test_viz.py +++ b/dimos/navigation/global_planner/test_viz.py @@ -15,29 +15,67 @@ from functools import partial import pickle +from dimos_generated.dimos_msgs.msg import RegionPointCloud2 +from dimos_generated.geometry_msgs.msg import Point, PointStamped, Pose, PoseStamped, Quaternion +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Header import numpy as np import rerun as rr from dimos.mapping.ray_tracing.viz import render_map_region -from dimos.msgs.geometry_msgs.PointStamped import PointStamped -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import time_from_seconds from dimos.navigation.global_planner import viz from dimos.navigation.global_planner.mls_planner.viz import render_surface_region -from dimos.visualization.rerun.bridge import is_keyed_by_seq +from dimos.visualization.rerun.bridge import is_keyed_by_region def test_goal_placeholder_is_not_drawn() -> None: - assert viz.render_goal(PointStamped(x=float("nan"), y=0.0, z=0.0)) is None - goal = viz.render_goal(PointStamped(x=1.0, y=2.0, z=0.0)) + assert ( + viz.render_goal( + PointStamped( + header=Header(stamp=time_from_seconds(0.0), frame_id=""), + point=Point(x=float("nan"), y=0.0, z=0.0), + ) + ) + is None + ) + goal = viz.render_goal( + PointStamped( + header=Header(stamp=time_from_seconds(0.0), frame_id=""), + point=Point(x=1.0, y=2.0, z=0.0), + ) + ) assert goal is not None assert goal.positions.as_arrow_array().to_pylist() == [[1.0, 2.0, 0.0]] def test_empty_path_keeps_the_last_one_drawn() -> None: - assert viz.render_path(Path(poses=[])) is None - path = viz.render_path(Path(poses=[PoseStamped(1.0, 0.0, 0.0), PoseStamped(2.0, 0.0, 0.0)])) + assert ( + viz.render_path(Path(header=Header(stamp=time_from_seconds(0.0), frame_id=""), poses=[])) + is None + ) + path = viz.render_path( + Path( + header=Header(stamp=time_from_seconds(0.0), frame_id=""), + poses=[ + PoseStamped( + header=Header(stamp=time_from_seconds(0.0), frame_id=""), + pose=Pose( + position=Point(x=1.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + PoseStamped( + header=Header(stamp=time_from_seconds(0.0), frame_id=""), + pose=Pose( + position=Point(x=2.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + ], + ) + ) assert path is not None lift = viz.PATH_Z_LIFT np.testing.assert_allclose( @@ -48,9 +86,9 @@ def test_empty_path_keeps_the_last_one_drawn() -> None: def test_keyed_renderers_keep_their_flag_through_the_pickle_round_trip() -> None: plain = pickle.loads(pickle.dumps(render_map_region)) bound = pickle.loads(pickle.dumps(partial(render_map_region, voxel_size=0.1))) - assert is_keyed_by_seq(plain) and is_keyed_by_seq(bound) - assert not is_keyed_by_seq(viz.render_path) - assert not is_keyed_by_seq(pickle.loads(pickle.dumps(partial(viz.render_path)))) + assert is_keyed_by_region(plain) and is_keyed_by_region(bound) + assert not is_keyed_by_region(viz.render_path) + assert not is_keyed_by_region(pickle.loads(pickle.dumps(partial(viz.render_path)))) def test_bridge_config_pickles_for_the_workers() -> None: @@ -63,7 +101,10 @@ def test_bridge_config_pickles_for_the_workers() -> None: regions = overrides["world/map_regions"] assert regions.func is render_map_region assert regions.keywords == {"voxel_size": 0.08, "height_range": viz.HEIGHT_RANGE} - region = PointCloud2.from_numpy(np.array([[1.0, 1.0, 0.0]], dtype=np.float32)) - region.seq = (4 << 16) | 7 + region = pointcloud_from_xyz( + np.array([[1.0, 1.0, 0.0]], dtype=np.float32), + header=Header(stamp=time_from_seconds(0.0), frame_id=""), + ) + region = RegionPointCloud2(region_id=(4 << 16) | 7, cloud=region) (path, _, static_entry) = regions(region)[0] assert path == "world/map_regions/4_7" and static_entry diff --git a/dimos/navigation/global_planner/viz.py b/dimos/navigation/global_planner/viz.py index 21f720bbd31..4699bf8e3c1 100644 --- a/dimos/navigation/global_planner/viz.py +++ b/dimos/navigation/global_planner/viz.py @@ -29,11 +29,11 @@ from dimos.navigation.global_planner.mls_planner.viz import planner_visual_override if TYPE_CHECKING: + from dimos_generated.geometry_msgs.msg import PointStamped + from dimos_generated.nav_msgs.msg import Path from numpy.typing import NDArray from rerun._baseclasses import Archetype - from dimos.msgs.geometry_msgs.PointStamped import PointStamped - from dimos.msgs.nav_msgs.Path import Path from dimos.visualization.rerun.bridge import VisualOverride # rerun is imported inside the renderers: it is heavy and only viewers load it. @@ -66,7 +66,12 @@ def render_path(msg: Path) -> Archetype | None: """The planned route. An empty path means no route, so the last one stays drawn.""" if len(msg.poses) == 0: return None - return path_strip(np.array([[p.x, p.y, p.z] for p in msg.poses], dtype=np.float32)) + return path_strip( + np.array( + [[p.pose.position.x, p.pose.position.y, p.pose.position.z] for p in msg.poses], + dtype=np.float32, + ) + ) def goal_point(xyz: tuple[float, float, float]) -> Archetype: @@ -77,9 +82,9 @@ def goal_point(xyz: tuple[float, float, float]) -> Archetype: def render_goal(msg: PointStamped) -> Archetype | None: """The active goal. NaN is the movement manager's placeholder for none.""" - if any(math.isnan(v) for v in (msg.x, msg.y, msg.z)): + if any(math.isnan(v) for v in (msg.point.x, msg.point.y, msg.point.z)): return None - return goal_point((msg.x, msg.y, msg.z)) + return goal_point((msg.point.x, msg.point.y, msg.point.z)) def robot_body_box(length: float, width: float, height: float) -> Archetype: diff --git a/dimos/navigation/go2/loop_closure/eval.py b/dimos/navigation/go2/loop_closure/eval.py index 99836f335cb..485d24314e3 100644 --- a/dimos/navigation/go2/loop_closure/eval.py +++ b/dimos/navigation/go2/loop_closure/eval.py @@ -31,17 +31,19 @@ import time from typing import Any +from dimos_generated.sensor_msgs.msg import Image +import numpy as np import typer from dimos.memory.store.sqlite import SqliteStore from dimos.memory.stream import Stream from dimos.memory.transform import QualityWindow, SpeedLimit from dimos.memory.type.observation import Observation -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.geometry import transform_matrix +from dimos.msgs.image import image_sharpness from dimos.navigation.go2.loop_closure.pgo import PGO from dimos.perception.fiducial.marker_transformer import DetectMarkers -from dimos.robot.unitree.go2.connection import _camera_info_static +from dimos.robot.unitree.go2.camera_calibration import front_camera_calibration from dimos.utils.data import get_data DEFAULT_DATASETS = [f"hk_village{i}" for i in range(1, 7)] @@ -68,7 +70,7 @@ def _eval_recording( ) -> tuple[float, float]: """Returns (pgo_time_s, spread_m) for one recording.""" db_path = get_data(f"{name}.db") - cam_info = _camera_info_static() + cam_info = front_camera_calibration() store = SqliteStore(path=str(db_path)) with store: @@ -86,7 +88,7 @@ def _eval_recording( smoothing_window=marker_smoothing, ) pipeline: Stream[Image] = color_image.transform( - QualityWindow(lambda img: img.sharpness, window=marker_quality_window) + QualityWindow(image_sharpness, window=marker_quality_window) ) if marker_max_speed > 0: pipeline = pipeline.transform( @@ -106,16 +108,13 @@ def _eval_recording( # PGO-correct each track's pose; group by marker_id. by_marker: dict[int, list[tuple[float, float, float]]] = {} for d in tracks: - raw_tf = Transform( - translation=d.data.center, - rotation=d.data.orientation, - frame_id="world", - child_frame_id=f"marker_{d.data.marker_id}", - ts=d.ts, + center = d.data.center + corrected = transform_matrix(graph.correction_at(d.ts).transform) @ np.array( + [center.x, center.y, center.z, 1.0] + ) + by_marker.setdefault(d.data.marker_id, []).append( + (float(corrected[0]), float(corrected[1]), float(corrected[2])) ) - corrected = graph.correct(raw_tf) - t = corrected.translation - by_marker.setdefault(d.data.marker_id, []).append((t.x, t.y, t.z)) spread = sum(_pairwise_sum(v) for v in by_marker.values()) return pgo_time, spread diff --git a/dimos/navigation/go2/loop_closure/pgo.py b/dimos/navigation/go2/loop_closure/pgo.py index b50fd60e4d0..7694b79f21d 100644 --- a/dimos/navigation/go2/loop_closure/pgo.py +++ b/dimos/navigation/go2/loop_closure/pgo.py @@ -48,16 +48,41 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal, TypedDict, TypeVar, cast +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header import numpy as np from scipy.spatial.transform import Rotation, Slerp from dimos.memory.transform import Transformer from dimos.memory.type.observation import Observation -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.geometry import ( + compose_transforms, + inverse_transform, + pose_from_matrix, + pose_matrix, + quaternion_from_matrix, + transform_from_matrix, + transform_matrix, +) +from dimos.msgs.pointcloud import ( + concatenate_clouds, + pointcloud_from_xyz, + pointcloud_stamps, + pointcloud_to_open3d, + pointcloud_xyz, + transform_cloud, + voxel_downsample_cloud, +) +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.protocol.service.spec import BaseConfig from dimos.utils.logging_config import setup_logger @@ -137,8 +162,8 @@ class Keyframe: """ ts: float - local: Transform - optimized: Transform + local: TransformStamped + optimized: TransformStamped @dataclass(frozen=True) @@ -151,8 +176,8 @@ class LoopClosure: tighter match. """ - source: Transform - target: Transform + source: TransformStamped + target: TransformStamped score: float @@ -168,48 +193,49 @@ class PoseGraph(Transformer[Any, Any]): keyframes: tuple[Keyframe, ...] = () loops: tuple[LoopClosure, ...] = () - def correct(self, pose: Transform) -> Transform: - """Map a raw-world pose into the drift-corrected frame at ``pose.ts``.""" - # Transform.__add__ composes: (T_corr + T_raw) applies T_corr after T_raw. - return self.correction_at(pose.ts) + pose + def correct(self, pose: TransformStamped) -> TransformStamped: + """Map a raw-world stamped transform into the drift-corrected frame.""" + result = compose_transforms(self.correction_at(to_seconds(pose.header.stamp)), pose) + result.header.stamp = pose.header.stamp + return result - def correction_at(self, ts: float) -> Transform: + def correction_at(self, ts: float) -> TransformStamped: """The raw ``world_corrected <- world_raw`` transform at ``ts``. Useful when applying the correction to non-pose data (e.g. point - clouds): ``pointcloud.transform(graph.correction_at(obs.ts))``. + clouds): ``transform_cloud(pointcloud, graph.correction_at(obs.ts))``. """ R, t = self._corrections(np.array([ts])) - return Transform( - translation=Vector3(t[0]), - rotation=Quaternion.from_rotation_matrix(R[0].as_matrix()), - frame_id=FRAME_WORLD_CORRECTED, + return TransformStamped( + header=Header(frame_id=FRAME_WORLD_CORRECTED, stamp=time_from_seconds(float(ts))), child_frame_id=FRAME_WORLD_RAW, - ts=float(ts), + transform=Transform( + translation=Vector3(x=t[0, 0], y=t[0, 1], z=t[0, 2]), + rotation=quaternion_from_matrix(R[0].as_matrix()), + ), ) def place(self, cloud: PointCloud2) -> PointCloud2: """Move each point by the correction for its capture time, anchored at the latest keyframe. Capture time is the cloud's per-point ``stamps`` when present, else - ``cloud.ts``. The latest keyframe's correction maps to identity, so + the cloud header timestamp. The latest keyframe's correction maps to identity, so current data stays where raw odometry puts it and older data moves. """ - pts = cloud.points_f32() + pts = pointcloud_xyz(cloud).astype(np.float32) if not len(pts): return cloud - stamps = cloud.stamps_f64() + stamps = pointcloud_stamps(cloud) if stamps is None: - stamps = np.full(len(pts), cloud.ts) + stamps = np.full(len(pts), to_seconds(cloud.header.stamp)) unique, inverse = np.unique(stamps, return_inverse=True) R, t = self._corrections(unique) now_R, now_t = self._corrections(np.array([np.inf])) anchor = now_R[0].inv() R, t = anchor * R, anchor.apply(t - now_t[0]) - return PointCloud2.from_numpy( + return pointcloud_from_xyz( R[inverse].apply(pts.astype(np.float64)) + t[inverse], - frame_id=cloud.frame_id, - timestamp=cloud.ts, + header=cloud.header, ) def __call__(self, upstream: Iterator[Observation[Any]]) -> Iterator[Observation[Any]]: @@ -219,8 +245,14 @@ def __call__(self, upstream: Iterator[Observation[Any]]) -> Iterator[Observation if ps is None: yield obs continue - raw_tf = Transform.from_pose(FRAME_BODY, ps) - yield obs.derive(data=obs.data, pose=self.correct(raw_tf)) + correction = self.correction_at(obs.ts) + corrected = PoseStamped( + header=Header(stamp=ps.header.stamp, frame_id=FRAME_WORLD_CORRECTED), + pose=pose_from_matrix( + transform_matrix(correction.transform) @ pose_matrix(ps.pose) + ), + ) + yield obs.derive(data=obs.data, pose=corrected) def _corrections(self, ts: np.ndarray) -> tuple[Rotation, np.ndarray]: """Vectorized drift correction: slerp rotation, lerp translation, clipped to the keyframes.""" @@ -240,9 +272,19 @@ def _drifts(self) -> tuple[np.ndarray, Slerp, np.ndarray]: # Per-keyframe drift: world_corrected <- body <- world_raw. ts_list = [kf.ts for kf in self.keyframes] - drifts = [(kf.optimized + kf.local.inverse()) for kf in self.keyframes] - quat_list = [tf.rotation.to_numpy() for tf in drifts] - t_list = [tf.translation.to_numpy() for tf in drifts] + drifts = [ + compose_transforms(kf.optimized, inverse_transform(kf.local)) for kf in self.keyframes + ] + quat_list = [ + [ + tf.transform.rotation.x, + tf.transform.rotation.y, + tf.transform.rotation.z, + tf.transform.rotation.w, + ] + for tf in drifts + ] + t_list = [transform_matrix(tf.transform)[:3, 3] for tf in drifts] # Slerp needs ≥2 keyframes; pad a len==1 list with a near-duplicate. if len(ts_list) == 1: @@ -290,7 +332,12 @@ def __call__( continue # Placeholder filter: zero translation OR uninitialized (all-zero) # quaternion. Identity rotation (qw=1) is valid and stays. - if pose.position.is_zero() or pose.orientation.is_zero(): + if np.allclose([pose.position.x, pose.position.y, pose.position.z], 0.0) or ( + pose.orientation.x == 0.0 + and pose.orientation.y == 0.0 + and pose.orientation.z == 0.0 + and pose.orientation.w == 0.0 + ): continue pgo.process(_obs_to_pose3(obs), obs.ts, obs.data) @@ -404,20 +451,23 @@ def process( ts: float, world_cloud: PointCloud2, ) -> None: - if len(world_cloud) == 0: + if world_cloud.width * world_cloud.height == 0: return if not self._is_keyframe(local_pose): return # Unregister: lift world-frame scan back into body frame using its # odom pose, so PGO can re-project it via the optimized pose later. - body_cloud = world_cloud.transform( - _pose3_to_transform( - local_pose.inverse(), - ts=ts, - frame_id=FRAME_BODY, - child_frame_id=FRAME_WORLD_RAW, - ) - ).voxel_downsample(self._cfg.submap_resolution) + body_cloud = voxel_downsample_cloud( + transform_cloud( + world_cloud, + TransformStamped( + header=Header(frame_id=FRAME_BODY, stamp=world_cloud.header.stamp), + child_frame_id=world_cloud.header.frame_id, + transform=transform_from_matrix(local_pose.inverse().matrix()), + ), + ), + self._cfg.submap_resolution, + ) self._add_keyframe(local_pose, ts, body_cloud) self._search_for_loops() self._smooth_and_update() @@ -539,26 +589,25 @@ def _get_submap(self, idx: int, half_range: int) -> PointCloud2: lo = max(0, idx - half_range) hi = min(len(self._key_poses) - 1, idx + half_range) if lo > hi: - return PointCloud2() - cloud = self._key_poses[lo].body_cloud.transform( - _pose3_to_transform( - self._key_poses[lo].optimized, - ts=self._key_poses[lo].timestamp, - frame_id=FRAME_WORLD_CORRECTED, - child_frame_id=FRAME_BODY, + return pointcloud_from_xyz( + np.empty((0, 3)), + header=Header(frame_id=FRAME_WORLD_CORRECTED, stamp=Time(sec=0, nanosec=0)), ) - ) - for i in range(lo + 1, hi + 1): - kp = self._key_poses[i] - cloud = cloud + kp.body_cloud.transform( - _pose3_to_transform( - kp.optimized, - ts=kp.timestamp, - frame_id=FRAME_WORLD_CORRECTED, + + def registered(kp: _KeyPose) -> PointCloud2: + return transform_cloud( + kp.body_cloud, + TransformStamped( + header=Header(frame_id=FRAME_WORLD_CORRECTED, stamp=kp.body_cloud.header.stamp), child_frame_id=FRAME_BODY, - ) + transform=transform_from_matrix(kp.optimized.matrix()), + ), ) - return cloud.voxel_downsample(self._cfg.submap_resolution) + + cloud = registered(self._key_poses[lo]) + for i in range(lo + 1, hi + 1): + cloud = concatenate_clouds(cloud, registered(self._key_poses[i])) + return voxel_downsample_cloud(cloud, self._cfg.submap_resolution) def _search_for_loops(self) -> None: if len(self._key_poses) < self._cfg.min_keyframes_for_loop_search: @@ -682,23 +731,24 @@ def _pose3_to_transform( ts: float, frame_id: str, child_frame_id: str, -) -> Transform: +) -> TransformStamped: """PGO-internal: build a Transform from a Pose3.""" t = np.asarray(pose.translation()) - return Transform( - translation=Vector3(float(t[0]), float(t[1]), float(t[2])), - rotation=Quaternion.from_rotation_matrix(pose.rotation().matrix()), - frame_id=frame_id, + return TransformStamped( + header=Header(frame_id=frame_id, stamp=time_from_seconds(ts)), child_frame_id=child_frame_id, - ts=ts, + transform=Transform( + translation=Vector3(x=float(t[0]), y=float(t[1]), z=float(t[2])), + rotation=quaternion_from_matrix(pose.rotation().matrix()), + ), ) -def _transform_to_pose3(tf: Transform) -> gtsam.Pose3: +def _transform_to_pose3(tf: TransformStamped) -> gtsam.Pose3: """PGO-internal: build a Pose3 from a Transform.""" import gtsam # type: ignore[import-not-found,import-untyped] - return gtsam.Pose3(tf.to_matrix()) + return gtsam.Pose3(transform_matrix(tf.transform)) def _icp( @@ -708,8 +758,8 @@ def _icp( max_dist: float = 1.0, tol: float = 1e-6, min_inliers: int = 10, - init: Transform | None = None, -) -> tuple[Transform, float]: + init: TransformStamped | None = None, +) -> tuple[TransformStamped, float]: """Point-to-plane ICP using Open3D's tensor pipeline. Returns ``(tf, fitness)`` where ``fitness`` is mean squared inlier @@ -720,11 +770,18 @@ def _icp( import open3d as o3d # type: ignore[import-untyped] import open3d.core as o3c # type: ignore[import-untyped] - if len(source) < min_inliers or len(target) < min_inliers: - return Transform.identity(), float("inf") + if source.width * source.height < min_inliers or target.width * target.height < min_inliers: + return TransformStamped( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + child_frame_id="", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), float("inf") - src_pcd = source.pointcloud_tensor - tgt_pcd = target.pointcloud_tensor + src_pcd = o3d.t.geometry.PointCloud.from_legacy(pointcloud_to_open3d(source)) + tgt_pcd = o3d.t.geometry.PointCloud.from_legacy(pointcloud_to_open3d(target)) # Normals on the target enable point-to-plane ICP — converges tighter # than point-to-point on indoor scenes (walls give unambiguous normals @@ -733,7 +790,7 @@ def _icp( device = src_pcd.device init_T = o3c.Tensor( - init.to_matrix() if init is not None else np.eye(4), + transform_matrix(init.transform) if init is not None else np.eye(4), dtype=o3c.float64, device=device, ) @@ -755,9 +812,20 @@ def _icp( ) if float(result.fitness) == 0.0: - return Transform.identity(), float("inf") + return TransformStamped( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + child_frame_id="", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), float("inf") - # Frames intentionally unlabeled: caller only reads .to_matrix(). - tf = Transform.from_matrix(result.transformation.numpy(), ts=source.ts) + # Frames intentionally unlabeled: caller only reads the transform matrix. + tf = TransformStamped( + header=Header(stamp=source.header.stamp, frame_id=""), + transform=transform_from_matrix(result.transformation.numpy()), + child_frame_id="", + ) rmse = float(result.inlier_rmse) return tf, rmse * rmse diff --git a/dimos/navigation/go2/loop_closure/pgo_auto.py b/dimos/navigation/go2/loop_closure/pgo_auto.py index 9e172c1f6f2..735bffccfbd 100644 --- a/dimos/navigation/go2/loop_closure/pgo_auto.py +++ b/dimos/navigation/go2/loop_closure/pgo_auto.py @@ -58,6 +58,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any, TypedDict, TypeVar +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Transform, TransformStamped, Vector3 from dimos_generated.sensor_msgs.msg import PointCloud2 from dimos_generated.std_msgs.msg import Header @@ -584,7 +585,8 @@ def _get_submap(self, idx: int, half_range: int) -> PointCloud2: hi = min(len(self._key_poses) - 1, idx + half_range) if lo > hi: return pointcloud_from_xyz( - np.empty((0, 3)), header=Header(frame_id=FRAME_WORLD_CORRECTED) + np.empty((0, 3)), + header=Header(frame_id=FRAME_WORLD_CORRECTED, stamp=Time(sec=0, nanosec=0)), ) def registered(kp: _KeyPose) -> PointCloud2: diff --git a/dimos/navigation/go2/loop_closure/test_pgo.py b/dimos/navigation/go2/loop_closure/test_pgo.py index 9f70d4625e3..1dba4848843 100644 --- a/dimos/navigation/go2/loop_closure/test_pgo.py +++ b/dimos/navigation/go2/loop_closure/test_pgo.py @@ -14,16 +14,27 @@ from __future__ import annotations +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from scipy.spatial.transform import Rotation from dimos.memory.store.memory import MemoryStore from dimos.memory.stream import Stream -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.geometry import quaternion_from_matrix, transform_matrix +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.navigation.go2.loop_closure.pgo import ( PGO, Keyframe, @@ -34,6 +45,27 @@ _pose3_to_transform, ) + +def _vector(*values): + values = values[0] if len(values) == 1 else values + return Vector3(x=float(values[0]), y=float(values[1]), z=float(values[2])) + + +def _quaternion(x, y, z, w): + return Quaternion(x=x, y=y, z=z, w=w) + + +def _stamped_transform(translation=None, rotation=None, ts=0.0, frame_id="", child_frame_id=""): + return TransformStamped( + header=Header(frame_id=frame_id, stamp=time_from_seconds(ts)), + child_frame_id=child_frame_id, + transform=Transform( + translation=translation if translation is not None else Vector3(x=0.0, y=0.0, z=0.0), + rotation=rotation if rotation is not None else Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), + ), + ) + + # TODO(PY311): drop — the mapping extra excludes gtsam-extended where it has no # wheels (py3.10 Linux), see pyproject. pytest.importorskip("gtsam") @@ -76,9 +108,9 @@ def test_observation_normalizes_transform_pose(self) -> None: """Constructing/deriving with pose=Transform should coerce to 7-tuple.""" from dimos.memory.type.observation import Observation - tf = Transform( - translation=Vector3(1.5, -2.0, 0.7), - rotation=Quaternion(0.1, 0.2, 0.3, 0.927), + tf = _stamped_transform( + translation=_vector(1.5, -2.0, 0.7), + rotation=_quaternion(0.1, 0.2, 0.3, 0.927), ts=1.0, ) obs: Observation[int] = Observation(id=0, ts=1.0, pose=tf, _data=0) @@ -92,9 +124,14 @@ def test_observation_normalizes_transform_pose(self) -> None: def test_observation_normalizes_posestamped(self) -> None: from dimos.memory.type.observation import Observation - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - ps = PoseStamped(ts=1.0, position=(1.0, 2.0, 3.0), orientation=(0.0, 0.0, 0.0, 1.0)) + ps = PoseStamped( + header=Header(stamp=time_from_seconds(1.0), frame_id=""), + pose=Pose( + position=Point(x=1.0, y=2.0, z=3.0), + orientation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), + ), + ) obs: Observation[int] = Observation(id=0, ts=1.0, pose=ps, _data=0) assert obs.pose_tuple == (1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 1.0) @@ -104,7 +141,7 @@ def test_obs_to_pose3_roundtrip(self) -> None: rng = np.random.default_rng(4) R = _random_R(rng) t = rng.uniform(-3, 3, size=3) - tf = Transform(translation=Vector3(t), rotation=Quaternion.from_rotation_matrix(R), ts=1.0) + tf = _stamped_transform(translation=_vector(t), rotation=quaternion_from_matrix(R), ts=1.0) obs: Observation[int] = Observation(id=0, ts=1.0, pose=tf, _data=0) p = _obs_to_pose3(obs) np.testing.assert_allclose(p.rotation().matrix(), R, atol=1e-9) @@ -118,8 +155,8 @@ def test_pose3_to_transform(self) -> None: t = rng.uniform(-3, 3, size=3) p = gtsam.Pose3(gtsam.Rot3(R), gtsam.Point3(t)) tf = _pose3_to_transform(p, ts=7.89, frame_id="world", child_frame_id="body") - np.testing.assert_allclose(tf.rotation.to_rotation_matrix(), R, atol=1e-10) - np.testing.assert_allclose(tf.translation.to_numpy(), t, atol=1e-10) + np.testing.assert_allclose(transform_matrix(tf.transform)[:3, :3], R, atol=1e-10) + np.testing.assert_allclose(transform_matrix(tf.transform)[:3, 3], t, atol=1e-10) def test_pose3_to_transform_with_frames(self) -> None: import gtsam @@ -129,10 +166,10 @@ def test_pose3_to_transform_with_frames(self) -> None: t = rng.uniform(-3, 3, size=3) p = gtsam.Pose3(gtsam.Rot3(R), gtsam.Point3(t)) tf = _pose3_to_transform(p, ts=1.0, frame_id="world_corrected", child_frame_id="body") - assert tf.frame_id == "world_corrected" + assert tf.header.frame_id == "world_corrected" assert tf.child_frame_id == "body" - np.testing.assert_allclose(tf.rotation.to_rotation_matrix(), R, atol=1e-10) - np.testing.assert_allclose(tf.translation.to_numpy(), t, atol=1e-10) + np.testing.assert_allclose(transform_matrix(tf.transform)[:3, :3], R, atol=1e-10) + np.testing.assert_allclose(transform_matrix(tf.transform)[:3, 3], t, atol=1e-10) def _make_lidar_stream(n_frames: int = 12, points_per_frame: int = 500) -> Stream[PointCloud2]: @@ -153,7 +190,10 @@ def _make_lidar_stream(n_frames: int = 12, points_per_frame: int = 500) -> Strea body = rng.uniform(-1, 1, size=(points_per_frame, 3)).astype(np.float32) world = (R_world @ body.T).T + np.array([i, 0, 0], dtype=np.float32) lidar.append( - PointCloud2.from_numpy(world.astype(np.float32)), + pointcloud_from_xyz( + world.astype(np.float32), + header=Header(frame_id="world_raw", stamp=time_from_seconds(float(i))), + ), ts=float(i), pose=(float(i), 0.0, 0.0, qx, qy, qz, qw), ) @@ -200,19 +240,31 @@ def test_apply_identity_corrections_preserves_poses(self) -> None: assert a == pytest.approx(b, abs=1e-6) -def _graph_with_drift_at(drifts: list[Transform]) -> PoseGraph: +def _graph_with_drift_at(drifts: list[TransformStamped]) -> PoseGraph: """PoseGraph whose drift correction equals each ``drifts[i]`` at ``drifts[i].ts``. Trick: drift = optimized + local^-1. With local=identity, drift==optimized. """ - identity = Vector3(0.0, 0.0, 0.0) - identity_rot = Quaternion(0.0, 0.0, 0.0, 1.0) + identity = _vector(0.0, 0.0, 0.0) + identity_rot = _quaternion(0.0, 0.0, 0.0, 1.0) return PoseGraph( keyframes=tuple( Keyframe( - ts=d.ts, - local=Transform(translation=identity, rotation=identity_rot, ts=d.ts), - optimized=Transform(translation=d.translation, rotation=d.rotation, ts=d.ts), + ts=to_seconds(d.header.stamp), + local=_stamped_transform( + translation=identity, + rotation=identity_rot, + ts=to_seconds(d.header.stamp), + frame_id="world_raw", + child_frame_id="body", + ), + optimized=_stamped_transform( + translation=d.transform.translation, + rotation=d.transform.rotation, + ts=to_seconds(d.header.stamp), + frame_id="world_corrected", + child_frame_id="body", + ), ) for d in drifts ) @@ -226,67 +278,77 @@ def test_empty_raises(self) -> None: def test_single_keyframe_returns_constant(self) -> None: R = Rotation.from_euler("z", np.pi / 4).as_matrix() - only = Transform( - translation=Vector3(1.0, 2.0, 3.0), - rotation=Quaternion.from_rotation_matrix(R), + only = _stamped_transform( + translation=_vector(1.0, 2.0, 3.0), + rotation=quaternion_from_matrix(R), ts=10.0, ) graph = _graph_with_drift_at([only]) for query_ts in (0.0, 10.0, 100.0): out = graph.correction_at(query_ts) - assert out.translation.x == pytest.approx(1.0, abs=1e-10) - assert out.translation.y == pytest.approx(2.0, abs=1e-10) - assert out.translation.z == pytest.approx(3.0, abs=1e-10) + assert out.transform.translation.x == pytest.approx(1.0, abs=1e-10) + assert out.transform.translation.y == pytest.approx(2.0, abs=1e-10) + assert out.transform.translation.z == pytest.approx(3.0, abs=1e-10) def test_out_of_range_clips_to_endpoints(self) -> None: - # Transform's ctor maps ts=0.0 -> time.time(); use ts>0 for determinism. - a = Transform(translation=Vector3(0.0, 0.0, 0.0), ts=1.0) - b = Transform(translation=Vector3(10.0, 0.0, 0.0), ts=11.0) + # Positive fixture stamps make the clipping interval explicit. + a = _stamped_transform(translation=_vector(0.0, 0.0, 0.0), ts=1.0) + b = _stamped_transform(translation=_vector(10.0, 0.0, 0.0), ts=11.0) graph = _graph_with_drift_at([a, b]) # Below range -> clipped to a - assert graph.correction_at(-5.0).translation.x == pytest.approx(0.0, abs=1e-10) + assert graph.correction_at(-5.0).transform.translation.x == pytest.approx(0.0, abs=1e-10) # Above range -> clipped to b - assert graph.correction_at(100.0).translation.x == pytest.approx(10.0, abs=1e-10) + assert graph.correction_at(100.0).transform.translation.x == pytest.approx(10.0, abs=1e-10) # In-range midpoint - assert graph.correction_at(6.0).translation.x == pytest.approx(5.0, abs=1e-10) + assert graph.correction_at(6.0).transform.translation.x == pytest.approx(5.0, abs=1e-10) def test_frozen(self) -> None: graph = PoseGraph() with pytest.raises(Exception): - graph.keyframes = (Keyframe(ts=0, local=Transform(), optimized=Transform()),) # type: ignore[misc] + graph.keyframes = ( + Keyframe(ts=0, local=_stamped_transform(), optimized=_stamped_transform()), + ) # type: ignore[misc] class TestPlace: def _graph(self) -> PoseGraph: - turn = Quaternion.from_rotation_matrix(Rotation.from_euler("z", np.pi / 2).as_matrix()) + turn = quaternion_from_matrix(Rotation.from_euler("z", np.pi / 2).as_matrix()) return _graph_with_drift_at( [ - Transform(translation=Vector3(0.0, 0.0, 0.0), ts=1.0), - Transform(translation=Vector3(10.0, 0.0, 0.0), rotation=turn, ts=11.0), + _stamped_transform(translation=_vector(0.0, 0.0, 0.0), ts=1.0), + _stamped_transform(translation=_vector(10.0, 0.0, 0.0), rotation=turn, ts=11.0), ] ) def test_latest_data_stays_put(self) -> None: - cloud = PointCloud2.from_numpy(np.array([[1.0, 2.0, 3.0]]), timestamp=11.0) + cloud = pointcloud_from_xyz( + np.array([[1.0, 2.0, 3.0]]), + header=Header(stamp=time_from_seconds(11.0), frame_id="world"), + ) np.testing.assert_allclose( - self._graph().place(cloud).points_f32(), [[1.0, 2.0, 3.0]], atol=1e-5 + pointcloud_xyz(self._graph().place(cloud)), [[1.0, 2.0, 3.0]], atol=1e-5 ) def test_old_data_moves_relative_to_the_latest_keyframe(self) -> None: # C(now)^-1 . C(1): undo the latest correction (turn 90 deg, +10 x). - cloud = PointCloud2.from_numpy(np.array([[10.0, 1.0, 0.0]]), timestamp=1.0) + cloud = pointcloud_from_xyz( + np.array([[10.0, 1.0, 0.0]]), + header=Header(stamp=time_from_seconds(1.0), frame_id="world"), + ) np.testing.assert_allclose( - self._graph().place(cloud).points_f32(), [[1.0, 0.0, 0.0]], atol=1e-5 + pointcloud_xyz(self._graph().place(cloud)), [[1.0, 0.0, 0.0]], atol=1e-5 ) def test_per_point_stamps_override_the_cloud_ts(self) -> None: - cloud = PointCloud2.from_numpy( + cloud = pointcloud_from_xyz( np.array([[10.0, 1.0, 0.0], [1.0, 2.0, 3.0]]), - timestamp=1.0, + header=Header(stamp=time_from_seconds(1.0), frame_id="world"), stamps=np.array([1.0, 11.0]), ) np.testing.assert_allclose( - self._graph().place(cloud).points_f32(), [[1.0, 0.0, 0.0], [1.0, 2.0, 3.0]], atol=1e-5 + pointcloud_xyz(self._graph().place(cloud)), + [[1.0, 0.0, 0.0], [1.0, 2.0, 3.0]], + atol=1e-5, ) @@ -299,14 +361,17 @@ def test_pure_translation_shifts_poses(self) -> None: lidar: Stream[PointCloud2] = mem.stream("lidar", PointCloud2) for i in range(3): lidar.append( - PointCloud2.from_numpy(np.zeros((1, 3), dtype=np.float32)), + pointcloud_from_xyz( + np.zeros((1, 3), dtype=np.float32), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ), ts=float(i + 1), pose=(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0), ) graph = _graph_with_drift_at( [ - Transform(translation=Vector3(5.0, 0.0, 0.0), ts=1.0), - Transform(translation=Vector3(5.0, 0.0, 0.0), ts=3.0), + _stamped_transform(translation=_vector(5.0, 0.0, 0.0), ts=1.0), + _stamped_transform(translation=_vector(5.0, 0.0, 0.0), ts=3.0), ] ) for obs in lidar.transform(graph): @@ -320,14 +385,17 @@ def test_passes_through_pose_none(self) -> None: mem = MemoryStore() lidar: Stream[PointCloud2] = mem.stream("lidar", PointCloud2) lidar.append( - PointCloud2.from_numpy(np.zeros((1, 3), dtype=np.float32)), + pointcloud_from_xyz( + np.zeros((1, 3), dtype=np.float32), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ), ts=1.0, pose=None, ) graph = _graph_with_drift_at( [ - Transform(translation=Vector3(5.0, 0.0, 0.0), ts=1.0), - Transform(translation=Vector3(5.0, 0.0, 0.0), ts=2.0), + _stamped_transform(translation=_vector(5.0, 0.0, 0.0), ts=1.0), + _stamped_transform(translation=_vector(5.0, 0.0, 0.0), ts=2.0), ] ) for obs in lidar.transform(graph): @@ -336,13 +404,13 @@ def test_passes_through_pose_none(self) -> None: class TestKeyframeType: def test_keyframe_is_frozen(self) -> None: - identity = Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), + identity = _stamped_transform( + translation=_vector(0.0, 0.0, 0.0), + rotation=_quaternion(0.0, 0.0, 0.0, 1.0), ts=1.0, ) kf = Keyframe(ts=1.0, local=identity, optimized=identity) with pytest.raises(Exception): kf.ts = 2.0 # type: ignore[misc] - assert isinstance(kf.local, Transform) - assert isinstance(kf.optimized, Transform) + assert isinstance(kf.local, TransformStamped) + assert isinstance(kf.optimized, TransformStamped) diff --git a/dimos/navigation/go2/loop_closure/test_pgo_auto.py b/dimos/navigation/go2/loop_closure/test_pgo_auto.py new file mode 100644 index 00000000000..c4f92b5844a --- /dev/null +++ b/dimos/navigation/go2/loop_closure/test_pgo_auto.py @@ -0,0 +1,108 @@ +# 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. + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.geometry_msgs.msg import Transform, TransformStamped, Vector3 +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest + +from dimos.memory.store.memory import MemoryStore +from dimos.memory.type.observation import Observation +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import time_from_seconds, to_seconds +from dimos.navigation.go2.loop_closure.pgo_auto import ( + FRAME_BODY, + FRAME_WORLD_CORRECTED, + FRAME_WORLD_RAW, + PGO, + Keyframe, + apply_corrections, + keyframes_to_corrections, + make_interpolator, +) + + +def _pose(x: float, frame: str, ts: float) -> TransformStamped: + return TransformStamped( + header=Header(frame_id=frame, stamp=time_from_seconds(ts)), + child_frame_id=FRAME_BODY, + transform=Transform( + translation=Vector3(x=x, y=0.0, z=0.0), rotation=quaternion_from_euler(0, 0, 0) + ), + ) + + +def test_generated_corrections_preserve_frames_time_and_nearest_neighbor() -> None: + with MemoryStore() as store: + keyframes = store.stream("keyframes", Keyframe) + keyframes.append( + Keyframe(0.0, _pose(1, FRAME_WORLD_RAW, 0), _pose(3, FRAME_WORLD_CORRECTED, 0)), ts=0 + ) + keyframes.append( + Keyframe(2.0, _pose(2, FRAME_WORLD_RAW, 2), _pose(6, FRAME_WORLD_CORRECTED, 2)), ts=2 + ) + corrections = keyframes_to_corrections(keyframes) + lookup = make_interpolator(corrections) + for ts, expected in [(-1, 2), (0.9, 2), (1.1, 4), (3, 4)]: + correction = cdr_decode(cdr_encode(lookup(ts)), TransformStamped) + assert correction.transform.translation.x == expected + assert correction.header.frame_id == FRAME_WORLD_CORRECTED + assert correction.child_frame_id == FRAME_WORLD_RAW + assert to_seconds(correction.header.stamp) == ts + stream = store.stream("values", str) + stream.append("payload", ts=0, pose=(1, 0, 0, 0, 0, 0, 1)) + stream.append("no pose", ts=2) + output = list(apply_corrections(stream, corrections)) + assert output[0].data == "payload" + assert output[0].pose_tuple == (3, 0, 0, 0, 0, 0, 1) + assert output[0].ts == 0 + assert output[1].pose is None + assert stream.first().pose_tuple == (1, 0, 0, 0, 0, 0, 1) + + +def test_empty_pgo_has_generated_identity_correction() -> None: + pytest.importorskip("gtsam") + graph = next(PGO()(iter(()))).data + assert not graph.keyframes and not graph.loops + correction = graph.correction_at(0) + assert to_seconds(correction.header.stamp) == 0 + assert correction.transform.rotation.w == 1 + result = graph.correct(_pose(5, FRAME_WORLD_RAW, 0)) + assert result.transform.translation.x == 5 + assert result.header.frame_id == FRAME_WORLD_CORRECTED + + +def test_pgo_consumes_generated_cloud_without_mutating_input() -> None: + pytest.importorskip("gtsam") + cloud = pointcloud_from_xyz( + np.array([[1.0, 0, 0], [1, 1, 0], [1, 0, 1]]), + header=Header(frame_id=FRAME_WORLD_RAW, stamp=time_from_seconds(12.25)), + ) + before = cdr_encode(cloud) + quaternion = quaternion_from_euler(0, 0, 0.1) + observation = Observation( + ts=12.25, + _data=cloud, + pose=(1, 0, 0, quaternion.x, quaternion.y, quaternion.z, quaternion.w), + ) + graph = next(PGO()(iter([observation]))).data + assert len(graph.keyframes) == 1 + assert graph.keyframes[0].ts == 12.25 + assert graph.keyframes[0].local.header.frame_id == FRAME_WORLD_RAW + correction = graph.correction_at(12.25) + assert correction.transform.translation.x == pytest.approx(0, abs=1e-6) + assert correction.transform.rotation.w == pytest.approx(1) + assert cdr_encode(cloud) == before diff --git a/dimos/navigation/go2/loop_closure/test_pgo_generated.py b/dimos/navigation/go2/loop_closure/test_pgo_generated.py new file mode 100644 index 00000000000..f1c833e9882 --- /dev/null +++ b/dimos/navigation/go2/loop_closure/test_pgo_generated.py @@ -0,0 +1,177 @@ +# 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. + +import math +import sys +from types import SimpleNamespace +from unittest.mock import MagicMock + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest + +from dimos.memory.type.observation import Observation +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.msgs.time import time_from_nanoseconds, time_from_seconds +from dimos.navigation.go2.loop_closure.pgo import ( + Keyframe, + PGOConfig, + PoseGraph, + _KeyPose, + _PGOState, +) + + +def test_generated_observation_pose_correction_applies_rotation_after_translation(): + local = TransformStamped( + header=Header(frame_id="world_raw", stamp=time_from_seconds(1.0)), + child_frame_id="body", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + optimized = TransformStamped( + header=Header(frame_id="world_corrected", stamp=time_from_seconds(1.0)), + child_frame_id="body", + transform=Transform( + translation=Vector3(x=5.0, y=0.0, z=0.0), + rotation=Quaternion(z=math.sqrt(0.5), w=math.sqrt(0.5), x=0.0, y=0.0), + ), + ) + graph = PoseGraph(keyframes=(Keyframe(ts=1.0, local=local, optimized=optimized),)) + source = Observation(id=1, ts=1.0, _data="payload", pose=(1.0, 2.0, 0.0, 0.0, 0.0, 0.0, 1.0)) + corrected = next(graph(iter([source]))) + assert corrected.data == "payload" and corrected.ts == source.ts + np.testing.assert_allclose( + corrected.pose_tuple, (3.0, 1.0, 0.0, 0.0, 0.0, math.sqrt(0.5), math.sqrt(0.5)), atol=1e-9 + ) + assert source.pose_tuple == (1.0, 2.0, 0.0, 0.0, 0.0, 0.0, 1.0) + + +def test_pose_graph_retains_poseless_observations_without_constructing_correction(): + source = Observation(id=1, ts=1.0, _data="payload") + assert list(PoseGraph()(iter([source]))) == [source] + + +def test_pgo_submap_places_generated_body_clouds_before_merging(monkeypatch): + class MatrixPose: + def __init__(self, x): + self.value = np.eye(4) + self.value[0, 3] = x + + def matrix(self): + return self.value + + solver = SimpleNamespace( + **{ + name: MagicMock(name=name) + for name in ( + "Pose3", + "ISAM2Params", + "ISAM2", + "NonlinearFactorGraph", + "Values", + ) + } + ) + monkeypatch.setitem(sys.modules, "gtsam", solver) + state = _PGOState(PGOConfig(submap_resolution=0.2)) + solver.ISAM2Params.return_value.setRelinearizeThreshold.assert_called_once_with(0.01) + assert solver.ISAM2Params.return_value.relinearizeSkip == 1 + clouds = [ + pointcloud_from_xyz( + np.array([[1.0, 0.0, 0.0]]), header=Header(frame_id="body", stamp=time_from_seconds(ts)) + ) + for ts in (1.0, 2.0) + ] + state._key_poses = [ + _KeyPose(local=MatrixPose(0), optimized=MatrixPose(offset), timestamp=ts, body_cloud=cloud) + for cloud, offset, ts in zip(clouds, (2.0, 4.0), (1.0, 2.0), strict=True) + ] + before = [cdr_encode(cloud) for cloud in clouds] + result = state._get_submap(0, 1) + np.testing.assert_allclose(pointcloud_xyz(result), [[3.0, 0.0, 0.0], [5.0, 0.0, 0.0]]) + assert result.header.frame_id == "world_corrected" + assert result.header.stamp == time_from_seconds(2.0) + assert [cdr_encode(cloud) for cloud in clouds] == before + + +@pytest.mark.parametrize( + "query, expected_x", [(-5.0, 0.0), (1.0, 0.0), (6.0, 5.0), (11.0, 10.0), (100.0, 10.0)] +) +def test_generated_correction_interpolates_and_clips_without_solver(query, expected_x): + keyframes = [] + for ts, x in ((1.0, 0.0), (11.0, 10.0)): + local = TransformStamped( + header=Header(frame_id="world_raw", stamp=time_from_seconds(ts)), + child_frame_id="body", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + optimized = TransformStamped( + header=Header(frame_id="world_corrected", stamp=time_from_seconds(ts)), + child_frame_id="body", + transform=Transform( + translation=Vector3(x=x, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + keyframes.append(Keyframe(ts=ts, local=local, optimized=optimized)) + graph = PoseGraph(keyframes=tuple(keyframes)) + value = graph.correction_at(query) + assert value.header.frame_id == "world_corrected" and value.child_frame_id == "world_raw" + assert value.header.stamp == time_from_seconds(query) + assert value.transform.translation.x == pytest.approx(expected_x) + assert cdr_decode(cdr_encode(value), TransformStamped) == value + + +def test_generated_correct_preserves_exact_source_stamp_and_rejects_bad_frame(): + local = TransformStamped( + header=Header(frame_id="world_raw", stamp=Time(sec=0, nanosec=0)), + child_frame_id="body", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + optimized = TransformStamped( + header=Header(frame_id="world_corrected", stamp=Time(sec=0, nanosec=0)), + child_frame_id="body", + transform=Transform( + translation=Vector3(x=5.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + graph = PoseGraph(keyframes=(Keyframe(ts=1.0, local=local, optimized=optimized),)) + source = TransformStamped( + header=Header(frame_id="world_raw", stamp=time_from_nanoseconds(1700000000123456789)), + child_frame_id="camera", + transform=Transform( + translation=Vector3(x=2.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + result = graph.correct(source) + assert result.header.stamp == source.header.stamp + assert result.header.frame_id == "world_corrected" and result.child_frame_id == "camera" + assert result.transform.translation.x == 7.0 and source.transform.translation.x == 2.0 + source.header.frame_id = "unrelated" + with pytest.raises(ValueError, match="Cannot compose frames"): + graph.correct(source) diff --git a/dimos/navigation/go2/loop_closure/test_pgo_map.py b/dimos/navigation/go2/loop_closure/test_pgo_map.py index 42de2edcefe..f9f5429b7ce 100644 --- a/dimos/navigation/go2/loop_closure/test_pgo_map.py +++ b/dimos/navigation/go2/loop_closure/test_pgo_map.py @@ -12,13 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.msgs.time import time_from_seconds from dimos.navigation.go2.loop_closure.pgo_map import PGOMap # TODO(PY311): drop. Skip before importing test_pgo, which skips mid-import without gtsam. @@ -46,35 +54,54 @@ def snapshot(self): def _wall(x, ts): ys, zs = np.meshgrid(np.arange(0.025, 1, VOXEL), np.arange(0.025, 1, VOXEL)) pts = np.stack([np.full(ys.size, x), ys.ravel(), zs.ravel()], axis=1) - return PointCloud2.from_numpy(pts, timestamp=ts) + return pointcloud_from_xyz(pts, header=Header(stamp=time_from_seconds(ts), frame_id="")) def _xs(world_map): - return {int(x) for x in np.floor(world_map.global_map().points_f32()[:, 0] / VOXEL)} + return {int(x) for x in np.floor(pointcloud_xyz(world_map.global_map())[:, 0] / VOXEL)} def test_rebuilds_on_loop_closure_and_respects_cooldown(): world_map = PGOMap(rebuild_cooldown_s=10.0) world_map._pgo = pgo = _FakePGO() - pose = PoseStamped(1.0, 0.0, 0.0) + pose = PoseStamped( + header=Header(stamp=time_from_seconds(0.0), frame_id=""), + pose=Pose( + position=Point(x=1.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + ) # the same wall seen twice, the first time 0.5 m off; no loop yet - assert not world_map.add(_wall(1.025, 100.0), pose) - assert not world_map.add(_wall(1.525, 101.0), pose) + assert not world_map.add(_wall(1.025, 100.0), pose.pose) + assert not world_map.add(_wall(1.525, 101.0), pose.pose) assert _xs(world_map) == {20, 30} # loop closes: the frame at t=100 was really 0.5 m further along x pgo.n_loops = 1 pgo.graph = _graph_with_drift_at( [ - Transform(translation=Vector3(0.5, 0.0, 0.0), ts=100.0), - Transform(translation=Vector3(0.0, 0.0, 0.0), ts=101.0), + TransformStamped( + header=Header(stamp=time_from_seconds(100.0), frame_id=""), + child_frame_id="", + transform=Transform( + translation=Vector3(0.5, 0.0, 0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + TransformStamped( + header=Header(stamp=time_from_seconds(101.0), frame_id=""), + child_frame_id="", + transform=Transform( + translation=Vector3(0.0, 0.0, 0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), ] ) - assert world_map.add(_wall(1.525, 102.0), pose) + assert world_map.add(_wall(1.525, 102.0), pose.pose) assert _xs(world_map) == {30} # a second loop inside the cooldown waits for it pgo.n_loops = 2 - assert not world_map.add(_wall(1.525, 105.0), pose) - assert world_map.add(_wall(1.525, 112.5), pose) + assert not world_map.add(_wall(1.525, 105.0), pose.pose) + assert world_map.add(_wall(1.525, 112.5), pose.pose) diff --git a/dimos/navigation/go2/replanning_a_star/controllers.py b/dimos/navigation/go2/replanning_a_star/controllers.py index 1b7f269d096..00327bf96bd 100644 --- a/dimos/navigation/go2/replanning_a_star/controllers.py +++ b/dimos/navigation/go2/replanning_a_star/controllers.py @@ -57,7 +57,7 @@ def advance(self, lookahead_point: NDArray[np.float64], current_odom: PoseStampe if distance < 1e-6: # Robot is coincidentally at the lookahead point; skip this cycle. - return Twist() + return Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) robot_yaw = quaternion_euler(current_odom.pose.orientation)[2] desired_yaw = np.arctan2(direction[1], direction[0]) @@ -74,8 +74,8 @@ def advance(self, lookahead_point: NDArray[np.float64], current_odom: PoseStampe linear_velocity = self._apply_min_velocity(linear_velocity, self._min_linear_velocity) return Twist( - linear=Vector3(x=linear_velocity), - angular=Vector3(z=angular_velocity), + linear=Vector3(x=linear_velocity, y=0.0, z=0.0), + angular=Vector3(z=angular_velocity, x=0.0, y=0.0), ) def rotate(self, yaw_error: float) -> Twist: @@ -108,8 +108,8 @@ def _angular_twist(self, angular_velocity: float) -> Twist: angular_velocity = 0.8 * np.sign(angular_velocity) return Twist( - linear=Vector3(), - angular=Vector3(z=angular_velocity), + linear=Vector3(x=0.0, y=0.0, z=0.0), + angular=Vector3(z=angular_velocity, x=0.0, y=0.0), ) diff --git a/dimos/navigation/go2/replanning_a_star/global_planner.py b/dimos/navigation/go2/replanning_a_star/global_planner.py index b1c39ee55fc..634992484dd 100644 --- a/dimos/navigation/go2/replanning_a_star/global_planner.py +++ b/dimos/navigation/go2/replanning_a_star/global_planner.py @@ -16,9 +16,10 @@ from threading import Event, RLock, Thread, current_thread import time +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, PoseStamped, Twist from dimos_generated.nav_msgs.msg import OccupancyGrid, Path -from dimos_generated.std_msgs.msg import Bool +from dimos_generated.std_msgs.msg import Bool, Header import numpy as np from reactivex import Subject from reactivex.disposable import CompositeDisposable @@ -162,7 +163,7 @@ def cancel_goal(self, *, but_will_try_again: bool = False, arrived: bool = False self._goal_reached = arrived self._replan_limiter.reset() - self.path.on_next(Path()) + self.path.on_next(Path(header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), poses=[])) self._local_planner.stop_planning() if not but_will_try_again: @@ -271,7 +272,7 @@ def _on_stopped_navigating(self, stop_message: StopMessage) -> None: def _handle_stop_message(self, stop_message: StopMessage) -> None: # Note, this runs in the monitoring thread. - self.path.on_next(Path()) + self.path.on_next(Path(header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), poses=[])) if stop_message == "arrived": logger.info("Arrived at goal.") diff --git a/dimos/navigation/go2/replanning_a_star/local_planner.py b/dimos/navigation/go2/replanning_a_star/local_planner.py index 56438270de3..4184a1e57ee 100644 --- a/dimos/navigation/go2/replanning_a_star/local_planner.py +++ b/dimos/navigation/go2/replanning_a_star/local_planner.py @@ -18,7 +18,7 @@ import traceback from typing import Literal, TypeAlias -from dimos_generated.geometry_msgs.msg import PoseStamped, Twist +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist, Vector3 from dimos_generated.nav_msgs.msg import OccupancyGrid, Path import numpy as np from reactivex import Subject @@ -130,7 +130,9 @@ def stop_planning(self) -> None: with self._lock: self._thread = None self._reset_state() - self.cmd_vel.on_next(Twist()) + self.cmd_vel.on_next( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) def get_state(self) -> NavigationState: with self._lock: @@ -157,7 +159,9 @@ def _thread_entrypoint(self) -> None: self.stopped_navigating.on_next("error") finally: self._reset_state() - self.cmd_vel.on_next(Twist()) + self.cmd_vel.on_next( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) def _change_state(self, new_state: PlannerState) -> None: if new_state == self._state: @@ -310,7 +314,7 @@ def _compute_final_rotation(self) -> Twist: logger.info("Final rotation complete, goal reached") with self._lock: self._change_state("arrived") - return Twist() + return Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) return self._controller.rotate(yaw_error) diff --git a/dimos/navigation/go2/replanning_a_star/min_cost_astar.py b/dimos/navigation/go2/replanning_a_star/min_cost_astar.py index 54f1f744d39..a81129879d2 100644 --- a/dimos/navigation/go2/replanning_a_star/min_cost_astar.py +++ b/dimos/navigation/go2/replanning_a_star/min_cost_astar.py @@ -83,7 +83,10 @@ def _reconstruct_path_from_coords( poses=[ PoseStamped( header=costmap.header, - pose=Pose(position=grid_to_world(costmap, coordinate), orientation=Quaternion(w=1)), + pose=Pose( + position=grid_to_world(costmap, coordinate), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), ) for coordinate in path_coords ], @@ -99,7 +102,9 @@ def min_cost_astar( use_cpp: bool = True, ) -> Path | None: cells = occupancy_view(costmap) - start_vector = world_to_grid(costmap, start if start is not None else Point()) + start_vector = world_to_grid( + costmap, start if start is not None else Point(x=0.0, y=0.0, z=0.0) + ) goal_vector = world_to_grid(costmap, goal) # Ignore sub-nanocell rotation roundoff before assigning a point to its cell. start_tuple = (math.floor(round(start_vector[0], 9)), math.floor(round(start_vector[1], 9))) diff --git a/dimos/navigation/go2/replanning_a_star/module.py b/dimos/navigation/go2/replanning_a_star/module.py index 7423dc5d91a..ba47fbcf5e6 100644 --- a/dimos/navigation/go2/replanning_a_star/module.py +++ b/dimos/navigation/go2/replanning_a_star/module.py @@ -99,7 +99,9 @@ def start(self) -> None: lambda pt: self._planner.handle_goal_request( PoseStamped( header=pt.header, - pose=Pose(position=pt.point, orientation=Quaternion(w=1)), + pose=Pose( + position=pt.point, orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), ) ) ) diff --git a/dimos/navigation/go2/replanning_a_star/test_generated_navigation.py b/dimos/navigation/go2/replanning_a_star/test_generated_navigation.py new file mode 100644 index 00000000000..340bdcc1417 --- /dev/null +++ b/dimos/navigation/go2/replanning_a_star/test_generated_navigation.py @@ -0,0 +1,216 @@ +# 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. + +import math +from threading import Event + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Twist, Vector3 +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid, Path +from dimos_generated.std_msgs.msg import Bool, Header +import numpy as np + +from dimos.core.global_config import GlobalConfig +from dimos.navigation.go2.replanning_a_star.global_planner import GlobalPlanner +from dimos.navigation.go2.replanning_a_star.local_planner import LocalPlanner +from dimos.navigation.go2.replanning_a_star.navigation_map import NavigationMap +from dimos.navigation.go2.replanning_a_star.path_clearance import PathClearance + + +def test_global_planner_reaches_goal_with_generated_cdr_feedback(): + header = Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id="map") + grid = OccupancyGrid( + header=header, + info=MapMetaData( + width=24, + height=24, + resolution=0.25, + origin=Pose( + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + position=Point(x=0.0, y=0.0, z=0.0), + ), + map_load_time=Time(sec=0, nanosec=0), + ), + data=np.zeros(24 * 24, dtype=np.int8), + ) + planner = GlobalPlanner(GlobalConfig()) + planner._local_planner._control_frequency = 100 + arrived = Event() + commands: list[Twist] = [] + paths: list[Path] = [] + position = [1.0, 1.0, 0.0] + + def update_odom() -> None: + message = PoseStamped( + header=header, + pose=Pose( + position=Point(x=position[0], y=position[1], z=0.0), + orientation=Quaternion( + z=math.sin(position[2] / 2), w=math.cos(position[2] / 2), x=0.0, y=0.0 + ), + ), + ) + planner.handle_odom(cdr_decode(cdr_encode(message), PoseStamped)) + + def command(message: Twist) -> None: + decoded = cdr_decode(cdr_encode(message), Twist) + commands.append(decoded) + position[0] += decoded.linear.x * math.cos(position[2]) * 0.1 + position[1] += decoded.linear.x * math.sin(position[2]) * 0.1 + position[2] += decoded.angular.z * 0.1 + update_odom() + + def goal_reached(message: Bool) -> None: + if cdr_decode(cdr_encode(message), Bool).data: + arrived.set() + + subscriptions = [ + planner.cmd_vel.subscribe(command), + planner.goal_reached.subscribe(goal_reached), + planner.path.subscribe(lambda value: paths.append(cdr_decode(cdr_encode(value), Path))), + ] + try: + planner.handle_global_costmap(cdr_decode(cdr_encode(grid), OccupancyGrid)) + update_odom() + planner.start() + planner.handle_goal_request( + PoseStamped( + header=header, + pose=Pose( + position=Point(x=3, y=1, z=0.0), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), + ) + ) + assert arrived.wait(5), "planner did not report arrival" + assert abs(position[0] - 3) < 0.25 + assert abs(position[1] - 1) < 0.1 + assert any(command.linear.x > 0 for command in commands) + published = next(path for path in paths if len(path.poses)) + assert published.header == header + assert all(pose.header == header for pose in published.poses) + finally: + planner.stop() + for subscription in subscriptions: + subscription.dispose() + assert commands[-1] == Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) + + +def test_clearance_mask_tracks_generated_map_geometry(): + path = Path( + poses=[ + PoseStamped( + pose=Pose( + position=Point(x=x, y=2, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + for x in [2, 3, 4] + ], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + clearance = PathClearance(GlobalConfig(robot_width=0.1), path) + grid = OccupancyGrid( + info=MapMetaData( + width=10, + height=10, + resolution=1, + origin=Pose( + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + position=Point(x=0.0, y=0.0, z=0.0), + ), + map_load_time=Time(sec=0, nanosec=0), + ), + data=np.asarray([0] * 100, dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + clearance.update_costmap(grid) + original = clearance.mask.copy() + assert original[2, 2] + grid.info.origin.position.x = 1 + shifted = clearance.mask.copy() + assert shifted[2, 1] and not original[2, 1] + grid.info.resolution = 0.5 + scaled = clearance.mask + assert scaled[4, 2] and not shifted[4, 2] + + +def test_cancel_joins_control_loop_and_ends_with_generated_stop(): + config = GlobalConfig() + maps = NavigationMap(config, "gradient") + maps.update( + OccupancyGrid( + info=MapMetaData( + width=10, + height=10, + resolution=1, + origin=Pose( + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + position=Point(x=0.0, y=0.0, z=0.0), + ), + map_load_time=Time(sec=0, nanosec=0), + ), + data=np.asarray([0] * 100, dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + ) + planner = LocalPlanner(config, maps, 0.2) + planner.handle_odom( + PoseStamped( + pose=Pose( + position=Point(x=1, y=1, z=0.0), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + ) + moving = Event() + commands = [] + + def command(message: Twist) -> None: + commands.append(cdr_decode(cdr_encode(message), Twist)) + if message.linear.x > 0: + moving.set() + + subscription = planner.cmd_vel.subscribe(command) + try: + planner.start_planning( + Path( + poses=[ + PoseStamped( + pose=Pose( + position=Point(x=x, y=1, z=0.0), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + for x in [1, 2, 3] + ], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + ) + thread = planner._thread + assert thread is not None + assert moving.wait(2) + planner.stop() + assert not thread.is_alive() + assert commands[-1] == Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) + finally: + planner.stop() + subscription.dispose() diff --git a/dimos/navigation/go2/replanning_a_star/test_generated_pipeline.py b/dimos/navigation/go2/replanning_a_star/test_generated_pipeline.py new file mode 100644 index 00000000000..bd9e60eec0e --- /dev/null +++ b/dimos/navigation/go2/replanning_a_star/test_generated_pipeline.py @@ -0,0 +1,130 @@ +# 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. + +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +import math + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid, Path +from dimos_generated.std_msgs.msg import Header +import numpy as np +import pytest + +from dimos.mapping.occupancy.path_mask import make_path_mask +from dimos.mapping.occupancy.path_resampling import simple_resample_path, smooth_resample_path +from dimos.msgs.geometry import quaternion_euler, quaternion_from_euler +from dimos.msgs.occupancy import grid_to_world, world_to_grid +from dimos.navigation.go2.replanning_a_star.min_cost_astar import _USE_CPP, min_cost_astar + + +@pytest.mark.parametrize("use_cpp", [False, True]) +@pytest.mark.parametrize("yaw", [0.0, math.pi / 2]) +def test_generated_grid_to_path_and_mask(use_cpp, yaw): + if use_cpp: + assert _USE_CPP, "native A* extension must be built for this integration test" + cells = np.zeros((30, 30), dtype=np.int8) + cells[15, 0:20] = 100 + header = Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id="map") + grid = OccupancyGrid( + header=header, + info=MapMetaData( + width=30, + height=30, + resolution=0.25, + origin=Pose( + position=Point(x=2, y=3, z=0.0), orientation=quaternion_from_euler(0, 0, yaw) + ), + map_load_time=Time(sec=0, nanosec=0), + ), + data=np.asarray(cells.ravel(), dtype=np.int8), + ) + received = cdr_decode(cdr_encode(grid), OccupancyGrid) + start = grid_to_world(received, (5, 5)) + goal = grid_to_world(received, (25, 25)) + path = min_cost_astar(received, goal, start, use_cpp=use_cpp) + assert path is not None + decoded = cdr_decode(cdr_encode(path), Path) + assert decoded.header == header + assert world_to_grid(received, decoded.poses[0].pose.position) == pytest.approx((5, 5)) + assert world_to_grid(received, decoded.poses[-1].pose.position) == pytest.approx((25, 25)) + for pose in decoded.poses: + assert pose.header == header + x, y = world_to_grid(received, pose.pose.position) + assert cells[round(y), round(x)] == 0 + mask = make_path_mask(received, decoded, 0.1) + assert mask.any() + assert not np.any(mask & (cells == 100)) + assert world_to_grid(grid, start) == pytest.approx((5, 5)) + assert world_to_grid(grid, goal) == pytest.approx((25, 25)) + + +@pytest.mark.parametrize("resample", [simple_resample_path, smooth_resample_path]) +def test_resampling_preserves_headers_orientations_and_input(resample): + header = Header(stamp=Time(sec=123, nanosec=456), frame_id="map") + path = Path( + header=header, + poses=[ + PoseStamped( + header=header, + pose=Pose( + position=Point(x=1, y=y, z=0.0), orientation=quaternion_from_euler(0, 0, 0) + ), + ) + for y in [1, 2, 3] + ], + ) + before = cdr_encode(path) + goal = Pose( + orientation=quaternion_from_euler(0, 0, -math.pi / 2), position=Point(x=0.0, y=0.0, z=0.0) + ) + actual = cdr_decode(cdr_encode(resample(path, goal, 0.1)), Path) + assert actual.header == header + assert len(actual.poses) > 3 + assert actual.poses[0].pose.position == path.poses[0].pose.position + assert actual.poses[-1].pose.position == path.poses[-1].pose.position + assert actual.poses[-1].pose.orientation == goal.orientation + for pose in list(actual.poses)[:-1]: + assert pose.header == header + assert quaternion_euler(pose.pose.orientation)[2] == pytest.approx(math.pi / 2) + assert cdr_encode(path) == before + + +@pytest.mark.parametrize("use_cpp", [False, True]) +def test_outside_start_and_goal_are_rejected(use_cpp): + grid = OccupancyGrid( + info=MapMetaData( + width=3, + height=3, + resolution=1, + origin=Pose( + orientation=quaternion_from_euler(0, 0, 0), position=Point(x=0.0, y=0.0, z=0.0) + ), + map_load_time=Time(sec=0, nanosec=0), + ), + data=np.asarray([0] * 9, dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + assert ( + min_cost_astar(grid, Point(x=1, y=1, z=0.0), Point(x=-0.1, y=0.0, z=0.0), use_cpp=use_cpp) + is None + ) + assert ( + min_cost_astar(grid, Point(x=-0.1, y=0.0, z=0.0), Point(x=1, y=1, z=0.0), use_cpp=use_cpp) + is None + ) + assert min_cost_astar(grid, Point(x=3, y=0.0, z=0.0), use_cpp=use_cpp) is None diff --git a/dimos/navigation/go2/replanning_a_star/test_global_planner.py b/dimos/navigation/go2/replanning_a_star/test_global_planner.py index 76c05a3ac82..68f6f0e8a08 100644 --- a/dimos/navigation/go2/replanning_a_star/test_global_planner.py +++ b/dimos/navigation/go2/replanning_a_star/test_global_planner.py @@ -12,12 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid +from dimos_generated.std_msgs.msg import Header import numpy as np from dimos.core.global_config import GlobalConfig -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid from dimos.navigation.go2.replanning_a_star.global_planner import GlobalPlanner @@ -29,16 +30,29 @@ def test_find_wide_path_with_start_inside_inflation() -> None: resolution = 0.05 grid = np.zeros((60, 60), dtype=np.int8) grid[20:40, 30] = 100 # wall at x=1.5m spanning y=1.0..2.0m - costmap = OccupancyGrid(grid=grid, resolution=resolution, origin=Pose(), frame_id="world") + costmap = OccupancyGrid( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + info=MapMetaData( + width=60, + height=60, + resolution=resolution, + origin=Pose( + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + position=Point(x=0.0, y=0.0, z=0.0), + ), + map_load_time=Time(sec=0, nanosec=0), + ), + data=np.asarray(grid.ravel(), dtype=np.int8), + ) planner = GlobalPlanner(GlobalConfig()) planner.handle_global_costmap(costmap) # 7 cm in front of the wall: within the inflation radius # (robot_width * 1.1 / 2 = 0.165m), so the start cell is engulfed. - robot = Vector3(1.43, 1.5, 0) + robot = Point(x=1.43, y=1.5, z=0.0) # On the other side of the wall; the path must round a wall end. - goal = Vector3(2.75, 1.5, 0) + goal = Point(x=2.75, y=1.5, z=0.0) path = planner._find_wide_path(goal, robot) diff --git a/dimos/navigation/go2/replanning_a_star/test_goal_validator.py b/dimos/navigation/go2/replanning_a_star/test_goal_validator.py index 73f1b1f111e..5d513429c51 100644 --- a/dimos/navigation/go2/replanning_a_star/test_goal_validator.py +++ b/dimos/navigation/go2/replanning_a_star/test_goal_validator.py @@ -12,42 +12,59 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid from dimos.navigation.go2.replanning_a_star.goal_validator import find_safe_goal from dimos.utils.data import get_data @pytest.fixture def costmap() -> OccupancyGrid: - return OccupancyGrid(np.load(get_data("occupancy_simple.npy"))) + cells = np.load(get_data("occupancy_simple.npy")) + return OccupancyGrid( + info=MapMetaData( + width=cells.shape[1], + height=cells.shape[0], + resolution=float(np.float32(0.05)), + origin=Pose( + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + position=Point(x=0.0, y=0.0, z=0.0), + ), + map_load_time=Time(sec=0, nanosec=0), + ), + data=np.asarray(cells.astype(np.int8).ravel(), dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) @pytest.mark.parametrize( "input_pos,expected_pos", [ - # Identical. - ((6.15, 10.0), (6.15, 10.0)), + # The requested point falls in the preceding cell at ROS float32 resolution. + ((6.15, 10.0), (6.10, 9.95)), # Very slightly off. - ((6.0, 10.0), (6.05, 10.0)), + ((6.0, 10.0), (6.05, 9.95)), # Don't pick a spot that's the closest, but is actually on the other side of the wall. ((5.0, 9.0), (5.85, 9.6)), ], ) def test_find_safe_goal(costmap, input_pos, expected_pos) -> None: - goal = Vector3(input_pos[0], input_pos[1], 0.0) + goal = Point(x=input_pos[0], y=input_pos[1], z=0.0) safe_goal = find_safe_goal( costmap, goal, algorithm="bfs_contiguous", - cost_threshold=CostValues.OCCUPIED, + cost_threshold=100, min_clearance=0.3, max_search_distance=5.0, connectivity_check_radius=0, ) - assert safe_goal == Vector3(expected_pos[0], expected_pos[1], 0.0) + assert safe_goal is not None + assert (safe_goal.x, safe_goal.y, safe_goal.z) == pytest.approx((*expected_pos, 0), abs=1e-6) diff --git a/dimos/navigation/go2/replanning_a_star/test_min_cost_astar.py b/dimos/navigation/go2/replanning_a_star/test_min_cost_astar.py index 81e36eb3b34..21e1a69e649 100644 --- a/dimos/navigation/go2/replanning_a_star/test_min_cost_astar.py +++ b/dimos/navigation/go2/replanning_a_star/test_min_cost_astar.py @@ -14,49 +14,69 @@ import time +import cv2 +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid +from dimos_generated.std_msgs.msg import Header import numpy as np -from open3d.geometry import PointCloud import pytest from dimos.mapping.occupancy.gradient import gradient, voronoi_gradient from dimos.mapping.occupancy.visualizations import visualize_occupancy_grid -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_view +from dimos.msgs.occupancy import grid_to_world, world_to_grid from dimos.navigation.go2.replanning_a_star.min_cost_astar import min_cost_astar from dimos.utils.data import get_data +def _grid(cells, resolution=0.05): + return OccupancyGrid( + info=MapMetaData( + width=cells.shape[1], + height=cells.shape[0], + resolution=float(np.float32(resolution)), + origin=Pose( + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + position=Point(x=0.0, y=0.0, z=0.0), + ), + map_load_time=Time(sec=0, nanosec=0), + ), + data=np.asarray(cells.astype(np.int8).ravel(), dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + + @pytest.fixture -def costmap() -> PointCloud: - return gradient(OccupancyGrid(np.load(get_data("occupancy_simple.npy"))), max_distance=1.5) +def costmap() -> OccupancyGrid: + return gradient(_grid(np.load(get_data("occupancy_simple.npy"))), max_distance=1.5) @pytest.fixture -def costmap_three_paths() -> PointCloud: - return voronoi_gradient(OccupancyGrid(np.load(get_data("three_paths.npy"))), max_distance=1.5) +def costmap_three_paths() -> OccupancyGrid: + return voronoi_gradient(_grid(np.load(get_data("three_paths.npy"))), max_distance=1.5) def test_astar(costmap) -> None: - start = Vector3(4.0, 2.0) - goal = Vector3(6.15, 10.0) - expected = Image.from_file(get_data("astar_min_cost.png")) + start = Point(x=4, y=2, z=0.0) + goal = Point(x=6.15, y=10, z=0.0) + expected = cv2.imread(str(get_data("astar_min_cost.png")), cv2.IMREAD_COLOR) path = min_cost_astar(costmap, goal, start, use_cpp=False) actual = visualize_occupancy_grid(costmap, "rainbow", path) - np.testing.assert_array_equal(actual.data, expected.data) + np.testing.assert_array_equal(image_view(actual), expected) def test_astar_corner(costmap_three_paths) -> None: - start = Vector3(2.8, 3.35) - goal = Vector3(6.35, 4.25) - expected = Image.from_file(get_data("astar_corner_min_cost.png")) + start = Point(x=2.8, y=3.35, z=0.0) + goal = Point(x=6.35, y=4.25, z=0.0) + expected = cv2.imread(str(get_data("astar_corner_min_cost.png")), cv2.IMREAD_COLOR) path = min_cost_astar(costmap_three_paths, goal, start, use_cpp=False) actual = visualize_occupancy_grid(costmap_three_paths, "rainbow", path) - np.testing.assert_array_equal(actual.data, expected.data) + np.testing.assert_array_equal(image_view(actual), expected) def test_astar_unknown_penalty_blocks_unknown_cells(costmap) -> None: @@ -69,10 +89,10 @@ def test_astar_unknown_penalty_blocks_unknown_cells(costmap) -> None: grid[10:90, 10] = 0 # left column grid[89, 10:90] = 0 # bottom row grid[10:90, 89] = 0 # right column - og = OccupancyGrid(grid, resolution=0.1) + og = _grid(grid, resolution=0.1) - start = og.grid_to_world((10, 10)) - goal = og.grid_to_world((89, 10)) + start = grid_to_world(og, (10, 10)) + goal = grid_to_world(og, (89, 10)) for use_cpp in [False, True]: path = min_cost_astar(og, goal, start, unknown_penalty=1.0, use_cpp=use_cpp) @@ -81,8 +101,8 @@ def test_astar_unknown_penalty_blocks_unknown_cells(costmap) -> None: continue # Verify no path cell lands on an unknown cell for pose in path.poses: - gp = og.world_to_grid((pose.position.x, pose.position.y)) - gx, gy = round(gp.x), round(gp.y) + gp = world_to_grid(og, pose.pose.position) + gx, gy = round(gp[0]), round(gp[1]) if 0 <= gx < 100 and 0 <= gy < 100: assert grid[gy, gx] != -1, ( f"Path traverses unknown cell at grid ({gx}, {gy}), use_cpp={use_cpp}" @@ -94,10 +114,10 @@ def test_astar_unknown_penalty_allows_with_low_penalty(costmap) -> None: grid = np.full((50, 50), -1, dtype=np.int8) # All unknown grid[5, 5] = 0 # start cell free grid[45, 45] = 0 # goal cell free - og = OccupancyGrid(grid, resolution=0.1) + og = _grid(grid, resolution=0.1) - start = og.grid_to_world((5, 5)) - goal = og.grid_to_world((45, 45)) + start = grid_to_world(og, (5, 5)) + goal = grid_to_world(og, (45, 45)) for use_cpp in [False, True]: path = min_cost_astar(og, goal, start, unknown_penalty=0.5, use_cpp=use_cpp) @@ -108,8 +128,8 @@ def test_astar_unknown_penalty_allows_with_low_penalty(costmap) -> None: def test_astar_python_and_cpp(costmap) -> None: - start = Vector3(4.0, 2.0, 0) - goal = Vector3(6.15, 10.0) + start = Point(x=4, y=2, z=0.0) + goal = Point(x=6.15, y=10, z=0.0) start_time = time.perf_counter() path_python = min_cost_astar(costmap, goal, start, use_cpp=False) @@ -130,7 +150,7 @@ def test_astar_python_and_cpp(costmap) -> None: # Assert that both implementations return almost identical points. np.testing.assert_allclose( - [(p.position.x, p.position.y) for p in path_python.poses], - [(p.position.x, p.position.y) for p in path_cpp.poses], + [(p.pose.position.x, p.pose.position.y) for p in path_python.poses], + [(p.pose.position.x, p.pose.position.y) for p in path_cpp.poses], atol=0.05001, ) diff --git a/dimos/navigation/local_planner/rust/Cargo.toml b/dimos/navigation/local_planner/rust/Cargo.toml index 6c753eeb247..c3d8e20297a 100644 --- a/dimos/navigation/local_planner/rust/Cargo.toml +++ b/dimos/navigation/local_planner/rust/Cargo.toml @@ -22,7 +22,7 @@ serde_json = { version = "1", features = ["float_roundtrip"] } pyo3 = { version = "0.25", features = ["extension-module", "abi3-py310"], optional = true } numpy = { version = "0.25", optional = true } dimos-module = { path = "../../../../native/rust/dimos-module", optional = true } -dimos-generated-messages = { path = "../../../../packages/dimos-generated/src/dimos_generated_schemas/package/rust", optional = true } +dimos-generated-messages = { path = "../../../../packages/dimos-generated/build/message-source/rust", optional = true } tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal", "time"], optional = true } tracing = { version = "0.1", optional = true } tracing-subscriber = { version = "0.3", features = ["env-filter"], optional = true } diff --git a/dimos/navigation/movement_manager/test_movement_manager.py b/dimos/navigation/movement_manager/test_movement_manager.py index 0c860e0d2ad..497e77b0500 100644 --- a/dimos/navigation/movement_manager/test_movement_manager.py +++ b/dimos/navigation/movement_manager/test_movement_manager.py @@ -19,6 +19,7 @@ import math import time +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import Point, PointStamped, Twist, Vector3 import pytest @@ -44,14 +45,16 @@ def _attach(module): captured = Captured() unsubs = [ module.cmd_vel.subscribe( - lambda msg: captured.cmd_vel.append(type(msg).decode(msg.encode())) + lambda msg: captured.cmd_vel.append(cdr_decode(cdr_encode(msg), type(msg))) ), module.stop_movement.subscribe( - lambda msg: captured.stop_movement.append(type(msg).decode(msg.encode())) + lambda msg: captured.stop_movement.append(cdr_decode(cdr_encode(msg), type(msg))) + ), + module.goal.subscribe( + lambda msg: captured.goal.append(cdr_decode(cdr_encode(msg), type(msg))) ), - module.goal.subscribe(lambda msg: captured.goal.append(type(msg).decode(msg.encode()))), module.way_point.subscribe( - lambda msg: captured.way_point.append(type(msg).decode(msg.encode())) + lambda msg: captured.way_point.append(cdr_decode(cdr_encode(msg), type(msg))) ), ] return captured, unsubs diff --git a/dimos/navigation/spec.py b/dimos/navigation/spec.py index 6f16a7c5964..24dc6235734 100644 --- a/dimos/navigation/spec.py +++ b/dimos/navigation/spec.py @@ -25,12 +25,12 @@ from typing import Protocol +from dimos_generated.geometry_msgs.msg import PointStamped, Twist +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.tf2_msgs.msg import TFMessage + from dimos.core.stream import IO, In, Out -from dimos.msgs.geometry_msgs.PointStamped import PointStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage class GlobalPlanner(Protocol): diff --git a/dimos/navigation/test_generated_bbox_navigation.py b/dimos/navigation/test_generated_bbox_navigation.py index 43ab745a57d..954d3d14b5f 100644 --- a/dimos/navigation/test_generated_bbox_navigation.py +++ b/dimos/navigation/test_generated_bbox_navigation.py @@ -15,9 +15,10 @@ from threading import Event from uuid import uuid4 +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import PoseStamped -from dimos_generated.sensor_msgs.msg import CameraInfo +from dimos_generated.sensor_msgs.msg import CameraInfo, RegionOfInterest from dimos_generated.std_msgs.msg import Header from dimos_generated.vision_msgs.msg import ( BoundingBox2D, @@ -26,10 +27,11 @@ Point2D, Pose2D, ) +import numpy as np import pytest +from dimos.agents.skills.visual_servoing.bbox_navigation import BBoxNavigationModule from dimos.core.transport import LCMTransport -from dimos.navigation.bbox_navigation import BBoxNavigationModule def test_generated_bbox_goal_preserves_header_and_projection(): @@ -41,22 +43,46 @@ def test_generated_bbox_goal_preserves_header_and_projection(): ready = Event() received = Event() unsubscribe = module.goal_request.subscribe( - lambda value: (goals.append(PoseStamped.decode(value.encode())), received.set()) + lambda value: ( + goals.append(cdr_decode(cdr_encode(value), PoseStamped)), + received.set(), + ) ) message = Detection2DArray( header=Header(frame_id="base_link", stamp=Time(sec=1700000000, nanosec=123456789)), - detections=[Detection2D(bbox=BoundingBox2D(center=Pose2D(position=Point2D(x=420, y=290))))], + detections=[ + Detection2D( + bbox=BoundingBox2D( + center=Pose2D(position=Point2D(x=420, y=290), theta=0.0), size_x=0.0, size_y=0.0 + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + results=[], + id="", + ) + ], ) try: module.start() module._on_detection(message) assert not goals - camera = CameraInfo(k=[500, 0, 320, 0, 250, 240, 0, 0, 1]) + camera = CameraInfo( + k=np.array([500, 0, 320, 0, 250, 240, 0, 0, 1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + height=0, + width=0, + distortion_model="", + d=np.array([], dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), + ) camera_unsubscribe = module.camera_info.subscribe(lambda _: ready.set()) - module.camera_info.transport.publish(CameraInfo.decode(camera.encode())) + module.camera_info.transport.publish(cdr_decode(cdr_encode(camera), CameraInfo)) assert ready.wait(2) camera_unsubscribe() - module.detection2d.transport.publish(Detection2DArray.decode(message.encode())) + module.detection2d.transport.publish(cdr_decode(cdr_encode(message), Detection2DArray)) assert received.wait(2) assert len(goals) == 1 goal = goals[0] @@ -64,7 +90,11 @@ def test_generated_bbox_goal_preserves_header_and_projection(): assert (goal.pose.position.x, goal.pose.position.y, goal.pose.position.z) == pytest.approx( (2, -0.4, -0.4) ) - module._on_detection(Detection2DArray()) + module._on_detection( + Detection2DArray( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), detections=[] + ) + ) assert len(goals) == 1 finally: unsubscribe() diff --git a/dimos/navigation/trajectory_follower/basic/module.py b/dimos/navigation/trajectory_follower/basic/module.py index 9b8ed5dd5c0..ed17e394482 100644 --- a/dimos/navigation/trajectory_follower/basic/module.py +++ b/dimos/navigation/trajectory_follower/basic/module.py @@ -19,7 +19,10 @@ import time from typing import Any -from dimos_lcm.std_msgs import Bool +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist, Vector3 +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Bool +from dimos_generated.tf2_msgs.msg import TFMessage import numpy as np from numpy.typing import NDArray from reactivex.disposable import Disposable @@ -28,11 +31,7 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import IO, In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.geometry import quaternion_euler from dimos.navigation import spec from dimos.utils.logging_config import setup_logger from dimos.utils.trigonometry import angle_diff @@ -103,7 +102,9 @@ def stop(self) -> None: self._stop_event.set() if self._thread is not None: self._thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) - self.nav_cmd_vel.publish(Twist()) + self.nav_cmd_vel.publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) super().stop() def _lookup_pose(self) -> PoseStamped | None: @@ -120,9 +121,13 @@ def _on_path(self, path: Path) -> None: if len(path.poses) == 0: with self._lock: self._waypoints = None - self.nav_cmd_vel.publish(Twist()) + self.nav_cmd_vel.publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) return - waypoints = np.array([[p.position.x, p.position.y] for p in path.poses], dtype=np.float32) + waypoints = np.array( + [[p.pose.position.x, p.pose.position.y] for p in path.poses], dtype=np.float32 + ) with self._lock: self._waypoints = waypoints @@ -130,7 +135,9 @@ def _on_stop(self, msg: Bool) -> None: if msg.data: with self._lock: self._waypoints = None - self.nav_cmd_vel.publish(Twist()) + self.nav_cmd_vel.publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) def _follow(self) -> None: period = 1.0 / self.config.control_frequency @@ -145,20 +152,22 @@ def _follow(self) -> None: self._stop_event.wait(max(0.0, period - elapsed)) def _step(self, pose: PoseStamped, waypoints: NDArray[np.float32]) -> None: - position = np.array([pose.position.x, pose.position.y], dtype=np.float32) + position = np.array([pose.pose.position.x, pose.pose.position.y], dtype=np.float32) if float(np.linalg.norm(waypoints[-1] - position)) < self.config.goal_tolerance: - self.nav_cmd_vel.publish(Twist()) + self.nav_cmd_vel.publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) with self._lock: if self._waypoints is waypoints: self._waypoints = None - self.goal_reached.publish(Bool(True)) + self.goal_reached.publish(Bool(data=True)) logger.info("Goal reached") return target = self._lookahead_point(waypoints, position) yaw_error = angle_diff( math.atan2(target[1] - position[1], target[0] - position[0]), - pose.orientation.euler[2], + quaternion_euler(pose.pose.orientation)[2], ) angular = max( @@ -166,7 +175,9 @@ def _step(self, pose: PoseStamped, waypoints: NDArray[np.float32]) -> None: min(self.config.max_angular, self.config.heading_gain * yaw_error), ) linear = self.config.speed * max(0.0, math.cos(yaw_error)) - self.nav_cmd_vel.publish(Twist(Vector3(linear, 0, 0), Vector3(0, 0, angular))) + self.nav_cmd_vel.publish( + Twist(linear=Vector3(x=linear, y=0.0, z=0.0), angular=Vector3(z=angular, x=0.0, y=0.0)) + ) def _lookahead_point( self, waypoints: NDArray[np.float32], position: NDArray[np.float32] diff --git a/dimos/navigation/trajectory_follower/basic/test_module.py b/dimos/navigation/trajectory_follower/basic/test_module.py index 9cf403aba33..9649b0411d1 100644 --- a/dimos/navigation/trajectory_follower/basic/test_module.py +++ b/dimos/navigation/trajectory_follower/basic/test_module.py @@ -12,9 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Twist, + Vector3, +) +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.std_msgs.msg import Bool, Header + from dimos.navigation.trajectory_follower.basic.module import BasicPathFollower, lookahead_distance from dimos.protocol.tf.tf import MultiTBuffer @@ -36,7 +48,7 @@ def get( time_tolerance: float | None = None, *, forward_tolerance: float = 0.0, - ) -> Transform | None: + ) -> TransformStamped | None: self.gets += 1 return super().get( parent_frame, @@ -50,23 +62,24 @@ def dispose(self) -> None: pass -def _mount() -> Transform: - return Transform( - translation=Vector3(0.0, 0.0, MOUNT_Z), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="base_link", +def _mount() -> TransformStamped: + return TransformStamped( + header=Header(frame_id="base_link", stamp=Time(sec=1, nanosec=0)), child_frame_id="mid360_link", - ts=1.0, + transform=Transform( + translation=Vector3(z=MOUNT_Z, x=0.0, y=0.0), + rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), ) -def _odom_edge() -> Transform: - return Transform( - translation=Vector3(1.0, 2.0, 3.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="odom", +def _odom_edge() -> TransformStamped: + return TransformStamped( + header=Header(frame_id="odom", stamp=Time(sec=1, nanosec=0)), child_frame_id="mid360_link", - ts=1.0, + transform=Transform( + translation=Vector3(x=1, y=2, z=3), rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), ) @@ -79,7 +92,7 @@ def test_lookup_pose_steers_from_the_tf_base_pose() -> None: try: pose = module._lookup_pose() assert pose is not None - assert abs(pose.position.z - (3.0 - MOUNT_Z)) < 1e-9 + assert abs(pose.pose.position.z - (3.0 - MOUNT_Z)) < 1e-9 finally: module.stop() @@ -121,3 +134,44 @@ def test_lookahead_scales_in_linear_region() -> None: def test_lookahead_clamped_at_ceiling() -> None: assert lookahead_distance(2.0, 1.5, 0.4, 1.5) == 1.5 + + +def test_generated_path_produces_velocity_and_arrival(): + module = BasicPathFollower() + commands = [] + arrivals = [] + module.nav_cmd_vel.subscribe( + lambda value: commands.append(cdr_decode(cdr_encode(value), Twist)) + ) + module.goal_reached.subscribe( + lambda value: arrivals.append(cdr_decode(cdr_encode(value), Bool)) + ) + try: + path = Path( + poses=[ + PoseStamped( + pose=Pose( + position=Point(x=x, y=0.0, z=0.0), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + for x in [0, 1, 2] + ], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + module._on_path(cdr_decode(cdr_encode(path), Path)) + waypoints = module._waypoints + assert waypoints is not None + module._step(path.poses[0], waypoints) + assert commands[-1].linear.x > 0 + assert not arrivals + module._step(path.poses[-1], waypoints) + assert commands[-1] == Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) + assert arrivals[-1].data + module._on_path(Path(header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), poses=[])) + assert module._waypoints is None + finally: + module.stop() diff --git a/dimos/navigation/trajectory_follower/fancy/rust/Cargo.toml b/dimos/navigation/trajectory_follower/fancy/rust/Cargo.toml index 032706f59e9..579f502af53 100644 --- a/dimos/navigation/trajectory_follower/fancy/rust/Cargo.toml +++ b/dimos/navigation/trajectory_follower/fancy/rust/Cargo.toml @@ -21,7 +21,7 @@ required-features = ["module"] dimos-local-planner = { path = "../../../local_planner/rust", default-features = false } serde = { version = "1", features = ["derive"], optional = true } dimos-module = { path = "../../../../../native/rust/dimos-module", optional = true } -dimos-generated-messages = { path = "../../../../../packages/dimos-generated/src/dimos_generated_schemas/package/rust", optional = true } +dimos-generated-messages = { path = "../../../../../packages/dimos-generated/build/message-source/rust", optional = true } tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal", "time"], optional = true } tracing = { version = "0.1", optional = true } tracing-subscriber = { version = "0.3", features = ["env-filter"], optional = true } diff --git a/dimos/perception/depth2depth_cloud/rust/Cargo.toml b/dimos/perception/depth2depth_cloud/rust/Cargo.toml index cb6d5b6a4d8..53e3bd79c98 100644 --- a/dimos/perception/depth2depth_cloud/rust/Cargo.toml +++ b/dimos/perception/depth2depth_cloud/rust/Cargo.toml @@ -33,7 +33,7 @@ depth_cloud = "sensor_msgs/msg/PointCloud2" [dependencies] dimos-module = { path = "../../../../native/rust/dimos-module" } -dimos-generated-messages = { path = "../../../../packages/dimos-generated/src/dimos_generated_schemas/package/rust" } +dimos-generated-messages = { path = "../../../../packages/dimos-generated/build/message-source/rust" } tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal"] } serde = { version = "1", features = ["derive"] } tracing = "0.1" diff --git a/dimos/perception/depth2depth_cloud/rust/flake.nix b/dimos/perception/depth2depth_cloud/rust/flake.nix index a504cd98b3f..5490cdeb64c 100644 --- a/dimos/perception/depth2depth_cloud/rust/flake.nix +++ b/dimos/perception/depth2depth_cloud/rust/flake.nix @@ -30,6 +30,8 @@ }; cudaPackages = if system == "aarch64-linux" then pkgs.cudaPackages_12_6 else pkgs.cudaPackages_12_8; + messageSource = import ../../../../native/message-source.nix { inherit pkgs; }; + src = pkgs.runCommand "depth2depth-cloud-src" {} '' mkdir -p $out/dimos/perception/depth2depth_cloud/rust cp -r ${./src} $out/dimos/perception/depth2depth_cloud/rust/src @@ -40,8 +42,8 @@ cp -r ${../../../../native/rust/dimos-module} $out/native/rust/dimos-module cp -r ${../../../../native/rust/dimos-module-macros} $out/native/rust/dimos-module-macros cp -r ${../../../../native/rust/dimos-lcm-transport} $out/native/rust/dimos-lcm-transport - mkdir -p $out/packages/dimos-generated/src/dimos_generated_schemas/package - cp -r ${../../../../packages/dimos-generated/src/dimos_generated_schemas/package/rust} $out/packages/dimos-generated/src/dimos_generated_schemas/package/rust + mkdir -p $out/packages/dimos-generated/build/message-source + cp -r ${messageSource.source}/rust $out/packages/dimos-generated/build/message-source/rust ''; generatedCargoNix = crate2nix.tools.${system}.generatedCargoNix { diff --git a/dimos/perception/detection/conftest.py b/dimos/perception/detection/conftest.py index ff065a0e1ea..16bb8ec8d6b 100644 --- a/dimos/perception/detection/conftest.py +++ b/dimos/perception/detection/conftest.py @@ -18,6 +18,7 @@ import json from typing import TypedDict +from dimos_cdrgen.registry import decode as cdr_decode from dimos_generated.geometry_msgs.msg import PoseStamped, TransformStamped from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 from dimos_generated.vision_msgs.msg import Detection2DArray @@ -78,10 +79,10 @@ def payload(name): assert hashlib.sha256(data).hexdigest() == stream["cdr_sha256"] return data - lidar_frame = PointCloud2.decode(payload("lidar")) - image_frame = Image.decode(payload("video")) + lidar_frame = cdr_decode(payload("lidar"), PointCloud2) + image_frame = cdr_decode(payload("video"), Image) image_frame.header.frame_id = "camera_optical" - odom_frame = PoseStamped.decode(payload("odom")) + odom_frame = cdr_decode(payload("odom"), PoseStamped) # Match the Go2 connection's configured global parent for this recording. odom_frame.header.frame_id = "world" diff --git a/dimos/perception/detection/detectors/grounding_dino.py b/dimos/perception/detection/detectors/grounding_dino.py index f025acf470d..38d0ed0d20e 100644 --- a/dimos/perception/detection/detectors/grounding_dino.py +++ b/dimos/perception/detection/detectors/grounding_dino.py @@ -18,12 +18,14 @@ from functools import cached_property +from dimos_generated.sensor_msgs.msg import Image import numpy as np from PIL import Image as PILImage import torch from dimos.models.base import HuggingFaceModel, HuggingFaceModelConfig -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_to_rgb +from dimos.msgs.time import to_seconds from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D @@ -83,7 +85,7 @@ def query_detections( class_id=label, confidence=float(row[label]), name=queries[label], - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, ) if det.is_valid(): @@ -107,7 +109,7 @@ def query_score_rows( from torchvision.ops.boxes import batched_nms # type: ignore[import-untyped] from transformers.image_transforms import center_to_corners_format - pil = PILImage.fromarray(image.to_rgb().data) + pil = PILImage.fromarray(image_to_rgb(image)) phrases = [" ".join(query.replace(".", " ").lower().split()) for query in queries] caption = " ".join(f"{phrase}." for phrase in phrases) with torch.inference_mode(): diff --git a/dimos/perception/detection/detectors/owlv2.py b/dimos/perception/detection/detectors/owlv2.py index b1cd5490803..8ba9a2d2002 100644 --- a/dimos/perception/detection/detectors/owlv2.py +++ b/dimos/perception/detection/detectors/owlv2.py @@ -19,12 +19,14 @@ from collections import OrderedDict from functools import cached_property +from dimos_generated.sensor_msgs.msg import Image import numpy as np from PIL import Image as PILImage import torch from dimos.models.base import HuggingFaceModel, HuggingFaceModelConfig -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_to_rgb +from dimos.msgs.time import to_seconds from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D @@ -111,7 +113,7 @@ def query_detections_batch( results are per-image, in input order. """ self.forwards += len(images) - pils = [PILImage.fromarray(image.to_rgb().data) for image in images] + pils = [PILImage.fromarray(image_to_rgb(image)) for image in images] with torch.inference_mode(), self._autocast(): inputs = self._processor( text=[queries] * len(pils), images=pils, return_tensors="pt" @@ -138,7 +140,7 @@ def query_detections_batch( class_id=int(label), confidence=float(score), name=queries[int(label)], - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, ) if det.is_valid(): @@ -172,7 +174,7 @@ def _frame_features( logit shift and scale, and clipped pixel ``(x1, y1, x2, y2)`` boxes. A miss runs the image tower once; a hit runs nothing. """ - key = (image.frame_id, image.ts) + key = (image.header.frame_id, to_seconds(image.header.stamp)) if key in self._features: self._features.move_to_end(key) return self._features[key] @@ -180,7 +182,7 @@ def _frame_features( from transformers.image_transforms import center_to_corners_format self.forwards += 1 - pil = PILImage.fromarray(image.to_rgb().data) + pil = PILImage.fromarray(image_to_rgb(image)) model = self._model head = model.class_head with torch.inference_mode(), self._autocast(): diff --git a/dimos/perception/detection/module2D.py b/dimos/perception/detection/module2D.py index 696b16ae2a3..671073e5e7b 100644 --- a/dimos/perception/detection/module2D.py +++ b/dimos/perception/detection/module2D.py @@ -117,7 +117,7 @@ def track(self, detections: ImageDetections2D) -> None: ) else: # No detection at this index - publish zero transform - position_3d = Vector3() + position_3d = Vector3(x=0.0, y=0.0, z=0.0) transforms.append( TransformStamped( @@ -125,7 +125,9 @@ def track(self, detections: ImageDetections2D) -> None: frame_id=sensor_frame.child_frame_id, stamp=detections.image.header.stamp ), child_frame_id=f"det_{index}", - transform=Transform(translation=position_3d, rotation=Quaternion(w=1)), + transform=Transform( + translation=position_3d, rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), ) ) diff --git a/dimos/perception/detection/module3D.py b/dimos/perception/detection/module3D.py index f3a53e531ab..d456e7662a3 100644 --- a/dimos/perception/detection/module3D.py +++ b/dimos/perception/detection/module3D.py @@ -166,7 +166,7 @@ def nav_vlm(self, question: str) -> PoseStamped | None: header=Header(stamp=detections.image.header.stamp, frame_id="world"), pose=Pose( position=Point(x=position.x, y=position.y, z=position.z), - orientation=Quaternion(w=1), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), ), ) diff --git a/dimos/perception/detection/person_tracker.py b/dimos/perception/detection/person_tracker.py index 0fb287856c4..4b8ace2c8a3 100644 --- a/dimos/perception/detection/person_tracker.py +++ b/dimos/perception/detection/person_tracker.py @@ -128,7 +128,8 @@ def track(self, detections2D: ImageDetections2D) -> None: pose_in_camera = PoseStamped( header=Header(frame_id="camera_link", stamp=detections2D.image.header.stamp), pose=Pose( - position=Point(x=vector.x, y=vector.y, z=vector.z), orientation=Quaternion(w=1) + position=Point(x=vector.x, y=vector.y, z=vector.z), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), ), ) diff --git a/dimos/perception/detection/reid/test_generated_module.py b/dimos/perception/detection/reid/test_generated_module.py index 1443f2d0c9d..6ed422aa767 100644 --- a/dimos/perception/detection/reid/test_generated_module.py +++ b/dimos/perception/detection/reid/test_generated_module.py @@ -59,7 +59,7 @@ def receive(value: object) -> None: header = Header(stamp=time_from_seconds(1.25), frame_id="camera") image = image_from_array(np.zeros((32, 32, 3), dtype=np.uint8), encoding="rgb8", header=header) try: - detections.on_next(Detection2DArray(header=header)) + detections.on_next(Detection2DArray(header=header, detections=[])) detections.on_next( Detection2DArray( header=header, @@ -68,8 +68,11 @@ def receive(value: object) -> None: header=header, id="7", bbox=BoundingBox2D( - center=Pose2D(position=Point2D(x=16, y=16)), size_x=8, size_y=8 + center=Pose2D(position=Point2D(x=16, y=16), theta=0.0), + size_x=8, + size_y=8, ), + results=[], ) ], ) diff --git a/dimos/perception/detection/type/detection2d/bbox.py b/dimos/perception/detection/type/detection2d/bbox.py index 301fc1b4bf8..734c1b0d0ec 100644 --- a/dimos/perception/detection/type/detection2d/bbox.py +++ b/dimos/perception/detection/type/detection2d/bbox.py @@ -18,6 +18,8 @@ import hashlib from typing import TYPE_CHECKING, Any +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseWithCovariance, Quaternion + if TYPE_CHECKING: from dimos_generated.sensor_msgs.msg import Image from typing_extensions import Self @@ -329,7 +331,14 @@ def to_ros_detection2d(self) -> ROSDetection2D: hypothesis=ObjectHypothesis( class_id=str(self.class_id), score=self.confidence, - ) + ), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), ) ], id=str(self.track_id), diff --git a/dimos/perception/detection/type/detection2d/point.py b/dimos/perception/detection/type/detection2d/point.py index 4042fd76ea0..da827474f51 100644 --- a/dimos/perception/detection/type/detection2d/point.py +++ b/dimos/perception/detection/type/detection2d/point.py @@ -17,6 +17,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseWithCovariance, Quaternion from dimos_generated.vision_msgs.msg import ( BoundingBox2D, Detection2D as ROSDetection2D, @@ -25,6 +26,7 @@ Point2D, Pose2D, ) +import numpy as np from dimos.msgs.image import image_from_array, image_view from dimos.perception.detection.type.detection2d.base import Detection2D @@ -89,7 +91,14 @@ def to_ros_detection2d(self) -> ROSDetection2D: hypothesis=ObjectHypothesis( class_id=str(self.class_id), score=self.confidence, - ) + ), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), ) ], id=str(self.track_id), diff --git a/dimos/perception/detection/type/detection2d/test_generated_bbox.py b/dimos/perception/detection/type/detection2d/test_generated_bbox.py index 56648f90c71..393b33a5fee 100644 --- a/dimos/perception/detection/type/detection2d/test_generated_bbox.py +++ b/dimos/perception/detection/type/detection2d/test_generated_bbox.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.std_msgs.msg import Header from dimos_generated.vision_msgs.msg import Detection2D, Detection2DArray @@ -47,7 +48,7 @@ def test_generated_bbox_image_and_wire_round_trip(): assert annotated.header == header np.testing.assert_array_equal(image_view(annotated)[39, 59], [0, 0, 250]) np.testing.assert_array_equal(image_view(image), pixels) - wire = Detection2D.decode(det.to_ros_detection2d().encode()) + wire = cdr_decode(cdr_encode(det.to_ros_detection2d()), Detection2D) assert wire.header == header assert wire.results[0].hypothesis.class_id == "2" restored = Detection2DBBox.from_ros_detection2d(wire, image=image) @@ -59,7 +60,7 @@ def test_generated_bbox_image_and_wire_round_trip(): def test_generated_detection_collection_round_trip(): pixels = np.zeros((40, 60, 3), dtype=np.uint8) image = image_from_array( - pixels, encoding="rgb8", header=Header(frame_id="camera", stamp=Time(nanosec=123)) + pixels, encoding="rgb8", header=Header(frame_id="camera", stamp=Time(nanosec=123, sec=0)) ) det = Detection2DBBox( bbox=(10, 10, 30, 30), @@ -71,7 +72,7 @@ def test_generated_detection_collection_round_trip(): image=image, ) collection = ImageDetections2D(image=image, detections=[det]) - wire = Detection2DArray.decode(collection.to_ros_detection2d_array().encode()) + wire = cdr_decode(cdr_encode(collection.to_ros_detection2d_array()), Detection2DArray) assert wire.header == image.header assert len(wire.detections) == 1 restored = ImageDetections2D.from_ros_detection2d_array(image, wire) @@ -87,14 +88,14 @@ def test_generated_detection_collection_round_trip(): def test_generated_point_detection_crop_and_cdr(): pixels = np.arange(100, dtype=np.uint8).reshape(10, 10) image = image_from_array( - pixels, encoding="mono8", header=Header(frame_id="camera", stamp=Time(nanosec=789)) + pixels, encoding="mono8", header=Header(frame_id="camera", stamp=Time(nanosec=789, sec=0)) ) point = Detection2DPoint(x=1, y=1, name="point", ts=0, image=image) assert point.is_valid() crop = point.cropped_image(padding=2) np.testing.assert_array_equal(image_view(crop), pixels[:3, :3]) assert crop.header == image.header - wire = Detection2D.decode(point.to_ros_detection2d().encode()) + wire = cdr_decode(cdr_encode(point.to_ros_detection2d()), Detection2D) assert wire.header == image.header assert wire.bbox.center.position.x == wire.bbox.center.position.y == 1 assert wire.bbox.size_x == wire.bbox.size_y == 0 diff --git a/dimos/perception/detection/type/detection3d/base.py b/dimos/perception/detection/type/detection3d/base.py index f3b95df4d50..f2f8b378ac5 100644 --- a/dimos/perception/detection/type/detection3d/base.py +++ b/dimos/perception/detection/type/detection3d/base.py @@ -18,7 +18,9 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING -from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.std_msgs.msg import Header from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox @@ -32,7 +34,14 @@ class Detection3D(Detection2DBBox): frame_id: str = "" transform: TransformStamped = field( - default_factory=lambda: TransformStamped(transform=Transform(rotation=Quaternion(w=1))) + default_factory=lambda: TransformStamped( + transform=Transform( + rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + translation=Vector3(x=0.0, y=0.0, z=0.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + child_frame_id="", + ) ) @classmethod diff --git a/dimos/perception/detection/type/detection3d/bbox.py b/dimos/perception/detection/type/detection3d/bbox.py index e12078effe7..94b47c885ce 100644 --- a/dimos/perception/detection/type/detection3d/bbox.py +++ b/dimos/perception/detection/type/detection3d/bbox.py @@ -18,10 +18,12 @@ import functools from typing import Any +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, Pose, PoseStamped, + PoseWithCovariance, Quaternion, TransformStamped, Vector3, @@ -33,6 +35,7 @@ ObjectHypothesis, ObjectHypothesisWithPose, ) +import numpy as np from dimos.msgs.geometry import point_distance from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox @@ -49,7 +52,7 @@ class Detection3DBBox(Detection2DBBox): size: Vector3 # Width, height, depth transform: TransformStamped | None = None # Camera to world transform frame_id: str = "" # Frame ID (e.g., "world", "map") - orientation: Quaternion = field(default_factory=lambda: Quaternion(w=1)) + orientation: Quaternion = field(default_factory=lambda: Quaternion(w=1, x=0.0, y=0.0, z=0.0)) @functools.cached_property def pose(self) -> PoseStamped: @@ -67,7 +70,18 @@ def pose(self) -> PoseStamped: def to_detection3d_msg(self) -> Detection3D: """Convert to ROS Detection3D message.""" - msg = Detection3D() + msg = Detection3D( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + results=[], + bbox=BoundingBox3D( + center=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + size=Vector3(x=0.0, y=0.0, z=0.0), + ), + id="", + ) msg.header = Header(stamp=self.image.header.stamp, frame_id=self.frame_id) # Results @@ -76,7 +90,14 @@ def to_detection3d_msg(self) -> Detection3D: hypothesis=ObjectHypothesis( class_id=str(self.class_id), score=self.confidence, - ) + ), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), ) ] diff --git a/dimos/perception/detection/type/detection3d/imageDetections3DPC.py b/dimos/perception/detection/type/detection3d/imageDetections3DPC.py index e45ba357aed..015fcfcd0f6 100644 --- a/dimos/perception/detection/type/detection3d/imageDetections3DPC.py +++ b/dimos/perception/detection/type/detection3d/imageDetections3DPC.py @@ -16,6 +16,7 @@ from typing import TYPE_CHECKING +from dimos.msgs.time import to_seconds from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC, lattice_quantum from dimos.perception.detection.type.imageDetections import ImageDetections @@ -59,8 +60,8 @@ def from_2d( points_2d, camera_info, world_to_optical_transform, - world_pointcloud.frame_id, - world_pointcloud.ts, + world_pointcloud.header.frame_id, + to_seconds(world_pointcloud.header.stamp), filters, splat_m=quantum / 2 if quantum is not None else None, ) diff --git a/dimos/perception/detection/type/detection3d/pointcloud.py b/dimos/perception/detection/type/detection3d/pointcloud.py index d1d2d6eb75b..239b816f91a 100644 --- a/dimos/perception/detection/type/detection3d/pointcloud.py +++ b/dimos/perception/detection/type/detection3d/pointcloud.py @@ -18,12 +18,29 @@ import functools from typing import TYPE_CHECKING, Any +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + TransformStamped, + Vector3, +) +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header import numpy as np - -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from numpy.typing import NDArray + +from dimos.msgs.geometry import inverse_transform, transform_matrix +from dimos.msgs.image import image_view +from dimos.msgs.pointcloud import ( + concatenate_clouds, + pointcloud_from_xyz, + pointcloud_xyz, + select_points, +) +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.perception.detection.type.detection3d.base import Detection3D from dimos.perception.detection.type.detection3d.pointcloud_filters import ( PointCloudFilter, @@ -33,9 +50,9 @@ ) if TYPE_CHECKING: - from dimos_lcm.sensor_msgs import CameraInfo + from dimos_generated.sensor_msgs.msg import CameraInfo, Image + import open3d as o3d - from dimos.msgs.sensor_msgs.Image import Image from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox @@ -65,11 +82,24 @@ def lattice_quantum(points: np.ndarray) -> float | None: @dataclass class Detection3DPC(Detection3D): - pointcloud: PointCloud2 = field(default_factory=PointCloud2) + pointcloud: PointCloud2 = field( + default_factory=lambda: PointCloud2( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + height=0, + width=0, + fields=[], + is_bigendian=False, + point_step=0, + row_step=0, + data=np.array([], dtype=np.uint8), + is_dense=False, + ) + ) @functools.cached_property def center(self) -> Vector3: - return Vector3(*self.pointcloud.center) + center = pointcloud_xyz(self.pointcloud).mean(axis=0) + return Vector3(x=float(center[0]), y=float(center[1]), z=float(center[2])) def __add__(self, other: Detection3DPC) -> Detection3DPC: """Union of two sightings of one object. @@ -88,7 +118,7 @@ def __add__(self, other: Detection3DPC) -> Detection3DPC: confidence=stronger.confidence, name=stronger.name, ts=later.ts, - pointcloud=self.pointcloud + other.pointcloud, + pointcloud=concatenate_clouds(self.pointcloud, other.pointcloud), transform=later.transform, frame_id=self.frame_id, ) @@ -101,36 +131,53 @@ def pose(self) -> PoseStamped: The pointcloud is already in world frame. """ return PoseStamped( - ts=self.ts, - frame_id=self.frame_id, - position=self.center, - orientation=(0.0, 0.0, 0.0, 1.0), # Identity quaternion + header=self.pointcloud.header, + pose=Pose( + position=Point(x=self.center.x, y=self.center.y, z=self.center.z), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), ) - def get_bounding_box(self): # type: ignore[no-untyped-def] + def _open3d(self) -> o3d.geometry.PointCloud: + import open3d as o3d + + return o3d.geometry.PointCloud(o3d.utility.Vector3dVector(pointcloud_xyz(self.pointcloud))) + + def get_bounding_box(self) -> o3d.geometry.AxisAlignedBoundingBox: """Get axis-aligned bounding box of the detection's pointcloud.""" - return self.pointcloud.axis_aligned_bounding_box + return self._open3d().get_axis_aligned_bounding_box() - def get_oriented_bounding_box(self): # type: ignore[no-untyped-def] + def get_oriented_bounding_box(self) -> o3d.geometry.OrientedBoundingBox: """Get oriented bounding box of the detection's pointcloud.""" - return self.pointcloud.oriented_bounding_box + return self._open3d().get_oriented_bounding_box() def get_bounding_box_dimensions(self) -> tuple[float, float, float]: """Get dimensions (width, height, depth) of the detection's bounding box.""" - return self.pointcloud.bounding_box_dimensions + extent = self.get_bounding_box().get_extent() + return float(extent[0]), float(extent[1]), float(extent[2]) def bounding_box_intersects(self, other: Detection3DPC) -> bool: """Check if this detection's bounding box intersects with another's.""" - return self.pointcloud.bounding_box_intersects(other.pointcloud) + first, second = self.get_bounding_box(), other.get_bounding_box() + return bool( + np.all(first.get_min_bound() <= second.get_max_bound()) + and np.all(second.get_min_bound() <= first.get_max_bound()) + ) def to_repr_dict(self) -> dict[str, Any]: # Calculate distance from camera # The pointcloud is in world frame, and transform gives camera position in world center_world = self.center # Camera position in world frame is the translation part of the transform - camera_pos = self.transform.translation + camera_pos = inverse_transform(self.transform).transform.translation # Use Vector3 subtraction and magnitude - distance = (center_world - camera_pos).magnitude() + distance = np.linalg.norm( + [ + center_world.x - camera_pos.x, + center_world.y - camera_pos.y, + center_world.z - camera_pos.z, + ] + ) parent_dict = super().to_repr_dict() # Remove bbox key if present @@ -139,7 +186,7 @@ def to_repr_dict(self) -> dict[str, Any]: return { **parent_dict, "dist": f"{distance:.2f}m", - "points": str(len(self.pointcloud)), + "points": str(self.pointcloud.width * self.pointcloud.height), } @classmethod @@ -148,7 +195,7 @@ def from_depth( det: Detection2DBBox, depth: Image, camera_info: CameraInfo, - world_to_optical_transform: Transform, + world_to_optical_transform: TransformStamped, filters: list[PointCloudFilter] | None = None, max_depth: float = 10.0, depth_gap: float = 0.1, @@ -163,13 +210,16 @@ def from_depth( segmentation mask is eroded to ``mask_scale`` of its size first, since the bleed lives on the mask boundary. """ + import cv2 + # no radius_outlier: dense depth clouds make radius search expensive, # and the depth-gap cluster above already drops disconnected points if filters is None: filters = [statistical()] - depth_m = np.asarray(depth.data, dtype=np.float32) - if depth.data.dtype == np.uint16: + pixels = image_view(depth) + depth_m = pixels.astype(np.float32) + if pixels.dtype.kind == "u" and pixels.dtype.itemsize == 2: depth_m *= 0.001 height, width = depth_m.shape[:2] @@ -177,8 +227,6 @@ def from_depth( if seg_mask is not None: pixel_mask = seg_mask > 0 if mask_scale < 1.0: - import cv2 - radius = float(np.sqrt(pixel_mask.sum() / np.pi)) erode_px = round((1.0 - mask_scale) * radius) if erode_px > 0: @@ -214,12 +262,16 @@ def from_depth( rows, cols, z = rows[keep], cols[keep], z[keep] break - fx, fy = camera_info.K[0], camera_info.K[4] - cx, cy = camera_info.K[2], camera_info.K[5] + fx, fy = camera_info.k[0], camera_info.k[4] + cx, cy = camera_info.k[2], camera_info.k[5] points_optical = np.column_stack(((cols - cx) * z / fx, (rows - cy) * z / fy, z)) - detection_pc = PointCloud2.from_numpy(points_optical, timestamp=det.ts).transform( - -world_to_optical_transform + optical_to_world = inverse_transform(world_to_optical_transform) + matrix = transform_matrix(optical_to_world.transform) + points_world = points_optical @ matrix[:3, :3].T + matrix[:3, 3] + detection_pc = pointcloud_from_xyz( + points_world, + header=Header(stamp=depth.header.stamp, frame_id=optical_to_world.header.frame_id), ) for filter_func in filters: @@ -228,7 +280,7 @@ def from_depth( return None detection_pc = result - if len(detection_pc.pointcloud.points) == 0: + if detection_pc.width * detection_pc.height == 0: return None return cls( @@ -241,7 +293,7 @@ def from_depth( ts=det.ts, pointcloud=detection_pc, transform=world_to_optical_transform, - frame_id=detection_pc.frame_id, + frame_id=detection_pc.header.frame_id, ) @staticmethod @@ -253,10 +305,10 @@ def project_pixels(points_camera: np.ndarray, camera_info: CameraInfo) -> np.nda an undistorted calibration falls through to the pinhole projection. Points beyond the model's monotonic radius get out-of-image pixels. """ - fx, fy = camera_info.K[0], camera_info.K[4] - cx, cy = camera_info.K[2], camera_info.K[5] + fx, fy = camera_info.k[0], camera_info.k[4] + cx, cy = camera_info.k[2], camera_info.k[5] coefficients = np.asarray( - camera_info.D if camera_info.D is not None else (), dtype=np.float64 + camera_info.d if camera_info.d is not None else (), dtype=np.float64 ) xy = points_camera[:, :2] / points_camera[:, 2:3] if coefficients.size == 0 or not np.any(coefficients): @@ -317,28 +369,42 @@ def project_pixels(points_camera: np.ndarray, camera_info: CameraInfo) -> np.nda def project_cloud( world_pointcloud: PointCloud2, camera_info: CameraInfo, - world_to_optical_transform: Transform, + world_to_optical_transform: TransformStamped, ) -> tuple[np.ndarray, np.ndarray]: + """Project a frame once for the existing batched detection API.""" + points, pixels, _ = Detection3DPC._project_cloud_indices( + world_pointcloud, camera_info, world_to_optical_transform + ) + return points, pixels + + @staticmethod + def _project_cloud_indices( + world_pointcloud: PointCloud2, + camera_info: CameraInfo, + world_to_optical_transform: TransformStamped, + ) -> tuple[np.ndarray, np.ndarray, NDArray[np.int64]]: """Project a world cloud through the camera once. Returns the world points that land inside the image and their pixel coordinates - the detection-independent half of ``from_2d``, shared by every detection of one frame via ``from_projection``. """ - world_points, _ = world_pointcloud.as_numpy() + world_points = pointcloud_xyz(world_pointcloud) + indices = np.arange(len(world_points), dtype=np.int64) # Project points to camera frame points_homogeneous = np.hstack([world_points, np.ones((world_points.shape[0], 1))]) - extrinsics_matrix = world_to_optical_transform.to_matrix() + extrinsics_matrix = transform_matrix(world_to_optical_transform.transform) points_camera = (extrinsics_matrix @ points_homogeneous.T).T # Filter out points behind the camera valid_mask = points_camera[:, 2] > 0 points_camera = points_camera[valid_mask] world_points = world_points[valid_mask] + indices = indices[valid_mask] if len(world_points) == 0: - return world_points, np.empty((0, 2)) + return world_points, np.empty((0, 2)), indices points_2d = Detection3DPC.project_pixels(points_camera, camera_info) @@ -349,7 +415,7 @@ def project_cloud( & (points_2d[:, 1] >= 0) & (points_2d[:, 1] < camera_info.height) ) - return world_points[in_image_mask], points_2d[in_image_mask] + return world_points[in_image_mask], points_2d[in_image_mask], indices[in_image_mask] @classmethod def from_projection( @@ -358,11 +424,13 @@ def from_projection( world_points: np.ndarray, points_2d: np.ndarray, camera_info: CameraInfo, - world_to_optical_transform: Transform, + world_to_optical_transform: TransformStamped, frame_id: str, timestamp: float, filters: list[PointCloudFilter] | None = None, splat_m: float | None = None, + *, + projected_cloud: PointCloud2 | None = None, ) -> Detection3DPC | None: """Create a Detection3D by selecting from a shared frame projection. @@ -399,9 +467,11 @@ def mask_hit(cols_f: np.ndarray, rows_f: np.ndarray) -> np.ndarray: in_det_mask = mask_hit(points_2d[:, 0], points_2d[:, 1]) if splat_m is not None: - camera = world_to_optical_transform.inverse().translation.to_numpy() + camera = transform_matrix(inverse_transform(world_to_optical_transform).transform)[ + :3, 3 + ] ranges = np.linalg.norm(world_points - camera, axis=1) - radius = splat_m * camera_info.K[0] / np.maximum(ranges, 1e-6) + radius = splat_m * camera_info.k[0] / np.maximum(ranges, 1e-6) for dx, dy in ((1, 0), (-1, 0), (0, 1), (0, -1)): in_det_mask |= mask_hit( points_2d[:, 0] + dx * radius, points_2d[:, 1] + dy * radius @@ -422,10 +492,13 @@ def mask_hit(cols_f: np.ndarray, rows_f: np.ndarray) -> np.ndarray: return None # Create initial pointcloud for this detection - initial_pc = PointCloud2.from_numpy( - detection_points, - frame_id=frame_id, - timestamp=timestamp, + initial_pc = ( + select_points(projected_cloud, in_det_mask) + if projected_cloud is not None + else pointcloud_from_xyz( + detection_points, + header=Header(frame_id=frame_id, stamp=time_from_seconds(timestamp)), + ) ) # Apply filters - each filter gets all arguments @@ -437,7 +510,7 @@ def mask_hit(cols_f: np.ndarray, rows_f: np.ndarray) -> np.ndarray: detection_pc = result # Final check for empty pointcloud - if len(detection_pc.pointcloud.points) == 0: + if detection_pc.width * detection_pc.height == 0: return None # Create Detection3D with filtered pointcloud @@ -460,7 +533,7 @@ def from_2d( # type: ignore[override] det: Detection2DBBox, world_pointcloud: PointCloud2, camera_info: CameraInfo, - world_to_optical_transform: Transform, + world_to_optical_transform: TransformStamped, # filters are to be adjusted based on the sensor noise characteristics if feeding # sensor data directly filters: list[PointCloudFilter] | None = None, @@ -473,9 +546,12 @@ def from_2d( # type: ignore[override] radius comes from the cloud's own lattice pitch, so every caller samples a grid-quantized source the same way. """ - world_points, points_2d = cls.project_cloud( + world_points, points_2d, indices = cls._project_cloud_indices( world_pointcloud, camera_info, world_to_optical_transform ) + keep = np.zeros(world_pointcloud.width * world_pointcloud.height, dtype=bool) + keep[indices] = True + projected_cloud = select_points(world_pointcloud, keep) quantum = lattice_quantum(world_points) return cls.from_projection( det, @@ -483,8 +559,9 @@ def from_2d( # type: ignore[override] points_2d, camera_info, world_to_optical_transform, - world_pointcloud.frame_id, - world_pointcloud.ts, + world_pointcloud.header.frame_id, + to_seconds(world_pointcloud.header.stamp), filters, splat_m=quantum / 2 if quantum is not None else None, + projected_cloud=projected_cloud, ) diff --git a/dimos/perception/detection/type/detection3d/pointcloud_filters.py b/dimos/perception/detection/type/detection3d/pointcloud_filters.py index 1c83be49d03..25e53e1b3e8 100644 --- a/dimos/perception/detection/type/detection3d/pointcloud_filters.py +++ b/dimos/perception/detection/type/detection3d/pointcloud_filters.py @@ -21,7 +21,7 @@ from dimos_generated.sensor_msgs.msg import CameraInfo, PointCloud2 import numpy as np -from dimos.msgs.geometry import inverse_transform +from dimos.msgs.geometry import inverse_transform, transform_matrix from dimos.msgs.pointcloud import pointcloud_xyz, select_points if TYPE_CHECKING: @@ -72,15 +72,14 @@ def range_cluster(gap: float = 0.3) -> PointCloudFilter: points a mask collects across a range discontinuity are background seen through or around the object, not the object. """ - import numpy as np def filter_func( - det: Detection2DBBox, pc: PointCloud2, ci: CameraInfo, tf: Transform + det: Detection2DBBox, pc: PointCloud2, ci: CameraInfo, tf: TransformStamped ) -> PointCloud2 | None: - points, _ = pc.as_numpy() + points = pointcloud_xyz(pc) if len(points) == 0: return None - camera = tf.inverse().translation.to_numpy() + camera = transform_matrix(inverse_transform(tf).transform)[:3, 3] ranges = np.linalg.norm(points - camera, axis=1) order = np.argsort(ranges) ranges_sorted = ranges[order] @@ -91,7 +90,7 @@ def filter_func( for start, end in zip(starts, ends, strict=False): if start <= median_idx < end: keep = order[start:end] - return PointCloud2.from_numpy(points[keep], frame_id=pc.frame_id, timestamp=pc.ts) + return _selected(pc, keep.tolist()) return pc return filter_func diff --git a/dimos/perception/detection/type/detection3d/test_generated_filters.py b/dimos/perception/detection/type/detection3d/test_generated_filters.py index 7ae5a158ec2..b3a679652be 100644 --- a/dimos/perception/detection/type/detection3d/test_generated_filters.py +++ b/dimos/perception/detection/type/detection3d/test_generated_filters.py @@ -14,9 +14,10 @@ from types import SimpleNamespace +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time -from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped -from dimos_generated.sensor_msgs.msg import CameraInfo, PointCloud2, PointField +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo, PointCloud2, PointField, RegionOfInterest from dimos_generated.std_msgs.msg import Header import numpy as np import pytest @@ -49,18 +50,37 @@ def test_filters_preserve_custom_point_records(make_filter): PointField(name=name, offset=i * 4, datatype=7 if i < 3 else 6, count=1) for i, name in enumerate(("x", "y", "z", "id")) ], - data=values.tobytes(), + data=np.frombuffer(values.tobytes(), dtype=np.uint8), + is_bigendian=False, + is_dense=False, ) transform = TransformStamped( - header=Header(frame_id="camera"), + header=Header(frame_id="camera", stamp=Time(sec=0, nanosec=0)), child_frame_id="world", - transform=Transform(rotation=Quaternion(w=1)), + transform=Transform( + rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), translation=Vector3(x=0.0, y=0.0, z=0.0) + ), ) result = make_filter()( - SimpleNamespace(), PointCloud2.decode(cloud.encode()), CameraInfo(), transform + SimpleNamespace(), + cdr_decode(cdr_encode(cloud), PointCloud2), + CameraInfo( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + height=0, + width=0, + distortion_model="", + d=np.array([], dtype=np.float64), + k=np.zeros(9, dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), + ), + transform, ) assert result is not None - result = PointCloud2.decode(result.encode()) + result = cdr_decode(cdr_encode(result), PointCloud2) assert result.header == cloud.header assert result.fields == cloud.fields records = pointcloud_view(result).reshape(-1) diff --git a/dimos/perception/detection/type/detection3d/test_generated_projection.py b/dimos/perception/detection/type/detection3d/test_generated_projection.py index ea903cb0a59..8a2aea7cf0d 100644 --- a/dimos/perception/detection/type/detection3d/test_generated_projection.py +++ b/dimos/perception/detection/type/detection3d/test_generated_projection.py @@ -14,6 +14,7 @@ from types import SimpleNamespace +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Quaternion, @@ -22,14 +23,14 @@ Twist, Vector3, ) -from dimos_generated.sensor_msgs.msg import CameraInfo +from dimos_generated.sensor_msgs.msg import CameraInfo, RegionOfInterest from dimos_generated.std_msgs.msg import Header import numpy as np import pytest +from dimos.agents.skills.visual_servoing.detection_navigation import DetectionNavigation from dimos.msgs.image import image_from_array from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_xyz -from dimos.navigation.visual_servoing.detection_navigation import DetectionNavigation from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC @@ -50,11 +51,25 @@ def test_depth_projection_preserves_header_and_transforms(): depth = image_from_array( np.full((4, 4), 2000, dtype=np.uint16), encoding="16UC1", header=header ) - camera = CameraInfo(width=4, height=4, k=[2, 0, 0, 0, 2, 0, 0, 0, 1]) + camera = CameraInfo( + width=4, + height=4, + k=np.array([2, 0, 0, 0, 2, 0, 0, 0, 1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + distortion_model="", + d=np.array([], dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), + ) tf = TransformStamped( - header=Header(frame_id="camera"), + header=Header(frame_id="camera", stamp=Time(sec=0, nanosec=0)), child_frame_id="world", - transform=Transform(translation=Vector3(x=-10), rotation=Quaternion(w=1)), + transform=Transform( + translation=Vector3(x=-10, y=0.0, z=0.0), rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), ) result = Detection3DPC.from_depth(detection(depth), depth, camera, tf, filters=[]) assert result is not None @@ -69,13 +84,28 @@ def test_depth_projection_preserves_header_and_transforms(): def test_cloud_projection_rejects_outside_and_behind_points(): cloud = pointcloud_from_xyz( - np.array([[0, 0, 2], [1, 1, 2], [20, 20, 2], [0, 0, -1]]), header=Header(frame_id="world") + np.array([[0, 0, 2], [1, 1, 2], [20, 20, 2], [0, 0, -1]]), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + ) + camera = CameraInfo( + width=4, + height=4, + k=np.array([2, 0, 0, 0, 2, 0, 0, 0, 1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + distortion_model="", + d=np.array([], dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), ) - camera = CameraInfo(width=4, height=4, k=[2, 0, 0, 0, 2, 0, 0, 0, 1]) tf = TransformStamped( - header=Header(frame_id="camera"), + header=Header(frame_id="camera", stamp=Time(sec=0, nanosec=0)), child_frame_id="world", - transform=Transform(rotation=Quaternion(w=1)), + transform=Transform( + rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), translation=Vector3(x=0.0, y=0.0, z=0.0) + ), ) result = Detection3DPC.from_2d(detection(None), cloud, camera, tf, filters=[]) assert result is not None @@ -85,13 +115,36 @@ def test_cloud_projection_rejects_outside_and_behind_points(): def test_generated_cloud_target_produces_cdr_navigation_twist(): points = np.array([[2.0 + i * 0.01, 0, 1] for i in range(40)]) - cloud = pointcloud_from_xyz(points, header=Header(frame_id="world")) - controller = DetectionNavigation(None, CameraInfo()) - target = controller._compute_robust_target_position(cloud, Vector3()) + cloud = pointcloud_from_xyz( + points, header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)) + ) + controller = DetectionNavigation( + None, + CameraInfo( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + height=0, + width=0, + distortion_model="", + d=np.array([], dtype=np.float64), + k=np.zeros(9, dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), + ), + ) + target = controller._compute_robust_target_position(cloud, Vector3(x=0.0, y=0.0, z=0.0)) assert target is not None assert 2.0 <= target.x < 2.1 - tf = TransformStamped(transform=Transform(rotation=Quaternion(w=1))) + tf = TransformStamped( + transform=Transform( + rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), translation=Vector3(x=0.0, y=0.0, z=0.0) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + child_frame_id="", + ) command = controller._compute_twist_from_3d(target, tf) - decoded = Twist.decode(command.encode()) + decoded = cdr_decode(cdr_encode(command), Twist) assert 0 < decoded.linear.x <= 0.5 assert decoded.angular.z == 0 diff --git a/dimos/perception/detection/type/detection3d/test_imageDetections3D.py b/dimos/perception/detection/type/detection3d/test_imageDetections3D.py index 85dcd1aabdf..84155abcad9 100644 --- a/dimos/perception/detection/type/detection3d/test_imageDetections3D.py +++ b/dimos/perception/detection/type/detection3d/test_imageDetections3D.py @@ -14,6 +14,7 @@ from __future__ import annotations +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import Quaternion, Vector3 from dimos_generated.sensor_msgs.msg import Image from dimos_generated.std_msgs.msg import Header @@ -90,7 +91,7 @@ def test_to_ros_detection3d_array_serializes_plain_bbox_results() -> None: assert len(msg.detections) == 1 assert len(msg.detections[0].results) == 1 - decoded = Detection3DArray.decode(msg.encode()) + decoded = cdr_decode(cdr_encode(msg), Detection3DArray) decoded_det = decoded.detections[0] assert len(decoded_det.results) == 1 assert decoded_det.results[0].hypothesis.class_id == "9" @@ -126,7 +127,7 @@ def test_to_ros_detection3d_array_preserves_marker_wire_identity() -> None: assert first.bbox.size.x == pytest.approx(0.16) assert first.bbox.size.z == pytest.approx(0.0) - decoded = Detection3DArray.decode(msg.encode()) + decoded = cdr_decode(cdr_encode(msg), Detection3DArray) assert decoded.header.frame_id == "world" assert len(decoded.detections) == 2 assert decoded.detections[1].id == "42" @@ -180,7 +181,7 @@ def test_marker_pose_and_array_keep_exact_source_stamp() -> None: pose = marker.pose message = marker.to_detection3d_msg() array = ImageDetections3D(image, [marker]).to_ros_detection3d_array() - decoded = Detection3DArray.decode(array.encode()) + decoded = cdr_decode(cdr_encode(array), Detection3DArray) assert pose.header.stamp == image.header.stamp assert message.header.stamp == image.header.stamp assert decoded.header.stamp == image.header.stamp diff --git a/dimos/perception/detection/type/detection3d/test_marker.py b/dimos/perception/detection/type/detection3d/test_marker.py index 6d9cb9e4e43..49052848458 100644 --- a/dimos/perception/detection/type/detection3d/test_marker.py +++ b/dimos/perception/detection/type/detection3d/test_marker.py @@ -13,6 +13,7 @@ # limitations under the License. import cv2 +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import Quaternion, Vector3 from dimos_generated.std_msgs.msg import Header from dimos_generated.vision_msgs.msg import Detection3D @@ -114,7 +115,7 @@ def test_marker_detection3d_msg_preserves_marker_identity_on_wire() -> None: assert msg.bbox.size.y == pytest.approx(size.y) assert msg.bbox.size.z == pytest.approx(size.z) - decoded = Detection3D.decode(msg.encode()) + decoded = cdr_decode(cdr_encode(msg), Detection3D) assert decoded.id == "42" assert len(decoded.results) == 1 assert decoded.results[0].hypothesis.class_id == "DICT_APRILTAG_36h11:42" diff --git a/dimos/perception/detection/type/detection3d/test_pointcloud.py b/dimos/perception/detection/type/detection3d/test_pointcloud.py index 74b28a09b38..55c0ea3feca 100644 --- a/dimos/perception/detection/type/detection3d/test_pointcloud.py +++ b/dimos/perception/detection/type/detection3d/test_pointcloud.py @@ -13,13 +13,16 @@ # limitations under the License. import cv2 +from dimos_generated.geometry_msgs.msg import TransformStamped +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2, RegionOfInterest +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.geometry import transform_from_matrix +from dimos.msgs.image import image_from_array +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.detection.type.detection2d.seg import Detection2DSeg @@ -31,15 +34,44 @@ PITCH = 0.05 K = [100.0, 0.0, 320.0, 0.0, 100.0, 240.0, 0.0, 0.0, 1.0] -CAMERA_INFO = CameraInfo(height=480, width=640, K=K) + + +def _camera(distortion=()): + return CameraInfo( + header=Header(stamp=time_from_seconds(1.0), frame_id="camera_optical"), + height=480, + width=640, + distortion_model="plumb_bob", + d=np.asarray(distortion, dtype=np.float64), + k=np.asarray(K, dtype=np.float64), + r=np.eye(3).reshape(-1), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), + ) + + +CAMERA_INFO = _camera() + + +def _identity(): + return TransformStamped( + header=Header(stamp=time_from_seconds(1.0), frame_id="camera_optical"), + child_frame_id="world", + transform=transform_from_matrix(np.eye(4)), + ) + + +def _cloud(points, ts=1.0): + return pointcloud_from_xyz(points, header=Header(stamp=time_from_seconds(ts), frame_id="world")) def _image() -> Image: - return Image( - data=np.zeros((480, 640, 3), dtype=np.uint8), - format=ImageFormat.BGR, - frame_id="camera_optical", - ts=1.0, + return image_from_array( + np.zeros((480, 640, 3), dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id="camera_optical", stamp=time_from_seconds(1.0)), ) @@ -57,7 +89,7 @@ def _sighting(ts: float, confidence: float, name: str, points: int) -> Detection confidence=confidence, name=name, ts=ts, - pointcloud=PointCloud2.from_numpy(np.full((points, 3), ts), frame_id="world", timestamp=ts), + pointcloud=_cloud(np.full((points, 3), ts), ts), frame_id="world", ) @@ -103,10 +135,10 @@ def test_detection3dpc(detection3dpc) -> None: # def test_point_cloud_properties(detection3dpc): """Test point cloud data and boundaries.""" - points, _ = detection3dpc.pointcloud.as_numpy() + points = pointcloud_xyz(detection3dpc.pointcloud) assert len(points) > 60 - assert detection3dpc.pointcloud.frame_id == "world", ( - f"Expected frame_id 'world', got '{detection3dpc.pointcloud.frame_id}'" + assert detection3dpc.pointcloud.header.frame_id == "world", ( + f"Expected frame_id 'world', got '{detection3dpc.pointcloud.header.frame_id}'" ) min_pt = np.min(points, axis=0) @@ -128,11 +160,11 @@ def test_detection3dpc(detection3dpc) -> None: # def test_detection_pose(detection3dpc): """Test detection pose and frame information.""" - assert detection3dpc.pose.x == pytest.approx(-3.327, abs=0.1) - assert detection3dpc.pose.y == pytest.approx(-0.202, abs=0.1) - assert detection3dpc.pose.z == pytest.approx(0.160, abs=0.1) - assert detection3dpc.pose.frame_id == "world", ( - f"Expected frame_id 'world', got '{detection3dpc.pose.frame_id}'" + assert detection3dpc.pose.pose.position.x == pytest.approx(-3.327, abs=0.1) + assert detection3dpc.pose.pose.position.y == pytest.approx(-0.202, abs=0.1) + assert detection3dpc.pose.pose.position.z == pytest.approx(0.160, abs=0.1) + assert detection3dpc.pose.header.frame_id == "world", ( + f"Expected frame_id 'world', got '{detection3dpc.pose.header.frame_id}'" ) @@ -160,7 +192,7 @@ def test_add() -> None: merged = strong + later - assert len(merged.pointcloud.as_numpy()[0]) == 10 + assert len(pointcloud_xyz(merged.pointcloud)) == 10 assert (merged.name, merged.confidence) == ("cup", 0.9) assert (merged.ts, merged.bbox) == (2.0, later.bbox) @@ -171,7 +203,7 @@ def test_add() -> None: def test_project_pixels(D: list[float]) -> None: # matches OpenCV with and without distortion points = np.array([[0.3, -0.2, 1.0], [-0.6, 0.4, 1.5]]) - info = CameraInfo(height=480, width=640, distortion_model="plumb_bob", D=D, K=K) + info = _camera(D) expected, _ = cv2.projectPoints( points, np.zeros(3), np.zeros(3), np.reshape(K, (3, 3)), np.array(D) ) @@ -181,9 +213,9 @@ def test_project_pixels(D: list[float]) -> None: def test_project_cloud_empty() -> None: # nothing in front of the camera - cloud = PointCloud2.from_numpy(np.array([[0.0, 0.0, -1.0]]), frame_id="world", timestamp=1.0) + cloud = _cloud(np.array([[0.0, 0.0, -1.0]])) - world_points, pixels = Detection3DPC.project_cloud(cloud, CAMERA_INFO, Transform.identity()) + world_points, pixels = Detection3DPC.project_cloud(cloud, CAMERA_INFO, _identity()) assert len(world_points) == len(pixels) == 0 @@ -199,21 +231,21 @@ def test_from_2d_splat() -> None: class_id=0, confidence=0.9, name="stripe", - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, mask=mask, ) - cloud = PointCloud2.from_numpy(_lattice(), frame_id="world", timestamp=1.0) + cloud = _cloud(_lattice()) lifted = ImageDetections3DPC.from_2d( ImageDetections2D(image=image, detections=[stripe]), cloud, CAMERA_INFO, - Transform.identity(), + _identity(), [], ) - columns = np.unique(np.round(lifted.detections[0].pointcloud.as_numpy()[0][:, 0], 3)) + columns = np.unique(np.round(pointcloud_xyz(lifted.detections[0].pointcloud)[:, 0], 3)) np.testing.assert_allclose(columns, [0.0, PITCH]) @@ -232,10 +264,10 @@ def test_range_cluster() -> None: ) def keep(points: np.ndarray) -> PointCloud2 | None: - cloud = PointCloud2.from_numpy(points, frame_id="world", timestamp=1.0) - return range_cluster()(det, cloud, CAMERA_INFO, Transform.identity()) + cloud = _cloud(points) + return range_cluster()(det, cloud, CAMERA_INFO, _identity()) kept = keep(np.vstack((near, far))) assert kept is not None - np.testing.assert_allclose(np.sort(kept.as_numpy()[0][:, 2]), near[:, 2], rtol=1e-6) + np.testing.assert_allclose(np.sort(pointcloud_xyz(kept)[:, 2]), near[:, 2], rtol=1e-6) assert keep(np.empty((0, 3))) is None diff --git a/dimos/perception/experimental/moduleDB.py b/dimos/perception/experimental/moduleDB.py index 4f5d6ddc266..e82bd1739d7 100644 --- a/dimos/perception/experimental/moduleDB.py +++ b/dimos/perception/experimental/moduleDB.py @@ -129,7 +129,7 @@ def to_pose(self) -> PoseStamped: ), pose=Pose( position=Point(x=self.center.x, y=self.center.y, z=self.center.z), - orientation=Quaternion(w=1), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), ), ) diff --git a/dimos/perception/experimental/object.py b/dimos/perception/experimental/object.py index 2c568201481..d11c0f6d764 100644 --- a/dimos/perception/experimental/object.py +++ b/dimos/perception/experimental/object.py @@ -19,10 +19,12 @@ from typing import TYPE_CHECKING, Any import uuid +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, Pose, PoseStamped, + PoseWithCovariance, Quaternion, TransformStamped, Vector3, @@ -158,7 +160,18 @@ def to_detection3d_msg(self) -> ROSDetection3D: """Convert to ROS Detection3D message.""" center, orientation, size = self._detection3d_bbox_components() - msg = ROSDetection3D() + msg = ROSDetection3D( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + results=[], + bbox=BoundingBox3D( + center=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + size=Vector3(x=0.0, y=0.0, z=0.0), + ), + id="", + ) msg.header = self.pose.header msg.id = self.object_id msg.results = [ @@ -166,7 +179,14 @@ def to_detection3d_msg(self) -> ROSDetection3D: hypothesis=ObjectHypothesis( class_id=self.name, score=self.confidence, - ) + ), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), ) ] msg.bbox = BoundingBox3D( @@ -339,7 +359,7 @@ def from_2d_to_list( aabb_extent = aabb.max_bound - aabb.min_bound center = Vector3(x=aabb_center[0], y=aabb_center[1], z=aabb_center[2]) sx, sy, sz = float(aabb_extent[0]), float(aabb_extent[1]), float(aabb_extent[2]) - orientation = Quaternion(w=1.0) + orientation = Quaternion(w=1.0, x=0.0, y=0.0, z=0.0) else: obb = pcd_filtered.get_oriented_bounding_box() center = Vector3(x=obb.center[0], y=obb.center[1], z=obb.center[2]) @@ -397,7 +417,7 @@ def aggregate_pointclouds(objects: list[Object]) -> PointCloud2: Combined PointCloud2 with all points colored by object (empty if no points). """ header = ( - Header() + Header(stamp=Time(sec=0, nanosec=0), frame_id="") if not objects else Header(frame_id=objects[0].frame_id, stamp=objects[0].pointcloud.header.stamp) ) diff --git a/dimos/perception/experimental/object_tracker.py b/dimos/perception/experimental/object_tracker.py index 045619a24dd..fa34e8c00fb 100644 --- a/dimos/perception/experimental/object_tracker.py +++ b/dimos/perception/experimental/object_tracker.py @@ -16,9 +16,11 @@ import time from typing import Any +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, Pose, + PoseWithCovariance, Quaternion, Transform, TransformStamped, @@ -30,11 +32,16 @@ # Generated detection values from dimos_generated.vision_msgs.msg import ( + BoundingBox2D, + BoundingBox3D, Detection2D, Detection2DArray, Detection3D, Detection3DArray, + ObjectHypothesis, ObjectHypothesisWithPose, + Point2D, + Pose2D, ) import numpy as np from numpy.typing import NDArray @@ -321,8 +328,12 @@ def _reset_tracking_state(self) -> None: self.tracking_frame_count = 0 # Reset frame counter # Publish empty detections to clear any visualizations - empty_2d = Detection2DArray(header=Header(), detections=[]) - empty_3d = Detection3DArray(header=Header(), detections=[]) + empty_2d = Detection2DArray( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), detections=[] + ) + empty_3d = Detection3DArray( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), detections=[] + ) self._latest_detection2d = empty_2d self._latest_detection3d = empty_3d self._detection_event.clear() @@ -433,12 +444,28 @@ def _process_tracking(self) -> None: height = float(y2 - y1) # Create Detection2D - detection_2d = Detection2D() + detection_2d = Detection2D( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + results=[], + bbox=BoundingBox2D( + center=Pose2D(position=Point2D(x=0.0, y=0.0), theta=0.0), size_x=0.0, size_y=0.0 + ), + id="", + ) detection_2d.id = "0" detection_2d.header = header # Create hypothesis - hypothesis = ObjectHypothesisWithPose() + hypothesis = ObjectHypothesisWithPose( + hypothesis=ObjectHypothesis(class_id="", score=0.0), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), + ) hypothesis.hypothesis.class_id = "tracked_object" hypothesis.hypothesis.score = 1.0 detection_2d.results = [hypothesis] @@ -450,7 +477,9 @@ def _process_tracking(self) -> None: detection_2d.bbox.size_x = width detection_2d.bbox.size_y = height - detection2darray = Detection2DArray() + detection2darray = Detection2DArray( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), detections=[] + ) detection2darray.header = header detection2darray.detections = [detection_2d] @@ -471,16 +500,21 @@ def _process_tracking(self) -> None: y_optical = (center_y - cy) * z_optical / fy # Create pose in optical frame - optical_pose = Pose() + optical_pose = Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ) optical_pose.position = Point(x=x_optical, y=y_optical, z=z_optical) - optical_pose.orientation = Quaternion(w=1.0) # Identity for now + optical_pose.orientation = Quaternion( + w=1.0, x=0.0, y=0.0, z=0.0 + ) # Identity for now # Convert to robot frame robot_pose = optical_to_robot_frame(optical_pose) # Calculate orientation: object facing towards camera (origin) yaw = yaw_towards_point(robot_pose.position) - euler = Vector3(z=yaw) # Only yaw, no roll/pitch + euler = Vector3(z=yaw, x=0.0, y=0.0) # Only yaw, no roll/pitch robot_pose.orientation = euler_to_quaternion(euler) # Estimate object size in meters @@ -489,7 +523,18 @@ def _process_tracking(self) -> None: size_z = 0.1 # Default depth size # Create Detection3D - detection_3d = Detection3D() + detection_3d = Detection3D( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + results=[], + bbox=BoundingBox3D( + center=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + size=Vector3(x=0.0, y=0.0, z=0.0), + ), + id="", + ) detection_3d.id = "0" detection_3d.header = header @@ -497,12 +542,17 @@ def _process_tracking(self) -> None: detection_3d.results = [hypothesis] # Create 3D bounding box with robot frame pose - detection_3d.bbox.center = Pose() + detection_3d.bbox.center = Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ) detection_3d.bbox.center.position = robot_pose.position detection_3d.bbox.center.orientation = robot_pose.orientation detection_3d.bbox.size = Vector3(x=size_x, y=size_y, z=size_z) - detection3darray = Detection3DArray() + detection3darray = Detection3DArray( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), detections=[] + ) detection3darray.header = header detection3darray.detections = [detection_3d] diff --git a/dimos/perception/experimental/object_tracker_2d.py b/dimos/perception/experimental/object_tracker_2d.py index abb4b413729..e5119593b3a 100644 --- a/dimos/perception/experimental/object_tracker_2d.py +++ b/dimos/perception/experimental/object_tracker_2d.py @@ -17,6 +17,7 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseWithCovariance, Quaternion from dimos_generated.sensor_msgs.msg import Image from dimos_generated.std_msgs.msg import Header @@ -277,7 +278,14 @@ def _process_tracking(self) -> None: ), results=[ ObjectHypothesisWithPose( - hypothesis=ObjectHypothesis(class_id="tracked_object", score=1.0) + hypothesis=ObjectHypothesis(class_id="tracked_object", score=1.0), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), ) ], ) diff --git a/dimos/perception/experimental/object_tracker_3d.py b/dimos/perception/experimental/object_tracker_3d.py index 5377f259874..3d92fa95164 100644 --- a/dimos/perception/experimental/object_tracker_3d.py +++ b/dimos/perception/experimental/object_tracker_3d.py @@ -14,9 +14,11 @@ # Generated detection values +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, Pose, + PoseWithCovariance, Quaternion, Transform, TransformStamped, @@ -26,9 +28,11 @@ from dimos_generated.std_msgs.msg import Header from dimos_generated.tf2_msgs.msg import TFMessage from dimos_generated.vision_msgs.msg import ( + BoundingBox3D, Detection2DArray, Detection3D, Detection3DArray, + ObjectHypothesis, ObjectHypothesisWithPose, ) import numpy as np @@ -210,16 +214,18 @@ def _create_detection3d_from_2d(self, detection2d: Detection2DArray) -> Detectio y_optical = (center_y - cy) * z_optical / fy # type: ignore[has-type] # Create pose in optical frame - optical_pose = Pose() + optical_pose = Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) optical_pose.position = Point(x=x_optical, y=y_optical, z=z_optical) - optical_pose.orientation = Quaternion(w=1.0) + optical_pose.orientation = Quaternion(w=1.0, x=0.0, y=0.0, z=0.0) # Convert to robot frame robot_pose = optical_to_robot_frame(optical_pose) # Calculate orientation: object facing towards camera yaw = yaw_towards_point(robot_pose.position) - euler = Vector3(z=yaw) + euler = Vector3(z=yaw, x=0.0, y=0.0) robot_pose.orientation = euler_to_quaternion(euler) # Estimate object size in meters @@ -229,23 +235,47 @@ def _create_detection3d_from_2d(self, detection2d: Detection2DArray) -> Detectio # Create Detection3D header = Header(stamp=detection2d.header.stamp, frame_id=self.config.frame_id) - detection_3d = Detection3D() + detection_3d = Detection3D( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + results=[], + bbox=BoundingBox3D( + center=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + size=Vector3(x=0.0, y=0.0, z=0.0), + ), + id="", + ) detection_3d.id = "0" detection_3d.header = header # Create hypothesis - hypothesis = ObjectHypothesisWithPose() + hypothesis = ObjectHypothesisWithPose( + hypothesis=ObjectHypothesis(class_id="", score=0.0), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), + ) hypothesis.hypothesis.class_id = "tracked_object" hypothesis.hypothesis.score = 1.0 detection_3d.results = [hypothesis] # Create 3D bounding box - detection_3d.bbox.center = Pose() + detection_3d.bbox.center = Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) detection_3d.bbox.center.position = robot_pose.position detection_3d.bbox.center.orientation = robot_pose.orientation detection_3d.bbox.size = Vector3(x=size_x, y=size_y, z=size_z) - detection3darray = Detection3DArray() + detection3darray = Detection3DArray( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), detections=[] + ) detection3darray.header = header detection3darray.detections = [detection_3d] diff --git a/dimos/perception/experimental/perceive_loop_skill.py b/dimos/perception/experimental/perceive_loop_skill.py index 65b2a300e6f..a81d2ebe885 100644 --- a/dimos/perception/experimental/perceive_loop_skill.py +++ b/dimos/perception/experimental/perceive_loop_skill.py @@ -14,21 +14,24 @@ from __future__ import annotations +import base64 from datetime import datetime, timezone import json import os from threading import RLock from typing import TYPE_CHECKING, Any +from dimos_generated.sensor_msgs.msg import Image + from dimos.agents.agent_spec import AgentSpec from dimos.agents.annotation import skill from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import In from dimos.models.vl.create import create -from dimos.msgs.sensor_msgs.Image import Image, sharpness_window +from dimos.msgs.image import image_sharpness, image_to_bgr, image_to_jpeg from dimos.utils.logging_config import setup_logger -from dimos.utils.reactive import backpressure +from dimos.utils.reactive import backpressure, quality_barrier from dimos.utils.threadpool import run_in_thread if TYPE_CHECKING: @@ -114,7 +117,9 @@ def look_out_for( model = self._vl_model.result() sharpest = backpressure( - sharpness_window(1.0 / self._period, self.color_image.pure_observable()) + self.color_image.pure_observable().pipe( + quality_barrier(image_sharpness, 1.0 / self._period) + ) ) model.start() self._model_started = True @@ -181,7 +186,7 @@ def _on_image(self, image: Image) -> None: continuation_context: dict[str, Any] = { "bbox": list(best.bbox), "label": best.name, - "image": image.to_base64(quality=70), + "image": base64.b64encode(image_to_jpeg(image, quality=70)).decode("ascii"), } logger.info( "Lookout matched, dispatching continuation", @@ -208,7 +213,7 @@ def _write_debug_image(image: Image, detections: ImageDetections2D[Detection2DBB import cv2 try: - debug_img = image.to_opencv().copy() + debug_img = image_to_bgr(image) for det in detections.detections: x1, y1, x2, y2 = (int(v) for v in det.bbox) cv2.rectangle(debug_img, (x1, y1), (x2, y2), (0, 255, 0), 2) diff --git a/dimos/perception/experimental/temporal_memory/temporal_memory.py b/dimos/perception/experimental/temporal_memory/temporal_memory.py index 24054dc648a..98c87c1d6ea 100644 --- a/dimos/perception/experimental/temporal_memory/temporal_memory.py +++ b/dimos/perception/experimental/temporal_memory/temporal_memory.py @@ -29,6 +29,9 @@ import time from typing import Any +from dimos_generated.dimos_msgs.msg import EntityMarker, EntityMarkers +from dimos_generated.geometry_msgs.msg import Point, PoseStamped +from dimos_generated.sensor_msgs.msg import Image from reactivex import Subject, interval from reactivex.disposable import Disposable @@ -38,9 +41,8 @@ from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out from dimos.models.vl.base import VlModel -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.Image import Image, sharpness_barrier -from dimos.msgs.visualization_msgs.EntityMarkers import EntityMarkers, Marker +from dimos.msgs.image import image_sharpness +from dimos.msgs.time import header_now from dimos.perception.experimental.temporal_memory.clip_filter import ( CLIP_AVAILABLE, adaptive_keyframes, @@ -57,6 +59,7 @@ from dimos.perception.experimental.temporal_memory.temporal_utils.helpers import is_scene_stale from dimos.perception.experimental.temporal_memory.window_analyzer import WindowAnalyzer from dimos.utils.logging_config import get_run_log_dir, setup_logger +from dimos.utils.reactive import quality_barrier try: from dimos.perception.experimental.temporal_memory.clip_filter import CLIPFrameFilter @@ -255,7 +258,7 @@ def _publish_entity_markers(self) -> None: if not all_entities: return - markers: list[Marker] = [] + markers: list[EntityMarker] = [] for e in all_entities: meta = e.get("metadata") or {} x = meta.get("world_x") @@ -264,18 +267,18 @@ def _publish_entity_markers(self) -> None: if x is None or y is None: continue markers.append( - Marker( + EntityMarker( entity_id=e["entity_id"], label=(e.get("descriptor") or "")[:40], entity_type=e.get("entity_type", "object"), - x=x, - y=y, - z=(z or 0.0) + 0.3, # Offset up so labels float above ground + position=Point(x=x, y=y, z=(z or 0.0) + 0.3), # Offset labels above ground ) ) if markers: - self.entity_markers.publish(EntityMarkers(markers=markers)) + self.entity_markers.publish( + EntityMarkers(header=header_now("world"), markers=markers) + ) logger.info(f"[temporal-memory] published {len(markers)} entity markers to Rerun") except Exception as e: logger.debug(f"entity marker publish error: {e}") @@ -302,7 +305,9 @@ def _on_frame(img: Image) -> None: ) self.register_disposable( - frame_subject.pipe(sharpness_barrier(self.config.fps)).subscribe(_on_frame) + frame_subject.pipe(quality_barrier(image_sharpness, self.config.fps)).subscribe( + _on_frame + ) ) unsub_image = self.color_image.subscribe(frame_subject.on_next) self.register_disposable(Disposable(unsub_image)) @@ -310,9 +315,9 @@ def _on_frame(img: Image) -> None: # Odometry tracking for entity world positioning (optional — # module works without it, entities just won't have world positions) def _on_odom(msg: PoseStamped) -> None: - self._robot_x = msg.position.x - self._robot_y = msg.position.y - self._robot_z = msg.position.z + self._robot_x = msg.pose.position.x + self._robot_y = msg.pose.position.y + self._robot_z = msg.pose.position.z self._odom_count += 1 if self.odom.transport is not None: diff --git a/dimos/perception/experimental/temporal_memory/test_temporal_memory_module.py b/dimos/perception/experimental/temporal_memory/test_temporal_memory_module.py index 79265931d1d..fac688a65cc 100644 --- a/dimos/perception/experimental/temporal_memory/test_temporal_memory_module.py +++ b/dimos/perception/experimental/temporal_memory/test_temporal_memory_module.py @@ -21,11 +21,18 @@ from typing import TYPE_CHECKING from unittest.mock import MagicMock, create_autospec, patch +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.dimos_msgs.msg import EntityMarker, EntityMarkers +from dimos_generated.geometry_msgs.msg import Point +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from dimos.models.vl.base import VlModel -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_from_array +from dimos.msgs.time import time_from_seconds from dimos.perception.experimental.temporal_memory.entity_graph_db import EntityGraphDB from dimos.perception.experimental.temporal_memory.frame_window_accumulator import ( Frame, @@ -39,6 +46,7 @@ extract_time_window, ) from dimos.utils.logging_config import setup_logger +from dimos.visualization.rerun.message_helpers import entity_points if TYPE_CHECKING: from pathlib import Path @@ -52,7 +60,7 @@ def _make_image(value: int = 128, shape: tuple[int, ...] = (64, 64, 3)) -> Image: data = np.full(shape, value, dtype=np.uint8) - return Image.from_numpy(data) + return image_from_array(data, encoding="rgb8") class TestFrameWindowAccumulator: @@ -61,7 +69,7 @@ def test_bounded_buffer(self) -> None: acc.set_start_time(0.0) for i in range(10): img = _make_image(i * 25) - img.ts = float(i) + img.header.stamp = time_from_seconds(float(i)) acc.add_frame(img, float(i)) assert acc.buffer_size == 5 assert acc.frame_count == 10 @@ -74,7 +82,7 @@ def test_window_extraction(self) -> None: # Add 3 frames for i in range(3): img = _make_image() - img.ts = float(i) + img.header.stamp = time_from_seconds(float(i)) acc.add_frame(img, float(i)) frames = acc.try_extract_window() assert frames is not None @@ -85,7 +93,7 @@ def test_stride_guard(self) -> None: acc.set_start_time(0.0) for i in range(3): img = _make_image() - img.ts = float(i) + img.header.stamp = time_from_seconds(float(i)) acc.add_frame(img, float(i)) # First extraction should succeed frames = acc.try_extract_window() @@ -104,7 +112,7 @@ def test_clear(self) -> None: acc = FrameWindowAccumulator(max_buffer_frames=50, window_s=1.0, stride_s=1.0, fps=1.0) acc.set_start_time(0.0) img = _make_image() - img.ts = 0.0 + img.header.stamp = time_from_seconds(0.0) acc.add_frame(img, 0.0) assert acc.buffer_size == 1 acc.clear() @@ -418,20 +426,29 @@ def test_publish_entity_markers(self, tmp_path: Path) -> None: ids = {m.entity_id for m in msg.markers} assert ids == {"E1", "E2"} e1 = next(m for m in msg.markers if m.entity_id == "E1") - assert e1.x == 1.0 - assert e1.y == 2.0 + assert e1.position.x == 1.0 + assert e1.position.y == 2.0 tm.stop() def test_markers_to_rerun(self) -> None: - from dimos.msgs.visualization_msgs.EntityMarkers import EntityMarkers, Marker - markers = EntityMarkers( markers=[ - Marker("E1", "person walking", "person", 1.0, 2.0, 0.3), - Marker("E2", "wooden table", "object", 3.0, 4.0, 0.3), - ] + EntityMarker( + entity_id="E1", + label="person walking", + entity_type="person", + position=Point(x=1.0, y=2.0, z=0.3), + ), + EntityMarker( + entity_id="E2", + label="wooden table", + entity_type="object", + position=Point(x=3.0, y=4.0, z=0.3), + ), + ], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) - archetype = markers.to_rerun() + archetype = entity_points(cdr_decode(cdr_encode(markers), EntityMarkers)) # Should return rr.Points3D import rerun as rr @@ -455,7 +472,7 @@ def test_analyze_window_calls_vlm(self) -> None: analyzer = WindowAnalyzer(mock_vlm) img = _make_image() - img.ts = 0.0 + img.header.stamp = time_from_seconds(0.0) frame = Frame(frame_index=0, timestamp_s=0.0, image=img) state_dict = {"entity_roster": [], "rolling_summary": ""} @@ -473,7 +490,7 @@ def test_analyze_window_vlm_error(self) -> None: analyzer = WindowAnalyzer(mock_vlm) img = _make_image() - img.ts = 0.0 + img.header.stamp = time_from_seconds(0.0) frame = Frame(frame_index=0, timestamp_s=0.0, image=img) result = analyzer.analyze_window([frame], {}, 0.0, 2.0) @@ -504,3 +521,23 @@ def test_answer_query(self) -> None: result = analyzer.answer_query("What is the answer?", {}, img) assert result is not None assert result.answer == "The answer is 42" + + +def test_scene_staleness_compares_pixels_not_padding(): + from dimos.perception.experimental.temporal_memory.frame_window_accumulator import Frame + from dimos.perception.experimental.temporal_memory.temporal_utils.helpers import is_scene_stale + + first = Image( + width=2, + height=2, + encoding="rgb8", + step=8, + data=np.asarray([10] * 6 + [0, 0] + [10] * 6 + [0, 0], dtype=np.uint8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + is_bigendian=0, + ) + last = cdr_decode(cdr_encode(first), Image) + last.data = np.asarray([10] * 6 + [255, 255] + [10] * 6 + [255, 255], dtype=np.uint8) + assert is_scene_stale([Frame(0, 0, first), Frame(1, 1, last)], stale_threshold=1) + last.data = np.asarray([30] * 6 + [255, 255] + [30] * 6 + [255, 255], dtype=np.uint8) + assert not is_scene_stale([Frame(0, 0, first), Frame(1, 1, last)], stale_threshold=1) diff --git a/dimos/perception/experimental/temporal_memory/window_analyzer.py b/dimos/perception/experimental/temporal_memory/window_analyzer.py index afb21c3ca48..cf0746d9875 100644 --- a/dimos/perception/experimental/temporal_memory/window_analyzer.py +++ b/dimos/perception/experimental/temporal_memory/window_analyzer.py @@ -35,8 +35,9 @@ from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: + from dimos_generated.sensor_msgs.msg import Image + from dimos.models.vl.base import VlModel - from dimos.msgs.sensor_msgs.Image import Image from dimos.perception.experimental.temporal_memory.frame_window_accumulator import Frame logger = setup_logger() diff --git a/dimos/perception/experimental/test_generated_tracker.py b/dimos/perception/experimental/test_generated_tracker.py index 3d8895afe80..b8f486bc196 100644 --- a/dimos/perception/experimental/test_generated_tracker.py +++ b/dimos/perception/experimental/test_generated_tracker.py @@ -14,6 +14,7 @@ from unittest.mock import MagicMock +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.sensor_msgs.msg import Image from dimos_generated.std_msgs.msg import Header from dimos_generated.tf2_msgs.msg import TFMessage @@ -45,7 +46,9 @@ def test_tracker_publishes_generated_detection_and_clears_failed_track(succeeded module.tracked_overlay = MagicMock() try: module._process_tracking() - result = Detection2DArray.decode(module.detection2darray.publish.call_args.args[0].encode()) + result = cdr_decode( + cdr_encode(module.detection2darray.publish.call_args.args[0]), Detection2DArray + ) assert result.header.frame_id == "camera" if not succeeded: assert not result.detections @@ -56,7 +59,7 @@ def test_tracker_publishes_generated_detection_and_clears_failed_track(succeeded assert (detection.bbox.center.position.x, detection.bbox.center.position.y) == (20, 20) assert (detection.bbox.size_x, detection.bbox.size_y) == (20, 16) assert detection.results[0].hypothesis.class_id == "tracked_object" - overlay = Image.decode(module.tracked_overlay.publish.call_args.args[0].encode()) + overlay = cdr_decode(cdr_encode(module.tracked_overlay.publish.call_args.args[0]), Image) assert overlay.header == result.header assert image_view(overlay).shape == (64, 64, 3) assert overlay.encoding == "rgb8" @@ -80,8 +83,10 @@ def test_tracker3d_generated_projection_preserves_depth_filter_and_source_stamp( Detection2D( header=header, bbox=BoundingBox2D( - center=Pose2D(position=Point2D(x=32, y=32)), size_x=20, size_y=16 + center=Pose2D(position=Point2D(x=32, y=32), theta=0.0), size_x=20, size_y=16 ), + results=[], + id="", ) ], ) @@ -91,12 +96,14 @@ def test_tracker3d_generated_projection_preserves_depth_filter_and_source_stamp( assert output is None module.tf.publish.assert_not_called() return - result = Detection3DArray.decode(output.encode()) + result = cdr_decode(cdr_encode(output), Detection3DArray) assert result.header == header box = result.detections[0].bbox assert (box.center.position.x, box.center.position.y, box.center.position.z) == (2, 0, 0) assert np.allclose([box.size.x, box.size.y, box.size.z], [0.4, 0.32, 0.1]) - transform = TFMessage.decode(module.tf.publish.call_args.args[0].encode()).transforms[0] + transform = cdr_decode( + cdr_encode(module.tf.publish.call_args.args[0]), TFMessage + ).transforms[0] assert transform.header == header assert transform.child_frame_id == "tracked_object" assert transform.transform.translation.x == 2.0 diff --git a/dimos/perception/experimental/test_object_generated.py b/dimos/perception/experimental/test_object_generated.py index be288947c6d..3054a5e5ed0 100644 --- a/dimos/perception/experimental/test_object_generated.py +++ b/dimos/perception/experimental/test_object_generated.py @@ -14,8 +14,9 @@ """Object geometry and deduplication use generated ROS values.""" +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time -from dimos_generated.geometry_msgs.msg import PoseStamped, Vector3 +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Vector3 from dimos_generated.sensor_msgs.msg import Image, PointCloud2 from dimos_generated.std_msgs.msg import Header from dimos_generated.vision_msgs.msg import Detection3D @@ -31,26 +32,48 @@ def _object(object_id, x, timestamp): return Object( object_id=object_id, name="cup", - center=Vector3(x=x), + center=Vector3(x=x, y=0.0, z=0.0), size=Vector3(x=0.1, y=0.2, z=0.3), pose=PoseStamped( - header=Header(frame_id="world", stamp=Time(sec=1700000000, nanosec=123456789)) + header=Header(frame_id="world", stamp=Time(sec=1700000000, nanosec=123456789)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + pointcloud=PointCloud2( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + height=0, + width=0, + fields=[], + is_bigendian=False, + point_step=0, + row_step=0, + data=np.array([], dtype=np.uint8), + is_dense=False, ), - pointcloud=PointCloud2(), bbox=(0.0, 0.0, 1.0, 1.0), track_id=-1, class_id=0, confidence=0.9, ts=timestamp, frame_id="world", - image=Image(), + image=Image( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + height=0, + width=0, + encoding="", + is_bigendian=0, + step=0, + data=np.array([], dtype=np.uint8), + ), ) def test_detection_cdr_preserves_source_pose_header_and_current_center(): obj = _object("cup-1", 0.2, 1.0) obj.set_center(Vector3(x=0.4, y=0.5, z=0.6)) - decoded = Detection3D.decode(obj.to_detection3d_msg().encode()) + decoded = cdr_decode(cdr_encode(obj.to_detection3d_msg()), Detection3D) assert decoded.header == obj.pose.header assert decoded.bbox.center == obj.pose.pose assert decoded.bbox.size == obj.size @@ -70,8 +93,8 @@ def test_spatial_deduplication_and_nearest_lookup_use_generated_centers(): assert database.add_objects([nearby]) == [first] assert first.center.x == 0.1 assert database.add_objects([distant]) == [distant] - assert database.find_nearest(Vector3(x=1.8)) is distant - assert database.find_nearest(Vector3(x=0.2)) is first + assert database.find_nearest(Vector3(x=1.8, y=0.0, z=0.0)) is distant + assert database.find_nearest(Vector3(x=0.2, y=0.0, z=0.0)) is first def test_generated_object_cloud_aggregation_and_dictionary_preserve_values(): @@ -89,7 +112,7 @@ def test_generated_object_cloud_aggregation_and_dictionary_preserve_values(): np.testing.assert_allclose(exported["pointcloud"][1], [[1.0, 128 / 255.0, 0.0]]) np.testing.assert_array_equal(exported["image"], pixels) merged = aggregate_pointclouds([obj]) - decoded = PointCloud2.decode(merged.encode()) + decoded = cdr_decode(cdr_encode(merged), PointCloud2) assert decoded.header.stamp == stamp and decoded.header.frame_id == "world" np.testing.assert_array_equal(pointcloud_xyz(decoded), [[1.0, 2.0, 3.0]]) assert pointcloud_rgb(decoded) is not None diff --git a/dimos/perception/experimental/test_object_scene_registration_temporal.py b/dimos/perception/experimental/test_object_scene_registration_temporal.py index 4aa1bade0e8..dfae40e4440 100644 --- a/dimos/perception/experimental/test_object_scene_registration_temporal.py +++ b/dimos/perception/experimental/test_object_scene_registration_temporal.py @@ -20,6 +20,7 @@ from typing import Any from unittest.mock import ANY, MagicMock, call +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 from dimos_generated.sensor_msgs.msg import Image from dimos_generated.std_msgs.msg import Header @@ -62,7 +63,7 @@ def _image(timestamp: float) -> Image: def module() -> Iterator[ObjectSceneRegistrationModule]: module = ObjectSceneRegistrationModule(target_frame="map") module._camera_info = camera_info_from_intrinsics( - 1, 1, 0, 0, 2, 2, header=Header(frame_id="camera") + 1, 1, 0, 0, 2, 2, header=Header(frame_id="camera", stamp=Time(sec=0, nanosec=0)) ) module._latest_scene_snapshot = None yield module @@ -146,9 +147,11 @@ def test_full_scene_pointcloud_uses_one_coherent_scene_snapshot( ) -> None: depth = _image(3.0) transform = TransformStamped( - header=Header(frame_id="map"), + header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), child_frame_id="camera", - transform=Transform(translation=Vector3(x=1), rotation=Quaternion(w=1)), + transform=Transform( + translation=Vector3(x=1, y=0.0, z=0.0), rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), ) module._tf = _FakeTF(transform) # type: ignore[assignment] module._latest_scene_snapshot = (depth, transform) @@ -454,9 +457,9 @@ def test_object_db_counts_each_source_frame_once(monkeypatch: Any) -> None: last_seen_ts=None, detections_count=1, ) - first.center = Vector3() + first.center = Vector3(x=0.0, y=0.0, z=0.0) duplicate = MagicMock(object_id="duplicate-id", track_id=-1, ts=4.0) - duplicate.center = Vector3() + duplicate.center = Vector3(x=0.0, y=0.0, z=0.0) observed = object_db.add_objects([first, duplicate]) @@ -465,7 +468,7 @@ def test_object_db_counts_each_source_frame_once(monkeypatch: Any) -> None: assert first.last_seen_ts == 1000.0 newer = MagicMock(object_id="newer-id", track_id=-1, ts=5.0) - newer.center = Vector3() + newer.center = Vector3(x=0.0, y=0.0, z=0.0) first.update_object.side_effect = lambda _: setattr(first, "detections_count", 2) now[0] = 1001.0 diff --git a/dimos/perception/experimental/test_spatial_memory.py b/dimos/perception/experimental/test_spatial_memory.py index 211644847cf..fcff498046e 100644 --- a/dimos/perception/experimental/test_spatial_memory.py +++ b/dimos/perception/experimental/test_spatial_memory.py @@ -17,13 +17,12 @@ import tempfile import time +from dimos_generated.geometry_msgs.msg import Vector3 import numpy as np import pytest from reactivex import of, operators as ops from reactivex.scheduler import ThreadPoolScheduler -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.perception.experimental.spatial_perception import SpatialMemory from dimos.stream.video_provider import VideoProvider @@ -109,7 +108,7 @@ def process_frame(frame): nonlocal frame_counter # Generate a unique position for this frame to ensure minimum distance threshold is met - pos = Pose(frame_counter * 0.5, frame_counter * 0.5, 0) + pos = Vector3(x=frame_counter * 0.5, y=frame_counter * 0.5, z=0.0) transform = {"position": pos, "timestamp": time.time()} frame_counter += 1 @@ -221,8 +220,16 @@ def flaky_get_embedding(image): spatial_memory.process_stream( of( - {"frame": frame, "position": Vector3(1000, 0, 0), "rotation": Vector3()}, - {"frame": frame, "position": Vector3(1001, 0, 0), "rotation": Vector3()}, + { + "frame": frame, + "position": Vector3(x=1000.0, y=0.0, z=0.0), + "rotation": Vector3(x=0.0, y=0.0, z=0.0), + }, + { + "frame": frame, + "position": Vector3(x=1001.0, y=0.0, z=0.0), + "rotation": Vector3(x=0.0, y=0.0, z=0.0), + }, ) ).subscribe(on_next=results.append, on_error=errors.append) diff --git a/dimos/perception/fiducial/blueprints/desk_marker_tf.py b/dimos/perception/fiducial/blueprints/desk_marker_tf.py index 2367589457c..125a65e11ff 100644 --- a/dimos/perception/fiducial/blueprints/desk_marker_tf.py +++ b/dimos/perception/fiducial/blueprints/desk_marker_tf.py @@ -20,6 +20,7 @@ import threading import time +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 from dimos_generated.sensor_msgs.msg import CameraInfo from dimos_generated.std_msgs.msg import Header @@ -65,7 +66,7 @@ def create_desk_camera_info( camera_info_yaml: str | Path = DEFAULT_DESK_CAMERA_INFO_YAML, ) -> CameraInfo: camera_info = camera_info_from_yaml( - camera_info_yaml, header=Header(frame_id=DESK_CAMERA_FRAME_ID) + camera_info_yaml, header=Header(frame_id=DESK_CAMERA_FRAME_ID, stamp=Time(sec=0, nanosec=0)) ) return camera_info @@ -138,7 +139,10 @@ def publish_static_chain(self) -> None: TransformStamped( header=Header(stamp=stamp, frame_id=self.config.world_frame), child_frame_id=self.config.base_frame, - transform=Transform(rotation=Quaternion(w=1)), + transform=Transform( + rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + translation=Vector3(x=0.0, y=0.0, z=0.0), + ), ), TransformStamped( header=Header(stamp=stamp, frame_id=self.config.base_frame), diff --git a/dimos/perception/fiducial/marker_detect.py b/dimos/perception/fiducial/marker_detect.py index c002a5ba4fe..a6d1975a43f 100644 --- a/dimos/perception/fiducial/marker_detect.py +++ b/dimos/perception/fiducial/marker_detect.py @@ -81,7 +81,7 @@ def detect_markers_in_image( child_frame_id=optical_frame, transform=world_T_optical.transform, ) - marker_size = Vector3(x=marker_length_m, y=marker_length_m) + marker_size = Vector3(x=marker_length_m, y=marker_length_m, z=0.0) detections: list[Detection3DMarker] = [] for corner_set, mid_arr in zip(corners, ids, strict=True): diff --git a/dimos/perception/fiducial/marker_pose.py b/dimos/perception/fiducial/marker_pose.py index 6e7c9b838ed..19e8509a176 100644 --- a/dimos/perception/fiducial/marker_pose.py +++ b/dimos/perception/fiducial/marker_pose.py @@ -42,7 +42,7 @@ def camera_info_to_cv_matrices( ) -> tuple[np.ndarray[Any, np.dtype[Any]], np.ndarray[Any, np.dtype[Any]]]: """Build OpenCV ``cameraMatrix`` and ``distCoeffs`` from ``CameraInfo``.""" k = np.array(camera_info.k, dtype=np.float64).reshape(3, 3) - d = np.array(camera_info.d if camera_info.d else [], dtype=np.float64).reshape(-1, 1) + d = np.array(camera_info.d, dtype=np.float64).reshape(-1, 1) return k, d diff --git a/dimos/perception/fiducial/marker_tf_module.py b/dimos/perception/fiducial/marker_tf_module.py index 573245a21bc..6e3b56b6ad0 100644 --- a/dimos/perception/fiducial/marker_tf_module.py +++ b/dimos/perception/fiducial/marker_tf_module.py @@ -90,7 +90,10 @@ def _process_detections(self, detections: Detection3DArray) -> None: TransformStamped( header=Header(stamp=stamp, frame_id=self.config.world_frame), child_frame_id=markers_parent, - transform=Transform(rotation=Quaternion(w=1)), + transform=Transform( + rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + translation=Vector3(x=0.0, y=0.0, z=0.0), + ), ) ] diff --git a/dimos/perception/fiducial/test_fixture_verification.py b/dimos/perception/fiducial/test_fixture_verification.py index 35047537e7b..6a2114aeda7 100644 --- a/dimos/perception/fiducial/test_fixture_verification.py +++ b/dimos/perception/fiducial/test_fixture_verification.py @@ -15,7 +15,7 @@ from __future__ import annotations import cv2 -from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 from dimos_generated.std_msgs.msg import Header import numpy as np import pytest @@ -168,7 +168,9 @@ def test_marker_tf_replay_synthetic_packed_board_publishes_twelve_markers( world_T_optical = TransformStamped( header=Header(stamp=header.stamp, frame_id="world"), child_frame_id="camera_optical", - transform=Transform(rotation=Quaternion(w=1)), + transform=Transform( + rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), translation=Vector3(x=0.0, y=0.0, z=0.0) + ), ) detections = detect_markers_in_image( image, diff --git a/dimos/perception/fiducial/test_generated_marker_detect.py b/dimos/perception/fiducial/test_generated_marker_detect.py index 178185eed28..d614d485a91 100644 --- a/dimos/perception/fiducial/test_generated_marker_detect.py +++ b/dimos/perception/fiducial/test_generated_marker_detect.py @@ -13,9 +13,10 @@ # limitations under the License. import cv2 +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 -from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, RegionOfInterest from dimos_generated.std_msgs.msg import Header from dimos_generated.vision_msgs.msg import Detection3D import numpy as np @@ -31,12 +32,26 @@ def test_generated_image_detects_real_aruco_and_composes_pose(): pixels[100:300, 100:300] = cv2.aruco.generateImageMarker(dictionary, 7, 200) header = Header(frame_id="camera", stamp=Time(sec=1700000000, nanosec=123456789)) image = image_from_array(pixels, encoding="mono8", header=header) - image = Image.decode(image.encode()) - camera = CameraInfo(width=400, height=400, k=[400, 0, 200, 0, 400, 200, 0, 0, 1]) + image = cdr_decode(cdr_encode(image), Image) + camera = CameraInfo( + width=400, + height=400, + k=np.array([400, 0, 200, 0, 400, 200, 0, 0, 1], dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + distortion_model="", + d=np.array([], dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), + ) tf = TransformStamped( - header=Header(frame_id="world"), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), child_frame_id="camera", - transform=Transform(translation=Vector3(x=2), rotation=Quaternion(w=1)), + transform=Transform( + translation=Vector3(x=2, y=0.0, z=0.0), rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), ) found = detect_markers_in_image( image, @@ -46,7 +61,7 @@ def test_generated_image_detects_real_aruco_and_composes_pose(): aruco_dictionary="DICT_4X4_50", ) assert len(found) == 1 - wire = Detection3D.decode(found[0].to_detection3d_msg().encode()) + wire = cdr_decode(cdr_encode(found[0].to_detection3d_msg()), Detection3D) assert wire.id == "7" assert wire.results[0].hypothesis.class_id == "DICT_4X4_50:7" assert wire.header.stamp == header.stamp diff --git a/dimos/perception/fiducial/test_helpers.py b/dimos/perception/fiducial/test_helpers.py index bf2449aa042..41d7a5add89 100644 --- a/dimos/perception/fiducial/test_helpers.py +++ b/dimos/perception/fiducial/test_helpers.py @@ -17,8 +17,9 @@ from __future__ import annotations import cv2 +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 -from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, RegionOfInterest from dimos_generated.std_msgs.msg import Header import numpy as np @@ -30,8 +31,15 @@ def camera_info(ts: float = 10.0) -> CameraInfo: info = CameraInfo( width=640, height=480, - k=[600, 0, 320, 0, 600, 240, 0, 0, 1], - header=Header(frame_id="camera_optical"), + k=np.array([600, 0, 320, 0, 600, 240, 0, 0, 1], dtype=np.float64), + header=Header(frame_id="camera_optical", stamp=Time(sec=0, nanosec=0)), + distortion_model="", + d=np.array([], dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), ) info.header.stamp = time_from_seconds(ts) return info diff --git a/dimos/perception/fiducial/test_marker_pose.py b/dimos/perception/fiducial/test_marker_pose.py index 31f64f6121e..55cffa8e660 100644 --- a/dimos/perception/fiducial/test_marker_pose.py +++ b/dimos/perception/fiducial/test_marker_pose.py @@ -13,7 +13,8 @@ # limitations under the License. import cv2 -from dimos_generated.sensor_msgs.msg import CameraInfo +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.sensor_msgs.msg import CameraInfo, RegionOfInterest from dimos_generated.std_msgs.msg import Header import numpy as np @@ -27,12 +28,42 @@ def test_camera_optical_frame_id_resolution() -> None: - info_named = CameraInfo(header=Header(frame_id="cam_info_optical")) - info_empty = CameraInfo() + info_named = CameraInfo( + header=Header(frame_id="cam_info_optical", stamp=Time(sec=0, nanosec=0)), + height=0, + width=0, + distortion_model="", + d=np.array([], dtype=np.float64), + k=np.zeros(9, dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), + ) + info_empty = CameraInfo( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + height=0, + width=0, + distortion_model="", + d=np.array([], dtype=np.float64), + k=np.zeros(9, dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), + ) pixels = np.zeros((480, 640, 3), dtype=np.uint8) - img_custom = image_from_array(pixels, encoding="bgr8", header=Header(frame_id="custom_optical")) + img_custom = image_from_array( + pixels, + encoding="bgr8", + header=Header(frame_id="custom_optical", stamp=Time(sec=0, nanosec=0)), + ) img_whitespace = image_from_array( - pixels, encoding="bgr8", header=Header(frame_id=" custom_optical ") + pixels, + encoding="bgr8", + header=Header(frame_id=" custom_optical ", stamp=Time(sec=0, nanosec=0)), ) img_empty = image_from_array(pixels, encoding="bgr8") diff --git a/dimos/perception/fiducial/test_marker_tf_integration.py b/dimos/perception/fiducial/test_marker_tf_integration.py index 43fa7fbe453..afd25de7d9f 100644 --- a/dimos/perception/fiducial/test_marker_tf_integration.py +++ b/dimos/perception/fiducial/test_marker_tf_integration.py @@ -19,7 +19,8 @@ import time import uuid -from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion, Vector3 +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseWithCovariance, Quaternion, Vector3 from dimos_generated.std_msgs.msg import Header from dimos_generated.tf2_msgs.msg import TFMessage from dimos_generated.vision_msgs.msg import ( @@ -29,6 +30,7 @@ ObjectHypothesis, ObjectHypothesisWithPose, ) +import numpy as np import pytest from dimos.core.transport import LCMTransport @@ -39,7 +41,18 @@ def _marker_detection_array(ts: float) -> Detection3DArray: - det = Detection3D() + det = Detection3D( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + results=[], + bbox=BoundingBox3D( + center=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + size=Vector3(x=0.0, y=0.0, z=0.0), + ), + id="", + ) det.header = Header(stamp=time_from_seconds(ts), frame_id="world") det.id = "11" det.results = [ @@ -47,7 +60,14 @@ def _marker_detection_array(ts: float) -> Detection3DArray: hypothesis=ObjectHypothesis( class_id="DICT_APRILTAG_36h11:11", score=1.0, - ) + ), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), ) ] det.bbox = BoundingBox3D( diff --git a/dimos/perception/fiducial/test_marker_tf_module.py b/dimos/perception/fiducial/test_marker_tf_module.py index 2c4aeb3c261..5d7a7546936 100644 --- a/dimos/perception/fiducial/test_marker_tf_module.py +++ b/dimos/perception/fiducial/test_marker_tf_module.py @@ -14,7 +14,9 @@ from unittest.mock import patch -from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion, Vector3 +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseWithCovariance, Quaternion, Vector3 from dimos_generated.std_msgs.msg import Header from dimos_generated.tf2_msgs.msg import TFMessage from dimos_generated.vision_msgs.msg import ( @@ -24,6 +26,7 @@ ObjectHypothesis, ObjectHypothesisWithPose, ) +import numpy as np import pytest from dimos.msgs.geometry import quaternion_from_euler @@ -45,7 +48,18 @@ def _detection_array( if orientation is None: orientation = Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) - det = Detection3D() + det = Detection3D( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + results=[], + bbox=BoundingBox3D( + center=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + size=Vector3(x=0.0, y=0.0, z=0.0), + ), + id="", + ) det.header = Header(stamp=time_from_seconds(ts), frame_id="world") det.id = marker_id det.results = [ @@ -53,7 +67,14 @@ def _detection_array( hypothesis=ObjectHypothesis( class_id=class_id, score=1.0, - ) + ), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), ) ] det.bbox = BoundingBox3D( @@ -188,7 +209,9 @@ def test_marker_namespace_prefix_child_frames() -> None: def test_marker_tf_wire_preserves_exact_timestamp(): module = MarkerTfModule() outputs = [] - unsubscribe = module.tf.subscribe(lambda msg: outputs.append(TFMessage.decode(msg.encode()))) + unsubscribe = module.tf.subscribe( + lambda msg: outputs.append(cdr_decode(cdr_encode(msg), TFMessage)) + ) detections = _detection_array(ts=0) detections.header.stamp.sec = 1700000000 detections.header.stamp.nanosec = 123456789 diff --git a/dimos/perception/localize/association.py b/dimos/perception/localize/association.py index 2b08b75898b..3f6cf6a2328 100644 --- a/dimos/perception/localize/association.py +++ b/dimos/perception/localize/association.py @@ -20,9 +20,13 @@ from typing import TYPE_CHECKING from uuid import uuid4 +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header import numpy as np -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.geometry import inverse_transform +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.msgs.time import time_from_seconds, to_seconds if TYPE_CHECKING: from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC @@ -56,9 +60,9 @@ class Groups: def add(self, det: Detection3DPC, radius: float, voxel: float) -> int: """Assign one sighting to its object, folding its cloud in.""" - points = np.asarray(det.pointcloud.pointcloud.points) + points = pointcloud_xyz(det.pointcloud) centroid = points.mean(axis=0) - camera = (-det.transform).translation + camera = inverse_transform(det.transform).transform.translation hit = next( (i for i, c in enumerate(self.centers) if np.linalg.norm(c - centroid) <= radius), -1 @@ -73,7 +77,7 @@ def add(self, det: Detection3DPC, radius: float, voxel: float) -> int: self.trace.append([]) self.ids.append(uuid4().hex[:8]) else: - self._fuse(hit, det.pointcloud.ts, points, centroid, voxel) + self._fuse(hit, to_seconds(det.pointcloud.header.stamp), points, centroid, voxel) self.members[hit].extend( (det.ts, float(det.confidence), *centroid, camera.x, camera.y, camera.z) @@ -89,13 +93,13 @@ def _fuse( weighted mean over every viewpoint rather than a mean of means, and the running center follows the same counts. """ - held = np.asarray(self.clouds[i].pointcloud.points) + held = pointcloud_xyz(self.clouds[i]) na, nb = self.raw[i], float(len(points)) self.centers[i] = (na * self.centers[i] + nb * centroid) / (na + nb) self.raw[i] = na + nb merged = np.vstack([held, points]) - ts = max(self.clouds[i].ts, ts) - frame = self.clouds[i].frame_id + ts = max(to_seconds(self.clouds[i].header.stamp), ts) + frame = self.clouds[i].header.frame_id if voxel > 0: w = np.ones(len(merged)) @@ -111,7 +115,9 @@ def _fuse( np.add.at(psum, inverse, merged * w[:, None]) merged = psum / wsum[:, None] self.weights[i] = wsum - self.clouds[i] = PointCloud2.from_numpy(merged, frame_id=frame, timestamp=ts) + self.clouds[i] = pointcloud_from_xyz( + merged, header=Header(frame_id=frame, stamp=time_from_seconds(ts)) + ) def rows(self, i: int) -> np.ndarray: return np.asarray(self.members[i]).reshape(-1, M_WIDTH) diff --git a/dimos/perception/localize/dandetect.py b/dimos/perception/localize/dandetect.py index 35398d44dae..4a9e6ccd465 100644 --- a/dimos/perception/localize/dandetect.py +++ b/dimos/perception/localize/dandetect.py @@ -36,12 +36,12 @@ if TYPE_CHECKING: from collections.abc import Callable + from dimos_generated.sensor_msgs.msg import Image import numpy as np from reactivex.abc import DisposableBase from dimos.memory.stream import Stream from dimos.models.embedding.base import EmbeddingModel - from dimos.msgs.sensor_msgs.Image import Image from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.localize.association import Groups from dimos.perception.localize.localize import LocalizeTrace diff --git a/dimos/perception/localize/demo_blueprints/xarm_feed.py b/dimos/perception/localize/demo_blueprints/xarm_feed.py index ef05babb4e3..84161720c70 100644 --- a/dimos/perception/localize/demo_blueprints/xarm_feed.py +++ b/dimos/perception/localize/demo_blueprints/xarm_feed.py @@ -36,14 +36,14 @@ import threading from typing import Any +from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.tf2_msgs.msg import TFMessage + from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import Out from dimos.memory.replay import ReplayStream from dimos.memory.store.sqlite import SqliteStore -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.utils.data import get_data from dimos.utils.logging_config import setup_logger diff --git a/dimos/perception/localize/localize.py b/dimos/perception/localize/localize.py index bcc6d0b737d..a4ba1159f8d 100644 --- a/dimos/perception/localize/localize.py +++ b/dimos/perception/localize/localize.py @@ -32,6 +32,10 @@ from dimos.memory.embed import EmbedImages from dimos.memory.transform import QualityWindow, peaks +from dimos.msgs.geometry import inverse_transform, vector_array +from dimos.msgs.image import image_brightness, image_sharpness +from dimos.msgs.pointcloud import pointcloud_xyz +from dimos.msgs.time import to_seconds from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.localize.association import ( @@ -121,15 +125,15 @@ def _lift( detections: ImageDetections2D[Any], rig: Rig, policy: LocalizePolicy, plane: SupportPlane | None ) -> list[Detection3DPC]: """Lift a frame's detections through its camera's geometry and gate them.""" - transform = rig.world_to_optical(detections.ts, detections.image.frame_id) + transform = rig.world_to_optical(detections.ts, detections.image.header.frame_id) lifted = rig.lift(detections, transform) if transform is not None else None if transform is None or lifted is None: return [] - camera = (-transform).translation.to_numpy() + camera = vector_array(inverse_transform(transform).transform.translation) valid: list[Detection3DPC] = [] for det3d in lifted: - points = np.asarray(det3d.pointcloud.pointcloud.points) + points = pointcloud_xyz(det3d.pointcloud) if len(points) < policy.min_points: continue if float((points.max(axis=0) - points.min(axis=0)).max()) > policy.max_object_extent_m: @@ -151,13 +155,13 @@ def _lift( def index_frames(rig: Rig, feed: Stream[Any, Any]) -> Stream[Any, Any]: """One camera's frames worth embedding: lit, sharpest per slot, world-posed, camera-tagged.""" return ( - feed.filter(lambda obs: obs.data.brightness > 0.1) - .transform(QualityWindow(lambda img: img.sharpness, window=1.0 / EMBED_HZ)) + feed.filter(lambda obs: image_brightness(obs.data) > 0.1) + .transform(QualityWindow(image_sharpness, window=1.0 / EMBED_HZ)) .map( lambda obs: obs.derive( data=obs.data, - pose=rig.camera_pose(obs.ts, obs.data.frame_id), - tags={**obs.tags, "camera": obs.data.frame_id}, + pose=rig.camera_pose(obs.ts, obs.data.header.frame_id), + tags={**obs.tags, "camera": obs.data.header.frame_id}, ) ) .filter(lambda obs: obs.pose is not None) @@ -183,7 +187,7 @@ def embed_index(rig: Rig, embedder: EmbeddingModel, t0: float, t1: float) -> Str feed.transform(EmbedImages(embedder)) .map( lambda obs: obs.derive( - data=IndexFrame(obs.ts, obs.data.frame_id, float(obs.data.sharpness)) + data=IndexFrame(obs.ts, obs.data.header.frame_id, image_sharpness(obs.data)) ) ) .materialize() @@ -319,7 +323,7 @@ def _detect(upstream: Iterator[Any]) -> Iterator[Any]: class_id=j, confidence=float(score), name=queries[j], - ts=img.ts, + ts=to_seconds(img.header.stamp), image=img, ) if det.is_valid(): diff --git a/dimos/perception/localize/module.py b/dimos/perception/localize/module.py index f113387ee53..884a4f13e30 100644 --- a/dimos/perception/localize/module.py +++ b/dimos/perception/localize/module.py @@ -38,8 +38,18 @@ import time from typing import Any, Literal -from dimos_lcm.geometry_msgs import Pose -from dimos_lcm.vision_msgs import BoundingBox3D, ObjectHypothesis, ObjectHypothesisWithPose +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseWithCovariance, Quaternion, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage +from dimos_generated.vision_msgs.msg import ( + BoundingBox3D, + Detection3D, + Detection3DArray, + ObjectHypothesis, + ObjectHypothesisWithPose, +) +import numpy as np from pydantic import Field from dimos.agents.annotation import skill @@ -49,15 +59,8 @@ from dimos.memory.blobstore.memory import MemoryBlobStore from dimos.memory.observationstore.memory import ListObservationStore from dimos.memory.store.memory import MemoryStore -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.std_msgs.Header import Header -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.msgs.vision_msgs.Detection3D import Detection3D -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.msgs.pointcloud import pointcloud_xyz +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.perception.localize.association import Groups from dimos.perception.localize.dandetect import DanDetector from dimos.perception.localize.rig import DEPTH_TOLERANCE, EMBED_HZ, Rig, RigCamera @@ -70,7 +73,7 @@ INDEX_STREAM = "color_image_embedded" GEOMETRY_STREAM = "geometry_memory" COLOR_CODEC = "jpeg" -GEOMETRY_CODEC = "lz4+lcm" +GEOMETRY_CODEC = "lz4+cdr" LIDAR_TOLERANCE = 0.2 # s - an embedded frame to its lidar scan @@ -153,11 +156,13 @@ def on_camera_info(info: CameraInfo) -> None: self._camera_seen.set() def on_geometry(msg: Image | PointCloud2) -> None: - self._geometry_feed.append(msg, ts=msg.ts) + self._geometry_feed.append(msg, ts=to_seconds(msg.header.stamp)) geometry = self.depth_image if self.config.geometry == "depth" else self.lidar self._unsubs = [ - self.color_image.subscribe(lambda img: self._color_feed.append(img, ts=img.ts)), + self.color_image.subscribe( + lambda img: self._color_feed.append(img, ts=to_seconds(img.header.stamp)) + ), geometry.subscribe(on_geometry), self.camera_info.subscribe(on_camera_info), ] @@ -202,7 +207,7 @@ def _warm(self) -> None: assert camera is not None depth = self._geometry_memory if self.config.geometry == "depth" else None self.rig = Rig( - cameras={camera.frame_id: RigCamera(info=camera, color=self.index, depth=depth)}, + cameras={camera.header.frame_id: RigCamera(info=camera, color=self.index, depth=depth)}, tf=self._tf_buffer, world_frame=self.config.world_frame, lidar=None if depth is not None else self._geometry_memory, @@ -285,30 +290,43 @@ def as_detection_array( latest = 0.0 for query, hits in zip(queries, results, strict=True): for hit in hits: - points = hit.point_cloud.as_numpy()[0] + points = pointcloud_xyz(hit.point_cloud) low, high = points.min(axis=0), points.max(axis=0) middle, extent = (low + high) / 2, high - low - center = Vector3(*(float(v) for v in middle)) - size = Vector3(*(max(float(v), 1e-3) for v in extent)) + center = Point(x=float(middle[0]), y=float(middle[1]), z=float(middle[2])) + size = Vector3( + x=max(float(extent[0]), 1e-3), + y=max(float(extent[1]), 1e-3), + z=max(float(extent[2]), 1e-3), + ) latest = max(latest, hit.last_seen_timestamp) boxes.append( Detection3D( - header=Header(hit.last_seen_timestamp, frame_id), + header=Header( + stamp=time_from_seconds(hit.last_seen_timestamp), frame_id=frame_id + ), id=hit.instance_id, results=[ ObjectHypothesisWithPose( - hypothesis=ObjectHypothesis(class_id=query, score=hit.semantic_score) + hypothesis=ObjectHypothesis(class_id=query, score=hit.semantic_score), + pose=PoseWithCovariance( + pose=Pose( + position=center, + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), ) ], - results_length=1, bbox=BoundingBox3D( - center=Pose(position=center, orientation=Quaternion(0.0, 0.0, 0.0, 1.0)), + center=Pose( + position=center, orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), size=size, ), ) ) return Detection3DArray( - detections_length=len(boxes), - header=Header(latest, frame_id), + header=Header(stamp=time_from_seconds(latest), frame_id=frame_id), detections=boxes, ) diff --git a/dimos/perception/localize/rig.py b/dimos/perception/localize/rig.py index f26c6a4692c..865ec9b764f 100644 --- a/dimos/perception/localize/rig.py +++ b/dimos/perception/localize/rig.py @@ -20,9 +20,21 @@ import math from typing import TYPE_CHECKING, Any +from dimos_generated.geometry_msgs.msg import PoseStamped, TransformStamped as Transform +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header import numpy as np from dimos.memory.tf import StreamTF +from dimos.msgs.geometry import inverse_transform, pose_from_transform +from dimos.msgs.image import image_view +from dimos.msgs.pointcloud import ( + pointcloud_from_rgbd, + pointcloud_from_xyz, + pointcloud_xyz, + transform_cloud, +) +from dimos.msgs.time import time_from_seconds from dimos.perception.detection.type.detection3d.imageDetections3DPC import ImageDetections3DPC from dimos.perception.detection.type.detection3d.pointcloud_filters import ( range_cluster, @@ -33,13 +45,7 @@ from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: - from dimos_lcm.sensor_msgs import CameraInfo - from dimos.memory.type.observation import Observation - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - from dimos.msgs.geometry_msgs.Transform import Transform - from dimos.msgs.sensor_msgs.Image import Image - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.localize.support_plane import SupportPlane from dimos.protocol.tf.tf import TFLookup @@ -64,7 +70,7 @@ def _tf_root(store: Any, camera: str) -> str: t0, t1 = stream.get_time_range() middle = (t0 + t1) / 2 parent = { - t.child_frame_id: t.frame_id + t.child_frame_id: t.header.frame_id for obs in stream.after(middle - TF_ROOT_WINDOW).before(middle + TF_ROOT_WINDOW) for t in obs.data.transforms } @@ -112,37 +118,36 @@ def from_store(cls, store: Any) -> Rig: """Every calibrated optical frame that records colour, with aligned depth or beside the store's one point cloud stream, is a camera, posed in the root of its tf chain.""" - from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo as CameraInfoMsg - from dimos.msgs.sensor_msgs.Image import Image as ImageMsg - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 as PointCloudMsg - names = store.list_streams() types = {name: store.stream(name).data_type for name in names} infos: dict[str, CameraInfo] = {} - for name in sorted(n for n in names if types[n] is CameraInfoMsg): + for name in sorted(n for n in names if types[n] is CameraInfo): if store.stream(name).count(): info = store.stream(name).first().data - infos.setdefault(info.frame_id, info) # sorted: a colour info beats its depth twin + infos.setdefault( + info.header.frame_id, info + ) # sorted: a colour info beats its depth twin images: dict[str, dict[str, str]] = {"color": {}, "depth": {}} for name in names: - if types[name] is not ImageMsg or not store.stream(name).count(): + if types[name] is not Image or not store.stream(name).count(): continue image = store.stream(name).first().data - if image.data.dtype == np.uint16 or image.data.dtype.kind == "f": + pixels = image_view(image) + if pixels.dtype == np.uint16 or pixels.dtype.kind == "f": kind = "depth" - elif image.data.ndim == 3: + elif pixels.ndim == 3: kind = "color" else: continue # a derived copy, such as an embedding index, holds at most its feed's frames - held = images[kind].get(image.frame_id) + held = images[kind].get(image.header.frame_id) if held is None or store.stream(name).count() > store.stream(held).count(): - images[kind][image.frame_id] = name + images[kind][image.header.frame_id] = name color, depth = images["color"], images["depth"] - clouds = [n for n in names if types[n] is PointCloudMsg and store.stream(n).count()] + clouds = [n for n in names if types[n] is PointCloud2 and store.stream(n).count()] if len(clouds) > 1 and any(f in color and f not in depth for f in infos): raise ValueError(f"several point cloud streams for a camera without depth: {clouds}") lidar = clouds[0] if clouds else None @@ -184,7 +189,7 @@ def world_to_optical(self, ts: float, frame: str) -> Transform | None: def camera_pose(self, ts: float, frame: str) -> PoseStamped | None: """World pose of a camera's optical frame at ts.""" transform = self.world_to_optical(ts, frame) - return (-transform).to_pose() if transform is not None else None + return pose_from_transform(inverse_transform(transform)) if transform is not None else None def frame_at(self, obs: Observation[Any]) -> Image: """The colour frame an index observation refers to. @@ -193,7 +198,6 @@ def frame_at(self, obs: Observation[Any]) -> Image: frame selection needs, so its pixels come back from its camera's colour stream. """ - from dimos.msgs.sensor_msgs.Image import Image if isinstance(obs.data, Image): return obs.data @@ -217,20 +221,20 @@ def lidar_map(self) -> PointCloud2 | None: Bound ``lidar`` to the window first; whatever moved during it is in the map wherever it was. """ - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 - if self._lidar_map is not None and self._lidar_map[0] == id(self.lidar): return self._lidar_map[1] parts: list[np.ndarray] = [] newest = 0.0 for obs in self.lidar: cloud: PointCloud2 = obs.data - if cloud.frame_id != self.world_frame: - transform = self.tf.get(self.world_frame, cloud.frame_id, obs.ts, TF_TOLERANCE) + if cloud.header.frame_id != self.world_frame: + transform = self.tf.get( + self.world_frame, cloud.header.frame_id, obs.ts, TF_TOLERANCE + ) if transform is None: continue - cloud = cloud.transform(transform) - parts.append(cloud.points_f32()) + cloud = transform_cloud(cloud, transform) + parts.append(pointcloud_xyz(cloud)) newest = max(newest, obs.ts) merged = None if parts: @@ -238,8 +242,9 @@ def lidar_map(self) -> PointCloud2 | None: cells = np.floor(points / LIDAR_CELL).astype(np.int64) + (1 << 20) keys = (cells[:, 0] << 42) | (cells[:, 1] << 21) | cells[:, 2] _, first = np.unique(keys, return_index=True) - merged = PointCloud2.from_numpy( - points[first], frame_id=self.world_frame, timestamp=newest + merged = pointcloud_from_xyz( + points[first], + header=Header(frame_id=self.world_frame, stamp=time_from_seconds(newest)), ) self._lidar_map = (id(self.lidar), merged) return merged @@ -248,7 +253,7 @@ def lift( self, detections: ImageDetections2D[Any], transform: Transform ) -> ImageDetections3DPC | None: """A frame's detections as world clouds, through its camera's depth or the lidar.""" - frame = detections.image.frame_id + frame = detections.image.header.frame_id camera = self.cameras[frame] if camera.depth is not None: depth = self.depth_at(detections.ts, frame) @@ -263,8 +268,6 @@ def lift( def backdrop(self, ts: float, frame: str, depth_trunc: float = 1.5) -> PointCloud2 | None: """A camera's own RGBD frame at ts as a world-frame cloud, for plane fits and rendering; the lidar map for a camera without depth.""" - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 - if self.cameras[frame].depth is None: return self.lidar_map() depth = self.depth_at(ts, frame) @@ -277,6 +280,9 @@ def backdrop(self, ts: float, frame: str, depth_trunc: float = 1.5) -> PointClou transform = self.world_to_optical(ts, frame) if transform is None: return None - return PointCloud2.from_rgbd( - color, depth, self.cameras[frame].info, depth_scale=0.001, depth_trunc=depth_trunc - ).transform(-transform) + return transform_cloud( + pointcloud_from_rgbd( + color, depth, self.cameras[frame].info, depth_scale=0.001, depth_trunc=depth_trunc + ), + inverse_transform(transform), + ) diff --git a/dimos/perception/localize/support_plane.py b/dimos/perception/localize/support_plane.py index c751724dd62..a919e5628ec 100644 --- a/dimos/perception/localize/support_plane.py +++ b/dimos/perception/localize/support_plane.py @@ -25,8 +25,9 @@ from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: + from dimos_generated.sensor_msgs.msg import Image + from dimos.memory.type.observation import Observation - from dimos.msgs.sensor_msgs.Image import Image from dimos.perception.localize.rig import Rig logger = setup_logger() diff --git a/dimos/perception/localize/test_rig.py b/dimos/perception/localize/test_rig.py index d7524681ecd..034d8763ed0 100644 --- a/dimos/perception/localize/test_rig.py +++ b/dimos/perception/localize/test_rig.py @@ -12,52 +12,117 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Canonical pose frames and camera extrinsics used by localization.""" +"""Camera frames, depth joins, and lidar registration used by localization.""" +from dimos_cdrgen.registry import decode, encode from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage import numpy as np import pytest +from dimos.memory.store.memory import MemoryStore from dimos.memory.type.observation import Observation -from dimos.msgs.geometry import transform_matrix +from dimos.msgs.image import image_from_array +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_xyz from dimos.msgs.time import time_from_seconds, to_seconds from dimos.perception.localize.rig import Rig +from dimos.protocol.tf.tf import MultiTBuffer from dimos.robot.unitree.go2.camera_calibration import front_camera_calibration -@pytest.mark.parametrize("query, expected_x", [(1.0, 2.0), (1.5, 3.0)]) -def test_pose_interpolation_preserves_world_frame_for_camera_composition(mocker, query, expected_x): - poses = mocker.Mock() - poses.at.return_value = [ - Observation(ts=1.0, pose=(2.0, 0.0, 0.0)), - Observation(ts=2.0, pose=(4.0, 0.0, 0.0)), - ] - mount = TransformStamped( - header=Header(stamp=time_from_seconds(0), frame_id="base_link"), - child_frame_id="camera_optical", +def _transform(parent: str, child: str, x: float, ts: float = 2.0) -> TransformStamped: + return TransformStamped( + header=Header(stamp=time_from_seconds(ts), frame_id=parent), + child_frame_id=child, transform=Transform( - translation=Vector3(x=0.5, y=0.0, z=0.0), + translation=Vector3(x=x, y=0.0, z=0.0), rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), ), ) - rig = Rig( - cameras={"camera_optical": front_camera_calibration()}, - color=None, - world_frame="world", - poses=poses, - mounts={"camera_optical": mount}, + + +@pytest.mark.parametrize("frame, expected_x", [("front", 2.5), ("rear", 1.5)]) +def test_camera_pose_composes_each_cameras_mount(frame: str, expected_x: float) -> None: + tf = MultiTBuffer() + tf.receive_transform( + _transform("world", "base_link", 2.0), + _transform("base_link", "front", 0.5), + _transform("base_link", "rear", -0.5), ) + rig = Rig(cameras={}, tf=tf) - pose = rig.pose_at(query) - camera = rig.world_to_optical(query) + pose = rig.camera_pose(2.0, frame) + camera = rig.world_to_optical(2.0, frame) + assert pose is not None and camera is not None assert pose.header.frame_id == "world" assert pose.pose.position.x == expected_x - assert to_seconds(pose.header.stamp) == query - assert camera.header.frame_id == "camera_optical" + assert to_seconds(pose.header.stamp) == 2.0 + assert camera.header.frame_id == frame assert camera.child_frame_id == "world" + assert camera.transform.translation.x == -expected_x + assert rig.camera_pose(3.0, frame) is None + + +@pytest.fixture +def memory(): + with MemoryStore() as store: + yield store + + +def test_from_store_finds_each_cameras_color_and_aligned_depth(memory) -> None: + info = front_camera_calibration() + for frame in ("front", "rear"): + info = decode(encode(info), CameraInfo) + info.header.frame_id = frame + memory.stream(f"{frame}_info", CameraInfo).append(info, ts=2.0) + # Wire data is flat; classification must inspect decoded pixels. + color = image_from_array( + np.full((2, 2, 3), 100, dtype=np.uint8), + encoding="rgb8", + header=Header(stamp=time_from_seconds(2.0), frame_id=frame), + ) + depth = image_from_array( + np.full((2, 2), 1000, dtype=np.uint16), + encoding="16UC1", + header=Header(stamp=time_from_seconds(2.0), frame_id=frame), + ) + memory.stream(f"{frame}_color", Image).append(decode(encode(color), Image), ts=2.0) + memory.stream(f"{frame}_depth", Image).append(decode(encode(depth), Image), ts=2.0) + memory.stream("tf", TFMessage).append( + TFMessage(transforms=[_transform("world", "front", 2.5), _transform("world", "rear", 1.5)]), + ts=2.0, + ) + + rig = Rig.from_store(memory) + + assert set(rig.cameras) == {"front", "rear"} + assert rig.world_frame == "world" + assert rig.time_range() == (2.0, 2.0) + assert rig.depth_at(2.0, "rear").encoding == "16UC1" + assert rig.frame_at(Observation(ts=2.0, _data=color)).header.frame_id == "rear" + assert rig.camera_pose(2.0, "front").pose.position.x == 2.5 + + +def test_lidar_map_registers_and_deduplicates_scans_in_world_frame() -> None: + tf = MultiTBuffer() + tf.receive_transform(_transform("world", "lidar", 10.0)) + cloud = pointcloud_from_xyz( + np.array([[0.0, 0.0, 0.0], [0.01, 0.0, 0.0], [0.10, 0.0, 0.0]]), + header=Header(stamp=time_from_seconds(2.0), frame_id="lidar"), + ) + rig = Rig( + cameras={}, tf=tf, lidar=[Observation(ts=2.0, _data=decode(encode(cloud), PointCloud2))] + ) + + merged = rig.lidar_map() + + assert merged is not None + assert merged.header.frame_id == "world" + assert to_seconds(merged.header.stamp) == 2.0 np.testing.assert_allclose( - transform_matrix(camera.transform)[:3, 3], [-expected_x - 0.5, 0.0, 0.0] + pointcloud_xyz(merged), [[10.0, 0.0, 0.0], [10.1, 0.0, 0.0]], atol=1e-6 ) - assert to_seconds(camera.header.stamp) == query + assert rig.lidar_map() is merged diff --git a/dimos/perception/localize/tools/tool_localize.py b/dimos/perception/localize/tools/tool_localize.py index 5c4493f0708..fe73b15b010 100644 --- a/dimos/perception/localize/tools/tool_localize.py +++ b/dimos/perception/localize/tools/tool_localize.py @@ -41,12 +41,21 @@ import sys from typing import Any +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import Image import numpy as np from dimos.memory.store.sqlite import SqliteStore from dimos.memory.transform import throttle -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.geometry import ( + inverse_transform, + pose_from_transform, + pose_matrix, + transform_matrix, + vector_array, +) +from dimos.msgs.image import image_from_array, image_to_bgr, image_to_rgb +from dimos.msgs.pointcloud import pointcloud_xyz, voxel_downsample_cloud from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC, lattice_quantum from dimos.perception.localize.dandetect import DanDetector from dimos.perception.localize.localize import LocalizeTrace, embed_index @@ -104,12 +113,13 @@ def render( with it. The 3D view hides those frames, and ``cameras`` unless there is a robot, until toggled on. """ + from matplotlib import colormaps import rerun as rr import rerun.blueprint as rrb from dimos.memory.vis.color import Color - from dimos.msgs.sensor_msgs.PointCloud2 import _get_colormap_lut from dimos.visualization.rerun.init import rerun_init + from dimos.visualization.rerun.message_helpers import camera_pinhole, image_archetype names = {frame: rr.escape_entity_path_part(frame) for frame in rig.cameras} robot = base_pose(rig, backdrop_ts) is not None @@ -136,6 +146,10 @@ def render( ) ) + def pose_archetype(pose: PoseStamped) -> Any: + matrix = pose_matrix(pose.pose) + return rr.Transform3D(translation=matrix[:3, 3], mat3x3=matrix[:3, :3]) + def at(ts: float) -> None: rr.set_time("ts", timestamp=ts) @@ -146,7 +160,7 @@ def cell_keys(points: np.ndarray) -> np.ndarray: # map cells claimed by a detected object, plus one cell of halo answer_points = np.vstack( [ - det.pointcloud.points_f32() + pointcloud_xyz(det.pointcloud) for _, trace in traces for members in trace.answers for det in members @@ -161,13 +175,13 @@ def carve(points: np.ndarray) -> np.ndarray: def image_colors(det: Any) -> np.ndarray: """Sample the sighting's own image at each cloud point's reprojection.""" - points = det.pointcloud.points_f32() - matrix = det.transform.to_matrix() + points = pointcloud_xyz(det.pointcloud) + matrix = transform_matrix(det.transform.transform) cam = points @ matrix[:3, :3].T + matrix[:3, 3] - pixels = Detection3DPC.project_pixels(cam, rig.cameras[det.image.frame_id].info) + pixels = Detection3DPC.project_pixels(cam, rig.cameras[det.image.header.frame_id].info) cols = np.round(pixels[:, 0]).astype(int) rows = np.round(pixels[:, 1]).astype(int) - rgb = det.image.to_rgb().data + rgb = image_to_rgb(det.image) height, width = rgb.shape[:2] return rgb[np.clip(rows, 0, height - 1), np.clip(cols, 0, width - 1)] @@ -178,13 +192,13 @@ def image_colors(det: Any) -> np.ndarray: for frame, name in names.items(): backdrop = rig.backdrop(backdrop_ts, frame) if frame not in through_lidar else None if backdrop is not None: - carved = carve(backdrop.voxel_downsample(MAP_CELL).points_f32()) + carved = carve(pointcloud_xyz(voxel_downsample_cloud(backdrop, MAP_CELL))) if len(carved): views[name] = (carved, POINT_SIZE) lidar_size = LIDAR_POINT_SIZE lidar_map = rig.lidar_map() if through_lidar else None if lidar_map is not None: - points = lidar_map.points_f32() + points = pointcloud_xyz(lidar_map) lidar_size = lattice_quantum(points) or LIDAR_POINT_SIZE carved = carve(points) if len(carved): @@ -193,7 +207,7 @@ def image_colors(det: Any) -> np.ndarray: if views: heights = np.concatenate([carved[:, 2] for carved, _ in views.values()]) low, high = float(heights.min()), float(heights.max()) - turbo = _get_colormap_lut("turbo") + turbo = (colormaps["turbo"](np.linspace(0, 1, 256))[:, :3] * 255).astype(np.uint8) for name, (carved, spacing) in views.items(): t = np.clip((carved[:, 2] - low) / (high - low + 1e-8), 0.0, 1.0) colors = turbo[(t * 255).astype(np.uint8)].astype(np.float32) @@ -208,7 +222,11 @@ def image_colors(det: Any) -> np.ndarray: # each camera's live feed + frustum tracking its capture pose along the # timeline; a robot's body box follows its base frame for frame, name in names.items(): - rr.log(f"cameras/{name}", rig.cameras[frame].info.to_rerun(), static=True) + rr.log( + f"cameras/{name}", + camera_pinhole(rig.cameras[frame].info, optical_frame=""), + static=True, + ) if robot: body = rr.Boxes3D(half_sizes=[ROBOT_HALF_SIZES], colors=[ROBOT_COLOR]) rr.log("robot", body, static=True) @@ -220,16 +238,16 @@ def image_colors(det: Any) -> np.ndarray: key=lambda obs: obs.ts, ) for obs in feed: - frame = obs.data.frame_id + frame = obs.data.header.frame_id pose = rig.camera_pose(obs.ts, frame) if pose is None: continue at(obs.ts) - rr.log(f"cameras/{names[frame]}/image", obs.data.to_rerun()) - rr.log(f"cameras/{names[frame]}", pose.to_rerun()) + rr.log(f"cameras/{names[frame]}/image", image_archetype(obs.data)) + rr.log(f"cameras/{names[frame]}", pose_archetype(pose)) base = base_pose(rig, obs.ts) if robot else None if base is not None: - rr.log("robot", base.to_rerun()) + rr.log("robot", pose_archetype(base)) n_instances = sum(len(trace.answers) for _, trace in traces) box_colors = [ @@ -247,29 +265,31 @@ def image_colors(det: Any) -> np.ndarray: for _, trace in traces: for obs in trace.detection_frames: image = obs.data.image - shots.setdefault((image.frame_id, obs.ts), (image, []))[1].extend(obs.data.detections) + shots.setdefault((image.header.frame_id, obs.ts), (image, []))[1].extend( + obs.data.detections + ) for (camera, ts), (image, dets) in sorted(shots.items(), key=lambda shot: shot[0][1]): if rig.camera_pose(ts, camera) is None: continue at(ts) marked = overlay(image, best_per_box(dets), label_colors) - rr.log(f"cameras/{names[camera]}/image", marked.to_rerun()) + rr.log(f"cameras/{names[camera]}/image", image_archetype(marked)) for query, trace in traces: root = f"detections/{query.replace(' ', '_')}" # a frozen frustum at each detection frame's capture pose, with this label's masks for i, obs in enumerate(trace.detection_frames): - camera = obs.data.image.frame_id + camera = obs.data.image.header.frame_id pose = rig.camera_pose(obs.ts, camera) if pose is None: continue at(obs.ts) frame = f"{root}/frames/{i}" - rr.log(frame, pose.to_rerun()) - rr.log(frame, rig.cameras[camera].info.to_rerun()) + rr.log(frame, pose_archetype(pose)) + rr.log(frame, camera_pinhole(rig.cameras[camera].info, optical_frame="")) marked = overlay(obs.data.image, obs.data.detections, label_colors) - rr.log(f"{frame}/image", marked.to_rerun()) + rr.log(f"{frame}/image", image_archetype(marked)) # the answers: per verified instance, every sighting's cloud colored # by its own image and grouped by camera, wrapped in a labeled @@ -279,13 +299,13 @@ def image_colors(det: Any) -> np.ndarray: rr.log( f"{root}/answer/{i}/{names[camera]}", rr.Points3D( - positions=np.vstack([det.pointcloud.points_f32() for det in sightings]), + positions=np.vstack([pointcloud_xyz(det.pointcloud) for det in sightings]), colors=np.vstack([image_colors(det) for det in sightings]), radii=size[camera] / 2, ), static=True, ) - positions = np.vstack([det.pointcloud.points_f32() for det in members]) + positions = np.vstack([pointcloud_xyz(det.pointcloud) for det in members]) aabb_min, aabb_max = positions.min(axis=0), positions.max(axis=0) score = max(det.confidence for det in members) rr.log( @@ -307,7 +327,7 @@ def overlay(image: Image, detections: list[Any], colors: dict[str, tuple[int, in its box and its label; a detection without a mask gets the box alone.""" import cv2 - canvas = image.to_opencv().copy() + canvas = image_to_bgr(image).copy() tint = canvas.copy() for det in detections: mask = getattr(det, "mask", None) @@ -331,7 +351,7 @@ def overlay(image: Image, detections: list[Any], colors: dict[str, tuple[int, in cv2.putText( canvas, label, (x1 + 2, y1 - 4), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1 ) - return Image.from_opencv(canvas, frame_id=image.frame_id, ts=image.ts) + return image_from_array(canvas, encoding="bgr8", header=image.header) def best_per_box(detections: list[Any]) -> list[Any]: @@ -348,14 +368,14 @@ def best_per_box(detections: list[Any]) -> list[Any]: def base_pose(rig: Rig, ts: float) -> PoseStamped | None: """World pose of the robot's base frame at ts; None when tf has none.""" transform = rig.tf.get(BASE_FRAME, rig.world_frame, ts, TF_TOLERANCE) - return (-transform).to_pose() if transform is not None else None + return pose_from_transform(inverse_transform(transform)) if transform is not None else None def by_camera(rig: Rig, sightings: list[Detection3DPC]) -> dict[str, list[Detection3DPC]]: """An instance's sightings per camera, in the rig's camera order.""" grouped: dict[str, list[Detection3DPC]] = {frame: [] for frame in rig.cameras} for det in sightings: - grouped[det.image.frame_id].append(det) + grouped[det.image.header.frame_id].append(det) return {frame: dets for frame, dets in grouped.items() if dets} @@ -365,8 +385,10 @@ def camera_line(rig: Rig, sightings: list[Detection3DPC]) -> str: means: list[np.ndarray] = [] parts: list[str] = [] for frame, dets in by_camera(rig, sightings).items(): - centroids = [np.asarray(det.pointcloud.pointcloud.points).mean(axis=0) for det in dets] - cameras = np.array([(-det.transform).translation.to_numpy() for det in dets]) + centroids = [pointcloud_xyz(det.pointcloud).mean(axis=0) for det in dets] + cameras = np.array( + [vector_array(inverse_transform(det.transform).transform.translation) for det in dets] + ) views = len(np.unique(np.round(cameras, 2), axis=0)) mean = np.mean(centroids, axis=0) means.append(mean) @@ -388,7 +410,7 @@ def report(query: str, hits: list[Localization], trace: LocalizeTrace, rig: Rig, print( f"hit {query!r} [{hit.instance_id}]: position=({x:.3f}, {y:.3f}, {z:.3f}) " f"frame={hit.frame_id} ts_offset={hit.last_seen_timestamp - lo:.1f}s " - f"points={len(hit.point_cloud)} views={hit.n_views} " + f"points={hit.point_cloud.width * hit.point_cloud.height} views={hit.n_views} " f"score={hit.semantic_score:.2f} margin={hit.ambiguity_margin:.2f}" ) if len(rig.cameras) > 1: diff --git a/dimos/porcelain/test_remote_module_source.py b/dimos/porcelain/test_remote_module_source.py index 52a3f8a1da3..f97bf40a7ed 100644 --- a/dimos/porcelain/test_remote_module_source.py +++ b/dimos/porcelain/test_remote_module_source.py @@ -22,6 +22,7 @@ from typing import Protocol import uuid +from dimos_generated.geometry_msgs.msg import Vector3 import pytest from reactivex.disposable import Disposable @@ -32,7 +33,6 @@ from dimos.core.module import Module from dimos.core.stream import IO, In, Out from dimos.core.transport import LCMTransport, pLCMTransport -from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.porcelain.dimos import Dimos from dimos.porcelain.module_handle import RemoteModuleProxy from dimos.porcelain.remote_module_source import RemoteModuleSource diff --git a/dimos/protocol/pubsub/benchmark/test_replay_bench.py b/dimos/protocol/pubsub/benchmark/test_replay_bench.py index 6d663567950..7823d8107d6 100644 --- a/dimos/protocol/pubsub/benchmark/test_replay_bench.py +++ b/dimos/protocol/pubsub/benchmark/test_replay_bench.py @@ -20,6 +20,7 @@ import time from types import SimpleNamespace +from dimos_cdrgen.registry import encode as cdr_encode from dimos_generated.sensor_msgs.msg import CompressedImage, Image from dimos_generated.std_msgs.msg import Header import numpy as np @@ -196,7 +197,13 @@ def inner(topic: str, typ: type): raw_wire = len(frame.encode()) jpeg_wire = len( - CompressedImage(header=frame.header, format="jpeg", data=image_to_jpeg(frame)).encode() + cdr_encode( + CompressedImage( + header=frame.header, + format="jpeg", + data=np.asarray(image_to_jpeg(frame), dtype=np.uint8), + ) + ) ) raw = _bench(inner(f"dimos/bench/{proto}_raw", Image), frame, n, raw_wire) diff --git a/dimos/protocol/pubsub/benchmark/testdata.py b/dimos/protocol/pubsub/benchmark/testdata.py index 0dfc28d8f3f..e8c32b7108e 100644 --- a/dimos/protocol/pubsub/benchmark/testdata.py +++ b/dimos/protocol/pubsub/benchmark/testdata.py @@ -18,10 +18,11 @@ import os from typing import TYPE_CHECKING, Any +from dimos_generated.sensor_msgs.msg import Image import numpy as np import pytest -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_from_array from dimos.protocol.pubsub.benchmark.type import Case from dimos.protocol.pubsub.impl.zenohpubsub import Topic as ZenohTopic, Zenoh from dimos.protocol.pubsub.spec import PubSub @@ -36,6 +37,11 @@ DDS_AVAILABLE = False from dimos.protocol.pubsub.impl.lcmpubsub import LCM, LCMPubSubBase, Topic as LCMTopic from dimos.protocol.pubsub.impl.memory import Memory +from dimos.protocol.pubsub.impl.shmpubsub import ( + BytesSharedMemory, + CDRSharedMemory, + PickleSharedMemory, +) def make_data_bytes(size: int) -> bytes: @@ -54,7 +60,7 @@ def make_data_image(size: int) -> Image: height = max(1, int(pixels**0.5)) width = pixels // height data = padded_data[: height * width * 3].reshape(height, width, 3) - return Image(data=data, format=ImageFormat.RGB) + return image_from_array(data, encoding="rgb8") testcases: list[Case[Any, Any]] = [] @@ -69,7 +75,7 @@ def lcm_pubsub_channel() -> Generator[LCM, None, None]: def lcm_msggen(size: int) -> tuple[LCMTopic, Image]: - topic = LCMTopic(topic="benchmark/lcm", lcm_type=Image) + topic = LCMTopic(topic="benchmark/lcm", msg_type=Image) return (topic, make_data_image(size)) @@ -122,6 +128,66 @@ def memory_msggen(size: int) -> tuple[str, Any]: # ) +@contextmanager +def shm_pickle_pubsub_channel() -> Generator[PickleSharedMemory, None, None]: + # 12MB capacity to handle benchmark sizes up to 10MB + shm_pubsub = PickleSharedMemory(prefer="cpu", default_capacity=12 * 1024 * 1024) + shm_pubsub.start() + yield shm_pubsub + shm_pubsub.stop() + + +def shm_msggen(size: int) -> tuple[str, Any]: + """Generate message for SharedMemory pubsub benchmark.""" + return ("benchmark/shm", make_data_image(size)) + + +testcases.append( + Case( + pubsub_context=shm_pickle_pubsub_channel, + msg_gen=shm_msggen, + ) +) + + +@contextmanager +def shm_bytes_pubsub_channel() -> Generator[BytesSharedMemory, None, None]: + """SharedMemory with raw bytes - no pickle overhead.""" + shm_pubsub = BytesSharedMemory(prefer="cpu", default_capacity=12 * 1024 * 1024) + shm_pubsub.start() + yield shm_pubsub + shm_pubsub.stop() + + +def shm_bytes_msggen(size: int) -> tuple[str, bytes]: + """Generate raw bytes for SharedMemory transport benchmark.""" + return ("benchmark/shm_bytes", make_data_bytes(size)) + + +testcases.append( + Case( + pubsub_context=shm_bytes_pubsub_channel, + msg_gen=shm_bytes_msggen, + ) +) + + +@contextmanager +def shm_cdr_pubsub_channel() -> Generator[CDRSharedMemory, None, None]: + """SharedMemory with CDR encoding - no pickle overhead.""" + shm_pubsub = CDRSharedMemory(prefer="cpu", default_capacity=12 * 1024 * 1024) + shm_pubsub.start() + yield shm_pubsub + shm_pubsub.stop() + + +testcases.append( + Case( + pubsub_context=shm_cdr_pubsub_channel, + msg_gen=lcm_msggen, # Reuse the LCM message generator + ) +) + if DDS_AVAILABLE: from cyclonedds.idl import IdlStruct from cyclonedds.idl.types import sequence, uint8 diff --git a/dimos/protocol/pubsub/benchmark/tool_replay_bench.py b/dimos/protocol/pubsub/benchmark/tool_replay_bench.py index da4796d0eee..34b36809b26 100644 --- a/dimos/protocol/pubsub/benchmark/tool_replay_bench.py +++ b/dimos/protocol/pubsub/benchmark/tool_replay_bench.py @@ -73,7 +73,9 @@ def broadcast(self, stream: Stream[Image] | None, msg: Image | CompressedImage) msg if isinstance(msg, CompressedImage) else CompressedImage( - header=msg.header, format="jpeg", data=image_to_jpeg(msg, quality=self.quality) + header=msg.header, + format="jpeg", + data=np.asarray(image_to_jpeg(msg, quality=self.quality), dtype=np.uint8), ) ) self.inner.broadcast(None, compressed) diff --git a/dimos/protocol/pubsub/encoders.py b/dimos/protocol/pubsub/encoders.py index 3ef7845d39e..1601962a6ac 100644 --- a/dimos/protocol/pubsub/encoders.py +++ b/dimos/protocol/pubsub/encoders.py @@ -20,6 +20,8 @@ import pickle from typing import TYPE_CHECKING, Generic, Protocol, TypeVar, cast +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode + from dimos.msgs.protocol import DimosMsg if TYPE_CHECKING: @@ -107,16 +109,16 @@ class CDREncoderMixin(PubSubEncoderMixin[TypedTopicProto, DimosMsg, bytes]): def encode(self, msg: DimosMsg | bytes, topic: TypedTopicProto) -> bytes: if isinstance(msg, bytes): return msg - if topic.msg_type is not None and msg.msg_name != topic.msg_type.msg_name: + if topic.msg_type is not None and msg.__msgtype__ != topic.msg_type.__msgtype__: raise ValueError( - f"Message {msg.msg_name} does not match topic type {topic.msg_type.msg_name}" + f"Message {msg.__msgtype__} does not match topic type {topic.msg_type.__msgtype__}" ) - return msg.encode() + return cdr_encode(msg) def decode(self, msg: bytes, topic: TypedTopicProto) -> DimosMsg: if topic.msg_type is None: raise DecodingError(f"Cannot decode: topic {topic.topic!r} has no msg_type") try: - return topic.msg_type.decode(msg) + return cast("DimosMsg", cdr_decode(msg, topic.msg_type.__msgtype__)) except ValueError as exc: raise DecodingError(f"Invalid CDR on topic {topic.topic!r}") from exc diff --git a/dimos/protocol/pubsub/impl/jpeg_lcm.py b/dimos/protocol/pubsub/impl/jpeg_lcm.py index bafb414a22f..7dc46e22092 100644 --- a/dimos/protocol/pubsub/impl/jpeg_lcm.py +++ b/dimos/protocol/pubsub/impl/jpeg_lcm.py @@ -21,6 +21,7 @@ from __future__ import annotations +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.sensor_msgs.msg import CompressedImage, Image from dimos.msgs.image import compressed_image_from_image, image_from_compressed @@ -34,7 +35,7 @@ class JpegEncoderMixin(PubSubEncoderMixin[TypedTopicProto, Image, bytes]): def encode(self, msg: Image, topic: TypedTopicProto) -> bytes: if topic.msg_type is not CompressedImage: raise ValueError("JPEG wire topics must declare generated CompressedImage") - return compressed_image_from_image(msg).encode() + return cdr_encode(compressed_image_from_image(msg)) def decode(self, msg: bytes, topic: TypedTopicProto) -> Image: if topic.topic == "LCM_SELF_TEST": @@ -43,7 +44,7 @@ def decode(self, msg: bytes, topic: TypedTopicProto) -> Image: raise DecodingError(f"Cannot decode: topic {topic.topic!r} has no msg_type") if topic.msg_type is not CompressedImage: raise DecodingError("JPEG wire topics must declare generated CompressedImage") - return image_from_compressed(CompressedImage.decode(msg)) + return image_from_compressed(cdr_decode(msg, CompressedImage)) class JpegLCM( # type: ignore[misc] diff --git a/dimos/protocol/pubsub/impl/jpeg_shm.py b/dimos/protocol/pubsub/impl/jpeg_shm.py new file mode 100644 index 00000000000..ee4a1f7be8f --- /dev/null +++ b/dimos/protocol/pubsub/impl/jpeg_shm.py @@ -0,0 +1,41 @@ +# Copyright 2025-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. + +from typing import Any + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.sensor_msgs.msg import CompressedImage, Image + +from dimos.msgs.image import compressed_image_from_image, image_from_compressed +from dimos.protocol.pubsub.encoders import PubSubEncoderMixin +from dimos.protocol.pubsub.impl.shmpubsub import SharedMemoryPubSubBase + + +class JpegSharedMemoryEncoderMixin(PubSubEncoderMixin[str, Image, bytes]): + def __init__(self, quality: int = 75, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.quality = quality + + def encode(self, msg: Any, _topic: str) -> bytes: + if not isinstance(msg, Image): + raise ValueError("Can only encode images.") + + return cdr_encode(compressed_image_from_image(msg, quality=self.quality)) + + def decode(self, msg: bytes, _topic: str) -> Image: + return image_from_compressed(cdr_decode(msg, CompressedImage)) + + +class JpegSharedMemory(JpegSharedMemoryEncoderMixin, SharedMemoryPubSubBase): # type: ignore[misc] + pass diff --git a/dimos/protocol/pubsub/impl/lcmpubsub.py b/dimos/protocol/pubsub/impl/lcmpubsub.py index 34f28d4ca03..2c4980ea3c3 100644 --- a/dimos/protocol/pubsub/impl/lcmpubsub.py +++ b/dimos/protocol/pubsub/impl/lcmpubsub.py @@ -54,7 +54,7 @@ def pattern(self) -> str: def __str__(self) -> str: if self.msg_type is None: return self.pattern - return f"{self.pattern}#{self.msg_type.msg_name}" + return f"{self.pattern}#{self.msg_type.__msgtype__}" @staticmethod def from_channel_str(channel: str, default_msg_type: type[DimosMsg] | None = None) -> Topic: diff --git a/dimos/protocol/pubsub/impl/rospubsub_conversion.py b/dimos/protocol/pubsub/impl/rospubsub_conversion.py index 16438801eda..b69498e579e 100644 --- a/dimos/protocol/pubsub/impl/rospubsub_conversion.py +++ b/dimos/protocol/pubsub/impl/rospubsub_conversion.py @@ -19,6 +19,8 @@ import importlib from typing import TYPE_CHECKING, Any, TypeVar, cast +from dimos_cdrgen.registry import decode, encode + try: import rclpy.serialization as ros_serialization except ImportError: @@ -34,9 +36,11 @@ def derive_ros_type(dimos_type: type[DimosMsg]) -> type[ROSMessage]: """Resolve the matching installed ROS type by canonical package/msg/Type name.""" - parts = dimos_type.msg_name.split("/") + parts = dimos_type.__msgtype__.split("/") if len(parts) != 3 or parts[1] != "msg" or not all(part.isidentifier() for part in parts): - raise ValueError(f"Invalid message name {dimos_type.msg_name!r}; expected package/msg/Type") + raise ValueError( + f"Invalid message name {dimos_type.__msgtype__!r}; expected package/msg/Type" + ) package, _, name = parts return cast("type[ROSMessage]", getattr(importlib.import_module(f"{package}.msg"), name)) @@ -46,8 +50,8 @@ def dimos_to_ros(msg: DimosMsg, ros_type: type[ROSMessage]) -> ROSMessage: if ros_serialization is None: raise ImportError("ROS message conversion requires rclpy; install and source ROS 2.") if derive_ros_type(type(msg)) is not ros_type: - raise TypeError(f"ROS target does not match {msg.msg_name}") - return cast("ROSMessage", ros_serialization.deserialize_message(msg.encode(), ros_type)) + raise TypeError(f"ROS target does not match {msg.__msgtype__}") + return cast("ROSMessage", ros_serialization.deserialize_message(encode(msg), ros_type)) def ros_to_dimos(msg: Any, dimos_type: type[MessageT]) -> MessageT: @@ -55,5 +59,7 @@ def ros_to_dimos(msg: Any, dimos_type: type[MessageT]) -> MessageT: if ros_serialization is None: raise ImportError("ROS message conversion requires rclpy; install and source ROS 2.") if type(msg) is not derive_ros_type(dimos_type): - raise TypeError(f"ROS source does not match {dimos_type.msg_name}") - return dimos_type.decode(ros_serialization.serialize_message(msg)) + raise TypeError(f"ROS source does not match {dimos_type.__msgtype__}") + return cast( + "MessageT", decode(ros_serialization.serialize_message(msg), dimos_type.__msgtype__) + ) diff --git a/dimos/protocol/pubsub/impl/shmpubsub.py b/dimos/protocol/pubsub/impl/shmpubsub.py new file mode 100644 index 00000000000..e6ea1044bec --- /dev/null +++ b/dimos/protocol/pubsub/impl/shmpubsub.py @@ -0,0 +1,353 @@ +# Copyright 2025-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. + + +from __future__ import annotations + +from dataclasses import dataclass +import hashlib +import os +import struct +import threading +import time +from typing import TYPE_CHECKING, Any +import uuid + +import numpy as np +import numpy.typing as npt + +from dimos.protocol.pubsub.encoders import CDREncoderMixin, PickleEncoderMixin +from dimos.protocol.pubsub.impl.lcmpubsub import Topic +from dimos.protocol.pubsub.shm.ipc_factory import CpuShmChannel, FrameChannel +from dimos.protocol.pubsub.spec import PubSub +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from collections.abc import Callable + +logger = setup_logger() + + +@dataclass +class SharedMemoryConfig: + prefer: str = "auto" # "auto" | "cpu" (DIMOS_IPC_BACKEND overrides), TODO: "cuda" + default_capacity: int = 3686400 # payload bytes (excludes 4-byte header) + close_channels_on_stop: bool = True + + +class SharedMemoryPubSubBase(PubSub[str, Any]): + """ + Pub/Sub over SharedMemory/CUDA-IPC, modeled after LCMPubSubBase but self-contained. + Wire format per topic/frame: [len:uint32_le] + payload bytes (padded to fixed capacity). + Features ported from Service: + - start()/stop() lifecycle + - one frame channel per topic + - per-topic fanout thread (reads from channel, invokes subscribers) + - CPU/CUDA backend selection (auto + env override) + - reconfigure(topic, capacity=...) + - drop initial empty frame; synchronous local delivery; echo suppression + """ + + # Double-buffered streaming retains the latest frame. + _channel_class: type[FrameChannel] = CpuShmChannel + + # Extra keyword args passed to _channel_class(...) at construction (empty for + # the streaming channel; ShmRPC sets e.g. {"slots": N}). Folded into the + # segment name so distinct layouts never mmap over one another. + _channel_kwargs: dict[str, Any] = {} + + # Per-topic state + # TODO: implement "is_cuda" below capacity, above cp + class _TopicState: + __slots__ = ( + "capacity", + "channel", + "cp", + "dtype", + "last_local_payload", + "last_seq", + "publish_buffer", + "publish_lock", + "shape", + "stop", + "subs", + "thread", + ) + + def __init__(self, channel, capacity: int, cp_mod) -> None: # type: ignore[no-untyped-def] + self.channel = channel + self.capacity = int(capacity) + self.shape = (self.capacity + 20,) # +20 for header: length(4) + uuid(16) + self.dtype = np.uint8 + self.subs: list[Callable[[bytes, str], None]] = [] + self.stop = threading.Event() + self.thread: threading.Thread | None = None + self.last_seq = 0 # start at 0 to avoid b"" on first poll + # TODO: implement an initializer variable for is_cuda once CUDA IPC is in + self.cp = cp_mod + self.last_local_payload: bytes | None = None + # Pre-allocated buffer to avoid allocation on every publish + self.publish_buffer: npt.NDArray[np.uint8] = np.zeros(self.shape, dtype=self.dtype) + # Lock for thread-safe publish buffer access + self.publish_lock = threading.Lock() + + # init / lifecycle + + def __init__( + self, + *, + prefer: str = "auto", + default_capacity: int = 3686400, + close_channels_on_stop: bool = True, + **_: Any, + ) -> None: + super().__init__() + self.config = SharedMemoryConfig( + prefer=prefer, + default_capacity=default_capacity, + close_channels_on_stop=close_channels_on_stop, + ) + self._topics: dict[str, SharedMemoryPubSubBase._TopicState] = {} + self._lock = threading.Lock() + self._publisher_id = uuid.uuid4().bytes + + def start(self) -> None: + pref = (self.config.prefer or "auto").lower() + backend = os.getenv("DIMOS_IPC_BACKEND", pref).lower() + logger.debug(f"SharedMemory PubSub starting (backend={backend})") + # No global thread needed; per-topic fanout starts on first subscribe. + + def stop(self) -> None: + with self._lock: + for _topic, st in list(self._topics.items()): + # stop fanout + try: + if st.thread: + st.stop.set() + st.thread.join(timeout=0.5) + st.thread = None + except Exception: + pass + # close/unlink channels if configured + if self.config.close_channels_on_stop: + try: + st.channel.close() + except Exception: + pass + self._topics.clear() + logger.debug("SharedMemory PubSub stopped.") + + # PubSub API (bytes on the wire) + + def publish(self, topic: str, message: bytes) -> None: + if not isinstance(message, bytes | bytearray | memoryview): + raise TypeError(f"publish expects bytes-like, got {type(message)!r}") + + st = self._ensure_topic(topic) + + # Normalize once + payload_bytes = bytes(message) + L = len(payload_bytes) + if L > st.capacity: + logger.error(f"Payload too large: {L} > capacity {st.capacity}") + raise ValueError(f"Payload too large: {L} > capacity {st.capacity}") + + # A stable publisher identity suppresses local echoes without retaining + # an unbounded set of messages that latest-frame delivery can skip. + message_id = self._publisher_id + + # Synchronous local delivery first (zero extra copies) + for cb in list(st.subs): + try: + cb(payload_bytes, topic) + except Exception: + logger.warn(f"Payload couldn't be pushed to topic: {topic}") + pass + + # Build host frame [len:4] + [uuid:16] + payload and publish + # We embed the message UUID in the frame for echo suppression + # Reuse pre-allocated buffer to avoid allocation overhead + # Lock to prevent concurrent threads from corrupting the shared buffer + with st.publish_lock: + host = st.publish_buffer + # Pack: length(4) + uuid(16) + payload + header = struct.pack(" Callable[[], None]: + """Subscribe a callback(message: bytes, topic). Returns unsubscribe.""" + st = self._ensure_topic(topic) + st.subs.append(callback) + if st.thread is None: + if isinstance(st.channel, CpuShmChannel): + # Frames already in the segment predate this subscriber. + st.last_seq = st.channel.current_seq() + st.thread = threading.Thread(target=self._fanout_loop, args=(topic, st), daemon=True) + st.thread.start() + + def _unsub() -> None: + try: + st.subs.remove(callback) + except ValueError: + pass + if not st.subs and st.thread: + st.stop.set() + st.thread.join(timeout=0.5) + st.thread = None + st.stop.clear() + + return _unsub + + # Capacity mgmt + + def reconfigure(self, topic: str, *, capacity: int) -> dict: # type: ignore[type-arg] + """Change payload capacity (bytes) for a topic; returns new descriptor.""" + st = self._ensure_topic(topic) + new_cap = int(capacity) + new_shape = (new_cap + 20,) # +20 for header: length(4) + uuid(16) + # Lock to ensure no publish is using the buffer while we replace it + with st.publish_lock: + desc = st.channel.reconfigure(new_shape, np.uint8) + st.capacity = new_cap + st.shape = new_shape + st.dtype = np.uint8 + st.last_seq = -1 + st.publish_buffer = np.zeros(new_shape, dtype=np.uint8) + return desc # type: ignore[no-any-return] + + # Internals + + def _ensure_topic(self, topic: str) -> _TopicState: + with self._lock: + st = self._topics.get(topic) + if st is not None: + return st + cap = int(self.config.default_capacity) + + def _names_for_topic(topic: str, capacity: int) -> tuple[str, str]: + # Python's SharedMemory requires names without a leading '/'. + # Fold the channel class name AND its layout kwargs (e.g. slots) + # into the hash so two pubsubs using different layouts (streaming + # CpuShmChannel vs RPC CpuShmQueue, or two ring sizes) on the same + # topic get distinct segments rather than mmapping incompatible + # layouts over each other. + # Use a short digest to avoid macOS shared memory name length limits. + layout = f"{self._channel_class.__name__}:{sorted(self._channel_kwargs.items())}" + h = hashlib.blake2b( + f"{topic}:{capacity}:{layout}".encode(), digest_size=8 + ).hexdigest() + return f"psm_{h}_data", f"psm_{h}_ctrl" + + data_name, ctrl_name = _names_for_topic(topic, cap) + ch = self._channel_class( + (cap + 20,), + np.uint8, + data_name=data_name, + ctrl_name=ctrl_name, + **self._channel_kwargs, + ) + st = SharedMemoryPubSubBase._TopicState(ch, cap, None) + self._topics[topic] = st + return st + + def _fanout_loop(self, topic: str, st: _TopicState) -> None: + while not st.stop.is_set(): + seq, _ts_ns, view = st.channel.read(last_seq=st.last_seq, require_new=True) + if view is None: + time.sleep(0.001) + continue + st.last_seq = seq + + host = np.array(view, copy=True) + + try: + # Read header: length(4) + uuid(16) + L = struct.unpack(" st.capacity + 16: + continue + + # Extract UUID + message_id = host[4:20].tobytes() + + # Extract actual payload (after removing the 16 bytes for uuid) + payload_len = L - 16 + payload = host[20 : 20 + payload_len].tobytes() + + if message_id == self._publisher_id: + continue + + except Exception: + continue + + for cb in list(st.subs): + try: + cb(payload, topic) + except Exception: + pass + + +BytesSharedMemory = SharedMemoryPubSubBase + + +class PickleSharedMemory( + PickleEncoderMixin[str, Any], + SharedMemoryPubSubBase, +): + """SharedMemory pubsub that transports arbitrary Python objects via pickle.""" + + ... + + +class TypedSharedMemoryPubSubBase(PubSub[Topic, Any]): + """SharedMemory pubsub that uses LCM Topic type, delegating to SharedMemoryPubSubBase.""" + + def __init__(self, **kwargs: Any) -> None: + super().__init__() + self._shm = SharedMemoryPubSubBase(**kwargs) + + def start(self) -> None: + self._shm.start() + + def stop(self) -> None: + self._shm.stop() + + def publish(self, topic: Topic, message: bytes) -> None: + self._shm.publish(str(topic), message) + + def subscribe( + self, topic: Topic, callback: Callable[[bytes, Topic], Any] + ) -> Callable[[], None]: + def wrapper(msg: bytes, _: str) -> None: + callback(msg, topic) + + return self._shm.subscribe(str(topic), wrapper) + + def reconfigure(self, topic: Topic, *, capacity: int) -> dict: # type: ignore[type-arg] + return self._shm.reconfigure(str(topic), capacity=capacity) + + +class CDRSharedMemory( # type: ignore[misc] + CDREncoderMixin, + TypedSharedMemoryPubSubBase, +): + """SharedMemory pubsub that uses CDR encoding (no pickle overhead).""" + + ... diff --git a/dimos/protocol/pubsub/impl/test_jpeg_cdr.py b/dimos/protocol/pubsub/impl/test_jpeg_cdr.py index b1827a2abf9..effabf17d4c 100644 --- a/dimos/protocol/pubsub/impl/test_jpeg_cdr.py +++ b/dimos/protocol/pubsub/impl/test_jpeg_cdr.py @@ -18,6 +18,7 @@ import pickle from types import SimpleNamespace +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import CompressedImage, Image from dimos_generated.std_msgs.msg import Header @@ -38,16 +39,16 @@ def test_jpeg_wire_is_standard_cdr_and_preserves_exact_image_header(codec): np.full((16, 16, 3), [20, 80, 140], dtype=np.uint8), encoding="rgb8", header=header ) topic = SimpleNamespace(topic="/image", msg_type=CompressedImage) - before = image.encode() + before = cdr_encode(image) encoded = codec.encode(image, topic) - wire = CompressedImage.decode(encoded) + wire = cdr_decode(encoded, CompressedImage) assert encoded[:4] == b"\x00\x01\x00\x00" assert wire.header == header assert wire.format == "rgb8; jpeg compressed bgr8" decoded = codec.decode(encoded, topic) assert decoded.encoding == "bgr8" and decoded.header == header np.testing.assert_allclose(image_view(decoded)[0, 0], [140, 80, 20], atol=3) - assert image.encode() == before + assert cdr_encode(image) == before def test_jpeg_lcm_requires_truthful_compressed_wire_type(): @@ -71,6 +72,18 @@ def test_jpeg_transport_wire_metadata_and_pickle_keep_separate_application_type( def test_compressed_image_rejects_invalid_data_and_unknown_format(): with pytest.raises(ValueError, match="invalid compressed"): - image_from_compressed(CompressedImage(format="jpeg", data=b"not jpeg")) + image_from_compressed( + CompressedImage( + format="jpeg", + data=np.frombuffer(b"not jpeg", dtype=np.uint8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + ) with pytest.raises(ValueError, match="unsupported compressed"): - image_from_compressed(CompressedImage(format="h264", data=b"not jpeg")) + image_from_compressed( + CompressedImage( + format="h264", + data=np.frombuffer(b"not jpeg", dtype=np.uint8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + ) diff --git a/dimos/protocol/pubsub/impl/test_lcmpubsub.py b/dimos/protocol/pubsub/impl/test_lcmpubsub.py index 92a827d0289..3acf3f8904d 100644 --- a/dimos/protocol/pubsub/impl/test_lcmpubsub.py +++ b/dimos/protocol/pubsub/impl/test_lcmpubsub.py @@ -15,6 +15,7 @@ from collections.abc import Iterator from typing import Any +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion, Vector3 from dimos_generated.std_msgs.msg import String import pytest @@ -60,7 +61,7 @@ def test_LCMPubSubBase_pubsub(lcm_pub_sub_base: LCMPubSubBase) -> None: test_message = String(data="test_data") lcm.subscribe(topic, collector) - lcm.publish(topic, test_message.encode()) + lcm.publish(topic, cdr_encode(test_message)) collector.wait() assert len(collector.results) == 1 @@ -69,7 +70,7 @@ def test_LCMPubSubBase_pubsub(lcm_pub_sub_base: LCMPubSubBase) -> None: received_topic = collector.results[0][1] assert isinstance(received_data, bytes) - assert String.decode(received_data).data == "test_data" + assert cdr_decode(received_data, String).data == "test_data" assert isinstance(received_topic, Topic) assert received_topic == topic @@ -91,7 +92,7 @@ def test_lcm_autodecoder_pubsub(lcm: LCM) -> None: received_topic = collector.results[0][1] assert isinstance(received_data, String) - assert received_data.encode() == test_message.encode() + assert cdr_encode(received_data) == cdr_encode(test_message) assert isinstance(received_topic, Topic) assert received_topic == topic @@ -102,7 +103,12 @@ def test_invalid_cdr_does_not_stop_the_subscriber(lcm: LCM) -> None: collector = CallbackCollector(1) lcm.subscribe(topic, collector) - lcm.publish(topic, b"invalid CDR") + # Exercise decode-error recovery with an actually truncated native payload; + # native rosbags does not validate arbitrary encapsulation bytes strictly. + truncated = cdr_encode(String(data="truncated"))[:4] + with pytest.raises(ValueError, match="Invalid CDR"): + cdr_decode(truncated, String) + lcm.publish(topic, truncated) lcm.publish(topic, String(data="after malformed payload")) collector.wait() @@ -111,7 +117,7 @@ def test_invalid_cdr_does_not_stop_the_subscriber(lcm: LCM) -> None: def test_different_message_types_cannot_be_published_on_a_typed_topic(lcm: LCM) -> None: with pytest.raises(ValueError, match="does not match"): - lcm.publish(Topic("/point", Vector3), Quaternion(w=1)) + lcm.publish(Topic("/point", Vector3), Quaternion(w=1, x=0.0, y=0.0, z=0.0)) def test_explicit_unknown_type_does_not_use_default_decoder() -> None: @@ -130,7 +136,7 @@ def test_invalid_lcm_channel_fails_before_sending(lcm: LCM, channel: str) -> Non test_msgs = [ (Vector3(x=1, y=2, z=3)), (Quaternion(x=1, y=2, z=3, w=4)), - (Pose(position=Point(x=1, y=2, z=3), orientation=Quaternion(w=1))), + (Pose(position=Point(x=1, y=2, z=3), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0))), ] @@ -151,7 +157,7 @@ def test_lcm_geometry_msgs_pubsub(test_message: Any, lcm: LCM) -> None: received_topic = collector.results[0][1] assert isinstance(received_data, test_message.__class__) - assert received_data.encode() == test_message.encode() + assert cdr_encode(received_data) == cdr_encode(test_message) assert isinstance(received_topic, Topic) assert received_topic == topic @@ -175,7 +181,7 @@ def test_lcm_geometry_msgs_autopickle_pubsub(test_message: Any, pickle_lcm: Pick received_topic = collector.results[0][1] assert isinstance(received_data, test_message.__class__) - assert received_data.encode() == test_message.encode() + assert cdr_encode(received_data) == cdr_encode(test_message) assert isinstance(received_topic, Topic) assert received_topic == topic diff --git a/dimos/protocol/pubsub/impl/test_rospubsub.py b/dimos/protocol/pubsub/impl/test_rospubsub.py index 48878767833..a08810ea77f 100644 --- a/dimos/protocol/pubsub/impl/test_rospubsub.py +++ b/dimos/protocol/pubsub/impl/test_rospubsub.py @@ -19,9 +19,18 @@ import uuid from dimos_generated.builtin_interfaces.msg import Time -from dimos_generated.geometry_msgs.msg import Point, PointStamped, PoseStamped, Twist, Vector3 +from dimos_generated.geometry_msgs.msg import ( + Point, + PointStamped, + Pose, + PoseStamped, + Quaternion, + Twist, + Vector3, +) from dimos_generated.sensor_msgs.msg import Image, PointCloud2, PointField from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from dimos.protocol.pubsub.impl.rospubsub import ROS_AVAILABLE, DimosROS, ROSTopic @@ -45,21 +54,51 @@ def nodes() -> Generator[tuple[DimosROS, DimosROS], None, None]: "original", [ Vector3(x=1.0, y=2.0, z=3.0), - Twist(linear=Vector3(x=1.0), angular=Vector3(z=-0.5)), - PoseStamped(header=Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id="map")), - PointStamped(header=Header(stamp=Time(sec=-1, nanosec=999999999)), point=Point(x=1.5)), - PointCloud2(), + Twist(linear=Vector3(x=1.0, y=0.0, z=0.0), angular=Vector3(z=-0.5, x=0.0, y=0.0)), + PoseStamped( + header=Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id="map"), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + PointStamped( + header=Header(stamp=Time(sec=-1, nanosec=999999999), frame_id=""), + point=Point(x=1.5, y=0.0, z=0.0), + ), + PointCloud2( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + height=0, + width=0, + fields=[], + is_bigendian=False, + point_step=0, + row_step=0, + data=np.array([], dtype=np.uint8), + is_dense=False, + ), PointCloud2( height=1, width=1, - fields=[PointField(name="x", datatype=7, count=1)], + fields=[PointField(name="x", datatype=7, count=1, offset=0)], point_step=4, row_step=4, - data=bytes([0, 0, 128, 63]), + data=np.frombuffer(bytes([0, 0, 128, 63]), dtype=np.uint8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + is_bigendian=False, + is_dense=False, + ), + Image( + height=1, + width=2, + encoding="rgb8", + step=6, + data=np.frombuffer(bytes([0, 1, 2, 253, 254, 255]), dtype=np.uint8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + is_bigendian=0, ), - Image(height=1, width=2, encoding="rgb8", step=6, data=bytes([0, 1, 2, 253, 254, 255])), ], - ids=lambda msg: msg.msg_name, + ids=lambda msg: msg.__msgtype__, ) def test_generated_pubsub(nodes, original): publisher, subscriber = nodes diff --git a/dimos/protocol/pubsub/impl/test_rospubsub_conversion.py b/dimos/protocol/pubsub/impl/test_rospubsub_conversion.py index c9d025e947f..4501e4bce91 100644 --- a/dimos/protocol/pubsub/impl/test_rospubsub_conversion.py +++ b/dimos/protocol/pubsub/impl/test_rospubsub_conversion.py @@ -14,6 +14,7 @@ """Conformance with installed ROS type support; skipped when ROS is absent.""" +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, @@ -21,6 +22,7 @@ Pose, PoseStamped, Quaternion, + Transform, TransformStamped, Vector3, ) @@ -30,9 +32,11 @@ Image, PointCloud2, PointField, + RegionOfInterest, ) from dimos_generated.std_msgs.msg import Header from dimos_generated.tf2_msgs.msg import TFMessage +import numpy as np import pytest from dimos.protocol.pubsub.impl import rospubsub_conversion as conversion @@ -44,24 +48,38 @@ def messages(): header = Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id="map/相機") return [ Point(x=1.25, y=-2.5, z=3.75), - PointStamped(header=header, point=Point(x=2.0)), - PoseStamped(header=header, pose=Pose(orientation=Quaternion(z=0.6, w=0.8))), + PointStamped(header=header, point=Point(x=2.0, y=0.0, z=0.0)), + PoseStamped( + header=header, + pose=Pose( + orientation=Quaternion(z=0.6, w=0.8, x=0.0, y=0.0), + position=Point(x=0.0, y=0.0, z=0.0), + ), + ), Image( header=header, height=1, width=2, encoding="rgb8", step=6, - data=bytes([0, 1, 2, 253, 254, 255]), + data=np.frombuffer(bytes([0, 1, 2, 253, 254, 255]), dtype=np.uint8), + is_bigendian=0, + ), + CompressedImage( + header=header, format="jpeg", data=np.frombuffer(bytes([0, 255, 127]), dtype=np.uint8) ), - CompressedImage(header=header, format="jpeg", data=bytes([0, 255, 127])), CameraInfo( header=header, width=640, height=480, distortion_model="plumb_bob", - d=[0.1, -0.2], - k=[1.0, 0.0, 3.0, 0.0, 2.0, 4.0, 0.0, 0.0, 1.0], + d=np.array([0.1, -0.2], dtype=np.float64), + k=np.array([1.0, 0.0, 3.0, 0.0, 2.0, 4.0, 0.0, 0.0, 1.0], dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), ), PointCloud2( header=header, @@ -71,21 +89,43 @@ def messages(): is_bigendian=True, point_step=8, row_step=12, - data=bytes([9, 8, 7, 6, 0x3F, 0x80, 0, 0, 5, 4, 3, 2]), + data=np.frombuffer(bytes([9, 8, 7, 6, 0x3F, 0x80, 0, 0, 5, 4, 3, 2]), dtype=np.uint8), + is_dense=False, + ), + PointCloud2( + header=header, + height=0, + width=0, + fields=[], + is_bigendian=False, + point_step=0, + row_step=0, + data=np.array([], dtype=np.uint8), + is_dense=False, + ), + TFMessage(transforms=[]), + TFMessage( + transforms=[ + TransformStamped( + header=header, + child_frame_id="camera", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ] ), - PointCloud2(header=header), - TFMessage(), - TFMessage(transforms=[TransformStamped(header=header, child_frame_id="camera")]), ] @requires_ros -@pytest.mark.parametrize("original", messages(), ids=lambda msg: msg.msg_name) +@pytest.mark.parametrize("original", messages(), ids=lambda msg: msg.__msgtype__) def test_generated_fields_survive_ros_serialization(original): ros_type = conversion.derive_ros_type(type(original)) ros_message = conversion.dimos_to_ros(original, ros_type) converted = conversion.ros_to_dimos(ros_message, type(original)) - assert converted == original + assert cdr_encode(converted) == cdr_encode(original) @requires_ros @@ -111,15 +151,15 @@ def test_matching_layout_does_not_allow_wrong_type(): from_type = conversion.derive_ros_type(Point) to_type = conversion.derive_ros_type(Vector3) with pytest.raises(TypeError, match="target does not match"): - conversion.dimos_to_ros(Point(), to_type) + conversion.dimos_to_ros(Point(x=0.0, y=0.0, z=0.0), to_type) with pytest.raises(TypeError, match="source does not match"): conversion.ros_to_dimos(from_type(), Vector3) def test_missing_ros_is_explicit_and_does_not_affect_generated_messages(monkeypatch): monkeypatch.setattr(conversion, "ros_serialization", None) - message = Point(x=1.0) - assert Point.decode(message.encode()) == message + message = Point(x=1.0, y=0.0, z=0.0) + assert cdr_decode(cdr_encode(message), Point) == message with pytest.raises(ImportError, match="requires rclpy"): conversion.dimos_to_ros(message, Point) with pytest.raises(ImportError, match="requires rclpy"): @@ -128,6 +168,6 @@ def test_missing_ros_is_explicit_and_does_not_affect_generated_messages(monkeypa @pytest.mark.parametrize("name", ["geometry_msgs.Point", "geometry_msgs/srv/Point", "../msg/Point"]) def test_invalid_type_name_fails_before_import(name): - invalid = type("Invalid", (), {"msg_name": name}) + invalid = type("Invalid", (), {"__msgtype__": name}) with pytest.raises(ValueError, match="expected package/msg/Type"): conversion.derive_ros_type(invalid) diff --git a/dimos/protocol/pubsub/impl/test_shmpubsub.py b/dimos/protocol/pubsub/impl/test_shmpubsub.py new file mode 100644 index 00000000000..58241e04893 --- /dev/null +++ b/dimos/protocol/pubsub/impl/test_shmpubsub.py @@ -0,0 +1,210 @@ +# Copyright 2025-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. + +from contextlib import ExitStack +from multiprocessing import get_context +import threading +from typing import Any +import uuid + +from dimos_generated.geometry_msgs.msg import Vector3 +import numpy as np +import pytest + +from dimos.msgs.image import image_from_array, image_view +from dimos.protocol.pubsub.impl.jpeg_shm import JpegSharedMemory +from dimos.protocol.pubsub.impl.lcmpubsub import Topic +from dimos.protocol.pubsub.impl.shmpubsub import ( + BytesSharedMemory, + CDRSharedMemory, + PickleSharedMemory, +) +from dimos.protocol.pubsub.registry import make_pubsub_transport + + +@pytest.fixture +def buses(): + with ExitStack() as cleanup: + + def make(cls=BytesSharedMemory): + bus = cls(prefer="cpu", default_capacity=4096) + bus.start() + cleanup.callback(bus.stop) + return bus + + yield make + + +@pytest.fixture +def topic() -> str: + return f"shm_test_{uuid.uuid4().hex}" + + +def test_new_subscriber_ignores_existing_frame_and_resubscribe_skips_gap(buses, topic): + owner, reader = buses(), buses() + owner.publish(topic, b"stale") + got = [] + arrived = threading.Event() + + def receive(message, _topic): + got.append(message) + arrived.set() + + unsubscribe = reader.subscribe(topic, receive) + assert not arrived.wait(0.05) + owner.publish(topic, b"fresh") + assert arrived.wait(2) + assert got == [b"fresh"] + unsubscribe() + arrived.clear() + owner.publish(topic, b"missed") + reader.subscribe(topic, receive) + assert not arrived.wait(0.05) + owner.publish(topic, b"second") + assert arrived.wait(2) + assert got == [b"fresh", b"second"] + + +@pytest.mark.parametrize("payload", [b"hello", b""]) +def test_bytes_between_instances(buses, topic, payload): + owner, reader = buses(), buses() + arrived = threading.Event() + got = [] + + def receive(message, channel): + got.append((message, channel)) + arrived.set() + + reader.subscribe(topic, receive) + owner.publish(topic, payload) + assert arrived.wait(2) + assert got == [(payload, topic)] + + +def test_generated_cdr_between_instances(buses, topic): + owner, reader = buses(CDRSharedMemory), buses(CDRSharedMemory) + typed_topic = Topic(topic, Vector3) + arrived = threading.Event() + got = [] + + def receive(message, channel): + got.append((message.x, message.y, message.z, channel.msg_type)) + arrived.set() + + reader.subscribe(typed_topic, receive) + owner.publish(typed_topic, Vector3(x=1.25, y=-2.5, z=3.75)) + assert arrived.wait(2) + assert got == [(1.25, -2.5, 3.75, Vector3)] + + +def test_jpeg_generated_image_between_instances(buses, topic): + owner, reader = buses(JpegSharedMemory), buses(JpegSharedMemory) + arrived = threading.Event() + got = [] + + def receive(message, _topic): + got.append(message) + arrived.set() + + reader.subscribe(topic, receive) + message = image_from_array(np.full((4, 6, 3), 128, dtype=np.uint8), encoding="rgb8") + owner.publish(topic, message) + assert arrived.wait(2) + assert got[0].encoding == "bgr8" + np.testing.assert_array_equal(image_view(got[0]), np.full((4, 6, 3), 128, dtype=np.uint8)) + + +def _publish_child(topic: str, pipe: Any) -> None: + bus = PickleSharedMemory(prefer="cpu", default_capacity=4096) + try: + bus.start() + bus.publish(topic, "stale") + pipe.send("ready") + assert pipe.recv() == "publish" + bus.publish(topic, {"value": [1, 2, 3]}) + pipe.send("published") + assert pipe.recv() == "stop" + finally: + bus.stop() + pipe.close() + + +@pytest.fixture +def publisher_process(topic): + context = get_context("fork") + parent, child = context.Pipe() + process = context.Process(target=_publish_child, args=(topic, child)) + process.start() + child.close() + try: + assert parent.poll(5) + assert parent.recv() == "ready" + yield process, parent + finally: + if process.is_alive(): + parent.send("stop") + process.join(5) + if process.is_alive(): + process.terminate() + process.join(5) + parent.close() + + +def test_pickle_across_processes(buses, topic, publisher_process): + process, pipe = publisher_process + reader = buses(PickleSharedMemory) + arrived = threading.Event() + got = [] + + def receive(message, _topic): + got.append(message) + arrived.set() + + reader.subscribe(topic, receive) + pipe.send("publish") + assert pipe.poll(5) + assert pipe.recv() == "published" + assert arrived.wait(2) + assert got == [{"value": [1, 2, 3]}] + pipe.send("stop") + process.join(5) + assert process.exitcode == 0 + + +@pytest.fixture +def uri_transports(topic): + with ExitStack() as cleanup: + transports = [ + make_pubsub_transport(f"shm:{topic}#geometry_msgs/msg/Vector3") for _ in range(2) + ] + for transport in transports: + transport.start() + cleanup.callback(transport.stop) + yield transports + + +def test_registry_cdr_roundtrip(uri_transports): + publisher, subscriber = uri_transports + arrived = threading.Event() + got = [] + + def receive(message): + got.append((message.x, message.y, message.z)) + arrived.set() + + unsubscribe = subscriber.subscribe(receive) + publisher.broadcast(None, Vector3(x=2.5, y=-3.5, z=4.5)) + assert arrived.wait(2) + assert got == [(2.5, -3.5, 4.5)] + unsubscribe() diff --git a/dimos/protocol/pubsub/impl/webrtc/providers/test_video_track.py b/dimos/protocol/pubsub/impl/webrtc/providers/test_video_track.py index 4a8abcf5eb6..76ab82933a1 100644 --- a/dimos/protocol/pubsub/impl/webrtc/providers/test_video_track.py +++ b/dimos/protocol/pubsub/impl/webrtc/providers/test_video_track.py @@ -18,11 +18,12 @@ from __future__ import annotations import asyncio +from dataclasses import replace import numpy as np import pytest -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_from_array pytest.importorskip("aiortc") @@ -34,12 +35,12 @@ async def test_16bit_frames_are_converted_to_rgb() -> None: track = CameraVideoTrack(asyncio.get_running_loop()) track.arm() - for fmt in (ImageFormat.GRAY16, ImageFormat.DEPTH16): - track.set_latest(Image(data=np.full((4, 6), 0x8000, np.uint16), format=fmt)) + for encoding in ("mono16", "16UC1"): + track.set_latest(image_from_array(np.full((4, 6), 0x8000, np.uint16), encoding=encoding)) frame = await asyncio.wait_for(track.recv(), 2) assert (frame.format.name, frame.width, frame.height) == ("rgb24", 6, 4) assert frame.to_ndarray()[0, 0].tolist() == [128, 128, 128] - track.set_latest(Image(data=np.zeros((4, 6, 3), np.uint8), format=ImageFormat.BGR)) + track.set_latest(image_from_array(np.zeros((4, 6, 3), np.uint8), encoding="bgr8")) frame = await asyncio.wait_for(track.recv(), 2) assert frame.format.name == "bgr24" @@ -49,7 +50,7 @@ async def test_odd_sized_frames_are_cropped_to_what_the_h264_encoder_opens() -> # opens, inside aiortc's sender task, so the encoder itself runs here. track = CameraVideoTrack(asyncio.get_running_loop()) track.arm() - track.set_latest(Image(data=np.zeros((47, 63, 3), np.uint8), format=ImageFormat.RGB)) + track.set_latest(image_from_array(np.zeros((47, 63, 3), np.uint8), encoding="rgb8")) frame = await asyncio.wait_for(track.recv(), 2) assert (frame.width, frame.height) == (62, 46) payloads, _timestamp = H264Encoder().encode(frame, force_keyframe=True) @@ -61,12 +62,15 @@ async def test_unconvertible_frames_are_skipped_not_raised() -> None: track.arm() pending = asyncio.ensure_future(track.recv()) for bad in ( - Image(data=np.zeros((4, 6), np.float32), format=ImageFormat.DEPTH), # no av format - Image(data=np.zeros((4, 6, 3), np.uint16), format=ImageFormat.GRAY16), # to_rgb() raises + image_from_array(np.zeros((4, 6), np.float32), encoding="32FC1"), # no av format + replace( + image_from_array(np.zeros((4, 6), np.uint16), encoding="mono16"), + step=1, # malformed wire metadata must not end the sender task + ), ): track.set_latest(bad) await asyncio.sleep(0.05) assert not pending.done() - track.set_latest(Image(data=np.zeros((4, 6), np.uint8), format=ImageFormat.GRAY)) + track.set_latest(image_from_array(np.zeros((4, 6), np.uint8), encoding="mono8")) frame = await asyncio.wait_for(pending, 2) assert frame.format.name == "gray" diff --git a/dimos/protocol/pubsub/impl/webrtc/providers/test_video_track_generated.py b/dimos/protocol/pubsub/impl/webrtc/providers/test_video_track_generated.py index 375a635b1cf..12b5b410506 100644 --- a/dimos/protocol/pubsub/impl/webrtc/providers/test_video_track_generated.py +++ b/dimos/protocol/pubsub/impl/webrtc/providers/test_video_track_generated.py @@ -17,6 +17,7 @@ import asyncio +from dimos_cdrgen.registry import encode as cdr_encode import numpy as np import pytest @@ -31,7 +32,7 @@ async def test_video_track_uses_generated_encoding_without_network(encoding, channels): pixels = np.full((4, 4) if channels == 1 else (4, 4, channels), 127, dtype=np.uint8) image = image_from_array(pixels, encoding=encoding) - original = image.encode() + original = cdr_encode(image) track = CameraVideoTrack(asyncio.get_running_loop()) try: track.arm() @@ -39,6 +40,6 @@ async def test_video_track_uses_generated_encoding_without_network(encoding, cha frame = await asyncio.wait_for(track.recv(), timeout=1) assert frame.width == 4 and frame.height == 4 assert frame.pts == 0 - assert image.encode() == original + assert cdr_encode(image) == original finally: track.stop() diff --git a/dimos/protocol/pubsub/impl/webrtc/providers/video_track.py b/dimos/protocol/pubsub/impl/webrtc/providers/video_track.py index eaa443b4367..553eb4243fe 100644 --- a/dimos/protocol/pubsub/impl/webrtc/providers/video_track.py +++ b/dimos/protocol/pubsub/impl/webrtc/providers/video_track.py @@ -27,19 +27,20 @@ from aiortc import RTCRtpSender from aiortc.mediastreams import VIDEO_CLOCK_RATE, VIDEO_TIME_BASE, VideoStreamTrack import av +from dimos_generated.sensor_msgs.msg import Image import numpy as np -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_to_rgb, image_view from dimos.utils.logging_config import setup_logger logger = setup_logger() _AV_FORMAT_MAP = { - ImageFormat.BGR: "bgr24", - ImageFormat.RGB: "rgb24", - ImageFormat.BGRA: "bgra", - ImageFormat.RGBA: "rgba", - ImageFormat.GRAY: "gray", + "bgr8": "bgr24", + "rgb8": "rgb24", + "bgra8": "bgra", + "rgba8": "rgba", + "mono8": "gray", } @@ -78,7 +79,7 @@ def __init__(self, loop: asyncio.AbstractEventLoop) -> None: self._armed = False self._first_mono: float | None = None self._new_frame = asyncio.Event() - self._dropped: set[ImageFormat] = set() + self._dropped: set[str] = set() def arm(self) -> None: """Discard buffered frames; start delivering from now. @@ -107,24 +108,26 @@ def _set() -> None: return def _to_av(self, img: Image) -> av.VideoFrame | None: - """None, logged once per source format, for what neither to_rgb() nor - av takes (DEPTH float32, a malformed array). Never raises: recv() runs - in aiortc's sender task, which ends the track on an exception.""" + """Drop unsupported or malformed frames without ending aiortc's sender task.""" try: - frame = img if img.format in _AV_FORMAT_MAP else img.to_rgb() - data = frame.data + if img.encoding in _AV_FORMAT_MAP: + data = image_view(img) + av_format = _AV_FORMAT_MAP[img.encoding] + else: + data = image_to_rgb(img) + av_format = "rgb24" h, w = data.shape[:2] if h % 2 or w % 2: # libx264 takes yuv420p, whose chroma planes are half-size: an # odd width or height fails avcodec_open2 in aiortc's sender # task, past this guard. The last row/column is cropped. data = np.ascontiguousarray(data[: h - h % 2, : w - w % 2]) - return av.VideoFrame.from_ndarray(data, format=_AV_FORMAT_MAP[frame.format]) + return av.VideoFrame.from_ndarray(data, format=av_format) except Exception as e: - if img.format not in self._dropped: - self._dropped.add(img.format) + if img.encoding not in self._dropped: + self._dropped.add(img.encoding) logger.warning( - "video track: dropping %s/%s frames: %s", img.format.value, img.data.dtype, e + "video track: dropping %s/%s frames: %s", img.encoding, img.data.dtype, e ) return None diff --git a/dimos/protocol/pubsub/impl/webrtc/test_transport.py b/dimos/protocol/pubsub/impl/webrtc/test_transport.py index c3bc29b1302..901ec3bcbae 100644 --- a/dimos/protocol/pubsub/impl/webrtc/test_transport.py +++ b/dimos/protocol/pubsub/impl/webrtc/test_transport.py @@ -19,6 +19,7 @@ from collections.abc import Callable import pickle +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import Point, TwistStamped, Vector3 from pydantic import ValidationError import pytest @@ -100,7 +101,7 @@ def test_typed_encode_decode() -> None: transport = MockTransport("cmd_unreliable", Point, name="typed") received: list[Point] = [] transport.subscribe(lambda msg: received.append(msg)) - transport.broadcast(None, Point(x=3.14)) + transport.broadcast(None, Point(x=3.14, y=0.0, z=0.0)) assert len(received) == 1 assert abs(received[0].x - 3.14) < 1e-9 @@ -114,8 +115,8 @@ def test_multiple_types_multiplexed() -> None: t1.subscribe(lambda msg: r1.append(msg)) t2.subscribe(lambda msg: r2.append(msg)) - t1.broadcast(None, Point(x=1.0)) - t2.broadcast(None, Vector3(x=2.0)) + t1.broadcast(None, Point(x=1.0, y=0.0, z=0.0)) + t2.broadcast(None, Vector3(x=2.0, y=0.0, z=0.0)) assert [m.x for m in r1] == [1.0] assert [m.x for m in r2] == [2.0] @@ -125,13 +126,15 @@ def test_generated_frames_declare_type_and_sequence() -> None: raw = MockTransport("data", name="framed") received: list[bytes] = [] raw.subscribe(received.append) - publisher.broadcast(None, Point(x=1)) - publisher.broadcast(None, Point(x=2)) + publisher.broadcast(None, Point(x=1, y=0.0, z=0.0)) + publisher.broadcast(None, Point(x=2, y=0.0, z=0.0)) frames = [decode_data_frame(data) for data in received] assert [frame.header.seq for frame in frames] == [0, 1] assert all(frame.header.ch == "data" for frame in frames) - assert all(frame.header.meta == {"type": Point.msg_name, "encoding": "cdr"} for frame in frames) - assert [Point.decode(frame.payload).x for frame in frames] == [1, 2] + assert all( + frame.header.meta == {"type": Point.__msgtype__, "encoding": "cdr"} for frame in frames + ) + assert [cdr_decode(frame.payload, Point).x for frame in frames] == [1, 2] @pytest.mark.parametrize("bad", ["type", "topic", "encoding", "truncated", "trailing", "raw"]) @@ -144,24 +147,24 @@ def test_invalid_frame_does_not_poison_subscription(bad: str) -> None: seq=1, ts=0, delivery="latest", - meta={"type": Point.msg_name, "encoding": "cdr"}, + meta={"type": Point.__msgtype__, "encoding": "cdr"}, ) if bad == "topic": header.ch = "other" elif bad in ("type", "encoding"): assert header.meta is not None header.meta[bad] = "unknown" - wire = encode_data_frame(header, Point(x=1).encode()) + wire = encode_data_frame(header, cdr_encode(Point(x=1, y=0.0, z=0.0))) if bad == "truncated": wire = wire[:-1] elif bad == "trailing": wire += b"extra" elif bad == "raw": - wire = Point(x=1).encode() + wire = cdr_encode(Point(x=1, y=0.0, z=0.0)) transport._config.provider().publish("cmd_unreliable", wire) assert received == [] - transport.broadcast(None, Point(x=2)) - assert received == [Point(x=2)] + transport.broadcast(None, Point(x=2, y=0.0, z=0.0)) + assert received == [Point(x=2, y=0.0, z=0.0)] # ─── Pickling + provider sharing ───────────────────────────────────── @@ -181,7 +184,7 @@ def test_pickle_roundtrip_preserves_everything() -> None: # Same config → same per-process provider, so the two halves interoperate. received: list[Point] = [] t2.subscribe(lambda msg: received.append(msg)) - t1.broadcast(None, Point(x=42.0)) + t1.broadcast(None, Point(x=42.0, y=0.0, z=0.0)) assert len(received) == 1 assert abs(received[0].x - 42.0) < 1e-9 diff --git a/dimos/protocol/pubsub/impl/webrtc/tool_blueprint_e2e.py b/dimos/protocol/pubsub/impl/webrtc/tool_blueprint_e2e.py index c00d5d53df3..5feef69b936 100644 --- a/dimos/protocol/pubsub/impl/webrtc/tool_blueprint_e2e.py +++ b/dimos/protocol/pubsub/impl/webrtc/tool_blueprint_e2e.py @@ -27,7 +27,9 @@ import os import time +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Twist, TwistStamped, Vector3 +from dimos_generated.std_msgs.msg import Header import pytest from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser @@ -54,7 +56,14 @@ class TwistSource(Module): @rpc def send(self, x: float) -> None: - self.cmd_webrtc.publish(TwistStamped(twist=Twist(linear=Vector3(x=x)))) + self.cmd_webrtc.publish( + TwistStamped( + twist=Twist( + linear=Vector3(x=x, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + ) class TwistSink(Module): diff --git a/dimos/protocol/pubsub/impl/webrtc/tool_broker_e2e.py b/dimos/protocol/pubsub/impl/webrtc/tool_broker_e2e.py index d82d7ae03d2..88598963cb1 100644 --- a/dimos/protocol/pubsub/impl/webrtc/tool_broker_e2e.py +++ b/dimos/protocol/pubsub/impl/webrtc/tool_broker_e2e.py @@ -37,6 +37,9 @@ import time import urllib.request +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.std_msgs.msg import Header import pytest from dimos.protocol.pubsub.impl.webrtc.providers.spec import WEBRTC_AVAILABLE @@ -231,15 +234,29 @@ def _on_back(payload: object) -> None: # Robot → operator: telemetry through the broker-bridged back channel. for i in range(10): - back_transport.broadcast(None, TwistStamped(twist=Twist(linear=Vector3(z=1.0 + i)))) + back_transport.broadcast( + None, + TwistStamped( + twist=Twist( + linear=Vector3(z=1.0 + i, x=0.0, y=0.0), + angular=Vector3(x=0.0, y=0.0, z=0.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ), + ) await asyncio.sleep(0.05) await _wait_for(lambda: bool(back_bytes), 5.0, "robot->operator telemetry") - back_msg = TwistStamped.decode(decode_data_frame(back_bytes[-1]).payload) + back_msg = cdr_decode(decode_data_frame(back_bytes[-1]).payload, TwistStamped) assert back_msg.twist.linear.z >= 1.0, back_msg.twist.linear sent = 0 for i in range(40): - msg = TwistStamped(twist=Twist(linear=Vector3(x=0.5), angular=Vector3(z=i * 0.01))) + msg = TwistStamped( + twist=Twist( + linear=Vector3(x=0.5, y=0.0, z=0.0), angular=Vector3(z=i * 0.01, x=0.0, y=0.0) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) ch.send( encode_data_frame( FrameHeader( @@ -247,9 +264,9 @@ def _on_back(payload: object) -> None: seq=i, ts=time.time(), delivery="latest", - meta={"type": TwistStamped.msg_name, "encoding": "cdr"}, + meta={"type": TwistStamped.__msgtype__, "encoding": "cdr"}, ), - msg.encode(), + cdr_encode(msg), ) ) sent += 1 diff --git a/dimos/protocol/pubsub/impl/zenohpubsub.py b/dimos/protocol/pubsub/impl/zenohpubsub.py index 9b6c677a9f0..ff3b426574a 100644 --- a/dimos/protocol/pubsub/impl/zenohpubsub.py +++ b/dimos/protocol/pubsub/impl/zenohpubsub.py @@ -94,7 +94,7 @@ def key_expr(self) -> str: Topic("dimos/data") -> "dimos/data" """ if self.msg_type is not None: - return f"{self.pattern}/{self.msg_type.msg_name}" + return f"{self.pattern}/{self.msg_type.__msgtype__}" return self.pattern diff --git a/dimos/protocol/pubsub/registry.py b/dimos/protocol/pubsub/registry.py index 84c882fff86..e3868c214f2 100644 --- a/dimos/protocol/pubsub/registry.py +++ b/dimos/protocol/pubsub/registry.py @@ -22,14 +22,15 @@ :[#] -- ````: registry key: ``zenoh``, ``pzenoh``, ``lcm``, ``jpeg_lcm``, ``plcm``. +- ````: registry key: ``zenoh``, ``pzenoh``, ``lcm``, ``jpeg_lcm``, ``plcm``, ``shm``, ``pshm``. - ````: channel/key, passed verbatim to the transport constructor. -- ````: optional ``module.ClassName`` resolved via - ``dimos.msgs.helpers.resolve_msg_type`` (e.g. ``sensor_msgs.Image``). +- ````: optional ``package/msg/Type`` resolved via + ``dimos.msgs.helpers.resolve_msg_type`` (e.g. ``sensor_msgs/msg/Image``). Typed protos (``zenoh``, ``lcm``, ``jpeg_lcm``) require a message type — either from the ``#``-suffix or the ``msg_type`` kwarg. Pickled protos (``pzenoh``, ``plcm``) -don't. Zenoh topics include the namespace, e.g. ``zenoh:dimos/color_image``; +don't. The ``shm`` proto uses CDR with a type and raw bytes without one. +Zenoh topics include the namespace, e.g. ``zenoh:dimos/color_image``; the typed transport appends the message type to the wire key. """ @@ -76,6 +77,20 @@ def _make_zenoh(topic: str, msg_type: type | None) -> Any: return ZenohTransport(topic, msg_type) +def _make_pshm(topic: str, msg_type: type | None) -> Any: + # pickled shared memory: same shape as plcm but over /dev/shm. + from dimos.core.transport import pSHMTransport + + return pSHMTransport(topic) + + +def _make_shm(topic: str, msg_type: type | None) -> Any: + # Explicit type selects CDR; an untyped channel carries raw bytes. + from dimos.core.transport import SHMTransport + + return SHMTransport(topic, msg_type) + + def _make_pzenoh(topic: str, msg_type: type | None) -> Any: from dimos.core.transport import pZenohTransport @@ -86,6 +101,8 @@ def _make_pzenoh(topic: str, msg_type: type | None) -> Any: "lcm": _make_lcm, "jpeg_lcm": _make_jpeg_lcm, "plcm": _make_plcm, + "shm": _make_shm, + "pshm": _make_pshm, "zenoh": _make_zenoh, "pzenoh": _make_pzenoh, } diff --git a/dimos/protocol/pubsub/shm/ipc_factory.py b/dimos/protocol/pubsub/shm/ipc_factory.py new file mode 100644 index 00000000000..5423bbcab95 --- /dev/null +++ b/dimos/protocol/pubsub/shm/ipc_factory.py @@ -0,0 +1,249 @@ +# Copyright 2025-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. + +from abc import ABC, abstractmethod +from multiprocessing.shared_memory import SharedMemory +import os +import time +from typing import Any +import weakref + +import numpy as np +from numpy.typing import DTypeLike + +from dimos.utils.logging_config import setup_logger +from dimos.utils.shm import ShmNotReadyError, attach_shm + +_UNLINK_ON_GC = os.getenv("DIMOS_IPC_UNLINK_ON_GC", "0").lower() not in ("0", "false", "no") + +logger = setup_logger() + + +class FrameChannel(ABC): + """Shared-memory IPC channel carrying frames behind a tiny control block. + + The double buffer retains the latest frame. Descriptors are JSON-safe; + attach() reconstructs a reader in another process. + """ + + @abstractmethod + def __init__( + self, + shape: tuple[int, ...], + dtype: DTypeLike = np.uint8, + *, + data_name: str | None = None, + ctrl_name: str | None = None, + ) -> None: + """Create (or attach by name) the channel's shared-memory segments.""" + ... + + @property + @abstractmethod + def device(self) -> str: # "cpu" or "cuda" + ... + + @property + @abstractmethod + def shape(self) -> tuple: ... # type: ignore[type-arg] + + @property + @abstractmethod + def dtype(self) -> np.dtype: ... + + @abstractmethod + def publish(self, frame, length: int | None = None) -> None: # type: ignore[no-untyped-def] + """Write into inactive buffer, then flip visible index (write control last). + + Args: + frame: The numpy array to publish + length: Optional length to copy (for variable-size messages). If None, copies full frame. + """ + ... + + @abstractmethod + def read(self, last_seq: int = -1, require_new: bool = True): # type: ignore[no-untyped-def] + """Return (seq:int, ts_ns:int, view-or-None).""" + ... + + @abstractmethod + def descriptor(self) -> dict: # type: ignore[type-arg] + """Tiny JSON-safe descriptor (names/handles/shape/dtype/device).""" + ... + + @classmethod + @abstractmethod + def attach(cls, desc: dict) -> "FrameChannel": # type: ignore[type-arg] + """Attach in another process.""" + ... + + @abstractmethod + def close(self) -> None: + """Detach resources (owner also unlinks manager if applicable).""" + ... + + +def _safe_unlink(name: str) -> None: + try: + shm = SharedMemory(name=name) + shm.unlink() # unlink() calls resource_tracker.unregister() + shm.close() + except FileNotFoundError: + pass + except Exception: + pass + + +def _create_or_open(name: str, size: int) -> tuple[SharedMemory, bool]: + """Create a named SHM segment (owner) or attach to an existing one (reader).""" + try: + return SharedMemory(create=True, size=size, name=name), True + except FileExistsError: + segment = attach_shm(name) + if segment.size != size: + segment.close() + raise ValueError(f"SHM {name!r} layout mismatch: expected {size} bytes") + return segment, False + + +class CpuShmChannel(FrameChannel): + def __init__( # type: ignore[no-untyped-def] + self, + shape, + dtype=np.uint8, + *, + data_name: str | None = None, + ctrl_name: str | None = None, + ) -> None: + self._shape = tuple(shape) + self._dtype = np.dtype(dtype) + self._nbytes = int(self._dtype.itemsize * np.prod(self._shape)) + + if data_name is None or ctrl_name is None: + # Fallback: random names (old behavior) -> always owner + self._shm_data = SharedMemory(create=True, size=2 * self._nbytes) + self._shm_ctrl = SharedMemory(create=True, size=24) + self._own_data = self._own_ctrl = True + else: + self._shm_data, self._own_data = _create_or_open(data_name, 2 * self._nbytes) + self._shm_ctrl, self._own_ctrl = _create_or_open(ctrl_name, 24) + + self._ctrl = np.ndarray((3,), dtype=np.int64, buffer=self._shm_ctrl.buf) + if self._own_ctrl: + self._ctrl[:] = 0 # initialize only once + + # only owners set unlink finalizers (beware cross-process timing) + self._finalizer_data = ( + weakref.finalize(self, _safe_unlink, self._shm_data.name) + if (_UNLINK_ON_GC and self._own_data) + else None + ) + self._finalizer_ctrl = ( + weakref.finalize(self, _safe_unlink, self._shm_ctrl.name) + if (_UNLINK_ON_GC and self._own_ctrl) + else None + ) + + @property + def device(self) -> str: + return "cpu" + + @property + def shape(self): # type: ignore[no-untyped-def] + return self._shape + + @property + def dtype(self): # type: ignore[no-untyped-def] + return self._dtype + + def publish(self, frame, length: int | None = None) -> None: # type: ignore[no-untyped-def] + assert isinstance(frame, np.ndarray) + assert frame.shape == self._shape and frame.dtype == self._dtype + active = int(self._ctrl[2]) + inactive = 1 - active + view = np.ndarray( + self._shape, + dtype=self._dtype, + buffer=self._shm_data.buf, + offset=inactive * self._nbytes, + ) + # Only copy actual payload length if specified, otherwise copy full frame + if length is not None and length < len(frame): + np.copyto(view[:length], frame[:length], casting="no") + else: + np.copyto(view, frame, casting="no") + ts = np.int64(time.time_ns()) + # Publish order: ts -> idx -> seq + self._ctrl[1] = ts + self._ctrl[2] = inactive + self._ctrl[0] += 1 + + def current_seq(self) -> int: + return int(self._ctrl[0]) + + def read(self, last_seq: int = -1, require_new: bool = True): # type: ignore[no-untyped-def] + for _ in range(3): + seq1 = int(self._ctrl[0]) + idx = int(self._ctrl[2]) + ts = int(self._ctrl[1]) + view = np.ndarray( + self._shape, dtype=self._dtype, buffer=self._shm_data.buf, offset=idx * self._nbytes + ) + if require_new and seq1 == last_seq: + return seq1, ts, None + snapshot = np.array(view, copy=True) + if seq1 == int(self._ctrl[0]): + return seq1, ts, snapshot + return last_seq, 0, None + + def descriptor(self): # type: ignore[no-untyped-def] + return { + "kind": "cpu", + "shape": self._shape, + "dtype": self._dtype.str, + "nbytes": self._nbytes, + "data_name": self._shm_data.name, + "ctrl_name": self._shm_ctrl.name, + } + + @classmethod + def attach(cls, desc: dict[str, Any]) -> "CpuShmChannel": + obj = object.__new__(cls) + obj._shape = tuple(desc["shape"]) + obj._dtype = np.dtype(desc["dtype"]) + obj._nbytes = int(desc["nbytes"]) + data_name = desc["data_name"] + ctrl_name = desc["ctrl_name"] + try: + obj._shm_data = attach_shm(data_name) + obj._shm_ctrl = attach_shm(ctrl_name) + except ShmNotReadyError as e: + raise ShmNotReadyError( + f"CPU IPC attach failed: control/data SHM never became ready " + f"(ctrl='{ctrl_name}', data='{data_name}'): {e}. " + f"Ensure the writer is running on the same host and the channel is alive." + ) from e + obj._ctrl = np.ndarray((3,), dtype=np.int64, buffer=obj._shm_ctrl.buf) + # attachments don’t own/unlink + obj._own_data = obj._own_ctrl = False + obj._finalizer_data = obj._finalizer_ctrl = None + return obj + + def close(self) -> None: + self._shm_ctrl.close() + self._shm_data.close() + if self._own_ctrl: + _safe_unlink(self._shm_ctrl.name) + if self._own_data: + _safe_unlink(self._shm_data.name) diff --git a/dimos/protocol/pubsub/shm/test_ipc_factory.py b/dimos/protocol/pubsub/shm/test_ipc_factory.py new file mode 100644 index 00000000000..db4b83b721e --- /dev/null +++ b/dimos/protocol/pubsub/shm/test_ipc_factory.py @@ -0,0 +1,90 @@ +# Copyright 2025-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. + +import os +import threading +import uuid + +import numpy as np +import pytest + +from dimos.protocol.pubsub.shm.ipc_factory import CpuShmChannel + + +@pytest.fixture +def channels(): + opened = [] + yield opened + for channel in opened: + channel.close() + + +def test_attach_reads_stable_snapshot(channels): + writer = CpuShmChannel((4,), np.uint8) + channels.append(writer) + reader = CpuShmChannel.attach(writer.descriptor()) + channels.append(reader) + writer.publish(np.array([1, 2, 3, 4], dtype=np.uint8)) + sequence, _, pixels = reader.read() + assert sequence == 1 + np.testing.assert_array_equal(pixels, [1, 2, 3, 4]) + writer.publish(np.array([5, 6, 7, 8], dtype=np.uint8)) + writer.publish(np.array([9, 10, 11, 12], dtype=np.uint8)) + np.testing.assert_array_equal(pixels, [1, 2, 3, 4]) + _, _, newest = reader.read(last_seq=sequence) + np.testing.assert_array_equal(newest, [9, 10, 11, 12]) + + +def test_concurrent_creation_waits_for_sized_segments(monkeypatch, channels): + original = os.ftruncate + creation_started = threading.Event() + allow_sizing = threading.Event() + + def delayed(fd, size): + creation_started.set() + assert allow_sizing.wait(2) + return original(fd, size) + + monkeypatch.setattr(os, "ftruncate", delayed) + tag = uuid.uuid4().hex + errors = [] + + def open_channel(): + try: + channels.append(CpuShmChannel((4,), data_name=f"{tag}_data", ctrl_name=f"{tag}_ctrl")) + except Exception as error: + errors.append(error) + + first = threading.Thread(target=open_channel) + second = threading.Thread(target=open_channel) + first.start() + assert creation_started.wait(2) + second.start() + allow_sizing.set() + first.join(3) + second.join(3) + assert not first.is_alive() and not second.is_alive() + assert errors == [] + assert len(channels) == 2 + channels[0].publish(np.array([4, 3, 2, 1], dtype=np.uint8)) + _, _, pixels = channels[1].read() + np.testing.assert_array_equal(pixels, [4, 3, 2, 1]) + + +def test_named_segment_layout_mismatch_is_rejected(channels): + tag = uuid.uuid4().hex + owner = CpuShmChannel((4,), data_name=f"{tag}_data", ctrl_name=f"{tag}_ctrl") + channels.append(owner) + with pytest.raises(ValueError, match="layout mismatch"): + CpuShmChannel((8,), data_name=f"{tag}_data", ctrl_name=f"{tag}_ctrl") diff --git a/dimos/protocol/pubsub/test_pattern_sub.py b/dimos/protocol/pubsub/test_pattern_sub.py index 10bc1d5290a..e84bf7fcf27 100644 --- a/dimos/protocol/pubsub/test_pattern_sub.py +++ b/dimos/protocol/pubsub/test_pattern_sub.py @@ -82,10 +82,12 @@ def lcm_bytes_context(url: str) -> Iterator[tuple[LCMPubSubBase, LCMPubSubBase]] pubsub_context=lcm_typed_context, topic_values=[ (Topic("/sensor/position", Vector3), Vector3(x=1, y=2, z=3)), - (Topic("/sensor/orientation", Quaternion), Quaternion(w=1)), + (Topic("/sensor/orientation", Quaternion), Quaternion(w=1, x=0.0, y=0.0, z=0.0)), ( Topic("/robot/arm", Pose), - Pose(position=Point(x=4, y=5, z=6), orientation=Quaternion(w=1)), + Pose( + position=Point(x=4, y=5, z=6), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), ), ], tags={"all", "glob", "regex"}, diff --git a/dimos/protocol/pubsub/test_registry.py b/dimos/protocol/pubsub/test_registry.py index 7b1d1e03598..02cdd146278 100644 --- a/dimos/protocol/pubsub/test_registry.py +++ b/dimos/protocol/pubsub/test_registry.py @@ -17,6 +17,9 @@ from contextlib import ExitStack import threading +from dimos_generated.geometry_msgs.msg import Vector3 +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest @@ -25,13 +28,15 @@ JpegLcmTransport, LCMTransport, PubSubTransport, + SHMTransport, ZenohTransport, pLCMTransport, + pSHMTransport, pZenohTransport, ) from dimos.core.transport_factory import make_transport -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_from_array, image_view +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.protocol.pubsub.registry import ( make_pubsub_transport, parse_pubsub_uri, @@ -52,7 +57,7 @@ def session_pool(monkeypatch): def test_supported_protos_includes_known_set() -> None: """Registry exposes the canonical proto names.""" - assert supported_protos() == ["jpeg_lcm", "lcm", "plcm", "pzenoh", "zenoh"] + assert supported_protos() == ["jpeg_lcm", "lcm", "plcm", "pshm", "pzenoh", "shm", "zenoh"] @pytest.mark.parametrize( @@ -63,16 +68,16 @@ def test_supported_protos_includes_known_set() -> None: ("zenoh:dimos/color_image", ("zenoh", "dimos/color_image", None)), ("pzenoh:foo/bar", ("pzenoh", "foo/bar", None)), ( - "lcm:/odom#nav_msgs.Odometry", - ("lcm", "/odom", "nav_msgs.Odometry"), + "lcm:/odom#nav_msgs/msg/Odometry", + ("lcm", "/odom", "nav_msgs/msg/Odometry"), ), ( - "jpeg_lcm:/color_image#sensor_msgs.Image", - ("jpeg_lcm", "/color_image", "sensor_msgs.Image"), + "jpeg_lcm:/color_image#sensor_msgs/msg/Image", + ("jpeg_lcm", "/color_image", "sensor_msgs/msg/Image"), ), ( - "zenoh:dimos/color_image#sensor_msgs.Image", - ("zenoh", "dimos/color_image", "sensor_msgs.Image"), + "zenoh:dimos/color_image#sensor_msgs/msg/Image", + ("zenoh", "dimos/color_image", "sensor_msgs/msg/Image"), ), ], ) @@ -97,7 +102,7 @@ def test_parse_pubsub_uri_rejects_malformed(uri: str) -> None: parse_pubsub_uri(uri) -@pytest.mark.parametrize("proto", ["shm", "pshm", "jpeg_shm"]) +@pytest.mark.parametrize("proto", ["jpeg_shm"]) def test_make_pubsub_transport_rejects_removed_protocols(proto) -> None: with pytest.raises(ValueError, match=f"Unsupported proto '{proto}'"): make_pubsub_transport(f"{proto}:color_image", msg_type=Image) @@ -131,7 +136,7 @@ def test_make_pubsub_transport_zenoh_matches_default_image_channel() -> None: default = make_transport("color_image", Image, g=GlobalConfig(transport="zenoh")) assert isinstance(t, ZenohTransport) - assert t.channel == default.channel == "dimos/color_image/sensor_msgs.Image" + assert t.channel == default.channel == "dimos/color_image/sensor_msgs/msg/Image" def test_make_pubsub_transport_pzenoh_ignores_msg_type() -> None: @@ -150,9 +155,13 @@ def test_make_pubsub_transport_typed_proto_without_msg_type_raises(proto) -> Non @pytest.mark.parametrize("proto", ["lcm", "jpeg_lcm", "zenoh"]) @pytest.mark.parametrize("msg_type", [None, Vector3]) def test_make_pubsub_transport_uri_suffix_resolves_and_overrides_msg_type(proto, msg_type) -> None: - t = make_pubsub_transport(f"{proto}:color_image#sensor_msgs.Image", msg_type=msg_type) + t = make_pubsub_transport(f"{proto}:color_image#sensor_msgs/msg/Image", msg_type=msg_type) - assert t.topic.lcm_type is Image + assert t.topic.msg_type.__msgtype__ == ( + "sensor_msgs/msg/CompressedImage" if proto == "jpeg_lcm" else "sensor_msgs/msg/Image" + ) + if proto == "jpeg_lcm": + assert t.image_type is Image @pytest.mark.parametrize("proto", ["lcm", "zenoh"]) @@ -185,11 +194,10 @@ def test_subscribe_pubsub_uri_returns_transport_and_unsubscribe(mocker) -> None: [("zenoh:dimos/registry_image", Image), ("pzenoh:dimos/registry_image", None)], ) def test_zenoh_uri_receives_image_from_standard_publisher(uri, msg_type, session_pool) -> None: - message = Image( - data=np.array([[[10, 20, 30], [40, 50, 60]]], dtype=np.uint8), - format=ImageFormat.RGB, - ts=12.5, - frame_id="camera", + message = image_from_array( + np.array([[[10, 20, 30], [40, 50, 60]]], dtype=np.uint8), + encoding="rgb8", + header=Header(stamp=time_from_seconds(12.5), frame_id="camera"), ) received: list[Image] = [] arrived = threading.Event() @@ -209,7 +217,20 @@ def on_image(image: Image) -> None: assert arrived.wait(timeout=5.0) assert len(received) == 1 - np.testing.assert_array_equal(received[0].data, message.data) - assert received[0].format == ImageFormat.RGB - assert received[0].ts == 12.5 - assert received[0].frame_id == "camera" + np.testing.assert_array_equal(image_view(received[0]), image_view(message)) + assert received[0].encoding == "rgb8" + assert to_seconds(received[0].header.stamp) == 12.5 + assert received[0].header.frame_id == "camera" + + +@pytest.mark.parametrize("msg_type", [None, Image]) +def test_make_pubsub_transport_shm(msg_type) -> None: + transport = make_pubsub_transport("shm:registry_image", msg_type=msg_type) + assert isinstance(transport, SHMTransport) + assert transport.channel == ( + "registry_image#sensor_msgs/msg/Image" if msg_type else "registry_image" + ) + + +def test_make_pubsub_transport_pshm() -> None: + assert isinstance(make_pubsub_transport("pshm:registry_image"), pSHMTransport) diff --git a/dimos/protocol/service/spec.py b/dimos/protocol/service/spec.py index 8121fd3da94..58079fbb922 100644 --- a/dimos/protocol/service/spec.py +++ b/dimos/protocol/service/spec.py @@ -15,6 +15,7 @@ from abc import ABC from typing import Any, ClassVar, get_type_hints +import numpy as np from pydantic import BaseModel from typing_extensions import Self @@ -24,6 +25,15 @@ class BaseConfig(BaseModel): model_config = {"arbitrary_types_allowed": True, "extra": "forbid"} + @classmethod + def __pydantic_init_subclass__(cls, **kwargs: Any) -> None: + super().__pydantic_init_subclass__(**kwargs) + if not cls.__pydantic_complete__: + # Native rosbags dataclasses keep these annotation names under + # TYPE_CHECKING. Resolve them through Pydantic's public API without + # altering their classes or generated module globals. + cls.model_rebuild(raise_errors=False, _types_namespace={"np": np, "ClassVar": ClassVar}) + class SessionConfig(BaseConfig): """One transport's session settings, as opposed to a module's own config.""" diff --git a/dimos/protocol/service/test_lcmservice.py b/dimos/protocol/service/test_lcmservice.py index 78f1145b0f8..f9fec74f242 100644 --- a/dimos/protocol/service/test_lcmservice.py +++ b/dimos/protocol/service/test_lcmservice.py @@ -340,7 +340,7 @@ def test_str_without_lcm_type(self) -> None: def test_str_with_lcm_type(self) -> None: mock_type = MagicMock() - mock_type.msg_name = "TestMessage" + mock_type.__msgtype__ = "TestMessage" topic = Topic(topic="my_topic", msg_type=mock_type) assert str(topic) == "my_topic#TestMessage" diff --git a/dimos/protocol/tf/static_tf_publisher.py b/dimos/protocol/tf/static_tf_publisher.py index 587fae4349b..74083cd669e 100644 --- a/dimos/protocol/tf/static_tf_publisher.py +++ b/dimos/protocol/tf/static_tf_publisher.py @@ -27,6 +27,7 @@ import asyncio import time +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Transform, TransformStamped, Vector3 from dimos_generated.std_msgs.msg import Header from dimos_generated.tf2_msgs.msg import TFMessage @@ -58,7 +59,7 @@ def frames_to_edge_transforms(frames: list[FrameSpec]) -> list[TransformStamped] continue transforms.append( TransformStamped( - header=Header(frame_id=parent), + header=Header(frame_id=parent, stamp=Time(sec=0, nanosec=0)), child_frame_id=name, transform=Transform( translation=Vector3(x=translation[0], y=translation[1], z=translation[2]), diff --git a/dimos/protocol/tf/test_static_tf_publisher.py b/dimos/protocol/tf/test_static_tf_publisher.py index d274f6bcc54..90316e4717e 100644 --- a/dimos/protocol/tf/test_static_tf_publisher.py +++ b/dimos/protocol/tf/test_static_tf_publisher.py @@ -16,6 +16,7 @@ import math +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.tf2_msgs.msg import TFMessage import pytest @@ -35,7 +36,7 @@ def test_frame_tree_skips_root_and_composes_rotated_mounts(): ("world", "body"), ("body", "camera"), ] - decoded = TFMessage.decode(TFMessage(transforms=edges).encode()) + decoded = cdr_decode(cdr_encode(TFMessage(transforms=edges)), TFMessage) combined = compose_transforms(*decoded.transforms) translation = combined.transform.translation assert (translation.x, translation.y, translation.z) == pytest.approx((1, 2, 3)) diff --git a/dimos/protocol/tf/test_tf.py b/dimos/protocol/tf/test_tf.py index cfc395d639d..cb672e3b1fd 100644 --- a/dimos/protocol/tf/test_tf.py +++ b/dimos/protocol/tf/test_tf.py @@ -141,7 +141,9 @@ def test_tf_main() -> None: random_t = TransformStamped( header=Header(frame_id="sensor", stamp=time_from_seconds(time.time())), child_frame_id="random_object", - transform=Transform(translation=Vector3(x=1, y=0, z=0)), + transform=Transform( + translation=Vector3(x=1, y=0, z=0), rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), ) # we broadcast our object location @@ -219,7 +221,10 @@ def test_get(self) -> None: transform = TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(base_time + i * 0.5)), child_frame_id="robot", - transform=Transform(translation=Vector3(x=float(i), y=0.0, z=0.0)), + transform=Transform( + translation=Vector3(x=float(i), y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) buffer.add(transform) @@ -245,7 +250,10 @@ def test_buffer_pruning(self) -> None: old_transform = TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(old_time)), child_frame_id="robot", - transform=Transform(translation=Vector3(x=1.0, y=0.0, z=0.0)), + transform=Transform( + translation=Vector3(x=1.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) buffer.add(old_transform) @@ -253,7 +261,10 @@ def test_buffer_pruning(self) -> None: recent_transform = TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(time.time())), child_frame_id="robot", - transform=Transform(translation=Vector3(x=2.0, y=0.0, z=0.0)), + transform=Transform( + translation=Vector3(x=2.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) buffer.add(recent_transform) @@ -272,13 +283,19 @@ def test_multiple_frame_pairs(self) -> None: transform1 = TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(time.time())), child_frame_id="robot1", - transform=Transform(translation=Vector3(x=1.0, y=0.0, z=0.0)), + transform=Transform( + translation=Vector3(x=1.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) transform2 = TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(time.time())), child_frame_id="robot2", - transform=Transform(translation=Vector3(x=2.0, y=0.0, z=0.0)), + transform=Transform( + translation=Vector3(x=2.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) ttbuffer.receive_transform(transform1, transform2) @@ -295,13 +312,19 @@ def test_graph(self) -> None: transform1 = TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(time.time())), child_frame_id="robot1", - transform=Transform(translation=Vector3(x=1.0, y=0.0, z=0.0)), + transform=Transform( + translation=Vector3(x=1.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) transform2 = TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(time.time())), child_frame_id="robot2", - transform=Transform(translation=Vector3(x=2.0, y=0.0, z=0.0)), + transform=Transform( + translation=Vector3(x=2.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) ttbuffer.receive_transform(transform1, transform2) @@ -318,7 +341,10 @@ def publish_after_delay() -> None: TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(base_time)), child_frame_id="robot", - transform=Transform(translation=Vector3(x=1.0, y=0.0, z=0.0)), + transform=Transform( + translation=Vector3(x=1.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) ) @@ -351,7 +377,10 @@ def test_forward_tolerance_fast_path_when_already_available(self) -> None: TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(base_time)), child_frame_id="robot", - transform=Transform(translation=Vector3(x=2.0, y=0.0, z=0.0)), + transform=Transform( + translation=Vector3(x=2.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) ) @@ -371,7 +400,10 @@ def test_forward_tolerance_wakes_on_chain_completion(self) -> None: TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(base_time)), child_frame_id="robot", - transform=Transform(translation=Vector3(x=1.0, y=0.0, z=0.0)), + transform=Transform( + translation=Vector3(x=1.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) ) @@ -381,7 +413,10 @@ def publish_after_delay() -> None: TransformStamped( header=Header(frame_id="robot", stamp=time_from_seconds(base_time)), child_frame_id="sensor", - transform=Transform(translation=Vector3(x=0.0, y=2.0, z=0.0)), + transform=Transform( + translation=Vector3(x=0.0, y=2.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) ) @@ -405,7 +440,10 @@ def test_get_transform_search_direct(self) -> None: transform = TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(base_time)), child_frame_id="robot", - transform=Transform(translation=Vector3(x=1.0, y=0.0, z=0.0)), + transform=Transform( + translation=Vector3(x=1.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) ttbuffer.receive_transform(transform) @@ -423,12 +461,18 @@ def test_get_transform_search_chain(self) -> None: transform1 = TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(base_time)), child_frame_id="robot", - transform=Transform(translation=Vector3(x=1.0, y=0.0, z=0.0)), + transform=Transform( + translation=Vector3(x=1.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) transform2 = TransformStamped( header=Header(frame_id="robot", stamp=time_from_seconds(base_time)), child_frame_id="sensor", - transform=Transform(translation=Vector3(x=0.0, y=2.0, z=0.0)), + transform=Transform( + translation=Vector3(x=0.0, y=2.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) ttbuffer.receive_transform(transform1, transform2) @@ -450,27 +494,42 @@ def test_get_transform_search_complex_chain(self) -> None: TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(base_time)), child_frame_id="base", - transform=Transform(translation=Vector3(x=1.0, y=0.0, z=0.0)), + transform=Transform( + translation=Vector3(x=1.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), TransformStamped( header=Header(frame_id="base", stamp=time_from_seconds(base_time)), child_frame_id="arm", - transform=Transform(translation=Vector3(x=0.0, y=1.0, z=0.0)), + transform=Transform( + translation=Vector3(x=0.0, y=1.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), TransformStamped( header=Header(frame_id="arm", stamp=time_from_seconds(base_time)), child_frame_id="hand", - transform=Transform(translation=Vector3(x=0.0, y=0.0, z=1.0)), + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=1.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(base_time)), child_frame_id="robot", - transform=Transform(translation=Vector3(x=2.0, y=0.0, z=0.0)), + transform=Transform( + translation=Vector3(x=2.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), TransformStamped( header=Header(frame_id="robot", stamp=time_from_seconds(base_time)), child_frame_id="sensor", - transform=Transform(translation=Vector3(x=0.0, y=2.0, z=0.0)), + transform=Transform( + translation=Vector3(x=0.0, y=2.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), ] @@ -493,12 +552,18 @@ def test_get_transform_search_no_path(self) -> None: transform1 = TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(base_time)), child_frame_id="robot", - transform=Transform(), + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) transform2 = TransformStamped( header=Header(frame_id="base", stamp=time_from_seconds(base_time)), child_frame_id="sensor", - transform=Transform(), + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) ttbuffer.receive_transform(transform1, transform2) @@ -514,12 +579,18 @@ def test_get_transform_search_with_time(self) -> None: old_transform = TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(base_time - 10.0)), child_frame_id="robot", - transform=Transform(translation=Vector3(x=1.0, y=0.0, z=0.0)), + transform=Transform( + translation=Vector3(x=1.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) new_transform = TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(base_time)), child_frame_id="robot", - transform=Transform(translation=Vector3(x=2.0, y=0.0, z=0.0)), + transform=Transform( + translation=Vector3(x=2.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) ttbuffer.receive_transform(old_transform, new_transform) @@ -545,22 +616,34 @@ def test_get_transform_search_shortest_path(self) -> None: TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(base_time)), child_frame_id="A", - transform=Transform(), + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), TransformStamped( header=Header(frame_id="A", stamp=time_from_seconds(base_time)), child_frame_id="B", - transform=Transform(), + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), TransformStamped( header=Header(frame_id="B", stamp=time_from_seconds(base_time)), child_frame_id="target", - transform=Transform(), + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(base_time)), child_frame_id="target", - transform=Transform(), + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), ] @@ -588,7 +671,10 @@ def test_string_representations(self) -> None: transform = TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(base_time + i * 0.1)), child_frame_id="robot", - transform=Transform(translation=Vector3(x=float(i), y=0.0, z=0.0)), + transform=Transform( + translation=Vector3(x=float(i), y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) buffer.add(transform) @@ -603,17 +689,26 @@ def test_string_representations(self) -> None: TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(base_time)), child_frame_id="robot1", - transform=Transform(), + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), TransformStamped( header=Header(frame_id="world", stamp=time_from_seconds(base_time + 0.5)), child_frame_id="robot2", - transform=Transform(), + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), TransformStamped( header=Header(frame_id="robot1", stamp=time_from_seconds(base_time + 1.0)), child_frame_id="sensor", - transform=Transform(), + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), ] @@ -640,7 +735,9 @@ def _t(parent: str, child: str, x: float, ts: float) -> TransformStamped: return TransformStamped( header=Header(frame_id=parent, stamp=time_from_seconds(ts)), child_frame_id=child, - transform=Transform(translation=Vector3(x=x, y=0.0, z=0.0)), + transform=Transform( + translation=Vector3(x=x, y=0.0, z=0.0), rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), ) @@ -855,10 +952,18 @@ def test_buffer_retains_distinct_nanoseconds_and_owns_its_snapshots(): first = TransformStamped( header=Header(frame_id="a", stamp=Time(sec=1700000000, nanosec=123456789)), child_frame_id="b", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) second = TransformStamped( header=Header(frame_id="a", stamp=Time(sec=1700000000, nanosec=123456790)), child_frame_id="b", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) buffer.add(first) buffer.add(second) @@ -879,7 +984,16 @@ def test_buffer_retains_distinct_nanoseconds_and_owns_its_snapshots(): def test_buffer_drops_out_of_order_data_older_than_retained_window(): buffer = TBuffer(buffer_size=1) for sec in (10, 11, 1): - buffer.add(TransformStamped(header=Header(stamp=Time(sec=sec)))) + buffer.add( + TransformStamped( + header=Header(stamp=Time(sec=sec, nanosec=0), frame_id=""), + child_frame_id="", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ) assert len(buffer) == 2 first = buffer.first() assert first is not None @@ -888,7 +1002,16 @@ def test_buffer_drops_out_of_order_data_older_than_retained_window(): def test_buffer_preserves_negative_stamp(): buffer = TBuffer() - buffer.add(TransformStamped(header=Header(stamp=Time(sec=-1, nanosec=500000000)))) + buffer.add( + TransformStamped( + header=Header(stamp=Time(sec=-1, nanosec=500000000), frame_id=""), + child_frame_id="", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ) result = buffer.get(time_point=-0.5, time_tolerance=0) assert result is not None assert result.header.stamp == Time(sec=-1, nanosec=500000000) diff --git a/dimos/protocol/tf/tf.py b/dimos/protocol/tf/tf.py index 87afac6b4cc..c6f4df2264d 100644 --- a/dimos/protocol/tf/tf.py +++ b/dimos/protocol/tf/tf.py @@ -15,7 +15,7 @@ # limitations under the License. from collections import deque -from copy import copy +from copy import deepcopy from functools import reduce import math import subprocess @@ -23,7 +23,13 @@ import time from typing import TYPE_CHECKING, Any, Protocol, cast, runtime_checkable -from dimos_generated.geometry_msgs.msg import PoseStamped, TransformStamped +from dimos_generated.geometry_msgs.msg import ( + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) from dimos_generated.std_msgs.msg import Header from dimos_generated.tf2_msgs.msg import TFMessage from sortedcontainers import SortedDict # type: ignore[import-untyped] @@ -71,21 +77,17 @@ def __len__(self) -> int: def add(self, transform: TransformStamped) -> None: stamp = to_nanoseconds(transform.header.stamp) - self._entries[stamp] = TransformStamped( - header=transform.header, - child_frame_id=transform.child_frame_id, - transform=transform.transform, - ) + self._entries[stamp] = deepcopy(transform) if math.isfinite(self.buffer_size): oldest = self._entries.peekitem(-1)[0] - round(self.buffer_size * 1_000_000_000) while self._entries and self._entries.peekitem(0)[0] < oldest: self._entries.popitem(0) def first(self) -> TransformStamped | None: - return copy(self._entries.peekitem(0)[1]) if self._entries else None + return deepcopy(self._entries.peekitem(0)[1]) if self._entries else None def last(self) -> TransformStamped | None: - return copy(self._entries.peekitem(-1)[1]) if self._entries else None + return deepcopy(self._entries.peekitem(-1)[1]) if self._entries else None def get( self, time_point: float | None = None, time_tolerance: float = 1.0 @@ -104,7 +106,7 @@ def get( nearest, result = min(candidates, key=lambda item: (abs(item[0] - stamp), -item[0])) if abs(nearest - stamp) / 1_000_000_000 > time_tolerance: return None - return copy(cast("TransformStamped", result)) + return deepcopy(cast("TransformStamped", result)) def __str__(self) -> str: first, last = self.first(), self.last() @@ -171,6 +173,10 @@ def get_transform( stamp=time_from_seconds(time_point if time_point is not None else time.time()), ), child_frame_id=child_frame, + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) # No explicit tolerance means "anything still buffered" — the buffer diff --git a/dimos/robot/deeprobotics/m20/blueprints/m20_kronknav.py b/dimos/robot/deeprobotics/m20/blueprints/m20_kronknav.py index e69e983e3cf..c124efddbf5 100644 --- a/dimos/robot/deeprobotics/m20/blueprints/m20_kronknav.py +++ b/dimos/robot/deeprobotics/m20/blueprints/m20_kronknav.py @@ -16,12 +16,13 @@ from typing import Any +from dimos_generated.foxglove_msgs.msg import CompressedVideo +from dimos_generated.sensor_msgs.msg import CameraInfo + from dimos.core.coordination.blueprints import autoconnect from dimos.core.global_config import global_config from dimos.core.transport import ZenohTransport from dimos.mapping.ray_tracing.module import RayTracingVoxelMap -from dimos.msgs.foxglove_msgs.CompressedVideo import CompressedVideo -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo from dimos.navigation.experimental.dannav.holonomic_tc.module import DanHolonomicTC from dimos.navigation.experimental.dannav.local_planner.module import DanLocalPlanner from dimos.navigation.global_planner.mls_planner.mls_planner_native import MLSPlannerNative @@ -38,6 +39,7 @@ PLANNING_HEIGHT_M, ROTATION_DIAMETER_M, ) +from dimos.visualization.rerun.message_helpers import camera_pinhole, video_archetype from dimos.visualization.vis_module import vis_module VOXEL_SIZE_M = 0.1 @@ -46,23 +48,15 @@ def _render_h265(msg: CompressedVideo) -> Any: - import rerun as rr - - return rr.VideoStream(codec=rr.VideoCodec.H265, sample=msg.data.tobytes()) + return video_archetype(msg) def _render_front_camera_info(msg: CameraInfo) -> Any: - return msg.to_rerun( - image_topic="world/front_camera", - optical_frame="front_camera_optical", - ) + return [("world/front_camera", camera_pinhole(msg))] def _render_rear_camera_info(msg: CameraInfo) -> Any: - return msg.to_rerun( - image_topic="world/rear_camera", - optical_frame="rear_camera_optical", - ) + return [("world/rear_camera", camera_pinhole(msg))] def _static_robot_body(rr: Any) -> list[Any]: diff --git a/dimos/robot/deeprobotics/m20/camera.py b/dimos/robot/deeprobotics/m20/camera.py index c7db9a6d594..4e645c19b6a 100644 --- a/dimos/robot/deeprobotics/m20/camera.py +++ b/dimos/robot/deeprobotics/m20/camera.py @@ -19,11 +19,13 @@ from typing import Any import av +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.foxglove_msgs.msg import CompressedVideo from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 from dimos_generated.sensor_msgs.msg import CameraInfo from dimos_generated.std_msgs.msg import Header from dimos_generated.tf2_msgs.msg import TFMessage +import numpy as np from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.core import rpc @@ -56,7 +58,7 @@ cy=_CAMERA_HEIGHT * 0.5, width=_CAMERA_WIDTH, height=_CAMERA_HEIGHT, - header=Header(frame_id="front_camera_optical"), + header=Header(frame_id="front_camera_optical", stamp=Time(sec=0, nanosec=0)), ) _REAR_CAMERA_INFO = camera_info_from_intrinsics( fx=_CAMERA_FOCAL_LENGTH, @@ -65,7 +67,7 @@ cy=_CAMERA_HEIGHT * 0.5, width=_CAMERA_WIDTH, height=_CAMERA_HEIGHT, - header=Header(frame_id="rear_camera_optical"), + header=Header(frame_id="rear_camera_optical", stamp=Time(sec=0, nanosec=0)), ) @@ -126,7 +128,7 @@ def _publish_camera_metadata(self) -> None: y=_FRONT_CAMERA_XYZ[1], z=_FRONT_CAMERA_XYZ[2], ), - rotation=Quaternion(w=1.0), + rotation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), ), ), TransformStamped( @@ -134,7 +136,9 @@ def _publish_camera_metadata(self) -> None: frame_id="front_camera_link", stamp=time_from_seconds(now) ), child_frame_id="front_camera_optical", - transform=Transform(rotation=_OPTICAL_ROT), + transform=Transform( + rotation=_OPTICAL_ROT, translation=Vector3(x=0.0, y=0.0, z=0.0) + ), ), TransformStamped( header=Header(frame_id="base_link", stamp=time_from_seconds(now)), @@ -145,7 +149,7 @@ def _publish_camera_metadata(self) -> None: y=_REAR_CAMERA_XYZ[1], z=_REAR_CAMERA_XYZ[2], ), - rotation=Quaternion(z=1.0), + rotation=Quaternion(z=1.0, x=0.0, y=0.0, w=1.0), ), ), TransformStamped( @@ -153,7 +157,9 @@ def _publish_camera_metadata(self) -> None: frame_id="rear_camera_link", stamp=time_from_seconds(now) ), child_frame_id="rear_camera_optical", - transform=Transform(rotation=_OPTICAL_ROT), + transform=Transform( + rotation=_OPTICAL_ROT, translation=Vector3(x=0.0, y=0.0, z=0.0) + ), ), ] ) @@ -188,7 +194,7 @@ def _relay(self, url: str, output: Out[CompressedVideo], camera_name: str) -> No continue output.publish( CompressedVideo( - data=bytes(packet), + data=np.frombuffer(bytes(packet), dtype=np.uint8), format="h265", frame_id="", timestamp=time_from_seconds(time.time()), diff --git a/dimos/robot/deeprobotics/m20/connection.py b/dimos/robot/deeprobotics/m20/connection.py index 2e9c076b1c6..066ad734eda 100644 --- a/dimos/robot/deeprobotics/m20/connection.py +++ b/dimos/robot/deeprobotics/m20/connection.py @@ -23,7 +23,7 @@ from threading import Condition, RLock from typing import Any, Literal -from dimos_generated.geometry_msgs.msg import PoseStamped, Twist +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist, Vector3 from dimos_generated.nav_msgs.msg import Odometry from dimos_generated.std_msgs.msg import Bool, Int32, UInt32 from pydantic import Field @@ -136,7 +136,9 @@ def move(self, twist: Twist, duration: float = 0.0) -> bool: @rpc def stop_movement(self) -> None: """Publish an immediate zero velocity without disabling future commands.""" - self.cmd_vel.publish(Twist()) + self.cmd_vel.publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) @rpc def standup(self) -> bool: diff --git a/dimos/robot/deeprobotics/m20/test_camera_messages.py b/dimos/robot/deeprobotics/m20/test_camera_messages.py index 89a9a8694c3..dfec6a4ceb9 100644 --- a/dimos/robot/deeprobotics/m20/test_camera_messages.py +++ b/dimos/robot/deeprobotics/m20/test_camera_messages.py @@ -15,6 +15,7 @@ from threading import Event from types import SimpleNamespace +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.sensor_msgs.msg import CameraInfo from dimos_generated.tf2_msgs.msg import TFMessage @@ -36,7 +37,7 @@ def test_camera_metadata_uses_generated_wire_values_without_connecting_streams() rear_camera_info=SimpleNamespace(publish=rear.append), ) M20CameraRelay._publish_camera_metadata(proxy) - tf = TFMessage.decode(transforms[0].encode()) + tf = cdr_decode(cdr_encode(transforms[0]), TFMessage) assert len(tf.transforms) == 4 assert [t.child_frame_id for t in tf.transforms] == [ "front_camera_link", @@ -44,7 +45,10 @@ def test_camera_metadata_uses_generated_wire_values_without_connecting_streams() "rear_camera_link", "rear_camera_optical", ] - first, second = CameraInfo.decode(front[0].encode()), CameraInfo.decode(rear[0].encode()) + first, second = ( + cdr_decode(cdr_encode(front[0]), CameraInfo), + cdr_decode(cdr_encode(rear[0]), CameraInfo), + ) assert first.header.frame_id == "front_camera_optical" assert second.header.frame_id == "rear_camera_optical" assert first.header.stamp == second.header.stamp == tf.transforms[0].header.stamp diff --git a/dimos/robot/deeprobotics/m20/test_connection_messages.py b/dimos/robot/deeprobotics/m20/test_connection_messages.py index 1c60f925b38..1586fb534d1 100644 --- a/dimos/robot/deeprobotics/m20/test_connection_messages.py +++ b/dimos/robot/deeprobotics/m20/test_connection_messages.py @@ -15,6 +15,7 @@ # Copyright 2026 Dimensional Inc. # SPDX-License-Identifier: Apache-2.0 +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, @@ -22,10 +23,12 @@ PoseWithCovariance, Quaternion, Twist, + TwistWithCovariance, Vector3, ) from dimos_generated.nav_msgs.msg import Odometry from dimos_generated.std_msgs.msg import Bool, Header +import numpy as np from dimos.robot.deeprobotics.m20.connection import M20Connection @@ -35,19 +38,34 @@ def test_generated_forwarding_preserves_readiness_pose_clock_and_zero_stop(monke velocities, poses = [], [] monkeypatch.setattr(module.cmd_vel, "publish", velocities.append) monkeypatch.setattr(module.odom, "publish", poses.append) - command = Twist(linear=Vector3(x=0.1), angular=Vector3(z=0.2)) + command = Twist(linear=Vector3(x=0.1, y=0.0, z=0.0), angular=Vector3(z=0.2, x=0.0, y=0.0)) try: assert module.move(command) is False module._on_command_ready(Bool(data=True)) assert module.move(command) is True source = Odometry( header=Header(frame_id="map", stamp=Time(sec=1, nanosec=123456789)), - pose=PoseWithCovariance(pose=Pose(position=Point(x=2), orientation=Quaternion(w=1))), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=2, y=0.0, z=0.0), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), + child_frame_id="", + twist=TwistWithCovariance( + twist=Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ), + covariance=np.zeros(36, dtype=np.float64), + ), ) module._on_odometry(source) assert poses[0].header == source.header and poses[0].pose == source.pose.pose module.stop_movement() - assert velocities[-1] == Twist() - assert Twist.decode(velocities[0].encode()) == command + assert velocities[-1] == Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) + assert cdr_decode(cdr_encode(velocities[0]), Twist) == command finally: module.stop() diff --git a/dimos/robot/diy/alfred/config.py b/dimos/robot/diy/alfred/config.py index 09c6295cd77..0a683a15a68 100644 --- a/dimos/robot/diy/alfred/config.py +++ b/dimos/robot/diy/alfred/config.py @@ -18,7 +18,7 @@ from pathlib import Path from typing import Any -from dimos_generated.geometry_msgs.msg import Point, Pose +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion DEFAULT_ADDRESS = "172.6.2.20:11323" ALFRED_URDF = Path(__file__).resolve().parent / "alfred.urdf" @@ -50,6 +50,9 @@ class AlfredConfig: d455_serial="260922302422", internal_odom_offsets={ # Mid-360 lidar: a bit forward, and a bit to the right of base center, above ground. - "mid360_link": Pose(position=Point(x=0.20, y=-0.20, z=0.30)), + "mid360_link": Pose( + position=Point(x=0.20, y=-0.20, z=0.30), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), }, ) diff --git a/dimos/robot/diy/alfred/effector_high_level.py b/dimos/robot/diy/alfred/effector_high_level.py index 07619405d43..2e553720045 100644 --- a/dimos/robot/diy/alfred/effector_high_level.py +++ b/dimos/robot/diy/alfred/effector_high_level.py @@ -178,7 +178,7 @@ async def move_velocity( self, x: float, y: float = 0.0, yaw: float = 0.0, duration: float = 0.0 ) -> str: """Move the Alfred at the given velocity for ``duration`` seconds.""" - twist = Twist(linear=Vector3(x=x, y=y), angular=Vector3(z=yaw)) + twist = Twist(linear=Vector3(x=x, y=y, z=0.0), angular=Vector3(z=yaw, x=0.0, y=0.0)) await self.move(twist, duration=duration) return f"Started moving with velocity=({x}, {y}, {yaw}) for {duration} seconds" @@ -213,7 +213,9 @@ async def _poll_wheel_odometry(self, client: portal.Client) -> None: # Odometry comes back in the inverted-Y frame move() sends into. y, yaw = -y, -yaw # The controller reports no velocity, so difference consecutive poses. - twist = Twist() + twist = Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) if previous is not None: last_ts, last_x, last_y, last_yaw = previous dt = ts - last_ts @@ -225,8 +227,8 @@ async def _poll_wheel_odometry(self, client: portal.Client) -> None: forward = math.cos(heading) * dx + math.sin(heading) * dy left = -math.sin(heading) * dx + math.cos(heading) * dy twist = Twist( - linear=Vector3(x=forward / dt, y=left / dt), - angular=Vector3(z=turn / dt), + linear=Vector3(x=forward / dt, y=left / dt, z=0.0), + angular=Vector3(z=turn / dt, x=0.0, y=0.0), ) previous = (ts, x, y, yaw) @@ -238,11 +240,14 @@ async def _poll_wheel_odometry(self, client: portal.Client) -> None: child_frame_id=self.config.base_frame_id, pose=PoseWithCovariance( pose=Pose( - position=Point(x=x, y=y), + position=Point(x=x, y=y, z=0.0), orientation=quaternion_from_euler(0.0, 0.0, yaw), - ) + ), + covariance=np.zeros(36, dtype=np.float64), + ), + twist=TwistWithCovariance( + twist=twist, covariance=np.zeros(36, dtype=np.float64) ), - twist=TwistWithCovariance(twist=twist), ) ) except asyncio.CancelledError: diff --git a/dimos/robot/diy/alfred/mount_tf.py b/dimos/robot/diy/alfred/mount_tf.py index 47962f957c0..75dfd26fe9b 100644 --- a/dimos/robot/diy/alfred/mount_tf.py +++ b/dimos/robot/diy/alfred/mount_tf.py @@ -25,6 +25,7 @@ from xml.etree import ElementTree +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Transform, TransformStamped, Vector3 from dimos_generated.std_msgs.msg import Header @@ -54,7 +55,7 @@ def mount_transforms() -> list[TransformStamped]: rpy = [float(value) for value in origin.attrib["rpy"].split()] transforms.append( TransformStamped( - header=Header(frame_id=parent.attrib["link"]), + header=Header(frame_id=parent.attrib["link"], stamp=Time(sec=0, nanosec=0)), child_frame_id=child_link, transform=Transform( translation=Vector3(x=translation[0], y=translation[1], z=translation[2]), diff --git a/dimos/robot/diy/alfred/test_generated_effector.py b/dimos/robot/diy/alfred/test_generated_effector.py index 2bc95b58835..d81168e2784 100644 --- a/dimos/robot/diy/alfred/test_generated_effector.py +++ b/dimos/robot/diy/alfred/test_generated_effector.py @@ -15,6 +15,7 @@ from types import SimpleNamespace from unittest.mock import MagicMock +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.nav_msgs.msg import Odometry import pytest @@ -48,7 +49,7 @@ def get_odometry(_request: object) -> SimpleNamespace: received: list[Odometry] = [] def publish(value: Odometry) -> None: - received.append(Odometry.decode(value.encode())) + received.append(cdr_decode(cdr_encode(value), Odometry)) if len(received) == 2: module._odometry_stop.set() diff --git a/dimos/robot/diy/alfred/test_mount_tf.py b/dimos/robot/diy/alfred/test_mount_tf.py index 758c3a1a37e..609f8f9a521 100644 --- a/dimos/robot/diy/alfred/test_mount_tf.py +++ b/dimos/robot/diy/alfred/test_mount_tf.py @@ -14,6 +14,7 @@ """URDF mount extraction emits generated transforms and excludes driver-owned frames.""" +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.tf2_msgs.msg import TFMessage from dimos.robot.diy.alfred import mount_tf @@ -34,7 +35,7 @@ def test_mounts_preserve_parent_and_translation_and_skip_imager_frames(tmp_path, monkeypatch.setattr(mount_tf, "ALFRED_URDF", urdf) edges = mount_tf.mount_transforms() - decoded = TFMessage.decode(TFMessage(transforms=edges).encode()) + decoded = cdr_decode(cdr_encode(TFMessage(transforms=edges)), TFMessage) assert len(decoded.transforms) == 1 edge = decoded.transforms[0] diff --git a/dimos/robot/drone/camera_module.py b/dimos/robot/drone/camera_module.py index 453a1e1ad02..2d181f5e5f1 100644 --- a/dimos/robot/drone/camera_module.py +++ b/dimos/robot/drone/camera_module.py @@ -22,9 +22,10 @@ from typing import Any from dimos_generated.geometry_msgs.msg import PoseStamped -from dimos_generated.sensor_msgs.msg import CameraInfo, Image +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, RegionOfInterest from dimos_generated.std_msgs.msg import Header from dimos_generated.tf2_msgs.msg import TFMessage +import numpy as np from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.core import rpc @@ -183,12 +184,13 @@ def _publish_camera_info(self, header: Header, shape: tuple[int, ...]) -> None: height=height, width=width, distortion_model="plumb_bob", - d=D, - k=K, - r=R, - p=P, + d=np.asarray(D, dtype=np.float64), + k=np.asarray(K, dtype=np.float64), + r=np.asarray(R, dtype=np.float64), + p=np.asarray(P, dtype=np.float64), binning_x=0, binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), ) self.camera_info.publish(msg) diff --git a/dimos/robot/drone/connection_module.py b/dimos/robot/drone/connection_module.py index e2074d287e4..25a1161f9cf 100644 --- a/dimos/robot/drone/connection_module.py +++ b/dimos/robot/drone/connection_module.py @@ -22,8 +22,10 @@ from typing import Any from dimos_generated.geometry_msgs.msg import ( + Point, Pose, PoseStamped, + Quaternion, Transform, TransformStamped, Twist, @@ -189,7 +191,10 @@ def _publish_tf(self, msg: PoseStamped) -> None: camera_link = TransformStamped( header=Header(stamp=time_from_seconds(time.time()), frame_id="base_link"), child_frame_id="camera_link", - transform=Transform(translation=Vector3(x=0.1, z=-0.05)), + transform=Transform( + translation=Vector3(x=0.1, z=-0.05, y=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) self.tf.publish(TFMessage(transforms=[base_link, camera_link])) @@ -225,7 +230,10 @@ def _telemetry_loop(self) -> None: # Publish default pose default_pose = PoseStamped( header=Header(stamp=time_from_seconds(time.time()), frame_id="world"), - pose=Pose(), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ) self._publish_tf(default_pose) logger.debug("Publishing default odometry") diff --git a/dimos/robot/drone/dji_video_stream.py b/dimos/robot/drone/dji_video_stream.py index 54229da7480..e169ea730a2 100644 --- a/dimos/robot/drone/dji_video_stream.py +++ b/dimos/robot/drone/dji_video_stream.py @@ -141,7 +141,9 @@ def _capture_loop(self) -> None: # Create Image message (RGB format - matches GStreamer pipeline output) img_msg = image_from_array( - frame, encoding="rgb8", header=Header(stamp=time_from_seconds(time.time())) + frame, + encoding="rgb8", + header=Header(stamp=time_from_seconds(time.time()), frame_id=""), ) # Publish diff --git a/dimos/robot/drone/drone_tracking_module.py b/dimos/robot/drone/drone_tracking_module.py index e22be472c9d..dfe0225fa11 100644 --- a/dimos/robot/drone/drone_tracking_module.py +++ b/dimos/robot/drone/drone_tracking_module.py @@ -270,9 +270,11 @@ def _visual_servoing_loop(self, tracker: Any, duration: float) -> None: # Publish velocity command via LCM if self.cmd_vel.transport: - twist = Twist() - twist.linear = Vector3(x=vx, y=vy) - twist.angular = Vector3() # No rotation for now + twist = Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) + twist.linear = Vector3(x=vx, y=vy, z=0.0) + twist.angular = Vector3(x=0.0, y=0.0, z=0.0) # No rotation for now self.cmd_vel.publish(twist) # Publish visualization if transport is set @@ -283,7 +285,7 @@ def _visual_servoing_loop(self, tracker: Any, duration: float) -> None: overlay_msg = image_from_array( overlay, encoding="bgr8", - header=Header(stamp=time_from_seconds(time.time())), + header=Header(stamp=time_from_seconds(time.time()), frame_id=""), ) self.tracking_overlay.publish(overlay_msg) @@ -308,9 +310,11 @@ def _visual_servoing_loop(self, tracker: Any, duration: float) -> None: finally: # Stop movement by publishing zero velocity if self.cmd_vel.transport: - stop_twist = Twist() - stop_twist.linear = Vector3() - stop_twist.angular = Vector3() + stop_twist = Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) + stop_twist.linear = Vector3(x=0.0, y=0.0, z=0.0) + stop_twist.angular = Vector3(x=0.0, y=0.0, z=0.0) self.cmd_vel.publish(stop_twist) self._tracking_active = False logger.info(f"Visual servoing loop ended after {frame_count} frames") @@ -382,9 +386,11 @@ def _stop_tracking(self) -> None: # Send stop command via LCM if self.cmd_vel.transport: - stop_twist = Twist() - stop_twist.linear = Vector3() - stop_twist.angular = Vector3() + stop_twist = Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) + stop_twist.linear = Vector3(x=0.0, y=0.0, z=0.0) + stop_twist.angular = Vector3(x=0.0, y=0.0, z=0.0) self.cmd_vel.publish(stop_twist) self._publish_status({"status": "stopped", "object": self._current_object}) diff --git a/dimos/robot/drone/test_drone.py b/dimos/robot/drone/test_drone.py index ce30006ea4f..ff562dee26f 100644 --- a/dimos/robot/drone/test_drone.py +++ b/dimos/robot/drone/test_drone.py @@ -23,6 +23,7 @@ import unittest from unittest.mock import MagicMock, patch +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import PoseStamped, Vector3 from dimos_generated.sensor_msgs.msg import Image from dimos_generated.std_msgs.msg import Header @@ -508,7 +509,9 @@ def test_video_stream_replay(self, mock_get_data, mock_replay) -> None: def subscribe_side_effect(callback) -> None: for frame in video_frames: - img = image_from_array(frame, encoding="bgr8", header=Header()) + img = image_from_array( + frame, encoding="bgr8", header=Header(stamp=Time(sec=0, nanosec=0), frame_id="") + ) callback(img) received_frames.append(img) @@ -520,7 +523,9 @@ def subscribe_side_effect(callback) -> None: def piped_subscribe(callback): for frame in video_frames: - img = image_from_array(frame, encoding="rgb8", header=Header()) # After format fix + img = image_from_array( + frame, encoding="rgb8", header=Header(stamp=Time(sec=0, nanosec=0), frame_id="") + ) # After format fix callback(img) piped_captured.append(img) diff --git a/dimos/robot/galaxea/r1pro/blueprints/basic/r1pro_coordinator.py b/dimos/robot/galaxea/r1pro/blueprints/basic/r1pro_coordinator.py index f9f60863659..713295be7e4 100644 --- a/dimos/robot/galaxea/r1pro/blueprints/basic/r1pro_coordinator.py +++ b/dimos/robot/galaxea/r1pro/blueprints/basic/r1pro_coordinator.py @@ -28,6 +28,17 @@ from collections.abc import Sequence from typing import Any +from dimos_generated.dimos_msgs.msg import MotorCommandArray +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import ( + CameraInfo, + CompressedImage, + Image, + Imu, + JointState, +) + from dimos.control.components import HardwareComponent, HardwareType, make_twist_base_joints from dimos.control.coordinator import ControlCoordinator, TaskConfig from dimos.control.tasks.trajectory_task.trajectory_task import joint_trajectory_task @@ -36,15 +47,6 @@ from dimos.core.transport import ZenohTransport from dimos.hardware.sensors.lidar.pointlio.module import PointLio from dimos.hardware.sensors.lidar.pointlio.pointlio_blueprints import mid360_for_pointlio -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.CompressedImage import CompressedImage -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.Imu import Imu -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.sensor_msgs.MotorCommandArray import MotorCommandArray from dimos.protocol.pubsub.impl.zenohpubsub import QOS_LATEST_WINS, Topic as ZenohTopic, Zenoh from dimos.robot.galaxea.r1pro.config import ( R1PRO_CHASSIS_LIDAR_HOST_IP, diff --git a/dimos/robot/galaxea/r1pro/blueprints/manipulation/test_r1pro_planar_preview.py b/dimos/robot/galaxea/r1pro/blueprints/manipulation/test_r1pro_planar_preview.py index a699e65d072..03e0afff527 100644 --- a/dimos/robot/galaxea/r1pro/blueprints/manipulation/test_r1pro_planar_preview.py +++ b/dimos/robot/galaxea/r1pro/blueprints/manipulation/test_r1pro_planar_preview.py @@ -27,9 +27,8 @@ from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser from dimos.manipulation.manipulation_module import ManipulationModule, ManipulationModuleConfig from dimos.manipulation.planning.spec.validation import prepare_robot_model -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint -from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState +from dimos.msgs.time import duration_from_seconds, header_now +from dimos.msgs.trajectory import TrajectoryState from dimos.robot.galaxea.r1pro.blueprints.basic.r1pro_coordinator import r1pro_control from dimos.robot.galaxea.r1pro.blueprints.manipulation.r1pro_planar_preview import ( r1pro_planar_preview, @@ -49,11 +48,22 @@ def _coordinator_config(blueprint) -> ControlCoordinatorConfig: def _forward(distance: float, duration: float) -> JointTrajectory: velocity = [distance / duration, 0.0, 0.0] return JointTrajectory( + header=header_now(), joint_names=["x", "y", "yaw"], points=[ - TrajectoryPoint(time_from_start=0.0, positions=[0.0, 0.0, 0.0], velocities=velocity), - TrajectoryPoint( - time_from_start=duration, positions=[distance, 0.0, 0.0], velocities=velocity + JointTrajectoryPoint( + time_from_start=duration_from_seconds(0.0), + positions=np.array([0.0, 0.0, 0.0], dtype=np.float64), + velocities=np.asarray(velocity, dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + JointTrajectoryPoint( + time_from_start=duration_from_seconds(duration), + positions=np.array([distance, 0.0, 0.0], dtype=np.float64), + velocities=np.asarray(velocity, dtype=np.float64), + accelerations=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ], ) @@ -96,7 +106,7 @@ def test_planar_preview_drives_its_mock_base_from_a_base_trajectory(wait_until) assert accepted.status is TrajectoryExecutionStatus.ACCEPTED wait_until( lambda: coordinator.task_invoke("base_trajectory", "get_status", {}).state - is TrajectoryState.COMPLETED, + == TrajectoryState.COMPLETED, timeout=5.0, message="the base trajectory never completed", ) diff --git a/dimos/robot/galaxea/r1pro/connection.py b/dimos/robot/galaxea/r1pro/connection.py index cb7d063e1f3..6eca96bd1ed 100644 --- a/dimos/robot/galaxea/r1pro/connection.py +++ b/dimos/robot/galaxea/r1pro/connection.py @@ -40,6 +40,23 @@ import time from typing import TYPE_CHECKING, Any +from dimos_generated.dimos_msgs.msg import MotorCommandArray +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + PoseWithCovariance, + Quaternion, + Twist, + TwistStamped, + TwistWithCovariance, + Vector3, +) +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import CompressedImage, Image, Imu, JointState +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage +import numpy as np from pydantic import Field from reactivex.disposable import Disposable @@ -51,16 +68,9 @@ from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out from dimos.hardware.whole_body.spec import VEL_STOP -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.sensor_msgs.CompressedImage import CompressedImage -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.Imu import Imu -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.sensor_msgs.MotorCommandArray import MotorCommandArray -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.geometry import transform_from_pose +from dimos.msgs.time import header_now, time_from_nanoseconds +from dimos.protocol.pubsub.impl.rospubsub_conversion import dimos_to_ros, ros_to_dimos from dimos.robot.galaxea.r1pro.joints import UPPER_BODY_JOINTS, coordinator_name from dimos.utils.logging_config import setup_logger @@ -115,18 +125,9 @@ def _make_qos() -> Any: ) -def _stamp_secs(msg: Any) -> float: - """Header stamp in seconds, falling back to wall clock if the driver left it 0.""" - ts = msg.header.stamp.sec + msg.header.stamp.nanosec * 1e-9 - return ts if ts > 0 else time.time() - - -def _ros_stamp_now() -> Any: - from builtin_interfaces.msg import Time as RosTime - - t = time.time() - sec = int(t) - return RosTime(sec=sec, nanosec=int((t - sec) * 1e9)) +def _stamp_nanoseconds(msg: Any) -> int: + """Preserve the driver's source stamp, including zero, as integer nanoseconds.""" + return int(msg.header.stamp.sec) * 1_000_000_000 + int(msg.header.stamp.nanosec) class R1ProConnectionConfig(ModuleConfig): @@ -206,9 +207,9 @@ def __init__(self, **kwargs: Any) -> None: self._latest_right_dq: list[float] = [0.0] * 7 self._latest_right_eff: list[float] = [0.0] * 7 self._torso_seen = False - self._ts_torso = 0.0 - self._ts_left = 0.0 - self._ts_right = 0.0 + self._ts_torso = 0 + self._ts_left = 0 + self._ts_right = 0 self._left_seen = False self._right_seen = False self._latest_imu_chassis: Imu | None = None @@ -218,7 +219,7 @@ def __init__(self, **kwargs: Any) -> None: self._odom_x = 0.0 self._odom_y = 0.0 self._odom_yaw = 0.0 - self._odom_last_ts: float | None = None + self._odom_last_ns: int | None = None # Sensor decode pipeline. self._sensor_stop = threading.Event() @@ -516,16 +517,14 @@ def sensor_stats(self) -> dict[str, Any]: # Control input handlers def _on_motor_command(self, msg: MotorCommandArray) -> None: - if msg.num_joints != _NUM_MOTORS: - logger.warning(f"Expected {_NUM_MOTORS} motor commands, got {msg.num_joints}; ignoring") + if any(len(values) != _NUM_MOTORS for values in (msg.q, msg.dq, msg.kp, msg.kd, msg.tau)): + logger.warning(f"Expected {_NUM_MOTORS} values in each motor-command array; ignoring") return - from sensor_msgs.msg import JointState as RosJointState - ros = self._ros if ros is None: return # pre-start / post-stop - stamp = _ros_stamp_now() + header = header_now() for topic, sl in ( (self._cmd_torso_topic, _TORSO_SLICE), @@ -534,13 +533,14 @@ def _on_motor_command(self, msg: MotorCommandArray) -> None: ): if topic is None: continue - cmd = RosJointState() - cmd.header.stamp = stamp - cmd.name = [""] - cmd.position = list(msg.q[sl]) - cmd.velocity = self._tracking_velocities(msg.dq[sl]) - cmd.effort = [0.0] - ros.publish(topic, cmd) + cmd = JointState( + header=header, + name=[""], + position=np.asarray(list(msg.q)[sl], dtype=np.float64), + velocity=np.asarray(self._tracking_velocities(list(msg.dq)[sl]), dtype=np.float64), + effort=np.array([0.0], dtype=np.float64), + ) + ros.publish(topic, dimos_to_ros(cmd, topic.ros_type)) def _tracking_velocities(self, dqs: list[float]) -> list[float]: """Map MotorCommand.dq to ROS tracking velocity (0/sentinel → configured).""" @@ -548,18 +548,17 @@ def _tracking_velocities(self, dqs: list[float]) -> list[float]: return [speed if (v == 0.0 or v == VEL_STOP) else float(v) for v in dqs] def _on_cmd_vel(self, msg: Twist) -> None: - from geometry_msgs.msg import TwistStamped - ros = self._ros if ros is None or self._speed_topic is None: return - - cmd = TwistStamped() - cmd.header.stamp = _ros_stamp_now() - cmd.twist.linear.x = msg.linear.x - cmd.twist.linear.y = msg.linear.y - cmd.twist.angular.z = msg.angular.z - ros.publish(self._speed_topic, cmd) + cmd = TwistStamped( + header=header_now(), + twist=Twist( + linear=Vector3(x=msg.linear.x, y=msg.linear.y, z=0.0), + angular=Vector3(z=msg.angular.z, x=0.0, y=0.0), + ), + ) + ros.publish(self._speed_topic, dimos_to_ros(cmd, self._speed_topic.ros_type)) # Control feedback callbacks (3 segments) @@ -568,7 +567,7 @@ def _on_feedback_torso(self, msg: Any, _topic: Any) -> None: self._copy_segment( msg, self._latest_torso_q, self._latest_torso_dq, self._latest_torso_eff ) - self._ts_torso = _stamp_secs(msg) + self._ts_torso = _stamp_nanoseconds(msg) self._torso_seen = True def _on_feedback_left(self, msg: Any, _topic: Any) -> None: @@ -576,7 +575,7 @@ def _on_feedback_left(self, msg: Any, _topic: Any) -> None: self._copy_segment( msg, self._latest_left_q, self._latest_left_dq, self._latest_left_eff ) - self._ts_left = _stamp_secs(msg) + self._ts_left = _stamp_nanoseconds(msg) self._left_seen = True def _on_feedback_right(self, msg: Any, _topic: Any) -> None: @@ -584,7 +583,7 @@ def _on_feedback_right(self, msg: Any, _topic: Any) -> None: self._copy_segment( msg, self._latest_right_q, self._latest_right_dq, self._latest_right_eff ) - self._ts_right = _stamp_secs(msg) + self._ts_right = _stamp_nanoseconds(msg) self._right_seen = True @staticmethod @@ -592,23 +591,23 @@ def _copy_segment( msg: Any, q_dst: list[float], dq_dst: list[float], eff_dst: list[float] ) -> None: n = min(len(msg.position), len(q_dst)) - q_dst[:n] = msg.position[:n] - if msg.velocity: + q_dst[:n] = list(msg.position)[:n] + if len(msg.velocity): nv = min(len(msg.velocity), len(dq_dst)) - dq_dst[:nv] = msg.velocity[:nv] - if msg.effort: + dq_dst[:nv] = list(msg.velocity)[:nv] + if len(msg.effort): ne = min(len(msg.effort), len(eff_dst)) - eff_dst[:ne] = msg.effort[:ne] + eff_dst[:ne] = list(msg.effort)[:ne] # Wheel odometry (integrated from executed chassis speed) def _on_chassis_speed(self, msg: Any, _topic: Any) -> None: - now = msg.header.stamp.sec + msg.header.stamp.nanosec * 1e-9 - if self._odom_last_ts is None: - self._odom_last_ts = now + now = _stamp_nanoseconds(msg) + if self._odom_last_ns is None: + self._odom_last_ns = now return - dt = now - self._odom_last_ts - self._odom_last_ts = now + dt = (now - self._odom_last_ns) / 1_000_000_000 + self._odom_last_ns = now if dt <= 0.0 or dt > 1.0: # Clock jump or first-tick anomaly — keep position, refresh ts. return @@ -620,28 +619,32 @@ def _on_chassis_speed(self, msg: Any, _topic: Any) -> None: self._odom_y += (sy * vx + cy * vy) * dt self._odom_yaw += wz * dt - from dimos.msgs.geometry_msgs.Pose import Pose - from dimos.msgs.geometry_msgs.Quaternion import Quaternion - from dimos.msgs.geometry_msgs.Vector3 import Vector3 - half = self._odom_yaw * 0.5 - position = Vector3(self._odom_x, self._odom_y, 0.0) - orientation = Quaternion(0.0, 0.0, math.sin(half), math.cos(half)) - frame_id = self.config.odom_frame_id + header = Header(stamp=time_from_nanoseconds(now), frame_id=self.config.odom_frame_id) base = self.config.frame_id - pose = PoseStamped(ts=now, frame_id=frame_id, position=position, orientation=orientation) + pose = PoseStamped( + header=header, + pose=Pose( + position=Point(x=self._odom_x, y=self._odom_y, z=0.0), + orientation=Quaternion(z=math.sin(half), w=math.cos(half), x=0.0, y=0.0), + ), + ) self.odometry.publish( Odometry( - ts=now, - frame_id=frame_id, + header=header, child_frame_id=base, - pose=Pose(position, orientation), - twist=Twist(Vector3(vx, vy, 0.0), Vector3(0.0, 0.0, wz)), + pose=PoseWithCovariance(pose=pose.pose, covariance=np.zeros(36, dtype=np.float64)), + twist=TwistWithCovariance( + twist=Twist( + linear=Vector3(x=vx, y=vy, z=0.0), angular=Vector3(z=wz, x=0.0, y=0.0) + ), + covariance=np.zeros(36, dtype=np.float64), + ), ) ) if self.config.publish_odom_tf: self.odom.publish(pose) - self.tf.publish(TFMessage(Transform.from_pose(base, pose))) + self.tf.publish(TFMessage(transforms=[transform_from_pose(pose, child_frame_id=base)])) # Aggregated motor_states publish loop @@ -685,12 +688,11 @@ def _publish_loop(self) -> None: if bootstrapped: self.motor_states.publish( JointState( - ts=ts, - frame_id=frame_id, + header=Header(stamp=time_from_nanoseconds(ts), frame_id=frame_id), name=R1PRO_UPPER_BODY_JOINTS, - position=positions, # type: ignore[arg-type] - velocity=velocities, - effort=efforts, + position=np.asarray(positions, dtype=np.float64), # type: ignore[arg-type] + velocity=np.asarray(velocities, dtype=np.float64), + effort=np.asarray(efforts, dtype=np.float64), ) ) if imu_chassis is not None: @@ -701,7 +703,7 @@ def _publish_loop(self) -> None: next_tick += period sleep_for = next_tick - time.perf_counter() if sleep_for > 0: - time.sleep(sleep_for) + self._stop_event.wait(sleep_for) else: next_tick = time.perf_counter() @@ -724,25 +726,14 @@ def _compressed_image_loop(self, stream: str, q: queue.Queue[Any]) -> None: continue last_pub = t0 try: - stamp = msg.header.stamp - ts = stamp.sec + stamp.nanosec * 1e-9 - out.publish( - CompressedImage( - data=bytes(msg.data), - format="jpeg", - frame_id=msg.header.frame_id or stream, - ts=ts if ts > 0 else time.time(), - ) - ) + out.publish(ros_to_dimos(msg, CompressedImage)) self._record_decode(stream, (time.perf_counter() - t0) * 1e3, ok=True) except Exception: self._record_decode(stream, (time.perf_counter() - t0) * 1e3, ok=False) logger.exception(f"R1Pro {stream} conversion error") def _convert_loop(self, stream: str, q: queue.Queue[Any], dimos_type: type) -> None: - """ros_to_dimos passthrough worker (depth images).""" - from dimos.protocol.pubsub.impl.rospubsub_conversion import ros_to_dimos - + """ros_to_dimos passthrough worker (depth images, lidar).""" out: Out[Any] = getattr(self, stream) while not self._sensor_stop.is_set(): try: @@ -761,8 +752,6 @@ def _convert_loop(self, stream: str, q: queue.Queue[Any], dimos_type: type) -> N def _imu_loop(self, stream: str, q: queue.Queue[Any]) -> None: """Store the latest converted IMU; re-emitted by the publish loop.""" - from dimos.protocol.pubsub.impl.rospubsub_conversion import ros_to_dimos - target_attr = "_latest_imu_chassis" if stream == "imu_chassis" else "_latest_imu_torso" while not self._sensor_stop.is_set(): try: diff --git a/dimos/robot/galaxea/r1pro/head_cameras.py b/dimos/robot/galaxea/r1pro/head_cameras.py index 7efc718b4ee..a7b4c89b829 100644 --- a/dimos/robot/galaxea/r1pro/head_cameras.py +++ b/dimos/robot/galaxea/r1pro/head_cameras.py @@ -32,15 +32,18 @@ from __future__ import annotations import threading -import time from typing import Any +from dimos_generated.sensor_msgs.msg import CameraInfo +import numpy as np + from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import Out from dimos.hardware.sensors.camera.utils.fsync_trigger import trigger_gmsl_cameras from dimos.hardware.sensors.camera.v4l2.module import V4L2Camera, V4L2CameraConfig -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo +from dimos.msgs.camera_info import camera_info_from_intrinsics +from dimos.msgs.time import header_now from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -116,16 +119,16 @@ def eye(side: str, frame_id: str) -> CameraInfo: k = storage.getNode(f"K_{side}").mat().ravel().tolist() # 14 coefficients; past the first eight (rational polynomial) they are the unused tilt terms. d = storage.getNode(f"D_{side}").mat().ravel().tolist()[:8] - return CameraInfo( - height=height, - width=width, - distortion_model="rational_polynomial", - D=d, - K=k, - R=[1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0], - P=[k[0], k[1], k[2], 0.0, k[3], k[4], k[5], 0.0, k[6], k[7], k[8], 0.0], - frame_id=frame_id, + info = camera_info_from_intrinsics( + k[0], k[4], k[2], k[5], width, height, header=header_now(frame_id) + ) + info.distortion_model = "rational_polynomial" + info.d = np.asarray(d, dtype=np.float64) + info.k = np.asarray(k, dtype=np.float64) + info.p = np.asarray( + [k[0], k[1], k[2], 0.0, k[3], k[4], k[5], 0.0, k[6], k[7], k[8], 0.0], dtype=np.float64 ) + return info return eye("left", HEAD_LEFT_FRAME), eye("right", HEAD_RIGHT_FRAME) @@ -180,6 +183,6 @@ def _run(self) -> None: return while not self._stop.is_set(): for out, info in ((self.left_info, left), (self.right_info, right)): - info.ts = time.time() + info.header = header_now(info.header.frame_id) out.publish(info) self._stop.wait(1.0 / self.config.publish_hz) diff --git a/dimos/robot/galaxea/r1pro/lio.py b/dimos/robot/galaxea/r1pro/lio.py index 1aef950ee6a..652e816ac92 100644 --- a/dimos/robot/galaxea/r1pro/lio.py +++ b/dimos/robot/galaxea/r1pro/lio.py @@ -18,15 +18,26 @@ from functools import cache +from dimos_generated.geometry_msgs.msg import ( + PoseStamped, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.std_msgs.msg import Header from reactivex.disposable import Disposable from dimos.core.core import rpc from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.msgs.geometry import ( + compose_transforms, + inverse_transform, + pose_from_transform, + quaternion_from_euler, + transform_from_odometry, +) +from dimos.msgs.time import time_from_seconds from dimos.protocol.tf.static_tf_publisher import StaticTfPublisher, StaticTfPublisherConfig from dimos.robot.assets.model import LoadedRobotModel from dimos.robot.galaxea.r1pro.config import R1PRO_MODEL @@ -46,41 +57,34 @@ def lio_model(lidar_frame: str, chassis_lidar_frame: str) -> LoadedRobotModel: return R1PRO_MODEL.with_fixed_frame(lidar_frame, chassis_lidar_frame).load() -def joint_transform(model: LoadedRobotModel, child_link: str) -> Transform: +def joint_transform(model: LoadedRobotModel, child_link: str) -> TransformStamped: """The fixed joint that places ``child_link`` under its parent.""" joint = next(j for j in model.joints if j.child_link == child_link) - return Transform( - translation=Vector3(*joint.origin_xyz), - rotation=Quaternion.from_euler(Vector3(*joint.origin_rpy)), - frame_id=joint.parent_link, + return TransformStamped( + header=Header(stamp=time_from_seconds(0.0), frame_id=joint.parent_link), child_frame_id=joint.child_link, + transform=Transform( + translation=Vector3(*joint.origin_xyz), + rotation=quaternion_from_euler(*joint.origin_rpy), + ), ) -def mount_transforms(config: R1ProLioConfig) -> list[Transform]: +def mount_transforms(config: R1ProLioConfig) -> list[TransformStamped]: """Rooted at Point-LIO's frame, plus the URDF's chassis-lidar edge.""" model = lio_model(config.lidar_frame, config.chassis_lidar_frame) base_to_chassis_lidar = joint_transform(model, config.chassis_lidar_frame) - assert base_to_chassis_lidar.frame_id == config.base_frame - base_to_lidar = base_to_chassis_lidar + joint_transform(model, config.lidar_frame) - return [-base_to_lidar, base_to_chassis_lidar] + assert base_to_chassis_lidar.header.frame_id == config.base_frame + base_to_lidar = compose_transforms( + base_to_chassis_lidar, joint_transform(model, config.lidar_frame) + ) + return [inverse_transform(base_to_lidar), base_to_chassis_lidar] -def base_pose(odometry: Odometry, lidar_to_base: Transform) -> PoseStamped: +def base_pose(odometry: Odometry, lidar_to_base: TransformStamped) -> PoseStamped: """Point-LIO's lidar pose carried through the mount to base_link.""" - odom_to_lidar = Transform( - translation=odometry.pose.position, - rotation=odometry.pose.orientation, - frame_id=odometry.frame_id, - ts=odometry.ts, - ) - odom_to_base = odom_to_lidar + lidar_to_base - return PoseStamped( - ts=odometry.ts, - frame_id=odometry.frame_id, - position=odom_to_base.translation, - orientation=odom_to_base.rotation, - ) + odom_to_lidar = transform_from_odometry(odometry) + return pose_from_transform(compose_transforms(odom_to_lidar, lidar_to_base)) class R1ProLio(StaticTfPublisher): @@ -91,7 +95,7 @@ class R1ProLio(StaticTfPublisher): odometry: In[Odometry] pose: Out[PoseStamped] - def transforms(self) -> list[Transform]: + def transforms(self) -> list[TransformStamped]: return mount_transforms(self.config) @rpc diff --git a/dimos/robot/galaxea/r1pro/test_connection_messages.py b/dimos/robot/galaxea/r1pro/test_connection_messages.py index fd1670c06d8..258ec89394a 100644 --- a/dimos/robot/galaxea/r1pro/test_connection_messages.py +++ b/dimos/robot/galaxea/r1pro/test_connection_messages.py @@ -12,13 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dataclasses import asdict import math +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.dimos_msgs.msg import MotorCommandArray from dimos_generated.geometry_msgs.msg import Twist, TwistStamped, Vector3 from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from dimos.hardware.whole_body.spec import VEL_STOP @@ -37,8 +40,8 @@ def robot(): def speed(sec, nanosec, *, vx=1.0, vy=0.0, wz=0.0): return TwistStamped( - header=Header(stamp=Time(sec=sec, nanosec=nanosec)), - twist=Twist(linear=Vector3(x=vx, y=vy), angular=Vector3(z=wz)), + header=Header(stamp=Time(sec=sec, nanosec=nanosec), frame_id=""), + twist=Twist(linear=Vector3(x=vx, y=vy, z=0.0), angular=Vector3(z=wz, x=0.0, y=0.0)), ) @@ -54,17 +57,15 @@ def test_integrated_pose_odometry_and_tf_preserve_nanoseconds(robot, monkeypatch assert pose.pose.position.y == pytest.approx(0.05) assert pose.pose.orientation.z == pytest.approx(math.sin(0.01)) assert odometry.pose.pose == pose.pose - assert odometry.twist.twist.linear == Vector3(x=1.0, y=0.5) - assert [edge.child_frame_id for edge in tf.transforms] == [ - robot.config.frame_id, - robot.config.lidar_frame_id, - ] - assert tf.transforms[1].header.frame_id == robot.config.frame_id - assert tf.transforms[1].transform.rotation.w == 1.0 + assert odometry.twist.twist.linear == Vector3(x=1.0, y=0.5, z=0.0) + assert [edge.child_frame_id for edge in tf.transforms] == [robot.config.frame_id] + assert tf.transforms[0].header.frame_id == robot.config.odom_frame_id for message in (pose, odometry, *tf.transforms): assert message.header.stamp == Time(sec=1700000000, nanosec=223456789) for message in (pose, odometry, tf): - assert type(message).decode(message.encode()) == message + np.testing.assert_equal( + asdict(cdr_decode(cdr_encode(message), type(message))), asdict(message) + ) @pytest.mark.parametrize("second_stamp", [(0, 0), (-1, 0), (2, 1)]) @@ -77,7 +78,9 @@ def test_clock_jump_does_not_move_base(robot, monkeypatch, second_stamp): assert robot._odom_x == 0.0 -@pytest.mark.parametrize("oldest", [Time(), Time(sec=1700000000, nanosec=123456788)]) +@pytest.mark.parametrize( + "oldest", [Time(sec=0, nanosec=0), Time(sec=1700000000, nanosec=123456788)] +) def test_joint_snapshot_uses_oldest_exact_stamp_and_all_segments(robot, monkeypatch, oldest): for callback, count, stamp, value in [ (robot._on_feedback_torso, 4, oldest, 1.0), @@ -86,10 +89,11 @@ def test_joint_snapshot_uses_oldest_exact_stamp_and_all_segments(robot, monkeypa ]: callback( JointState( - header=Header(stamp=stamp), - position=[value] * count, - velocity=[0.5] * count, - effort=[0.2] * count, + header=Header(stamp=stamp, frame_id=""), + position=np.asarray([value] * count, dtype=np.float64), + velocity=np.asarray([0.5] * count, dtype=np.float64), + effort=np.asarray([0.2] * count, dtype=np.float64), + name=[], ), None, ) @@ -106,7 +110,7 @@ def collect(message): assert list(state.velocity) == [0.5] * 18 assert list(state.name) == connection.R1PRO_UPPER_BODY_JOINTS assert state.header.stamp == oldest - assert JointState.decode(state.encode()) == state + np.testing.assert_equal(asdict(cdr_decode(cdr_encode(state), JointState)), asdict(state)) @pytest.fixture() @@ -114,7 +118,9 @@ def command_sink(robot, monkeypatch, mocker): ros = mocker.Mock(spec=RawROS) monkeypatch.setattr(robot, "_ros", ros) monkeypatch.setattr(connection, "dimos_to_ros", lambda msg, _: msg) - monkeypatch.setattr(connection, "header_now", lambda: Header(stamp=Time(sec=9, nanosec=123))) + monkeypatch.setattr( + connection, "header_now", lambda: Header(stamp=Time(sec=9, nanosec=123), frame_id="") + ) for name in ("_cmd_torso_topic", "_cmd_left_topic", "_cmd_right_topic", "_speed_topic"): monkeypatch.setattr(robot, name, RawROSTopic(name, JointState)) return ros.publish @@ -122,7 +128,12 @@ def command_sink(robot, monkeypatch, mocker): def motor_command(): return MotorCommandArray( - q=list(range(18)), dq=[0.0, VEL_STOP, 0.2] * 6, kp=[0.0] * 18, kd=[0.0] * 18, tau=[0.0] * 18 + q=np.asarray(list(range(18)), dtype=np.float64), + dq=np.asarray([0.0, VEL_STOP, 0.2] * 6, dtype=np.float64), + kp=np.asarray([0.0] * 18, dtype=np.float64), + kd=np.asarray([0.0] * 18, dtype=np.float64), + tau=np.asarray([0.0] * 18, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) @@ -152,6 +163,6 @@ def test_base_command_only_forwards_supported_axes(robot, command_sink): ) message = command_sink.call_args.args[1] assert message == TwistStamped( - header=Header(stamp=Time(sec=9, nanosec=123)), - twist=Twist(linear=Vector3(x=1.0, y=2.0), angular=Vector3(z=6.0)), + header=Header(stamp=Time(sec=9, nanosec=123), frame_id=""), + twist=Twist(linear=Vector3(x=1.0, y=2.0, z=0.0), angular=Vector3(z=6.0, x=0.0, y=0.0)), ) diff --git a/dimos/robot/galaxea/r1pro/test_head_cameras.py b/dimos/robot/galaxea/r1pro/test_head_cameras.py index dceb9189216..89c8f5572b0 100644 --- a/dimos/robot/galaxea/r1pro/test_head_cameras.py +++ b/dimos/robot/galaxea/r1pro/test_head_cameras.py @@ -16,7 +16,9 @@ from pathlib import Path -from dimos.msgs.sensor_msgs.CompressedImage import CompressedImage +from dimos_generated.sensor_msgs.msg import CompressedImage +import numpy as np + from dimos.robot.galaxea.r1pro.blueprints.basic.r1pro_coordinator import r1pro_control from dimos.robot.galaxea.r1pro.connection import R1ProConnection from dimos.robot.galaxea.r1pro.head_cameras import ( @@ -86,12 +88,13 @@ def test_each_eye_gets_its_own_intrinsics_from_the_factory_calibration(tmp_path: left, right = head_camera_infos(str(path)) - assert (left.width, left.height, left.frame_id) == (1920, 1536, "camera_head_left_link") - assert right.frame_id == "camera_head_right_link" - assert left.K[0] == 1012.5 and right.K[2] == 958.7 + assert (left.width, left.height, left.header.frame_id) == (1920, 1536, "camera_head_left_link") + assert right.header.frame_id == "camera_head_right_link" + assert left.k[0] == 1012.5 and right.k[2] == 958.7 # Eight rational-polynomial coefficients; the six zero tilt terms are dropped. - assert left.distortion_model == "rational_polynomial" and len(left.D) == 8 - assert left.P[:3] == left.K[:3] and left.P[3] == 0.0 + assert left.distortion_model == "rational_polynomial" and len(left.d) == 8 + np.testing.assert_array_equal(left.p[:3], left.k[:3]) + assert left.p[3] == 0.0 def test_r1pro_control_publishes_both_eyes_intrinsics() -> None: diff --git a/dimos/robot/manipulators/common/coordinators.py b/dimos/robot/manipulators/common/coordinators.py index 19da3227047..1dbbc65b667 100644 --- a/dimos/robot/manipulators/common/coordinators.py +++ b/dimos/robot/manipulators/common/coordinators.py @@ -23,11 +23,11 @@ from __future__ import annotations +from dimos_generated.geometry_msgs.msg import PoseStamped, TwistStamped +from dimos_generated.tf2_msgs.msg import TFMessage + from dimos.control.coordinator import ControlCoordinator from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped -from dimos.msgs.tf2_msgs.TFMessage import TFMessage class ArmPoseCoordinator(ControlCoordinator): diff --git a/dimos/robot/manipulators/dual_openyam/test_blueprints.py b/dimos/robot/manipulators/dual_openyam/test_blueprints.py index ab7129ed98a..746276b9a15 100644 --- a/dimos/robot/manipulators/dual_openyam/test_blueprints.py +++ b/dimos/robot/manipulators/dual_openyam/test_blueprints.py @@ -14,9 +14,11 @@ from typing import Any, cast -from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Float32, Header +import numpy as np import pytest from pytest_mock import MockerFixture @@ -72,14 +74,32 @@ def test_mock_webxr_coordinator_commands_both_arms_and_grippers( DualOpenYamPinkPoseTargetSolver, "frame_poses", return_value={ - "left_grasp_frame": PoseStamped(pose=Pose(position=Point(x=0.5, y=0.2, z=0.4))), - "right_grasp_frame": PoseStamped(pose=Pose(position=Point(x=0.5, y=-0.2, z=0.4))), + "left_grasp_frame": PoseStamped( + pose=Pose( + position=Point(x=0.5, y=0.2, z=0.4), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ), + "right_grasp_frame": PoseStamped( + pose=Pose( + position=Point(x=0.5, y=-0.2, z=0.4), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ), }, ) mocker.patch.object( DualOpenYamPinkPoseTargetSolver, "step", - return_value=JointState(name=DUAL_OPENYAM_ARM_JOINTS, position=[0.01] * 12), + return_value=JointState( + name=DUAL_OPENYAM_ARM_JOINTS, + position=np.asarray([0.01] * 12, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), ) mocker.patch.object(TickLoop, "start") coordinator = DualOpenYamCoordinator(publish_joint_state=False, **kwargs) @@ -97,15 +117,21 @@ def test_mock_webxr_coordinator_commands_both_arms_and_grippers( coordinator._dispatch( "left_cartesian_command", PoseStamped( - header=Header(frame_id=DUAL_OPENYAM_WEBXR_TASK_NAME), - pose=Pose(position=Point(x=1.0)), + header=Header(frame_id=DUAL_OPENYAM_WEBXR_TASK_NAME, stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=1.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), ) coordinator._dispatch( "right_cartesian_command", PoseStamped( - header=Header(frame_id=DUAL_OPENYAM_WEBXR_TASK_NAME), - pose=Pose(position=Point(x=-1.0)), + header=Header(frame_id=DUAL_OPENYAM_WEBXR_TASK_NAME, stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=-1.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), ) assert coordinator._tick_loop is not None diff --git a/dimos/robot/manipulators/dual_openyam/test_teleop_ik.py b/dimos/robot/manipulators/dual_openyam/test_teleop_ik.py index 104e24945fa..7815e91b3f8 100644 --- a/dimos/robot/manipulators/dual_openyam/test_teleop_ik.py +++ b/dimos/robot/manipulators/dual_openyam/test_teleop_ik.py @@ -14,8 +14,10 @@ """Objective tests for Dual OpenYAM WebXR teleoperation.""" +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest @@ -54,7 +56,10 @@ def test_solver_uses_nominal_posture_without_manipulability() -> None: solver = _solver() seed = JointState( name=DUAL_OPENYAM_ARM_JOINTS, - position=DUAL_OPENYAM_HOME_JOINTS, + position=np.asarray(DUAL_OPENYAM_HOME_JOINTS, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) targets = solver.frame_poses(seed, _TARGET_FRAMES) @@ -71,7 +76,10 @@ def test_webxr_solver_matches_a1z_target_tracking_speed() -> None: solver = _solver() state = JointState( name=DUAL_OPENYAM_ARM_JOINTS, - position=DUAL_OPENYAM_HOME_JOINTS, + position=np.asarray(DUAL_OPENYAM_HOME_JOINTS, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) initial = solver.frame_poses(state, _TARGET_FRAMES) targets = { diff --git a/dimos/robot/manipulators/openarm/test_openarm_teleop.py b/dimos/robot/manipulators/openarm/test_openarm_teleop.py index c242f34773f..1610ab23d28 100644 --- a/dimos/robot/manipulators/openarm/test_openarm_teleop.py +++ b/dimos/robot/manipulators/openarm/test_openarm_teleop.py @@ -16,7 +16,8 @@ from typing import Any, cast -from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.sensor_msgs.msg import JointState from dimos_generated.std_msgs.msg import Float32, Header import numpy as np @@ -161,10 +162,18 @@ def test_openarm_webxr_commands_both_arms_and_grippers_through_coordinator( "frame_poses", return_value={ "openarm_left_grasp_frame": PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(x=0.5, y=0.2, z=0.4)) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.5, y=0.2, z=0.4), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), "openarm_right_grasp_frame": PoseStamped( - header=Header(frame_id=""), pose=Pose(position=Point(x=0.5, y=-0.2, z=0.4)) + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.5, y=-0.2, z=0.4), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), }, ) @@ -173,7 +182,10 @@ def test_openarm_webxr_commands_both_arms_and_grippers_through_coordinator( "step", return_value=JointState( name=list(OPENARM_ARM_JOINTS), - position=[0.01] * len(OPENARM_ARM_JOINTS), + position=np.asarray([0.01] * len(OPENARM_ARM_JOINTS), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ) mocker.patch.object(TickLoop, "start") @@ -204,15 +216,21 @@ def test_openarm_webxr_commands_both_arms_and_grippers_through_coordinator( coordinator._dispatch( "left_cartesian_command", PoseStamped( - header=Header(frame_id=OPENARM_WEBXR_TASK_NAME), - pose=Pose(position=Point(x=1.0, y=0.0, z=0.0)), + header=Header(frame_id=OPENARM_WEBXR_TASK_NAME, stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=1.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), ) coordinator._dispatch( "right_cartesian_command", PoseStamped( - header=Header(frame_id=OPENARM_WEBXR_TASK_NAME), - pose=Pose(position=Point(x=-1.0, y=0.0, z=0.0)), + header=Header(frame_id=OPENARM_WEBXR_TASK_NAME, stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=-1.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), ) @@ -251,7 +269,13 @@ def test_openarm_teleop_pink_objective_uses_robot_specific_tuning() -> None: lm_damping=0.01, gain=0.25, ) - seed = JointState(name=OPENARM_ARM_JOINTS, position=[0.0] * len(OPENARM_ARM_JOINTS)) + seed = JointState( + name=OPENARM_ARM_JOINTS, + position=np.asarray([0.0] * len(OPENARM_ARM_JOINTS), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) solver = OpenArmPinkPoseTargetSolver(_solver_config(model, frames, config)) targets = solver.frame_poses(seed, frames) @@ -273,7 +297,12 @@ def test_openarm_teleop_pink_objective_uses_robot_specific_tuning() -> None: moved_seed = JointState( name=OPENARM_ARM_JOINTS, - position=np.tile([0.1, -0.1, 0.2, 0.4, 0.1, -0.1, 0.2], 2).tolist(), + position=np.asarray( + np.tile([0.1, -0.1, 0.2, 0.4, 0.1, -0.1, 0.2], 2).tolist(), dtype=np.float64 + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) solver.reset() solver.step(targets, moved_seed, 0.01) @@ -297,11 +326,17 @@ def test_openarm_bimanual_pink_steps_from_canonical_zero_with_bounded_updates() gain=0.25, ) ik = OpenArmPinkPoseTargetSolver(_solver_config(model, frames, config)) - seed = JointState(name=OPENARM_ARM_JOINTS, position=[0.0] * len(OPENARM_ARM_JOINTS)) + seed = JointState( + name=OPENARM_ARM_JOINTS, + position=np.asarray([0.0] * len(OPENARM_ARM_JOINTS), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) initial = ik.frame_poses(seed, frames) targets = { name: PoseStamped( - header=Header(frame_id=pose.header.frame_id), + header=Header(frame_id=pose.header.frame_id, stamp=Time(sec=0, nanosec=0)), pose=Pose( position=Point( x=pose.pose.position.x, y=pose.pose.position.y, z=pose.pose.position.z + 0.01 diff --git a/dimos/robot/manipulators/openyam/test_teleop_ik.py b/dimos/robot/manipulators/openyam/test_teleop_ik.py index fcc0cf193d3..e1dd2846374 100644 --- a/dimos/robot/manipulators/openyam/test_teleop_ik.py +++ b/dimos/robot/manipulators/openyam/test_teleop_ik.py @@ -14,8 +14,11 @@ """Objective tests for single-arm OpenYAM Quest teleoperation.""" +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Pose, PoseStamped, Quaternion from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from scipy.spatial.transform import Rotation @@ -52,7 +55,13 @@ def _solver( @pytest.mark.self_hosted def test_solver_weights_large_joints_above_wrist_joints() -> None: solver = _solver() - state = JointState(name=OPENYAM_ARM_JOINTS, position=OPENYAM_HOME_JOINTS) + state = JointState( + name=OPENYAM_ARM_JOINTS, + position=np.asarray(OPENYAM_HOME_JOINTS, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) targets = solver.frame_poses(state, (_TARGET_FRAME,)) assert solver.step(targets, state, 0.01) is not None @@ -67,7 +76,13 @@ def _orientation_target_motion( solver_type: type[PinkPoseTargetSolver], ) -> tuple[float, float]: solver = _solver(solver_type) - initial = JointState(name=OPENYAM_ARM_JOINTS, position=OPENYAM_HOME_JOINTS) + initial = JointState( + name=OPENYAM_ARM_JOINTS, + position=np.asarray(OPENYAM_HOME_JOINTS, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) pose = solver.frame_poses(initial, (_TARGET_FRAME,))[_TARGET_FRAME] q = pose.pose.orientation rotated = ( diff --git a/dimos/robot/manipulators/xarm/blueprints/grasp.py b/dimos/robot/manipulators/xarm/blueprints/grasp.py index 5376657f7f1..92b10a93edf 100644 --- a/dimos/robot/manipulators/xarm/blueprints/grasp.py +++ b/dimos/robot/manipulators/xarm/blueprints/grasp.py @@ -27,6 +27,18 @@ from __future__ import annotations +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.std_msgs.msg import Header + from dimos.control.coordinator import TaskConfig from dimos.core.coordination.blueprints import Blueprint, autoconnect from dimos.core.global_config import global_config @@ -40,10 +52,6 @@ RoboPlanPointCloudSelfFilter, ) from dimos.mapping.ray_tracing.module import RayTracingVoxelMap -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.perception.experimental.object_scene_registration import ObjectSceneRegistrationModule from dimos.robot.manipulators.common.blueprints import coordinator, trajectory_task from dimos.robot.manipulators.xarm.config import ( @@ -106,11 +114,13 @@ # resolves into world, and every cloud the camera produces is silently unusable. # link7 is a frame the model already publishes, so ManipulationModule emits the # whole chain from one loop at one rate. -XARM_WRIST_CAMERA_TRANSFORM = Transform( - translation=Vector3(x=0.06693724, y=-0.0309563, z=0.00691482), - rotation=Quaternion(0.70513398, 0.00535696, 0.70897578, -0.01052180), # xyzw - frame_id="link7", +XARM_WRIST_CAMERA_TRANSFORM = TransformStamped( + header=Header(frame_id="link7", stamp=Time(sec=0, nanosec=0)), child_frame_id="camera_link", + transform=Transform( + translation=Vector3(x=0.06693724, y=-0.0309563, z=0.00691482), + rotation=Quaternion(x=0.70513398, y=0.00535696, z=0.70897578, w=-0.01052180), + ), ) @@ -119,7 +129,13 @@ # of the 12 cm pedestal data/xarm7 uses. Inheriting that offset would put the # planning model 12 cm above the arm MuJoCo simulates. _model = make_xarm7_sim_robot_config( - base_pose=PoseStamped(frame_id="world"), + base_pose=PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), tf_extra_links=[], ) _hardware = make_xarm7_sim_hardware(XARM_GRASP_SCENE_PATH, home_joints=XARM_GRASP_SCAN_JOINTS) @@ -127,7 +143,13 @@ _model = make_xarm7_model_config( add_gripper=True, gripper_hardware_id="arm", - base_pose=PoseStamped(frame_id="world"), + base_pose=PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), tf_extra_links=["link7"], ) _hardware = xarm7_hardware("arm", gripper=True) diff --git a/dimos/robot/manipulators/xarm/config.py b/dimos/robot/manipulators/xarm/config.py index cbbc6be7bae..44c23c9a42b 100644 --- a/dimos/robot/manipulators/xarm/config.py +++ b/dimos/robot/manipulators/xarm/config.py @@ -20,6 +20,10 @@ from pathlib import Path from typing import Any +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.std_msgs.msg import Header + from dimos.control.components import ( HardwareComponent, HardwareType, @@ -28,8 +32,6 @@ from dimos.hardware.spec import JointLimits from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.robot.assets.model import RobotModel from dimos.robot.assets.source import RobotDescriptionSource from dimos.robot.manipulators._modeling import joint_names @@ -66,7 +68,12 @@ # The sim scene stands the arm on a pedestal: xarm7.xml mounts link_base at # z=0.12. Place the planning model to match, or the planner solves poses 12cm # below the arm it is driving and every grasp closes on air. -XARM7_SIM_BASE_POSE = PoseStamped(frame_id="world", position=Vector3(z=0.12)) +XARM7_SIM_BASE_POSE = PoseStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(z=0.12, x=0.0, y=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), +) def make_xarm7_sim_robot_config( @@ -320,7 +327,15 @@ def make_xarm_model_config( package_paths=XARM_PACKAGE_PATHS, xacro_args=xacro_args, ).with_default_joint_acceleration_limit(2.0), - base_pose=base_pose if base_pose is not None else PoseStamped(), + base_pose=base_pose + if base_pose is not None + else PoseStamped( + header=Header(frame_id="", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), joint_names=model_joint_names, base_link=f"{prefix}link_base", planning_groups=[ diff --git a/dimos/robot/raw_robot_bridge.py b/dimos/robot/raw_robot_bridge.py index bd010ee85f0..0550c0fc51d 100644 --- a/dimos/robot/raw_robot_bridge.py +++ b/dimos/robot/raw_robot_bridge.py @@ -36,6 +36,10 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist, TwistStamped, Vector3 +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, JointState, PointCloud2 +from dimos_generated.std_msgs.msg import Float32, Header +from dimos_generated.tf2_msgs.msg import TFMessage import numpy as np from PIL import Image as PILImage from scipy.spatial.transform import Rotation @@ -55,15 +59,10 @@ RAW_MAX_LINEAR_MPS, RAW_TOPIC_PREFIX, ) -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.std_msgs.Float32 import Float32 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.geometry import quaternion_array, vector_array +from dimos.msgs.image import image_to_rgb, image_view +from dimos.msgs.pointcloud import pointcloud_xyz +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -118,27 +117,36 @@ def close(self) -> None: def jpeg_bytes(image: Image, quality: int = 90) -> bytes: buf = io.BytesIO() - PILImage.fromarray(image.to_rgb().as_numpy()).save(buf, format="JPEG", quality=quality) + PILImage.fromarray(image_to_rgb(image)).save(buf, format="JPEG", quality=quality) return buf.getvalue() def depth_f32(image: Image) -> bytes: """Metric optical-axis depth without quantization or image compression.""" - depth = image.as_numpy() - if image.format != ImageFormat.DEPTH or depth.ndim != 2 or depth.dtype.kind != "f": + depth = image_view(image) + if image.encoding not in {"32FC1", "64FC1"} or depth.ndim != 2 or depth.dtype.kind != "f": raise ValueError("Raw depth requires a 2D floating-point DEPTH image in metres") return np.ascontiguousarray(depth, dtype=" bytes: - points, _ = cloud.as_numpy() + points = pointcloud_xyz(cloud) return np.ascontiguousarray(points, dtype=" str: - p, q = pose.position, pose.orientation + p, q = pose.pose.position, pose.pose.orientation return json.dumps( - {"t": pose.ts, "x": p.x, "y": p.y, "z": p.z, "qx": q.x, "qy": q.y, "qz": q.z, "qw": q.w} + { + "t": to_seconds(pose.header.stamp), + "x": p.x, + "y": p.y, + "z": p.z, + "qx": q.x, + "qy": q.y, + "qz": q.z, + "qw": q.w, + } ) @@ -252,14 +260,18 @@ def start(self) -> None: cfg = self.config self._topics = RawTopics(cfg.endpoint, cfg.prefix, listen=True) q = cfg.jpeg_quality - self.color_image.subscribe(lambda img: self._put("camera/jpeg", jpeg_bytes(img, q), img.ts)) + self.color_image.subscribe( + lambda img: self._put("camera/jpeg", jpeg_bytes(img, q), to_seconds(img.header.stamp)) + ) self.depth_image.subscribe(self._on_depth) - self.lidar.subscribe(lambda cloud: self._put("lidar/xyz_f32", xyz_f32(cloud), cloud.ts)) + self.lidar.subscribe( + lambda cloud: self._put("lidar/xyz_f32", xyz_f32(cloud), to_seconds(cloud.header.stamp)) + ) self.odom.subscribe(lambda pose: self._put("odom/json", odom_json(pose))) self.camera_info.subscribe( lambda info: self._put( "camera_info/json", - json.dumps({"width": info.width, "height": info.height, "K": info.K}), + json.dumps({"width": info.width, "height": info.height, "K": list(info.k)}), ) ) self.coordinator_joint_state.subscribe(self._on_joint_state) @@ -280,8 +292,17 @@ def stop(self) -> None: self._stop.set() if self._drive_thread is not None: self._drive_thread.join(timeout=1.0) - self.cmd_vel.publish(Twist()) - self.ee_twist_command.publish(TwistStamped()) + self.cmd_vel.publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) + self.ee_twist_command.publish( + TwistStamped( + header=Header(stamp=time_from_seconds(time.time()), frame_id=""), + twist=Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ), + ) + ) self._topics.close() self._topics = None super().stop() @@ -312,17 +333,37 @@ def _drive(self) -> None: while not self._stop.wait(1.0 / self.config.drive_hz): vx, vy, wz = self._base.current() if (vx, vy, wz) != (0.0, 0.0, 0.0): - self.cmd_vel.publish(Twist(linear=(vx, vy, 0.0), angular=(0.0, 0.0, wz))) + self.cmd_vel.publish( + Twist(linear=Vector3(x=vx, y=vy, z=0.0), angular=Vector3(x=0.0, y=0.0, z=wz)) + ) moving["base"] = True elif moving["base"]: - self.cmd_vel.publish(Twist()) + self.cmd_vel.publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) moving["base"] = False arm = self._arm.current() if any(arm): - self.ee_twist_command.publish(TwistStamped(linear=arm[:3], angular=arm[3:])) + self.ee_twist_command.publish( + TwistStamped( + header=Header(stamp=time_from_seconds(time.time()), frame_id=""), + twist=Twist( + linear=Vector3(x=arm[0], y=arm[1], z=arm[2]), + angular=Vector3(x=arm[3], y=arm[4], z=arm[5]), + ), + ) + ) moving["arm"] = True elif moving["arm"]: - self.ee_twist_command.publish(TwistStamped()) + self.ee_twist_command.publish( + TwistStamped( + header=Header(stamp=time_from_seconds(time.time()), frame_id=""), + twist=Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), + angular=Vector3(x=0.0, y=0.0, z=0.0), + ), + ) + ) moving["arm"] = False def _on_depth(self, image: Image) -> None: @@ -334,15 +375,15 @@ def _on_depth(self, image: Image) -> None: logger.warning("Raw depth frames rejected", error=str(exc)) return info = { - "t": image.ts, + "t": to_seconds(image.header.stamp), "width": image.width, "height": image.height, "dtype": " None: now = time.monotonic() @@ -354,7 +395,7 @@ def _on_joint_state(self, state: JointState) -> None: velocities = dict(zip(state.name, state.velocity, strict=False)) joints = [n for n in state.name if n != gripper_joint] message: dict[str, Any] = { - "t": state.ts, + "t": to_seconds(state.header.stamp), "joint_names": joints, "positions": [positions[n] for n in joints], "velocities": [velocities.get(n, 0.0) for n in joints], @@ -372,17 +413,17 @@ def _on_joint_state(self, state: JointState) -> None: payload = json.dumps(message, allow_nan=False) except ValueError: return # a non-finite joint value would make invalid JSON - self._put("arm/state/json", payload, state.ts) + self._put("arm/state/json", payload, to_seconds(state.header.stamp)) def _on_tf(self, message: TFMessage) -> None: for transform in message.transforms: child = transform.child_frame_id if child not in (self.config.camera_frame, self.config.ee_frame): continue - q = transform.rotation.to_numpy() + q = quaternion_array(transform.transform.rotation) pose = { - "frame": transform.frame_id, - "xyz": transform.translation.to_numpy().tolist(), + "frame": transform.header.frame_id, + "xyz": vector_array(transform.transform.translation).tolist(), "quaternion_xyzw": Rotation.from_quat(q).as_quat().tolist(), } if child == self.config.ee_frame: @@ -390,4 +431,8 @@ def _on_tf(self, message: TFMessage) -> None: self._ee_pose = pose self._ee_seen = time.monotonic() else: - self._put("camera_pose/json", json.dumps({"t": transform.ts, **pose}), transform.ts) + self._put( + "camera_pose/json", + json.dumps({"t": to_seconds(transform.header.stamp), **pose}), + to_seconds(transform.header.stamp), + ) diff --git a/dimos/robot/test_raw_robot_bridge.py b/dimos/robot/test_raw_robot_bridge.py index e3d78146e41..a70c9192640 100644 --- a/dimos/robot/test_raw_robot_bridge.py +++ b/dimos/robot/test_raw_robot_bridge.py @@ -18,18 +18,27 @@ import time from unittest.mock import MagicMock +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.sensor_msgs.msg import JointState +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage import numpy as np from PIL import Image as PILImage import pytest -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.core.transport_factory import rpc_backend +from dimos.msgs.geometry import vector_array +from dimos.msgs.image import image_from_array +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import time_from_seconds from dimos.robot.raw_robot_bridge import ( Deadman, RawRobotBridge, @@ -43,6 +52,33 @@ ENDPOINT = "tcp/127.0.0.1:17448" +@pytest.fixture(autouse=True) +def offline_rpc(mocker): + for method in ("start", "serve_module_rpc", "stop"): + mocker.patch.object(rpc_backend(), method) + + +def _transform(translation=None, rotation=None, frame_id="", child_frame_id="", ts=0.0): + return TransformStamped( + header=Header(stamp=time_from_seconds(ts), frame_id=frame_id), + child_frame_id=child_frame_id, + transform=Transform( + translation=translation if translation is not None else Vector3(0.0, 0.0, 0.0), + rotation=rotation if rotation is not None else Quaternion(0.0, 0.0, 0.0, 1.0), + ), + ) + + +def _joint_state(name, position, velocity=(), ts=0.0): + return JointState( + header=Header(stamp=time_from_seconds(ts), frame_id=""), + name=name, + position=np.asarray(position, dtype=np.float64), + velocity=np.asarray(velocity, dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + + def wait_for(condition, timeout_s: float = 5.0) -> None: # type: ignore[no-untyped-def] deadline = time.monotonic() + timeout_s while not condition(): @@ -72,7 +108,12 @@ def test_client_receives_sensor_topics_and_robot_receives_commands( robot.subscribe("cmd_vel/json", lambda b, _ts: commands.append(b)) time.sleep(0.5) # subscriptions propagate over the link - pose = PoseStamped(position=(1.0, 2.0, 0.5), orientation=(0.0, 0.0, 0.0, 1.0), ts=3.5) + pose = PoseStamped( + header=Header(stamp=time_from_seconds(3.5), frame_id=""), + pose=Pose( + position=Point(x=1, y=2, z=0.5), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), + ) robot.put("odom/json", odom_json(pose)) robot.put("camera/jpeg", b"\xff\xd8jpeg", ts=4.25) client.put("cmd_vel/json", json.dumps({"vx": 0.3, "vy": 0.0, "wz": 0.1, "t": 1.0})) @@ -87,15 +128,33 @@ def test_client_receives_sensor_topics_and_robot_receives_commands( def test_encoders_are_plain_formats() -> None: frame = np.zeros((6, 8, 3), dtype=np.uint8) frame[:, :, 0] = 200 - decoded = PILImage.open(io.BytesIO(jpeg_bytes(Image.from_numpy(frame, ts=1.0)))) + decoded = PILImage.open( + io.BytesIO( + jpeg_bytes( + image_from_array( + frame, encoding="rgb8", header=Header(stamp=time_from_seconds(1.0), frame_id="") + ) + ) + ) + ) assert decoded.format == "JPEG" and decoded.size == (8, 6) points = np.arange(12, dtype=np.float32).reshape(4, 3) - raw = xyz_f32(PointCloud2.from_numpy(points, timestamp=2.0)) + raw = xyz_f32( + pointcloud_from_xyz(points, header=Header(stamp=time_from_seconds(2.0), frame_id="")) + ) assert np.frombuffer(raw, dtype=" None: bridge._arm.set(json.dumps({"vx": 0.05, "vz": 9, "wz": -9, "t": 60})) _drive(bridge, 2) twists = [c.args[0] for c in bridge.ee_twist_command.publish.call_args_list] - assert [(*t.linear.to_numpy(), *t.angular.to_numpy()) for t in twists] == [ + assert [(*vector_array(t.twist.linear), *vector_array(t.twist.angular)) for t in twists] == [ (0.05, 0, 0.1, 0, 0, -0.5) ] * 2 bridge._arm.until = 0.0 # deadline passed _drive(bridge, 2) last = bridge.ee_twist_command.publish.call_args_list[2:] - assert len(last) == 1 and not any(last[0].args[0].linear.to_numpy()) + assert len(last) == 1 and not any(vector_array(last[0].args[0].twist.linear)) bridge.cmd_vel.publish.assert_not_called() # the base deadman is independent @@ -193,17 +252,17 @@ def test_gripper_opening_passes_through(bridge: RawRobotBridge) -> None: def test_state_reports_joints_tcp_pose_and_normalized_gripper(bridge: RawRobotBridge) -> None: - tcp = Transform( + tcp = _transform( translation=Vector3(0.5, 0.0, 0.3), rotation=Quaternion(1, 0, 0, 0), frame_id="world", child_frame_id="link_tcp", ts=1.0, ) - hidden = Transform(frame_id="world", child_frame_id="cup", ts=1.0) - bridge._on_tf(TFMessage(tcp, hidden)) + hidden = _transform(frame_id="world", child_frame_id="cup", ts=1.0) + bridge._on_tf(TFMessage(transforms=[tcp, hidden])) bridge._on_joint_state( - JointState(name=["j1", "arm/gripper"], position=[0.1, 0.425], velocity=[0.5, 0.0], ts=7.0) + _joint_state(name=["j1", "arm/gripper"], position=[0.1, 0.425], velocity=[0.5, 0.0], ts=7.0) ) (state,) = _published(bridge, "arm/state/json") assert state == { @@ -221,7 +280,7 @@ def test_unconfigured_robot_reports_plain_joints() -> None: module = RawRobotBridge() module._topics = MagicMock() try: - module._on_joint_state(JointState(name=["a", "b"], position=[1.0, 2.0], ts=3.0)) + module._on_joint_state(_joint_state(name=["a", "b"], position=[1.0, 2.0], ts=3.0)) (state,) = _published(module, "arm/state/json") assert state == { "t": 3.0, @@ -229,7 +288,9 @@ def test_unconfigured_robot_reports_plain_joints() -> None: "positions": [1.0, 2.0], "velocities": [0.0, 0.0], } - module._on_tf(TFMessage(Transform(frame_id="world", child_frame_id="link_tcp"))) + module._on_tf( + TFMessage(transforms=[_transform(frame_id="world", child_frame_id="link_tcp")]) + ) assert len(module._topics.put.call_args_list) == 1 finally: module._topics = None @@ -237,17 +298,23 @@ def test_unconfigured_robot_reports_plain_joints() -> None: def test_camera_pose_and_depth_keep_capture_timestamps(bridge: RawRobotBridge) -> None: - camera = Transform( + camera = _transform( translation=Vector3(1, 2, 3), rotation=Quaternion(0, 0, 0, 1), frame_id="world", child_frame_id="wrist_optical", ts=5.0, ) - bridge._on_tf(TFMessage(camera)) + bridge._on_tf(TFMessage(transforms=[camera])) (pose,) = _published(bridge, "camera_pose/json") assert pose == {"t": 5.0, "frame": "world", "xyz": [1, 2, 3], "quaternion_xyzw": [0, 0, 0, 1]} - bridge._on_depth(Image(data=np.ones((2, 3), dtype=np.float32), format=ImageFormat.DEPTH, ts=6)) + bridge._on_depth( + image_from_array( + np.ones((2, 3), dtype=np.float32), + encoding="32FC1", + header=Header(stamp=time_from_seconds(6), frame_id=""), + ) + ) info, depth = bridge._topics.put.call_args_list[-2:] assert info.args[0] == "camera/depth_info/json" and json.loads(info.args[1])["dtype"] == " None: source = np.array([[0.123456, 0.5, 1.25], [np.nan, np.inf, 0]], dtype=">f4")[:, ::-1] - encoded = depth_f32(Image(data=source, format=ImageFormat.DEPTH)) + encoded = depth_f32( + image_from_array( + source, encoding="32FC1", header=Header(stamp=time_from_seconds(0), frame_id="") + ) + ) np.testing.assert_allclose( np.frombuffer(encoded, dtype=" None: - bridge._on_joint_state(JointState(name=["j1"], position=[float("nan")], ts=1.0)) + bridge._on_joint_state(_joint_state(name=["j1"], position=[float("nan")], ts=1.0)) assert not _published(bridge, "arm/state/json") def test_rejected_depth_warns_once(bridge: RawRobotBridge, mocker) -> None: # type: ignore[no-untyped-def] warn = mocker.patch("dimos.robot.raw_robot_bridge.logger.warning") for _ in range(3): - bridge._on_depth(Image(data=np.zeros((2, 3, 3), dtype=np.uint8), format=ImageFormat.RGB)) + bridge._on_depth( + image_from_array( + np.zeros((2, 3, 3), dtype=np.uint8), + encoding="rgb8", + header=Header(stamp=time_from_seconds(0), frame_id=""), + ) + ) assert warn.call_count == 1 diff --git a/dimos/robot/unitree/b1/connection.py b/dimos/robot/unitree/b1/connection.py index 864deb59661..3cf36493a6e 100644 --- a/dimos/robot/unitree/b1/connection.py +++ b/dimos/robot/unitree/b1/connection.py @@ -23,17 +23,15 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import PoseStamped, TwistStamped +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.std_msgs.msg import Int32 +from dimos_generated.tf2_msgs.msg import TFMessage from reactivex.disposable import Disposable from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.std_msgs.Int32 import Int32 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.robot.unitree.b1.b1_command import B1Command from dimos.utils.logging_config import setup_logger @@ -190,7 +188,7 @@ def handle_twist_stamped(self, twist_stamped: TwistStamped) -> None: This is called automatically when messages arrive on cmd_vel input. """ # Extract Twist from TwistStamped - twist = Twist(linear=twist_stamped.linear, angular=twist_stamped.angular) + twist = twist_stamped.twist logger.debug( f"Received cmd_vel: linear=({twist.linear.x:.3f}, {twist.linear.y:.3f}, {twist.linear.z:.3f}), angular=({twist.angular.x:.3f}, {twist.angular.y:.3f}, {twist.angular.z:.3f})" @@ -308,10 +306,8 @@ def _publish_odom_pose(self, msg: Odometry) -> None: """ if self.odom_pose: pose_stamped = PoseStamped( - ts=msg.ts, - frame_id=msg.frame_id, - position=msg.pose.pose.position, - orientation=msg.pose.pose.orientation, + header=msg.header, + pose=msg.pose.pose, ) self.odom_pose.publish(pose_stamped) diff --git a/dimos/robot/unitree/b1/joystick_module.py b/dimos/robot/unitree/b1/joystick_module.py index 7c9e578895d..0af54f0bd2a 100644 --- a/dimos/robot/unitree/b1/joystick_module.py +++ b/dimos/robot/unitree/b1/joystick_module.py @@ -78,7 +78,9 @@ def stop(self) -> None: self.pygame_ready = False # Send stop command - stop_twist = Twist() + stop_twist = Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) stop_twist_stamped = TwistStamped( header=Header(stamp=time_from_seconds(time.time()), frame_id="base_link"), twist=Twist(linear=stop_twist.linear, angular=stop_twist.angular), @@ -118,19 +120,19 @@ def _pygame_loop(self) -> None: # Mode changes - publish to mode_out for connection module if event.key == pygame.K_0: self.current_mode = 0 - mode_msg = Int32() + mode_msg = Int32(data=0) mode_msg.data = 0 self.mode_out.publish(mode_msg) print("Mode: IDLE") elif event.key == pygame.K_1: self.current_mode = 1 - mode_msg = Int32() + mode_msg = Int32(data=0) mode_msg.data = 1 self.mode_out.publish(mode_msg) print("Mode: STAND") elif event.key == pygame.K_2: self.current_mode = 2 - mode_msg = Int32() + mode_msg = Int32(data=0) mode_msg.data = 2 self.mode_out.publish(mode_msg) print("Mode: WALK") @@ -138,11 +140,14 @@ def _pygame_loop(self) -> None: self.keys_held.clear() # Send IDLE mode for emergency stop self.current_mode = 0 - mode_msg = Int32() + mode_msg = Int32(data=0) mode_msg.data = 0 self.mode_out.publish(mode_msg) # Also send zero twist - stop_twist = Twist() + stop_twist = Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), + angular=Vector3(x=0.0, y=0.0, z=0.0), + ) stop_twist.linear = Vector3(x=0, y=0, z=0) stop_twist.angular = Vector3(x=0, y=0, z=0) stop_twist_stamped = TwistStamped( @@ -161,7 +166,7 @@ def _pygame_loop(self) -> None: self.keys_held.discard(event.key) # Generate Twist message from held keys - twist = Twist() + twist = Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) twist.linear = Vector3(x=0, y=0, z=0) twist.angular = Vector3(x=0, y=0, z=0) diff --git a/dimos/robot/unitree/b1/test_cdr_forwarding.py b/dimos/robot/unitree/b1/test_cdr_forwarding.py index 0456fa1e8b0..f2207cd50c8 100644 --- a/dimos/robot/unitree/b1/test_cdr_forwarding.py +++ b/dimos/robot/unitree/b1/test_cdr_forwarding.py @@ -15,6 +15,7 @@ from types import SimpleNamespace from unittest.mock import Mock +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, @@ -24,10 +25,12 @@ Quaternion, Twist, TwistStamped, + TwistWithCovariance, Vector3, ) from dimos_generated.nav_msgs.msg import Odometry from dimos_generated.std_msgs.msg import Header, Int32 +import numpy as np import pytest from dimos.robot.unitree.b1.connection import B1ConnectionModule, MockB1ConnectionModule, RobotMode @@ -38,13 +41,21 @@ def test_b1_odometry_forwarding_preserves_generated_pose_and_exact_header(): source = Odometry( header=Header(frame_id="map", stamp=Time(sec=1700000000, nanosec=123456789)), pose=PoseWithCovariance( - pose=Pose(position=Point(x=1, y=2, z=3), orientation=Quaternion(w=1)) + pose=Pose( + position=Point(x=1, y=2, z=3), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), + covariance=np.zeros(36, dtype=np.float64), + ), + child_frame_id="", + twist=TwistWithCovariance( + twist=Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + covariance=np.zeros(36, dtype=np.float64), ), ) received = [] proxy = SimpleNamespace(odom_pose=SimpleNamespace(publish=received.append)) B1ConnectionModule._publish_odom_pose(proxy, source) - result = PoseStamped.decode(received[0].encode()) + result = cdr_decode(cdr_encode(received[0]), PoseStamped) assert result.header == source.header assert result.pose == source.pose.pose @@ -53,19 +64,19 @@ def test_b1_odometry_forwarding_preserves_generated_pose_and_exact_header(): def test_b1_generated_command_mapping_preserves_mode_clamps_and_source(mode): module = MockB1ConnectionModule() try: - module.handle_mode(Int32.decode(Int32(data=mode).encode())) + module.handle_mode(cdr_decode(cdr_encode(Int32(data=mode)), Int32)) message = TwistStamped( header=Header(frame_id="base_link", stamp=Time(sec=1700000000, nanosec=123456789)), twist=Twist(linear=Vector3(x=2, y=-2, z=2), angular=Vector3(x=2, y=-2, z=4)), ) - before = message.encode() - module.handle_twist_stamped(TwistStamped.decode(before)) + before = cdr_encode(message) + module.handle_twist_stamped(cdr_decode(before, TwistStamped)) command = module._current_cmd assert module.current_mode == mode and command.mode == mode assert command.lx == -1 and command.ly == 1 assert command.rx == (1 if mode == RobotMode.WALK else -1) assert command.ry == (0 if mode == RobotMode.WALK else -1) - assert message.encode() == before + assert cdr_encode(message) == before assert module.socket is None finally: module.stop() @@ -80,8 +91,10 @@ def test_b1_joystick_stop_publishes_generated_zero_without_starting_loop(monkeyp monkeypatch.setattr("dimos.robot.unitree.b1.joystick_module.time.time", lambda: 1700000000.25) module.stop() assert len(received) == 1 - value = TwistStamped.decode(received[0].encode()) + value = cdr_decode(cdr_encode(received[0]), TwistStamped) assert value.header.frame_id == "base_link" assert value.header.stamp == Time(sec=1700000000, nanosec=250000000) - assert value.twist == Twist() + assert value.twist == Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) thread.join.assert_called_once() diff --git a/dimos/robot/unitree/b1/test_connection.py b/dimos/robot/unitree/b1/test_connection.py index df2f97db70b..557cf3ff319 100644 --- a/dimos/robot/unitree/b1/test_connection.py +++ b/dimos/robot/unitree/b1/test_connection.py @@ -27,13 +27,15 @@ import threading import time +from dimos_generated.geometry_msgs.msg import Twist, TwistStamped, Vector3 +from dimos_generated.std_msgs.msg import Header, Int32 import pytest +from dimos.msgs.time import time_from_seconds + _IS_MACOS = sys.platform == "darwin" -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.std_msgs.Int32 import Int32 + from dimos.robot.unitree.b1.connection import MockB1ConnectionModule @@ -52,10 +54,8 @@ def test_watchdog_actually_zeros_commands(self) -> None: # Send a forward command twist_stamped = TwistStamped( - ts=time.time(), - frame_id="base_link", - linear=Vector3(1.0, 0, 0), - angular=Vector3(0, 0, 0), + header=Header(stamp=time_from_seconds(time.time()), frame_id="base_link"), + twist=Twist(linear=Vector3(x=1.0, y=0, z=0), angular=Vector3(x=0, y=0, z=0)), ) conn.handle_twist_stamped(twist_stamped) @@ -100,10 +100,8 @@ def test_watchdog_resets_on_new_command(self) -> None: # Send commands in rapid succession — each resets the 200ms watchdog for val in [1.0, 0.8, 0.6, 0.5]: twist = TwistStamped( - ts=time.time(), - frame_id="base_link", - linear=Vector3(val, 0, 0), - angular=Vector3(0, 0, 0), + header=Header(stamp=time_from_seconds(time.time()), frame_id="base_link"), + twist=Twist(linear=Vector3(x=val, y=0, z=0), angular=Vector3(x=0, y=0, z=0)), ) conn.handle_twist_stamped(twist) time.sleep(0.02) # 20ms between commands, well under timeout @@ -134,10 +132,8 @@ def test_watchdog_thread_efficiency(self) -> None: # Send many commands rapidly (would create many Timer threads in old implementation) for i in range(50): twist = TwistStamped( - ts=time.time(), - frame_id="base_link", - linear=Vector3(i * 0.01, 0, 0), - angular=Vector3(0, 0, 0), + header=Header(stamp=time_from_seconds(time.time()), frame_id="base_link"), + twist=Twist(linear=Vector3(x=i * 0.01, y=0, z=0), angular=Vector3(x=0, y=0, z=0)), ) conn.handle_twist_stamped(twist) time.sleep(0.01) # 100Hz command rate @@ -177,10 +173,8 @@ def blocking_send_loop() -> None: try: # Send command twist = TwistStamped( - ts=time.time(), - frame_id="base_link", - linear=Vector3(1.0, 0, 0), - angular=Vector3(0, 0, 0), + header=Header(stamp=time_from_seconds(time.time()), frame_id="base_link"), + twist=Twist(linear=Vector3(x=1.0, y=0, z=0), angular=Vector3(x=0, y=0, z=0)), ) conn.handle_twist_stamped(twist) assert conn._current_cmd.ly == 1.0 @@ -220,10 +214,8 @@ def test_continuous_commands_prevent_timeout(self) -> None: commands_sent = 0 while time.time() - start < 1.0: twist = TwistStamped( - ts=time.time(), - frame_id="base_link", - linear=Vector3(0.5, 0, 0), - angular=Vector3(0, 0, 0), + header=Header(stamp=time_from_seconds(time.time()), frame_id="base_link"), + twist=Twist(linear=Vector3(x=0.5, y=0, z=0), angular=Vector3(x=0, y=0, z=0)), ) conn.handle_twist_stamped(twist) commands_sent += 1 @@ -259,10 +251,8 @@ def test_watchdog_timing_accuracy(self) -> None: # Send command and record time start_time = time.time() twist = TwistStamped( - ts=time.time(), - frame_id="base_link", - linear=Vector3(1.0, 0, 0), - angular=Vector3(0, 0, 0), + header=Header(stamp=time_from_seconds(time.time()), frame_id="base_link"), + twist=Twist(linear=Vector3(x=1.0, y=0, z=0), angular=Vector3(x=0, y=0, z=0)), ) conn.handle_twist_stamped(twist) @@ -303,10 +293,8 @@ def test_mode_changes_with_watchdog(self) -> None: # Send walk command twist = TwistStamped( - ts=time.time(), - frame_id="base_link", - linear=Vector3(1.0, 0, 0), - angular=Vector3(0, 0, 0), + header=Header(stamp=time_from_seconds(time.time()), frame_id="base_link"), + twist=Twist(linear=Vector3(x=1.0, y=0, z=0), angular=Vector3(x=0, y=0, z=0)), ) conn.handle_twist_stamped(twist) assert conn.current_mode == 2 @@ -320,7 +308,7 @@ def test_mode_changes_with_watchdog(self) -> None: assert conn._current_cmd.ly == 0.0 # Watchdog zeroed it # Now change mode to STAND - mode_msg = Int32() + mode_msg = Int32(data=0) mode_msg.data = 1 # STAND conn.handle_mode(mode_msg) assert conn.current_mode == 1 @@ -350,10 +338,8 @@ def test_watchdog_stops_movement_when_commands_stop(self) -> None: # Simulate sending movement commands for a while for _i in range(5): twist = TwistStamped( - ts=time.time(), - frame_id="base_link", - linear=Vector3(1.0, 0, 0), - angular=Vector3(0, 0, 0.5), # Forward and turning + header=Header(stamp=time_from_seconds(time.time()), frame_id="base_link"), + twist=Twist(linear=Vector3(x=1.0, y=0, z=0), angular=Vector3(x=0, y=0, z=0.5)), ) conn.handle_twist_stamped(twist) time.sleep(0.05) # Send at 20Hz @@ -381,10 +367,8 @@ def test_watchdog_stops_movement_when_commands_stop(self) -> None: # Verify recovery works - send new command twist = TwistStamped( - ts=time.time(), - frame_id="base_link", - linear=Vector3(0.5, 0, 0), - angular=Vector3(0, 0, 0), + header=Header(stamp=time_from_seconds(time.time()), frame_id="base_link"), + twist=Twist(linear=Vector3(x=0.5, y=0, z=0), angular=Vector3(x=0, y=0, z=0)), ) conn.handle_twist_stamped(twist) @@ -421,10 +405,10 @@ def test_rapid_command_thread_safety(self) -> None: def send_commands(thread_id) -> None: for _i in range(10): twist = TwistStamped( - ts=time.time(), - frame_id="base_link", - linear=Vector3(thread_id * 0.1, 0, 0), - angular=Vector3(0, 0, 0), + header=Header(stamp=time_from_seconds(time.time()), frame_id="base_link"), + twist=Twist( + linear=Vector3(x=thread_id * 0.1, y=0, z=0), angular=Vector3(x=0, y=0, z=0) + ), ) conn.handle_twist_stamped(twist) time.sleep(0.01) diff --git a/dimos/robot/unitree/connection.py b/dimos/robot/unitree/connection.py index 784948b6ad8..f4e411ee6f7 100644 --- a/dimos/robot/unitree/connection.py +++ b/dimos/robot/unitree/connection.py @@ -14,12 +14,13 @@ import asyncio from dataclasses import dataclass -import functools import json import threading import time -from typing import Any, TypeAlias, TypeVar +from typing import Any, TypeAlias +from dimos_generated.geometry_msgs.msg import PoseStamped, TransformStamped, Twist +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 import numpy as np from numpy.typing import NDArray from reactivex import operators as ops @@ -38,19 +39,16 @@ from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.resource import Resource -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.geometry import transform_from_pose +from dimos.msgs.image import image_from_array +from dimos.msgs.time import header_now, time_from_nanoseconds from dimos.robot.unitree.go2.constants import CMD_VEL_TIMEOUT from dimos.robot.unitree.type.lidar import ( RawLidarMsg, pointcloud2_from_webrtc_lidar, ) from dimos.robot.unitree.type.lowstate import LowStateMsg -from dimos.robot.unitree.type.odometry import Odometry -from dimos.types.timestamped import Timestamped +from dimos.robot.unitree.type.odometry import RawOdometryMessage, pose_from_webrtc_odometry from dimos.utils.decorators.decorators import simple_mcache from dimos.utils.logging_config import setup_logger from dimos.utils.reactive import backpressure, callback_to_observable @@ -61,14 +59,6 @@ logger = setup_logger() -_T = TypeVar("_T", bound=Timestamped) - - -def time_is_now(x: _T) -> _T: - x.ts = time.time() - return x - - @dataclass class SerializableVideoFrame: """Pickleable wrapper for av.VideoFrame with all metadata""" @@ -284,34 +274,29 @@ def raw_lidar_stream(self) -> Observable[RawLidarMsg]: return backpressure(self.unitree_sub_stream(RTC_TOPIC["ULIDAR_ARRAY"])) @simple_mcache - def raw_odom_stream(self) -> Observable[Pose]: + def raw_odom_stream(self) -> Observable[RawOdometryMessage]: return backpressure(self.unitree_sub_stream(RTC_TOPIC["ROBOTODOM"])) @simple_mcache def lidar_stream(self) -> Observable[PointCloud2]: - return backpressure( - self.raw_lidar_stream().pipe( - ops.map(pointcloud2_from_webrtc_lidar), - ops.map(time_is_now), - # repair_stale_ts(), - ) - ) + def convert(raw: RawLidarMsg) -> PointCloud2: + return pointcloud2_from_webrtc_lidar(raw, stamp=time_from_nanoseconds(time.time_ns())) + + return backpressure(self.raw_lidar_stream().pipe(ops.map(convert))) @simple_mcache - def tf_stream(self) -> Observable[Transform]: - base_link = functools.partial(Transform.from_pose, "base_link") - return backpressure(self.odom_stream().pipe(ops.map(base_link))) + def tf_stream(self) -> Observable[TransformStamped]: + def convert(pose: PoseStamped) -> TransformStamped: + return transform_from_pose(pose, child_frame_id="base_link") + + return backpressure(self.odom_stream().pipe(ops.map(convert))) @simple_mcache - def odom_stream(self) -> Observable[Pose]: - return backpressure( - self.raw_odom_stream().pipe( - ops.map( - Odometry.from_msg, - ), - ops.map(time_is_now), - ) - ) + def odom_stream(self) -> Observable[PoseStamped]: + def convert(raw: RawOdometryMessage) -> PoseStamped: + return pose_from_webrtc_odometry(raw, header=header_now("world")) + + return backpressure(self.raw_odom_stream().pipe(ops.map(convert))) @simple_mcache def video_stream(self) -> Observable[Image]: @@ -319,14 +304,13 @@ def video_stream(self) -> Observable[Image]: self.raw_video_stream().pipe( ops.filter(lambda frame: frame is not None), ops.map( - lambda frame: Image.from_numpy( + lambda frame: image_from_array( # np.ascontiguousarray(frame.to_ndarray("rgb24")), frame.to_ndarray(format="rgb24"), # type: ignore[attr-defined] - format=ImageFormat.RGB, # Frame is RGB24, not BGR - frame_id="camera_optical", + encoding="rgb8", + header=header_now("camera_optical"), ), ), - ops.map(time_is_now), ) ) diff --git a/dimos/robot/unitree/dimsim_connection.py b/dimos/robot/unitree/dimsim_connection.py index 1b40e3ce614..e74941aef4a 100644 --- a/dimos/robot/unitree/dimsim_connection.py +++ b/dimos/robot/unitree/dimsim_connection.py @@ -17,20 +17,25 @@ import threading from typing import Any +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Twist, + Vector3, +) +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage from reactivex import Observable, Subject from dimos.core.global_config import GlobalConfig from dimos.core.transport import PubSubTransport from dimos.core.transport_factory import make_transport -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.camera_info import camera_info_from_fov +from dimos.msgs.geometry import transform_from_pose from dimos.simulation.dimsim.dimsim_process import DimSimProcess from dimos.utils.logging_config import setup_logger @@ -45,12 +50,12 @@ class DimSimConnection: - camera_info_static: CameraInfo = CameraInfo.from_fov( - fov_deg=_FOV_DEG, + camera_info_static: CameraInfo = camera_info_from_fov( + fov_degrees=_FOV_DEG, width=_WIDTH, height=_HEIGHT, axis="horizontal", - frame_id="camera_optical", + header=Header(frame_id="camera_optical", stamp=Time(sec=0, nanosec=0)), ) def __init__(self, global_config: GlobalConfig) -> None: @@ -139,41 +144,35 @@ def publish_request(self, topic: str, data: dict[str, Any]) -> dict[Any, Any]: def _handle_odom(self, msg: PoseStamped) -> None: self._first_odom.set() - self._tf_transport.publish(TFMessage(*_odom_to_tf(msg))) + self._tf_transport.publish(TFMessage(transforms=_odom_to_tf(msg))) -def _odom_to_tf(odom: PoseStamped) -> list[Transform]: - """Build transform chain from odometry pose. - - Transform tree: world -> base_link -> {camera_link -> camera_optical, lidar_link} - """ - camera_link = Transform( - translation=Vector3(0.3, 0.0, 0.0), # camera 30cm forward - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="base_link", - child_frame_id="camera_link", - ts=odom.ts, - ) - - camera_optical = Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(-0.5, 0.5, -0.5, 0.5), - frame_id="camera_link", - child_frame_id="camera_optical", - ts=odom.ts, - ) - - lidar_link = Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="base_link", - child_frame_id="lidar_link", - ts=odom.ts, - ) - +def _odom_to_tf(odom: PoseStamped) -> list[TransformStamped]: + """Build world → base_link → camera/lidar transforms at the source pose stamp.""" return [ - Transform.from_pose("base_link", odom), - camera_link, - camera_optical, - lidar_link, + transform_from_pose(odom, child_frame_id="base_link"), + TransformStamped( + header=Header(stamp=odom.header.stamp, frame_id="base_link"), + child_frame_id="camera_link", + transform=Transform( + translation=Vector3(x=0.3, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + TransformStamped( + header=Header(stamp=odom.header.stamp, frame_id="camera_link"), + child_frame_id="camera_optical", + transform=Transform( + rotation=Quaternion(x=-0.5, y=0.5, z=-0.5, w=0.5), + translation=Vector3(x=0.0, y=0.0, z=0.0), + ), + ), + TransformStamped( + header=Header(stamp=odom.header.stamp, frame_id="base_link"), + child_frame_id="lidar_link", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), ] diff --git a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py index 0d71bb043de..a990c04e0ed 100644 --- a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py @@ -49,6 +49,11 @@ from pathlib import Path from typing import Any, cast +from dimos_generated.dimos_msgs.msg import MotorCommandArray +from dimos_generated.geometry_msgs.msg import Twist +from dimos_generated.nav_msgs.msg import Odometry, Path as NavPath +from dimos_generated.sensor_msgs.msg import Imu, JointState + from dimos.control.components import HardwareComponent, HardwareType from dimos.control.coordinator import TaskConfig from dimos.control.tasks.g1_groot_wbc_task.g1_groot_wbc_task import ( @@ -69,11 +74,7 @@ from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.mapping.costmapper import CostMapper from dimos.mapping.pointclouds.occupancy import HeightCostConfig -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Path import Path as NavPath -from dimos.msgs.sensor_msgs.Imu import Imu -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.sensor_msgs.MotorCommandArray import MotorCommandArray +from dimos.msgs.geometry import pose_matrix, quaternion_from_matrix from dimos.navigation.go2.replanning_a_star.module import ReplanningAStarPlanner from dimos.navigation.movement_manager.movement_manager import MovementManager from dimos.robot.unitree.g1.config import G1 @@ -87,6 +88,7 @@ from dimos.robot.unitree.g1.teleop_ik import G1PinkPoseTargetSolver from dimos.simulation.scene_assets.spec import ScenePackage from dimos.utils.data import LfsPath +from dimos.visualization.rerun.message_helpers import navigation_archetype from dimos.visualization.rerun.scene_package import scene_package_static_entities from dimos.visualization.vis_module import vis_module @@ -377,7 +379,7 @@ def _g1_groot_rerun_blueprint() -> Any: def _g1_nav_path(path: NavPath) -> Any: - return path.to_rerun(z_offset=0.3) + return navigation_archetype(path, z_offset=0.3) # Mesh root: sim roots under the /odom transform; real hw under the LIO's @@ -417,20 +419,17 @@ def _g1_pelvis_to_mid360() -> Any: return _g1_pelvis_mid360_cache[0] -def _g1_real_odometry_root(odom: Any) -> Any: +def _g1_real_odometry_root(odom: Odometry) -> Any: """Robot-mesh root: pelvis pose from the LIO's mid360 odometry (rest offset).""" import numpy as np import rerun as rr - from dimos.msgs.geometry_msgs.Quaternion import Quaternion - - t_world_mid360 = np.eye(4) + t_world_mid360 = pose_matrix(odom.pose.pose) # The MID-360 is mounted upside down (the URDF doesn't carry the flip): # un-roll by Rx(pi) == diag(1, -1, -1). - t_world_mid360[:3, :3] = odom.orientation.to_rotation_matrix() @ np.diag([1.0, -1.0, -1.0]) - t_world_mid360[:3, 3] = (odom.x, odom.y, odom.z) + t_world_mid360[:3, :3] = t_world_mid360[:3, :3] @ np.diag([1.0, -1.0, -1.0]) t_world_pelvis = t_world_mid360 @ np.linalg.inv(_g1_pelvis_to_mid360()) - q = Quaternion.from_rotation_matrix(t_world_pelvis[:3, :3]) + q = quaternion_from_matrix(t_world_pelvis[:3, :3]) return rr.Transform3D( translation=t_world_pelvis[:3, 3].tolist(), rotation=rr.Quaternion(xyzw=[q.x, q.y, q.z, q.w]), diff --git a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_teleop.py b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_teleop.py index ee17461b365..62f2f991cab 100644 --- a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_teleop.py +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_teleop.py @@ -50,6 +50,10 @@ from datetime import datetime from typing import Any +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import Image, JointState +from dimos_generated.std_msgs.msg import String + from dimos.constants import STATE_DIR from dimos.core.coordination.blueprints import Blueprint, autoconnect from dimos.core.global_config import global_config @@ -59,10 +63,6 @@ from dimos.manipulation.manipulation_module import ManipulationModule from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig from dimos.memory.module import Recorder -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.std_msgs.String import String from dimos.robot.unitree.g1.blueprints.basic.unitree_g1_groot_wbc import ( _G1GrootCoordinator, _unitree_g1_groot_wbc_core, diff --git a/dimos/robot/unitree/g1/blueprints/primitive/test_unitree_g1_vis.py b/dimos/robot/unitree/g1/blueprints/primitive/test_unitree_g1_vis.py index 8386659f8a6..6da019b2164 100644 --- a/dimos/robot/unitree/g1/blueprints/primitive/test_unitree_g1_vis.py +++ b/dimos/robot/unitree/g1/blueprints/primitive/test_unitree_g1_vis.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, PointStamped from dimos_generated.std_msgs.msg import Header import numpy as np @@ -26,8 +28,11 @@ def test_global_map_uses_generated_xyz_and_height_colors(): - cloud = pointcloud_from_xyz(np.array([[1.0, 2.0, 0.0], [3.0, 4.0, 2.0]]), header=Header()) - result = _global_map_colors(type(cloud).decode(cloud.encode())) + cloud = pointcloud_from_xyz( + np.array([[1.0, 2.0, 0.0], [3.0, 4.0, 2.0]]), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + result = _global_map_colors(cdr_decode(cdr_encode(cloud), type(cloud))) np.testing.assert_array_equal( result.positions.as_arrow_array().to_pylist(), [[1, 2, 0], [3, 4, 2]] ) @@ -36,8 +41,10 @@ def test_global_map_uses_generated_xyz_and_height_colors(): @pytest.mark.parametrize("render", [_goal_colors, _waypoint_colors]) def test_generated_stamped_targets_keep_position_and_reject_nonfinite(render): - target = PointStamped(header=Header(frame_id="world"), point=Point(x=1, y=2, z=3)) - result = render(PointStamped.decode(target.encode())) + target = PointStamped( + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), point=Point(x=1, y=2, z=3) + ) + result = render(cdr_decode(cdr_encode(target), PointStamped)) np.testing.assert_allclose(result.positions.as_arrow_array().to_pylist(), [[1, 2, 3.3]]) target.point.x = float("nan") assert render(target) is None diff --git a/dimos/robot/unitree/g1/blueprints/primitive/unitree_g1_primitive_no_nav.py b/dimos/robot/unitree/g1/blueprints/primitive/unitree_g1_primitive_no_nav.py index 57eae88c635..100e6a6ab36 100644 --- a/dimos/robot/unitree/g1/blueprints/primitive/unitree_g1_primitive_no_nav.py +++ b/dimos/robot/unitree/g1/blueprints/primitive/unitree_g1_primitive_no_nav.py @@ -17,7 +17,17 @@ from typing import Any -from dimos_lcm.sensor_msgs import CameraInfo +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + PoseStamped, + Transform, + TransformStamped, + Twist, + Vector3, +) +from dimos_generated.nav_msgs.msg import Odometry, Path +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 +from dimos_generated.std_msgs.msg import Bool, Header from dimos.core.coordination.blueprints import autoconnect from dimos.core.global_config import global_config @@ -27,28 +37,17 @@ from dimos.hardware.sensors.camera.zed import compat as zed from dimos.mapping.costmapper import CostMapper from dimos.mapping.voxels.module import VoxelGridMapper -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.std_msgs.Bool import Bool +from dimos.msgs.geometry import quaternion_from_euler from dimos.navigation.experimental.frontier_exploration.wavefront_frontier_goal_selector import ( WavefrontFrontierExplorer, ) from dimos.robot.unitree.g1.g1_rerun import g1_costmap +from dimos.visualization.rerun.message_helpers import camera_pinhole from dimos.visualization.vis_module import vis_module -def _convert_camera_info(camera_info: Any) -> Any: - return camera_info.to_rerun( - image_topic="/world/color_image", - optical_frame="camera_optical", - ) +def _convert_camera_info(camera_info: CameraInfo) -> Any: + return camera_pinhole(camera_info) def _static_base_link(rr: Any) -> list[Any]: @@ -109,11 +108,13 @@ def _create_webcam() -> Webcam: _camera = ( autoconnect( CameraModule.blueprint( - transform=Transform( - translation=Vector3(0.05, 0.0, 0.6), # height of camera on G1 robot - rotation=Quaternion.from_euler(Vector3(0.0, 0.2, 0.0)), - frame_id="sensor", + transform=TransformStamped( + header=Header(frame_id="sensor", stamp=Time(sec=0, nanosec=0)), child_frame_id="camera_link", + transform=Transform( + translation=Vector3(x=0.05, z=0.6, y=0.0), + rotation=quaternion_from_euler(0.0, 0.2, 0.0), + ), ), hardware=_create_webcam, ), diff --git a/dimos/robot/unitree/g1/config.py b/dimos/robot/unitree/g1/config.py index 4a423fc835e..802dc529fa7 100644 --- a/dimos/robot/unitree/g1/config.py +++ b/dimos/robot/unitree/g1/config.py @@ -20,7 +20,7 @@ from pathlib import Path from typing import Any -from dimos_generated.geometry_msgs.msg import Point, Pose +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion @dataclass(frozen=True) @@ -41,6 +41,8 @@ class G1Config: width_clearance=0.6, internal_odom_offsets={ # Mid-360 lidar: 1.2 m above ground. - "mid360_link": Pose(position=Point(x=0.0, y=0.0, z=1.2)), + "mid360_link": Pose( + position=Point(x=0.0, y=0.0, z=1.2), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), }, ) diff --git a/dimos/robot/unitree/g1/debug/demo_arrow_control_cmd_vel.py b/dimos/robot/unitree/g1/debug/demo_arrow_control_cmd_vel.py index fab4326f17e..6b156f3d4a4 100644 --- a/dimos/robot/unitree/g1/debug/demo_arrow_control_cmd_vel.py +++ b/dimos/robot/unitree/g1/debug/demo_arrow_control_cmd_vel.py @@ -20,6 +20,7 @@ import traceback from typing import Any +from dimos_cdrgen.registry import encode as cdr_encode from dimos_generated.geometry_msgs.msg import Twist, Vector3 import lcm @@ -30,7 +31,7 @@ def publish_twist(lc: lcm.LCM, twist: Twist) -> None: - lc.publish(CMD_VEL_CHANNEL, twist.encode()) + lc.publish(CMD_VEL_CHANNEL, cdr_encode(twist)) def draw_ui(stdscr: Any, state_text: str = "Not connected") -> None: diff --git a/dimos/robot/unitree/g1/effectors/high_level/dds_sdk.py b/dimos/robot/unitree/g1/effectors/high_level/dds_sdk.py index 72849d2f9c4..859c5464c66 100644 --- a/dimos/robot/unitree/g1/effectors/high_level/dds_sdk.py +++ b/dimos/robot/unitree/g1/effectors/high_level/dds_sdk.py @@ -296,7 +296,7 @@ def move_velocity( self, x: float, y: float = 0.0, yaw: float = 0.0, duration: float = 0.0 ) -> str: """Move the robot at the given velocity for ``duration`` seconds.""" - twist = Twist(linear=Vector3(x=x, y=y), angular=Vector3(z=yaw)) + twist = Twist(linear=Vector3(x=x, y=y, z=0.0), angular=Vector3(z=yaw, x=0.0, y=0.0)) self.move(twist, duration=duration) return f"Started moving with velocity=({x}, {y}, {yaw}) for {duration} seconds" diff --git a/dimos/robot/unitree/g1/effectors/high_level/webrtc.py b/dimos/robot/unitree/g1/effectors/high_level/webrtc.py index 24acc0f20db..f70efa5e61e 100644 --- a/dimos/robot/unitree/g1/effectors/high_level/webrtc.py +++ b/dimos/robot/unitree/g1/effectors/high_level/webrtc.py @@ -115,7 +115,7 @@ def move_velocity( args = { "x": 0.5, "y": 0.0, "yaw": 0.0, "duration": 2.0 } move_velocity(**args) """ - twist = Twist(linear=Vector3(x=x, y=y), angular=Vector3(z=yaw)) + twist = Twist(linear=Vector3(x=x, y=y, z=0.0), angular=Vector3(z=yaw, x=0.0, y=0.0)) self.move(twist, duration=duration) return f"Started moving with velocity=({x}, {y}, {yaw}) for {duration} seconds" diff --git a/dimos/robot/unitree/g1/g1_tf_publisher.py b/dimos/robot/unitree/g1/g1_tf_publisher.py index aba566175ea..5c79f813f1a 100644 --- a/dimos/robot/unitree/g1/g1_tf_publisher.py +++ b/dimos/robot/unitree/g1/g1_tf_publisher.py @@ -26,6 +26,7 @@ import time from typing import Any +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Transform, TransformStamped, Vector3 from dimos_generated.std_msgs.msg import Header from dimos_generated.tf2_msgs.msg import TFMessage @@ -61,7 +62,7 @@ def _mount( yaw: float = 0.0, ) -> TransformStamped: return TransformStamped( - header=Header(frame_id=parent), + header=Header(frame_id=parent, stamp=Time(sec=0, nanosec=0)), child_frame_id=child, transform=Transform( translation=Vector3(x=xyz[0], y=xyz[1], z=xyz[2]), diff --git a/dimos/robot/unitree/g1/mujoco_sim.py b/dimos/robot/unitree/g1/mujoco_sim.py index b2e910e4747..6e50c6803f5 100644 --- a/dimos/robot/unitree/g1/mujoco_sim.py +++ b/dimos/robot/unitree/g1/mujoco_sim.py @@ -12,11 +12,23 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy import threading from threading import Thread import time from typing import Any +from dimos_generated.geometry_msgs.msg import ( + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Twist, + Vector3, +) +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage from pydantic import Field from reactivex.disposable import Disposable @@ -24,18 +36,10 @@ from dimos.core.core import rpc from dimos.core.module import ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.geometry import transform_from_pose +from dimos.msgs.time import time_from_nanoseconds from dimos.robot.unitree.g1.connection import G1ConnectionBase from dimos.robot.unitree.mujoco_connection import MujocoConnection -from dimos.robot.unitree.type.odometry import Odometry as SimOdometry from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -68,7 +72,7 @@ def start(self) -> None: self.connection.start() self.register_disposable(Disposable(self.cmd_vel.subscribe(self.move))) - self.register_disposable(self.connection.odom_stream().subscribe(self._publish_sim_odom)) + self.register_disposable(self.connection.odom_stream().subscribe(self._publish_tf)) self.register_disposable(self.connection.lidar_stream().subscribe(self.lidar.publish)) self.register_disposable(self.connection.video_stream().subscribe(self.color_image.publish)) @@ -89,52 +93,42 @@ def stop(self) -> None: def _publish_camera_info_loop(self) -> None: assert self.connection is not None - info = self.connection.camera_info_static while not self._stop_event.is_set(): + info = copy.deepcopy(self.connection.camera_info_static) + info.header.stamp = time_from_nanoseconds(time.time_ns()) self.camera_info.publish(info) self._stop_event.wait(1.0) def _publish_tf(self, msg: PoseStamped) -> None: self.odom.publish(msg) - - base_link = Transform.from_pose("base_link", msg) - - # Publish camera_link and camera_optical transforms - camera_link = Transform( - translation=Vector3(0.05, 0.0, 0.6), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="base_link", - child_frame_id="camera_link", - ts=time.time(), - ) - - camera_optical = Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(-0.5, 0.5, -0.5, 0.5), - frame_id="camera_link", - child_frame_id="camera_optical", - ts=time.time(), - ) - - map_to_world = Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="map", - child_frame_id="world", - ts=time.time(), - ) - - self.tf.publish(TFMessage(base_link, camera_link, camera_optical, map_to_world)) - - def _publish_sim_odom(self, msg: SimOdometry) -> None: - self._publish_tf( - PoseStamped( - ts=msg.ts, - frame_id=msg.frame_id, - position=msg.position, - orientation=msg.orientation, - ) - ) + transforms = [ + transform_from_pose(msg, child_frame_id="base_link"), + TransformStamped( + header=Header(stamp=msg.header.stamp, frame_id="base_link"), + child_frame_id="camera_link", + transform=Transform( + translation=Vector3(x=0.05, z=0.6, y=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + TransformStamped( + header=Header(stamp=msg.header.stamp, frame_id="camera_link"), + child_frame_id="camera_optical", + transform=Transform( + rotation=Quaternion(x=-0.5, y=0.5, z=-0.5, w=0.5), + translation=Vector3(x=0.0, y=0.0, z=0.0), + ), + ), + TransformStamped( + header=Header(stamp=msg.header.stamp, frame_id="map"), + child_frame_id="world", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + ] + self.tf.publish(TFMessage(transforms=transforms)) @rpc def move(self, twist: Twist, duration: float = 0.0) -> None: diff --git a/dimos/robot/unitree/g1/skill_container.py b/dimos/robot/unitree/g1/skill_container.py index 68f8327fe5e..01ef92f1adc 100644 --- a/dimos/robot/unitree/g1/skill_container.py +++ b/dimos/robot/unitree/g1/skill_container.py @@ -89,7 +89,7 @@ def move(self, x: float, y: float = 0.0, yaw: float = 0.0, duration: float = 0.0 duration: How long to move (seconds) """ - twist = Twist(linear=Vector3(x=x, y=y), angular=Vector3(z=yaw)) + twist = Twist(linear=Vector3(x=x, y=y, z=0.0), angular=Vector3(z=yaw, x=0.0, y=0.0)) self._connection.move(twist, duration=duration) return f"Started moving with velocity=({x}, {y}, {yaw}) for {duration} seconds" diff --git a/dimos/robot/unitree/g1/test_g1_manipulation_integration.py b/dimos/robot/unitree/g1/test_g1_manipulation_integration.py index 60ab74a5705..0abc4e244fc 100644 --- a/dimos/robot/unitree/g1/test_g1_manipulation_integration.py +++ b/dimos/robot/unitree/g1/test_g1_manipulation_integration.py @@ -14,6 +14,9 @@ """Self-hosted integration coverage for G1 full-body manipulation planning.""" +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest pytest.importorskip("roboplan.core") @@ -44,7 +47,10 @@ def test_full_body_state_supports_collision_checked_arm_only_plan() -> None: world.finalize() full_body_state = JointState( name=list(config.joint_names), - position=[0.0] * len(config.joint_names), + position=np.asarray([0.0] * len(config.joint_names), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) world.sync_from_joint_state(full_body_state) @@ -53,10 +59,21 @@ def test_full_body_state_supports_collision_checked_arm_only_plan() -> None: selection = world._planning_groups.select(("left_arm", "right_arm")) arm_joint_names = [*G1_LEFT_ARM_JOINTS, *G1_RIGHT_ARM_JOINTS] - start = JointState(name=arm_joint_names, position=[0.0] * len(arm_joint_names)) + start = JointState( + name=arm_joint_names, + position=np.asarray([0.0] * len(arm_joint_names), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) goal = JointState( name=arm_joint_names, - position=[*G1_READY_JOINTS["left_arm"], *G1_READY_JOINTS["right_arm"]], + position=np.array( + [*G1_READY_JOINTS["left_arm"], *G1_READY_JOINTS["right_arm"]], dtype=np.float64 + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ) result = roboplan_planner_module.RoboPlanPlanner( diff --git a/dimos/robot/unitree/g1/test_g1_teleop.py b/dimos/robot/unitree/g1/test_g1_teleop.py index 1dc65f46338..68e6f382b77 100644 --- a/dimos/robot/unitree/g1/test_g1_teleop.py +++ b/dimos/robot/unitree/g1/test_g1_teleop.py @@ -17,6 +17,9 @@ from typing import Any, cast import xml.etree.ElementTree as ET +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import Image, JointState +from dimos_generated.std_msgs.msg import Header, String import numpy as np import pytest @@ -32,10 +35,7 @@ from dimos.imitation.collection.episode import EpisodeStatus from dimos.manipulation.planning.spec.validation import prepare_robot_model from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.std_msgs.String import String +from dimos.msgs.time import time_from_seconds from dimos.robot.unitree.g1.blueprints.basic.unitree_g1_groot_wbc import ( _G1_TELEOP_MODEL, _G1GrootCoordinator, @@ -160,7 +160,20 @@ async def test_g1_collection_accepts_blueprint_config_and_records_without_world_ } assert recorder.config.record_tf is False for name in ports: - assert await recorder._resolve_pose(name, JointState(ts=1.0), 1.0) is None + assert ( + await recorder._resolve_pose( + name, + JointState( + header=Header(stamp=time_from_seconds(1.0), frame_id=""), + name=[], + position=np.array([], dtype=np.float64), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ), + 1.0, + ) + is None + ) finally: recorder.stop() diff --git a/dimos/robot/unitree/g1/test_g1_tf_publisher.py b/dimos/robot/unitree/g1/test_g1_tf_publisher.py index df689179879..04c45225155 100644 --- a/dimos/robot/unitree/g1/test_g1_tf_publisher.py +++ b/dimos/robot/unitree/g1/test_g1_tf_publisher.py @@ -16,6 +16,7 @@ import math +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Transform, TransformStamped, Vector3 from dimos_generated.std_msgs.msg import Header @@ -71,10 +72,11 @@ def test_mid360_frame_is_upside_down() -> None: def test_flipped_level_sensor_yields_level_base() -> None: """A standing robot reports a flipped sensor pose. base_link must come out level, below it.""" live = TransformStamped( - header=Header(frame_id="world"), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), child_frame_id="mid360_link", transform=Transform( - translation=Vector3(z=LIDAR_HEIGHT), rotation=torso_to_mid360().transform.rotation + translation=Vector3(z=LIDAR_HEIGHT, x=0.0, y=0.0), + rotation=torso_to_mid360().transform.rotation, ), ) base = _buffer(live=live).get("world", "base_link") diff --git a/dimos/robot/unitree/g1/test_wholebody_connection.py b/dimos/robot/unitree/g1/test_wholebody_connection.py index 4a6067981e9..9babd8fc421 100644 --- a/dimos/robot/unitree/g1/test_wholebody_connection.py +++ b/dimos/robot/unitree/g1/test_wholebody_connection.py @@ -15,11 +15,14 @@ from __future__ import annotations from collections.abc import Iterator +from dataclasses import asdict from types import SimpleNamespace +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.dimos_msgs.msg import MotorCommandArray from dimos_generated.std_msgs.msg import Header +import numpy as np from pydantic import ValidationError import pytest @@ -72,11 +75,12 @@ def _wire(connection, soft_start_seconds): def _command(): return MotorCommandArray( - q=[1.0] * _NUM_MOTORS, - dq=[0.0] * _NUM_MOTORS, - kp=[100.0] * _NUM_MOTORS, - kd=[5.0] * _NUM_MOTORS, - tau=[8.0] * _NUM_MOTORS, + q=np.asarray([1.0] * _NUM_MOTORS, dtype=np.float64), + dq=np.asarray([0.0] * _NUM_MOTORS, dtype=np.float64), + kp=np.asarray([100.0] * _NUM_MOTORS, dtype=np.float64), + kd=np.asarray([5.0] * _NUM_MOTORS, dtype=np.float64), + tau=np.asarray([8.0] * _NUM_MOTORS, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) @@ -121,7 +125,16 @@ def test_soft_start_disabled_passes_through(connection: G1WholeBodyConnection): def test_wrong_joint_count_is_dropped(connection: G1WholeBodyConnection): publisher = _wire(connection, soft_start_seconds=0.0) - connection._on_motor_command(MotorCommandArray(q=[0.0] * 5)) + connection._on_motor_command( + MotorCommandArray( + q=np.asarray([0.0] * 5, dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + dq=np.array([], dtype=np.float64), + kp=np.array([], dtype=np.float64), + kd=np.array([], dtype=np.float64), + tau=np.array([], dtype=np.float64), + ) + ) assert publisher.frames == [] @@ -156,7 +169,9 @@ def test_generated_feedback_has_one_exact_header_and_ros_quaternion(connection, assert imu[0].angular_velocity.x == 1.0 assert imu[0].linear_acceleration.z == 6.0 for message in [joints[0], imu[0]]: - assert type(message).decode(message.encode()) == message + np.testing.assert_equal( + asdict(cdr_decode(cdr_encode(message), type(message))), asdict(message) + ) header.stamp.nanosec = 0 assert joints[0].header.stamp.nanosec == 123456789 diff --git a/dimos/robot/unitree/g1/wholebody_connection.py b/dimos/robot/unitree/g1/wholebody_connection.py index 2139ae6c3eb..f10fe2e8be4 100644 --- a/dimos/robot/unitree/g1/wholebody_connection.py +++ b/dimos/robot/unitree/g1/wholebody_connection.py @@ -21,6 +21,7 @@ from __future__ import annotations +from copy import deepcopy from dataclasses import dataclass import threading from threading import Thread @@ -31,6 +32,7 @@ from dimos_generated.geometry_msgs.msg import Quaternion, Vector3 from dimos_generated.sensor_msgs.msg import Imu, JointState from dimos_generated.std_msgs.msg import Header +import numpy as np from pydantic import Field from reactivex.disposable import Disposable @@ -76,7 +78,10 @@ def _imu_from_unitree_wxyz( orientation=Quaternion(x=x, y=y, z=z, w=_w), angular_velocity=Vector3(x=gyroscope[0], y=gyroscope[1], z=gyroscope[2]), linear_acceleration=Vector3(x=accelerometer[0], y=accelerometer[1], z=accelerometer[2]), - header=header, + header=deepcopy(header), + orientation_covariance=np.zeros(9, dtype=np.float64), + angular_velocity_covariance=np.zeros(9, dtype=np.float64), + linear_acceleration_covariance=np.zeros(9, dtype=np.float64), ) @@ -337,11 +342,11 @@ def _snapshot_motor_imu(self) -> G1LowStateSnapshot | None: def _publish_motor_state_and_imu(self, header: Header, sample: G1LowStateSnapshot) -> None: self.motor_states.publish( JointState( - header=header, - name=G1_JOINT_NAMES, - position=sample.positions, - velocity=sample.velocities, - effort=sample.efforts, + header=deepcopy(header), + name=list(G1_JOINT_NAMES), + position=np.asarray(sample.positions, dtype=np.float64), + velocity=np.asarray(sample.velocities, dtype=np.float64), + effort=np.asarray(sample.efforts, dtype=np.float64), ) ) # Unitree reports quaternions as (w,x,y,z); Imu/Quaternion stores (x,y,z,w). diff --git a/dimos/robot/unitree/go2/base.py b/dimos/robot/unitree/go2/base.py index b4280e35bc4..035d51522de 100644 --- a/dimos/robot/unitree/go2/base.py +++ b/dimos/robot/unitree/go2/base.py @@ -30,13 +30,13 @@ from dimos_generated.nav_msgs.msg import Odometry from dimos_generated.sensor_msgs.msg import CameraInfo, PointCloud2 from dimos_generated.std_msgs.msg import String +from dimos_generated.tf2_msgs.msg import TFMessage from pydantic import Field, field_validator from reactivex.disposable import Disposable from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import In, Out -from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.protocol.tf.static_tf_publisher import StaticTfPublisherConfig from dimos.robot.unitree.go2.camera_calibration import front_camera_calibration from dimos.robot.unitree.go2.go2_mid360_static_transforms import ( diff --git a/dimos/robot/unitree/go2/blueprints/basic/go2_sim.py b/dimos/robot/unitree/go2/blueprints/basic/go2_sim.py index 22b271a06ea..46bafc69f2e 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/go2_sim.py +++ b/dimos/robot/unitree/go2/blueprints/basic/go2_sim.py @@ -29,11 +29,10 @@ from dimos.visualization.vis_module import vis_module if TYPE_CHECKING: + from dimos_generated.dimos_msgs.msg import LineSegments3D from rerun._baseclasses import Archetype from rerun.blueprint import Blueprint - from dimos.msgs.nav_msgs.LineSegments3D import LineSegments3D - def _scene_lines(scene: LineSegments3D) -> Archetype: return scene.to_rerun(radii=0.01) diff --git a/dimos/robot/unitree/go2/blueprints/basic/test_unitree_go2_joystick_record.py b/dimos/robot/unitree/go2/blueprints/basic/test_unitree_go2_joystick_record.py index 995421c7289..7103af8035b 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/test_unitree_go2_joystick_record.py +++ b/dimos/robot/unitree/go2/blueprints/basic/test_unitree_go2_joystick_record.py @@ -12,7 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos.msgs.sensor_msgs.Joy import Joy +from dimos_generated.sensor_msgs.msg import Joy +import numpy as np + +from dimos.msgs.time import header_now from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_joystick_record import ( record_rerun_config, ) @@ -20,8 +23,22 @@ def test_keys_text_logs_each_change_once() -> None: keys = record_rerun_config["visual_override"]["world/joystick"] - w_shift = Joy(buttons=[1, 0, 0, 0, 0, 0, 1, 0, 0]) - lines = [keys(w_shift), keys(w_shift), keys(Joy(buttons=[0] * 9))] + w_shift = Joy( + header=header_now(), + axes=np.array([], dtype=np.float32), + buttons=np.array([1, 0, 0, 0, 0, 0, 1, 0, 0], dtype=np.int32), + ) + lines = [ + keys(w_shift), + keys(w_shift), + keys( + Joy( + header=header_now(), + axes=np.array([], dtype=np.float32), + buttons=np.zeros(9, dtype=np.int32), + ) + ), + ] assert [ [entry.text.as_arrow_array().to_pylist()[0] for _, entry in line] for line in lines ] == [ diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py index 2a0cbe8830c..ea7752c1048 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py @@ -29,17 +29,14 @@ if TYPE_CHECKING: from types import ModuleType + from dimos_generated.dimos_msgs.msg import LineSegments3D + from dimos_generated.geometry_msgs.msg import PoseArray, PoseStamped, Twist + from dimos_generated.nav_msgs.msg import OccupancyGrid + from dimos_generated.sensor_msgs.msg import CameraInfo, PointCloud2 + from dimos_generated.std_msgs.msg import Float32 from rerun._baseclasses import Archetype from rerun.blueprint import Blueprint - from dimos.msgs.geometry_msgs.PoseArray import PoseArray - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - from dimos.msgs.geometry_msgs.Twist import Twist - from dimos.msgs.nav_msgs.LineSegments3D import LineSegments3D - from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid - from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 - from dimos.msgs.std_msgs.Float32 import Float32 from dimos.visualization.rerun.bridge import RerunData, RerunMulti diff --git a/dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py b/dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py index 40ef01fe3c6..51998f86432 100644 --- a/dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py +++ b/dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py @@ -15,6 +15,10 @@ from pathlib import Path +from dimos_generated.geometry_msgs.msg import Pose, PoseStamped +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 +from dimos_generated.vision_msgs.msg import Detection3DArray + from dimos.core.coordination.blueprints import autoconnect from dimos.core.stream import In from dimos.core.transport import LCMTransport @@ -22,12 +26,6 @@ from dimos.mapping.relocalization.go2.module import Go2Relocalization from dimos.mapping.voxels.module import VoxelGridMapper from dimos.memory.module import Recorder, RecorderConfig, pose_setter_for -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray from dimos.navigation.experimental.frontier_exploration.wavefront_frontier_goal_selector import ( WavefrontFrontierExplorer, ) @@ -66,7 +64,7 @@ class Go2Memory(Recorder): @pose_setter_for("odom") async def _odom_pose(self, msg: PoseStamped) -> Pose | None: - self._last_odom_pose = msg + self._last_odom_pose = msg.pose return self._last_odom_pose @pose_setter_for("lidar") diff --git a/dimos/robot/unitree/go2/blueprints/teleop_webxr.py b/dimos/robot/unitree/go2/blueprints/teleop_webxr.py index ffc23319a65..cec5824f472 100644 --- a/dimos/robot/unitree/go2/blueprints/teleop_webxr.py +++ b/dimos/robot/unitree/go2/blueprints/teleop_webxr.py @@ -15,9 +15,10 @@ """WebXR velocity teleoperation and headset video for Go2.""" +from dimos_generated.geometry_msgs.msg import Twist + from dimos.core.coordination.blueprints import autoconnect from dimos.core.transport import LCMTransport -from dimos.msgs.geometry_msgs.Twist import Twist from dimos.robot.unitree.go2.connection import GO2Connection from dimos.teleop.webxr.extensions import Go2TeleopModule diff --git a/dimos/robot/unitree/go2/blueprints/test_tf_topology.py b/dimos/robot/unitree/go2/blueprints/test_tf_topology.py index ddec0be5a42..ec2253e72b0 100644 --- a/dimos/robot/unitree/go2/blueprints/test_tf_topology.py +++ b/dimos/robot/unitree/go2/blueprints/test_tf_topology.py @@ -22,12 +22,14 @@ mid360_link so it never writes the frame PointLio owns. """ +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.std_msgs.msg import Header import pytest from dimos.core.coordination.blueprints import Blueprint from dimos.hardware.sensors.lidar.pointlio.module import PointLio from dimos.mapping.relocalization.module import RelocalizationModule -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_mid360_record import ( unitree_go2_mid360_record, ) @@ -57,7 +59,12 @@ def _tf_children_by_publisher(blueprint: Blueprint) -> dict[str, set[str]]: """Child frames each tf publisher the blueprint actually enables will write.""" - odom = PoseStamped(ts=1.0, frame_id="go2_odom") + odom = PoseStamped( + header=Header(frame_id="go2_odom", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + ) children: dict[str, set[str]] = {} for atom in blueprint.active_blueprints: if atom.module is GO2Connection and atom.kwargs.get("publish_tf", True): diff --git a/dimos/robot/unitree/go2/camera_calibration.py b/dimos/robot/unitree/go2/camera_calibration.py index 7a79c10b353..73371f173d1 100644 --- a/dimos/robot/unitree/go2/camera_calibration.py +++ b/dimos/robot/unitree/go2/camera_calibration.py @@ -16,6 +16,7 @@ from importlib import resources +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import CameraInfo from dimos_generated.std_msgs.msg import Header @@ -25,4 +26,6 @@ def front_camera_calibration() -> CameraInfo: resource = resources.files("dimos.robot.unitree.go2").joinpath("front_camera_720.yaml") with resources.as_file(resource) as path: - return camera_info_from_yaml(path, header=Header(frame_id="camera_optical")) + return camera_info_from_yaml( + path, header=Header(frame_id="camera_optical", stamp=Time(sec=0, nanosec=0)) + ) diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index 9ef0c858425..7aad7bbba58 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -19,6 +19,7 @@ import time from typing import Any, Protocol +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( PoseStamped, Quaternion, @@ -118,14 +119,20 @@ def _prefixed(prefix: str | None, name: str) -> str: # TODO we need a standardized way to specify this for all cameras in dimos BASE_TO_OPTICAL = compose_transforms( TransformStamped( - header=Header(frame_id="base_link"), + header=Header(frame_id="base_link", stamp=Time(sec=0, nanosec=0)), child_frame_id="camera_link", - transform=Transform(translation=Vector3(x=0.3)), + transform=Transform( + translation=Vector3(x=0.3, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), TransformStamped( - header=Header(frame_id="camera_link"), + header=Header(frame_id="camera_link", stamp=Time(sec=0, nanosec=0)), child_frame_id="camera_optical", - transform=Transform(rotation=Quaternion(x=-0.5, y=0.5, z=-0.5, w=0.5)), + transform=Transform( + rotation=Quaternion(x=-0.5, y=0.5, z=-0.5, w=0.5), + translation=Vector3(x=0.0, y=0.0, z=0.0), + ), ), ) @@ -377,7 +384,7 @@ def __init__(self, **kwargs: Any) -> None: if self.config.frame_id_prefix and self.camera_info_static.header.frame_id: # Copy so the class-level default is not mutated. - self.camera_info_static = copy.copy(self.camera_info_static) + self.camera_info_static = copy.deepcopy(self.camera_info_static) self.camera_info_static.header.frame_id = _prefixed( self.config.frame_id_prefix, self.camera_info_static.header.frame_id ) @@ -453,17 +460,23 @@ def _odom_to_tf(cls, odom: PoseStamped, prefix: str = "") -> list[TransformStamp TransformStamped( header=Header(stamp=odom.header.stamp, frame_id=_prefixed(prefix, "base_link")), child_frame_id=_prefixed(prefix, "camera_link"), - transform=Transform(translation=Vector3(x=0.3)), + transform=Transform( + translation=Vector3(x=0.3, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), TransformStamped( header=Header(stamp=odom.header.stamp, frame_id=_prefixed(prefix, "camera_link")), child_frame_id=_prefixed(prefix, "camera_optical"), - transform=Transform(rotation=Quaternion(x=-0.5, y=0.5, z=-0.5, w=0.5)), + transform=Transform( + rotation=Quaternion(x=-0.5, y=0.5, z=-0.5, w=0.5), + translation=Vector3(x=0.0, y=0.0, z=0.0), + ), ), ] def _publish_tf(self, msg: PoseStamped) -> None: - message = copy.copy(msg) + message = copy.deepcopy(msg) message.header.frame_id = self.config.odom_frame_id if self.config.publish_tf: transforms = self._odom_to_tf(message, prefix=self.config.frame_id_prefix or "") @@ -473,7 +486,7 @@ def _publish_tf(self, msg: PoseStamped) -> None: def publish_camera_info(self) -> None: while True: - message = copy.copy(self.camera_info_static) + message = copy.deepcopy(self.camera_info_static) message.header.stamp = time_from_nanoseconds(time.time_ns()) self.camera_info.publish(message) time.sleep(1.0) diff --git a/dimos/robot/unitree/go2/dds/cli/render.py b/dimos/robot/unitree/go2/dds/cli/render.py index 0eae19fbb10..88572ed88fc 100644 --- a/dimos/robot/unitree/go2/dds/cli/render.py +++ b/dimos/robot/unitree/go2/dds/cli/render.py @@ -34,6 +34,7 @@ import subprocess from typing import TYPE_CHECKING, Any, cast +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, Pose, @@ -145,7 +146,9 @@ def step(state: Any, obs: Observation[Any]) -> tuple[Any, TwistStamped]: vel = vel + a_world * (obs.ts - prev) twist = TwistStamped( header=Header(frame_id=WORLD, stamp=time_from_seconds(obs.ts)), - twist=Twist(linear=Vector3(x=vel[0], y=vel[1], z=vel[2])), + twist=Twist( + linear=Vector3(x=vel[0], y=vel[1], z=vel[2]), angular=Vector3(x=0.0, y=0.0, z=0.0) + ), ) return (vel, obs.ts), twist @@ -160,14 +163,17 @@ def integrate_position(state: Any, obs: Observation[Any]) -> tuple[Any, PoseStam pos = pos + np.array([v.x, v.y, v.z]) * (obs.ts - prev) pose = PoseStamped( header=Header(frame_id=WORLD, stamp=time_from_seconds(obs.ts)), - pose=Pose(position=Point(x=pos[0], y=pos[1], z=pos[2]), orientation=Quaternion(w=1)), + pose=Pose( + position=Point(x=pos[0], y=pos[1], z=pos[2]), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), ) return (pos, obs.ts), pose def accumulate_path(upstream: Iterator[Observation[PoseStamped]]) -> Iterator[Observation[Path]]: """transform: yield the growing nav_msgs/Path as each pose streams in.""" - path = Path(header=Header(frame_id=WORLD)) + path = Path(header=Header(frame_id=WORLD, stamp=Time(sec=0, nanosec=0)), poses=[]) for obs in upstream: path = Path(header=obs.data.header, poses=[*path.poses, obs.data]) yield obs.derive(data=path) diff --git a/dimos/robot/unitree/go2/dds/cli/test_render_generated.py b/dimos/robot/unitree/go2/dds/cli/test_render_generated.py index 864a5565d67..d6472fdb346 100644 --- a/dimos/robot/unitree/go2/dds/cli/test_render_generated.py +++ b/dimos/robot/unitree/go2/dds/cli/test_render_generated.py @@ -18,6 +18,8 @@ import math from types import SimpleNamespace +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Vector3 from dimos_generated.sensor_msgs.msg import Imu from dimos_generated.std_msgs.msg import Header @@ -39,8 +41,13 @@ def test_world_acceleration_and_integrators_keep_world_axes_and_clock(): imu = Imu( - orientation=Quaternion(z=math.sqrt(0.5), w=math.sqrt(0.5)), - linear_acceleration=Vector3(x=1, z=9.81), + orientation=Quaternion(z=math.sqrt(0.5), w=math.sqrt(0.5), x=0.0, y=0.0), + linear_acceleration=Vector3(x=1, z=9.81, y=0.0), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + orientation_covariance=np.zeros(9, dtype=np.float64), + angular_velocity=Vector3(x=0.0, y=0.0, z=0.0), + angular_velocity_covariance=np.zeros(9, dtype=np.float64), + linear_acceleration_covariance=np.zeros(9, dtype=np.float64), ) observations = [Observation(ts=ts, _data=imu) for ts in (2.0, 3.0)] np.testing.assert_allclose(world_accel(observations[0]), [0, 1, 9.81], atol=1e-12) @@ -65,7 +72,9 @@ def test_path_accumulation_keeps_previous_snapshot_and_exact_pose_headers(): poses = [ PoseStamped( header=Header(frame_id="world", stamp=time_from_seconds(ts)), - pose=Pose(position=Point(x=ts), orientation=Quaternion(w=1)), + pose=Pose( + position=Point(x=ts, y=0.0, z=0.0), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), ) for ts in (1.0, 2.0) ] @@ -76,7 +85,7 @@ def test_path_accumulation_keeps_previous_snapshot_and_exact_pose_headers(): ) assert [len(obs.data.poses) for obs in result] == [1, 2] assert result[1].data.header == poses[1].header - assert result[0].data.poses[0].encode() == poses[0].encode() + assert cdr_encode(result[0].data.poses[0]) == cdr_encode(poses[0]) def test_sportmode_wxyz_is_converted_to_generated_xyzw(): @@ -87,14 +96,14 @@ def test_sportmode_wxyz_is_converted_to_generated_xyzw(): assert pose.header.stamp == time_from_seconds(4.25) q = pose.pose.orientation assert (q.x, q.y, q.z, q.w) == (-0.5, 0.5, -0.5, 0.5) - assert PoseStamped.decode(pose.encode()).pose.position.y == 2.0 + assert cdr_decode(cdr_encode(pose), PoseStamped).pose.position.y == 2.0 def test_measured_lidar_mount_retains_coordinates_and_cloud_stamp(): cloud = pointcloud_from_xyz( np.array([[1.0, 0.0, 0.0]]), header=Header(frame_id="lidar", stamp=time_from_seconds(12.25)) ) - original = cloud.encode() + original = cdr_encode(cloud) transformed = transform_cloud(cloud, LIDAR_TO_BASE) matrix = transform_matrix(LIDAR_TO_BASE.transform) np.testing.assert_allclose(matrix[:3, :3], EXT_R, atol=1e-6) @@ -103,4 +112,4 @@ def test_measured_lidar_mount_retains_coordinates_and_cloud_stamp(): ) assert transformed.header.stamp == cloud.header.stamp assert transformed.header.frame_id == "base_link" - assert cloud.encode() == original + assert cdr_encode(cloud) == original diff --git a/dimos/robot/unitree/go2/dds/extrinsics.py b/dimos/robot/unitree/go2/dds/extrinsics.py index 10d5002c601..122696bf791 100644 --- a/dimos/robot/unitree/go2/dds/extrinsics.py +++ b/dimos/robot/unitree/go2/dds/extrinsics.py @@ -24,6 +24,7 @@ [yaw -57.9, pitch -8.1, roll -167.3] (agrees to a few degrees). """ +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 from dimos_generated.std_msgs.msg import Header import numpy as np @@ -50,7 +51,7 @@ # Normalize the measured six-decimal rotation with the same SciPy conversion # as the previous numeric helper, then publish generated frame-labelled values. LIDAR_TO_BASE = TransformStamped( - header=Header(frame_id="base_link"), + header=Header(frame_id="base_link", stamp=Time(sec=0, nanosec=0)), child_frame_id="lidar", transform=Transform( translation=Vector3(x=EXT_T[0], y=EXT_T[1], z=EXT_T[2]), @@ -58,7 +59,7 @@ ), ) BASE_TO_CAMERA = TransformStamped( - header=Header(frame_id="base_link"), + header=Header(frame_id="base_link", stamp=Time(sec=0, nanosec=0)), child_frame_id="camera_optical", transform=Transform( translation=Vector3(x=CAM_T[0], y=CAM_T[1], z=CAM_T[2]), diff --git a/dimos/robot/unitree/go2/dds/ros.py b/dimos/robot/unitree/go2/dds/ros.py index 4d7c263a2e8..b9ea5d53d90 100644 --- a/dimos/robot/unitree/go2/dds/ros.py +++ b/dimos/robot/unitree/go2/dds/ros.py @@ -21,6 +21,7 @@ from dataclasses import dataclass +from dimos_cdrgen.registry import decode as cdr_decode from dimos_generated.foxglove_msgs.msg import CompressedVideo from dimos_generated.geometry_msgs.msg import Quaternion from dimos_generated.nav_msgs.msg import Odometry @@ -32,26 +33,26 @@ def decode_imu(buf: bytes) -> Imu: - message = Imu.decode(buf) + message = cdr_decode(buf, Imu) firmware = message.orientation message.orientation = Quaternion(x=firmware.y, y=firmware.z, z=firmware.w, w=firmware.x) return message def decode_odometry(buf: bytes) -> Odometry: - return Odometry.decode(buf) + return cdr_decode(buf, Odometry) def decode_pointcloud2(buf: bytes) -> PointCloud2: - return PointCloud2.decode(buf) + return cdr_decode(buf, PointCloud2) def decode_compressed_image(buf: bytes) -> CompressedImage: - return CompressedImage.decode(buf) + return cdr_decode(buf, CompressedImage) def decode_compressed_video(buf: bytes) -> CompressedVideo: - return CompressedVideo.decode(buf) + return cdr_decode(buf, CompressedVideo) # unitree_go/HeightMap (rt/utlidar/height_map_array) diff --git a/dimos/robot/unitree/go2/dds/rust/Cargo.toml b/dimos/robot/unitree/go2/dds/rust/Cargo.toml index 629a0033ea9..7f835efb11a 100644 --- a/dimos/robot/unitree/go2/dds/rust/Cargo.toml +++ b/dimos/robot/unitree/go2/dds/rust/Cargo.toml @@ -17,7 +17,7 @@ path = "src/bin/go2_dds.rs" # Our forks, vendored: the robot has no GitHub access. See vendor/*/Cargo.toml. cyclonedds-rs = { path = "vendor/cyclonedds-rs", default-features = false } dimos-module = { path = "../../../../../../native/rust/dimos-module" } -dimos-generated-messages = { path = "../../../../../../packages/dimos-generated/src/dimos_generated_schemas/package/rust" } +dimos-generated-messages = { path = "../../../../../../packages/dimos-generated/build/message-source/rust" } serde = { version = "1", features = ["derive"] } serde-big-array = "0.5" tokio = { version = "1", features = ["rt-multi-thread", "macros"] } diff --git a/dimos/robot/unitree/go2/dds/test_blueprints.py b/dimos/robot/unitree/go2/dds/test_blueprints.py index f061e598180..12f783dda59 100644 --- a/dimos/robot/unitree/go2/dds/test_blueprints.py +++ b/dimos/robot/unitree/go2/dds/test_blueprints.py @@ -16,10 +16,12 @@ from typing import Any +from dimos_generated.sensor_msgs.msg import CompressedImage +import numpy as np import rerun as rr -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.CompressedImage import CompressedImage +from dimos.msgs.camera_info import camera_info_from_intrinsics +from dimos.msgs.time import header_now from dimos.robot.unitree.go2.dds.blueprints import CAMERA_ENTITY, go2_dds_nav_viewer @@ -32,9 +34,15 @@ def test_both_encodings_land_on_the_pane() -> None: kwargs = _rerun_kwargs(go2_dds_nav_viewer) overrides = kwargs["visual_override"] assert kwargs["blueprint"]().root_container is not None - info = CameraInfo(width=640, height=480, frame_id="camera_optical") + info = camera_info_from_intrinsics( + 500.0, 500.0, 320.0, 240.0, 640, 480, header=header_now("camera_optical") + ) [(pinhole_path, _)] = overrides["world/camera_info"](info) - jpeg = CompressedImage(data=b"\xff\xd8", format="jpeg", frame_id="camera_optical") + jpeg = CompressedImage( + data=np.array([255, 216], dtype=np.uint8), + format="jpeg", + header=header_now("camera_optical"), + ) [(image_path, archetype)] = overrides["world/image"](jpeg) assert pinhole_path == image_path == CAMERA_ENTITY assert isinstance(archetype, rr.EncodedImage) diff --git a/dimos/robot/unitree/go2/dds/test_generated_codecs.py b/dimos/robot/unitree/go2/dds/test_generated_codecs.py index 44930920549..2e39b2b1993 100644 --- a/dimos/robot/unitree/go2/dds/test_generated_codecs.py +++ b/dimos/robot/unitree/go2/dds/test_generated_codecs.py @@ -14,9 +14,18 @@ import struct +from dimos_cdrgen.registry import encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.foxglove_msgs.msg import CompressedVideo -from dimos_generated.geometry_msgs.msg import Quaternion +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseWithCovariance, + Quaternion, + Twist, + TwistWithCovariance, + Vector3, +) from dimos_generated.nav_msgs.msg import Odometry from dimos_generated.sensor_msgs.msg import CompressedImage, Imu, PointCloud2, PointField from dimos_generated.std_msgs.msg import Header @@ -37,14 +46,25 @@ def test_odometry_preserves_independent_cdr_covariances_and_exact_stamp(little_e source = Odometry( header=Header(stamp=Time(sec=42, nanosec=123456789), frame_id="odom"), child_frame_id="base_link", + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), + twist=TwistWithCovariance( + twist=Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + covariance=np.zeros(36, dtype=np.float64), + ), ) source.pose.pose.position.x = 1.25 - source.pose.covariance = list(range(36)) - source.twist.covariance = list(range(36, 72)) + source.pose.covariance = np.arange(36, dtype=np.float64) + source.twist.covariance = np.arange(36, 72, dtype=np.float64) reference = get_typestore(Stores.ROS2_JAZZY) - independent = reference.deserialize_cdr(source.encode(), source.msg_name) + independent = reference.deserialize_cdr(cdr_encode(source), source.__msgtype__) payload = bytes( - reference.serialize_cdr(independent, source.msg_name, little_endian=little_endian) + reference.serialize_cdr(independent, source.__msgtype__, little_endian=little_endian) ) result = ros.decode_odometry(payload) assert result.header.stamp.nanosec == 123456789 @@ -67,21 +87,29 @@ def test_pointcloud_preserves_big_endian_fields_and_row_padding(): for i, n in enumerate(["x", "y", "z"]) ] + [PointField(name="intensity", offset=24, datatype=4, count=1)], - data=struct.pack(">dddH", 1.25, -2.5, 3.0, 1234) + b"padding!", + data=np.frombuffer( + struct.pack(">dddH", 1.25, -2.5, 3.0, 1234) + b"padding!", dtype=np.uint8 + ), + is_dense=False, ) - result = ros.decode_pointcloud2(source.encode()) - assert result.encode() == source.encode() + result = ros.decode_pointcloud2(cdr_encode(source)) + assert cdr_encode(result) == cdr_encode(source) assert result.header.stamp.nanosec == 123456789 np.testing.assert_array_equal(pointcloud_xyz(result), [[1.25, -2.5, 3.0]]) def test_unitree_imu_reorders_firmware_quaternion_without_losing_metadata(): source = Imu( - header=Header(stamp=Time(nanosec=123456789), frame_id="imu"), + header=Header(stamp=Time(nanosec=123456789, sec=0), frame_id="imu"), orientation=Quaternion(x=0.5, y=0.1, z=0.2, w=0.3), + orientation_covariance=np.zeros(9, dtype=np.float64), + angular_velocity=Vector3(x=0.0, y=0.0, z=0.0), + angular_velocity_covariance=np.zeros(9, dtype=np.float64), + linear_acceleration=Vector3(x=0.0, y=0.0, z=0.0), + linear_acceleration_covariance=np.zeros(9, dtype=np.float64), ) - source.orientation_covariance = list(range(9)) - result = ros.decode_imu(source.encode()) + source.orientation_covariance = np.arange(9, dtype=np.float64) + result = ros.decode_imu(cdr_encode(source)) q = result.orientation assert (q.x, q.y, q.z, q.w) == (0.1, 0.2, 0.3, 0.5) assert list(result.orientation_covariance) == list(range(9)) @@ -90,16 +118,23 @@ def test_unitree_imu_reorders_firmware_quaternion_without_losing_metadata(): def test_compressed_image_stays_a_compressed_generated_message(): source = CompressedImage( - header=Header(frame_id="camera"), format="jpeg", data=b"encoded-packet" + header=Header(frame_id="camera", stamp=Time(sec=0, nanosec=0)), + format="jpeg", + data=np.frombuffer(b"encoded-packet", dtype=np.uint8), ) - result = GO2_CODECS["rt/frontvideo"].decode(source.encode()) + result = GO2_CODECS["rt/frontvideo"].decode(cdr_encode(source)) assert isinstance(result, CompressedImage) - assert result.encode() == source.encode() + assert cdr_encode(result) == cdr_encode(source) def test_h264_emission_preserves_packet_stamp_and_pixel_layout(mocker): stamp = Time(sec=42, nanosec=123456789) - packet = CompressedVideo(timestamp=stamp, frame_id="camera", format="h264", data=b"packet") + packet = CompressedVideo( + timestamp=stamp, + frame_id="camera", + format="h264", + data=np.frombuffer(b"packet", dtype=np.uint8), + ) source = Observation(id=7, ts=99.0, data_type=CompressedVideo, _data=packet) frame = mocker.Mock() pixels = np.arange(18, dtype=np.uint8).reshape(2, 3, 3) diff --git a/dimos/robot/unitree/go2/test_connection.py b/dimos/robot/unitree/go2/test_connection.py index 0698e12f94d..208ab210039 100644 --- a/dimos/robot/unitree/go2/test_connection.py +++ b/dimos/robot/unitree/go2/test_connection.py @@ -20,17 +20,22 @@ from collections.abc import Callable, Iterator import threading -import time from types import SimpleNamespace from unittest.mock import MagicMock +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import CameraInfo, RegionOfInterest +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage +import numpy as np import pytest from pytest_mock import MockerFixture from reactivex import Subject from dimos.core.global_config import GlobalConfig -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo +from dimos.msgs.time import to_nanoseconds from dimos.robot.unitree.go2 import connection as go2_conn from dimos.robot.unitree.go2.connection import ConnectionConfig, GO2Connection @@ -137,11 +142,16 @@ def test_connection_config_aes_key_defaults_from_global_config() -> None: def test_odom_to_tf_unprefixed_by_default() -> None: - odom = PoseStamped(ts=1.0, frame_id="world") + odom = PoseStamped( + header=Header(stamp=Time(sec=1, nanosec=123456789), frame_id="world"), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + ) base, camera_link, camera_optical = GO2Connection._odom_to_tf(odom) - assert (base.frame_id, base.child_frame_id) == ("world", "base_link") - assert (camera_link.frame_id, camera_link.child_frame_id) == ("base_link", "camera_link") - assert (camera_optical.frame_id, camera_optical.child_frame_id) == ( + assert (base.header.frame_id, base.child_frame_id) == ("world", "base_link") + assert (camera_link.header.frame_id, camera_link.child_frame_id) == ("base_link", "camera_link") + assert (camera_optical.header.frame_id, camera_optical.child_frame_id) == ( "camera_link", "camera_optical", ) @@ -173,14 +183,30 @@ def test_publish_tf_off_keeps_odometry_on_its_port( ) -> None: """Turning tf off hands the base_link edge to another publisher, not the odom port.""" conn = connection(publish_tf=False) - conn._publish_tf(PoseStamped(ts=1.0, frame_id="ignored")) + conn._publish_tf( + PoseStamped( + header=Header(stamp=Time(sec=1, nanosec=0), frame_id="ignored"), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ) assert conn.tf.publish.call_count == 0 assert conn.odom.publish.call_count == 1 def test_publish_tf_on_by_default(connection: Callable[[bool], GO2Connection]) -> None: conn = connection(publish_tf=True) - conn._publish_tf(PoseStamped(ts=1.0, frame_id="ignored")) + conn._publish_tf( + PoseStamped( + header=Header(stamp=Time(sec=1, nanosec=0), frame_id="ignored"), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + ) assert conn.tf.publish.call_count == 1 assert conn.odom.publish.call_count == 1 @@ -195,34 +221,79 @@ def test_camera_info_is_restamped_on_each_publish( """A frozen stamp collapses a recording's camera_info onto one instant before the run.""" conn = connection(publish_tf=True) conn.camera_info = MagicMock() - conn.camera_info_static = CameraInfo(frame_id="camera_optical", ts=1.0) + conn.camera_info_static = CameraInfo( + header=Header(frame_id="camera_optical", stamp=Time(sec=1, nanosec=0)), + height=0, + width=0, + distortion_model="", + d=np.array([], dtype=np.float64), + k=np.zeros(9, dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), + ) def sleep(_seconds: float) -> None: if conn.camera_info.publish.call_count >= 2: raise _StopLoopError mocker.patch.object(go2_conn.time, "sleep", side_effect=sleep) - before = time.time() + mocker.patch.object( + go2_conn.time, "time_ns", side_effect=[1700000000123456789, 1700000000123456790] + ) with pytest.raises(_StopLoopError): conn.publish_camera_info() published = [call.args[0] for call in conn.camera_info.publish.call_args_list] - assert [info.ts >= before for info in published] == [True, True] - assert [info.frame_id for info in published] == [conn.camera_info_static.frame_id] * 2 - assert conn.camera_info_static.ts == 1.0 + assert [to_nanoseconds(info.header.stamp) for info in published] == [ + 1700000000123456789, + 1700000000123456790, + ] + assert [info.header.frame_id for info in published] == [ + conn.camera_info_static.header.frame_id + ] * 2 + assert conn.camera_info_static.header.stamp == Time(sec=1, nanosec=0) def test_odom_to_tf_prefixed() -> None: """.namespace() sets frame_id_prefix: robot-local frames get prefixed, the odom parent frame stays global so all robots hang off one tree root.""" - odom = PoseStamped(ts=1.0, frame_id="world") + odom = PoseStamped( + header=Header(stamp=Time(sec=1, nanosec=123456789), frame_id="world"), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + ) base, camera_link, camera_optical = GO2Connection._odom_to_tf(odom, prefix="robot0") - assert (base.frame_id, base.child_frame_id) == ("world", "robot0/base_link") - assert (camera_link.frame_id, camera_link.child_frame_id) == ( + assert (base.header.frame_id, base.child_frame_id) == ("world", "robot0/base_link") + assert (camera_link.header.frame_id, camera_link.child_frame_id) == ( "robot0/base_link", "robot0/camera_link", ) - assert (camera_optical.frame_id, camera_optical.child_frame_id) == ( + assert (camera_optical.header.frame_id, camera_optical.child_frame_id) == ( "robot0/camera_link", "robot0/camera_optical", ) + + +def test_published_tf_copies_source_pose_and_keeps_exact_stamp(connection): + conn = connection(publish_tf=True) + source = PoseStamped( + header=Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id="device_odom"), + pose=Pose( + position=Point(x=1.25, y=-2.5, z=0.3), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + conn._publish_tf(source) + message = conn.odom.publish.call_args.args[0] + tf = cdr_decode(cdr_encode(conn.tf.publish.call_args.args[0]), TFMessage) + assert source.header.frame_id == "device_odom" + assert message.header.frame_id == "go2_odom" + assert all(edge.header.stamp == source.header.stamp for edge in tf.transforms) + assert tf.transforms[0].header.frame_id == "go2_odom" + assert tf.transforms[0].transform.translation.x == 1.25 + message.pose.position.x = 99 + assert source.pose.position.x == 1.25 diff --git a/dimos/robot/unitree/go2/test_replay_benchmark.py b/dimos/robot/unitree/go2/test_replay_benchmark.py index 631093636a5..c5a77acf562 100644 --- a/dimos/robot/unitree/go2/test_replay_benchmark.py +++ b/dimos/robot/unitree/go2/test_replay_benchmark.py @@ -22,12 +22,11 @@ import sys import threading +from dimos_generated.geometry_msgs.msg import PoseStamped +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 import pytest from dimos.core.transport_factory import make_transport -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 REPLAY_DB = os.environ.get("DIMOS_BENCH_REPLAY_DB", "go2_short") RUN_TIMEOUT = 420.0 # spawn -> exit deadline, inside the test timeout diff --git a/dimos/robot/unitree/keyboard_teleop.py b/dimos/robot/unitree/keyboard_teleop.py index 99b2eb4fa5b..d767be8467f 100644 --- a/dimos/robot/unitree/keyboard_teleop.py +++ b/dimos/robot/unitree/keyboard_teleop.py @@ -19,6 +19,10 @@ import time from typing import Any +from dimos_generated.geometry_msgs.msg import Twist, Vector3 +from dimos_generated.sensor_msgs.msg import Joy +from dimos_generated.std_msgs.msg import Float32, Int8 +import numpy as np import pygame from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT @@ -31,11 +35,7 @@ from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import Out -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.Joy import Joy -from dimos.msgs.std_msgs.Float32 import Float32 -from dimos.msgs.std_msgs.Int8 import Int8 +from dimos.msgs.time import header_now from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -160,9 +160,11 @@ def start(self) -> None: @rpc def stop(self) -> None: - stop_twist = Twist() - stop_twist.linear = Vector3(0, 0, 0) - stop_twist.angular = Vector3(0, 0, 0) + stop_twist = Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) + stop_twist.linear = Vector3(x=0.0, y=0.0, z=0.0) + stop_twist.angular = Vector3(x=0.0, y=0.0, z=0.0) self.cmd_vel.publish(stop_twist) self._stop_event.set() @@ -191,14 +193,23 @@ def _pygame_loop(self) -> None: self._keys_held.add(event.key) self._keys_down.add(event.key) if event.key in _JOY_KEY_SET: - self.joystick.publish(Joy(buttons=joy_buttons(self._keys_down))) + self.joystick.publish( + Joy( + header=header_now(), + axes=np.array([], dtype=np.float32), + buttons=np.asarray(joy_buttons(self._keys_down), dtype=np.int32), + ) + ) if event.key == pygame.K_SPACE: # Emergency stop - clear all keys and send zero twist self._keys_held.clear() - stop_twist = Twist() - stop_twist.linear = Vector3(0, 0, 0) - stop_twist.angular = Vector3(0, 0, 0) + stop_twist = Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), + angular=Vector3(x=0.0, y=0.0, z=0.0), + ) + stop_twist.linear = Vector3(x=0.0, y=0.0, z=0.0) + stop_twist.angular = Vector3(x=0.0, y=0.0, z=0.0) self.cmd_vel.publish(stop_twist) self._estop_at = time.monotonic() logger.warning("EMERGENCY STOP!") @@ -206,11 +217,11 @@ def _pygame_loop(self) -> None: # ESC quits self._stop_event.set() elif event.key == pygame.K_RETURN: - self.operator_command.publish(Int8(GATE_ADVANCE)) + self.operator_command.publish(Int8(data=GATE_ADVANCE)) elif event.key == pygame.K_k: - self.operator_command.publish(Int8(GATE_SKIP)) + self.operator_command.publish(Int8(data=GATE_SKIP)) elif event.key == pygame.K_BACKSPACE: - self.operator_command.publish(Int8(GATE_QUIT)) + self.operator_command.publish(Int8(data=GATE_QUIT)) elif pygame.K_0 <= event.key <= pygame.K_9: # 0 → 0.0 m, 1 → 0.1 m, …, 9 → 0.9 m corridor half-width. self.e_max.publish(Float32(data=(event.key - pygame.K_0) * 0.1)) @@ -219,12 +230,18 @@ def _pygame_loop(self) -> None: self._keys_held.discard(event.key) self._keys_down.discard(event.key) if event.key in _JOY_KEY_SET: - self.joystick.publish(Joy(buttons=joy_buttons(self._keys_down))) + self.joystick.publish( + Joy( + header=header_now(), + axes=np.array([], dtype=np.float32), + buttons=np.asarray(joy_buttons(self._keys_down), dtype=np.int32), + ) + ) # Generate Twist message from held keys - twist = Twist() - twist.linear = Vector3(0, 0, 0) - twist.angular = Vector3(0, 0, 0) + twist = Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + twist.linear = Vector3(x=0.0, y=0.0, z=0.0) + twist.angular = Vector3(x=0.0, y=0.0, z=0.0) # Movement keys (WASD/QE) — guarded by disable_movement so the # window can run as a pure e_max slider (0-9 keys stay live in @@ -272,7 +289,13 @@ def _pygame_loop(self) -> None: buttons = joy_buttons(self._keys_down) if any(buttons): - self.joystick.publish(Joy(buttons=buttons)) + self.joystick.publish( + Joy( + header=header_now(), + axes=np.array([], dtype=np.float32), + buttons=np.asarray(buttons, dtype=np.int32), + ) + ) self._update_display(twist) # Maintain control loop rate diff --git a/dimos/robot/unitree/mujoco_connection.py b/dimos/robot/unitree/mujoco_connection.py index 0478cb14e0e..0e774c4c2ec 100644 --- a/dimos/robot/unitree/mujoco_connection.py +++ b/dimos/robot/unitree/mujoco_connection.py @@ -31,6 +31,7 @@ from typing import Any, TypeVar import weakref +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Twist from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 from dimos_generated.std_msgs.msg import Header @@ -111,7 +112,7 @@ def cleanup_on_exit( width=VIDEO_WIDTH, height=VIDEO_HEIGHT, axis="vertical", - header=Header(frame_id="camera_optical"), + header=Header(frame_id="camera_optical", stamp=Time(sec=0, nanosec=0)), ) def _launch(self) -> None: diff --git a/dimos/robot/unitree/test_connection.py b/dimos/robot/unitree/test_connection.py index 041ac8b5621..b8dc7e3d6b1 100644 --- a/dimos/robot/unitree/test_connection.py +++ b/dimos/robot/unitree/test_connection.py @@ -23,6 +23,7 @@ from typing import Any from unittest.mock import ANY, AsyncMock, MagicMock, call +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import PoseStamped, TransformStamped, Twist, Vector3 from dimos_generated.sensor_msgs.msg import Image, PointCloud2 import numpy as np @@ -120,8 +121,8 @@ def test_move_api_toggle_sends_selected_wire_command( driver = _stub_driver() monkeypatch.setattr(conn_mod, "LegionConnection", MagicMock(return_value=driver)) twist = Twist( - linear=Vector3(x=1.5, y=-0.4), - angular=Vector3(z=0.8), + linear=Vector3(x=1.5, y=-0.4, z=0.0), + angular=Vector3(z=0.8, x=0.0, y=0.0), ) connection = UnitreeWebRTCConnection(ip="10.0.0.99", **connection_options) @@ -220,11 +221,11 @@ def fail(error): message = received[0] assert (message.header.stamp.sec, message.header.stamp.nanosec) == (1700000000, 123456789) if kind == "lidar": - decoded = PointCloud2.decode(message.encode()) + decoded = cdr_decode(cdr_encode(message), PointCloud2) np.testing.assert_array_equal(pointcloud_xyz(decoded), values) assert decoded.header.frame_id == "world" else: - decoded = Image.decode(message.encode()) + decoded = cdr_decode(cdr_encode(message), Image) np.testing.assert_array_equal(image_view(decoded), pixels) assert decoded.header.frame_id == "camera_optical" finally: @@ -259,7 +260,7 @@ def fail(error): assert not errors assert len(received) == 1 cls = TransformStamped if as_tf else PoseStamped - message = cls.decode(received[0].encode()) + message = cdr_decode(cdr_encode(received[0]), cls) assert message.header.frame_id == "world" assert (message.header.stamp.sec, message.header.stamp.nanosec) == (1700000000, 123456789) if as_tf: diff --git a/dimos/robot/unitree/test_generated_commands.py b/dimos/robot/unitree/test_generated_commands.py index 12ee1901377..31dde6d15a0 100644 --- a/dimos/robot/unitree/test_generated_commands.py +++ b/dimos/robot/unitree/test_generated_commands.py @@ -18,7 +18,8 @@ from types import SimpleNamespace from unittest.mock import Mock -from dimos_generated.geometry_msgs.msg import Twist +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.geometry_msgs.msg import Twist, Vector3 from dimos_generated.std_msgs.msg import Float32, Int8 import pytest @@ -58,16 +59,18 @@ def test_keyboard_generated_operator_gates_and_slider_without_window(monkeypatch monkeypatch.setattr(keyboard.pygame.event, "get", lambda: events) try: module._pygame_loop() - assert [Int8.decode(value.encode()).data for value in received] == [ + assert [cdr_decode(cdr_encode(value), Int8).data for value in received] == [ GATE_ADVANCE, GATE_SKIP, GATE_QUIT, ] - assert Float32.decode(slider[0].encode()).data == pytest.approx(0.7) + assert cdr_decode(cdr_encode(slider[0]), Float32).data == pytest.approx(0.7) assert not motion # disable_movement keeps this an operator-only interface finally: module.stop() - assert len(motion) == 1 and Twist.decode(motion[0].encode()) == Twist() + assert len(motion) == 1 and cdr_decode(cdr_encode(motion[0]), Twist) == Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) def test_g1_skill_constructs_generated_twist_at_captured_rpc_boundary(): @@ -79,7 +82,7 @@ def test_g1_skill_constructs_generated_twist_at_captured_rpc_boundary(): gateway.move.assert_called_once() args, kwargs = gateway.move.call_args assert type(args[0]) is Twist - value = Twist.decode(args[0].encode()) + value = cdr_decode(cdr_encode(args[0]), Twist) assert value.linear.x == 0.25 and value.linear.y == -0.125 assert value.angular.z == 0.5 and value.linear.z == 0 assert kwargs == {"duration": 2} diff --git a/dimos/robot/unitree/test_mujoco_connection.py b/dimos/robot/unitree/test_mujoco_connection.py index 80080cac25b..5f836b675a4 100644 --- a/dimos/robot/unitree/test_mujoco_connection.py +++ b/dimos/robot/unitree/test_mujoco_connection.py @@ -18,6 +18,9 @@ import threading from typing import Any +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.geometry_msgs.msg import Point, PoseStamped, Quaternion + # Imported before subprocess.Popen is patched: mujoco's own import spawns a subprocess. import mujoco # noqa: F401 import numpy as np @@ -26,6 +29,8 @@ from pytest import MonkeyPatch from dimos.core.global_config import GlobalConfig +from dimos.msgs.pointcloud import pointcloud_xyz +from dimos.msgs.time import to_seconds from dimos.robot.unitree import mujoco_connection from dimos.robot.unitree.mujoco_connection import MujocoConnection @@ -190,6 +195,22 @@ def test_stop_terminates_child_before_closing_pumped_output( assert process.terminations == 1 +def test_odometry_reads_generated_pose_and_converts_wxyz(quiet_connection, monkeypatch): + monkeypatch.setattr( + quiet_connection.shm_data, + "read_odom", + lambda: ((np.array([1.25, -2.5, 0.3]), np.array([0.8, 0.0, 0.0, 0.6]), -0.5), 1), + ) + message = quiet_connection.get_odom_message() + assert message is not None + decoded = cdr_decode(cdr_encode(message), PoseStamped) + assert decoded.header.frame_id == "world" + assert (decoded.header.stamp.sec, decoded.header.stamp.nanosec) == (-1, 500000000) + assert decoded.pose.position == Point(x=1.25, y=-2.5, z=0.3) + assert decoded.pose.orientation == Quaternion(z=0.6, w=0.8, x=0.0, y=0.0) + assert quiet_connection.get_odom_message() is None + + def test_start_waits_on_the_process_launched_at_construction(monkeypatch: MonkeyPatch) -> None: launches: list[_QuietProcess] = [] @@ -250,8 +271,8 @@ def test_get_lidar_message_builds_a_point_cloud_once_per_frame( message = quiet_connection.get_lidar_message() assert message is not None - assert message.frame_id == "world" - assert message.ts == 3.5 - np.testing.assert_allclose(message.points_f32(), points) + assert message.header.frame_id == "world" + assert to_seconds(message.header.stamp) == 3.5 + np.testing.assert_allclose(pointcloud_xyz(message), points) assert quiet_connection.get_lidar_message() is None # same frame again diff --git a/dimos/robot/unitree/test_simulation_tf.py b/dimos/robot/unitree/test_simulation_tf.py index 9ffb8d97eff..5636b083eaf 100644 --- a/dimos/robot/unitree/test_simulation_tf.py +++ b/dimos/robot/unitree/test_simulation_tf.py @@ -12,11 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time -from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped -from dimos_generated.sensor_msgs.msg import CameraInfo +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import CameraInfo, RegionOfInterest from dimos_generated.std_msgs.msg import Header from dimos_generated.tf2_msgs.msg import TFMessage +import numpy as np import pytest from dimos.core.global_config import GlobalConfig @@ -29,12 +31,14 @@ def source_pose(): return PoseStamped( header=Header(frame_id="world", stamp=Time(sec=1700000000, nanosec=123456789)), - pose=Pose(position=Point(x=1.25, y=2.5, z=0.7)), + pose=Pose( + position=Point(x=1.25, y=2.5, z=0.7), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), ) def test_dimsim_tf_preserves_source_pose_and_stamp(source_pose): - message = TFMessage.decode(TFMessage(transforms=_odom_to_tf(source_pose)).encode()) + message = cdr_decode(cdr_encode(TFMessage(transforms=_odom_to_tf(source_pose))), TFMessage) assert [(edge.header.frame_id, edge.child_frame_id) for edge in message.transforms] == [ ("world", "base_link"), ("base_link", "camera_link"), @@ -61,7 +65,7 @@ def g1_connection(mocker): def test_g1_simulator_tf_uses_pose_time_for_all_mounts(g1_connection, source_pose): g1_connection._publish_tf(source_pose) - message = TFMessage.decode(g1_connection.tf.publish.call_args.args[0].encode()) + message = cdr_decode(cdr_encode(g1_connection.tf.publish.call_args.args[0]), TFMessage) assert [(edge.header.frame_id, edge.child_frame_id) for edge in message.transforms] == [ ("world", "base_link"), ("base_link", "camera_link"), @@ -74,14 +78,28 @@ def test_g1_simulator_tf_uses_pose_time_for_all_mounts(g1_connection, source_pos def test_g1_camera_info_stamps_each_publication_without_mutating_template(g1_connection, mocker): - template = CameraInfo(header=Header(frame_id="camera_optical"), width=640, height=288) + template = CameraInfo( + header=Header(frame_id="camera_optical", stamp=Time(sec=0, nanosec=0)), + width=640, + height=288, + distortion_model="", + d=np.array([], dtype=np.float64), + k=np.zeros(9, dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), + ) g1_connection.connection.camera_info_static = template mocker.patch("dimos.robot.unitree.g1.mujoco_sim.time.time_ns", return_value=1700000000123456789) mocker.patch.object( g1_connection._stop_event, "wait", side_effect=lambda _: g1_connection._stop_event.set() ) g1_connection._publish_camera_info_loop() - message = CameraInfo.decode(g1_connection.camera_info.publish.call_args.args[0].encode()) + message = cdr_decode( + cdr_encode(g1_connection.camera_info.publish.call_args.args[0]), CameraInfo + ) assert message.header.stamp == Time(sec=1700000000, nanosec=123456789) assert message.header.frame_id == "camera_optical" - assert template.header.stamp == Time() + assert template.header.stamp == Time(sec=0, nanosec=0) diff --git a/dimos/robot/unitree/type/odometry.py b/dimos/robot/unitree/type/odometry.py index 8240d414b89..72ffc779eba 100644 --- a/dimos/robot/unitree/type/odometry.py +++ b/dimos/robot/unitree/type/odometry.py @@ -11,6 +11,7 @@ # 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. +from copy import deepcopy from typing import Literal, TypedDict from dimos_generated.builtin_interfaces.msg import Time @@ -82,7 +83,7 @@ def pose_from_webrtc_odometry( source_header = data["header"] pose = data["pose"] return GeneratedPoseStamped( - header=header + header=deepcopy(header) if header is not None else GeneratedHeader( stamp=Time(**source_header["stamp"]), frame_id=source_header["frame_id"] diff --git a/dimos/robot/unitree/type/test_lidar.py b/dimos/robot/unitree/type/test_lidar.py index e92bbf19135..9d8faecddee 100644 --- a/dimos/robot/unitree/type/test_lidar.py +++ b/dimos/robot/unitree/type/test_lidar.py @@ -16,6 +16,7 @@ import itertools from typing import cast +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import PointCloud2 from dimos_generated.std_msgs.msg import Header @@ -34,7 +35,17 @@ def _stamped(seconds: float) -> PointCloud2: - return PointCloud2(header=Header(stamp=time_from_seconds(seconds))) + return PointCloud2( + header=Header(stamp=time_from_seconds(seconds), frame_id=""), + height=0, + width=0, + fields=[], + is_bigendian=False, + point_step=0, + row_step=0, + data=np.array([], dtype=np.uint8), + is_dense=False, + ) @pytest.mark.self_hosted @@ -115,7 +126,7 @@ def test_repair_stale_ts_calibration_boundary_one_differs() -> None: def test_webrtc_xyz_conversion_preserves_points_and_explicit_time(override): raw = {"data": {"stamp": -0.5, "data": {"points": np.array([[1.25, 2.5, 3.75]])}}} cloud = pointcloud2_from_webrtc_lidar(cast("RawLidarMsg", raw), stamp=override) - decoded = PointCloud2.decode(cloud.encode()) + decoded = cdr_decode(cdr_encode(cloud), PointCloud2) np.testing.assert_array_equal(pointcloud_xyz(decoded), [[1.25, 2.5, 3.75]]) assert decoded.header.frame_id == "world" assert decoded.header.stamp == ( diff --git a/dimos/robot/unitree/type/test_pose_conversion.py b/dimos/robot/unitree/type/test_pose_conversion.py index a8da638faed..5c1a7571c06 100644 --- a/dimos/robot/unitree/type/test_pose_conversion.py +++ b/dimos/robot/unitree/type/test_pose_conversion.py @@ -14,6 +14,7 @@ import copy +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import PoseStamped from dimos_generated.std_msgs.msg import Header @@ -27,7 +28,7 @@ def test_device_pose_preserves_original_ros_header(seconds, nanoseconds): raw = copy.deepcopy(raw_odometry_msg_sample) raw["data"]["header"]["stamp"] = {"sec": seconds, "nanosec": nanoseconds} message = pose_from_webrtc_odometry(raw) - decoded = PoseStamped.decode(message.encode()) + decoded = cdr_decode(cdr_encode(message), PoseStamped) assert decoded.header == Header(frame_id="odom", stamp=Time(sec=seconds, nanosec=nanoseconds)) assert decoded.pose.position.x == 5.961965 assert decoded.pose.position.y == -2.916958 diff --git a/dimos/robot/unitree/unitree_skill_container.py b/dimos/robot/unitree/unitree_skill_container.py index de4f6385ed4..6ad3d3f684b 100644 --- a/dimos/robot/unitree/unitree_skill_container.py +++ b/dimos/robot/unitree/unitree_skill_container.py @@ -206,7 +206,7 @@ def _goal_pose( """Where move_to sends the robot, in the world frame.""" roll, pitch, yaw = quaternion_euler(current.pose.orientation) if relative: - translated = translate_pose_local(current.pose, Vector3(x=x, y=y)) + translated = translate_pose_local(current.pose, Vector3(x=x, y=y, z=0.0)) position = translated.position heading = yaw + math.radians(degrees or 0.0) else: diff --git a/dimos/simulation/dimsim/cdr_schemas.py b/dimos/simulation/dimsim/cdr_schemas.py index 44d4f858f30..4ee6402b1b2 100644 --- a/dimos/simulation/dimsim/cdr_schemas.py +++ b/dimos/simulation/dimsim/cdr_schemas.py @@ -19,6 +19,7 @@ import json +from dimos_cdrgen.registry import schema as cdr_schema from dimos_generated.geometry_msgs.msg import PoseStamped, Twist from dimos_generated.sensor_msgs.msg import Image, PointCloud2 @@ -28,7 +29,9 @@ def schema_source() -> str: - definitions = {t.msg_name: t.schema for t in (Image, PointCloud2, PoseStamped, Twist)} + definitions = { + t.__msgtype__: cdr_schema(t.__msgtype__) for t in (Image, PointCloud2, PoseStamped, Twist) + } return ( "// Generated from the canonical .msg definitions; do not edit by hand.\n" "export const schemas: Record = " diff --git a/dimos/simulation/engines/mujoco_engine.py b/dimos/simulation/engines/mujoco_engine.py index 7ba9c24c220..b609f6828c6 100644 --- a/dimos/simulation/engines/mujoco_engine.py +++ b/dimos/simulation/engines/mujoco_engine.py @@ -723,15 +723,15 @@ def read_joint_efforts(self) -> list[float]: return self.joint_efforts def write_joint_command(self, command: JointState) -> None: - if command.position: + if len(command.position): self._command_mode = "position" self._set_position_targets(list(command.position)) return - if command.velocity: + if len(command.velocity): self._command_mode = "velocity" self._set_velocity_targets(list(command.velocity)) return - if command.effort: + if len(command.effort): self._command_mode = "effort" self._set_effort_targets(list(command.effort)) return diff --git a/dimos/simulation/engines/mujoco_sim_module.py b/dimos/simulation/engines/mujoco_sim_module.py index 94bb3ca7a57..af4e7ab37ba 100644 --- a/dimos/simulation/engines/mujoco_sim_module.py +++ b/dimos/simulation/engines/mujoco_sim_module.py @@ -32,6 +32,19 @@ import time from typing import Any +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, Imu, JointState, PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage import mujoco import numpy as np from numpy.typing import NDArray @@ -43,16 +56,11 @@ from dimos.core.module import Module, ModuleConfig from dimos.core.stream import Out from dimos.hardware.sensors.camera.spec import DepthCameraConfig, DepthCameraHardware -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.sensor_msgs.Imu import Imu -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.camera_info import camera_info_from_intrinsics, camera_info_with_stamp +from dimos.msgs.geometry import quaternion_from_matrix +from dimos.msgs.image import image_from_array +from dimos.msgs.pointcloud import pointcloud_from_rgbd, pointcloud_from_xyz, voxel_downsample_cloud +from dimos.msgs.time import time_from_seconds from dimos.simulation.engines.mujoco_engine import ( CameraConfig, CameraFrame, @@ -96,24 +104,27 @@ def _pose_matrix( def _transform_from_matrix( matrix: NDArray[np.float64], *, frame_id: str, child_frame_id: str, ts: float -) -> Transform: - return Transform( - translation=Vector3( - float(matrix[0, 3]), - float(matrix[1, 3]), - float(matrix[2, 3]), - ), - rotation=Quaternion.from_rotation_matrix(matrix[:3, :3]), - frame_id=frame_id, +) -> TransformStamped: + return TransformStamped( + header=Header(stamp=time_from_seconds(ts), frame_id=frame_id), child_frame_id=child_frame_id, - ts=ts, + transform=Transform( + translation=Vector3( + x=float(matrix[0, 3]), y=float(matrix[1, 3]), z=float(matrix[2, 3]) + ), + rotation=quaternion_from_matrix(matrix[:3, :3]), + ), ) -def _default_identity_transform() -> Transform: - return Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), +def _default_identity_transform() -> TransformStamped: + return TransformStamped( + transform=Transform( + rotation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), + translation=Vector3(x=0.0, y=0.0, z=0.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + child_frame_id="", ) @@ -127,11 +138,13 @@ def _imu_from_mujoco_wxyz( ) -> Imu: w, x, y, z = quaternion return Imu( - orientation=Quaternion(x, y, z, w), - angular_velocity=Vector3(*gyroscope), - linear_acceleration=Vector3(*accelerometer), - frame_id=frame_id, - ts=ts, + header=Header(stamp=time_from_seconds(ts), frame_id=frame_id), + orientation=Quaternion(x=x, y=y, z=z, w=w), + angular_velocity=Vector3(x=gyroscope[0], y=gyroscope[1], z=gyroscope[2]), + linear_acceleration=Vector3(x=accelerometer[0], y=accelerometer[1], z=accelerometer[2]), + orientation_covariance=np.zeros(9, dtype=np.float64), + angular_velocity_covariance=np.zeros(9, dtype=np.float64), + linear_acceleration_covariance=np.zeros(9, dtype=np.float64), ) @@ -166,11 +179,27 @@ def pre_step(self, engine: MujocoEngine) -> None: if shm.read_command_mode() == CMD_MODE_PD_TAU: self._latest_pd_pos_target = pos_cmd else: - engine.write_joint_command(JointState(position=pos_cmd.tolist())) + engine.write_joint_command( + JointState( + position=np.asarray(pos_cmd.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) vel_cmd = shm.read_velocity_command(dof) if vel_cmd is not None: - engine.write_joint_command(JointState(velocity=vel_cmd.tolist())) + engine.write_joint_command( + JointState( + velocity=np.asarray(vel_cmd.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + position=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), + ) + ) kp_cmd = shm.read_kp_command(dof) if kp_cmd is not None: @@ -195,7 +224,15 @@ def pre_step(self, engine: MujocoEngine) -> None: + self._latest_pd_kd * (-dq) + tau_ff ) - engine.write_joint_command(JointState(effort=tau.tolist())) + engine.write_joint_command( + JointState( + effort=np.asarray(tau.tolist(), dtype=np.float64), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + name=[], + position=np.array([], dtype=np.float64), + velocity=np.array([], dtype=np.float64), + ) + ) if self._gripper_idx is not None: gripper_cmd = shm.read_gripper_command() @@ -261,7 +298,7 @@ class MujocoSimModuleConfig(ModuleConfig, DepthCameraConfig): height: int = 480 fps: int = 15 base_frame_id: str = "link7" - base_transform: Transform | None = Field(default_factory=_default_identity_transform) + base_transform: TransformStamped | None = Field(default_factory=_default_identity_transform) align_depth_to_color: bool = True enable_color: bool = True enable_depth: bool = True @@ -394,7 +431,7 @@ def get_color_camera_info(self) -> CameraInfo | None: base = self._camera_info_base if base is None: return None - return base.with_ts(self._camera_info_ts()) + return camera_info_with_stamp(base, time_from_seconds(self._camera_info_ts())) @rpc def get_depth_camera_info(self) -> CameraInfo | None: @@ -402,7 +439,7 @@ def get_depth_camera_info(self) -> CameraInfo | None: base = self._camera_info_base if base is None: return None - return base.with_ts(self._camera_info_ts()) + return camera_info_with_stamp(base, time_from_seconds(self._camera_info_ts())) @rpc def get_depth_scale(self) -> float: @@ -796,15 +833,18 @@ def _publish_shm_and_lcm(self, engine: MujocoEngine) -> None: ] # (w, x, y, z) per MuJoCo convention self.odom.publish( PoseStamped( - ts=time.time(), - frame_id="world", - position=Vector3(float(base_pos[0]), float(base_pos[1]), float(base_pos[2])), - orientation=Quaternion( - float(base_quat[1]), - float(base_quat[2]), - float(base_quat[3]), - float(base_quat[0]), - ), # PoseStamped uses x,y,z,w + header=Header(stamp=time_from_seconds(time.time()), frame_id="world"), + pose=Pose( + position=Point( + x=float(base_pos[0]), y=float(base_pos[1]), z=float(base_pos[2]) + ), + orientation=Quaternion( + x=float(base_quat[1]), + y=float(base_quat[2]), + z=float(base_quat[3]), + w=float(base_quat[0]), + ), + ), ) ) @@ -861,14 +901,14 @@ def _build_camera_info(self) -> None: fovy_rad = math.radians(fovy_deg) fy = h / (2.0 * math.tan(fovy_rad / 2.0)) fx = fy # square pixels - camera_info = CameraInfo.from_intrinsics( + camera_info = camera_info_from_intrinsics( fx=fx, fy=fy, cx=w / 2.0, cy=h / 2.0, width=w, height=h, - frame_id=self._color_optical_frame, + header=Header(frame_id=self._color_optical_frame, stamp=Time(sec=0, nanosec=0)), ) with self._state_lock: self._camera_info_base = camera_info @@ -916,20 +956,18 @@ def _publish_loop(self) -> None: self._latest_frame_ts = ts if self.config.enable_color: - color_img = Image( - data=frame.rgb, - format=ImageFormat.RGB, - frame_id=self._color_optical_frame, - ts=ts, + color_img = image_from_array( + frame.rgb, + encoding="rgb8", + header=Header(frame_id=self._color_optical_frame, stamp=time_from_seconds(ts)), ) self.color_image.publish(color_img) if self.config.enable_depth: - depth_img = Image( - data=frame.depth, - format=ImageFormat.DEPTH, - frame_id=self._color_optical_frame, - ts=ts, + depth_img = image_from_array( + frame.depth, + encoding="32FC1", + header=Header(frame_id=self._color_optical_frame, stamp=time_from_seconds(ts)), ) self.depth_image.publish(depth_img) @@ -954,16 +992,7 @@ def _publish_camera_info(self) -> None: if base is None: return ts = self._camera_info_ts() - info = CameraInfo( - height=base.height, - width=base.width, - distortion_model=base.distortion_model, - D=base.D, - K=base.K, - P=base.P, - frame_id=base.frame_id, - ts=ts, - ) + info = camera_info_with_stamp(base, time_from_seconds(ts)) self.camera_info.publish(info) self.depth_camera_info.publish(info) @@ -984,35 +1013,37 @@ def _publish_tf(self, ts: float, frame: CameraFrame | None) -> None: self.tf.publish( TFMessage( - _transform_from_matrix( - optical_transform, - frame_id=parent_frame, - child_frame_id=self._color_optical_frame, - ts=ts, - ), - _transform_from_matrix( - optical_transform, - frame_id=parent_frame, - child_frame_id=self._depth_optical_frame, - ts=ts, - ), - _transform_from_matrix( - camera_transform, - frame_id=parent_frame, - child_frame_id=self._camera_link, - ts=ts, - ), - *self._body_transforms(), + transforms=[ + _transform_from_matrix( + optical_transform, + frame_id=parent_frame, + child_frame_id=self._color_optical_frame, + ts=ts, + ), + _transform_from_matrix( + optical_transform, + frame_id=parent_frame, + child_frame_id=self._depth_optical_frame, + ts=ts, + ), + _transform_from_matrix( + camera_transform, + frame_id=parent_frame, + child_frame_id=self._camera_link, + ts=ts, + ), + *self._body_transforms(), + ], ) ) - def _body_transforms(self) -> list[Transform]: + def _body_transforms(self) -> list[TransformStamped]: """World poses of ``tracked_bodies``; a name missing from the model is skipped after one warning.""" engine = self._engine if engine is None or not self.config.tracked_bodies: return [] ts = time.time() - transforms: list[Transform] = [] + transforms: list[TransformStamped] = [] for name in self.config.tracked_bodies: pose = engine.get_body_pose(name) if pose is None: @@ -1022,12 +1053,15 @@ def _body_transforms(self) -> list[Transform]: continue position, (qx, qy, qz, qw) = pose transforms.append( - Transform( - translation=Vector3(float(position[0]), float(position[1]), float(position[2])), - rotation=Quaternion(float(qx), float(qy), float(qz), float(qw)), - frame_id="world", + TransformStamped( + header=Header(stamp=time_from_seconds(ts), frame_id="world"), child_frame_id=name, - ts=ts, + transform=Transform( + translation=Vector3( + float(position[0]), float(position[1]), float(position[2]) + ), + rotation=Quaternion(float(qx), float(qy), float(qz), float(qw)), + ), ) ) return transforms @@ -1047,25 +1081,27 @@ def _generate_pointcloud(self) -> None: if frame is None: return try: - color_img = Image( - data=frame.rgb, - format=ImageFormat.RGB, - frame_id=self._color_optical_frame, - ts=frame.timestamp, + color_img = image_from_array( + frame.rgb, + encoding="rgb8", + header=Header( + frame_id=self._color_optical_frame, stamp=time_from_seconds(frame.timestamp) + ), ) - depth_img = Image( - data=frame.depth, - format=ImageFormat.DEPTH, - frame_id=self._color_optical_frame, - ts=frame.timestamp, + depth_img = image_from_array( + frame.depth, + encoding="32FC1", + header=Header( + frame_id=self._color_optical_frame, stamp=time_from_seconds(frame.timestamp) + ), ) - pcd = PointCloud2.from_rgbd( - color_image=color_img, - depth_image=depth_img, - camera_info=camera_info, + pcd = pointcloud_from_rgbd( + color=color_img, + depth=depth_img, + calibration=camera_info, depth_scale=1.0, ) - pcd = pcd.voxel_downsample(0.005) + pcd = voxel_downsample_cloud(pcd, 0.005) self.pointcloud.publish(pcd) except Exception as exc: logger.error("Pointcloud generation error", error=str(exc)) @@ -1093,12 +1129,11 @@ def _generate_mujoco_lidar_pointcloud(self) -> None: return try: - pcd = PointCloud2.from_numpy( + pcd = pointcloud_from_xyz( np.vstack(all_points), - frame_id="world", - timestamp=latest_ts or time.time(), + header=Header(frame_id="world", stamp=time_from_seconds(latest_ts or time.time())), ) - pcd = pcd.voxel_downsample(self.config.mujoco_lidar_voxel_size) + pcd = voxel_downsample_cloud(pcd, self.config.mujoco_lidar_voxel_size) self.pointcloud.publish(pcd) except Exception as exc: logger.error("MuJoCo lidar pointcloud generation error", error=str(exc)) diff --git a/dimos/simulation/engines/test_mujoco_sim_module.py b/dimos/simulation/engines/test_mujoco_sim_module.py index 27bf4598fb5..dc04c8dabeb 100644 --- a/dimos/simulation/engines/test_mujoco_sim_module.py +++ b/dimos/simulation/engines/test_mujoco_sim_module.py @@ -20,6 +20,8 @@ from typing import Any from unittest.mock import MagicMock +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import PoseStamped from dimos_generated.sensor_msgs.msg import Imu from dimos_generated.std_msgs.msg import Header @@ -129,7 +131,7 @@ def cleanup(self) -> None: module._publish_shm_and_lcm(_FakeEngine) assert events == ["joint_state", "imu", "ready"] - pose = PoseStamped.decode(module.odom.publish.call_args.args[0].encode()) + pose = cdr_decode(cdr_encode(module.odom.publish.call_args.args[0]), PoseStamped) assert pose.header.frame_id == "world" assert (pose.pose.position.x, pose.pose.position.y, pose.pose.position.z) == ( 0.0, @@ -137,7 +139,7 @@ def cleanup(self) -> None: 0.75, ) assert pose.pose.orientation.w == 1.0 - imu = Imu.decode(module.imu.publish.call_args.args[0].encode()) + imu = cdr_decode(cdr_encode(module.imu.publish.call_args.args[0]), Imu) assert imu.orientation.w == 1.0 assert (imu.angular_velocity.x, imu.angular_velocity.y, imu.angular_velocity.z) == ( 0.1, @@ -415,7 +417,7 @@ def test_publish_loop_stamps_messages_with_frame_timestamp() -> None: fy=1.0, cx=0.5, cy=0.5, - header=Header(frame_id="wrist_camera_color_frame"), + header=Header(frame_id="wrist_camera_color_frame", stamp=Time(sec=0, nanosec=0)), ) color: list[Any] = [] depth: list[Any] = [] @@ -451,7 +453,7 @@ def test_camera_info_falls_back_to_wall_clock_before_first_frame() -> None: fy=1.0, cx=0.5, cy=0.5, - header=Header(frame_id="wrist_camera_color_frame"), + header=Header(frame_id="wrist_camera_color_frame", stamp=Time(sec=0, nanosec=0)), ) assert module._latest_frame_ts is None assert module._camera_info_ts() == pytest.approx(time.time(), abs=5.0) diff --git a/dimos/simulation/go2_sim/test_world.py b/dimos/simulation/go2_sim/test_world.py index 6e1ad5ff0b0..b3613ed9810 100644 --- a/dimos/simulation/go2_sim/test_world.py +++ b/dimos/simulation/go2_sim/test_world.py @@ -14,12 +14,12 @@ from __future__ import annotations +from dimos_generated.dimos_msgs.msg import Contact import mujoco import numpy as np from numpy.typing import NDArray import pytest -from dimos.msgs.sim_msgs.Contacts import Contact from dimos.simulation.go2_legged.policy import OnnxGo2Policy from dimos.simulation.go2_legged.robot import CONTROL_DT from dimos.simulation.go2_sim.world import ( @@ -91,21 +91,21 @@ def test_standing_robot_touches_the_floor_only_with_its_feet(sim: Go2Sim) -> Non for _ in range(25): sim.tick(STILL) contacts = sim.contacts() - assert Contact("foot", "floor") in contacts + assert Contact(part="foot", kind="floor") in contacts assert {c.kind for c in contacts} == {"floor"} assert "trunk" not in {c.part for c in contacts} def test_contacts_name_the_scene_kind_and_the_robot_part(sim: Go2Sim) -> None: sim.reset(-0.05, 3.0, sim.scene.start[2], 0.0) - assert Contact("trunk", "wall") in sim.contacts() + assert Contact(part="trunk", kind="wall") in sim.contacts() clutter = max( (b for b in sim.scene.boxes if b.kind == "clutter"), key=lambda b: min(b.half[:2]) ) sim.reset(clutter.center[0], clutter.center[1], clutter.center[2] + clutter.half[2], 0.0) for _ in range(25): sim.tick(STILL) - assert Contact("foot", "clutter") in sim.contacts() + assert Contact(part="foot", kind="clutter") in sim.contacts() def test_sensor_velocity_matches_the_motion(sim: Go2Sim) -> None: diff --git a/dimos/simulation/go2_sim/world.py b/dimos/simulation/go2_sim/world.py index d356f409af3..5959ad7088b 100644 --- a/dimos/simulation/go2_sim/world.py +++ b/dimos/simulation/go2_sim/world.py @@ -26,6 +26,22 @@ from threading import Event, Thread import time +from dimos_generated.dimos_msgs.msg import Contact, Contacts, LineSegment3D, LineSegments3D +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + PoseWithCovariance, + Transform, + TransformStamped, + Twist, + TwistWithCovariance, + Vector3, +) +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage import mujoco import mujoco.viewer import numpy as np @@ -38,17 +54,10 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.LineSegments3D import LineSegments3D -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.sim_msgs.Contacts import Contact, Contacts, Part -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.contacts import Part +from dimos.msgs.geometry import quaternion_from_array, transform_matrix, vector_from_array +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import header_now, time_from_seconds from dimos.protocol.tf.static_tf_publisher import frames_to_edge_transforms from dimos.robot.unitree.go2.constants import CMD_VEL_TIMEOUT from dimos.robot.unitree.go2.go2_mid360_static_transforms import FRAMES @@ -108,7 +117,9 @@ def _mount() -> tuple[NDArray[np.float64], NDArray[np.float64]]: """base_link -> mid360_link as a translation and rotation matrix, from the rig's frame tree.""" edges = {t.child_frame_id: t for t in frames_to_edge_transforms(FRAMES)} - matrix = edges["front_camera"].to_matrix() @ edges["mid360_link"].to_matrix() + matrix = transform_matrix(edges["front_camera"].transform) @ transform_matrix( + edges["mid360_link"].transform + ) return matrix[:3, 3].copy(), matrix[:3, :3].copy() @@ -255,14 +266,14 @@ def sensor_velocity(self) -> tuple[NDArray[np.float64], NDArray[np.float64]]: def contacts(self) -> list[Contact]: """Every (robot part, scene kind) pair in contact right now, sorted.""" m, d = self.model, self.data - found: set[Contact] = set() + found: set[tuple[Part, str]] = set() for i in range(d.ncon): c = d.contact[i] for scene_geom, robot_geom in ((c.geom1, c.geom2), (c.geom2, c.geom1)): kind = self._kind.get(scene_geom) if kind is not None and m.geom_bodyid[robot_geom] != 0: - found.add(Contact(self._part(robot_geom), kind)) - return sorted(found) + found.add((self._part(robot_geom), kind)) + return [Contact(part=part, kind=kind) for part, kind in sorted(found)] def _part(self, geom: int) -> Part: if geom == self._lidar_geom: @@ -384,7 +395,7 @@ def _open_viewer(self, sim: Go2Sim) -> mujoco.viewer.Handle | None: def _run(self) -> None: # open3d loads on the first cloud, about a second - PointCloud2.from_numpy(np.zeros((0, 3), np.float32), frame_id=SENSOR_FRAME_ID) + pointcloud_from_xyz(np.zeros((0, 3), np.float32), header=header_now(SENSOR_FRAME_ID)) viewer = self._open_viewer(self._sim) try: self._simulate(self._sim, viewer) @@ -406,16 +417,32 @@ def _simulate(self, sim: Go2Sim, viewer: mujoco.viewer.Handle | None) -> None: self._publish_poses(sim, stamp) contacts = sim.contacts() if contacts != last_contacts: - self.contacts.publish(Contacts(contacts, ts=stamp)) + self.contacts.publish( + Contacts( + header=Header(stamp=time_from_seconds(stamp), frame_id=""), + contacts=contacts, + ) + ) last_contacts = contacts if frame is not None: self.lidar.publish( - PointCloud2.from_numpy(frame.points, frame_id=SENSOR_FRAME_ID, timestamp=stamp) + pointcloud_from_xyz( + frame.points, + header=Header(stamp=time_from_seconds(stamp), frame_id=SENSOR_FRAME_ID), + ) ) if sim.t >= next_scene_publish: self.scene.publish( LineSegments3D( - ts=stamp, frame_id=ODOM_FRAME_ID, segments=scene_edges(sim.scene) + header=Header(stamp=time_from_seconds(stamp), frame_id=ODOM_FRAME_ID), + segments=[ + LineSegment3D( + start=Point(x=float(a[0]), y=float(a[1]), z=float(a[2])), + end=Point(x=float(c[0]), y=float(c[1]), z=float(c[2])), + weight=1.0, + ) + for a, c in scene_edges(sim.scene) + ], ) ) next_scene_publish = sim.t + SCENE_PUBLISH_DT @@ -430,34 +457,52 @@ def _publish_poses(self, sim: Go2Sim, stamp: float) -> None: position, rotation = sim.sensor_pose() linear, angular = sim.sensor_velocity() q = Rotation.from_matrix(rotation).as_quat() - sensor = Pose(*map(float, position), *map(float, q)) + sensor = Pose( + position=Point(x=float(position[0]), y=float(position[1]), z=float(position[2])), + orientation=quaternion_from_array(q), + ) + header = Header(stamp=time_from_seconds(stamp), frame_id=ODOM_FRAME_ID) self.odometry.publish( Odometry( - ts=stamp, - frame_id=ODOM_FRAME_ID, + header=header, child_frame_id=SENSOR_FRAME_ID, - pose=sensor, - twist=Twist(Vector3(linear.tolist()), Vector3(angular.tolist())), + pose=PoseWithCovariance(pose=sensor, covariance=np.zeros(36, dtype=np.float64)), + twist=TwistWithCovariance( + twist=Twist( + linear=vector_from_array(linear), angular=vector_from_array(angular) + ), + covariance=np.zeros(36, dtype=np.float64), + ), ) ) self.tf.publish( TFMessage( - Transform( - translation=Vector3(sensor.position), - rotation=Quaternion(sensor.orientation), - frame_id=ODOM_FRAME_ID, - child_frame_id=SENSOR_FRAME_ID, - ts=stamp, - ) + transforms=[ + TransformStamped( + header=header, + child_frame_id=SENSOR_FRAME_ID, + transform=Transform( + translation=Vector3( + x=sensor.position.x, y=sensor.position.y, z=sensor.position.z + ), + rotation=sensor.orientation, + ), + ) + ] ) ) base_position, base_rotation = sim.base_pose() base_q = Rotation.from_matrix(base_rotation).as_quat() self.ground_truth.publish( PoseStamped( - *map(float, base_position), - *map(float, base_q), - ts=stamp, - frame_id=ODOM_FRAME_ID, + header=header, + pose=Pose( + position=Point( + x=float(base_position[0]), + y=float(base_position[1]), + z=float(base_position[2]), + ), + orientation=quaternion_from_array(base_q), + ), ) ) diff --git a/dimos/simulation/habitat/connection.py b/dimos/simulation/habitat/connection.py index 01dcb079d8a..e8556a70517 100644 --- a/dimos/simulation/habitat/connection.py +++ b/dimos/simulation/habitat/connection.py @@ -19,17 +19,15 @@ from pathlib import Path from typing import Any +from dimos_generated.geometry_msgs.msg import Twist +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 +from dimos_generated.tf2_msgs.msg import TFMessage +from dimos_generated.vision_msgs.msg import Detection3DArray from pydantic import Field from dimos.core.native_module import LogFormat, NativeModule, NativeModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray from dimos.simulation.habitat.constants import HABITAT_ROOT, HM3D_EXAMPLE_DATASET_CONFIG from dimos.simulation.habitat.server import HabitatProp diff --git a/dimos/simulation/habitat/nix/install.sh b/dimos/simulation/habitat/nix/install.sh index e4450a9873a..3368327695e 100755 --- a/dimos/simulation/habitat/nix/install.sh +++ b/dimos/simulation/habitat/nix/install.sh @@ -26,7 +26,7 @@ fi if [ ! -f "$ROOT/build/message-codegen/install/lib/cmake/fastcdr/fastcdr-config.cmake" ]; then bash "$ROOT/scripts/setup_message_codegen.sh" fi -(cd "$ROOT" && .venv/bin/python -m dimos.message_codegen.generate --package \ +(cd "$ROOT" && .venv/bin/python -m dimos_cdrgen.generate --package \ --output "$OUT/messages" \ --type geometry_msgs/msg/Twist --type nav_msgs/msg/Odometry \ --type sensor_msgs/msg/CameraInfo --type sensor_msgs/msg/Image \ diff --git a/dimos/simulation/habitat/server.py b/dimos/simulation/habitat/server.py index 9dc64d53140..0113f388eef 100644 --- a/dimos/simulation/habitat/server.py +++ b/dimos/simulation/habitat/server.py @@ -16,7 +16,7 @@ """Habitat-sim native: drives a scene from Twist, publishes dimos messages on zenoh. Runs under python 3.9 in its own env, so: no dimos imports, no 3.10+ syntax -(``test_server.py`` enforces both). ``dimos_lcm`` provides the encoders. +(``test_server.py`` enforces both). ``dimos_generated`` provides the standalone CDR encoders. Reads one JSON line on stdin: ``topics`` (port -> zenoh key), ``config``, ``session``. """ @@ -33,24 +33,36 @@ import time from typing import Any, NamedTuple -from dimos_lcm.geometry_msgs.Pose import Pose as LCMPose -from dimos_lcm.geometry_msgs.Quaternion import Quaternion -from dimos_lcm.geometry_msgs.Transform import Transform -from dimos_lcm.geometry_msgs.TransformStamped import TransformStamped -from dimos_lcm.geometry_msgs.Twist import Twist as LCMTwist -from dimos_lcm.geometry_msgs.Vector3 import Vector3 -from dimos_lcm.nav_msgs.Odometry import Odometry as LCMOdometry -from dimos_lcm.sensor_msgs.CameraInfo import CameraInfo as LCMCameraInfo -from dimos_lcm.sensor_msgs.Image import Image as LCMImage -from dimos_lcm.sensor_msgs.PointCloud2 import PointCloud2 as LCMPointCloud2 -from dimos_lcm.sensor_msgs.PointField import PointField -from dimos_lcm.std_msgs.Header import Header -from dimos_lcm.tf2_msgs.TFMessage import TFMessage as LCMTFMessage -from dimos_lcm.vision_msgs.BoundingBox3D import BoundingBox3D -from dimos_lcm.vision_msgs.Detection3D import Detection3D -from dimos_lcm.vision_msgs.Detection3DArray import Detection3DArray as LCMDetection3DArray -from dimos_lcm.vision_msgs.ObjectHypothesis import ObjectHypothesis -from dimos_lcm.vision_msgs.ObjectHypothesisWithPose import ObjectHypothesisWithPose +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseWithCovariance, + Quaternion, + Transform, + TransformStamped, + Twist, + TwistWithCovariance, + Vector3, +) +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import ( + CameraInfo, + Image, + PointCloud2, + PointField, + RegionOfInterest, +) +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage +from dimos_generated.vision_msgs.msg import ( + BoundingBox3D, + Detection3D, + Detection3DArray, + ObjectHypothesis, + ObjectHypothesisWithPose, +) import numpy as np import numpy.typing as npt import zenoh @@ -76,14 +88,21 @@ def warn(msg: str) -> None: def _stamp(header: Header, ts: float) -> None: - header.seq = 0 header.stamp.sec = int(ts) - header.stamp.nsec = int((ts - int(ts)) * 1e9) + header.stamp.nanosec = int((ts - int(ts)) * 1e9) def image_msg(array: np.ndarray, encoding: str, frame_id: str, ts: float) -> bytes: - m = LCMImage() - m.header = Header() + m = Image( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + height=0, + width=0, + encoding="", + is_bigendian=0, + step=0, + data=np.array([], dtype=np.uint8), + ) + m.header = Header(stamp=Time(sec=0, nanosec=0), frame_id="") _stamp(m.header, ts) m.header.frame_id = frame_id m.height, m.width = int(array.shape[0]), int(array.shape[1]) @@ -92,35 +111,45 @@ def image_msg(array: np.ndarray, encoding: str, frame_id: str, ts: float) -> byt channels = 1 if array.ndim == 2 else array.shape[2] m.step = m.width * array.dtype.itemsize * channels view = memoryview(np.ascontiguousarray(array)).cast("B") - m.data_length = len(view) - m.data = view - return bytes(m.lcm_encode()) + m.data = np.frombuffer(view, dtype=np.uint8) + return bytes(cdr_encode(m)) def camera_info_msg(k: dict[str, float], frame_id: str, ts: float) -> bytes: - m = LCMCameraInfo() - m.header = Header() + m = CameraInfo( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + height=0, + width=0, + distortion_model="", + d=np.array([], dtype=np.float64), + k=np.zeros(9, dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), + ) + m.header = Header(stamp=Time(sec=0, nanosec=0), frame_id="") _stamp(m.header, ts) m.header.frame_id = frame_id m.width, m.height = int(k["width"]), int(k["height"]) m.distortion_model = "plumb_bob" - m.D = [0.0] * 5 - m.D_length = 5 + m.d = np.zeros(5, dtype=np.float64) fx, fy, cx, cy = k["fx"], k["fy"], k["cx"], k["cy"] - m.K = [fx, 0.0, cx, 0.0, fy, cy, 0.0, 0.0, 1.0] - m.R = [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0] - m.P = [fx, 0.0, cx, 0.0, 0.0, fy, cy, 0.0, 0.0, 0.0, 1.0, 0.0] + m.k = np.array([fx, 0.0, cx, 0.0, fy, cy, 0.0, 0.0, 1.0], dtype=np.float64) + m.r = np.array([1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0], dtype=np.float64) + m.p = np.array([fx, 0.0, cx, 0.0, 0.0, fy, cy, 0.0, 0.0, 0.0, 1.0, 0.0], dtype=np.float64) m.binning_x = m.binning_y = 0 - return bytes(m.lcm_encode()) + return bytes(cdr_encode(m)) def _xyzrgb_fields() -> list[Any]: fields = [] for i, name in enumerate(["x", "y", "z"]): - f = PointField() + f = PointField(name="", offset=0, datatype=0, count=0) f.name, f.offset, f.datatype, f.count = name, i * 4, 7, 1 fields.append(f) - f = PointField() + f = PointField(name="", offset=0, datatype=0, count=0) f.name, f.offset, f.datatype, f.count = "rgb", 12, 7, 1 fields.append(f) return fields @@ -128,12 +157,21 @@ def _xyzrgb_fields() -> list[Any]: def cloud_msg(points: np.ndarray, colors: np.ndarray, frame_id: str, ts: float) -> bytes: """xyz + packed-rgb PointCloud2, matching dimos PointCloud2.lcm_encode's layout.""" - m = LCMPointCloud2() - m.header = Header() + m = PointCloud2( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + height=0, + width=0, + fields=[], + is_bigendian=False, + point_step=0, + row_step=0, + data=np.array([], dtype=np.uint8), + is_dense=False, + ) + m.header = Header(stamp=Time(sec=0, nanosec=0), frame_id="") _stamp(m.header, ts) m.header.frame_id = frame_id m.fields = _xyzrgb_fields() - m.fields_length = 4 m.point_step = 16 m.is_bigendian = False m.is_dense = True @@ -142,9 +180,8 @@ def cloud_msg(points: np.ndarray, colors: np.ndarray, frame_id: str, ts: float) if len(points) == 0: m.height = 0 m.row_step = 0 - m.data_length = 0 - m.data = b"" - return bytes(m.lcm_encode()) + m.data = np.array([], dtype=np.uint8) + return bytes(cdr_encode(m)) # ROS convention: rgb is a float32 whose bytes are [padding, r, g, b]. rgb_u32 = ( @@ -155,9 +192,8 @@ def cloud_msg(points: np.ndarray, colors: np.ndarray, frame_id: str, ts: float) data = np.column_stack([points.astype(np.float32), rgb_u32.view(np.float32)]).astype(np.float32) view = memoryview(np.ascontiguousarray(data)).cast("B") m.row_step = m.point_step * m.width - m.data_length = len(view) - m.data = view - return bytes(m.lcm_encode()) + m.data = np.frombuffer(view, dtype=np.uint8) + return bytes(cdr_encode(m)) def odometry_msg( @@ -168,8 +204,22 @@ def odometry_msg( child_frame_id: str, ts: float, ) -> bytes: - m = LCMOdometry() - m.header = Header() + m = Odometry( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + child_frame_id="", + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), + twist=TwistWithCovariance( + twist=Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + covariance=np.zeros(36, dtype=np.float64), + ), + ) + m.header = Header(stamp=Time(sec=0, nanosec=0), frame_id="") _stamp(m.header, ts) m.header.frame_id = frame_id m.child_frame_id = child_frame_id @@ -177,28 +227,38 @@ def odometry_msg( p.x, p.y, p.z = [float(v) for v in position] q = m.pose.pose.orientation q.x, q.y, q.z, q.w = [float(v) for v in quat_xyzw] - m.pose.covariance = [0.0] * 36 + m.pose.covariance = np.zeros(36, dtype=np.float64) m.twist.twist.linear.x = float(twist[0]) m.twist.twist.linear.y = float(twist[1]) m.twist.twist.angular.z = float(twist[2]) - m.twist.covariance = [0.0] * 36 - return bytes(m.lcm_encode()) + m.twist.covariance = np.zeros(36, dtype=np.float64) + return bytes(cdr_encode(m)) def tf_msg(links: list[tuple[str, str, Any, Any]], ts: float) -> bytes: """links: (parent, child, translation xyz, rotation xyzw).""" - m = LCMTFMessage() + m = TFMessage(transforms=[]) out = [] for parent, child, xyz, quat in links: - t = TransformStamped() - # Fresh nested messages: the generated bindings share defaults across instances. - t.header = Header() + t = TransformStamped( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + child_frame_id="", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + # Each edge owns independent generated nested values. + t.header = Header(stamp=Time(sec=0, nanosec=0), frame_id="") _stamp(t.header, ts) t.header.frame_id = parent t.child_frame_id = child - transform = Transform() - transform.translation = Vector3() - transform.rotation = Quaternion() + transform = Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ) + transform.translation = Vector3(x=0.0, y=0.0, z=0.0) + transform.rotation = Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) (transform.translation.x, transform.translation.y, transform.translation.z) = ( float(v) for v in xyz ) @@ -211,8 +271,7 @@ def tf_msg(links: list[tuple[str, str, Any, Any]], ts: float) -> bytes: t.transform = transform out.append(t) m.transforms = out - m.transforms_length = len(out) - return bytes(m.lcm_encode()) + return bytes(cdr_encode(m)) def _pixels( @@ -281,30 +340,34 @@ def visible_objects( def objects_msg(objects: Sequence[VisibleObject], ts: float) -> bytes: - m = LCMDetection3DArray() - m.header = Header() - _stamp(m.header, ts) - m.header.frame_id = "world" - m.detections = [] + def header() -> Header: + value = Header(stamp=Time(sec=0, nanosec=0), frame_id="world") + _stamp(value, ts) + return value + + detections = [] for label, center, size in objects: - # Generated constructors share default nested objects: build every level fresh. - d = Detection3D() - d.header = Header() - _stamp(d.header, ts) - d.header.frame_id = "world" - d.results = [ - ObjectHypothesisWithPose(hypothesis=ObjectHypothesis(class_id=label, score=1.0)) - ] - d.results_length = 1 cx, cy, cz = (float(v) for v in center) sx, sy, sz = (float(v) for v in size) - d.bbox = BoundingBox3D( - center=LCMPose(position=Vector3(x=cx, y=cy, z=cz), orientation=Quaternion(w=1.0)), - size=Vector3(x=sx, y=sy, z=sz), + pose = Pose( + position=Point(x=cx, y=cy, z=cz), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ) + detections.append( + Detection3D( + header=header(), + id=label, + results=[ + ObjectHypothesisWithPose( + hypothesis=ObjectHypothesis(class_id=label, score=1.0), + pose=PoseWithCovariance( + pose=pose, covariance=np.zeros(36, dtype=np.float64) + ), + ) + ], + bbox=BoundingBox3D(center=pose, size=Vector3(x=sx, y=sy, z=sz)), + ) ) - m.detections.append(d) - m.detections_length = len(m.detections) - return bytes(m.lcm_encode()) + return bytes(cdr_encode(Detection3DArray(header=header(), detections=detections))) class MotionType(str, Enum): @@ -499,7 +562,7 @@ def main() -> None: def on_cmd_vel(sample: Any) -> None: try: - t = LCMTwist.lcm_decode(bytes(sample.payload.to_bytes())) + t = cdr_decode(bytes(sample.payload.to_bytes()), Twist) cmd["vx"], cmd["vy"], cmd["wz"] = t.linear.x, t.linear.y, t.angular.z cmd["ts"] = time.time() except Exception as exc: diff --git a/dimos/simulation/habitat/test_server.py b/dimos/simulation/habitat/test_server.py index bf45ef89e3e..b59382f706e 100644 --- a/dimos/simulation/habitat/test_server.py +++ b/dimos/simulation/habitat/test_server.py @@ -12,23 +12,25 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""The native encodes with dimos_lcm; dimos decodes. These pin that contract. +"""The native encodes with generated CDR; dimos decodes. These pin that contract. -server.py runs under python 3.9 in nix/env, but its encoders are pure dimos_lcm +server.py runs under python 3.9 in nix/env, but its encoders are standalone generated CDR and import fine here, so the wire format is testable without habitat installed. """ import ast from pathlib import Path +from dimos_cdrgen.registry import decode as cdr_decode +from dimos_generated.nav_msgs.msg import Odometry +from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 +from dimos_generated.tf2_msgs.msg import TFMessage +from dimos_generated.vision_msgs.msg import Detection3DArray import numpy as np +from rosbags.typesys import Stores, get_typestore -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.msgs.image import image_view +from dimos.msgs.pointcloud import pointcloud_xyz from dimos.simulation.habitat import server SERVER_PY = Path(server.__file__) @@ -36,21 +38,21 @@ def test_color_image_round_trip(): rgb = np.arange(8 * 12 * 3, dtype=np.uint8).reshape(8, 12, 3) - out = Image.lcm_decode(server.image_msg(rgb, "rgb8", "camera_optical", 1.5)) - np.testing.assert_array_equal(out.data, rgb) - assert out.frame_id == "camera_optical" + out = cdr_decode(server.image_msg(rgb, "rgb8", "camera_optical", 1.5), Image) + np.testing.assert_array_equal(image_view(out), rgb) + assert out.header.frame_id == "camera_optical" def test_depth_image_round_trip(): depth = np.linspace(0.2, 4.5, 8 * 12, dtype=np.float32).reshape(8, 12) - out = Image.lcm_decode(server.image_msg(depth, "32FC1", "camera_optical", 1.5)) - np.testing.assert_allclose(out.data, depth, atol=1e-6) + out = cdr_decode(server.image_msg(depth, "32FC1", "camera_optical", 1.5), Image) + np.testing.assert_allclose(image_view(out), depth, atol=1e-6) def test_camera_info_round_trip(): k = {"fx": 320.0, "fy": 320.0, "cx": 320.0, "cy": 180.0, "width": 640.0, "height": 360.0} - out = CameraInfo.lcm_decode(server.camera_info_msg(k, "camera_optical", 1.5)) - matrix = out.get_K_matrix() + out = cdr_decode(server.camera_info_msg(k, "camera_optical", 1.5), CameraInfo) + matrix = np.asarray(out.k).reshape(3, 3) assert (out.width, out.height) == (640, 360) np.testing.assert_allclose([matrix[0, 0], matrix[1, 1]], [320.0, 320.0]) np.testing.assert_allclose([matrix[0, 2], matrix[1, 2]], [320.0, 180.0]) @@ -59,31 +61,33 @@ def test_camera_info_round_trip(): def test_cloud_round_trip_keeps_points_and_colour(): pts = np.array([[1.0, 2.0, 3.0], [-4.0, 0.5, 0.25]], dtype=np.float32) colors = np.array([[255, 0, 0], [0, 128, 64]], dtype=np.uint8) - out = PointCloud2.lcm_decode(server.cloud_msg(pts, colors, "world", 1.5)) - got = out.points_f32() - assert out.frame_id == "world" + out = cdr_decode(server.cloud_msg(pts, colors, "world", 1.5), PointCloud2) + got = pointcloud_xyz(out) + assert out.header.frame_id == "world" np.testing.assert_allclose(np.sort(got, axis=0), np.sort(pts, axis=0), atol=1e-5) def test_empty_cloud_is_valid(): - out = PointCloud2.lcm_decode( - server.cloud_msg(np.zeros((0, 3), np.float32), np.zeros((0, 3), np.uint8), "world", 1.0) + out = cdr_decode( + server.cloud_msg(np.zeros((0, 3), np.float32), np.zeros((0, 3), np.uint8), "world", 1.0), + PointCloud2, ) - assert len(out) == 0 + assert out.width == 0 and len(out.data) == 0 def test_odometry_round_trip(): pos = np.array([1.0, -2.0, 0.5]) quat = np.array([0.0, 0.0, 0.3826834, 0.9238795]) # 45 deg yaw - out = Odometry.lcm_decode( - server.odometry_msg(pos, quat, (0.4, 0.0, 0.2), "world", "base_link", 1.5) + out = cdr_decode( + server.odometry_msg(pos, quat, (0.4, 0.0, 0.2), "world", "base_link", 1.5), + Odometry, ) np.testing.assert_allclose( [out.pose.pose.position.x, out.pose.pose.position.y, out.pose.pose.position.z], pos, atol=1e-6, ) - assert out.frame_id == "world" + assert out.header.frame_id == "world" assert out.child_frame_id == "base_link" @@ -95,7 +99,23 @@ def test_objects_round_trip_keeps_each_object(): ], 1.5, ) - assert Detection3DArray.lcm_decode(msg).to_json() == [ + out = cdr_decode(msg, Detection3DArray) + assert out.header.frame_id == "world" + assert [(d.id, d.results[0].hypothesis.class_id) for d in out.detections] == [ + ("chair", "chair"), + ("table", "table"), + ] + assert [ + dict( + label=d.id, + score=d.results[0].hypothesis.score, + position=dict( + x=d.bbox.center.position.x, y=d.bbox.center.position.y, z=d.bbox.center.position.z + ), + size=dict(x=d.bbox.size.x, y=d.bbox.size.y, z=d.bbox.size.z), + ) + for d in out.detections + ] == [ { "label": "chair", "score": 1.0, @@ -128,14 +148,14 @@ def test_tf_round_trip_carries_the_optical_link(): ("base_link", "camera", (0.0, 0.0, 0.45), (0.0, 0.0, 0.0, 1.0)), ("camera", "camera_optical", (0.0, 0.0, 0.0), server.frames.OPTICAL_QUAT_XYZW), ] - out = TFMessage.lcm_decode(server.tf_msg(links, 1.5)) - pairs = {(t.frame_id, t.child_frame_id) for t in out.transforms} + out = cdr_decode(server.tf_msg(links, 1.5), TFMessage) + pairs = {(t.header.frame_id, t.child_frame_id) for t in out.transforms} assert ("camera", "camera_optical") in pairs assert len(out.transforms) == 3 def test_each_tf_link_keeps_its_own_pose(): - """Regression: the generated bindings share nested defaults between instances. + """Regression: TF edges must retain independent nested transform values. Mutating ``Transform().translation`` in place made all three links serialize the last one's pose, which showed up as the whole map pointing up. Only a @@ -147,28 +167,42 @@ def test_each_tf_link_keeps_its_own_pose(): ("camera", "camera_optical", (0.0, 0.0, 0.0), server.frames.OPTICAL_QUAT_XYZW), ] by_child = { - t.child_frame_id: t for t in TFMessage.lcm_decode(server.tf_msg(links, 1.5)).transforms + t.child_frame_id: t for t in cdr_decode(server.tf_msg(links, 1.5), TFMessage).transforms } base = by_child["base_link"] np.testing.assert_allclose( - [base.translation.x, base.translation.y, base.translation.z], [1.0, 2.0, 3.0] + [base.transform.translation.x, base.transform.translation.y, base.transform.translation.z], + [1.0, 2.0, 3.0], ) np.testing.assert_allclose( - [base.rotation.x, base.rotation.y, base.rotation.z, base.rotation.w], [0, 0, 0, 1] + [ + base.transform.rotation.x, + base.transform.rotation.y, + base.transform.rotation.z, + base.transform.rotation.w, + ], + [0, 0, 0, 1], ) cam = by_child["camera"] np.testing.assert_allclose( - [cam.translation.x, cam.translation.y, cam.translation.z], [0.0, 0.0, 0.45] + [cam.transform.translation.x, cam.transform.translation.y, cam.transform.translation.z], + [0.0, 0.0, 0.45], ) opt = by_child["camera_optical"] np.testing.assert_allclose( - [opt.translation.x, opt.translation.y, opt.translation.z], [0.0, 0.0, 0.0] + [opt.transform.translation.x, opt.transform.translation.y, opt.transform.translation.z], + [0.0, 0.0, 0.0], ) np.testing.assert_allclose( - [opt.rotation.x, opt.rotation.y, opt.rotation.z, opt.rotation.w], + [ + opt.transform.rotation.x, + opt.transform.rotation.y, + opt.transform.rotation.z, + opt.transform.rotation.w, + ], server.frames.OPTICAL_QUAT_XYZW, ) @@ -182,11 +216,16 @@ def test_tf_chain_puts_the_camera_looking_forward(): ("camera", "camera_optical", (0.0, 0.0, 0.0), server.frames.OPTICAL_QUAT_XYZW), ] by_child = { - t.child_frame_id: t for t in TFMessage.lcm_decode(server.tf_msg(links, 1.5)).transforms + t.child_frame_id: t for t in cdr_decode(server.tf_msg(links, 1.5), TFMessage).transforms } def rot(t): - x, y, z, w = t.rotation.x, t.rotation.y, t.rotation.z, t.rotation.w + x, y, z, w = ( + t.transform.rotation.x, + t.transform.rotation.y, + t.transform.rotation.z, + t.transform.rotation.w, + ) return np.array( [ [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)], @@ -262,3 +301,56 @@ def test_every_frames_attribute_used_by_server_exists(): } missing = sorted(name for name in used if not hasattr(frames_mod, name)) assert not missing, f"server.py calls frames.{missing} which frames.py does not define" + + +def test_habitat_payloads_decode_with_independent_ros_types(): + types = get_typestore(Stores.ROS2_HUMBLE) + packets = [ + ( + "sensor_msgs/msg/Image", + server.image_msg(np.full((2, 3, 3), 7, dtype=np.uint8), "rgb8", "camera", 12.25), + ), + ( + "sensor_msgs/msg/CameraInfo", + server.camera_info_msg( + {"width": 3, "height": 2, "fx": 10, "fy": 11, "cx": 1, "cy": 1}, "camera", 12.25 + ), + ), + ( + "sensor_msgs/msg/PointCloud2", + server.cloud_msg( + np.array([[1, 2, 3]], dtype=np.float32), + np.array([[255, 128, 64]], dtype=np.uint8), + "world", + 12.25, + ), + ), + ( + "nav_msgs/msg/Odometry", + server.odometry_msg( + np.array([1, 2, 3]), + np.array([0, 0, 0, 1]), + (0.4, 0.2, 0.3), + "world", + "base_link", + 12.25, + ), + ), + ( + "tf2_msgs/msg/TFMessage", + server.tf_msg([("world", "base_link", (1, 2, 3), (0, 0, 0, 1))], 12.25), + ), + ] + decoded = {} + for typename, packet in packets: + message = types.deserialize_cdr(packet, typename) + header = ( + message.transforms[0].header if typename == "tf2_msgs/msg/TFMessage" else message.header + ) + assert (header.stamp.sec, header.stamp.nanosec) == (12, 250000000) + decoded[typename] = message + assert len(decoded["sensor_msgs/msg/Image"].data) == 18 + assert decoded["sensor_msgs/msg/CameraInfo"].k[0] == 10.0 + assert decoded["sensor_msgs/msg/PointCloud2"].row_step == 16 + assert decoded["nav_msgs/msg/Odometry"].pose.pose.position.z == 3.0 + assert decoded["tf2_msgs/msg/TFMessage"].transforms[0].transform.translation.x == 1.0 diff --git a/dimos/simulation/mujoco/direct_cmd_vel_explorer.py b/dimos/simulation/mujoco/direct_cmd_vel_explorer.py index 3b012c4092f..22f10726f19 100644 --- a/dimos/simulation/mujoco/direct_cmd_vel_explorer.py +++ b/dimos/simulation/mujoco/direct_cmd_vel_explorer.py @@ -76,7 +76,9 @@ def _normalize_angle(angle: float) -> float: def _stop(self) -> None: assert self._cmd_vel is not None - self._cmd_vel.broadcast(None, Twist(linear=Vector3(), angular=Vector3())) + self._cmd_vel.broadcast( + None, Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) def _drive_to(self, target_x: float, target_y: float) -> None: """Pursuit controller: steer toward the target while driving forward.""" @@ -98,7 +100,9 @@ def _drive_to(self, target_x: float, target_y: float) -> None: assert self._cmd_vel is not None self._cmd_vel.broadcast( None, - Twist(linear=Vector3(x=linear), angular=Vector3(z=angular)), + Twist( + linear=Vector3(x=linear, y=0.0, z=0.0), angular=Vector3(z=angular, x=0.0, y=0.0) + ), ) self._stop() diff --git a/dimos/simulation/mujoco/model.py b/dimos/simulation/mujoco/model.py index ff52a8e4444..58829e757ef 100644 --- a/dimos/simulation/mujoco/model.py +++ b/dimos/simulation/mujoco/model.py @@ -19,12 +19,14 @@ from typing import Any import xml.etree.ElementTree as ET +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.std_msgs.msg import Header import mujoco import numpy as np from dimos.core.global_config import GlobalConfig from dimos.mapping.occupancy.extrude_occupancy import generate_mujoco_scene -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid +from dimos.msgs.occupancy import occupancy_from_file from dimos.simulation.mujoco.input_controller import InputController from dimos.simulation.mujoco.menagerie import menagerie_path from dimos.simulation.mujoco.policy import G1OnnxController, Go1OnnxController, OnnxController @@ -149,7 +151,9 @@ def _add_person_object(root: ET.Element) -> None: def load_scene_xml(config: GlobalConfig) -> str: if config.mujoco_room_from_occupancy: path = Path(config.mujoco_room_from_occupancy) - return generate_mujoco_scene(OccupancyGrid.from_path(path)) + return generate_mujoco_scene( + occupancy_from_file(path, header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0))) + ) mujoco_room = config.mujoco_room or "office_555" if mujoco_room == "office_555": diff --git a/dimos/simulation/mujoco/person_on_track.py b/dimos/simulation/mujoco/person_on_track.py index ec63ceebc6c..9a311493d7a 100644 --- a/dimos/simulation/mujoco/person_on_track.py +++ b/dimos/simulation/mujoco/person_on_track.py @@ -12,11 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.geometry_msgs.msg import Pose import mujoco from dimos.core.transport import PubSubTransport from dimos.core.transport_factory import make_transport -from dimos.msgs.geometry_msgs.Pose import Pose class PersonPositionController: diff --git a/dimos/simulation/mujoco/test_generated_control_consumers.py b/dimos/simulation/mujoco/test_generated_control_consumers.py index bd54c2e62db..a579d88c10c 100644 --- a/dimos/simulation/mujoco/test_generated_control_consumers.py +++ b/dimos/simulation/mujoco/test_generated_control_consumers.py @@ -14,6 +14,8 @@ from unittest.mock import MagicMock +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Twist from dimos_generated.std_msgs.msg import Header import numpy as np @@ -36,18 +38,26 @@ def test_pursuit_generated_pose_preserves_steering_and_stop( poses = iter( [ PoseStamped( - header=Header(frame_id="world"), - pose=Pose(position=Point(), orientation=quaternion_from_euler(0.0, 0.0, heading)), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=quaternion_from_euler(0.0, 0.0, heading), + ), ), PoseStamped( - header=Header(frame_id="world"), - pose=Pose(position=Point(x=1.0), orientation=Quaternion(w=1.0)), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), + pose=Pose( + position=Point(x=1.0, y=0.0, z=0.0), + orientation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), + ), ), ] ) monkeypatch.setattr(explorer, "_wait_for_pose", lambda: next(poses)) explorer._drive_to(1.0, 0.0) - commands = [Twist.decode(call.args[1].encode()) for call in output.broadcast.call_args_list] + commands = [ + cdr_decode(cdr_encode(call.args[1]), Twist) for call in output.broadcast.call_args_list + ] assert len(commands) == 2 assert commands[0].linear.x == expected_linear assert commands[0].angular.z == expected_angular @@ -64,7 +74,7 @@ def test_person_track_generated_pose_preserves_position_and_heading( ) publisher = PersonTrackPublisher([(1.0, 2.0), (1.0, 3.0), (2.0, 3.0)]) publisher.tick() - pose = Pose.decode(output.broadcast.call_args.args[1].encode()) + pose = cdr_decode(cdr_encode(output.broadcast.call_args.args[1]), Pose) assert (pose.position.x, pose.position.y, pose.position.z) == (1.0, 2.0, 0.0) assert np.isclose(pose.orientation.z, np.sin(np.pi / 4)) assert np.isclose(pose.orientation.w, np.cos(np.pi / 4)) diff --git a/dimos/simulation/scenes/procedural.py b/dimos/simulation/scenes/procedural.py index 69211db9821..a5ba7253e4d 100644 --- a/dimos/simulation/scenes/procedural.py +++ b/dimos/simulation/scenes/procedural.py @@ -26,7 +26,7 @@ import numpy as np from numpy.typing import ArrayLike, NDArray -from dimos.msgs.sim_msgs.Contacts import Kind +from dimos.msgs.contacts import Kind Vec3 = tuple[float, float, float] Rect = tuple[float, float, float, float] diff --git a/dimos/teleop/hosted/arm_command.py b/dimos/teleop/hosted/arm_command.py index 04d1a272016..06fea2a3bfd 100644 --- a/dimos/teleop/hosted/arm_command.py +++ b/dimos/teleop/hosted/arm_command.py @@ -24,6 +24,7 @@ import time from typing import Any +from dimos_cdrgen.registry import decode as cdr_decode from dimos_generated.geometry_msgs.msg import PoseStamped, Twist, TwistStamped, Vector3 from dimos_generated.std_msgs.msg import Float32, Header from reactivex.disposable import Disposable @@ -111,7 +112,7 @@ def _on_cmd_raw(self, data: Any) -> None: def _on_pose_bytes(self, data: bytes) -> None: """Controller pose → robot frame; stale/future/out-of-order dropped.""" - msg = PoseStamped.decode(data) + msg = cdr_decode(data, PoseStamped) try: hand = self._resolve_hand(msg.header.frame_id) except ValueError: @@ -143,7 +144,7 @@ def _on_twist_bytes(self, data: bytes) -> None: """Browser keyboard EE-twist → coordinator's eef_twist task.""" if self._estopped: return - msg = TwistStamped.decode(data) + msg = cdr_decode(data, TwistStamped) ts = to_seconds(msg.header.stamp) if not math.isfinite(ts): return @@ -167,7 +168,7 @@ def _on_twist_bytes(self, data: bytes) -> None: scale = self._translation_scale self.ee_twist_command.publish( TwistStamped( - header=Header(stamp=msg.header.stamp), + header=Header(stamp=msg.header.stamp, frame_id=""), twist=Twist( linear=Vector3( x=msg.twist.linear.x * scale, diff --git a/dimos/teleop/hosted/blueprints/cloudflare.py b/dimos/teleop/hosted/blueprints/cloudflare.py index cadd369279d..177f7cb2a9a 100644 --- a/dimos/teleop/hosted/blueprints/cloudflare.py +++ b/dimos/teleop/hosted/blueprints/cloudflare.py @@ -26,6 +26,9 @@ from __future__ import annotations +from dimos_generated.geometry_msgs.msg import Twist, TwistStamped +from dimos_generated.sensor_msgs.msg import Image + from dimos.core.coordination.blueprints import autoconnect from dimos.core.transport import ( CloudflareAudioTransport, @@ -36,9 +39,6 @@ from dimos.hardware.sensors.camera.realsense.camera import RealSenseCamera from dimos.mapping.costmapper import CostMapper from dimos.mapping.voxels.module import VoxelGridMapper -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped -from dimos.msgs.sensor_msgs.Image import Image from dimos.navigation.go2.replanning_a_star.module import ReplanningAStarPlanner from dimos.navigation.movement_manager.movement_manager import MovementManager from dimos.robot.manipulators.xarm.blueprints.teleop import ( diff --git a/dimos/teleop/hosted/camera_mux.py b/dimos/teleop/hosted/camera_mux.py index 7d5dc07dd23..a12ce7c7aeb 100644 --- a/dimos/teleop/hosted/camera_mux.py +++ b/dimos/teleop/hosted/camera_mux.py @@ -23,6 +23,7 @@ from __future__ import annotations from collections.abc import Callable +from copy import deepcopy import json import threading import time @@ -160,7 +161,9 @@ def _composite(self) -> Image | None: image_from_array( np.hstack(tiles), encoding=imgs[0].encoding, - header=Header(stamp=latest.header.stamp, frame_id="camera_mux"), + header=Header( + stamp=deepcopy(latest.header.stamp), frame_id="camera_mux" + ), ) ) ) @@ -181,7 +184,7 @@ def _even_dims(img: Image) -> Image: return img data = data[: h - (h % 2), : w - (w % 2)] return image_from_array( - np.ascontiguousarray(data), encoding=img.encoding, header=img.header + np.ascontiguousarray(data), encoding=img.encoding, header=deepcopy(img.header) ) def _downscale(self, img: Image) -> Image: @@ -196,7 +199,7 @@ def _downscale(self, img: Image) -> Image: if w <= max_w: return img out = cv2.resize(image_view(img), (max_w, max(1, int(h * max_w / w)))) - return image_from_array(out, encoding=img.encoding, header=img.header) + return image_from_array(out, encoding=img.encoding, header=deepcopy(img.header)) def _stamp(self, img: Image) -> Image: """Append (not overwrite) a bottom strip encoding capture time as B/W @@ -228,7 +231,7 @@ def _stamp(self, img: Image) -> Image: if bit: strip[:, i * s : (i + 1) * s] = 255 out = np.vstack([data, strip]) - return image_from_array(out, encoding=img.encoding, header=img.header) + return image_from_array(out, encoding=img.encoding, header=deepcopy(img.header)) def _set_cam_selection(self, data: bytes) -> None: """camera_select kind → filter to known cams, republish immediately so diff --git a/dimos/teleop/hosted/go2_command.py b/dimos/teleop/hosted/go2_command.py index 5ac70951f22..9250d34fef1 100644 --- a/dimos/teleop/hosted/go2_command.py +++ b/dimos/teleop/hosted/go2_command.py @@ -365,7 +365,9 @@ def _handle_nav_goal(self, msg: dict[str, Any]) -> None: return pose = PoseStamped( header=header_now("world"), - pose=Pose(position=Point(x=x, y=y), orientation=Quaternion(w=1)), + pose=Pose( + position=Point(x=x, y=y, z=0.0), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), ) try: self.goal_request.publish(pose) @@ -383,7 +385,7 @@ def _handle_nav_cancel(self, nonce: Any) -> None: def _cancel_nav(self) -> None: try: - msg = Bool() + msg = Bool(data=False) msg.data = True self.stop_movement.publish(msg) except Exception: diff --git a/dimos/teleop/hosted/hosted_stats.py b/dimos/teleop/hosted/hosted_stats.py index 3b8bdd66fdb..dcb73b889c3 100644 --- a/dimos/teleop/hosted/hosted_stats.py +++ b/dimos/teleop/hosted/hosted_stats.py @@ -28,6 +28,7 @@ import time from typing import Any +from dimos_cdrgen.registry import decode as cdr_decode from dimos_generated.dimos_msgs.msg import VideoStats from dimos_generated.geometry_msgs.msg import TwistStamped from reactivex.disposable import Disposable @@ -126,7 +127,7 @@ def _on_cmd_raw(self, data: Any) -> None: if isinstance(data, str): data = data.encode() try: - cmd = TwistStamped.decode(data) + cmd = cdr_decode(data, TwistStamped) source_time = to_seconds(cmd.header.stamp) except Exception: return diff --git a/dimos/teleop/hosted/test_arm_command.py b/dimos/teleop/hosted/test_arm_command.py index fa194ceed3c..f42f2f3394b 100644 --- a/dimos/teleop/hosted/test_arm_command.py +++ b/dimos/teleop/hosted/test_arm_command.py @@ -30,7 +30,16 @@ from typing import Any from unittest.mock import MagicMock, patch -from dimos_generated.geometry_msgs.msg import PoseStamped, Twist, TwistStamped, Vector3 +from dimos_cdrgen.registry import encode as cdr_encode +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Twist, + TwistStamped, + Vector3, +) from dimos_generated.std_msgs.msg import Header import pytest @@ -81,9 +90,9 @@ def _frame(channel, message): seq=0, ts=time.time(), delivery="latest", - meta={"encoding": "cdr", "type": message.msg_name}, + meta={"encoding": "cdr", "type": message.__msgtype__}, ), - message.encode(), + cdr_encode(message), ) @@ -93,7 +102,11 @@ def _pose_bytes(frame_id: str, ts: float | None = None) -> bytes: PoseStamped( header=Header( stamp=time_from_seconds(time.time() if ts is None else ts), frame_id=frame_id - ) + ), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), ), ) @@ -105,7 +118,9 @@ def _twist_bytes(x: float = 0.1, angular_x: float = 0.0, ts: float | None = None header=Header( stamp=time_from_seconds(time.time() if ts is None else ts), frame_id="eef_twist_arm" ), - twist=Twist(linear=Vector3(x=x), angular=Vector3(x=angular_x)), + twist=Twist( + linear=Vector3(x=x, y=0.0, z=0.0), angular=Vector3(x=angular_x, y=0.0, z=0.0) + ), ), ) diff --git a/dimos/teleop/hosted/test_camera_mux.py b/dimos/teleop/hosted/test_camera_mux.py index 35f8e811b63..591c5f0b6ba 100644 --- a/dimos/teleop/hosted/test_camera_mux.py +++ b/dimos/teleop/hosted/test_camera_mux.py @@ -24,6 +24,8 @@ import json from types import SimpleNamespace +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import Image from dimos_generated.std_msgs.msg import Header import numpy as np @@ -215,10 +217,10 @@ def test_generated_crop_scale_stamp_preserve_zero_source_header() -> None: source = image_from_array( np.zeros((721, 1280, 3), np.uint8), encoding="bgr8", - header=Header(frame_id="camera_optical"), + header=Header(frame_id="camera_optical", stamp=Time(sec=0, nanosec=0)), ) - mux._on_cam("cam1", Image.decode(source.encode())) - output = Image.decode(mux.published[0].encode()) + mux._on_cam("cam1", cdr_decode(cdr_encode(source), Image)) + output = cdr_decode(cdr_encode(mux.published[0]), Image) assert output.width % 2 == 0 and output.height % 2 == 0 assert output.header.frame_id == "camera_optical" assert to_nanoseconds(output.header.stamp) == 0 diff --git a/dimos/teleop/hosted/test_go2_command.py b/dimos/teleop/hosted/test_go2_command.py index d604791ac92..9519601a609 100644 --- a/dimos/teleop/hosted/test_go2_command.py +++ b/dimos/teleop/hosted/test_go2_command.py @@ -29,6 +29,7 @@ from typing import Any from unittest.mock import MagicMock +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import PoseStamped, Twist, TwistStamped, Vector3 from dimos_generated.std_msgs.msg import Header import pytest @@ -64,7 +65,10 @@ def module(monkeypatch: pytest.MonkeyPatch) -> Iterator[Go2CommandModule]: def _twist(ts: float, *, vx: float = 0.3) -> TwistStamped: """A drive frame at time ``ts`` (vx=0.3 moving, vx=0 idle-joystick).""" - t = TwistStamped(header=Header(stamp=time_from_seconds(ts)), twist=Twist(linear=Vector3(x=vx))) + t = TwistStamped( + header=Header(stamp=time_from_seconds(ts), frame_id=""), + twist=Twist(linear=Vector3(x=vx, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + ) return t @@ -191,9 +195,9 @@ def test_drive_drops_invalid_nanoseconds(module: Go2CommandModule) -> None: def test_drive_orders_frames_one_nanosecond_apart(module: Go2CommandModule) -> None: first = _twist(time.time() - 0.1) stamp_ns = to_nanoseconds(first.header.stamp) - second = TwistStamped.decode(first.encode()) + second = cdr_decode(cdr_encode(first), TwistStamped) second.header.stamp = time_from_nanoseconds(stamp_ns + 1) - module._on_cmd_vel_in(TwistStamped.decode(first.encode())) + module._on_cmd_vel_in(cdr_decode(cdr_encode(first), TwistStamped)) module._on_cmd_vel_in(second) assert module.tele_cmd_vel.publish.call_count == 2 assert module._last_cmd_ns == stamp_ns + 1 @@ -390,7 +394,7 @@ def test_nav_goal_publishes_and_acks( (pose,) = module.goal_request.publish.call_args.args assert pose.pose.position.x == pytest.approx(2.5) - assert PoseStamped.decode(pose.encode()).header.frame_id == "world" + assert cdr_decode(cdr_encode(pose), PoseStamped).header.frame_id == "world" assert acks == [(11, True)] diff --git a/dimos/teleop/hosted/test_hosted_stats.py b/dimos/teleop/hosted/test_hosted_stats.py index 78ec00b76d2..262770eef13 100644 --- a/dimos/teleop/hosted/test_hosted_stats.py +++ b/dimos/teleop/hosted/test_hosted_stats.py @@ -21,6 +21,7 @@ from types import SimpleNamespace from unittest.mock import MagicMock +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import Twist, TwistStamped, Vector3 from dimos_generated.std_msgs.msg import Header import pytest @@ -44,9 +45,12 @@ def module(monkeypatch: pytest.MonkeyPatch) -> Iterator[HostedStatsModule]: def _cmd(vx: float = 0.3, ts: float = 123.0) -> bytes: - return TwistStamped( - header=Header(stamp=time_from_seconds(ts)), twist=Twist(linear=Vector3(x=vx)) - ).encode() + return cdr_encode( + TwistStamped( + header=Header(stamp=time_from_seconds(ts), frame_id=""), + twist=Twist(linear=Vector3(x=vx, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + ) + ) def test_telemetry_without_go2(monkeypatch: pytest.MonkeyPatch) -> None: @@ -100,7 +104,7 @@ def test_state_json_ignores_foreign_kind(module: HostedStatsModule) -> None: def test_cmd_raw_rejects_invalid_ros_nanoseconds(module: HostedStatsModule) -> None: - message = TwistStamped.decode(_cmd()) + message = cdr_decode(_cmd(), TwistStamped) message.header.stamp.nanosec = 1_000_000_000 - module._on_cmd_raw(message.encode()) + module._on_cmd_raw(cdr_encode(message)) module.cmd_vel_stamped.publish.assert_not_called() diff --git a/dimos/teleop/hosted/test_map_compress.py b/dimos/teleop/hosted/test_map_compress.py index 3dde277789b..96bdc1b772a 100644 --- a/dimos/teleop/hosted/test_map_compress.py +++ b/dimos/teleop/hosted/test_map_compress.py @@ -24,6 +24,8 @@ from unittest.mock import MagicMock import cv2 +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid from dimos_generated.std_msgs.msg import Header @@ -61,10 +63,20 @@ def _published_json(mock: MagicMock, msg_type: str) -> dict[str, Any] | None: def _occupancy(grid: Any, resolution: float = 0.1) -> OccupancyGrid: cells = np.asarray(grid, dtype=np.int8) message = OccupancyGrid( - info=MapMetaData(width=cells.shape[1], height=cells.shape[0], resolution=resolution), - data=cells.ravel(), + info=MapMetaData( + width=cells.shape[1], + height=cells.shape[0], + resolution=resolution, + map_load_time=Time(sec=0, nanosec=0), + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.asarray(cells.ravel(), dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) - return OccupancyGrid.decode(message.encode()) + return cdr_decode(cdr_encode(message), OccupancyGrid) def test_costmap_encodes_and_publishes_map(module: MapCompressModule) -> None: @@ -119,7 +131,7 @@ def test_block_max_preserves_obstacle_when_coarsening(module: MapCompressModule) def test_odom_publishes_planar_pose(module: MapCompressModule) -> None: q = quaternion_from_euler(0.0, 0.0, math.pi / 2) # yaw = 90° pose = PoseStamped( - header=Header(stamp=time_from_seconds(123)), + header=Header(stamp=time_from_seconds(123), frame_id=""), pose=Pose(position=Point(x=1.5, y=-2, z=0.3), orientation=q), ) module._on_odom(pose) @@ -132,14 +144,34 @@ def test_odom_publishes_planar_pose(module: MapCompressModule) -> None: def test_empty_costmap_publishes_nothing(module: MapCompressModule) -> None: - module._on_costmap(OccupancyGrid()) # no-arg = empty 1D grid; must be skipped + module._on_costmap( + OccupancyGrid( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + info=MapMetaData( + map_load_time=Time(sec=0, nanosec=0), + resolution=0.0, + width=0, + height=0, + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([], dtype=np.int8), + ) + ) # no-arg = empty 1D grid; must be skipped assert _published_json(module.map_out, "map") is None def test_odom_degenerate_quaternion_does_not_raise(module: MapCompressModule) -> None: # A zero quaternion makes to_euler() (scipy) raise; _on_odom runs inside an # RxPY subscriber, so it must drop the frame, not kill the odom stream. - pose = PoseStamped(pose=Pose(orientation=Quaternion(w=0))) + pose = PoseStamped( + pose=Pose( + orientation=Quaternion(w=0, x=0.0, y=0.0, z=0.0), position=Point(x=0.0, y=0.0, z=0.0) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) module._on_odom(pose) # must not raise assert _published_json(module.map_out, "odom") is None @@ -153,7 +185,22 @@ def test_oversized_map_dropped(module: MapCompressModule) -> None: def test_malformed_grid_does_not_break_next_frame(module): - module._on_costmap(OccupancyGrid(info=MapMetaData(width=2, height=2, resolution=0.1), data=[0])) + module._on_costmap( + OccupancyGrid( + info=MapMetaData( + width=2, + height=2, + resolution=0.1, + map_load_time=Time(sec=0, nanosec=0), + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([0], dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + ) module.map_out.publish.assert_not_called() module._on_costmap(_occupancy([[0, 100]])) assert _published_json(module.map_out, "map")["w"] == 2 diff --git a/dimos/teleop/keyboard/test_keyboard_teleop_module.py b/dimos/teleop/keyboard/test_keyboard_teleop_module.py index 1a1ba1a680c..150c6bcc785 100644 --- a/dimos/teleop/keyboard/test_keyboard_teleop_module.py +++ b/dimos/teleop/keyboard/test_keyboard_teleop_module.py @@ -16,6 +16,7 @@ from collections.abc import Iterator +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import TwistStamped import pytest @@ -50,7 +51,7 @@ def test_publish_twist_emits_unaddressed_twist_stamped( module._publish_twist(linear=(0.1, 0.2, 0.3), angular=(0.4, 0.5, 0.6)) - msg = TwistStamped.decode(publish.call_args.args[0].encode()) + msg = cdr_decode(cdr_encode(publish.call_args.args[0]), TwistStamped) assert isinstance(msg, TwistStamped) assert msg.header.frame_id == "" # no task-name address in the payload assert [msg.twist.linear.x, msg.twist.linear.y, msg.twist.linear.z] == [0.1, 0.2, 0.3] @@ -62,7 +63,7 @@ def test_publish_twist_defaults_to_zero_twist(module: KeyboardTeleopModule, mock module._publish_twist() - msg = TwistStamped.decode(publish.call_args.args[0].encode()) + msg = cdr_decode(cdr_encode(publish.call_args.args[0]), TwistStamped) assert [msg.twist.linear.x, msg.twist.linear.y, msg.twist.linear.z] == [0.0, 0.0, 0.0] assert [msg.twist.angular.x, msg.twist.angular.y, msg.twist.angular.z] == [0.0, 0.0, 0.0] @@ -98,7 +99,7 @@ def test_final_key_release_publishes_zero_velocity(module: KeyboardTeleopModule, assert held == set() assert publish.call_count == 1 - msg = TwistStamped.decode(publish.call_args.args[0].encode()) + msg = cdr_decode(cdr_encode(publish.call_args.args[0]), TwistStamped) assert [msg.twist.linear.x, msg.twist.linear.y, msg.twist.linear.z] == [0.0, 0.0, 0.0] assert [msg.twist.angular.x, msg.twist.angular.y, msg.twist.angular.z] == [0.0, 0.0, 0.0] @@ -112,7 +113,7 @@ def test_keyup_preserves_remaining_motion_key(module: KeyboardTeleopModule, mock assert held == {keyboard_mod.pygame.K_a} assert publish.call_count == 1 - msg = TwistStamped.decode(publish.call_args.args[0].encode()) + msg = cdr_decode(cdr_encode(publish.call_args.args[0]), TwistStamped) assert [msg.twist.linear.x, msg.twist.linear.y, msg.twist.linear.z] == [0.0, 0.05, 0.0] diff --git a/dimos/teleop/phone/phone_teleop_module.py b/dimos/teleop/phone/phone_teleop_module.py index 2417bcf529c..654c78e9ffe 100644 --- a/dimos/teleop/phone/phone_teleop_module.py +++ b/dimos/teleop/phone/phone_teleop_module.py @@ -27,6 +27,7 @@ import time from typing import Any +from dimos_cdrgen.registry import decode as cdr_decode, schema as cdr_schema from dimos_generated.geometry_msgs.msg import Twist, TwistStamped, Vector3 from dimos_generated.std_msgs.msg import Bool, Header from fastapi import WebSocket, WebSocketDisconnect @@ -108,7 +109,10 @@ async def teleop_index() -> HTMLResponse: @self._web_server.app.get("/teleop/schema") async def command_schemas() -> dict[str, dict[str, str]]: return { - channel: {"type": message_type.msg_name, "definition": message_type.schema} + channel: { + "type": message_type.__msgtype__, + "definition": cdr_schema(message_type.__msgtype__), + } for channel, message_type in self._command_types.items() } @@ -142,7 +146,7 @@ def _dispatch_binary_message(self, data: bytes) -> bool: if ( message_type is None or metadata.get("encoding") != "cdr" - or metadata.get("type") != message_type.msg_name + or metadata.get("type") != message_type.__msgtype__ ): return False self._decoders[frame.header.ch](frame.payload) @@ -181,13 +185,13 @@ def _disengage(self) -> None: def _on_sensors_bytes(self, data: bytes) -> None: """Decode generated CDR bytes into TwistStamped and update sensor state.""" - msg = TwistStamped.decode(data) + msg = cdr_decode(data, TwistStamped) with self._lock: self._current_sensors = msg def _on_button_bytes(self, data: bytes) -> None: """Decode generated CDR bytes into Bool and update button state.""" - msg = Bool.decode(data) + msg = cdr_decode(data, Bool) with self._lock: self._teleop_button = bool(msg.data) @@ -279,7 +283,8 @@ def _get_output_twist(self) -> TwistStamped | None: x=current.twist.linear.x - initial.twist.linear.x, y=current.twist.linear.y - initial.twist.linear.y, z=current.twist.linear.z - initial.twist.linear.z, - ) + ), + angular=Vector3(x=0.0, y=0.0, z=0.0), ) # Handle yaw wraparound (linear.z = yaw, 0-360 degrees) diff --git a/dimos/teleop/phone/test_phone_cdr.py b/dimos/teleop/phone/test_phone_cdr.py index e426374ffc1..7595def569c 100644 --- a/dimos/teleop/phone/test_phone_cdr.py +++ b/dimos/teleop/phone/test_phone_cdr.py @@ -16,6 +16,11 @@ from collections.abc import Iterator +from dimos_cdrgen.registry import ( + decode as cdr_decode, + encode as cdr_encode, + schema as cdr_schema, +) from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Twist, TwistStamped, Vector3 from dimos_generated.std_msgs.msg import Bool, Header @@ -35,9 +40,9 @@ def frame(channel: str, message: DimosMsg, *, type_name: str | None = None) -> b seq=1, ts=0, delivery="latest", - meta={"type": type_name or message.msg_name, "encoding": "cdr"}, + meta={"type": type_name or message.__msgtype__, "encoding": "cdr"}, ), - message.encode(), + cdr_encode(message), ) @@ -54,13 +59,27 @@ def test_websocket_and_schema_endpoint_use_generated_contract(module: PhoneTeleo with TestClient(module._web_server.app) as client: schemas = client.get("/teleop/schema").json() assert schemas["sensors"] == { - "type": TwistStamped.msg_name, - "definition": TwistStamped.schema, + "type": TwistStamped.__msgtype__, + "definition": cdr_schema(TwistStamped.__msgtype__), + } + assert schemas["button"] == { + "type": Bool.__msgtype__, + "definition": cdr_schema(Bool.__msgtype__), } - assert schemas["button"] == {"type": Bool.msg_name, "definition": Bool.schema} with client.websocket_connect("/ws") as ws: ws.send_bytes(frame("button", Bool(data=True))) - ws.send_bytes(frame("sensors", TwistStamped(header=Header(frame_id="phone")))) + ws.send_bytes( + frame( + "sensors", + TwistStamped( + header=Header(frame_id="phone", stamp=Time(sec=0, nanosec=0)), + twist=Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), + angular=Vector3(x=0.0, y=0.0, z=0.0), + ), + ), + ) + ) assert module._teleop_button is True assert module._current_sensors.header.frame_id == "phone" @@ -69,7 +88,7 @@ def test_websocket_and_schema_endpoint_use_generated_contract(module: PhoneTeleo "wire", [ b"", - Bool(data=True).encode(), + cdr_encode(Bool(data=True)), frame("unknown", Bool(data=True)), frame("button", Bool(data=True), type_name="std_msgs/msg/Int8"), frame("sensors", Bool(data=True)), @@ -87,16 +106,19 @@ def test_wrong_identity_or_legacy_unframed_bytes_are_rejected( def test_generated_control_math_preserves_exact_stamp_and_wraps_yaw( module: PhoneTeleopModule, ) -> None: - initial = TwistStamped(twist=Twist(linear=Vector3(z=350))) + initial = TwistStamped( + twist=Twist(linear=Vector3(z=350, x=0.0, y=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) current = TwistStamped( - header=Header(stamp=Time(sec=1700000000, nanosec=123456789)), - twist=Twist(linear=Vector3(x=15, y=-30, z=10), angular=Vector3(z=60)), + header=Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id=""), + twist=Twist(linear=Vector3(x=15, y=-30, z=10), angular=Vector3(z=60, x=0.0, y=0.0)), ) - module._on_sensors_bytes(initial.encode()) + module._on_sensors_bytes(cdr_encode(initial)) assert module._engage() - module._on_sensors_bytes(current.encode()) + module._on_sensors_bytes(cdr_encode(current)) output = module._get_output_twist() - decoded = TwistStamped.decode(output.encode()) + decoded = cdr_decode(cdr_encode(output), TwistStamped) assert decoded.header.stamp.nanosec == 123456789 assert decoded.header.frame_id == "phone" assert decoded.twist.linear.x == 1.0 @@ -110,8 +132,13 @@ def test_ground_robot_extension_publishes_generated_twist() -> None: outputs: list[Twist] = [] unsubscribe = module.cmd_vel.subscribe(outputs.append) try: - module._publish_msg(TwistStamped(twist=Twist(linear=Vector3(x=1, y=2, z=3)))) - result = Twist.decode(outputs[0].encode()) + module._publish_msg( + TwistStamped( + twist=Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=0.0, y=0.0, z=0.0)), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + ) + result = cdr_decode(cdr_encode(outputs[0]), Twist) assert result.linear.x == 1 assert result.linear.y == 2 assert result.linear.z == 0 diff --git a/dimos/teleop/utils/test_report.py b/dimos/teleop/utils/test_report.py index 44fd519099d..da4967c3747 100644 --- a/dimos/teleop/utils/test_report.py +++ b/dimos/teleop/utils/test_report.py @@ -18,7 +18,15 @@ from pathlib import Path from dimos_generated.dimos_msgs.msg import VideoStats -from dimos_generated.geometry_msgs.msg import PoseStamped, TwistStamped +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Twist, + TwistStamped, + Vector3, +) from dimos_generated.std_msgs.msg import Header, UInt32 import pytest @@ -28,7 +36,10 @@ def stamped(stamp_ns: int) -> TwistStamped: - return TwistStamped(header=Header(stamp=time_from_nanoseconds(stamp_ns))) + return TwistStamped( + header=Header(stamp=time_from_nanoseconds(stamp_ns), frame_id=""), + twist=Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + ) def test_epoch_nanoseconds_and_out_of_order_source_frames() -> None: @@ -55,12 +66,42 @@ def test_generated_sqlite_recording_produces_report(tmp_path: Path) -> None: commands = store.stream("cmd_vel_stamped", TwistStamped) poses = store.stream("left_controller_output", PoseStamped) for index in range(3): - header = Header(stamp=time_from_nanoseconds(base + index * 20_000_000)) - commands.append(TwistStamped(header=header), ts=100 + index) - poses.append(PoseStamped(header=header), ts=100 + index) + header = Header(stamp=time_from_nanoseconds(base + index * 20_000_000), frame_id="") + commands.append( + TwistStamped( + header=header, + twist=Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ), + ), + ts=100 + index, + ) + poses.append( + PoseStamped( + header=header, + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + ts=100 + index, + ) store.stream("teleop_buttons", UInt32).append(UInt32(data=3), ts=101) store.stream("video_stats", VideoStats).append( - VideoStats(header=header, width=640, height=480, frames_dropped=2**32 + 1), ts=102 + VideoStats( + header=header, + width=640, + height=480, + frames_dropped=2**32 + 1, + fps=0.0, + kbps=0.0, + loss_pct=0.0, + jitter_buffer_ms=0.0, + decode_ms=0.0, + freezes=0, + e2e_latency_ms=0.0, + ), + ts=102, ) report = json.loads(generate_report(path).read_text()) assert report["duration_s"] == 0.04 diff --git a/dimos/teleop/utils/test_video_stats.py b/dimos/teleop/utils/test_video_stats.py index 77db6b78810..b622a602282 100644 --- a/dimos/teleop/utils/test_video_stats.py +++ b/dimos/teleop/utils/test_video_stats.py @@ -14,6 +14,11 @@ """CDR video telemetry carries declared fields, not a positional Joy payload.""" +from dimos_cdrgen.registry import ( + decode as cdr_decode, + encode as cdr_encode, + schema as cdr_schema, +) from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.dimos_msgs.msg import VideoStats from dimos_generated.std_msgs.msg import Header @@ -33,10 +38,13 @@ def test_cdr_and_sqlite_preserve_exact_counters_and_source_header(tmp_path) -> N loss_pct=2.1, frames_dropped=2**32 + 1, freezes=2**32 + 2, + jitter_buffer_ms=0.0, + decode_ms=0.0, + e2e_latency_ms=0.0, ) - assert VideoStats.decode(expected.encode()) == expected - assert VideoStats.msg_name == "dimos_msgs/msg/VideoStats" - assert "uint64 frames_dropped" in VideoStats.schema + assert cdr_decode(cdr_encode(expected), VideoStats) == expected + assert VideoStats.__msgtype__ == "dimos_msgs/msg/VideoStats" + assert "uint64 frames_dropped" in cdr_schema(VideoStats.__msgtype__) path = tmp_path / "video.db" with SqliteStore(path=str(path)) as store: stream = store.stream("video_stats", VideoStats, codec="cdr") diff --git a/dimos/teleop/webxr/extensions.py b/dimos/teleop/webxr/extensions.py index d3fe150e82d..5492a95c040 100644 --- a/dimos/teleop/webxr/extensions.py +++ b/dimos/teleop/webxr/extensions.py @@ -319,7 +319,9 @@ class Go2TeleopModule(WebXRTeleopModule): cmd_vel: Out[Twist] def _publish_safe_command(self) -> None: - self.cmd_vel.publish(Twist()) + self.cmd_vel.publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) def _deadzone(self, v: float) -> float: return 0.0 if abs(v) < self.config.deadzone else v @@ -336,9 +338,9 @@ def _on_joy_bytes(self, data: bytes) -> bool: with self._lock: left = self._controllers.get(Hand.LEFT) right = self._controllers.get(Hand.RIGHT) - twist = Twist() - twist.linear = Vector3() - twist.angular = Vector3() + twist = Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + twist.linear = Vector3(x=0.0, y=0.0, z=0.0) + twist.angular = Vector3(x=0.0, y=0.0, z=0.0) if left is not None: twist.linear.x = -self._deadzone(left.thumbstick.y) * self.config.linear_speed twist.linear.y = -self._deadzone(left.thumbstick.x) * self.config.linear_speed diff --git a/dimos/teleop/webxr/module.py b/dimos/teleop/webxr/module.py index 141eb7ec18f..f3e176a42d1 100644 --- a/dimos/teleop/webxr/module.py +++ b/dimos/teleop/webxr/module.py @@ -31,28 +31,35 @@ import time from typing import Any, TypeVar -from dimos_lcm.geometry_msgs import PoseStamped as LCMPoseStamped -from dimos_lcm.sensor_msgs import Joy as LCMJoy +from dimos_cdrgen.registry import decode as cdr_decode, schema as cdr_schema +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos_generated.sensor_msgs.msg import Joy +from dimos_generated.std_msgs.msg import String, UInt32 from fastapi import WebSocket, WebSocketDisconnect from fastapi.responses import HTMLResponse from fastapi.staticfiles import StaticFiles from pydantic import Field, ValidationError from reactivex.disposable import Disposable +from scipy.spatial.transform import Rotation from dimos.constants import DIMOS_PROJECT_ROOT from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out from dimos.imitation.collection.episode import EpisodeStatus -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.Joy import Joy -from dimos.msgs.std_msgs.String import String +from dimos.msgs.geometry import quaternion_from_matrix +from dimos.msgs.protocol import DimosMsg from dimos.teleop.utils.teleop_transforms import webxr_to_robot from dimos.teleop.webxr.body_tracking import BodyTrackingMode, BodyTrackingSnapshot # Hand is re-exported for callers; it lives in controller_types. from dimos.teleop.webxr.controller_types import Buttons, Hand, WebXRControllerState from dimos.utils.logging_config import setup_logger +from dimos.web.relay_bridge.protocol import ( + ProtocolError, + decode_data_frame, + peek_data_frame_lengths, +) from dimos.web.robot_web_interface import RobotWebInterface logger = setup_logger() @@ -103,7 +110,7 @@ class WebXRTeleopModule(Module): Outputs: - left_controller_output: PoseStamped (output pose for left hand) - right_controller_output: PoseStamped (output pose for right hand) - - teleop_buttons: Buttons (button states for both controllers) + - teleop_buttons: UInt32 (packed button states for both controllers) - body_tracking: named body-joint poses in their WebXR reference space """ @@ -112,9 +119,9 @@ class WebXRTeleopModule(Module): # Outputs: delta poses for each controller left_controller_output: Out[PoseStamped] right_controller_output: Out[PoseStamped] - teleop_buttons: Out[Buttons] - button_pressed: Out[Buttons] - button_released: Out[Buttons] + teleop_buttons: Out[UInt32] + button_pressed: Out[UInt32] + button_released: Out[UInt32] status: In[String] body_tracking: Out[BodyTrackingSnapshot] @@ -145,11 +152,8 @@ def __init__(self, **kwargs: Any) -> None: self._web_server: RobotWebInterface | None = None self._web_server_thread: threading.Thread | None = None - # Fingerprint-based message dispatch table - self._decoders: dict[bytes, Any] = { - LCMPoseStamped._get_packed_fingerprint(): self._on_pose_bytes, - LCMJoy._get_packed_fingerprint(): self._on_joy_bytes, - } + self._command_types: dict[str, type[DimosMsg]] = {"pose": PoseStamped, "joy": Joy} + self._decoders = {"pose": self._on_pose_bytes, "joy": self._on_joy_bytes} # Tracked here so subclasses can push from non-asyncio threads. # _clients_lock guards add/discard/snapshot of the set across the @@ -177,6 +181,16 @@ async def teleop_index() -> HTMLResponse: async def teleop_config() -> dict[str, Any]: return self._webxr_client_config() + @self._web_server.app.get("/teleop/schema") + async def command_schemas() -> dict[str, dict[str, str]]: + return { + channel: { + "type": message_type.__msgtype__, + "definition": cdr_schema(message_type.__msgtype__), + } + for channel, message_type in self._command_types.items() + } + if STATIC_DIR.is_dir(): self._web_server.app.mount( "/static", StaticFiles(directory=str(STATIC_DIR)), name="teleop_static" @@ -233,13 +247,23 @@ def _webxr_client_config(self) -> dict[str, Any]: } def _dispatch_binary_message(self, data: bytes) -> bool: - fingerprint = data[:8] - decoder = self._decoders.get(fingerprint) - if decoder is None: - logger.warning("Unknown WebXR message fingerprint", fingerprint=fingerprint.hex()) + try: + lengths = peek_data_frame_lengths(data) + if lengths is None or lengths[2] != len(data): + return False + frame = decode_data_frame(data) + message_type = self._command_types.get(frame.header.ch) + metadata = frame.header.meta or {} + if ( + message_type is None + or metadata.get("encoding") != "cdr" + or metadata.get("type") != message_type.__msgtype__ + ): + return False + return self._decoders[frame.header.ch](frame.payload) is not False + except (ProtocolError, ValueError): + logger.warning("Dropping malformed WebXR command frame") return False - decoder(data) - return True def _dispatch_text_message(self, payload: str) -> bool: try: @@ -444,23 +468,23 @@ def _resolve_hand(frame_id: str) -> Hand: raise ValueError(f"Unexpected frame_id: {frame_id!r}, expected 'left' or 'right'") def _on_pose_bytes(self, data: bytes) -> None: - """Decode LCM bytes into PoseStamped, transform to robot frame.""" - msg = PoseStamped.lcm_decode(data) - hand = self._resolve_hand(msg.frame_id) + """Decode CDR bytes into PoseStamped, transform to robot frame.""" + msg = cdr_decode(data, PoseStamped) + hand = self._resolve_hand(msg.header.frame_id) robot_pose = webxr_to_robot(msg, is_left_controller=(hand == Hand.LEFT)) with self._lock: self._current_poses[hand] = robot_pose self._last_pose_update[hand] = time.monotonic() def _on_joy_bytes(self, data: bytes) -> bool: - """Decode LCM bytes into Joy, parse into WebXRControllerState.""" - msg = Joy.lcm_decode(data) - hand = self._resolve_hand(msg.frame_id) + """Decode CDR bytes into Joy, parse into WebXRControllerState.""" + msg = cdr_decode(data, Joy) + hand = self._resolve_hand(msg.header.frame_id) try: controller = WebXRControllerState.from_joy(msg, is_left=(hand == Hand.LEFT)) except ValueError: logger.warning( - f"Malformed Joy for {hand.name}: axes={len(msg.axes or [])}, buttons={len(msg.buttons or [])}" + f"Malformed Joy for {hand.name}: axes={len(msg.axes)}, buttons={len(msg.buttons)}" ) with self._lock: self._controllers[hand] = None @@ -595,12 +619,36 @@ def _get_output_pose(self, hand: Hand) -> PoseStamped | None: if current_pose is None or initial_pose is None: return None - delta = current_pose - initial_pose + current = current_pose.pose + initial = initial_pose.pose + rotation = ( + Rotation.from_quat( + [ + current.orientation.x, + current.orientation.y, + current.orientation.z, + current.orientation.w, + ] + ) + * Rotation.from_quat( + [ + initial.orientation.x, + initial.orientation.y, + initial.orientation.z, + initial.orientation.w, + ] + ).inv() + ) return PoseStamped( - position=delta.position * self._translation_scale, - orientation=delta.orientation, - ts=current_pose.ts, - frame_id=current_pose.frame_id, + header=current_pose.header, + pose=Pose( + position=Point( + x=(current.position.x - initial.position.x) * self._translation_scale, + y=(current.position.y - initial.position.y) * self._translation_scale, + z=(current.position.z - initial.position.z) * self._translation_scale, + ), + orientation=quaternion_from_matrix(rotation.as_matrix()), + ), ) def _set_translation_scale(self, translation_scale: float) -> None: @@ -635,7 +683,7 @@ def _publish_button_state( def _publish_buttons(self, buttons: Buttons) -> None: """Publish raw state and stable digital button edges.""" - self.teleop_buttons.publish(buttons) + self.teleop_buttons.publish(buttons.to_message()) now = time.monotonic() observed = buttons.data & Buttons.DIGITAL_MASK pressed = 0 @@ -661,9 +709,9 @@ def _publish_buttons(self, buttons: Buttons) -> None: self._debounced_buttons &= ~mask released |= mask if pressed: - self.button_pressed.publish(Buttons(pressed)) + self.button_pressed.publish(UInt32(data=pressed)) if released: - self.button_released.publish(Buttons(released)) + self.button_released.publish(UInt32(data=released)) def _release_all_buttons(self) -> None: """Emit release edges without debounce when input ownership disappears.""" @@ -671,4 +719,4 @@ def _release_all_buttons(self) -> None: self._debounced_buttons = 0 self._button_candidates.clear() if released: - self.button_released.publish(Buttons(released)) + self.button_released.publish(UInt32(data=released)) diff --git a/dimos/teleop/webxr/test_module.py b/dimos/teleop/webxr/test_module.py index 8e49963736b..bdd200fe9f8 100644 --- a/dimos/teleop/webxr/test_module.py +++ b/dimos/teleop/webxr/test_module.py @@ -16,18 +16,20 @@ from collections.abc import Awaitable, Callable, Iterator import json import logging -from types import SimpleNamespace from typing import Any, cast +from dimos_cdrgen.registry import encode as cdr_encode, schema as cdr_schema +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.sensor_msgs.msg import Joy +from dimos_generated.std_msgs.msg import Header, String, UInt32 from fastapi import FastAPI from fastapi.testclient import TestClient +import numpy as np import pytest import pytest_mock from dimos.imitation.collection.episode import EpisodeStatus -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.Joy import Joy -from dimos.msgs.std_msgs.String import String from dimos.msgs.time import time_from_seconds from dimos.teleop.webxr.body_tracking import BodyTrackingSnapshot from dimos.teleop.webxr.controller_types import ( @@ -38,6 +40,30 @@ ) from dimos.teleop.webxr.extensions import ArmTeleopModule, Go2TeleopModule, HandTeleopModule from dimos.teleop.webxr.module import WebXRTeleopModule, _ws_send_text +from dimos.web.relay_bridge.protocol import FrameHeader, encode_data_frame + + +def _pose(*, ts=0.0, frame_id="", position=(0.0, 0.0, 0.0)): + return PoseStamped( + header=Header(stamp=time_from_seconds(ts), frame_id=frame_id), + pose=Pose( + position=Point(x=position[0], y=position[1], z=position[2]), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) + + +def _command_frame(channel, message): + return encode_data_frame( + FrameHeader( + ch=channel, + seq=0, + ts=0.0, + delivery="latest", + meta={"encoding": "cdr", "type": message.__msgtype__}, + ), + cdr_encode(message), + ) @pytest.mark.parametrize("state", ["unavailable", "empty", "tracking"]) @@ -186,8 +212,12 @@ def test_unknown_joy_controller_identity_is_rejected( module: WebXRTeleopModule, mocker: pytest_mock.MockerFixture ) -> None: mocker.patch( - "dimos.teleop.webxr.module.Joy.lcm_decode", - return_value=SimpleNamespace(frame_id="unknown"), + "dimos.teleop.webxr.module.cdr_decode", + return_value=Joy( + header=Header(frame_id="unknown", stamp=Time(sec=0, nanosec=0)), + axes=np.array([], dtype=np.float32), + buttons=np.array([], dtype=np.int32), + ), ) with pytest.raises(ValueError, match="Unexpected frame_id"): @@ -226,7 +256,7 @@ def test_episode_status_is_cached_and_broadcast( broadcast = mocker.patch.object(module, "_broadcast_text") mocker.patch("dimos.teleop.webxr.module.time.time", return_value=165.5) - module._on_episode_status(String(_episode_status().to_json())) + module._on_episode_status(String(data=_episode_status().to_json())) assert module._latest_episode_status == _episode_status() payload = json.loads(broadcast.call_args.args[0]) @@ -271,7 +301,7 @@ def test_control_client_disconnect_clears_state( module: WebXRTeleopModule, mocker: pytest_mock.MockerFixture ) -> None: first = mocker.MagicMock() - published: list[Buttons] = [] + published: list[UInt32] = [] module.teleop_buttons.subscribe(published.append) pose = mocker.MagicMock(spec=PoseStamped) assert module._client_connected(first) is True @@ -386,7 +416,7 @@ def test_stale_controller_input_disengages_hand( def test_stop_publishes_safe_button_state( module: WebXRTeleopModule, mocker: pytest_mock.MockerFixture ) -> None: - published: list[Buttons] = [] + published: list[UInt32] = [] module.teleop_buttons.subscribe(published.append) module._controllers[Hand.RIGHT] = WebXRControllerState(primary=True) module._is_engaged[Hand.RIGHT] = True @@ -419,9 +449,9 @@ def test_button_edges_are_debounced_and_not_repeated_while_held( module._publish_buttons(Buttons()) pressed.assert_called_once() - assert pressed.call_args.args[0].right_primary + assert Buttons(data=pressed.call_args.args[0].data).right_primary released.assert_called_once() - assert released.call_args.args[0].right_primary + assert Buttons(data=released.call_args.args[0].data).right_primary def test_button_edges_can_contain_simultaneous_digital_buttons( @@ -440,7 +470,7 @@ def test_button_edges_can_contain_simultaneous_digital_buttons( module._publish_buttons(buttons) publish.assert_called_once() - edge = publish.call_args.args[0] + edge = Buttons(data=publish.call_args.args[0].data) assert edge.left_secondary assert edge.right_primary @@ -469,7 +499,7 @@ def test_disconnect_releases_debounced_buttons_immediately( module._release_all_buttons() publish.assert_called_once() - assert publish.call_args.args[0].right_primary + assert Buttons(data=publish.call_args.args[0].data).right_primary assert module._debounced_buttons == 0 @@ -562,13 +592,12 @@ def test_go2_accepts_pico_six_button_joystick( module = Go2TeleopModule() publish = mocker.patch.object(module.cmd_vel, "publish") joy = Joy( - ts=1.0, - frame_id="left", - axes=[0.25, -0.75, 0.0, 0.0], - buttons=[0, 0, 0, 0, 0, 0], + header=Header(stamp=time_from_seconds(1.0), frame_id="left"), + axes=np.array([0.25, -0.75, 0.0, 0.0], dtype=np.float32), + buttons=np.array([0, 0, 0, 0, 0, 0], dtype=np.int32), ) try: - assert module._on_joy_bytes(joy.lcm_encode()) is True + assert module._on_joy_bytes(cdr_encode(joy)) is True twist = publish.call_args.args[0] assert twist.linear.x == pytest.approx(0.75 * module.config.linear_speed) @@ -584,14 +613,13 @@ def test_go2_rejects_short_controller_packet_safely( module = Go2TeleopModule() publish = mocker.patch.object(module.cmd_vel, "publish") joy = Joy( - ts=1.0, - frame_id="left", - axes=[0.25, -0.75, 0.0, 0.0], - buttons=[0, 0, 0, 0, 0], + header=Header(stamp=time_from_seconds(1.0), frame_id="left"), + axes=np.array([0.25, -0.75, 0.0, 0.0], dtype=np.float32), + buttons=np.array([0, 0, 0, 0, 0], dtype=np.int32), ) module._controllers[Hand.LEFT] = WebXRControllerState(thumbstick=ThumbstickState(y=-1.0)) try: - assert module._on_joy_bytes(joy.lcm_encode()) is False + assert module._on_joy_bytes(cdr_encode(joy)) is False assert module._controllers[Hand.LEFT] is None publish.assert_called_once() @@ -609,8 +637,12 @@ def test_go2_malformed_joy_clears_stale_state_and_publishes_zero_velocity( module = Go2TeleopModule() publish = mocker.patch.object(module.cmd_vel, "publish") mocker.patch( - "dimos.teleop.webxr.module.Joy.lcm_decode", - return_value=SimpleNamespace(frame_id="left", axes=[], buttons=[]), + "dimos.teleop.webxr.module.cdr_decode", + return_value=Joy( + header=Header(frame_id="left", stamp=Time(sec=0, nanosec=0)), + axes=np.array([], dtype=np.float32), + buttons=np.array([], dtype=np.int32), + ), ) module._controllers[Hand.LEFT] = WebXRControllerState(thumbstick=ThumbstickState(y=-1.0)) try: @@ -649,8 +681,12 @@ def test_go2_unknown_controller_identity_publishes_zero_velocity( module = Go2TeleopModule() publish = mocker.patch.object(module.cmd_vel, "publish") mocker.patch( - "dimos.teleop.webxr.module.Joy.lcm_decode", - return_value=SimpleNamespace(frame_id="unknown"), + "dimos.teleop.webxr.module.cdr_decode", + return_value=Joy( + header=Header(frame_id="unknown", stamp=Time(sec=0, nanosec=0)), + axes=np.array([], dtype=np.float32), + buttons=np.array([], dtype=np.int32), + ), ) module._controllers[Hand.LEFT] = WebXRControllerState(thumbstick=ThumbstickState(y=-1.0)) try: @@ -712,9 +748,9 @@ def test_binary_pose_dispatch_remains_on_existing_decoder( mocker: pytest_mock.MockerFixture, ) -> None: body_publish = mocker.patch.object(module.body_tracking, "publish") - pose = PoseStamped(ts=1.0, frame_id="left", position=[1.0, 2.0, 3.0]) + pose = _pose(ts=1.0, frame_id="left", position=[1.0, 2.0, 3.0]) - accepted = module._dispatch_binary_message(pose.lcm_encode()) + accepted = module._dispatch_binary_message(_command_frame("pose", pose)) assert accepted assert module._current_poses[Hand.LEFT] is not None @@ -730,23 +766,20 @@ def test_unknown_binary_message_is_dropped( accepted = module._dispatch_binary_message(b"unknown-message") assert not accepted - warning.assert_called_once_with( - "Unknown WebXR message fingerprint", - fingerprint=b"unknown-".hex(), - ) + warning.assert_called_once_with("Dropping malformed WebXR command frame") def test_translation_scale_changes_pose_delta(module: WebXRTeleopModule) -> None: - module._initial_poses[Hand.RIGHT] = PoseStamped(position=[1.0, 2.0, 3.0]) - module._current_poses[Hand.RIGHT] = PoseStamped(position=[1.2, 1.5, 4.0]) + module._initial_poses[Hand.RIGHT] = _pose(position=[1.0, 2.0, 3.0]) + module._current_poses[Hand.RIGHT] = _pose(position=[1.2, 1.5, 4.0]) module._set_translation_scale(2.0) output = module._get_output_pose(Hand.RIGHT) assert output is not None - assert output.position.x == pytest.approx(0.4) - assert output.position.y == pytest.approx(-1.0) - assert output.position.z == pytest.approx(2.0) + assert output.pose.position.x == pytest.approx(0.4) + assert output.pose.position.y == pytest.approx(-1.0) + assert output.pose.position.z == pytest.approx(2.0) @pytest.mark.parametrize("translation_scale", [0.0, -1.0, float("inf")]) @@ -762,9 +795,9 @@ def test_translation_scale_must_be_positive_and_finite( def test_arm_teleop_publishes_absolute_controller_pose() -> None: module = ArmTeleopModule() try: - pose = PoseStamped(frame_id="left", position=[1.0, 2.0, 3.0]) + pose = _pose(frame_id="left", position=[1.0, 2.0, 3.0]) module._current_poses[Hand.LEFT] = pose - module._initial_poses[Hand.LEFT] = PoseStamped(position=[0.5, 0.5, 0.5]) + module._initial_poses[Hand.LEFT] = _pose(position=[0.5, 0.5, 0.5]) assert module._get_output_pose(Hand.LEFT) is pose finally: @@ -775,7 +808,7 @@ def test_arm_teleop_publishes_pose_without_face_button_engagement() -> None: module = ArmTeleopModule() try: module._controllers[Hand.RIGHT] = WebXRControllerState(is_left=False) - module._current_poses[Hand.RIGHT] = PoseStamped(frame_id="right") + module._current_poses[Hand.RIGHT] = _pose(frame_id="right") module._handle_engage() @@ -836,8 +869,10 @@ def test_hand_teleop_pinch_toggles_engagement(mocker: pytest_mock.MockerFixture) assert module._is_engaged[Hand.RIGHT] module._publish_button_state(None, module._controllers[Hand.RIGHT]) - assert publish.call_args.args[0].right_grip - assert publish.call_args.args[0].right_trigger_analog == pytest.approx(1.0) + assert Buttons(data=publish.call_args.args[0].data).right_grip + assert Buttons(data=publish.call_args.args[0].data).right_trigger_analog == pytest.approx( + 1.0 + ) module._handle_engage() @@ -846,17 +881,54 @@ def test_hand_teleop_pinch_toggles_engagement(mocker: pytest_mock.MockerFixture) module._controllers[Hand.RIGHT] = WebXRControllerState(is_left=False, primary=False) module._handle_engage() module._publish_button_state(None, module._controllers[Hand.RIGHT]) - assert publish.call_args.args[0].right_grip + assert Buttons(data=publish.call_args.args[0].data).right_grip module._controllers[Hand.RIGHT] = WebXRControllerState(is_left=False, primary=True) module._handle_engage() assert not module._is_engaged[Hand.RIGHT] module._publish_button_state(None, module._controllers[Hand.RIGHT]) - assert not publish.call_args.args[0].right_grip + assert not Buttons(data=publish.call_args.args[0].data).right_grip finally: module.stop() +def test_command_schema_route_matches_generated_wire_types(module, mocker): + app = _setup_test_app(module, mocker) + with TestClient(app) as client: + response = client.get("/teleop/schema") + assert response.status_code == 200 + assert response.json() == { + "pose": { + "type": PoseStamped.__msgtype__, + "definition": cdr_schema(PoseStamped.__msgtype__), + }, + "joy": {"type": Joy.__msgtype__, "definition": cdr_schema(Joy.__msgtype__)}, + } + + +@pytest.mark.parametrize( + "invalid", ["wrong_type", "wrong_encoding", "unknown_channel", "trailing", "truncated"] +) +def test_command_dispatch_rejects_invalid_frame_before_control_handler(module, mocker, invalid): + handler = mocker.Mock(return_value=True) + module._decoders["pose"] = handler + channel = "unknown" if invalid == "unknown_channel" else "pose" + metadata = { + "encoding": "pickle" if invalid == "wrong_encoding" else "cdr", + "type": Joy.__msgtype__ if invalid == "wrong_type" else PoseStamped.__msgtype__, + } + frame = encode_data_frame( + FrameHeader(ch=channel, seq=0, ts=0.0, delivery="latest", meta=metadata), + cdr_encode(_pose()), + ) + if invalid == "trailing": + frame += b"extra" + elif invalid == "truncated": + frame = frame[:-1] + assert module._dispatch_binary_message(frame) is False + handler.assert_not_called() + + @pytest.mark.parametrize( "document", [ @@ -868,8 +940,8 @@ def test_hand_teleop_pinch_toggles_engagement(mocker: pytest_mock.MockerFixture) ) def test_invalid_episode_document_keeps_the_previous_hud_state(module, mocker, document): broadcast = mocker.patch.object(module, "_broadcast_text") - module._on_episode_status(String(_episode_status().to_json())) + module._on_episode_status(String(data=_episode_status().to_json())) broadcast.reset_mock() - module._on_episode_status(String(document)) + module._on_episode_status(String(data=document)) assert module._latest_episode_status == _episode_status() broadcast.assert_not_called() diff --git a/dimos/utils/test_transform_utils.py b/dimos/utils/test_transform_utils.py index 94597a1f10b..55b5d0bb2f4 100644 --- a/dimos/utils/test_transform_utils.py +++ b/dimos/utils/test_transform_utils.py @@ -12,11 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, Pose, PoseStamped, Quaternion, + Transform, TransformStamped, Vector3, ) @@ -208,9 +210,16 @@ def test_rotation_transform(self) -> None: def test_transform_with_transform_object(self) -> None: pose = Pose(position=Point(x=1, y=0, z=0), orientation=Quaternion(x=0, y=0, z=0, w=1)) - pose = PoseStamped(header=Header(frame_id="base"), pose=pose) - - transform = TransformStamped() + pose = PoseStamped(header=Header(frame_id="base", stamp=Time(sec=0, nanosec=0)), pose=pose) + + transform = TransformStamped( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + child_frame_id="", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) transform.header.frame_id = "world" transform.child_frame_id = "base" transform.transform.translation = Vector3(x=2, y=3, z=4) @@ -225,9 +234,16 @@ def test_transform_with_transform_object(self) -> None: def test_transform_frame_mismatch_raises(self) -> None: pose = Pose(position=Point(x=1, y=0, z=0), orientation=Quaternion(x=0, y=0, z=0, w=1)) - pose = PoseStamped(header=Header(frame_id="base"), pose=pose) - - transform = TransformStamped() + pose = PoseStamped(header=Header(frame_id="base", stamp=Time(sec=0, nanosec=0)), pose=pose) + + transform = TransformStamped( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + child_frame_id="", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) transform.header.frame_id = "world" transform.child_frame_id = "different_frame" transform.transform.translation = Vector3(x=2, y=3, z=4) diff --git a/dimos/utils/test_twist.py b/dimos/utils/test_twist.py index 8288d9cd6e4..4b198206df7 100644 --- a/dimos/utils/test_twist.py +++ b/dimos/utils/test_twist.py @@ -12,7 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Twist, TwistStamped, Vector3 +from dimos_generated.std_msgs.msg import Header import numpy as np from dimos.utils.transform_utils import twist_to_numpy @@ -20,7 +23,7 @@ def test_twist_to_numpy_orders_linear_then_angular(): twist = Twist(linear=Vector3(x=1, y=2, z=3), angular=Vector3(x=4, y=5, z=6)) - result = twist_to_numpy(Twist.decode(twist.encode())) + result = twist_to_numpy(cdr_decode(cdr_encode(twist), Twist)) assert result.dtype == np.float64 np.testing.assert_array_equal(result, [1, 2, 3, 4, 5, 6]) result[0] = 99 @@ -29,7 +32,8 @@ def test_twist_to_numpy_orders_linear_then_angular(): def test_stamped_twist_uses_explicit_nested_twist(): stamped = TwistStamped( - twist=Twist(linear=Vector3(x=-1, z=1), angular=Vector3(x=0.1, y=0.2, z=0.3)) + twist=Twist(linear=Vector3(x=-1, z=1, y=0.0), angular=Vector3(x=0.1, y=0.2, z=0.3)), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) result = twist_to_numpy(stamped.twist) np.testing.assert_array_equal(result, [-1, 0, 1, 0.1, 0.2, 0.3]) diff --git a/dimos/utils/testing/test_moment.py b/dimos/utils/testing/test_moment.py index d2f03a0f5aa..cac43eb9498 100644 --- a/dimos/utils/testing/test_moment.py +++ b/dimos/utils/testing/test_moment.py @@ -14,6 +14,7 @@ from pathlib import Path import time +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import PoseStamped, TransformStamped from dimos_generated.sensor_msgs.msg import CameraInfo, Image, PointCloud2 from dimos_generated.tf2_msgs.msg import TFMessage @@ -47,7 +48,7 @@ def transforms(self) -> list[TransformStamped]: # we just make sure to change timestamps so that we can jump # back and forth through time and the viewer doesn't get confused - odom = PoseStamped.decode(self.odom.value.encode()) + odom = cdr_decode(cdr_encode(self.odom.value), PoseStamped) odom.header.stamp = time_from_seconds(time.time()) return connection.GO2Connection._odom_to_tf(odom) @@ -58,7 +59,9 @@ def publish(self) -> None: t.dispose() tf_transport.stop() - camera_info = CameraInfo.decode(connection.GO2Connection.camera_info_static.encode()) + camera_info = cdr_decode( + cdr_encode(connection.GO2Connection.camera_info_static), CameraInfo + ) camera_info.header.stamp = time_from_seconds(time.time()) camera_info_transport: LCMTransport[CameraInfo] = LCMTransport("/camera_info", CameraInfo) camera_info_transport.publish(camera_info) diff --git a/dimos/utils/testing/test_sqlite_moment.py b/dimos/utils/testing/test_sqlite_moment.py index db5ea3463d8..59445bc5504 100644 --- a/dimos/utils/testing/test_sqlite_moment.py +++ b/dimos/utils/testing/test_sqlite_moment.py @@ -15,6 +15,7 @@ from itertools import islice from pathlib import Path +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.geometry_msgs.msg import Point import pytest @@ -59,7 +60,7 @@ def test_sensor_moment_publishes_selected_cdr_value_and_clears_missing_seek(reco moment.seek(1.0) moment.publish() (published,) = transport.publish.call_args.args - assert Point.decode(published.encode()).x == 1 + assert cdr_decode(cdr_encode(published), Point).x == 1 moment.seek(10.0) assert moment.value is None moment.publish() diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index 64ec840d5e6..b5efa957823 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -57,6 +57,7 @@ from dimos.core.core import rpc from dimos.core.global_config import global_config from dimos.core.module import Module, ModuleConfig +from dimos.msgs.time import message_header from dimos.protocol.pubsub.impl.lcmpubsub import LCM from dimos.protocol.pubsub.impl.zenohpubsub import Topic as ZenohTopic, Zenoh from dimos.protocol.pubsub.patterns import Glob, pattern_matches @@ -121,7 +122,7 @@ class RerunEntry(NamedTuple): RerunMulti: TypeAlias = "list[tuple[str, Archetype] | RerunEntry]" RerunData: TypeAlias = "Archetype | RerunMulti" -# A region cell packed into a header seq: i in the high half, j in the low half, +# A region cell packed into a signed region ID: i in the high half, j in the low half, # both signed 16 bit. pack_cell in mls_planner/rust/src/region_viz.rs owns this # layout, the ray tracer copies it. CELL_SHIFT = 16 @@ -129,10 +130,10 @@ class RerunEntry(NamedTuple): CELL_SIGN = 0x8000 -def region_entity(base: str, seq: int) -> str: - """The entity of the region cell a native packed into a header seq.""" - i = seq >> CELL_SHIFT - j = ((seq & CELL_MASK) ^ CELL_SIGN) - CELL_SIGN +def region_entity(base: str, region_id: int) -> str: + """The entity of the region cell a native packed into a region ID.""" + i = region_id >> CELL_SHIFT + j = ((region_id & CELL_MASK) ^ CELL_SIGN) - CELL_SIGN return f"{base}/{i}_{j}" @@ -149,9 +150,9 @@ def region_entity(base: str, seq: int) -> str: @runtime_checkable class KeyedRenderer(Protocol): - """A renderer whose entities are keyed by the message's header seq.""" + """A renderer whose entities are keyed by the message's region ID.""" - keyed_by_seq: bool + keyed_by_region: bool def __call__(self, *args: object, **kwargs: object) -> object: ... @@ -159,17 +160,17 @@ def __call__(self, *args: object, **kwargs: object) -> object: ... RendererT = TypeVar("RendererT", bound=Callable[..., object]) -def keyed_by_seq(renderer: RendererT) -> RendererT: +def keyed_by_region(renderer: RendererT) -> RendererT: """Mark a renderer whose topic carries one message per key, taken unconflated.""" - renderer.keyed_by_seq = True # type: ignore[attr-defined] + renderer.keyed_by_region = True # type: ignore[attr-defined] return renderer -def is_keyed_by_seq(renderer: object) -> bool: +def is_keyed_by_region(renderer: object) -> bool: """Whether a renderer, or the function under a partial, is marked keyed.""" while isinstance(renderer, functools.partial): renderer = renderer.func - return isinstance(renderer, KeyedRenderer) and renderer.keyed_by_seq + return isinstance(renderer, KeyedRenderer) and renderer.keyed_by_region def is_rerun_multi(data: Any) -> TypeGuard[RerunMulti]: @@ -316,17 +317,18 @@ class Config(ModuleConfig): class RerunBridgeModule(Module): """Bridge that logs messages from pubsubs to Rerun. - Spawns its own Rerun viewer and subscribes to all topics on each provided - pubsub. Any message that has a to_rerun() method is automatically logged. + Spawns its own Rerun viewer and subscribes to all topics on each provided + pubsub. Any message that has a to_rerun() method is automatically logged. - Example: - from dimos.protocol.pubsub.impl.lcmpubsub import LCM + Example: + from dimos.msgs.time import message_header + from dimos.protocol.pubsub.impl.lcmpubsub import LCM - lcm = LCM() - bridge = RerunBridgeModule(pubsubs=[lcm]) - bridge.start() - # All messages with to_rerun() are now logged to Rerun - bridge.stop() + lcm = LCM() + bridge = RerunBridgeModule(pubsubs=[lcm]) + bridge.start() + # All messages with to_rerun() are now logged to Rerun + bridge.stop() """ config: Config @@ -426,12 +428,12 @@ def composed(msg: Any) -> RerunData | None: return composed def _keyed_topics(self) -> list[str]: - """Channels whose renderer keys entities by seq, so every message must arrive.""" + """Channels whose renderer keys entities by region ID, so every message must arrive.""" prefix = f"{self.config.entity_prefix}/" return [ str(path).removeprefix(prefix) for path, renderer in self.config.visual_override.items() - if is_keyed_by_seq(renderer) + if is_keyed_by_region(renderer) ] def _get_entity_path(self, topic: Any) -> str: @@ -486,7 +488,7 @@ def _on_message(self, msg: Any, topic: Any) -> None: rr.log(entry[0], entry[1], static=isinstance(entry, RerunEntry) and entry.static) # Pin the base entity to the message frame once, so child entities # follow it the way a single archetype does below. - frame_id = getattr(msg, "frame_id", None) + frame_id = getattr(message_header(msg), "frame_id", None) if frame_id and self._frame_attached.get(entity_path) != frame_id: rr.log(entity_path, rr.Transform3D(parent_frame=f"tf#/{frame_id}")) self._frame_attached[entity_path] = frame_id @@ -500,7 +502,7 @@ def _on_message(self, msg: Any, topic: Any) -> None: frame_id = ( msg.frame_id if isinstance(msg, CompressedVideo) - else getattr(getattr(msg, "header", None), "frame_id", None) + else getattr(message_header(msg), "frame_id", None) ) if frame_id and self._frame_attached.get(entity_path) != frame_id: rr.log(entity_path, rr.Transform3D(parent_frame=f"tf#/{frame_id}")) diff --git a/dimos/visualization/rerun/test_bridge_topics.py b/dimos/visualization/rerun/test_bridge_topics.py index 3d7171cd7dd..c592d860e6d 100644 --- a/dimos/visualization/rerun/test_bridge_topics.py +++ b/dimos/visualization/rerun/test_bridge_topics.py @@ -26,19 +26,24 @@ from typing import Any from unittest.mock import patch +from dimos_generated.dimos_msgs.msg import LineSegments3D, RegionLineSegments3D, RegionPointCloud2 +from dimos_generated.sensor_msgs.msg import PointCloud2 +import numpy as np import pytest from pytest_mock import MockerFixture import rerun as rr from dimos.core.coordination.module_coordinator import ModuleCoordinator from dimos.core.global_config import GlobalConfig -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.core.transport_factory import rpc_backend +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import header_now from dimos.protocol.pubsub.impl.zenohpubsub import Topic as ZenohTopic, Zenoh from dimos.visualization.rerun.bridge import ( RerunBridgeModule, RerunEntry, is_rerun_multi, - keyed_by_seq, + keyed_by_region, region_entity, ) @@ -62,8 +67,10 @@ def subscribe_all(self, callback: Any, unconflated: Collection[str] = ()) -> Any @pytest.fixture -def make_bridge() -> Iterator[Callable[..., RerunBridgeModule]]: +def make_bridge(mocker: MockerFixture) -> Iterator[Callable[..., RerunBridgeModule]]: """Build real bridge modules. Constructing one opens a transport, so each is stopped again.""" + for method in ("start", "serve_module_rpc", "stop"): + mocker.patch.object(rpc_backend(), method) bridges: list[RerunBridgeModule] = [] def make(**config: Any) -> RerunBridgeModule: @@ -95,7 +102,7 @@ def test_no_topics_keeps_the_firehose(make_bridge: Callable[..., RerunBridgeModu def test_the_firehose_takes_keyed_renderers_unconflated( make_bridge: Callable[..., RerunBridgeModule], ) -> None: - @keyed_by_seq + @keyed_by_region def by_cell(msg: PointCloud2) -> None: return None @@ -129,8 +136,9 @@ def test_a_leaf_stack_does_not_claim_the_coordinator_name() -> None: assert coordinator._coordinator_rpc is None +@pytest.mark.parametrize("message_type", [PointCloud2, RegionPointCloud2, RegionLineSegments3D]) def test_multi_entries_log_their_static_flag_and_pin_the_frame_once( - mocker: MockerFixture, make_bridge: Callable[..., RerunBridgeModule] + mocker: MockerFixture, make_bridge: Callable[..., RerunBridgeModule], message_type ) -> None: cells = [ RerunEntry("world/surface_map/0_0", rr.Points3D([[0.0, 0.0, 0.0]]), static=True), @@ -138,10 +146,18 @@ def test_multi_entries_log_their_static_flag_and_pin_the_frame_once( ] assert is_rerun_multi(cells) log = mocker.patch("rerun.log") - topic = ZenohTopic(topic="dimos/surface_map", lcm_type=PointCloud2) + topic = ZenohTopic(topic="dimos/surface_map", msg_type=message_type) bridge = make_bridge(visual_override={"world/surface_map": lambda m: cells}) entity = bridge._get_entity_path(topic) - msg = PointCloud2(frame_id="odom") + cloud = pointcloud_from_xyz(np.empty((0, 3)), header=header_now("odom")) + if message_type is RegionPointCloud2: + msg = RegionPointCloud2(region_id=-196603, cloud=cloud) + elif message_type is RegionLineSegments3D: + msg = RegionLineSegments3D( + region_id=-196603, lines=LineSegments3D(header=cloud.header, segments=[]) + ) + else: + msg = cloud bridge._on_message(msg, topic) bridge._on_message(msg, topic) calls = [(c.args[0], c.kwargs.get("static", False)) for c in log.call_args_list] @@ -151,6 +167,6 @@ def test_multi_entries_log_their_static_flag_and_pin_the_frame_once( assert len(pins) == 1 and isinstance(pins[0].args[1], rr.Transform3D) -def test_region_cells_unpack_from_the_seq_the_planner_packs() -> None: +def test_region_cells_unpack_from_the_signed_id_the_planner_packs() -> None: assert region_entity("world/surface_map", (-3 << 16) | (5 & 0xFFFF)) == "world/surface_map/-3_5" assert region_entity("world/node_edges", (7 << 16) | (-2 & 0xFFFF)) == "world/node_edges/7_-2" diff --git a/dimos/visualization/rerun/test_detection3d_bridge.py b/dimos/visualization/rerun/test_detection3d_bridge.py index 0bf6cfd2006..8184e9d9a9b 100644 --- a/dimos/visualization/rerun/test_detection3d_bridge.py +++ b/dimos/visualization/rerun/test_detection3d_bridge.py @@ -15,8 +15,9 @@ from dataclasses import dataclass from unittest.mock import patch +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time -from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion, Vector3 +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseWithCovariance, Quaternion, Vector3 from dimos_generated.std_msgs.msg import Header from dimos_generated.vision_msgs.msg import ( BoundingBox3D, @@ -25,6 +26,7 @@ ObjectHypothesis, ObjectHypothesisWithPose, ) +import numpy as np import rerun as rr from dimos.visualization.rerun.bridge import RerunBridgeModule @@ -37,26 +39,44 @@ class Topic: def _detection_array() -> Detection3DArray: - det = Detection3D() - det.header = Header(stamp=Time(sec=10), frame_id="world") + det = Detection3D( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + results=[], + bbox=BoundingBox3D( + center=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + size=Vector3(x=0.0, y=0.0, z=0.0), + ), + id="", + ) + det.header = Header(stamp=Time(sec=10, nanosec=0), frame_id="world") det.id = "4" det.results = [ ObjectHypothesisWithPose( hypothesis=ObjectHypothesis( class_id="DICT_APRILTAG_36h11:4", score=1.0, - ) + ), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), ) ] det.bbox = BoundingBox3D( center=Pose( position=Point(x=1.0, y=2.0, z=3.0), - orientation=Quaternion(w=1.0), + orientation=Quaternion(w=1.0, x=0.0, y=0.0, z=0.0), ), - size=Vector3(x=0.1, y=0.1), + size=Vector3(x=0.1, y=0.1, z=0.0), ) return Detection3DArray( - header=Header(stamp=Time(sec=10), frame_id="world"), + header=Header(stamp=Time(sec=10, nanosec=0), frame_id="world"), detections=[det], ) @@ -68,7 +88,7 @@ def test_detection3darray_bridge_attaches_topic_entity_to_message_frame() -> Non try: with patch("rerun.log") as mock_log: bridge._on_message( - Detection3DArray.decode(_detection_array().encode()), + cdr_decode(cdr_encode(_detection_array()), Detection3DArray), Topic("/marker_detection/detections"), ) finally: @@ -85,9 +105,11 @@ def test_detection3darray_bridge_attaches_topic_entity_to_message_frame() -> Non def test_generated_detection_boxes_preserve_geometry_and_labels() -> None: - boxes = detection_boxes(Detection3DArray.decode(_detection_array().encode())) + boxes = detection_boxes(cdr_decode(cdr_encode(_detection_array()), Detection3DArray)) assert boxes.centers.as_arrow_array().to_pylist() == [[1, 2, 3]] assert boxes.labels.as_arrow_array().to_pylist() == ["DICT_APRILTAG_36h11:4 id=4"] assert boxes.quaternions.as_arrow_array().to_pylist() == [[0, 0, 0, 1]] - empty = detection_boxes(Detection3DArray()) + empty = detection_boxes( + Detection3DArray(header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), detections=[]) + ) assert empty.centers.as_arrow_array().to_pylist() == [] diff --git a/dimos/visualization/rerun/test_generated_camera_bridge.py b/dimos/visualization/rerun/test_generated_camera_bridge.py index 2f9693fac5d..f269ad82b0d 100644 --- a/dimos/visualization/rerun/test_generated_camera_bridge.py +++ b/dimos/visualization/rerun/test_generated_camera_bridge.py @@ -15,7 +15,9 @@ from types import SimpleNamespace from unittest.mock import patch -from dimos_generated.sensor_msgs.msg import CameraInfo, CompressedImage, Image +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.sensor_msgs.msg import CameraInfo, CompressedImage, Image, RegionOfInterest from dimos_generated.std_msgs.msg import Header import numpy as np import pytest @@ -35,16 +37,33 @@ def test_generated_calibration_pairs_with_image(camera_first: bool, compressed: bridge._min_intervals = {} bridge._camera_infos = {} bridge._image_entities = set() - header = Header(frame_id="optical") + header = Header(frame_id="optical", stamp=Time(sec=0, nanosec=0)) image = image_from_array(np.zeros((4, 6, 3), dtype=np.uint8), encoding="rgb8", header=header) message = ( - CompressedImage(header=header, format="jpeg", data=image_to_jpeg(image)) + CompressedImage( + header=header, format="jpeg", data=np.frombuffer(image_to_jpeg(image), dtype=np.uint8) + ) if compressed else image ) - message = type(message).decode(message.encode()) - info = CameraInfo(header=header, width=6, height=4, k=[10, 0, 3, 0, 11, 2, 0, 0, 1]) - pairs = [(CameraInfo.decode(info.encode()), "/camera_info"), (message, "/image")] + message = cdr_decode(cdr_encode(message), type(message)) + info = CameraInfo( + header=header, + width=6, + height=4, + k=np.array([10, 0, 3, 0, 11, 2, 0, 0, 1], dtype=np.float64), + distortion_model="", + d=np.array([], dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), + ) + pairs = [ + (cdr_decode(cdr_encode(info), CameraInfo), "/camera_info"), + (message, "/image"), + ] if not camera_first: pairs.reverse() try: @@ -65,14 +84,24 @@ def test_generated_calibration_pairs_with_image(camera_first: bool, compressed: ) def test_generated_depth_units(encoding, dtype, meter) -> None: message = image_from_array(np.ones((2, 3), dtype=dtype), encoding=encoding) - result = image_archetype(Image.decode(message.encode())) + result = image_archetype(cdr_decode(cdr_encode(message), Image)) assert isinstance(result, rr.DepthImage) assert result.meter.as_arrow_array().to_pylist() == [meter] def test_camera_override_parents_frustum_without_mutating_header(): info = CameraInfo( - header=Header(frame_id="original"), width=6, height=4, k=[10, 0, 3, 0, 11, 2, 0, 0, 1] + header=Header(frame_id="original", stamp=Time(sec=0, nanosec=0)), + width=6, + height=4, + k=np.array([10, 0, 3, 0, 11, 2, 0, 0, 1], dtype=np.float64), + distortion_model="", + d=np.array([], dtype=np.float64), + r=np.zeros(9, dtype=np.float64), + p=np.zeros(12, dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), ) result = camera_pinhole(info, optical_frame="d435_color_optical_frame") assert result.parent_frame.as_arrow_array().to_pylist() == ["tf#/d435_color_optical_frame"] diff --git a/dimos/visualization/rerun/test_generated_cloud_bridge.py b/dimos/visualization/rerun/test_generated_cloud_bridge.py index 04d17f57b7d..1d02584ac8e 100644 --- a/dimos/visualization/rerun/test_generated_cloud_bridge.py +++ b/dimos/visualization/rerun/test_generated_cloud_bridge.py @@ -16,6 +16,8 @@ from types import SimpleNamespace from unittest.mock import patch +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.sensor_msgs.msg import PointCloud2, PointField from dimos_generated.std_msgs.msg import Header import numpy as np @@ -41,17 +43,18 @@ def test_padded_big_endian_cloud_keeps_color_alignment_and_frame() -> None: + b"pad!" ) cloud = PointCloud2( - header=Header(frame_id="lidar"), + header=Header(frame_id="lidar", stamp=Time(sec=0, nanosec=0)), height=2, width=1, point_step=16, row_step=20, is_bigendian=True, fields=fields, - data=data, + data=np.frombuffer(data, dtype=np.uint8), + is_dense=False, ) - decoded = PointCloud2.decode(cloud.encode()) - before = decoded.encode() + decoded = cdr_decode(cdr_encode(cloud), PointCloud2) + before = cdr_encode(decoded) bridge = RerunBridgeModule() bridge._min_intervals = {} try: @@ -66,14 +69,19 @@ def test_padded_big_endian_cloud_keeps_color_alignment_and_frame() -> None: ] finally: bridge.stop() - assert decoded.encode() == before + assert cdr_encode(decoded) == before def test_height_colors_and_empty_cloud() -> None: - cloud = pointcloud_from_xyz(np.array([[0, 0, 0], [0, 0, 2]], dtype=np.float32), header=Header()) + cloud = pointcloud_from_xyz( + np.array([[0, 0, 0], [0, 0, 2]], dtype=np.float32), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) rendered = cloud_archetype(cloud) assert len(set(rendered.colors.as_arrow_array().to_pylist())) == 2 - empty = pointcloud_from_xyz(np.empty((0, 3)), header=Header()) + empty = pointcloud_from_xyz( + np.empty((0, 3)), header=Header(stamp=Time(sec=0, nanosec=0), frame_id="") + ) assert cloud_archetype(empty).positions.as_arrow_array().to_pylist() == [] @@ -82,23 +90,23 @@ def test_cloud_modes_filter_height_without_losing_rgb_alignment(mode): cloud = pointcloud_from_xyz_rgb( np.array([[0, 0, -1], [1, 2, 3]], dtype=np.float32), np.array([[255, 0, 0], [0x12, 0x34, 0x56]], dtype=np.uint8), - header=Header(frame_id="map"), + header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), ) - before = cloud.encode() + before = cdr_encode(cloud) rendered = cloud_archetype(cloud, mode=mode, bottom_cutoff=0, voxel_size=0.2) positions = rendered.centers if mode == "boxes" else rendered.positions assert positions.as_arrow_array().to_pylist() == [[1, 2, 3]] assert rendered.colors.as_arrow_array().to_pylist() == [0x123456FF] if mode == "points": assert rendered.radii.as_arrow_array().to_pylist() == [-2.0] - assert cloud.encode() == before + assert cdr_encode(cloud) == before @pytest.mark.parametrize("mode", ["points", "boxes"]) def test_explicit_cloud_colors_follow_finite_and_height_filters(mode): cloud = pointcloud_from_xyz( np.array([[0, 0, -1], [np.nan, 0, 1], [1, 2, 3]], dtype=np.float32), - header=Header(), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) rendered = cloud_archetype( cloud, @@ -113,6 +121,9 @@ def test_explicit_cloud_colors_follow_finite_and_height_filters(mode): @pytest.mark.parametrize("colors", [[256, 0, 0], [-1, 0, 0], [0.5, 0, 0], [[1, 2]]]) def test_explicit_cloud_colors_reject_invalid_channels(colors): - cloud = pointcloud_from_xyz(np.array([[0, 0, 1]], dtype=np.float32), header=Header()) + cloud = pointcloud_from_xyz( + np.array([[0, 0, 1]], dtype=np.float32), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) with pytest.raises(ValueError, match="color"): cloud_archetype(cloud, colors=colors) diff --git a/dimos/visualization/rerun/test_generated_entity_bridge.py b/dimos/visualization/rerun/test_generated_entity_bridge.py index 70819bec2fe..ab7a6265912 100644 --- a/dimos/visualization/rerun/test_generated_entity_bridge.py +++ b/dimos/visualization/rerun/test_generated_entity_bridge.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_cdrgen.registry import encode as cdr_encode, schema as cdr_schema from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.dimos_msgs.msg import EntityMarker, EntityMarkers from dimos_generated.geometry_msgs.msg import Point @@ -37,14 +38,16 @@ def test_generated_entity_schema_independently_decodes_labels_positions_and_stam ], ) independent = get_typestore(Stores.ROS2_HUMBLE) - independent.register(get_types_from_msg(EntityMarkers.schema, EntityMarkers.msg_name)) - result = independent.deserialize_cdr(message.encode(), EntityMarkers.msg_name) + independent.register( + get_types_from_msg(cdr_schema(EntityMarkers.__msgtype__), EntityMarkers.__msgtype__) + ) + result = independent.deserialize_cdr(cdr_encode(message), EntityMarkers.__msgtype__) assert result.header.frame_id == "world" assert (result.header.stamp.sec, result.header.stamp.nanosec) == (1700000000, 999999999) (marker,) = result.markers assert (marker.entity_id, marker.label, marker.entity_type) == ("E1", "机器人", "person") assert (marker.position.x, marker.position.y, marker.position.z) == (1, 2, 0.3) - assert resolve_msg_type(EntityMarkers.msg_name) is EntityMarkers + assert resolve_msg_type(EntityMarkers.__msgtype__) is EntityMarkers def test_generated_entity_viewer_helper_preserves_points_labels_and_radius() -> None: @@ -59,7 +62,8 @@ def test_generated_entity_viewer_helper_preserves_points_labels_and_radius() -> EntityMarker( entity_id="E2", label="table", entity_type="object", position=Point(x=3, y=4, z=0.3) ), - ] + ], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) archetype = entity_points(message) assert isinstance(archetype, rr.Points3D) diff --git a/dimos/visualization/rerun/test_generated_navigation_bridge.py b/dimos/visualization/rerun/test_generated_navigation_bridge.py index e5bd3bb016a..6e1dd2c02cf 100644 --- a/dimos/visualization/rerun/test_generated_navigation_bridge.py +++ b/dimos/visualization/rerun/test_generated_navigation_bridge.py @@ -15,6 +15,8 @@ from types import SimpleNamespace from unittest.mock import patch +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, PointStamped, @@ -22,9 +24,13 @@ PoseStamped, PoseWithCovariance, Quaternion, + Twist, + TwistWithCovariance, + Vector3, ) from dimos_generated.nav_msgs.msg import Odometry, Path from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest import rerun as rr @@ -36,16 +42,26 @@ @pytest.mark.parametrize("kind", ["pose", "odom", "point", "path"]) def test_navigation_bridge_preserves_coordinates_and_frame(kind: str) -> None: - header = Header(frame_id="map") - pose = Pose(position=Point(x=1, y=2, z=3), orientation=Quaternion(w=1)) + header = Header(frame_id="map", stamp=Time(sec=0, nanosec=0)) + pose = Pose(position=Point(x=1, y=2, z=3), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0)) messages = { "pose": PoseStamped(header=header, pose=pose), - "odom": Odometry(header=header, pose=PoseWithCovariance(pose=pose)), + "odom": Odometry( + header=header, + pose=PoseWithCovariance(pose=pose, covariance=np.zeros(36, dtype=np.float64)), + child_frame_id="", + twist=TwistWithCovariance( + twist=Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ), + covariance=np.zeros(36, dtype=np.float64), + ), + ), "point": PointStamped(header=header, point=pose.position), "path": Path(header=header, poses=[PoseStamped(header=header, pose=pose)]), } value = messages[kind] - decoded = type(value).decode(value.encode()) + decoded = cdr_decode(cdr_encode(value), type(value)) bridge = RerunBridgeModule() bridge._min_intervals = {} try: @@ -66,15 +82,33 @@ def test_navigation_bridge_preserves_coordinates_and_frame(kind: str) -> None: assert output.strips.as_arrow_array().to_pylist() == [[[1, 2, 3.5]]] finally: bridge.stop() - assert decoded.encode() == value.encode() + assert cdr_encode(decoded) == cdr_encode(value) def test_empty_generated_path_clears_geometry() -> None: - assert navigation_archetype(Path()).strips.as_arrow_array().to_pylist() == [] + assert ( + navigation_archetype( + Path(header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), poses=[]) + ) + .strips.as_arrow_array() + .to_pylist() + == [] + ) def test_path_override_preserves_configured_display_height(): - value = Path(poses=[PoseStamped(pose=Pose(position=Point(x=1, y=2, z=3)))]) + value = Path( + poses=[ + PoseStamped( + pose=Pose( + position=Point(x=1, y=2, z=3), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + ], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) result = navigation_archetype(value, z_offset=0.3) assert result.strips.as_arrow_array().to_pylist() == [[[1, 2, pytest.approx(3.3)]]] assert value.poses[0].pose.position.z == 3 diff --git a/dimos/visualization/rerun/test_generated_occupancy_bridge.py b/dimos/visualization/rerun/test_generated_occupancy_bridge.py index f90d3993682..c89c94f46fb 100644 --- a/dimos/visualization/rerun/test_generated_occupancy_bridge.py +++ b/dimos/visualization/rerun/test_generated_occupancy_bridge.py @@ -16,6 +16,8 @@ from types import SimpleNamespace from unittest.mock import patch +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid from dimos_generated.std_msgs.msg import Header @@ -30,19 +32,20 @@ def test_rotated_occupancy_cdr_mesh_and_texture() -> None: grid = OccupancyGrid( - header=Header(frame_id="map"), + header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), info=MapMetaData( width=2, height=2, resolution=0.5, origin=Pose( position=Point(x=10, y=20, z=3), - orientation=Quaternion(z=math.sqrt(0.5), w=math.sqrt(0.5)), + orientation=Quaternion(z=math.sqrt(0.5), w=math.sqrt(0.5), x=0.0, y=0.0), ), + map_load_time=Time(sec=0, nanosec=0), ), - data=[0, 100, -1, 50], + data=np.array([0, 100, -1, 50], dtype=np.int8), ) - decoded = OccupancyGrid.decode(grid.encode()) + decoded = cdr_decode(cdr_encode(grid), OccupancyGrid) bridge = RerunBridgeModule() bridge._min_intervals = {} try: @@ -60,10 +63,45 @@ def test_rotated_occupancy_cdr_mesh_and_texture() -> None: ] finally: bridge.stop() - assert decoded.encode() == grid.encode() + assert cdr_encode(decoded) == cdr_encode(grid) def test_empty_and_invalid_occupancy_data() -> None: - assert occupancy_mesh(OccupancyGrid()).vertex_positions.as_arrow_array().to_pylist() == [] + assert ( + occupancy_mesh( + OccupancyGrid( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + info=MapMetaData( + map_load_time=Time(sec=0, nanosec=0), + resolution=0.0, + width=0, + height=0, + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([], dtype=np.int8), + ) + ) + .vertex_positions.as_arrow_array() + .to_pylist() + == [] + ) with pytest.raises(ValueError, match="dimensions"): - occupancy_mesh(OccupancyGrid(info=MapMetaData(width=2, height=2), data=[0])) + occupancy_mesh( + OccupancyGrid( + info=MapMetaData( + width=2, + height=2, + map_load_time=Time(sec=0, nanosec=0), + resolution=0.0, + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([0], dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + ) diff --git a/dimos/visualization/rerun/test_generated_tf_bridge.py b/dimos/visualization/rerun/test_generated_tf_bridge.py index de90ea2ba5f..948043944e8 100644 --- a/dimos/visualization/rerun/test_generated_tf_bridge.py +++ b/dimos/visualization/rerun/test_generated_tf_bridge.py @@ -15,6 +15,8 @@ from types import SimpleNamespace from unittest.mock import patch +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 from dimos_generated.std_msgs.msg import Header from dimos_generated.tf2_msgs.msg import TFMessage @@ -30,11 +32,13 @@ def test_generated_tf_reaches_bridge_with_or_without_axes(axes: float) -> None: bridge = RerunBridgeModule(tf_axes=axes) bridge._min_intervals = {} edge = TransformStamped( - header=Header(frame_id="world"), + header=Header(frame_id="world", stamp=Time(sec=0, nanosec=0)), child_frame_id="base_link", - transform=Transform(translation=Vector3(x=3), rotation=Quaternion(w=1)), + transform=Transform( + translation=Vector3(x=3, y=0.0, z=0.0), rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), ) - message = TFMessage.decode(TFMessage(transforms=[edge]).encode()) + message = cdr_decode(cdr_encode(TFMessage(transforms=[edge])), TFMessage) try: with patch("rerun.log") as log: bridge._on_message(message, SimpleNamespace(name="/tf")) @@ -52,7 +56,14 @@ def test_generated_tf_reaches_bridge_with_or_without_axes(axes: float) -> None: def test_generated_tree_reparents_and_clears_old_axes() -> None: tree = TfFrameTree(root="frames") - edge = TransformStamped(header=Header(frame_id="first"), child_frame_id="child") + edge = TransformStamped( + header=Header(frame_id="first", stamp=Time(sec=0, nanosec=0)), + child_frame_id="child", + transform=Transform( + translation=Vector3(x=0.0, y=0.0, z=0.0), + rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ) with patch("rerun.log") as log: tree.update([edge]) edge.header.frame_id = "second" diff --git a/dimos/visualization/rerun/test_websocket_server.py b/dimos/visualization/rerun/test_websocket_server.py index c771ac3caee..12f7c241710 100644 --- a/dimos/visualization/rerun/test_websocket_server.py +++ b/dimos/visualization/rerun/test_websocket_server.py @@ -21,7 +21,8 @@ import threading from typing import Any -from dimos_generated.geometry_msgs.msg import PointStamped, Twist +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.geometry_msgs.msg import PointStamped, Twist, Vector3 import pytest import websockets.asyncio.client as ws_client @@ -131,7 +132,7 @@ def test_click_publishes_point_stamped( unsub() assert len(received) == 1 - point = PointStamped.decode(received[0].encode()) + point = cdr_decode(cdr_encode(received[0]), PointStamped) assert point.point.x == pytest.approx(1.5) assert point.point.y == pytest.approx(2.5) assert point.point.z == pytest.approx(0.0) @@ -171,7 +172,9 @@ def test_stop_publishes_zero_twist( unsub() assert len(received) == 1 - assert Twist.decode(received[0].encode()) == Twist() + assert cdr_decode(cdr_encode(received[0]), Twist) == Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) def test_invalid_json_does_not_crash(server: RerunWebSocketServer) -> None: @@ -189,7 +192,9 @@ async def _send_bad() -> None: try: asyncio.run(_send_bad()) assert done.wait(timeout=2.0) - assert received == [Twist()] + assert received == [ + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ] finally: unsubscribe() diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index c7ac89963c2..42199dcbd5a 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -223,4 +223,6 @@ def _num(v: Any) -> float: ) elif msg_type == "stop": - self.tele_cmd_vel.publish(Twist()) + self.tele_cmd_vel.publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) diff --git a/dimos/web/cdr_codec.py b/dimos/web/cdr_codec.py index 62d83e3770a..1cd0752c36d 100644 --- a/dimos/web/cdr_codec.py +++ b/dimos/web/cdr_codec.py @@ -18,17 +18,18 @@ import re from typing import Any +from dimos_cdrgen.registry import encode, message_types, schema + CDR_V1_SUFFIX = ".cdr.v1" _TYPE_NAME = re.compile(r"^[A-Za-z][A-Za-z0-9_]*/msg/[A-Za-z][A-Za-z0-9_]*$") def cdr_type_name(message_type: type[Any]) -> str | None: - name = getattr(message_type, "msg_name", None) + name = getattr(message_type, "__msgtype__", None) if ( isinstance(name, str) and _TYPE_NAME.fullmatch(name) - and isinstance(getattr(message_type, "schema", None), str) - and callable(getattr(message_type, "encode", None)) + and message_types().get(name) is message_type ): return name return None @@ -49,7 +50,7 @@ def export_schema(message_type: type[Any]) -> dict[str, str]: name = cdr_type_name(message_type) if name is None: raise ValueError(f"{message_type.__qualname__} has no generated CDR schema") - return {"type": name, "definition": message_type.schema} + return {"type": name, "definition": schema(name)} def check_cdr_params(params: Mapping[str, Any]) -> None: @@ -70,5 +71,4 @@ def encode_cdr_v1(msg: Any, params: Mapping[str, Any]) -> bytes: raise ValueError( f"message type does not match declared channel type {params['cdr']['type']}" ) - payload: bytes = msg.encode() - return payload + return encode(msg) diff --git a/dimos/web/cockpit.py b/dimos/web/cockpit.py index 22dbe908c44..cc4d61d7479 100644 --- a/dimos/web/cockpit.py +++ b/dimos/web/cockpit.py @@ -46,13 +46,13 @@ else: from typing_extensions import Self -from dimos_lcm.std_msgs import Bool +from dimos_generated.geometry_msgs.msg import PointStamped +from dimos_generated.nav_msgs.msg import Path +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Bool -from dimos.msgs.geometry_msgs.PointStamped import PointStamped -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.web.codecs import is_generic_lcm_encoding -from dimos.web.lcm_codec import default_encoding, export_schema, schema_class_for +from dimos.web.cdr_codec import default_encoding, export_schema +from dimos.web.codecs import is_generic_cdr_encoding from dimos.web.relay_bridge.manifest import ( MANIFEST_VERSION, MAX_MANIFEST_ID_LEN, @@ -166,9 +166,9 @@ class Channel: `encoding` names a codec: a registered @web_encoder (dimos.web.codecs) whose message type must match `message_type`, the generic "json.v1" for JSON-shaped types and dataclasses (rx) / JSON scalars, lists and dicts - (tx decode), or `.lcm.v1` for rx DimOS messages with a - dimos_lcm schema (the frame is `lcm_encode()`, the schema rides - params["lcm"]). None picks the default: `.lcm.v1` for rx + (tx decode), or `.cdr.v1` for rx DimOS messages with a + generated schema (the frame is `encode()`, the schema rides + params["cdr"]). None picks the default: `.cdr.v1` for rx DimOS messages, "json.v1" otherwise. A dir="tx" channel is a generic browser publish input: it must declare @@ -278,10 +278,10 @@ def _request_of(channel: Channel) -> ChannelRequest: # Deep plain copy: the manifest and specs must not alias the (frozen) # authoring record's nested values. params = _thaw_params(channel.params or {}) - if is_generic_lcm_encoding(channel.encoding): - if "lcm" in params: - raise ValueError("params key 'lcm' is reserved for the LCM schema") - params["lcm"] = export_schema(schema_class_for(channel.message_type)) + if is_generic_cdr_encoding(channel.encoding): + if "cdr" in params: + raise ValueError("params key 'cdr' is reserved for the CDR schema") + params["cdr"] = export_schema(channel.message_type) return ChannelRequest( channel.stream, channel.dir, diff --git a/dimos/web/codecs.py b/dimos/web/codecs.py index f9a6c905da2..a9f9d4d9a71 100644 --- a/dimos/web/codecs.py +++ b/dimos/web/codecs.py @@ -37,7 +37,7 @@ def decode_goal(value: dict[str, Any], context: PublishContext) -> PoseStamped: return `bytes`, an `EncodedPayload` when the frame needs header meta, or None to skip the sample. Two codecs apply without registration: json.v1 for JSON-shaped message types, and `.cdr.v1` (dimos/web/cdr_codec.py) -for generated messages (the frame is `msg.encode()`, +for generated messages (the frame is registry-encoded CDR, the schema rides the channel's params["cdr"]). Anything else potentially large (images, arrays, bytes) needs an explicit @web_encoder. """ @@ -461,7 +461,7 @@ def resolve_encoder(encoding: str, message_type: type[Any]) -> EncoderDef: # issubclass() is unavailable, and importing dimos.msgs here would # defeat this module's import-lightness. Some wire messages (Image) # are dataclasses too - the marker must win. - is_dimos_msg = callable(getattr(message_type, "encode", None)) + is_dimos_msg = isinstance(getattr(message_type, "__msgtype__", None), str) plain_dataclass = dataclasses.is_dataclass(message_type) and not is_dimos_msg if message_type not in _JSON_V1_TYPES and not plain_dataclass: raise ValueError( diff --git a/dimos/web/relay_bridge/builtin_codecs.py b/dimos/web/relay_bridge/builtin_codecs.py index 267931348e1..9cd22d1ec1f 100644 --- a/dimos/web/relay_bridge/builtin_codecs.py +++ b/dimos/web/relay_bridge/builtin_codecs.py @@ -27,16 +27,18 @@ from typing import Any import zlib -from dimos_lcm.std_msgs import Bool +from dimos_generated.geometry_msgs.msg import Point, PointStamped, PoseStamped +from dimos_generated.nav_msgs.msg import OccupancyGrid, Path +from dimos_generated.sensor_msgs.msg import Image, PointCloud2 +from dimos_generated.std_msgs.msg import Bool import numpy as np from dimos.mapping.voxels.keys import KEY_OFFSET, pack_indices, unpack_keys -from dimos.msgs.geometry_msgs.PointStamped import PointStamped -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid, block_max_reduce -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.geometry import yaw +from dimos.msgs.image import image_to_jpeg +from dimos.msgs.occupancy import block_max_reduce, occupancy_view +from dimos.msgs.pointcloud import pointcloud_xyz +from dimos.msgs.time import header_now, to_seconds from dimos.utils.generic import finite_number from dimos.web.codecs import EncodedPayload, web_decoder, web_encoder @@ -56,9 +58,9 @@ def _check_jpeg_params(params: Mapping[str, Any]) -> None: @web_encoder("jpeg.v1", check_params=_check_jpeg_params) def encode_jpeg(msg: Image, params: Mapping[str, Any]) -> EncodedPayload: - # TurboJPEG via the message's own encoder (handles BGR/RGB/gray inputs). + # OpenCV via the image conversion helper (handles BGR/RGB/gray inputs). return EncodedPayload( - msg.to_jpeg_bytes(quality=params.get("quality", _DEFAULT_JPEG_QUALITY)), + image_to_jpeg(msg, quality=params.get("quality", _DEFAULT_JPEG_QUALITY)), {"w": msg.width, "h": msg.height}, ) @@ -75,11 +77,11 @@ def decode_text(msg: str) -> str: @web_encoder("pose.json.v1") def encode_pose(msg: PoseStamped) -> bytes: pose = { - "x": msg.position.x, - "y": msg.position.y, - "z": msg.position.z, - "yaw": msg.yaw, - "ts": msg.ts, + "x": msg.pose.position.x, + "y": msg.pose.position.y, + "z": msg.pose.position.z, + "yaw": yaw(msg.pose.orientation), + "ts": to_seconds(msg.header.stamp), } return json.dumps(pose, separators=(",", ":")).encode() @@ -87,7 +89,7 @@ def encode_pose(msg: PoseStamped) -> bytes: @web_encoder("path.json.v1") def encode_path(msg: Path) -> bytes: # Empty paths must reach the viewer to clear the overlay. - points = [[round(p.x, 3), round(p.y, 3)] for p in msg.poses] + points = [[round(p.pose.position.x, 3), round(p.pose.position.y, 3)] for p in msg.poses] return json.dumps(points, separators=(",", ":"), allow_nan=False).encode() @@ -96,7 +98,8 @@ def decode_point(msg: dict[str, Any]) -> PointStamped: if not isinstance(msg, dict): raise ValueError(f"point.json.v1 wants an object, got {type(msg).__name__}") return PointStamped( - finite_number(msg.get("x"), "x"), finite_number(msg.get("y"), "y"), frame_id="world" + header=header_now("world"), + point=Point(x=finite_number(msg.get("x"), "x"), y=finite_number(msg.get("y"), "y"), z=0.0), ) @@ -120,10 +123,10 @@ def decode_bool(msg: bool) -> Bool: @web_encoder("costmap.zlib.v1") def encode_costmap(msg: OccupancyGrid) -> EncodedPayload | None: - grid = msg.grid + grid = occupancy_view(msg) if grid.size == 0: return None # mapper still warming up; nothing to draw - res = msg.resolution + res = msg.info.resolution side = max(grid.shape) if side > _COSTMAP_MAX_SIDE: factor = -(-side // _COSTMAP_MAX_SIDE) @@ -134,12 +137,12 @@ def encode_costmap(msg: OccupancyGrid) -> EncodedPayload | None: # int8 -1 is byte 0xff and 0..100 are byte-identical, so the raw buffer # already is the wire payload - no mask/astype/tobytes copies. cells = np.ascontiguousarray(grid) - origin = msg.origin + origin = msg.info.origin meta = { "w": w, "h": h, "res": res, - "origin": [origin.position.x, origin.position.y, origin.yaw], + "origin": [origin.position.x, origin.position.y, yaw(origin.orientation)], } return EncodedPayload(zlib.compress(cells, _COSTMAP_ZLIB_LEVEL), meta) @@ -183,7 +186,7 @@ def _voxel_keys(idx: np.ndarray) -> np.ndarray: @web_encoder("voxels.zlib.v1", check_params=_check_voxel_params) def encode_voxels(msg: PointCloud2, params: Mapping[str, Any]) -> EncodedPayload: res = float(params.get("res", _DEFAULT_VOXEL_RES)) - pts = msg.points_f32() + pts = pointcloud_xyz(msg).astype(np.float32) pts = pts[np.isfinite(pts).all(axis=1)] idx = np.floor(pts / res).astype(np.int64) # Indices live in the key's 21-bit fields (+-52 km at 5 cm); beyond that diff --git a/dimos/web/relay_bridge/gen_cdr_fixtures.py b/dimos/web/relay_bridge/gen_cdr_fixtures.py index 3bbabe2efce..fb6370fca6b 100644 --- a/dimos/web/relay_bridge/gen_cdr_fixtures.py +++ b/dimos/web/relay_bridge/gen_cdr_fixtures.py @@ -18,13 +18,15 @@ import json from typing import Any +from dimos_cdrgen.definitions import Definitions +from dimos_cdrgen.registry import encode as cdr_encode from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.dimos_msgs.msg import LineSegment3D, LineSegments3D -from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Vector3 from dimos_generated.sensor_msgs.msg import Image, Imu, JointState from dimos_generated.std_msgs.msg import Empty, Header, Int64, String +import numpy as np -from dimos.message_codegen.definitions import Definitions from dimos.web.cdr_codec import CDR_V1_SUFFIX, export_schema from dimos.web.relay_bridge.locate import find_web_dir @@ -37,14 +39,20 @@ def build_messages() -> list[tuple[str, Any]]: PoseStamped( header=header, pose=Pose( - position=Point(x=1.5, y=-2.5, z=0.25), orientation=Quaternion(z=0.6, w=0.8) + position=Point(x=1.5, y=-2.5, z=0.25), + orientation=Quaternion(z=0.6, w=0.8, x=0.0, y=0.0), ), ), ), ( "custom_segments", LineSegments3D( - header=header, segments=[LineSegment3D(start=Point(x=1), end=Point(y=2), weight=4)] + header=header, + segments=[ + LineSegment3D( + start=Point(x=1, y=0.0, z=0.0), end=Point(y=2, x=0.0, z=0.0), weight=4 + ) + ], ), ), ( @@ -55,11 +63,32 @@ def build_messages() -> list[tuple[str, Any]]: height=1, encoding="rgb8", step=6, - data=[0, 1, 2, 253, 254, 255], + data=np.array([0, 1, 2, 253, 254, 255], dtype=np.uint8), + is_bigendian=0, + ), + ), + ( + "imu", + Imu( + header=header, + orientation_covariance=np.asarray(range(9), dtype=np.float64), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + angular_velocity=Vector3(x=0.0, y=0.0, z=0.0), + angular_velocity_covariance=np.zeros(9, dtype=np.float64), + linear_acceleration=Vector3(x=0.0, y=0.0, z=0.0), + linear_acceleration_covariance=np.zeros(9, dtype=np.float64), + ), + ), + ( + "joints", + JointState( + header=header, + name=["joint_a", "joint_b"], + position=np.array([1.5, -2.5], dtype=np.float64), + velocity=np.array([], dtype=np.float64), + effort=np.array([], dtype=np.float64), ), ), - ("imu", Imu(header=header, orientation_covariance=range(9))), - ("joints", JointState(header=header, name=["joint_a", "joint_b"], position=[1.5, -2.5])), ("int64", Int64(data=-9007199254740993)), ("empty", Empty()), ("string", String(data="robot 🦾 café")), @@ -67,16 +96,20 @@ def build_messages() -> list[tuple[str, Any]]: def build_vectors() -> list[dict[str, Any]]: - definitions = {message.name: message for message in Definitions([]).resolve()} + definitions = {message.name: message for message in Definitions([], installed=True).resolve()} def walk(msg: Any) -> dict[str, Any]: result: dict[str, Any] = {} - for field in definitions[msg.msg_name].fields: + for field in definitions[msg.__msgtype__].fields: value = getattr(msg, field.name) convert = ( walk if field.type.nested - else (str if field.type.name in ("int64", "uint64") else lambda x: x) + else ( + str + if field.type.name in ("int64", "uint64") + else lambda x: x.item() if isinstance(x, np.generic) else x + ) ) result[field.name] = ( [convert(item) for item in value] if field.type.is_array else convert(value) @@ -86,11 +119,11 @@ def walk(msg: Any) -> dict[str, Any]: return [ { "name": name, - "encoding": msg.msg_name + CDR_V1_SUFFIX, + "encoding": msg.__msgtype__ + CDR_V1_SUFFIX, "schema": export_schema(type(msg)), "value": walk(msg), - "payload_b64": base64.b64encode(msg.encode()).decode(), - "big_endian_b64": base64.b64encode(msg.encode(little_endian=False)).decode(), + "payload_b64": base64.b64encode(cdr_encode(msg)).decode(), + "big_endian_b64": base64.b64encode(cdr_encode(msg, little_endian=False)).decode(), } for name, msg in build_messages() ] diff --git a/dimos/web/relay_bridge/gen_costmap_fixtures.py b/dimos/web/relay_bridge/gen_costmap_fixtures.py index 2c303de5156..7e392f6b452 100644 --- a/dimos/web/relay_bridge/gen_costmap_fixtures.py +++ b/dimos/web/relay_bridge/gen_costmap_fixtures.py @@ -32,8 +32,10 @@ import math from typing import Any +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid +from dimos_generated.std_msgs.msg import Header import numpy as np from dimos.web.relay_bridge.builtin_codecs import encode_costmap @@ -43,12 +45,20 @@ def grid_msg(rows: list[list[int]], res: float, x: float, y: float, yaw: float) -> OccupancyGrid: """Fixture-shaped OccupancyGrid; also used by test_costmap_encoding.py.""" origin = Pose( - position=Point(x=x, y=y), orientation=Quaternion(z=math.sin(yaw / 2), w=math.cos(yaw / 2)) + position=Point(x=x, y=y, z=0.0), + orientation=Quaternion(z=math.sin(yaw / 2), w=math.cos(yaw / 2), x=0.0, y=0.0), ) grid = np.array(rows, dtype=np.int8) return OccupancyGrid( - info=MapMetaData(width=grid.shape[1], height=grid.shape[0], resolution=res, origin=origin), - data=grid.ravel(), + info=MapMetaData( + width=grid.shape[1], + height=grid.shape[0], + resolution=res, + origin=origin, + map_load_time=Time(sec=0, nanosec=0), + ), + data=np.asarray(grid.ravel(), dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) diff --git a/dimos/web/relay_bridge/gen_voxel_fixtures.py b/dimos/web/relay_bridge/gen_voxel_fixtures.py index 81b43de726a..163a50e9646 100644 --- a/dimos/web/relay_bridge/gen_voxel_fixtures.py +++ b/dimos/web/relay_bridge/gen_voxel_fixtures.py @@ -31,17 +31,21 @@ import json from typing import Any +from dimos_generated.sensor_msgs.msg import PointCloud2 +from dimos_generated.std_msgs.msg import Header import numpy as np -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import time_from_seconds from dimos.web.relay_bridge.builtin_codecs import encode_voxels from dimos.web.relay_bridge.locate import find_web_dir def cloud_msg(points: list[list[float]]) -> PointCloud2: """Fixture-shaped PointCloud2; also used by test_voxel_encoding.py.""" - return PointCloud2.from_numpy( - np.array(points, dtype=np.float32).reshape(-1, 3), frame_id="world", timestamp=1752576000.5 + return pointcloud_from_xyz( + np.array(points, dtype=np.float32).reshape(-1, 3), + header=Header(stamp=time_from_seconds(1752576000.5), frame_id="world"), ) diff --git a/dimos/web/relay_bridge/relay_bridge_module.py b/dimos/web/relay_bridge/relay_bridge_module.py index 5b7d724e8b0..41c02d10f97 100644 --- a/dimos/web/relay_bridge/relay_bridge_module.py +++ b/dimos/web/relay_bridge/relay_bridge_module.py @@ -53,15 +53,14 @@ from typing import TYPE_CHECKING, Any, TypeVar import webbrowser +from dimos_generated.geometry_msgs.msg import PoseStamped, Twist +from dimos_generated.nav_msgs.msg import OccupancyGrid +from dimos_generated.sensor_msgs.msg import Image from reactivex.disposable import Disposable from dimos.core.coordination.blueprints import Blueprint, autoconnect from dimos.core.module import Module from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.sensor_msgs.Image import Image from dimos.utils.logging_config import setup_logger from dimos.web.codecs import EncodedPayload, encoder_definition diff --git a/dimos/web/relay_bridge/rtc_publisher.py b/dimos/web/relay_bridge/rtc_publisher.py index 799a0b3a086..19ca6522866 100644 --- a/dimos/web/relay_bridge/rtc_publisher.py +++ b/dimos/web/relay_bridge/rtc_publisher.py @@ -33,8 +33,8 @@ RTCPeerConnection, RTCSessionDescription, ) +from dimos_generated.sensor_msgs.msg import Image -from dimos.msgs.sensor_msgs.Image import Image from dimos.protocol.pubsub.impl.webrtc.providers.sdp import propagate_bundle_candidates from dimos.protocol.pubsub.impl.webrtc.providers.spec import wait_connected from dimos.protocol.pubsub.impl.webrtc.providers.video_track import ( diff --git a/dimos/web/relay_bridge/teleop.py b/dimos/web/relay_bridge/teleop.py index 8fea1de4271..19a4233815b 100644 --- a/dimos/web/relay_bridge/teleop.py +++ b/dimos/web/relay_bridge/teleop.py @@ -24,8 +24,8 @@ import math import time -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos_generated.geometry_msgs.msg import Twist, Vector3 + from dimos.utils.generic import finite_number from dimos.utils.logging_config import setup_logger from dimos.web.relay_bridge.protocol import ( @@ -122,7 +122,7 @@ def on_twist(self, msg: WireTwist) -> None: self._zero("release") return self._driving = True - self._publish(Twist(linear=Vector3(vx, vy, 0.0), angular=Vector3(0.0, 0.0, wz))) + self._publish(Twist(linear=Vector3(x=vx, y=vy, z=0.0), angular=Vector3(x=0.0, y=0.0, z=wz))) def on_stop(self, msg: WireStop) -> None: """E-stop: unconditional zero, even from idle - it must also cancel @@ -195,4 +195,6 @@ def _zero(self, reason: str, *, force: bool = False) -> None: return self._driving = False logger.warning(f"relay bridge teleop: zero twist ({reason})") - self._publish(Twist.zero()) + self._publish( + Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) diff --git a/dimos/web/relay_bridge/test_costmap_encoding.py b/dimos/web/relay_bridge/test_costmap_encoding.py index 43b207d1540..7eeae3b349f 100644 --- a/dimos/web/relay_bridge/test_costmap_encoding.py +++ b/dimos/web/relay_bridge/test_costmap_encoding.py @@ -23,8 +23,11 @@ from typing import Any import zlib -from dimos_generated.geometry_msgs.msg import Point, Pose +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, Quaternion from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest @@ -60,7 +63,7 @@ def test_encoder_reproduces_golden_vector(vec: dict[str, Any]) -> None: def test_encoder_handles_wire_decoded_grid() -> None: """The bridge derives yaw from the quaternion after a CDR roundtrip.""" msg = grid_msg([[0, 100], [-1, 50]], 0.05, -1.25, 2.5, 0.5) - decoded = OccupancyGrid.decode(msg.encode()) + decoded = cdr_decode(cdr_encode(msg), OccupancyGrid) encoded = encode_costmap(decoded) assert encoded is not None payload, meta = encoded.payload, encoded.meta @@ -69,7 +72,25 @@ def test_encoder_handles_wire_decoded_grid() -> None: def test_empty_grid_encodes_to_none() -> None: - assert encode_costmap(OccupancyGrid()) is None + assert ( + encode_costmap( + OccupancyGrid( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + info=MapMetaData( + map_load_time=Time(sec=0, nanosec=0), + resolution=0.0, + width=0, + height=0, + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([], dtype=np.int8), + ) + ) + is None + ) def test_oversized_grid_is_downsampled_within_budget() -> None: @@ -81,9 +102,13 @@ def test_oversized_grid_is_downsampled_within_budget() -> None: width=rows.shape[1], height=rows.shape[0], resolution=0.05, - origin=Pose(position=Point(x=1, y=2)), + origin=Pose( + position=Point(x=1, y=2, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + map_load_time=Time(sec=0, nanosec=0), ), - data=rows.ravel(), + data=np.asarray(rows.ravel(), dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) encoded = encode_costmap(msg) assert encoded is not None @@ -100,8 +125,18 @@ def test_oversized_grid_is_downsampled_within_budget() -> None: def test_grid_at_exactly_max_side_is_not_downsampled() -> None: rows = np.full((4, 2048), 7, dtype=np.int8) msg = OccupancyGrid( - info=MapMetaData(width=rows.shape[1], height=rows.shape[0], resolution=0.05), - data=rows.ravel(), + info=MapMetaData( + width=rows.shape[1], + height=rows.shape[0], + resolution=0.05, + map_load_time=Time(sec=0, nanosec=0), + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.asarray(rows.ravel(), dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) encoded = encode_costmap(msg) assert encoded is not None @@ -113,8 +148,18 @@ def test_grid_at_exactly_max_side_is_not_downsampled() -> None: def test_non_square_oversized_grid_uses_ceil_factor() -> None: rows = np.zeros((100, 4100), dtype=np.int8) msg = OccupancyGrid( - info=MapMetaData(width=rows.shape[1], height=rows.shape[0], resolution=0.05), - data=rows.ravel(), + info=MapMetaData( + width=rows.shape[1], + height=rows.shape[0], + resolution=0.05, + map_load_time=Time(sec=0, nanosec=0), + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.asarray(rows.ravel(), dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) encoded = encode_costmap(msg) assert encoded is not None diff --git a/dimos/web/relay_bridge/test_map_codecs.py b/dimos/web/relay_bridge/test_map_codecs.py index c901f1c2429..c936e06b0b0 100644 --- a/dimos/web/relay_bridge/test_map_codecs.py +++ b/dimos/web/relay_bridge/test_map_codecs.py @@ -16,9 +16,10 @@ import math -from dimos_generated.geometry_msgs.msg import Point, PointStamped, Pose, PoseStamped +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, PointStamped, Pose, PoseStamped, Quaternion from dimos_generated.nav_msgs.msg import Path -from dimos_generated.std_msgs.msg import Bool +from dimos_generated.std_msgs.msg import Bool, Header import pytest from dimos.web.codecs import resolve_decoder, resolve_encoder @@ -26,18 +27,29 @@ def _pose(x: float, y: float) -> PoseStamped: - return PoseStamped(pose=Pose(position=Point(x=x, y=y))) + return PoseStamped( + pose=Pose( + position=Point(x=x, y=y, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) def test_path_encodes_a_rounded_xy_polyline() -> None: - path = Path(poses=[_pose(1.23456, -2.0), _pose(0.0004, 3.5)]) + path = Path( + poses=[_pose(1.23456, -2.0), _pose(0.0004, 3.5)], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) assert encode_path(path) == b"[[1.235,-2.0],[0.0,3.5]]" assert resolve_encoder("path.json.v1", Path).encode is encode_path def test_empty_path_encodes_as_a_clear() -> None: # The planner publishes Path() on cancel and arrival: the overlay must go. - assert encode_path(Path()) == b"[]" + assert ( + encode_path(Path(header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), poses=[])) + == b"[]" + ) def test_point_decodes_a_click() -> None: diff --git a/dimos/web/relay_bridge/test_relay_bridge_authoring.py b/dimos/web/relay_bridge/test_relay_bridge_authoring.py index a870b56f05b..0ee7144ff66 100644 --- a/dimos/web/relay_bridge/test_relay_bridge_authoring.py +++ b/dimos/web/relay_bridge/test_relay_bridge_authoring.py @@ -27,6 +27,8 @@ import struct from typing import Any +from dimos_cdrgen.registry import encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, PointStamped, @@ -166,8 +168,11 @@ def _encode_boom(msg: NavPath) -> bytes: header=Header(stamp=time_from_seconds(1), frame_id="world"), poses=[ PoseStamped( - header=Header(stamp=time_from_seconds(1)), - pose=Pose(position=Point(x=1.5, y=-2.5), orientation=Quaternion(w=1)), + header=Header(stamp=time_from_seconds(1), frame_id=""), + pose=Pose( + position=Point(x=1.5, y=-2.5, z=0.0), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), ) ], ) @@ -300,15 +305,19 @@ def test_cdr_channel_ships_generated_encode_bytes(monkeypatch) -> None: assert wait_until(lambda: pose.subscribers) module._min_interval = {"pose": 0.0} msg = PoseStamped( - header=Header(stamp=time_from_seconds(2)), - pose=Pose(position=Point(x=1, y=2), orientation=Quaternion(w=1)), + header=Header(stamp=time_from_seconds(2), frame_id=""), + pose=Pose( + position=Point(x=1, y=2, z=0.0), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), ) pose.publish(msg) assert wait_until(lambda: clients[0].frames) - assert clients[0].frames[0] == ("pose", msg.encode(), "reliable", None) + assert clients[0].frames[0] == ("pose", cdr_encode(msg), "reliable", None) # A sample of another type has a different declared message name: dropped and # logged, never sent to a browser that compiled the PoseStamped schema. - pose.publish(Pose(position=Point(x=1, y=2), orientation=Quaternion(w=1))) + pose.publish( + Pose(position=Point(x=1, y=2, z=0.0), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0)) + ) assert wait_until(lambda: bool(exceptions)) assert "pose" in exceptions[0] flush_loop(module) @@ -341,8 +350,11 @@ def test_encoder_failure_is_isolated_and_rate_limited(monkeypatch) -> None: bad.publish(_NAV_PATH) pose.publish( PoseStamped( - header=Header(stamp=time_from_seconds(2)), - pose=Pose(position=Point(x=1, y=2), orientation=Quaternion(w=1)), + header=Header(stamp=time_from_seconds(2), frame_id=""), + pose=Pose( + position=Point(x=1, y=2, z=0.0), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), ) ) assert wait_until(lambda: clients[0].frames) @@ -428,7 +440,7 @@ def start_with(spec: RuntimeChannelSpec, match: str) -> None: start_with(replace(good, dir="tx"), "does not match its compiled runtime spec") start_with( replace(good, message_type=Twist), - "message type Twist does not match the RelayBridgeModule port type PoseStamped", + f"message type {Twist.__qualname__} does not match the RelayBridgeModule port type {PoseStamped.__qualname__}", ) @@ -782,8 +794,10 @@ def test_publish_frame_with_unusable_meta_is_dropped(monkeypatch) -> None: def _nav_path(*xy: tuple[float, float]) -> NavPath: poses = [ PoseStamped( - header=Header(stamp=time_from_seconds(1)), - pose=Pose(position=Point(x=x, y=y), orientation=Quaternion(w=1)), + header=Header(stamp=time_from_seconds(1), frame_id=""), + pose=Pose( + position=Point(x=x, y=y, z=0.0), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), ) for x, y in xy ] @@ -836,7 +850,7 @@ def test_map2d_path_replays_and_paces(map_bridge) -> None: assert json.loads(offers[0][0]) == [[1.5, -2.5]] # A planner burst must not leave the viewer stuck on the empty path. - path.publish(NavPath()) + path.publish(NavPath(header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), poses=[])) path.publish(_nav_path((0.25, 0.5), (1.0, 1.0), (1.75, 0.5))) assert wait_until(lambda: len(offers) == 3) assert [json.loads(payload) for payload, _ in offers[1:]] == [ diff --git a/dimos/web/relay_bridge/test_relay_bridge_e2e.py b/dimos/web/relay_bridge/test_relay_bridge_e2e.py index 77628f74f70..ecf4efa803b 100644 --- a/dimos/web/relay_bridge/test_relay_bridge_e2e.py +++ b/dimos/web/relay_bridge/test_relay_bridge_e2e.py @@ -43,15 +43,16 @@ from cryptography.hazmat.primitives import hashes, serialization from cryptography.hazmat.primitives.asymmetric import ec from cryptography.x509.oid import NameOID +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Twist, Vector3 +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid +from dimos_generated.std_msgs.msg import Header import numpy as np import pytest from dimos.core.transport import pLCMTransport -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_from_array +from dimos.msgs.time import time_from_seconds from dimos.web.relay_bridge.config import RelayBridgeConfig, default_manifest from dimos.web.relay_bridge.e2e_support import ( arm_teleop, @@ -73,12 +74,25 @@ pytestmark = pytest.mark.self_hosted ROBOT_ID = "bridge-e2e" -POSE = PoseStamped(ts=42.5, position=[1.5, -2.5, 0.25], orientation=[0.0, 0.0, 0.0, 1.0]) +POSE = PoseStamped( + header=Header(stamp=time_from_seconds(42.5), frame_id=""), + pose=Pose( + position=Point(x=1.5, y=-2.5, z=0.25), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), +) COSTMAP_GRID = OccupancyGrid( - grid=np.array([[-1, 0, 50], [100, 0, -1]], dtype=np.int8), - resolution=0.05, - origin=Pose(-1.25, 2.5, 0.0), - ts=42.5, + data=np.array([-1, 0, 50, 100, 0, -1], dtype=np.int8), + info=MapMetaData( + width=3, + height=2, + resolution=0.05, + origin=Pose( + position=Point(x=-1.25, y=2.5, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + map_load_time=Time(sec=0, nanosec=0), + ), + header=Header(stamp=time_from_seconds(42.5), frame_id=""), ) COSTMAP_CELLS = bytes([255, 0, 50, 100, 0, 255]) _LISTENER = """ @@ -104,7 +118,7 @@ def __init__(self) -> None: self.stop = threading.Event() arr = np.zeros((48, 64, 3), dtype=np.uint8) arr[:, :, 1] = np.linspace(0, 255, 64, dtype=np.uint8) - self._image = Image.from_numpy(arr) + self._image = image_from_array(arr, encoding="rgb8") self._thread = threading.Thread(target=self._run, daemon=True) def start(self) -> None: @@ -353,10 +367,18 @@ def test_costmap_replay_reflects_publishes_while_unwatched(bridge: RelayBridgeMo # staleness on the wire), not the replay time. _wait_costmap_unsubscribed(bridge) grid_b = OccupancyGrid( - grid=np.array([[100, 100, 100], [0, 0, -1]], dtype=np.int8), - resolution=0.05, - origin=Pose(-1.25, 2.5, 0.0), - ts=43.0, + data=np.array([100, 100, 100, 0, 0, -1], dtype=np.int8), + info=MapMetaData( + width=3, + height=2, + resolution=0.05, + origin=Pose( + position=Point(x=-1.25, y=2.5, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + map_load_time=Time(sec=0, nanosec=0), + ), + header=Header(stamp=time_from_seconds(43), frame_id=""), ) t0 = time.time() publisher = _CostmapPublisher(grid_b) @@ -365,14 +387,14 @@ def test_costmap_replay_reflects_publishes_while_unwatched(bridge: RelayBridgeMo deadline = time.monotonic() + 10 while time.monotonic() < deadline: cached = bridge._last_msg.get("global_costmap") - if cached is not None and cached[0].grid[0, 0] == 100: + if cached is not None and cached[0].data[0] == 100: break time.sleep(0.05) finally: publisher.close() t1 = time.time() cached = bridge._last_msg.get("global_costmap") - assert cached is not None and cached[0].grid[0, 0] == 100, "cache never saw grid B" + assert cached is not None and cached[0].data[0] == 100, "cache never saw grid B" frame = asyncio.run(_watch_costmap(bridge, expected_cells=bytes([100, 100, 100, 0, 0, 255]))) assert t0 <= frame.header.ts <= t1 @@ -487,7 +509,16 @@ async def flow() -> None: stop = asyncio.Event() driver = asyncio.create_task(_drive(viewer, 0, stop)) - await _until(lambda: any(not t.is_zero() for t in twists), "no twist published") + await _until( + lambda: any( + t + != Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) + for t in twists + ), + "no twist published", + ) # E-stop: the stop datagram publishes an unconditional zero. stop.set() @@ -495,7 +526,10 @@ async def flow() -> None: marker = len(twists) viewer.send_control(WireStop(seq=seq + 1, ts=time.time())) await _until( - lambda: len(twists) > marker and twists[-1].is_zero(), "no zero after stop" + lambda: len(twists) > marker + and twists[-1] + == Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + "no zero after stop", ) asyncio.run(flow()) @@ -540,14 +574,25 @@ async def flow() -> None: stop = asyncio.Event() driver = asyncio.create_task(_drive(viewer, 0, stop)) - await _until(lambda: any(not t.is_zero() for t in twists), "no twist published") + await _until( + lambda: any( + t + != Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) + for t in twists + ), + "no twist published", + ) marker = len(twists) assert relay._process is not None killed_at = time.monotonic() relay._process.kill() await _until( - lambda: len(twists) > marker and twists[-1].is_zero(), + lambda: len(twists) > marker + and twists[-1] + == Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), "no zero after relay kill (deadman broken?)", timeout=5.0, ) @@ -590,13 +635,27 @@ async def held() -> bool: # The holder drives; the bystander's twists are gated out. second.send_control(WireTwist(vx=9.0, vy=0.0, wz=0.0, seq=1, ts=time.time())) holder.send_control(WireTwist(vx=0.5, vy=0.0, wz=0.0, seq=1, ts=time.time())) - await _until(lambda: any(not t.is_zero() for t in twists), "no twist published") + await _until( + lambda: any( + t + != Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), + angular=Vector3(x=0.0, y=0.0, z=0.0), + ) + for t in twists + ), + "no twist published", + ) assert all(t.linear.x != 9.0 for t in twists) # Holder disconnected mid-drive: the relay releases the lease and # sends teleop_stop, so the bridge zeroes without waiting for the # watchdog, and the second viewer can now arm and drive. - await _until(lambda: twists[-1].is_zero(), "no zero after holder disconnect") + await _until( + lambda: twists[-1] + == Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + "no zero after holder disconnect", + ) await arm_teleop(second) marker = len(twists) second.send_control(WireTwist(vx=0.25, vy=0.0, wz=0.0, seq=1, ts=time.time())) diff --git a/dimos/web/relay_bridge/test_relay_bridge_module.py b/dimos/web/relay_bridge/test_relay_bridge_module.py index 05252316293..a3e6b3f5637 100644 --- a/dimos/web/relay_bridge/test_relay_bridge_module.py +++ b/dimos/web/relay_bridge/test_relay_bridge_module.py @@ -37,17 +37,18 @@ from typing import Any import zlib +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion, Twist, Vector3 +from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid +from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header import numpy as np from pydantic import ValidationError import pytest from dimos.core.global_config import global_config -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_from_array +from dimos.msgs.time import time_from_seconds from dimos.simulation.mujoco.constants import VIDEO_FPS from dimos.web.relay_bridge import builtin_codecs, relay_bridge_module, teleop from dimos.web.relay_bridge.config import RelayBridgeConfig, default_manifest, resolve_robot_info @@ -200,7 +201,12 @@ def test_encode_paths_and_max_hz_gate(bridge) -> None: lambda: image_transport(module).subscribers and odom_transport(module).subscribers ) - pose = PoseStamped(ts=42.5, position=[1.5, -2.5, 0.25], orientation=[0.0, 0.0, 0.0, 1.0]) + pose = PoseStamped( + header=Header(stamp=time_from_seconds(42.5), frame_id=""), + pose=Pose( + position=Point(x=1.5, y=-2.5, z=0.25), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), + ) odom_transport(module).publish(pose) assert module.encoded["odom"] == 1 assert wait_until(lambda: len(client.frames) == 1) @@ -209,7 +215,7 @@ def test_encode_paths_and_max_hz_gate(bridge) -> None: decoded = json.loads(payload) assert decoded == {"x": 1.5, "y": -2.5, "z": 0.25, "yaw": 0.0, "ts": 42.5} - image = Image.from_numpy(np.zeros((8, 12, 3), dtype=np.uint8)) + image = image_from_array(np.zeros((8, 12, 3), dtype=np.uint8), encoding="rgb8") image_transport(module).publish(image) assert module.encoded["color_image"] == 1 assert wait_until(lambda: client.writers["color_image"].offers) @@ -255,14 +261,14 @@ def test_no_jpeg_encode_while_unsubscribed(bridge, monkeypatch) -> None: # independent of the module.encoded bookkeeping. module, clients = bridge calls = {"n": 0} - real = Image.to_jpeg_bytes + real = builtin_codecs.image_to_jpeg def spy(self: Image, quality: int = 75) -> bytes: calls["n"] += 1 return real(self, quality=quality) - monkeypatch.setattr(Image, "to_jpeg_bytes", spy) - image = Image.from_numpy(np.zeros((8, 12, 3), dtype=np.uint8)) + monkeypatch.setattr(builtin_codecs, "image_to_jpeg", spy) + image = image_from_array(np.zeros((8, 12, 3), dtype=np.uint8), encoding="rgb8") image_transport(module).publish(image) image_transport(module).publish(image) flush_loop(module) @@ -278,10 +284,18 @@ def spy(self: Image, quality: int = 75) -> bytes: # Covers every value class of the wire contract: -1 unknown -> 255, 0 free, # graded cost, 100 lethal. COSTMAP_GRID = OccupancyGrid( - grid=np.array([[-1, 0, 50], [100, 0, -1]], dtype=np.int8), - resolution=0.05, - origin=Pose(-1.25, 2.5, 0.0), - ts=42.5, + data=np.array([-1, 0, 50, 100, 0, -1], dtype=np.int8), + info=MapMetaData( + width=3, + height=2, + resolution=0.05, + origin=Pose( + position=Point(x=-1.25, y=2.5, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + map_load_time=Time(sec=0, nanosec=0), + ), + header=Header(stamp=time_from_seconds(42.5), frame_id=""), ) COSTMAP_CELLS = bytes([255, 0, 50, 100, 0, 255]) @@ -308,7 +322,7 @@ def test_costmap_encode_roundtrip_and_meta(costmap_bridge) -> None: payload, meta = client.writers["global_costmap"].offers[0] assert zlib.decompress(payload) == COSTMAP_CELLS assert meta is not None - assert (meta["w"], meta["h"], meta["res"]) == (3, 2, 0.05) + assert (meta["w"], meta["h"], meta["res"]) == pytest.approx((3, 2, 0.05)) assert meta["origin"][:2] == [-1.25, 2.5] assert meta["origin"][2] == pytest.approx(0.0, abs=1e-12) @@ -318,7 +332,22 @@ def test_costmap_empty_grid_is_skipped(costmap_bridge) -> None: push(module, clients[0], Subs(chs=["global_costmap"], n=1)) assert wait_until(lambda: len(costmap_transport(module).subscribers) == 2) - costmap_transport(module).publish(OccupancyGrid()) + costmap_transport(module).publish( + OccupancyGrid( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + info=MapMetaData( + map_load_time=Time(sec=0, nanosec=0), + resolution=0.0, + width=0, + height=0, + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([], dtype=np.int8), + ) + ) flush_loop(module) assert module.encoded["global_costmap"] == 0 assert clients[0].writers["global_costmap"].offers == [] @@ -424,10 +453,18 @@ def test_costmap_replay_uses_message_published_while_unsubscribed(costmap_bridge push(module, client, Subs(chs=[], n=2)) assert wait_until(lambda: len(costmap_transport(module).subscribers) == 1) grid_b = OccupancyGrid( - grid=np.array([[100, 100, 100], [0, 0, 0]], dtype=np.int8), - resolution=0.05, - origin=Pose(-1.25, 2.5, 0.0), - ts=43.0, + data=np.array([100, 100, 100, 0, 0, 0], dtype=np.int8), + info=MapMetaData( + width=3, + height=2, + resolution=0.05, + origin=Pose( + position=Point(x=-1.25, y=2.5, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + map_load_time=Time(sec=0, nanosec=0), + ), + header=Header(stamp=time_from_seconds(43), frame_id=""), ) costmap_transport(module).publish(grid_b) @@ -441,7 +478,22 @@ def test_costmap_replay_uses_message_published_while_unsubscribed(costmap_bridge def test_costmap_empty_cached_grid_is_not_replayed(costmap_bridge) -> None: module, clients = costmap_bridge client = clients[0] - costmap_transport(module).publish(OccupancyGrid()) + costmap_transport(module).publish( + OccupancyGrid( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + info=MapMetaData( + map_load_time=Time(sec=0, nanosec=0), + resolution=0.0, + width=0, + height=0, + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([], dtype=np.int8), + ) + ) push(module, client, Subs(chs=["global_costmap"], n=1)) assert wait_until(lambda: len(costmap_transport(module).subscribers) == 2) session = module._session @@ -527,9 +579,11 @@ def blocking_encode(msg: PoseStamped) -> bytes: target=odom_transport(module).publish, args=( PoseStamped( - ts=42.5, - position=[1.5, -2.5, 0.25], - orientation=[0.0, 0.0, 0.0, 1.0], + header=Header(stamp=time_from_seconds(42.5), frame_id=""), + pose=Pose( + position=Point(x=1.5, y=-2.5, z=0.25), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), ), ), ) @@ -700,7 +754,7 @@ def test_unsubscribe_failure_does_not_skip_other_cleanup_or_leak_into_new_sessio assert wait_until(lambda: len(clients) == 2) assert old_session.unsubs == {} - color.publish(Image.from_numpy(np.zeros((8, 12, 3), dtype=np.uint8))) + color.publish(image_from_array(np.zeros((8, 12, 3), dtype=np.uint8), encoding="rgb8")) flush_loop(module) assert clients[1].writers["color_image"].offers == [] @@ -767,16 +821,18 @@ def test_start_with_authored_manifest_rates_and_quality(monkeypatch) -> None: (image_spec,) = module._channel_specs assert image_spec.params["quality"] == 33 qualities: list[int] = [] - real = Image.to_jpeg_bytes + real = builtin_codecs.image_to_jpeg def spy(self: Image, quality: int = 75) -> bytes: qualities.append(quality) return real(self, quality=quality) - monkeypatch.setattr(Image, "to_jpeg_bytes", spy) + monkeypatch.setattr(builtin_codecs, "image_to_jpeg", spy) push(module, clients[0], Subs(chs=["color_image"], n=1)) assert wait_until(lambda: image_transport(module).subscribers) - image_transport(module).publish(Image.from_numpy(np.zeros((8, 12, 3), dtype=np.uint8))) + image_transport(module).publish( + image_from_array(np.zeros((8, 12, 3), dtype=np.uint8), encoding="rgb8") + ) assert wait_until(lambda: qualities == [33]) finally: stop_module(module) @@ -1339,7 +1395,9 @@ def test_teleop_twist_publishes_geometry_twist(teleop_bridge) -> None: module, clients, twists = teleop_bridge push(module, clients[0], wire_twist(0.4, 0.2, -0.5, seq=1)) assert wait_until(lambda: len(twists) == 1) - assert twists[0] == Twist(linear=Vector3(0.4, 0.2, 0.0), angular=Vector3(0.0, 0.0, -0.5)) + assert twists[0] == Twist( + linear=Vector3(x=0.4, y=0.2, z=0.0), angular=Vector3(z=-0.5, x=0.0, y=0.0) + ) def test_teleop_clamps_to_boost_bounds(teleop_bridge) -> None: @@ -1347,7 +1405,9 @@ def test_teleop_clamps_to_boost_bounds(teleop_bridge) -> None: push(module, clients[0], wire_twist(100.0, -100.0, -100.0, seq=1)) assert wait_until(lambda: len(twists) == 1) # maxLinear 0.8 * boost 2.0; maxAngular 1.0 * boost 2.0. - assert twists[0] == Twist(linear=Vector3(1.6, -1.6, 0.0), angular=Vector3(0.0, 0.0, -2.0)) + assert twists[0] == Twist( + linear=Vector3(x=1.6, y=-1.6, z=0.0), angular=Vector3(z=-2.0, x=0.0, y=0.0) + ) def test_teleop_seq_guard_drops_stale_within_live_stream(teleop_bridge) -> None: @@ -1371,7 +1431,9 @@ def test_teleop_watchdog_deadline_and_high_water_survives_silence(teleop_bridge) # wait_until poll and thread scheduling). assert wait_until(lambda: len(twists) == 2) elapsed = time.monotonic() - started - assert twists[1].is_zero() + assert twists[1] == Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) deadline = _TELEOP_TEST_WATCHDOG_MS / 1000 + teleop._TELEOP_POLL_S + 0.5 assert elapsed < deadline, f"deadman zero took {elapsed:.3f}s (deadline {deadline:.3f}s)" time.sleep(3 * _TELEOP_TEST_WATCHDOG_MS / 1000) @@ -1401,7 +1463,9 @@ def test_teleop_zero_only_on_release_edge(teleop_bridge) -> None: assert wait_until(lambda: len(twists) == 2) settle(module) assert len(twists) == 2 - assert twists[1].is_zero() + assert twists[1] == Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) def test_teleop_stop_is_unconditional(teleop_bridge) -> None: @@ -1411,7 +1475,10 @@ def test_teleop_stop_is_unconditional(teleop_bridge) -> None: push(module, clients[0], WireStop(seq=1, ts=time.time(), gen=1)) push(module, clients[0], WireStop(seq=2, ts=time.time(), gen=1)) assert wait_until(lambda: len(twists) == 2) - assert all(t.is_zero() for t in twists) + assert all( + t == Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + for t in twists + ) def test_teleop_stop_blocks_stale_reordered_twist(teleop_bridge) -> None: @@ -1433,7 +1500,9 @@ def test_teleop_lease_end_zeroes_once_and_resets_seq(teleop_bridge) -> None: assert wait_until(lambda: len(twists) == 1) push(module, clients[0], WireTeleopStop(gen=1)) assert wait_until(lambda: len(twists) == 2) - assert twists[1].is_zero() + assert twists[1] == Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) push(module, clients[0], WireTeleopStop(gen=1)) # repeat: dead on the gen gate settle(module) assert len(twists) == 2 @@ -1448,7 +1517,11 @@ def test_teleop_session_drop_zeroes(teleop_bridge) -> None: push(module, clients[0], wire_twist(0.5, 0.0, 0.0, seq=1)) assert wait_until(lambda: len(twists) == 1) kill_session(module, clients[0]) - assert wait_until(lambda: len(twists) == 2 and twists[1].is_zero()) + assert wait_until( + lambda: len(twists) == 2 + and twists[1] + == Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) assert wait_until(lambda: len(clients) == 2) # supervisor reconnected settle(module) assert len(twists) == 2 @@ -1459,7 +1532,11 @@ def test_teleop_lease_end_blocks_stale_gen_twist_permanently(teleop_bridge) -> N push(module, clients[0], wire_twist(0.5, 0.0, 0.0, seq=50)) assert wait_until(lambda: len(twists) == 1) push(module, clients[0], WireTeleopStop(gen=1)) - assert wait_until(lambda: len(twists) == 2 and twists[1].is_zero()) + assert wait_until( + lambda: len(twists) == 2 + and twists[1] + == Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) # Past the watchdog window, delayed twists from the released lease must # stay dead regardless of seq: a reordered pre-stop command must never # restart the robot after a stop. @@ -1494,7 +1571,11 @@ def test_teleop_start_adopts_new_gen_and_zeroes_lost_stop(teleop_bridge) -> None # The lease changed hands but its teleop_stop datagram was lost: the # next grant's announcement stops the robot and voids the old lease. push(module, clients[0], WireTeleopStart(gen=2)) - assert wait_until(lambda: len(twists) == 2 and twists[1].is_zero()) + assert wait_until( + lambda: len(twists) == 2 + and twists[1] + == Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) push(module, clients[0], wire_twist(0.5, 0.0, 0.0, seq=51)) # old lease settle(module) assert len(twists) == 2 @@ -1515,7 +1596,11 @@ def test_teleop_estop_does_not_lower_high_water(teleop_bridge) -> None: # A stale reordered e-stop still zeroes (safe direction) but must not # lower the high-water and let the superseded twist 11 re-apply. push(module, clients[0], WireStop(seq=10, ts=time.time(), gen=1)) - assert wait_until(lambda: len(twists) == 2 and twists[1].is_zero()) + assert wait_until( + lambda: len(twists) == 2 + and twists[1] + == Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)) + ) push(module, clients[0], wire_twist(0.9, 0.0, 0.0, seq=11)) settle(module) assert len(twists) == 2 @@ -1669,7 +1754,9 @@ def blocking_encode(msg: OccupancyGrid) -> bytes: assert wait_until(lambda: len(twists) == 2) elapsed = time.monotonic() - started assert not release_encode.is_set() - assert twists[1].is_zero() + assert twists[1] == Twist( + linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0) + ) deadline = _TELEOP_TEST_WATCHDOG_MS / 1000 + teleop._TELEOP_POLL_S + 0.5 assert elapsed < deadline, f"deadman zero took {elapsed:.3f}s during a replay encode" push(module, clients[0], wire_twist(0.3, 0.0, 0.0, seq=2)) @@ -1739,14 +1826,22 @@ def encode(msg: OccupancyGrid) -> bytes: if calls["n"] == 1: replay_started.set() assert release_replay.wait(timeout=5.0) - return msg.grid.tobytes() + return np.asarray(msg.data, dtype=np.int8).tobytes() swap_costmap_encoder(module, encode) grid_b = OccupancyGrid( - grid=np.array([[100, 100, 100], [0, 0, 0]], dtype=np.int8), - resolution=0.05, - origin=Pose(-1.25, 2.5, 0.0), - ts=43.0, + header=Header(stamp=time_from_seconds(43.0), frame_id=""), + info=MapMetaData( + map_load_time=Time(sec=0, nanosec=0), + resolution=0.05, + width=3, + height=2, + origin=Pose( + position=Point(x=-1.25, y=2.5, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([100, 100, 100, 0, 0, 0], dtype=np.int8), ) writer = clients[0].writers["global_costmap"] try: @@ -1770,7 +1865,7 @@ def encode(msg: OccupancyGrid) -> bytes: flush_loop(module) # B only: A's encode finished on a moved cache and never reached the # mailbox, and B came from the cache, not from a live encode. - assert writer.offers == [(grid_b.grid.tobytes(), None)] + assert writer.offers == [(np.asarray(grid_b.data, dtype=np.int8).tobytes(), None)] assert calls["n"] == 2 assert module.encoded["global_costmap"] == 0 ts = writer.tss[0] diff --git a/dimos/web/relay_bridge/test_relay_bridge_rtc.py b/dimos/web/relay_bridge/test_relay_bridge_rtc.py index 9dc8f3cb8a7..8ed0e0029fa 100644 --- a/dimos/web/relay_bridge/test_relay_bridge_rtc.py +++ b/dimos/web/relay_bridge/test_relay_bridge_rtc.py @@ -26,7 +26,7 @@ import numpy as np import pytest -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_from_array from dimos.web.relay_bridge import relay_bridge_module from dimos.web.relay_bridge.e2e_support import stop_module from dimos.web.relay_bridge.manifest import TRACK_ENCODING @@ -102,7 +102,7 @@ def test_rtc_ice_offers_the_answer_connects_and_frames_feed_the_track(monkeypatc # never a JPEG writer. push(module, client, Subs(chs=["color_image"], n=1)) assert wait_until(lambda: bool(image_transport(module).subscribers)) - image = Image(data=np.zeros((4, 4, 3), np.uint8), format=ImageFormat.RGB) + image = image_from_array(np.zeros((4, 4, 3), np.uint8), encoding="rgb8") image_transport(module).publish(image) assert wait_until(lambda: publishers[0].fed == [("color_image", image)]) assert module.encoded["color_image"] == 1 @@ -157,7 +157,7 @@ def test_rtc_feed_gap_past_the_track_lifetime_is_reported_to_the_relay(monkeypat assert wait_until(lambda: publishers[0].accepted == "a") push(module, client, Subs(chs=["color_image"], n=1)) assert wait_until(lambda: bool(image_transport(module).subscribers)) - image = Image(data=np.zeros((4, 4, 3), np.uint8), format=ImageFormat.RGB) + image = image_from_array(np.zeros((4, 4, 3), np.uint8), encoding="rgb8") image_transport(module).publish(image) assert wait_until(lambda: len(publishers[0].fed) == 1) flush_loop(module) @@ -235,7 +235,7 @@ def test_rtc_peer_loss_after_connect_is_torn_down_and_reoffered(monkeypatch) -> fail_peer(module, publishers[0]) assert wait_until(lambda: publishers[0].closed == 1 and len(_offers(client)) == 2) # Frames fed while the peer is down are dropped, not counted. - image = Image(data=np.zeros((4, 4, 3), np.uint8), format=ImageFormat.RGB) + image = image_from_array(np.zeros((4, 4, 3), np.uint8), encoding="rgb8") image_transport(module).publish(image) flush_loop(module) assert publishers[1].fed == [] and module.encoded["color_image"] == 0 diff --git a/dimos/web/relay_bridge/test_rtc_publisher.py b/dimos/web/relay_bridge/test_rtc_publisher.py index 81508fdf872..aa8ab79dc5d 100644 --- a/dimos/web/relay_bridge/test_rtc_publisher.py +++ b/dimos/web/relay_bridge/test_rtc_publisher.py @@ -26,7 +26,7 @@ import numpy as np import pytest -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_from_array from dimos.web.relay_bridge.protocol import RtcTrack from dimos.web.relay_bridge.relay_bridge_module import RTC_AVAILABLE @@ -56,7 +56,7 @@ async def test_feed_reports_a_media_gap_past_the_track_lifetime() -> None: from dimos.web.relay_bridge.rtc_publisher import TRACK_GC_S, RtcPublisher publisher = RtcPublisher(["color_image"]) - image = Image(data=np.zeros((48, 64, 3), np.uint8), format=ImageFormat.RGB) + image = image_from_array(np.zeros((48, 64, 3), np.uint8), encoding="rgb8") try: await publisher.start([]) assert publisher.feed("color_image", image) is False # nothing before it @@ -79,7 +79,7 @@ async def test_frames_fed_from_a_thread_reach_the_peer() -> None: received: asyncio.Queue[object] = asyncio.Queue() peer.on("track", received.put_nowait) stop = threading.Event() - image = Image(data=np.full((48, 64, 3), 200, np.uint8), format=ImageFormat.RGB) + image = image_from_array(np.full((48, 64, 3), 200, np.uint8), encoding="rgb8") def producer() -> None: # The transport thread's cadence: a frame every 50 ms. diff --git a/dimos/web/relay_bridge/test_voxel_encoding.py b/dimos/web/relay_bridge/test_voxel_encoding.py index 36be154ca35..8de990d088f 100644 --- a/dimos/web/relay_bridge/test_voxel_encoding.py +++ b/dimos/web/relay_bridge/test_voxel_encoding.py @@ -25,10 +25,11 @@ from typing import Any import zlib +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.sensor_msgs.msg import PointCloud2 import numpy as np import pytest -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.web.codecs import resolve_encoder from dimos.web.relay_bridge import builtin_codecs from dimos.web.relay_bridge.builtin_codecs import _VOXEL_RECORD, encode_voxels @@ -66,9 +67,9 @@ def test_encoder_reproduces_golden_vector(vec: dict[str, Any]) -> None: def test_encoder_handles_wire_decoded_cloud() -> None: - """The bridge encodes clouds that arrived over LCM.""" + """The bridge encodes CDR clouds that arrived over a transport.""" msg = cloud_msg([[0.025, 0.025, 0.025], [1.025, -0.475, 0.125]]) - encoded = encode_voxels(PointCloud2.lcm_decode(msg.lcm_encode()), {"res": 0.05}) + encoded = encode_voxels(cdr_decode(cdr_encode(msg), PointCloud2), {"res": 0.05}) assert encoded.meta == {"res": 0.05, "n": 2, "chunks": 2} assert unpack(encoded.payload, encoded.meta) == [[0, 0, 0], [20, -10, 2]] @@ -79,7 +80,7 @@ def test_encoder_handles_wire_decoded_cloud() -> None: def test_empty_cloud_encodes_to_an_empty_frame() -> None: # An empty cloud is a frame, not a skip: a filtered cloud can empty out # and a producer can clear its map, and the panel must clear with it. - encoded = encode_voxels(PointCloud2(), {}) + encoded = encode_voxels(cloud_msg([]), {}) assert (encoded.payload, encoded.meta) == EMPTY_FRAME assert unpack(encoded.payload, encoded.meta) == [] nan = float("nan") diff --git a/dimos/web/test_cdr_codec.py b/dimos/web/test_cdr_codec.py index 04d10157ea1..540d5e81888 100644 --- a/dimos/web/test_cdr_codec.py +++ b/dimos/web/test_cdr_codec.py @@ -15,12 +15,17 @@ """Generated messages advertise the same schema to recordings and browsers.""" import base64 +from dataclasses import asdict import json import subprocess import sys -from dimos_generated.geometry_msgs.msg import Point, PoseStamped +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time +from dimos_generated.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from dimos.web.cdr_codec import check_cdr_params, default_encoding, encode_cdr_v1, export_schema @@ -42,10 +47,15 @@ def test_schema_and_declared_type() -> None: assert "MSG: std_msgs/Header" in schema["definition"] params = {"cdr": schema} check_cdr_params(params) - msg = PoseStamped() - assert encode_cdr_v1(msg, params) == msg.encode() + msg = PoseStamped( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + ) + assert encode_cdr_v1(msg, params) == cdr_encode(msg) with pytest.raises(ValueError, match="declared channel type"): - encode_cdr_v1(Point(), params) + encode_cdr_v1(Point(x=0.0, y=0.0, z=0.0), params) for bad in [{}, {"cdr": {}}, {"cdr": {"type": "p.T", "definition": ""}}]: with pytest.raises(ValueError, match="params"): check_cdr_params(bad) @@ -57,8 +67,12 @@ def test_browser_fixtures_match_generated_codecs() -> None: expected = json.loads((find_web_dir() / "shared/fixtures/cdr_frames.json").read_text()) assert {"vectors": build_vectors()} == expected for (_, msg), vector in zip(build_messages(), expected["vectors"], strict=True): - assert type(msg).decode(base64.b64decode(vector["payload_b64"])) == msg - assert type(msg).decode(base64.b64decode(vector["big_endian_b64"])) == msg + np.testing.assert_equal( + asdict(cdr_decode(base64.b64decode(vector["payload_b64"]), type(msg))), asdict(msg) + ) + np.testing.assert_equal( + asdict(cdr_decode(base64.b64decode(vector["big_endian_b64"]), type(msg))), asdict(msg) + ) def test_codec_import_is_light() -> None: diff --git a/dimos/web/test_cockpit.py b/dimos/web/test_cockpit.py index 409be85d4ef..365ade1f2e9 100644 --- a/dimos/web/test_cockpit.py +++ b/dimos/web/test_cockpit.py @@ -20,30 +20,21 @@ import subprocess import sys -from dimos_lcm import ( - geometry_msgs as lcm_geometry_msgs, - visualization_msgs as lcm_visualization_msgs, -) -from dimos_lcm.std_msgs import Bool +from dimos_cdrgen.registry import schema as cdr_schema +from dimos_generated.dimos_msgs.msg import MotorCommandArray +from dimos_generated.geometry_msgs.msg import PointStamped, PoseArray, PoseStamped, Transform, Twist +from dimos_generated.nav_msgs.msg import OccupancyGrid, Odometry, Path +from dimos_generated.sensor_msgs.msg import Image, Imu, JointState, PointCloud2 +from dimos_generated.std_msgs.msg import Bool +from dimos_generated.tf2_msgs.msg import TFMessage +from dimos_generated.trajectory_msgs.msg import JointTrajectory +from dimos_generated.visualization_msgs.msg import MarkerArray from langchain_core.messages import BaseMessage import pytest from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser from dimos.core.coordination.blueprints import autoconnect -from dimos.msgs.geometry_msgs.PointStamped import PointStamped -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.Imu import Imu -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.sensor_msgs.MotorCommandArray import MotorCommandArray -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.web.cdr_codec import encode_cdr_v1, export_schema from dimos.web.cockpit import ( Channel, ChannelRequest, @@ -59,7 +50,6 @@ cockpit, ) from dimos.web.codecs import EncodedPayload, decode_json_v1, encode_json_v1, web_encoder -from dimos.web.lcm_codec import encode_lcm_v1, export_schema from dimos.web.relay_bridge.audio_codec import AudioChunk, decode_audio_chunk from dimos.web.relay_bridge.builtin_codecs import ( decode_bool, @@ -849,10 +839,12 @@ def test_explicit_channel_conflicts_with_panel_raise(channel: Channel) -> None: def test_builtin_stream_type_and_table_mismatches() -> None: - with pytest.raises(ValueError, match="does not match the bridge port type PoseStamped"): + with pytest.raises( + ValueError, match=f"does not match the bridge port type {PoseStamped.__qualname__}" + ): cockpit(channels=[Channel("odom", Twist, encoding="pose.json.v1")]) with pytest.raises( - ValueError, match="'odom' encodes pose.json.v1, not geometry_msgs.PoseStamped.lcm.v1" + ValueError, match="'odom' encodes pose.json.v1, not geometry_msgs/msg/PoseStamped.cdr.v1" ): cockpit(channels=[Channel("odom", PoseStamped)]) with pytest.raises(ValueError, match="'odom' delivers reliable, not latest"): @@ -871,109 +863,108 @@ def test_json_v1_requires_explicit_codec_for_large_types() -> None: def test_image_has_no_default_encoding() -> None: # Raw pixels are never the default; the error points at the codec to use. with pytest.raises( - ValueError, match=r"sensor_msgs\.Image has no default web encoding.*jpeg\.v1" + ValueError, match=r"sensor_msgs/msg/Image has no default web encoding.*jpeg\.v1" ): Channel("snapshot", Image) def test_channel_default_encoding_by_type() -> None: - assert Channel("pose", PoseStamped).encoding == "geometry_msgs.PoseStamped.lcm.v1" - # Transform's wire format is a one-element TFMessage, and so is its id. - assert Channel("tf", Transform).encoding == "tf2_msgs.TFMessage.lcm.v1" + assert Channel("pose", PoseStamped).encoding == "geometry_msgs/msg/PoseStamped.cdr.v1" + # Transform and TFMessage are distinct generated messages. + assert Channel("tf", Transform).encoding == "geometry_msgs/msg/Transform.cdr.v1" assert Channel("note", dict).encoding == "json.v1" assert Channel("ops_note", _OpsNote).encoding == "json.v1" assert Channel("pose_in", PoseStamped, dir="tx", publish="shared").encoding == "json.v1" assert Channel("odom", PoseStamped, encoding="pose.json.v1").encoding == "pose.json.v1" -def test_lcm_channel_blueprint() -> None: +def test_cdr_channel_blueprint() -> None: blueprint = cockpit(channels=[Channel("pose", PoseStamped, max_hz=20.0)]) (atom,) = blueprint.blueprints manifest = atom.kwargs["manifest"] (channel,) = manifest["channels"] - schema = export_schema(lcm_geometry_msgs.PoseStamped) - assert channel["encoding"] == "geometry_msgs.PoseStamped.lcm.v1" - assert channel["params"] == {"lcm": schema} - assert schema["type"] == "geometry_msgs.PoseStamped" and schema["fp"] == "6a82696458c279a0" - assert "std_msgs.Header" in schema["structs"] + schema = export_schema(PoseStamped) + assert channel["encoding"] == "geometry_msgs/msg/PoseStamped.cdr.v1" + assert channel["params"] == {"cdr": schema} + assert schema["type"] == "geometry_msgs/msg/PoseStamped" + assert "MSG: std_msgs/Header" in schema["definition"] assert parse_manifest(manifest).model_dump() == manifest (spec,) = atom.kwargs["channels"] - assert spec.encoder is encode_lcm_v1 and spec.encoder_takes_params is True - assert spec.params["lcm"]["fp"] == "6a82696458c279a0" + assert spec.encoder is encode_cdr_v1 and spec.encoder_takes_params is True + assert spec.params["cdr"]["definition"] == cdr_schema(PoseStamped.__msgtype__) assert any( s.name == "pose" and s.type is PoseStamped and s.direction == "in" for s in atom.streams ) restored = pickle.loads(pickle.dumps(blueprint)) - assert restored.blueprints[0].kwargs["channels"][0].encoder is encode_lcm_v1 + assert restored.blueprints[0].kwargs["channels"][0].encoder is encode_cdr_v1 def test_generated_class_channel_uses_its_canonical_name() -> None: - # The generated class's own msg_name is bare ("MarkerArray"); the id and - # the schema use the package-qualified name, default or explicit. - markers = lcm_visualization_msgs.MarkerArray + # Default and explicit encodings both use the generated qualified name. + markers = MarkerArray blueprint = cockpit( channels=[ Channel("markers", markers), - Channel("markers_too", markers, encoding="visualization_msgs.MarkerArray.lcm.v1"), + Channel("markers_too", markers, encoding="visualization_msgs/msg/MarkerArray.cdr.v1"), ] ) (atom,) = blueprint.blueprints for wire in atom.kwargs["manifest"]["channels"]: - assert wire["encoding"] == "visualization_msgs.MarkerArray.lcm.v1" - assert wire["params"]["lcm"]["type"] == "visualization_msgs.MarkerArray" - assert all(spec.encoder is encode_lcm_v1 for spec in atom.kwargs["channels"]) + assert wire["encoding"] == "visualization_msgs/msg/MarkerArray.cdr.v1" + assert wire["params"]["cdr"]["type"] == "visualization_msgs/msg/MarkerArray" + assert all(spec.encoder is encode_cdr_v1 for spec in atom.kwargs["channels"]) -def test_lcm_schema_joins_user_params_in_the_request_only() -> None: +def test_cdr_schema_joins_user_params_in_the_request_only() -> None: channel = Channel("pose", PoseStamped, params={"note": "x"}) assert dict(channel.params) == {"note": "x"} (wire,) = cockpit(channels=[channel]).blueprints[0].kwargs["manifest"]["channels"] - assert wire["params"]["note"] == "x" and wire["params"]["lcm"]["fp"] == "6a82696458c279a0" - with pytest.raises(ValueError, match="'pose': params key 'lcm' is reserved"): - cockpit(channels=[Channel("pose", PoseStamped, params={"lcm": {}})]) + assert wire["params"]["note"] == "x" and wire["params"]["cdr"]["definition"] == cdr_schema( + PoseStamped.__msgtype__ + ) + with pytest.raises(ValueError, match="'pose': params key 'cdr' is reserved"): + cockpit(channels=[Channel("pose", PoseStamped, params={"cdr": {}})]) -@pytest.mark.parametrize( - ("channel", "match"), - [ - ( - Channel("pose", PoseStamped, encoding="nav_msgs.Odometry.lcm.v1"), - r"'pose': encoding 'nav_msgs\.Odometry\.lcm\.v1' encodes nav_msgs\.Odometry, " - r"not PoseStamped", - ), - ( - Channel("note", dict, encoding="geometry_msgs.PoseStamped.lcm.v1"), - r"'note': builtins\.dict is not a DimOS message", - ), - (Channel("traj", JointTrajectory), r"'traj': .*declares its own LCM fingerprint"), - (Channel("motors", MotorCommandArray), r"'motors': .*dimos_lcm.*@web_encoder"), - (Channel("poses", bytes), r"'poses': .*not supported by json\.v1"), - ], - ids=["id_type_mismatch", "lcm_id_on_dict", "foreign_fingerprint", "no_schema", "no_lcm"], -) -def test_lcm_encoding_errors(channel: Channel, match: str) -> None: - with pytest.raises(ValueError, match=match): - cockpit(channels=[channel]) +def test_cdr_encoding_errors() -> None: + with pytest.raises( + ValueError, match=f"encodes nav_msgs/msg/Odometry, not {PoseStamped.__qualname__}" + ): + cockpit(channels=[Channel("pose", PoseStamped, encoding="nav_msgs/msg/Odometry.cdr.v1")]) + with pytest.raises(ValueError, match="no generated CDR schema"): + cockpit(channels=[Channel("note", dict, encoding="geometry_msgs/msg/PoseStamped.cdr.v1")]) + + +def test_all_generated_types_have_default_cdr_encoders() -> None: + for message_type in (JointTrajectory, MotorCommandArray, PoseArray): + (atom,) = cockpit(channels=[Channel("message", message_type)]).blueprints + (channel,) = atom.kwargs["manifest"]["channels"] + assert channel["params"]["cdr"]["type"] == message_type.__msgtype__ + + +def test_untyped_bytes_have_no_default_json_encoder() -> None: + with pytest.raises(ValueError, match=r"'poses': .*not supported by json\.v1"): + cockpit(channels=[Channel("poses", bytes)]) -@web_encoder("t.ck.lcm.v1") -def _encode_note_lcm(msg: _OpsNote) -> bytes: +@web_encoder("t.ck.cdr.v1") +def _encode_note_cdr(msg: _OpsNote) -> bytes: return msg.text.encode() -def test_user_registered_lcm_v1_encoding_wins() -> None: - # An explicit @web_encoder with an *.lcm.v1 id is its own codec: no schema +def test_user_registered_cdr_v1_encoding_wins() -> None: + # An explicit @web_encoder with an *.cdr.v1 id is its own codec: no schema # injected, no reserved params key. - channel = Channel("ops_note", _OpsNote, encoding="t.ck.lcm.v1", params={"lcm": 1}) + channel = Channel("ops_note", _OpsNote, encoding="t.ck.cdr.v1", params={"cdr": 1}) (atom,) = cockpit(channels=[channel]).blueprints (wire,) = atom.kwargs["manifest"]["channels"] - assert wire["params"] == {"lcm": 1} + assert wire["params"] == {"cdr": 1} (spec,) = atom.kwargs["channels"] - assert spec.encoder is _encode_note_lcm and spec.encoder_takes_params is False + assert spec.encoder is _encode_note_cdr and spec.encoder_takes_params is False -def test_lcm_channels_fit_the_control_payload_cap() -> None: - # Every LCM channel adds its schema (150 B to 1.5 KB) to the hello frame. +def test_cdr_channels_fit_the_control_payload_cap() -> None: + # Every CDR channel adds its schema (150 B to 1.5 KB) to the hello frame. channels = [ Channel("odometry", Odometry), Channel("lidar", PointCloud2, max_hz=5.0), @@ -989,7 +980,7 @@ def test_lcm_channels_fit_the_control_payload_cap() -> None: manifest=cockpit(layout=GO2_LAYOUT, channels=channels).blueprints[0].kwargs["manifest"], ) size = len(encode_datagram(hello)) - assert size <= MAX_CONTROL_PAYLOAD_BYTES, f"hello with six LCM channels is {size} B" + assert size <= MAX_CONTROL_PAYLOAD_BYTES, f"hello with six CDR channels is {size} B" def test_unregistered_custom_encoding_names_the_decorator() -> None: diff --git a/dimos/web/test_codecs.py b/dimos/web/test_codecs.py index b71f36b25ea..1922689e2a8 100644 --- a/dimos/web/test_codecs.py +++ b/dimos/web/test_codecs.py @@ -321,7 +321,9 @@ def test_resolve_cdr_v1_family() -> None: definition = resolve_encoder("geometry_msgs/msg/PoseStamped.cdr.v1", PoseStamped) assert definition.encode is encode_cdr_v1 and definition.takes_params is True assert definition.check_params is check_cdr_params - with pytest.raises(ValueError, match="encodes nav_msgs/msg/Odometry, not PoseStamped"): + with pytest.raises( + ValueError, match=f"encodes nav_msgs/msg/Odometry, not {PoseStamped.__qualname__}" + ): resolve_encoder("nav_msgs/msg/Odometry.cdr.v1", PoseStamped) with pytest.raises(ValueError, match="encodes x, not dict"): resolve_encoder("x.cdr.v1", dict) diff --git a/dimos/web/websocket_vis/test_websocket_vis_module.py b/dimos/web/websocket_vis/test_websocket_vis_module.py index 4d7c29b35dc..2ca7567ef72 100644 --- a/dimos/web/websocket_vis/test_websocket_vis_module.py +++ b/dimos/web/websocket_vis/test_websocket_vis_module.py @@ -19,6 +19,8 @@ import math import zlib +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.geometry_msgs.msg import ( Point, Pose, @@ -28,6 +30,8 @@ TwistStamped, ) from dimos_generated.nav_msgs.msg import MapMetaData, OccupancyGrid, Path +from dimos_generated.std_msgs.msg import Header +import numpy as np import pytest from dimos.msgs.geometry import quaternion_from_euler @@ -47,10 +51,14 @@ def module(): def test_click_and_velocity_generate_cdr_messages(module): points, velocities, stamped = [], [], [] subscriptions = [ - module.goal_request.subscribe(lambda msg: points.append(PoseStamped.decode(msg.encode()))), - module.tele_cmd_vel.subscribe(lambda msg: velocities.append(Twist.decode(msg.encode()))), + module.goal_request.subscribe( + lambda msg: points.append(cdr_decode(cdr_encode(msg), PoseStamped)) + ), + module.tele_cmd_vel.subscribe( + lambda msg: velocities.append(cdr_decode(cdr_encode(msg), Twist)) + ), module.movecmd_stamped.subscribe( - lambda msg: stamped.append(TwistStamped.decode(msg.encode())) + lambda msg: stamped.append(cdr_decode(cdr_encode(msg), TwistStamped)) ), ] try: @@ -66,7 +74,9 @@ def test_click_and_velocity_generate_cdr_messages(module): ) ) assert points[0].header.frame_id == "world" - assert points[0].pose == Pose(position=Point(x=2.5, y=-1), orientation=Quaternion(w=1)) + assert points[0].pose == Pose( + position=Point(x=2.5, y=-1, z=0.0), orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ) assert velocities[0].linear.x == 0.5 assert velocities[0].angular.z == 0.8 assert stamped[0].twist == velocities[0] @@ -77,22 +87,30 @@ def test_click_and_velocity_generate_cdr_messages(module): def test_generated_pose_path_and_costmap_produce_browser_state(module): - pose = PoseStamped(pose=Pose(position=Point(x=2, y=3, z=4))) - module._on_robot_pose(PoseStamped.decode(pose.encode())) - module._on_path(Path(poses=[pose])) + pose = PoseStamped( + pose=Pose( + position=Point(x=2, y=3, z=4), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + module._on_robot_pose(cdr_decode(cdr_encode(pose), PoseStamped)) + module._on_path(Path(poses=[pose], header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""))) grid = OccupancyGrid( info=MapMetaData( width=2, height=2, resolution=1, origin=Pose( - position=Point(x=5, y=-2), orientation=quaternion_from_euler(0, 0, math.pi / 2) + position=Point(x=5, y=-2, z=0.0), + orientation=quaternion_from_euler(0, 0, math.pi / 2), ), + map_load_time=Time(sec=0, nanosec=0), ), - data=[100, 0, 0, -1], + data=np.array([100, 0, 0, -1], dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) - wire = grid.encode() - module._on_global_costmap(OccupancyGrid.decode(wire)) + wire = cdr_encode(grid) + module._on_global_costmap(cdr_decode(wire, OccupancyGrid)) assert module.vis_state["robot_pose"] == {"type": "vector", "c": [2, 3, 4]} assert module.vis_state["path"] == {"type": "path", "points": [[2, 3]]} costmap = module.vis_state["costmap"] @@ -101,4 +119,4 @@ def test_generated_pose_path_and_costmap_produce_browser_state(module): assert costmap["resolution"] == 1 assert costmap["grid"]["shape"] == [2, 2] assert list(zlib.decompress(base64.b64decode(costmap["grid"]["data"]))) == [100, 100, 100, 255] - assert grid.encode() == wire + assert cdr_encode(grid) == wire diff --git a/dimos/web/websocket_vis/websocket_vis_module.py b/dimos/web/websocket_vis/websocket_vis_module.py index ab659d6225b..a9e064162d4 100644 --- a/dimos/web/websocket_vis/websocket_vis_module.py +++ b/dimos/web/websocket_vis/websocket_vis_module.py @@ -24,11 +24,20 @@ import asyncio from pathlib import Path as FilePath import threading -import time from typing import TYPE_CHECKING, Any import webbrowser -from dimos_lcm.std_msgs import Bool +from dimos_generated.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Twist, + TwistStamped, + Vector3, +) +from dimos_generated.nav_msgs.msg import OccupancyGrid, Path +from dimos_generated.std_msgs.msg import Bool from reactivex.disposable import Disposable from dimos.utils.data import get_data @@ -48,12 +57,9 @@ from dimos.mapping.models import LatLon from dimos.mapping.occupancy.gradient import gradient from dimos.mapping.occupancy.inflation import simple_inflate -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.geometry import quaternion_euler +from dimos.msgs.occupancy import occupancy_view +from dimos.msgs.time import header_now from dimos.utils.logging_config import setup_logger from dimos.web.websocket_vis.optimized_costmap import OptimizedCostmapEncoder @@ -92,13 +98,13 @@ class WebsocketVisModule(Module): config: WebsocketConfig - # LCM inputs + # Typed inputs odom: In[PoseStamped] gps_location: In[LatLon] path: In[Path] global_costmap: In[OccupancyGrid] - # LCM outputs + # Typed outputs goal_request: Out[PoseStamped] gps_goal: Out[LatLon] explore_cmd: Out[Bool] @@ -296,13 +302,17 @@ async def connect(sid, environ) -> None: # type: ignore[no-untyped-def] @self.sio.event # type: ignore[untyped-decorator] async def click(sid, position) -> None: # type: ignore[no-untyped-def] goal = PoseStamped( - position=(position[0], position[1], 0), - orientation=(0, 0, 0, 1), # Default orientation - frame_id="world", + header=header_now("world"), + pose=Pose( + position=Point(x=position[0], y=position[1], z=0.0), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), ) self.goal_request.publish(goal) logger.info( - "Click goal published", x=round(goal.position.x, 3), y=round(goal.position.y, 3) + "Click goal published", + x=round(goal.pose.position.x, 3), + y=round(goal.pose.position.y, 3), ) @self.sio.event # type: ignore[untyped-decorator] @@ -343,30 +353,18 @@ async def clear_gps_goals(sid: str) -> None: @self.sio.event # type: ignore[untyped-decorator] async def move_command(sid: str, data: dict[str, Any]) -> None: - # Publish Twist if transport is configured - if self.tele_cmd_vel and self.tele_cmd_vel.transport: - twist = Twist( - linear=Vector3(data["linear"]["x"], data["linear"]["y"], data["linear"]["z"]), - angular=Vector3( - data["angular"]["x"], data["angular"]["y"], data["angular"]["z"] - ), - ) - self.tele_cmd_vel.publish(twist) - - # Publish TwistStamped if transport is configured - if self.movecmd_stamped and self.movecmd_stamped.transport: - twist_stamped = TwistStamped( - ts=time.time(), - frame_id="base_link", - linear=Vector3(data["linear"]["x"], data["linear"]["y"], data["linear"]["z"]), - angular=Vector3( - data["angular"]["x"], data["angular"]["y"], data["angular"]["z"] - ), - ) - self.movecmd_stamped.publish(twist_stamped) + twist = Twist( + linear=Vector3(**data["linear"]), + angular=Vector3(**data["angular"]), + ) + self.tele_cmd_vel.publish(twist) + self.movecmd_stamped.publish(TwistStamped(header=header_now("base_link"), twist=twist)) def _on_robot_pose(self, msg: PoseStamped) -> None: - pose_data = {"type": "vector", "c": [msg.position.x, msg.position.y, msg.position.z]} + pose_data = { + "type": "vector", + "c": [msg.pose.position.x, msg.pose.position.y, msg.pose.position.z], + } self.vis_state["robot_pose"] = pose_data self._emit("robot_pose", pose_data) @@ -376,7 +374,7 @@ def _on_gps_location(self, msg: LatLon) -> None: self._emit("gps_location", pose_data) def _on_path(self, msg: Path) -> None: - points = [[pose.position.x, pose.position.y] for pose in msg.poses] + points = [[pose.pose.position.x, pose.pose.position.y] for pose in msg.poses] path_data = {"type": "path", "points": points} self.vis_state["path"] = path_data self._emit("path", path_data) @@ -389,17 +387,17 @@ def _on_global_costmap(self, msg: OccupancyGrid) -> None: def _process_costmap(self, costmap: OccupancyGrid) -> dict[str, Any]: """Convert OccupancyGrid to visualization format.""" costmap = gradient(simple_inflate(costmap, 0.1), max_distance=1.0) - grid_data = self.costmap_encoder.encode_costmap(costmap.grid) + grid_data = self.costmap_encoder.encode_costmap(occupancy_view(costmap)) return { "type": "costmap", "grid": grid_data, "origin": { "type": "vector", - "c": [costmap.origin.position.x, costmap.origin.position.y, 0], + "c": [costmap.info.origin.position.x, costmap.info.origin.position.y, 0], }, - "resolution": costmap.resolution, - "origin_theta": 0, # Assuming no rotation for now + "resolution": costmap.info.resolution, + "origin_theta": quaternion_euler(costmap.info.origin.orientation)[2], } def _emit(self, event: str, data: Any) -> None: diff --git a/docs/development/message-source-packages.md b/docs/development/message-source-packages.md index 7b26bd085fd..1845d368d10 100644 --- a/docs/development/message-source-packages.md +++ b/docs/development/message-source-packages.md @@ -15,7 +15,7 @@ from dimos_cdrgen.source_package import setup_arguments ``` There is no DimOS compatibility distribution. Existing consumers of -`dimos_message_build` or generic `dimos.message_codegen` APIs must change their +`dimos_message_project` or generic `dimos.message_codegen` APIs must change their imports and dependency declarations to `dimos_cdrgen` and `dimos-cdrgen`. DimOS retains two adapters: diff --git a/docs/development/messages-in-repository.md b/docs/development/messages-in-repository.md index ddf01a34a0d..0aa28ae9c79 100644 --- a/docs/development/messages-in-repository.md +++ b/docs/development/messages-in-repository.md @@ -2,7 +2,7 @@ Use this tutorial when changing dimOS's own message definitions. Work in a checkout of the CDR proposal; this is not a main-branch or published-release API. -The separate `packages/dimos-generated` package owns checked-in generated Python source. +The separate `packages/dimos-generated` package owns schema inputs and generated Python source. The root dimOS package consumes it and does not generate it in `setup.py`. Prepare Python 3.12 and uv for installation. Message maintainers additionally @@ -15,7 +15,7 @@ uv sync --group tests --group message-codegen --frozen ``` `uv sync` installs the built-in message package as editable Python source from -`packages/dimos-generated/src`; no message-native extension is built. It consumes +`packages/dimos-generated/build/message-source/python`; no message-native extension is built. It consumes the local packages in `[tool.uv.sources]`. **Plain `pip install .` at the dimOS root does not read those uv sources.** It needs matching `dimos-generated` and `dimos-message-project` diff --git a/docs/usage/lcm.md b/docs/usage/lcm.md index 3739723f46e..398bba1b986 100644 --- a/docs/usage/lcm.md +++ b/docs/usage/lcm.md @@ -15,7 +15,7 @@ local-message user story and the three language examples. ```python session=cdr_transport_demo ansi=false from dimos_generated.geometry_msgs.msg import Vector3 -from dimos_message_build.registry import encode, decode +from dimos_cdrgen.registry import encode, decode message = Vector3(x=1.0, y=2.0, z=3.0) payload = encode(message) diff --git a/docs/usage/transports/index.md b/docs/usage/transports/index.md index bb9aecf895a..039fe333e47 100644 --- a/docs/usage/transports/index.md +++ b/docs/usage/transports/index.md @@ -5,7 +5,7 @@ Transports connect **module streams** across **process boundaries** and/or **net * **Module**: a running component (e.g., camera, mapping, nav). * **Stream**: a unidirectional flow of messages owned by a module (one broadcaster → many receivers). * **Topic**: the name/identifier used by a transport or pubsub backend. -* **Message**: payload carried on a stream (often `dimos.msgs.*`, but can be bytes / images / pointclouds / etc.). +* **Message**: payload carried on a stream (often `dimos_generated.*`, but can be bytes / images / pointclouds / etc.). Each edge in the graph is a **transported stream** (potentially different protocols). Each node is a **module**: @@ -29,7 +29,7 @@ So: treat the API as uniform, but pick a backend whose semantics match the task. ## Choosing a backend -* `zenoh`: is our default. Reliable delivery and the same typed message model through `LCM` binary encoding of messages +* `zenoh`: is our default. Reliable delivery and the same typed message model through CDR binary encoding of generated ROS messages * `lcm`: the legacy path, opt-in. Fast and simple, but UDP multicast is best-effort. At the CLI level, you can select the stream transport globally with: @@ -180,6 +180,7 @@ from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out from dimos.core.transport import LCMTransport from dimos_generated.sensor_msgs.msg import Image +from dimos_generated.builtin_interfaces.msg import Time from dimos_generated.std_msgs.msg import Header from dimos.msgs.image import image_from_array, image_view @@ -200,7 +201,7 @@ class TickerCameraModule(Module): img = image_from_array( np.zeros((480, 640, 3), dtype=np.uint8), encoding="rgb8", - header=Header(frame_id="synthetic"), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id="synthetic"), ) self.color_image.publish(img) @@ -252,6 +253,26 @@ Received: (480, 640, 3) 12:36:08.388 [inf][/coordination/python_worker.py] Worker module stopped. module=TickerCameraModule module_id=0 worker_id=0 12:36:08.394 [inf][ation/worker_manager_python.py] All workers shut down ``` +```results + @abstractmethod + def publish(self, topic: TopicT, message: MsgT) -> None: + """Publish a message to a topic.""" + ... + + @abstractmethod + def subscribe( + self, topic: TopicT, callback: Callable[[MsgT, TopicT], None] + ) -> Callable[[], None]: + """Subscribe to a topic with a callback. returns unsubscribe function + + The unsubscribe function must not block waiting for an in-flight + callback (callers may hold an event loop the backend needs for + progress), must be callable from within the callback itself, and once + it returns no further deliveries start (a callback already executing + may still complete). + """ + ... +``` See [Modules](/docs/usage/modules.md) for more on module architecture. @@ -269,11 +290,11 @@ dimos lcmspy # deprecated alias for: dimos spy --transport lcm ![dimos spy](../assets/lcmspy.png) -`dimos topic echo /topic` listens on typed channels like `/topic#pkg.Msg` and decodes automatically: +`dimos topic echo /topic` listens on typed channels like `/topic#pkg/msg/Msg` and decodes automatically: ```sh skip -Listening on /camera/rgb (inferring from typed LCM channels like '/camera/rgb#pkg.Msg')... (Ctrl+C to stop) -Image(shape=(480, 640, 3), format=RGB, dtype=uint8, dev=cpu, ts=2026-01-24 20:28:59) +Listening on /camera/rgb (inferring from typed LCM channels like '/camera/rgb#pkg/msg/Msg')... (Ctrl+C to stop) +Image(header=Header(...), height=480, width=640, encoding='rgb8', ...) ``` ## Implementing a transport @@ -290,11 +311,11 @@ Your `Transport.__init__` args can be anything meaningful for your backend: * a filesystem path * a Redis channel -Encoding is an implementation detail, but we encourage using LCM-compatible message types when possible. +Encoding is an implementation detail, but we encourage using generated ROS message types with CDR serialization when possible. ### Encoding helpers -Many of our message types provide `lcm_encode` / `lcm_decode` for compact, language-agnostic binary encoding (often faster than pickle). For details, see [LCM](/docs/usage/lcm.md). +Use `encode(message)` and `decode(payload, message_type.__msgtype__)` from `dimos_cdrgen.registry` for CDR serialization of generated message dataclasses. The same serialization works over LCM, Zenoh, and typed shared memory. ## PubSub transports @@ -312,27 +333,10 @@ print(inspect.getsource(PubSub.subscribe)) ``` ```results - @abstractmethod - def publish(self, topic: TopicT, message: MsgT) -> None: - """Publish a message to a topic.""" - ... - - @abstractmethod - def subscribe( - self, topic: TopicT, callback: Callable[[MsgT, TopicT], None] - ) -> Callable[[], None]: - """Subscribe to a topic with a callback. returns unsubscribe function - - The unsubscribe function must not block waiting for an in-flight - callback (callers may hold an event loop the backend needs for - progress), must be callable from within the callback itself, and once - it returns no further deliveries start (a callback already executing - may still complete). - """ - ... +Received velocity: x=1.0, y=0.0, z=0.5 ``` -Topic/message types are flexible: bytes, JSON, or our ROS-compatible [LCM](/docs/usage/lcm.md) types. We also have pickle-based transports for arbitrary Python objects. +Topic/message types are flexible: bytes, JSON, or our generated ROS message types. We also have pickle-based transports for arbitrary Python objects. ### LCM (UDP multicast) @@ -350,7 +354,7 @@ received = [] topic = Topic("/robot/velocity", Vector3) lcm.subscribe(topic, lambda msg, t: received.append(msg)) -lcm.publish(topic, Vector3(x=1.0, z=0.5)) +lcm.publish(topic, Vector3(x=1.0, y=0.0, z=0.5)) import time time.sleep(0.1) @@ -360,13 +364,39 @@ lcm.stop() ``` ```results -Received velocity: x=1.0, y=0.0, z=0.5 +Received: [{'data': [1, 2, 3]}] ``` ### Zenoh See [Zenoh](#zenoh) for backend selection, connection settings, and per-topic QoS. +### Shared memory (IPC) + +Shared memory only works on the **same machine**. Use `SHMTransport(topic, Image)` or a `shm:topic#sensor_msgs/msg/Image` URI for typed CDR messages. An untyped `shm:topic` carries raw bytes; `pshm:topic` carries pickled Python objects. + +```python +from dimos.protocol.pubsub.impl.shmpubsub import PickleSharedMemory + +shm = PickleSharedMemory(prefer="cpu") +shm.start() + +received = [] +unsubscribe = shm.subscribe("test/topic", lambda msg, topic: received.append(msg)) +shm.publish("test/topic", {"data": [1, 2, 3]}) + +import time +time.sleep(0.1) + +print(f"Received: {received}") +unsubscribe() +shm.stop() +``` + +```results +Received: [{'data': [1, 2, 3]}] +``` + ### DDS Transport For network communication, DDS uses the Data Distribution Service (DDS) protocol: @@ -398,7 +428,9 @@ dds.stop() ``` ```results -Received: [SensorReading(value=22.5)] +Received 2 messages: + {'temperature': 22.5} + {'temperature': 23.0} ``` ## A minimal transport: `Memory` @@ -441,10 +473,10 @@ Transports often need to serialize messages before sending and deserialize after | Mixin | Encoding | Use case | |----------------------|-----------------|------------------------------------| | `PickleEncoderMixin` | Python pickle | Any Python object, Python-only | -| `LCMEncoderMixin` | LCM binary | Cross-language (C/C++/Python/Go/…) | +| `CDREncoderMixin` | CDR binary | Cross-language (C/C++/Python/Go/…) | | `JpegEncoderMixin` | JPEG compressed | Image data, reduces bandwidth | -`LCMEncoderMixin` is especially useful: you can use LCM message definitions with *any* transport (not just UDP multicast). See [LCM](/docs/usage/lcm.md) for details. +`CDREncoderMixin` serializes generated ROS message dataclasses over any compatible pubsub backend. ### Creating a custom mixin diff --git a/docs/web/bridge.md b/docs/web/bridge.md index b8f96aa91b7..73623debe0c 100644 --- a/docs/web/bridge.md +++ b/docs/web/bridge.md @@ -98,7 +98,7 @@ arrow from B.w - (0, 0.3in) to M.e - (0, 0.3in) "publishes" below | `encoding` | A codec id. `None` picks the default (see [Encodings](#encodings)). | | `delivery` | `reliable` or `latest`. | | `max_hz` | The rate cap. For a tx channel it is the publish rate the relay accepts, per viewer and per robot. | -| `params` | Extra JSON for the encoder and the browser, shipped in the manifest. `params["lcm"]` is reserved for the LCM schema. | +| `params` | Extra JSON for the encoder and the browser, shipped in the manifest. `params["cdr"]` is reserved for the ROS2 schema. | | `publish` | `none` (the default) or `shared`. A tx channel must say `shared`: any viewer may publish on it. `exclusive` is not implemented. | | `required_scope` | Carried in the manifest for a relay with scopes. Nothing checks it today. | | `paced` | rx only. Space sends instead of sampling. | @@ -130,12 +130,12 @@ Behind a relay with Cloudflare video ([Relay hosting](/docs/web/relay_hosting.md Two more need no registration: - `json.v1`: JSON scalars, lists, dicts and plain dataclasses on rx. On tx, scalars, lists and dicts only (a dataclass built from untrusted browser JSON needs an explicit decoder). -- `..lcm.v1`, for example `geometry_msgs.PoseStamped.lcm.v1`: any dimOS message with an LCM schema, rx only. The frame is the message's `lcm_encode()` bytes and the manifest carries the schema, so the browser decodes it into a plain object with no registration. A bulk message costs its full size per frame (a `PointCloud2` is 16 bytes per point), so set `max_hz` accordingly or write an encoder that sends less. `voxels.zlib.v1`, the Map3D encoding, is one: it sends a cloud as voxel occupancy bits, about 0.3 bytes per voxel. +- `/msg/.cdr.v1`, for example `geometry_msgs/msg/PoseStamped.cdr.v1`: any dimOS message with an ROS2 schema, rx only. The frame is the message's `registry.encode(message)` bytes and the manifest carries the schema, so the browser decodes it into a plain object with no registration. A bulk message costs its full size per frame (a `PointCloud2` is 16 bytes per point), so set `max_hz` accordingly or write an encoder that sends less. `voxels.zlib.v1`, the Map3D encoding, is one: it sends a cloud as voxel occupancy bits, about 0.3 bytes per voxel. -When `encoding` is not given, an rx dimOS message gets its LCM encoding and everything else (every tx channel included) gets `json.v1`. `Image` has no default: use `jpeg.v1` or an encoder of your own. The built-in names keep their codecs: `Channel("odom", PoseStamped)` in `cockpit(channels=[...])` raises, because `odom` is `pose.json.v1`. Any other name takes the default: +When `encoding` is not given, an rx dimOS message gets its CDR encoding and everything else (every tx channel included) gets `json.v1`. `Image` has no default: use `jpeg.v1` or an encoder of your own. The built-in names keep their codecs: `Channel("odom", PoseStamped)` in `cockpit(channels=[...])` raises, because `odom` is `pose.json.v1`. Any other name takes the default: ```python -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos_generated.geometry_msgs.msg import PoseStamped from dimos.web.cockpit import Channel print(Channel("pose", PoseStamped).encoding) @@ -143,12 +143,12 @@ print(Channel("health", dict).encoding) ``` ```results -geometry_msgs.PoseStamped.lcm.v1 +geometry_msgs/msg/PoseStamped.cdr.v1 json.v1 ``` ```python -from dimos.msgs.sensor_msgs.Image import Image +from dimos_generated.sensor_msgs.msg import Image from dimos.web.cockpit import Channel try: @@ -168,14 +168,14 @@ sensor_msgs.Image has no default web encoding (raw pixel buffers; use encoding=' ```python skip import struct -from dimos.msgs.nav_msgs.Path import Path +from dimos_generated.nav_msgs.msg import Path from dimos.web.cockpit import Channel, cockpit from dimos.web.codecs import EncodedPayload, web_encoder @web_encoder("path.points.v1") def encode_path_points(msg: Path) -> EncodedPayload: - payload = b"".join(struct.pack(" +#include // Copyright 2026 Dimensional Inc. // SPDX-License-Identifier: Apache-2.0 #include "dimos/native.hpp" -#include + using namespace dimos::native; using dimos_msgs::msg::LineSegments3D; diff --git a/examples/native-modules/cpp/src/ping.cpp b/examples/native-modules/cpp/src/ping.cpp index 7647b29b056..3643b6a76a5 100644 --- a/examples/native-modules/cpp/src/ping.cpp +++ b/examples/native-modules/cpp/src/ping.cpp @@ -1,3 +1,4 @@ +#include // Copyright 2026 Dimensional Inc. // SPDX-License-Identifier: Apache-2.0 @@ -6,7 +7,7 @@ #include #include "dimos/native.hpp" -#include + using dimos::native::Builder; using dimos::native::Config; diff --git a/examples/native-modules/cpp/src/pong.cpp b/examples/native-modules/cpp/src/pong.cpp index 212caaf088d..055e1b7dd10 100644 --- a/examples/native-modules/cpp/src/pong.cpp +++ b/examples/native-modules/cpp/src/pong.cpp @@ -1,3 +1,4 @@ +#include // Copyright 2026 Dimensional Inc. // SPDX-License-Identifier: Apache-2.0 @@ -5,7 +6,7 @@ #include #include "dimos/native.hpp" -#include + using dimos::native::Builder; using dimos::native::Config; diff --git a/examples/native-modules/rust/Cargo.toml b/examples/native-modules/rust/Cargo.toml index f70865ec9c1..e68552709ea 100644 --- a/examples/native-modules/rust/Cargo.toml +++ b/examples/native-modules/rust/Cargo.toml @@ -25,7 +25,7 @@ path = "src/cdr_relay.rs" [dependencies] dimos-module = { path = "../../../native/rust/dimos-module" } -dimos-generated-messages = { path = "../../../dimos/message_codegen" } +dimos-generated-messages = { path = "../../../packages/dimos-generated/build/message-source/rust" } tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] } serde = { version = "1", features = ["derive"] } tracing = "0.1" diff --git a/examples/native-modules/rust/src/cdr_relay.rs b/examples/native-modules/rust/src/cdr_relay.rs index 37164709bfc..53ade5ee62e 100644 --- a/examples/native-modules/rust/src/cdr_relay.rs +++ b/examples/native-modules/rust/src/cdr_relay.rs @@ -15,7 +15,9 @@ // Copyright 2026 Dimensional Inc. // SPDX-License-Identifier: Apache-2.0 -use dimos_generated_messages::{dimos_msgs::msg::LineSegments3D, sensor_msgs::msg::Image}; +use dimos_generated_messages::{ + dimos_msgs::msg::line_segments3_d::LineSegments3D, sensor_msgs::msg::image::Image, +}; use dimos_module::{cdr, run_with_transport, Input, Module, Output}; #[derive(Module)] diff --git a/examples/native-modules/rust/src/ping.rs b/examples/native-modules/rust/src/ping.rs index 25bd35e7d57..082d20b4367 100644 --- a/examples/native-modules/rust/src/ping.rs +++ b/examples/native-modules/rust/src/ping.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use dimos_generated_messages::geometry_msgs::msg::{Twist, Vector3}; +use dimos_generated_messages::geometry_msgs::msg::{twist::Twist, vector3::Vector3}; use dimos_module::cdr; use dimos_module::{run_with_transport, Input, Module, Output}; use tokio::time::{interval, Duration}; diff --git a/examples/native-modules/rust/src/pong.rs b/examples/native-modules/rust/src/pong.rs index ac92bb4cb77..5f3c5b77507 100644 --- a/examples/native-modules/rust/src/pong.rs +++ b/examples/native-modules/rust/src/pong.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use dimos_generated_messages::geometry_msgs::msg::{Twist, Vector3}; +use dimos_generated_messages::geometry_msgs::msg::{twist::Twist, vector3::Vector3}; use dimos_module::cdr; use dimos_module::{native_config, run_with_transport, Input, Module, Output}; diff --git a/examples/native-modules/rust_ping_pong.py b/examples/native-modules/rust_ping_pong.py index fb4e94cdd3c..af2b7d47d58 100644 --- a/examples/native-modules/rust_ping_pong.py +++ b/examples/native-modules/rust_ping_pong.py @@ -26,11 +26,12 @@ import argparse +from dimos_generated.geometry_msgs.msg import Twist + from dimos.core.coordination.blueprints import Blueprint, autoconnect from dimos.core.coordination.module_coordinator import ModuleCoordinator from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.Twist import Twist from dimos.protocol.service.zenohservice import ZenohConfig # The crate is a workspace member, so cargo builds into the repo-root target dir. diff --git a/examples/native-modules/rust_tf.py b/examples/native-modules/rust_tf.py index c18406cd3f3..c0c234a947b 100644 --- a/examples/native-modules/rust_tf.py +++ b/examples/native-modules/rust_tf.py @@ -27,22 +27,28 @@ import math import time +from dimos_generated.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos_generated.std_msgs.msg import Header +from dimos_generated.tf2_msgs.msg import TFMessage + from dimos.core.coordination.blueprints import autoconnect from dimos.core.coordination.module_coordinator import ModuleCoordinator from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.stream import IO -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.time import time_from_nanoseconds +from dimos.protocol.tf.static_tf_publisher import StaticTfPublisher, frames_to_edge_transforms +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() # The crate is a workspace member, so cargo builds into the repo-root target dir. _BUILD = "cargo build --release" class TfProducer(Module): - """Publishes a time-varying a -> b -> c transform chain onto /tf.""" + """Publishes the time-varying a -> b edge onto /tf.""" tf: IO[TFMessage] @@ -58,21 +64,30 @@ async def _publish_loop(self) -> None: start = time.time() while self._running: t = time.time() - start - now = time.time() - self.tfbuffer.publish( - Transform( - translation=Vector3(0.0, math.cos(t), math.sin(t)), - frame_id="a", - child_frame_id="b", - ts=now, - ), - Transform( - translation=Vector3(1.0, 0.0, 0.0), - frame_id="b", - child_frame_id="c", - ts=now, - ), + stamp = time_from_nanoseconds(time.time_ns()) + self.tf.publish( + TFMessage( + transforms=[ + TransformStamped( + header=Header(frame_id="a", stamp=stamp), + child_frame_id="b", + transform=Transform( + translation=Vector3(y=math.cos(t), z=math.sin(t), x=0.0), + rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), + ), + ] + ) ) + composed = self.tfbuffer.get("a", "d", warn=False) + if composed is not None: + position = composed.transform.translation + logger.info( + "Python TF lookup", + x=position.x, + y=position.y, + z=position.z, + ) await asyncio.sleep(0.1) @rpc @@ -81,6 +96,13 @@ def stop(self) -> None: super().stop() +class StaticMount(StaticTfPublisher): + """Exercise the production periodic publisher for the fixed b -> c mount.""" + + def transforms(self) -> list[TransformStamped]: + return frames_to_edge_transforms([("c", "b", (1.0, 0.0, 0.0), (0.0, 0.0, 0.0))]) + + class TfListenerConfig(NativeModuleConfig): executable: str = "../../../target/release/tf_listener" build_command: str = _BUILD @@ -114,6 +136,9 @@ class TfBroadcasterModule(NativeModule): if __name__ == "__main__": bp = autoconnect( - TfProducer.blueprint(), TfBroadcasterModule.blueprint(), TfListenerModule.blueprint() + TfProducer.blueprint(), + StaticMount.blueprint(), + TfBroadcasterModule.blueprint(), + TfListenerModule.blueprint(), ).global_config(viewer="none") ModuleCoordinator.build(bp).loop() diff --git a/misc/optimize_patrol/optimize_candidates_child.py b/misc/optimize_patrol/optimize_candidates_child.py index c53f801c98a..ae23b4c33f3 100644 --- a/misc/optimize_patrol/optimize_candidates_child.py +++ b/misc/optimize_patrol/optimize_candidates_child.py @@ -29,7 +29,8 @@ from dimos.mapping.occupancy.path_resampling import smooth_resample_path from dimos.mapping.pointclouds.occupancy import height_cost_occupancy from dimos.mapping.pointclouds.util import read_pointcloud -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import header_now from dimos.navigation.experimental.patrolling.create_patrol_router import create_patrol_router from dimos.navigation.experimental.patrolling.routers.coverage_patrol_router import ( CoveragePatrolRouter, @@ -143,7 +144,7 @@ def main() -> None: args = parser.parse_args() data = read_pointcloud(get_data("big_office.ply")) - cloud = PointCloud2.from_numpy(np.asarray(data.points), frame_id="") + cloud = pointcloud_from_xyz(np.asarray(data.points), header=header_now()) occupancy_grid = height_cost_occupancy(cloud) costmap = gradient(occupancy_grid, max_distance=1.5) diff --git a/misc/optimize_patrol/optimize_patrol_router_child.py b/misc/optimize_patrol/optimize_patrol_router_child.py index 0d92f6824c1..1b600e5a6f1 100644 --- a/misc/optimize_patrol/optimize_patrol_router_child.py +++ b/misc/optimize_patrol/optimize_patrol_router_child.py @@ -23,7 +23,8 @@ from dimos.mapping.occupancy.path_resampling import smooth_resample_path from dimos.mapping.pointclouds.occupancy import height_cost_occupancy from dimos.mapping.pointclouds.util import read_pointcloud -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import header_now from dimos.navigation.experimental.patrolling.create_patrol_router import create_patrol_router from dimos.navigation.experimental.patrolling.routers.visitation_history import VisitationHistory from dimos.navigation.experimental.patrolling.utilities import point_to_pose_stamped @@ -127,7 +128,7 @@ def main() -> None: args = parser.parse_args() data = read_pointcloud(get_data("big_office.ply")) - cloud = PointCloud2.from_numpy(np.asarray(data.points), frame_id="") + cloud = pointcloud_from_xyz(np.asarray(data.points), header=header_now()) occupancy_grid = height_cost_occupancy(cloud) costmap = gradient(occupancy_grid, max_distance=1.5) diff --git a/native/cpp/dimos_nativeConfig.cmake b/native/cpp/dimos_nativeConfig.cmake new file mode 100644 index 00000000000..e0c6667716b --- /dev/null +++ b/native/cpp/dimos_nativeConfig.cmake @@ -0,0 +1,7 @@ +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 +# Source-package entry point: compilation happens in the consumer's build tree. +if(NOT TARGET dimos_native::dimos_native) + add_subdirectory("${CMAKE_CURRENT_LIST_DIR}" "${CMAKE_BINARY_DIR}/dimos_native" EXCLUDE_FROM_ALL) + add_library(dimos_native::dimos_native ALIAS dimos_native) +endif() diff --git a/native/cpp/include/dimos/native/module.hpp b/native/cpp/include/dimos/native/module.hpp index a39cfaef521..47c37c21910 100644 --- a/native/cpp/include/dimos/native/module.hpp +++ b/native/cpp/include/dimos/native/module.hpp @@ -277,14 +277,14 @@ class Builder { /// Handlers run serialized, so they touch module state without locks. template void input(const std::string& port, void (Self::*handler)(const T&), Self* self, - DecodeFn decode = cdr_decode) { + DecodeFn decode = [](const uint8_t* data, std::size_t size) { return cdr_decode(data, size); }) { input(port, std::move(decode), [self, handler](T msg) { (self->*handler)(msg); }); } /// publish() hands off to a per-channel worker, so it never blocks. template - Output output(const std::string& port, EncodeFn encode = cdr_encode) { + Output output(const std::string& port, EncodeFn encode = [](const T& message) { return cdr_encode(message); }) { std::string topic = topic_for(port); auto queue = std::make_shared(topic); publish_queues_.push_back(queue); diff --git a/native/cpp/messages.nix b/native/cpp/messages.nix index 5c5c1a9ee04..17685858c16 100644 --- a/native/cpp/messages.nix +++ b/native/cpp/messages.nix @@ -1,26 +1,71 @@ -# Built-in message package generated from the same pinned inputs as Cargo/Python. +# Build product schemas using the published generator and pinned source archives. +# All downloads happen before entering the build sandbox. { pkgs }: let - fastcdr = pkgs.stdenv.mkDerivation { - pname = "fastcdr"; - version = "2.4.0"; - src = pkgs.fetchurl { - url = "https://github.com/eProsima/Fast-CDR/archive/refs/tags/v2.4.0.tar.gz"; - sha256 = "79d8466107dd6b7d1defe961c4aa31735038937cf9dd1175cf6b0da0df2209ab"; - }; - nativeBuildInputs = [ pkgs.cmake ]; - cmakeFlags = [ "-DBUILD_TESTING=OFF" "-DBUILD_SHARED_LIBS=OFF" "-DCMAKE_POSITION_INDEPENDENT_CODE=ON" ]; + messageSource = import ../message-source.nix { inherit pkgs; }; + inherit (messageSource) generator python; + lock = builtins.fromJSON (builtins.readFile "${generator}/dimos_cdrgen/native_sources.json"); + archivePins = builtins.fromJSON (builtins.readFile ./source-archives.json); + archives = builtins.mapAttrs (name: source: + assert archivePins.${name}.revision == source.revision; + pkgs.fetchurl { + url = "https://codeload.github.com/${pkgs.lib.removePrefix "https://github.com/" (pkgs.lib.removeSuffix ".git" source.url)}/tar.gz/${source.revision}"; + sha256 = archivePins.${name}.sha256; + } + ) lock.repositories; + fastcdr = pkgs.fetchurl { + url = lock.fastcdr.url; + sha256 = lock.fastcdr.sha256; + }; + unpack = name: archive: '' + mkdir -p upstream/${name} + tar -xf ${archive} --strip-components=1 -C upstream/${name} + chmod -R u+w upstream/${name} + ''; + support = pkgs.stdenv.mkDerivation { + pname = "dimos-native-message-support"; + version = "0.1.0"; + dontUnpack = true; + nativeBuildInputs = [ pkgs.cmake python ]; + configurePhase = '' + runHook preConfigure + mkdir source + cp ${generator}/dimos_cdrgen/templates/native-support.cmake source/CMakeLists.txt + cp ${generator}/dimos_cdrgen/templates/native-toolchain.cmake source/native-toolchain.cmake + cp ${generator}/dimos_cdrgen/native_sources.json source/native_sources.json + ${pkgs.lib.concatStringsSep "\n" (pkgs.lib.mapAttrsToList unpack archives)} + ${unpack "fastcdr" fastcdr} + export PIP_NO_INDEX=1 PIP_DISABLE_PIP_VERSION_CHECK=1 + cmake -S source -B build \ + -DPython3_EXECUTABLE=${python}/bin/python3 \ + -DDIMOS_SUPPORT_PREFIX="$out" \ + -DFETCHCONTENT_FULLY_DISCONNECTED=ON \ + -DFETCHCONTENT_SOURCE_DIR_FASTCDR="$PWD/upstream/fastcdr" \ + ${pkgs.lib.concatStringsSep " " (pkgs.lib.mapAttrsToList (name: _: "-DFETCHCONTENT_SOURCE_DIR_${pkgs.lib.toUpper name}=\"$PWD/upstream/${name}\"") archives)} + runHook postConfigure + ''; + buildPhase = "cmake --build build --parallel $NIX_BUILD_CORES"; + # ExternalProject installs each prerequisite as part of its build. + installPhase = "test -d $out/share/rosidl_generator_cpp"; }; in pkgs.stdenv.mkDerivation { pname = "dimos-generated-messages"; version = "0.1.0"; - src = pkgs.lib.fileset.toSource { - root = ../..; - fileset = pkgs.lib.fileset.unions [ ../../dimos/__init__.py ../../dimos/message_codegen ]; - }; - nativeBuildInputs = [ pkgs.cmake pkgs.python3 ]; - propagatedBuildInputs = [ fastcdr ]; - postPatch = "python3 -m dimos.message_codegen.generate --output generated"; - cmakeDir = "../generated/cpp"; - cmakeFlags = [ "-DDIMOS_BUILD_PYTHON=OFF" ]; + src = "${messageSource.source}/cpp"; + nativeBuildInputs = [ pkgs.cmake python ]; + propagatedBuildInputs = [ support ]; + # Installed ament CMake exports require Python even for C++ consumers. + propagatedNativeBuildInputs = [ python ]; + setupHook = pkgs.writeText "dimos-message-python-hook" '' + addToSearchPath PYTHONPATH "${support}/${pkgs.python3.sitePackages}" + ''; + preConfigure = '' + export PYTHONPATH="${support}/${pkgs.python3.sitePackages}:$PYTHONPATH" + export AMENT_PREFIX_PATH="${support}:$out" + cmakeFlagsArray+=("-DCMAKE_PREFIX_PATH=${support};$out") + ''; + cmakeFlags = [ + "-DPython3_EXECUTABLE=${python}/bin/python3" + "-DDIMOS_RUNTIME_PATHS=${support}/lib" + ]; } diff --git a/native/cpp/source-archives.json b/native/cpp/source-archives.json new file mode 100644 index 00000000000..46c66801107 --- /dev/null +++ b/native/cpp/source-archives.json @@ -0,0 +1,38 @@ +{ + "ament_package": { + "revision": "ee3c5eda4fbc4df2c7a135d0ac1385e4a53fa0ee", + "sha256": "19dd626283a9f5dcfed5039ecd0ab149daa5e658a970bf607ed5bd053d54e629" + }, + "ament_cmake": { + "revision": "5741cff5b9f83253bf3521bd8f44108fde3504ad", + "sha256": "7172066ce00a4521205b74f70bfeb975345789ce1a69901e7a8da9dc45e22fcb" + }, + "ament_cmake_ros": { + "revision": "5b77a227b940a25be5224d78d37591c8baccd756", + "sha256": "d063e9cc5c68d12d33597f0ac5bb9568710848ad08741b019a5436eb52c43988" + }, + "rcutils": { + "revision": "26413034bab69688475f230ea7723e534ec3f8cd", + "sha256": "e71a0fd5fe8c42ae2f92ec1076d2fdf3e753034f9200fa7a4f14bc93f90e22b2" + }, + "rmw": { + "revision": "8cb94a87cf3f4049abe2d563af893517bf336930", + "sha256": "d0d43d3f19b125b1abc1e26c412a674cca4d20ab88bd9ba5496cf5a336d25f8e" + }, + "rosidl_dynamic_typesupport": { + "revision": "cb8c54d12c678daa70410bb0626ed8626f561e45", + "sha256": "c9d30f03eaad34657ce94ffe95faec445d46dccdb2372d66d06d382abc504a76" + }, + "ament_index": { + "revision": "2ccdcb8fd19a364fc793273b5793459207a440b9", + "sha256": "54ca269d31fc6b8f650ce262f5781f8805cc7c8750467f8dfaf6a16cbd9f8dbd" + }, + "rosidl": { + "revision": "85fa592b698b0f665e3120f48fac0d35e2f7d8a4", + "sha256": "ba5369bca27297f606bf3396a9e066140d07ce754d8659159148102b089e9510" + }, + "rosidl_typesupport_fastrtps": { + "revision": "b883555055cce17982a22172a42ef557ef40bac5", + "sha256": "4bb21020ab28a249890339301efb8e200af8330a6dee6e13ecd9388981cba83d" + } +} diff --git a/native/cpp/tests/CMakeLists.txt b/native/cpp/tests/CMakeLists.txt index 1a4c203eae5..26c77c11360 100644 --- a/native/cpp/tests/CMakeLists.txt +++ b/native/cpp/tests/CMakeLists.txt @@ -35,6 +35,6 @@ include(${doctest_SOURCE_DIR}/scripts/cmake/doctest.cmake) doctest_discover_tests(dimos_native_tests) # Exercise the production lidar message builder without sensor/PCL dependencies. -add_executable(dimos_lidar_cdr_demo ../../../examples/message-codegen/demo_lidar.cpp) -target_link_libraries(dimos_lidar_cdr_demo PRIVATE dimos_generated::messages) +add_executable(dimos_lidar_cdr_demo lidar_demo.cpp) +target_link_libraries(dimos_lidar_cdr_demo PRIVATE dimos_native) add_test(NAME lidar_generated_cdr COMMAND dimos_lidar_cdr_demo) diff --git a/native/cpp/tests/test_cdr_codec.cpp b/native/cpp/tests/test_cdr_codec.cpp index f3d487369a6..0d6d5d01990 100644 --- a/native/cpp/tests/test_cdr_codec.cpp +++ b/native/cpp/tests/test_cdr_codec.cpp @@ -1,6 +1,9 @@ // Copyright 2026 Dimensional Inc. // SPDX-License-Identifier: Apache-2.0 +#include +#include +#include #include #include "dimos/native/cdr_codec.hpp" @@ -18,7 +21,7 @@ TEST_CASE("default CDR codecs preserve nested generated fields") { CHECK(bytes[0] == 0); CHECK(bytes[1] == 1); CHECK(cdr_decode(bytes.data(), bytes.size()) == message); - const auto big = dimos::cdr::encode(message, false); + const auto big = dimos::native::cdr_encode(message, false); CHECK(cdr_decode(big.data(), big.size()) == message); } diff --git a/native/cpp/tests/test_module.cpp b/native/cpp/tests/test_module.cpp index a7df43bee41..9a4ba84a971 100644 --- a/native/cpp/tests/test_module.cpp +++ b/native/cpp/tests/test_module.cpp @@ -1,6 +1,7 @@ // Copyright 2026 Dimensional Inc. // SPDX-License-Identifier: Apache-2.0 +#include #include #include diff --git a/native/message-source.nix b/native/message-source.nix new file mode 100644 index 00000000000..a8ee88deebc --- /dev/null +++ b/native/message-source.nix @@ -0,0 +1,53 @@ +# Generate the product catalog from pinned inputs before native builds enter their sandbox. +{ pkgs }: +let + generatorWheel = pkgs.fetchurl { + url = "https://files.pythonhosted.org/packages/a3/9c/086701c3249ed2caabe83cea008c3e08bd6063c48a1df852fbca6aadff50/dimos_cdrgen-0.1.0-py3-none-any.whl"; + sha256 = "6731e8900f7e61af419b8c77599587f32677c9d64f27ac15f5e78be4b76bc011"; + }; + generator = pkgs.runCommand "dimos-cdrgen-0.1.0-source" { nativeBuildInputs = [ pkgs.unzip ]; } '' + mkdir -p $out + unzip -q ${generatorWheel} -d $out + ''; + python = pkgs.python3.withPackages (p: [ + p.pip p.setuptools p.wheel p.empy p.lark p.catkin-pkg p.pyyaml + p.rosbags p.numpy p.requests + ]); + product = pkgs.lib.fileset.toSource { + root = ../packages/dimos-generated; + fileset = pkgs.lib.fileset.unions [ + ../packages/dimos-generated/pyproject.toml + ../packages/dimos-generated/catalog.json + ../packages/dimos-generated/schema-sources.json + ../packages/dimos-generated/custom + ]; + }; + schemaLock = builtins.fromJSON (builtins.readFile ../packages/dimos-generated/schema-sources.json); + schemaArchives = map (source: { + name = builtins.hashString "sha256" source.url + ".tar.gz"; + archive = pkgs.fetchurl { inherit (source) url sha256; }; + }) schemaLock.sources; + source = pkgs.runCommand "dimos-generated-message-source" { nativeBuildInputs = [ python ]; } '' + cp -r ${product} product + chmod -R u+w product + mkdir -p cache/archives + ${pkgs.lib.concatMapStringsSep "\n" (item: "cp ${item.archive} cache/archives/${item.name}") schemaArchives} + export PYTHONPATH="${generator}:$PYTHONPATH" + python3 - <<'PYTHON' +from pathlib import Path +import shutil +from dimos_cdrgen.source_inputs import prepare_sources +from dimos_cdrgen.source_package import prepare +project = Path("product") +prepare_sources(project / "schema-sources.json", project / "interfaces", Path("cache/archives"), offline=True) +shutil.copytree(project / "custom/dimos_msgs/msg", project / "interfaces/dimos_msgs/msg") +shutil.copyfile(project / "custom/dimos_msgs/LICENSE", project / "interfaces/licenses/dimos_msgs-LICENSE") +prepare(project) +PYTHON + cp -r product/build/message-source $out + test -f $out/cpp/CMakeLists.txt + test -f $out/rust/Cargo.toml + ''; +in { + inherit generator python source; +} diff --git a/native/rust/dimos-lcm-transport/Cargo.toml b/native/rust/dimos-lcm-transport/Cargo.toml new file mode 100644 index 00000000000..f83c98802d0 --- /dev/null +++ b/native/rust/dimos-lcm-transport/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "dimos-lcm-transport" +version = "0.1.0" +edition = "2021" +description = "Schema-independent LCM UDP multicast transport for DimOS" +license = "Apache-2.0" + +[dependencies] +byteorder = "1" +socket2 = { version = "0.5", features = ["all"] } +tokio = { version = "1.51.1", features = ["net", "rt-multi-thread", "macros", "time"] } diff --git a/native/rust/dimos-lcm-transport/README.md b/native/rust/dimos-lcm-transport/README.md new file mode 100644 index 00000000000..c4860be61bb --- /dev/null +++ b/native/rust/dimos-lcm-transport/README.md @@ -0,0 +1,37 @@ +# DimOS raw LCM transport + +This crate publishes and receives opaque bytes over LCM UDP multicast. It has no +message definitions, fingerprints, generated message crate, or CDR dependency. +Typed consumers choose their own codec above this transport. + +The initial implementation was extracted from Dimensional's `dimos-lcm` repository, +`tools/rust/lcm`, commit `04d78e8622500244123ba9cefa4c51b4cb454549`, matching the +existing workspace lockfile. That crate declares Apache-2.0 in its Cargo manifest. +The source is maintained here so adding a message no longer depends on that +repository. Original transport behavior and its unit tests are retained; local +changes are reviewed in the DimOS history. + +Source: https://github.com/dimensionalOS/dimos-lcm/tree/04d78e8622500244123ba9cefa4c51b4cb454549/tools/rust/lcm + +Run the raw transport checks with: + +```bash +cargo test -p dimos-lcm-transport +cargo build -p dimos-lcm-transport --example interop +.venv/bin/python examples/message-codegen/demo_transport.py \ + --executable target/debug/examples/interop +``` + +The demo chooses a local UDP port, starts the native receiver, waits for its ready +signal, and exchanges both short and fragmented binary payloads with Python's LCM +binding. It terminates and reaps the native process even on failure. + +The receiver rejects malformed offsets, lengths, fragment numbers, conflicting +headers, overlaps, holes, and invalid channel UTF-8. Repeated fragments do not +advance completion; out-of-order fragments remain supported. Limits are 64 MiB +per message, 128 MiB of incomplete payload storage, and 64 incomplete messages +per transport instance. Fragment bookkeeping is also bounded by these limits and +the protocol's 16-bit fragment count. Incomplete messages expire after five +seconds and are purged when another fragment arrives. Publishing above 64 MiB +returns `InvalidInput` before any datagrams are sent. These are raw transport +limits and do not depend on the payload's encoding. diff --git a/native/rust/dimos-lcm-transport/examples/interop.rs b/native/rust/dimos-lcm-transport/examples/interop.rs new file mode 100644 index 00000000000..efc4c9e9576 --- /dev/null +++ b/native/rust/dimos-lcm-transport/examples/interop.rs @@ -0,0 +1,46 @@ +// 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. + +use dimos_lcm_transport::{Lcm, LcmOptions}; +use std::io::{self, Write}; +use std::time::Duration; + +#[tokio::main] +async fn main() -> Result<(), Box> { + let port = std::env::args() + .nth(1) + .ok_or("usage: interop ")? + .parse()?; + let transport = Lcm::with_options(LcmOptions { + port, + ttl: 0, + ..Default::default() + }) + .await?; + println!("READY"); + io::stdout().flush()?; + for _ in 0..3 { + let message = tokio::time::timeout(Duration::from_secs(10), async { + loop { + let message = transport.recv().await?; + if message.channel == "INPUT" { + return Ok::<_, io::Error>(message); + } + } + }) + .await??; + transport.publish("REPLY", &message.data).await?; + } + Ok(()) +} diff --git a/native/rust/dimos-lcm-transport/src/lib.rs b/native/rust/dimos-lcm-transport/src/lib.rs new file mode 100644 index 00000000000..1bbaef79375 --- /dev/null +++ b/native/rust/dimos-lcm-transport/src/lib.rs @@ -0,0 +1,41 @@ +// Copyright 2025-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. + +//! Pure Rust LCM (Lightweight Communications and Marshalling) transport. +//! +//! Provides UDP multicast publish/subscribe for LCM messages. +//! No system LCM library required. +//! +//! # Example +//! +//! ```no_run +//! use dimos_lcm_transport::Lcm; +//! +//! #[tokio::main] +//! async fn main() { +//! let mut lcm = Lcm::new().await.unwrap(); +//! +//! // Publish +//! let data = vec![1, 2, 3]; +//! lcm.publish("EXAMPLE", &data).await.unwrap(); +//! +//! // Receive +//! let msg = lcm.recv().await.unwrap(); +//! println!("{}: {} bytes", msg.channel, msg.data.len()); +//! } +//! ``` + +mod transport; + +pub use transport::{Lcm, LcmOptions, ReceivedMessage}; diff --git a/native/rust/dimos-lcm-transport/src/transport.rs b/native/rust/dimos-lcm-transport/src/transport.rs new file mode 100644 index 00000000000..54482f69524 --- /dev/null +++ b/native/rust/dimos-lcm-transport/src/transport.rs @@ -0,0 +1,661 @@ +// Copyright 2025-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. + +use byteorder::{BigEndian, ByteOrder}; +use socket2::{Domain, Protocol, Socket, Type}; +use std::collections::HashMap; +use std::io; +use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4}; +use std::sync::atomic::{AtomicU32, Ordering}; +use std::sync::Mutex; +use std::time::{Duration, Instant}; +use tokio::net::UdpSocket; + +const MAGIC_SHORT: u32 = 0x4c433032; // "LC02" +const MAGIC_LONG: u32 = 0x4c433033; // "LC03" +const SHORT_HEADER_SIZE: usize = 8; +const FRAGMENT_HEADER_SIZE: usize = 20; +const MAX_DATAGRAM_SIZE: usize = 65507; +// Bound incomplete UDP messages independently of sender-supplied lengths. +const MAX_MESSAGE_SIZE: usize = 64 * 1024 * 1024; +const MAX_REASSEMBLY_BYTES: usize = 128 * 1024 * 1024; +const MAX_INCOMPLETE_MESSAGES: usize = 64; +const REASSEMBLY_TIMEOUT: Duration = Duration::from_secs(5); + +/// Default LCM multicast group address. +pub const DEFAULT_MULTICAST_GROUP: Ipv4Addr = Ipv4Addr::new(239, 255, 76, 67); +/// Default LCM multicast port. +pub const DEFAULT_PORT: u16 = 7667; + +static SEQ: AtomicU32 = AtomicU32::new(0); + +struct FragmentBuffer { + channel: Option, + fragments: Vec>, + received: usize, + received_bytes: usize, + created: Instant, + data: Vec, +} + +/// Configuration for an LCM transport instance. +#[derive(Debug, Clone)] +pub struct LcmOptions { + /// Multicast group address (default: 239.255.76.67). + pub multicast_group: Ipv4Addr, + /// UDP port (default: 7667). + pub port: u16, + /// Multicast TTL (default: 1). + pub ttl: u32, + /// Network interface to bind to (default: any). + pub interface: Ipv4Addr, +} + +impl Default for LcmOptions { + fn default() -> Self { + Self { + multicast_group: DEFAULT_MULTICAST_GROUP, + port: DEFAULT_PORT, + ttl: 1, + interface: Ipv4Addr::UNSPECIFIED, + } + } +} + +/// A received LCM message. +#[derive(Debug, Clone)] +pub struct ReceivedMessage { + /// Channel name. + pub channel: String, + /// Encoded message payload. + pub data: Vec, +} + +/// Returns the first and subsequent payload sizes, and the number of fragments +fn fragment_params(msg_size: usize, channel_len: usize) -> (usize, usize, usize) { + let first_payload_size = MAX_DATAGRAM_SIZE - FRAGMENT_HEADER_SIZE - channel_len - 1; + let subsequent_payload_size = MAX_DATAGRAM_SIZE - FRAGMENT_HEADER_SIZE; + let num_fragments = if msg_size <= first_payload_size { + 1 + } else { + 1 + msg_size + .saturating_sub(first_payload_size) + .div_ceil(subsequent_payload_size) + }; + (first_payload_size, subsequent_payload_size, num_fragments) +} + +/// Pure Rust LCM UDP multicast transport. +pub struct Lcm { + socket: UdpSocket, + multicast_addr: SocketAddrV4, + reassembly: Mutex>, +} + +impl Lcm { + /// Create a new LCM transport with default options. + pub async fn new() -> io::Result { + Self::with_options(LcmOptions::default()).await + } + + /// Create a new LCM transport with custom options. + pub async fn with_options(opts: LcmOptions) -> io::Result { + let s2 = Socket::new(Domain::IPV4, Type::DGRAM, Some(Protocol::UDP))?; + s2.set_reuse_address(true)?; + // Sensor messages arrive as bursts of large fragments. Request enough + // queue space to reassemble an image without relying on kernel defaults. + s2.set_recv_buffer_size(4 * 1024 * 1024)?; + #[cfg(not(target_os = "windows"))] + s2.set_reuse_port(true)?; + + let bind_addr = SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, opts.port); + s2.bind(&bind_addr.into())?; + + let std_socket: std::net::UdpSocket = s2.into(); + std_socket.set_nonblocking(true)?; + let socket = UdpSocket::from_std(std_socket)?; + + socket.join_multicast_v4(opts.multicast_group, opts.interface)?; + socket.set_multicast_ttl_v4(opts.ttl)?; + + Ok(Self { + socket, + multicast_addr: SocketAddrV4::new(opts.multicast_group, opts.port), + reassembly: Mutex::new(HashMap::new()), + }) + } + + /// Publish encoded message data on the given channel. + pub async fn publish(&self, channel: &str, data: &[u8]) -> io::Result<()> { + let channel_bytes = channel.as_bytes(); + if channel_bytes.len() > 63 || channel_bytes.contains(&0) { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "LCM channel must fit in 63 bytes and contain no NUL", + )); + } + if data.len() > MAX_MESSAGE_SIZE { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "LCM message exceeds the 64 MiB transport limit", + )); + } + let total = SHORT_HEADER_SIZE + channel_bytes.len() + 1 + data.len(); + let seqno = SEQ.fetch_add(1, Ordering::Relaxed); + + if total > MAX_DATAGRAM_SIZE { + self.publish_fragmented(channel_bytes, data, seqno).await + } else { + self.publish_small(channel_bytes, data, seqno).await + } + } + + async fn publish_small(&self, channel_bytes: &[u8], data: &[u8], seqno: u32) -> io::Result<()> { + let total = SHORT_HEADER_SIZE + channel_bytes.len() + 1 + data.len(); + let mut buf = vec![0u8; total]; + + BigEndian::write_u32(&mut buf[0..4], MAGIC_SHORT); + BigEndian::write_u32(&mut buf[4..8], seqno); + + buf[SHORT_HEADER_SIZE..SHORT_HEADER_SIZE + channel_bytes.len()] + .copy_from_slice(channel_bytes); + // null terminator already 0 from vec![0u8; ..] + let payload_start = SHORT_HEADER_SIZE + channel_bytes.len() + 1; + buf[payload_start..].copy_from_slice(data); + + self.socket.send_to(&buf, self.multicast_addr).await?; + Ok(()) + } + + async fn publish_fragmented( + &self, + channel_bytes: &[u8], + data: &[u8], + seqno: u32, + ) -> io::Result<()> { + let msg_size = data.len(); + let (first_payload_size, subsequent_payload_size, num_fragments) = + fragment_params(msg_size, channel_bytes.len()); + + let mut payload_offset = 0; + + for fragment_no in 0..num_fragments { + let is_first = fragment_no == 0; + let channel_size = if is_first { channel_bytes.len() + 1 } else { 0 }; + let max_payload = if is_first { + first_payload_size + } else { + subsequent_payload_size + }; + let payload_len = (msg_size - payload_offset).min(max_payload); + + let datagram_size = FRAGMENT_HEADER_SIZE + channel_size + payload_len; + let mut buf = vec![0u8; datagram_size]; + + BigEndian::write_u32(&mut buf[0..4], MAGIC_LONG); + BigEndian::write_u32(&mut buf[4..8], seqno); + BigEndian::write_u32(&mut buf[8..12], msg_size as u32); + BigEndian::write_u32(&mut buf[12..16], payload_offset as u32); + BigEndian::write_u16(&mut buf[16..18], fragment_no as u16); + BigEndian::write_u16(&mut buf[18..20], num_fragments as u16); + + let mut offset = FRAGMENT_HEADER_SIZE; + + if is_first { + buf[offset..offset + channel_bytes.len()].copy_from_slice(channel_bytes); + // null terminator already 0 + offset += channel_bytes.len() + 1; + } + + buf[offset..offset + payload_len] + .copy_from_slice(&data[payload_offset..payload_offset + payload_len]); + + self.socket.send_to(&buf, self.multicast_addr).await?; + payload_offset += payload_len; + } + + Ok(()) + } + + /// Receive one LCM message asynchronously. + /// + /// Waits until a complete message arrives, reassembling fragments if necessary. + pub async fn recv(&self) -> io::Result { + let mut buf = vec![0u8; MAX_DATAGRAM_SIZE]; + loop { + let (n, sender) = self.socket.recv_from(&mut buf).await?; + let pkt = &buf[..n]; + + if pkt.len() < 4 { + continue; + } + let magic = BigEndian::read_u32(&pkt[0..4]); + + if magic == MAGIC_SHORT { + if let Some(msg) = Self::decode_small(pkt)? { + return Ok(msg); + } + } else if magic == MAGIC_LONG { + if let Some(msg) = self.process_fragment(sender, pkt)? { + return Ok(msg); + } + } + // Unknown magic or incomplete fragment — wait for the next datagram + } + } + + fn process_fragment( + &self, + sender: SocketAddr, + buf: &[u8], + ) -> io::Result> { + if buf.len() < FRAGMENT_HEADER_SIZE { + return Ok(None); + } + + let seqno = BigEndian::read_u32(&buf[4..8]); + let total_size = BigEndian::read_u32(&buf[8..12]) as usize; + let fragment_offset = BigEndian::read_u32(&buf[12..16]) as usize; + let fragment_no = BigEndian::read_u16(&buf[16..18]); + let num_fragments = BigEndian::read_u16(&buf[18..20]); + + let mut offset = FRAGMENT_HEADER_SIZE; + + if total_size == 0 + || total_size > MAX_MESSAGE_SIZE + || num_fragments == 0 + || fragment_no >= num_fragments + || usize::from(num_fragments) > total_size + || (fragment_no == 0 && fragment_offset != 0) + { + return Ok(None); + } + + // Only fragment zero carries the channel. Reject invalid UTF-8 rather + // than routing a replacement-character channel to the wrong subscriber. + let channel = if fragment_no == 0 { + let Some(pos) = buf[offset..].iter().position(|&b| b == 0) else { + return Ok(None); + }; + if pos > 63 { + return Ok(None); + } + let Ok(ch) = std::str::from_utf8(&buf[offset..offset + pos]) else { + return Ok(None); + }; + offset += pos + 1; + Some(ch.to_owned()) + } else { + None + }; + let payload = &buf[offset..]; + let Some(end) = fragment_offset.checked_add(payload.len()) else { + return Ok(None); + }; + if payload.is_empty() || end > total_size { + return Ok(None); + } + + let key = (sender, seqno); + let now = Instant::now(); + let mut reassembly = self.reassembly.lock().unwrap(); + reassembly.retain(|_, entry| now.duration_since(entry.created) < REASSEMBLY_TIMEOUT); + if !reassembly.contains_key(&key) { + let used: usize = reassembly.values().map(|entry| entry.data.len()).sum(); + if reassembly.len() >= MAX_INCOMPLETE_MESSAGES + || total_size > MAX_REASSEMBLY_BYTES - used + { + return Ok(None); + } + } + let entry = reassembly.entry(key).or_insert_with(|| FragmentBuffer { + channel: None, + fragments: vec![None; usize::from(num_fragments)], + received: 0, + received_bytes: 0, + created: now, + data: vec![0u8; total_size], + }); + if entry.data.len() != total_size || entry.fragments.len() != usize::from(num_fragments) { + reassembly.remove(&key); + return Ok(None); + } + if let Some(range) = entry.fragments[usize::from(fragment_no)] { + // Repeated datagrams do not advance completion. Conflicting repeats + // invalidate this message instead of silently replacing its bytes. + if range != (fragment_offset, end) + || entry.data[fragment_offset..end] != *payload + || (fragment_no == 0 && entry.channel != channel) + { + reassembly.remove(&key); + } + return Ok(None); + } + if entry.received_bytes + payload.len() > total_size { + reassembly.remove(&key); + return Ok(None); + } + if channel.is_some() { + entry.channel = channel; + } + entry.data[fragment_offset..end].copy_from_slice(payload); + entry.fragments[usize::from(fragment_no)] = Some((fragment_offset, end)); + entry.received += 1; + entry.received_bytes += payload.len(); + + if entry.received == entry.fragments.len() { + let complete = reassembly.remove(&key).unwrap(); + if complete.received_bytes != total_size { + return Ok(None); + } + let mut ranges: Vec<_> = complete.fragments.into_iter().flatten().collect(); + ranges.sort_unstable(); + let mut expected = 0; + for (start, end) in ranges { + if start != expected { + return Ok(None); // overlap or hole + } + expected = end; + } + if expected == total_size { + if let Some(channel) = complete.channel { + return Ok(Some(ReceivedMessage { + channel, + data: complete.data, + })); + } + } + } + Ok(None) + } + + fn decode_small(buf: &[u8]) -> io::Result> { + if buf.len() < SHORT_HEADER_SIZE || BigEndian::read_u32(&buf[0..4]) != MAGIC_SHORT { + return Ok(None); + } + let channel_start = SHORT_HEADER_SIZE; + let channel_end = match buf[channel_start..].iter().position(|&b| b == 0) { + Some(pos) => channel_start + pos, + None => return Ok(None), + }; + if channel_end - channel_start > 63 { + return Ok(None); + } + let Ok(channel) = std::str::from_utf8(&buf[channel_start..channel_end]) else { + return Ok(None); + }; + let channel = channel.to_owned(); + let data = buf[channel_end + 1..].to_vec(); + Ok(Some(ReceivedMessage { channel, data })) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn invalid_channels_are_rejected_before_sending() { + let lcm = Lcm::with_options(LcmOptions { + port: 0, + ..Default::default() + }) + .await + .unwrap(); + for channel in ["x".repeat(64), "é".repeat(32), "invalid\0channel".into()] { + let error = lcm.publish(&channel, b"payload").await.unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::InvalidInput); + } + } + + async fn receiver() -> Lcm { + Lcm::with_options(LcmOptions { + port: 0, + ..Default::default() + }) + .await + .unwrap() + } + + fn fragment(seq: u32, size: u32, offset: u32, index: u16, count: u16, data: &[u8]) -> Vec { + let mut packet = vec![0; FRAGMENT_HEADER_SIZE]; + BigEndian::write_u32(&mut packet[0..4], MAGIC_LONG); + BigEndian::write_u32(&mut packet[4..8], seq); + BigEndian::write_u32(&mut packet[8..12], size); + BigEndian::write_u32(&mut packet[12..16], offset); + BigEndian::write_u16(&mut packet[16..18], index); + BigEndian::write_u16(&mut packet[18..20], count); + if index == 0 { + packet.extend_from_slice(b"CHAN\0"); + } + packet.extend_from_slice(data); + packet + } + + fn accept(lcm: &Lcm, packet: &[u8]) -> Option { + lcm.process_fragment("127.0.0.1:1234".parse().unwrap(), packet) + .unwrap() + } + + #[tokio::test] + async fn reordered_and_repeated_fragments_complete_once_with_exact_bytes() { + let lcm = receiver().await; + let last = fragment(1, 6, 4, 2, 3, b"ef"); + assert!(accept(&lcm, &last).is_none()); + assert!(accept(&lcm, &last).is_none()); + assert!(accept(&lcm, &fragment(1, 6, 0, 0, 3, b"ab")).is_none()); + let message = accept(&lcm, &fragment(1, 6, 2, 1, 3, b"cd")).unwrap(); + assert_eq!(message.channel, "CHAN"); + assert_eq!(message.data, b"abcdef"); + assert!(lcm.reassembly.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn invalid_headers_and_ranges_do_not_allocate_reassembly() { + let lcm = receiver().await; + for packet in [ + fragment(1, u32::MAX, 0, 0, 2, b"ab"), + fragment(1, 4, u32::MAX, 1, 2, b"cd"), + fragment(1, 4, 3, 1, 2, b"cd"), + fragment(1, 4, 0, 0, 0, b"ab"), + fragment(1, 4, 0, 2, 2, b"ab"), + fragment(1, 1, 0, 0, 2, b"a"), + fragment(1, 4, 1, 0, 2, b"ab"), + fragment(1, 4, 0, 0, 2, b""), + ] { + assert!(accept(&lcm, &packet).is_none()); + assert!(lcm.reassembly.lock().unwrap().is_empty()); + } + } + + #[tokio::test] + async fn conflicting_metadata_and_repeated_payloads_invalidate_message() { + for conflicting in [ + fragment(1, 5, 2, 1, 2, b"cd"), + fragment(1, 4, 2, 1, 3, b"cd"), + fragment(1, 4, 0, 0, 2, b"zz"), + ] { + let lcm = receiver().await; + assert!(accept(&lcm, &fragment(1, 4, 0, 0, 2, b"ab")).is_none()); + assert!(accept(&lcm, &conflicting).is_none()); + assert!(lcm.reassembly.lock().unwrap().is_empty()); + } + } + + #[tokio::test] + async fn overlaps_and_holes_never_produce_a_message() { + for second in [ + fragment(1, 4, 1, 1, 2, b"cd"), // overlap plus a hole + fragment(1, 4, 3, 1, 2, b"d"), // uncovered byte + fragment(1, 4, 1, 1, 2, b"cde"), // too many bytes + ] { + let lcm = receiver().await; + assert!(accept(&lcm, &fragment(1, 4, 0, 0, 2, b"ab")).is_none()); + assert!(accept(&lcm, &second).is_none()); + assert!(lcm.reassembly.lock().unwrap().is_empty()); + } + } + + #[tokio::test] + async fn incomplete_messages_are_bounded_and_expired_slots_are_reused() { + let lcm = receiver().await; + for seq in 0..MAX_INCOMPLETE_MESSAGES as u32 { + assert!(accept(&lcm, &fragment(seq, 4, 0, 0, 2, b"ab")).is_none()); + } + assert!(accept(&lcm, &fragment(100, 4, 0, 0, 2, b"ab")).is_none()); + assert_eq!( + lcm.reassembly.lock().unwrap().len(), + MAX_INCOMPLETE_MESSAGES + ); + { + let mut buffers = lcm.reassembly.lock().unwrap(); + for buffer in buffers.values_mut() { + buffer.created = Instant::now() - REASSEMBLY_TIMEOUT; + } + } + assert!(accept(&lcm, &fragment(100, 4, 0, 0, 2, b"ab")).is_none()); + assert_eq!(lcm.reassembly.lock().unwrap().len(), 1); + assert_eq!( + accept(&lcm, &fragment(100, 4, 2, 1, 2, b"cd")) + .unwrap() + .data, + b"abcd" + ); + } + + #[tokio::test] + async fn incomplete_payload_storage_is_bounded_across_senders() { + let lcm = receiver().await; + for seq in 0..2 { + assert!(accept(&lcm, &fragment(seq, MAX_MESSAGE_SIZE as u32, 0, 0, 2, b"a")).is_none()); + } + assert!(accept(&lcm, &fragment(3, 4, 0, 0, 2, b"ab")).is_none()); + let buffers = lcm.reassembly.lock().unwrap(); + assert_eq!(buffers.len(), 2); + assert_eq!( + buffers + .values() + .map(|entry| entry.data.len()) + .sum::(), + MAX_REASSEMBLY_BYTES + ); + } + + #[tokio::test] + async fn recv_continues_after_malformed_datagrams() { + let lcm = receiver().await; + let sender = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let destination = + SocketAddrV4::new(Ipv4Addr::LOCALHOST, lcm.socket.local_addr().unwrap().port()); + for packet in [ + fragment(1, 4, u32::MAX, 1, 2, b"cd"), + make_small_packet(&[0xff], b"invalid channel"), + make_small_packet(b"GOOD", b"intact"), + ] { + sender.send_to(&packet, destination).await.unwrap(); + } + let message = tokio::time::timeout(Duration::from_secs(2), lcm.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(message.channel, "GOOD"); + assert_eq!(message.data, b"intact"); + } + + #[test] + fn invalid_short_channels_are_dropped() { + for channel in [vec![b'x'; 64], vec![0xff]] { + assert!(Lcm::decode_small(&make_small_packet(&channel, b"data")) + .unwrap() + .is_none()); + } + } + + const TEST_CHANNEL_LEN: usize = 13; // "/test_channel" + const FIRST_PAYLOAD: usize = MAX_DATAGRAM_SIZE - FRAGMENT_HEADER_SIZE - TEST_CHANNEL_LEN - 1; + const SUBSEQUENT_PAYLOAD: usize = MAX_DATAGRAM_SIZE - FRAGMENT_HEADER_SIZE; + + #[test] + fn fragment_count_fits_in_one() { + let (_, _, n) = fragment_params(FIRST_PAYLOAD, TEST_CHANNEL_LEN); + assert_eq!(n, 1); + } + + #[test] + fn fragment_count_spills_into_two() { + let (_, _, n) = fragment_params(FIRST_PAYLOAD + 1, TEST_CHANNEL_LEN); + assert_eq!(n, 2); + } + + #[test] + fn fragment_count_spills_into_three() { + let (_, _, n) = fragment_params(FIRST_PAYLOAD + SUBSEQUENT_PAYLOAD + 1, TEST_CHANNEL_LEN); + assert_eq!(n, 3); + } + + #[test] + fn fragment_count_one_mb() { + let (_, _, n) = fragment_params(1024 * 1024, TEST_CHANNEL_LEN); + assert_eq!(n, 17); + } + + fn make_small_packet(channel: &[u8], payload: &[u8]) -> Vec { + let mut buf = vec![0u8; SHORT_HEADER_SIZE + channel.len() + 1 + payload.len()]; + BigEndian::write_u32(&mut buf[0..4], MAGIC_SHORT); + BigEndian::write_u32(&mut buf[4..8], 0); + buf[SHORT_HEADER_SIZE..SHORT_HEADER_SIZE + channel.len()].copy_from_slice(channel); + buf[SHORT_HEADER_SIZE + channel.len() + 1..].copy_from_slice(payload); + buf + } + + #[test] + fn decode_small_known_good() { + let buf = make_small_packet(b"CHAN", &[1, 2, 3]); + let msg = Lcm::decode_small(&buf).unwrap().unwrap(); + assert_eq!(msg.channel, "CHAN"); + assert_eq!(msg.data, [1u8, 2, 3]); + } + + #[test] + fn decode_small_empty_payload() { + let buf = make_small_packet(b"CHAN", &[]); + let msg = Lcm::decode_small(&buf).unwrap().unwrap(); + assert_eq!(msg.channel, "CHAN"); + assert!(msg.data.is_empty()); + } + + #[test] + fn decode_small_wrong_magic() { + let mut buf = make_small_packet(b"CHAN", &[1, 2, 3]); + BigEndian::write_u32(&mut buf[0..4], 0xDEADBEEF); + assert!(Lcm::decode_small(&buf).unwrap().is_none()); + } + + #[test] + fn decode_small_truncated() { + // Shorter than SHORT_HEADER_SIZE + let buf = vec![0x4C, 0x43, 0x30, 0x32, 0x00]; + assert!(Lcm::decode_small(&buf).unwrap().is_none()); + } + + #[test] + fn decode_small_missing_null_terminator() { + // Valid header but channel bytes have no null terminator + let mut buf = vec![0u8; SHORT_HEADER_SIZE + 4]; + BigEndian::write_u32(&mut buf[0..4], MAGIC_SHORT); + BigEndian::write_u32(&mut buf[4..8], 0); + buf[SHORT_HEADER_SIZE..SHORT_HEADER_SIZE + 4].copy_from_slice(b"CHAN"); + assert!(Lcm::decode_small(&buf).unwrap().is_none()); + } +} diff --git a/native/rust/dimos-module-macros/src/lib.rs b/native/rust/dimos-module-macros/src/lib.rs index 46d822af9d3..d01140dad60 100644 --- a/native/rust/dimos-module-macros/src/lib.rs +++ b/native/rust/dimos-module-macros/src/lib.rs @@ -428,7 +428,7 @@ struct PortDecl { } /// The last `::`-separated segment of the single generic argument of -/// `Input`/`Output`, e.g. `Input` -> `PointCloud2`. +/// `Input`/`Output`, e.g. `Input` -> `PointCloud2`. fn port_payload_type(ty: &Type) -> Option { let Type::Path(p) = ty else { return None }; let segment = p.path.segments.last()?; @@ -880,8 +880,10 @@ mod tests { #[test] fn payload_type_is_the_last_generic_segment() { - let ty: syn::Type = - parse_str("Input").unwrap(); + let ty: syn::Type = parse_str( + "Input", + ) + .unwrap(); assert_eq!(port_payload_type(&ty).as_deref(), Some("PointCloud2")); let ty: syn::Type = parse_str("Output").unwrap(); assert_eq!(port_payload_type(&ty).as_deref(), Some("Path")); diff --git a/native/rust/dimos-module/Cargo.toml b/native/rust/dimos-module/Cargo.toml index 3d0982e4767..2a22f64d0ca 100644 --- a/native/rust/dimos-module/Cargo.toml +++ b/native/rust/dimos-module/Cargo.toml @@ -6,9 +6,10 @@ description = "Rust native module SDK for dimos NativeModule framework" license = "Apache-2.0" [dependencies] +re_cdr = "=0.1.0" dimos-lcm-transport = { path = "../dimos-lcm-transport", version = "0.1.0" } dimos-module-macros = { version = "=0.1.0", path = "../dimos-module-macros" } -dimos-generated-messages = { version = "0.1.0", path = "../../../dimos/message_codegen" } +dimos-generated-messages = { version = "0.1.0", path = "../../../packages/dimos-generated/build/message-source/rust" } nalgebra = "0.35.0" rayon = "1" tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync", "time", "signal", "io-std", "io-util"] } diff --git a/native/rust/dimos-module/src/cdr.rs b/native/rust/dimos-module/src/cdr.rs index 9749b6e0c17..049edb4a09e 100644 --- a/native/rust/dimos-module/src/cdr.rs +++ b/native/rust/dimos-module/src/cdr.rs @@ -15,37 +15,91 @@ // Copyright 2026 Dimensional Inc. // SPDX-License-Identifier: Apache-2.0 -//! Generated CDR codecs for native module ports. -use dimos_generated_messages::codec::Message; +//! XCDR1 transport framing around the upstream Serde CDR codec. +use re_cdr::{BigEndian, LittleEndian}; +use serde::{de::DeserializeOwned, Serialize}; use std::io; -pub fn encode(message: &T) -> io::Result> { - message.encode() +pub fn encode(message: &T) -> io::Result> { + encode_endian(message, true) } -pub fn decode(bytes: &[u8]) -> io::Result { - T::decode(bytes) +pub fn encode_endian(message: &T, little_endian: bool) -> io::Result> { + let body = if little_endian { + re_cdr::to_vec::<_, LittleEndian>(message) + } else { + re_cdr::to_vec::<_, BigEndian>(message) + } + .map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error.to_string()))?; + let mut bytes = Vec::with_capacity(4 + body.len()); + bytes.extend_from_slice(&[0, u8::from(little_endian), 0, 0]); + bytes.extend(body); + Ok(bytes) +} + +pub fn decode(bytes: &[u8]) -> io::Result { + if bytes.len() < 4 || bytes[0] != 0 || bytes[1] > 1 || bytes[2] != 0 || bytes[3] != 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "Expected plain CDR/XCDR1 encapsulation", + )); + } + let (value, consumed) = if bytes[1] == 1 { + re_cdr::from_bytes::(&bytes[4..]) + } else { + re_cdr::from_bytes::(&bytes[4..]) + } + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error.to_string()))?; + if consumed != bytes.len() - 4 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "Trailing bytes after CDR message", + )); + } + Ok(value) } #[cfg(test)] mod tests { use super::*; - use dimos_generated_messages::{geometry_msgs::msg::PoseStamped, std_msgs::msg::Int32}; + use dimos_generated_messages::{ + builtin_interfaces::msg::time::Time, + geometry_msgs::msg::{ + point::Point, pose::Pose, pose_stamped::PoseStamped, quaternion::Quaternion, + }, + std_msgs::msg::{header::Header, int32::Int32}, + }; #[test] fn nested_fields_and_both_byte_orders() { - let mut message = PoseStamped::default(); - message.header.frame_id = "map".into(); - message.header.stamp.sec = 1700000000; - message.header.stamp.nanosec = 123456789; - message.pose.position.x = 1.5; - message.pose.orientation.w = 1.0; + let message = PoseStamped { + header: Header { + frame_id: "map".into(), + stamp: Time { + sec: 1700000000, + nanosec: 123456789, + }, + }, + pose: Pose { + position: Point { + x: 1.5, + y: 0.0, + z: 0.0, + }, + orientation: Quaternion { + x: 0.0, + y: 0.0, + z: 0.0, + w: 1.0, + }, + }, + }; assert_eq!( decode::(&encode(&message).unwrap()).unwrap(), message ); assert_eq!( - decode::(&message.encode_endian(false).unwrap()).unwrap(), + decode::(&encode_endian(&message, false).unwrap()).unwrap(), message ); } diff --git a/native/rust/dimos-module/src/pointcloud.rs b/native/rust/dimos-module/src/pointcloud.rs index 58a07e04fb2..02f7024e2a5 100644 --- a/native/rust/dimos-module/src/pointcloud.rs +++ b/native/rust/dimos-module/src/pointcloud.rs @@ -12,95 +12,233 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Reading a `PointCloud2` off the wire. Every wire count is signed and untrusted, so each -//! is range-checked before it indexes; a malformed cloud is an `ExtractError`, never a panic. +//! Geometry extraction kept separate from generated message values and codecs. +use std::io; -use lcm_msgs::sensor_msgs::{PointCloud2, PointField}; +use dimos_generated_messages::sensor_msgs::msg::{ + point_cloud2::PointCloud2, point_field::PointField, +}; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct ExtractError(pub &'static str); - -impl std::fmt::Display for ExtractError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.write_str(self.0) - } +fn invalid(reason: &str) -> io::Error { + io::Error::new(io::ErrorKind::InvalidData, reason) } -impl std::error::Error for ExtractError {} - -fn count(v: i32, what: &'static str) -> Result { - usize::try_from(v).map_err(|_| ExtractError(what)) -} - -/// The finite xyz points of a little-endian float32 cloud, in order. -pub fn extract_xyz(msg: &PointCloud2) -> Result, ExtractError> { - let mut offsets: [Option; 3] = [None; 3]; - for f in &msg.fields { - if f.datatype != PointField::FLOAT32 as u8 { - continue; +/// Read finite XYZ points in row order, respecting field offsets, row padding, +/// float32/float64 coordinates and byte order. Other fields are left untouched. +pub fn xyz_points(message: &PointCloud2) -> io::Result> { + let mut coordinates = Vec::with_capacity(3); + let step = message.point_step as usize; + for name in ["x", "y", "z"] { + let mut matching = message.fields.iter().filter(|field| field.name == name); + let field = matching + .next() + .ok_or_else(|| invalid("missing XYZ field"))?; + if matching.next().is_some() || field.count != 1 { + return Err(invalid("XYZ fields must be unique and scalar")); } - let slot = match f.name.as_str() { - "x" => 0, - "y" => 1, - "z" => 2, - _ => continue, + let size = match field.datatype { + PointField::FLOAT32 => 4, + PointField::FLOAT64 => 8, + _ => return Err(invalid("XYZ fields must be float32 or float64")), }; - offsets[slot] = Some(count(f.offset, "negative field offset")?); + let offset = field.offset as usize; + if offset.checked_add(size).is_none_or(|end| end > step) { + return Err(invalid("XYZ field exceeds point_step")); + } + coordinates.push((offset, size)); + } + let width = message.width as usize; + let height = message.height as usize; + let row_step = message.row_step as usize; + let row_bytes = width + .checked_mul(step) + .ok_or_else(|| invalid("row size overflow"))?; + let total = height + .checked_mul(row_step) + .ok_or_else(|| invalid("data size overflow"))?; + if step == 0 || row_step < row_bytes || total != message.data.len() { + return Err(invalid("point cloud dimensions/strides do not match data")); } - let [Some(xo), Some(yo), Some(zo)] = offsets else { - return Err(ExtractError("missing a float32 x/y/z field")); + let count = width + .checked_mul(height) + .ok_or_else(|| invalid("point count overflow"))?; + let read = |base: usize, (offset, size): (usize, usize)| -> f32 { + let bytes = &message.data[base + offset..base + offset + size]; + match (size, message.is_bigendian) { + (4, false) => f32::from_le_bytes(bytes.try_into().unwrap()), + (4, true) => f32::from_be_bytes(bytes.try_into().unwrap()), + (8, false) => f64::from_le_bytes(bytes.try_into().unwrap()) as f32, + (8, true) => f64::from_be_bytes(bytes.try_into().unwrap()) as f32, + _ => unreachable!("field size checked above"), + } }; - if msg.is_bigendian { - return Err(ExtractError("big-endian point data not supported")); + let mut points = Vec::with_capacity(count); + for row in 0..height { + for column in 0..width { + let base = row * row_step + column * step; + let (x, y, z) = ( + read(base, coordinates[0]), + read(base, coordinates[1]), + read(base, coordinates[2]), + ); + if x.is_finite() && y.is_finite() && z.is_finite() { + points.push((x, y, z)); + } + } } + Ok(points) +} - let step = count(msg.point_step, "negative point_step")?; - if step == 0 { - return Err(ExtractError("point_step is 0")); - } - let n = count(msg.width, "negative width")? - .checked_mul(count(msg.height, "negative height")?) - .ok_or(ExtractError("width*height overflows"))?; - let needed = n - .checked_mul(step) - .ok_or(ExtractError("width*height*point_step overflows"))?; - if msg.data.len() < needed { - return Err(ExtractError( - "data buffer shorter than width*height*point_step", - )); +#[cfg(test)] +mod tests { + use super::*; + + fn cloud(big: bool) -> PointCloud2 { + let mut data = vec![0; 64]; + for row in 0..2 { + let x = row as f32 + 1.0; + let z = row as f64 + 3.0; + data[row * 32..row * 32 + 4].copy_from_slice(&if big { + x.to_be_bytes() + } else { + x.to_le_bytes() + }); + data[row * 32 + 8..row * 32 + 16].copy_from_slice(&if big { + z.to_be_bytes() + } else { + z.to_le_bytes() + }); + data[row * 32 + 16..row * 32 + 20].copy_from_slice(&if big { + 2f32.to_be_bytes() + } else { + 2f32.to_le_bytes() + }); + } + PointCloud2 { + width: 1, + height: 2, + point_step: 24, + row_step: 32, + data, + is_bigendian: big, + fields: vec![ + PointField { + name: "z".into(), + offset: 8, + datatype: PointField::FLOAT64, + count: 1, + }, + PointField { + name: "x".into(), + offset: 0, + datatype: PointField::FLOAT32, + count: 1, + }, + PointField { + name: "y".into(), + offset: 16, + datatype: PointField::FLOAT32, + count: 1, + }, + ], + header: dimos_generated_messages::std_msgs::msg::header::Header { + stamp: dimos_generated_messages::builtin_interfaces::msg::time::Time { + sec: 0, + nanosec: 0, + }, + frame_id: String::new(), + }, + is_dense: false, + } } - for off in [xo, yo, zo] { - if off.checked_add(4).is_none_or(|end| end > step) { - return Err(ExtractError( - "xyz field offsets do not fit within point_step", - )); + + #[test] + fn organized_padded_mixed_precision_both_byte_orders() { + for big in [false, true] { + assert_eq!( + xyz_points(&cloud(big)).unwrap(), + vec![(1., 2., 3.), (2., 2., 4.)] + ); } } - let mut out = Vec::with_capacity(n); - for base in (0..needed).step_by(step) { - let x = read_f32_le(&msg.data, base + xo); - let y = read_f32_le(&msg.data, base + yo); - let z = read_f32_le(&msg.data, base + zo); - if x.is_finite() && y.is_finite() && z.is_finite() { - out.push([x, y, z]); + #[test] + fn invalid_layouts_fail_before_reading() { + let mut bad = Vec::new(); + let mut message = cloud(false); + message.row_step = 23; + bad.push(message); + let mut message = cloud(false); + message.data.pop(); + bad.push(message); + let mut message = cloud(false); + message.fields[0].offset = 20; + bad.push(message); + let mut message = cloud(false); + message.fields[0].count = 2; + bad.push(message); + let mut message = cloud(false); + message.fields[0].datatype = PointField::UINT32; + bad.push(message); + let mut message = cloud(false); + message.fields.push(message.fields[0].clone()); + bad.push(message); + let mut message = cloud(false); + message.fields.remove(0); + bad.push(message); + for message in bad { + assert_eq!( + xyz_points(&message).unwrap_err().kind(), + io::ErrorKind::InvalidData + ); } } - Ok(out) + + #[test] + fn nonfinite_coordinates_are_skipped() { + let mut message = cloud(false); + message.data[..4].copy_from_slice(&f32::NAN.to_le_bytes()); + assert_eq!(xyz_points(&message).unwrap(), vec![(2., 2., 4.)]); + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ExtractError(pub &'static str); + +impl std::fmt::Display for ExtractError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.0) + } } -#[inline] -fn read_f32_le(buf: &[u8], off: usize) -> f32 { - let bytes: [u8; 4] = buf[off..off + 4] - .try_into() - .expect("bounds checked by caller"); - f32::from_le_bytes(bytes) +impl std::error::Error for ExtractError {} + +/// Read the little-endian float32 XYZ layout expected by native mapping modules. +/// Preserve their array-valued API while sharing the checked row/stride reader. +pub fn extract_xyz(message: &PointCloud2) -> Result, ExtractError> { + if message.is_bigendian { + return Err(ExtractError("big-endian point data not supported")); + } + for name in ["x", "y", "z"] { + if !message + .fields + .iter() + .any(|field| field.name == name && field.datatype == PointField::FLOAT32) + { + return Err(ExtractError("missing a float32 x/y/z field")); + } + } + xyz_points(message) + .map(|points| points.into_iter().map(|(x, y, z)| [x, y, z]).collect()) + .map_err(|_| ExtractError("invalid point cloud fields, dimensions or strides")) } #[cfg(test)] -mod tests { +mod extraction_tests { use super::*; - use lcm_msgs::std_msgs::Header; + use dimos_generated_messages::{ + builtin_interfaces::msg::time::Time, std_msgs::msg::header::Header, + }; /// An xyz float32 cloud, the shape the raytracer publishes. fn cloud_of(points: &[[f32; 3]]) -> PointCloud2 { @@ -110,20 +248,23 @@ mod tests { data.extend_from_slice(&v.to_le_bytes()); } } - let field = |name: &str, off: i32| PointField { + let field = |name: &str, off: u32| PointField { name: name.into(), offset: off, datatype: PointField::FLOAT32, count: 1, }; PointCloud2 { - header: Header::default(), + header: Header { + stamp: Time { sec: 0, nanosec: 0 }, + frame_id: String::new(), + }, height: 1, - width: points.len() as i32, + width: points.len() as u32, fields: vec![field("x", 0), field("y", 4), field("z", 8)], is_bigendian: false, point_step: 12, - row_step: 12 * points.len() as i32, + row_step: 12 * points.len() as u32, data, is_dense: true, } @@ -155,24 +296,24 @@ mod tests { #[test] fn malformed_metadata_is_an_error_not_a_panic() { - // every wire count is signed; none of these may reach an index + // ROS counts are unsigned; oversized metadata must not reach an index. for (what, tweak) in [ ( - "negative width", - (|c: &mut PointCloud2| c.width = -1) as fn(&mut PointCloud2), + "oversized width", + (|c: &mut PointCloud2| c.width = u32::MAX) as fn(&mut PointCloud2), ), - ("negative height", |c| c.height = -1), - ("negative step", |c| c.point_step = -12), - ("negative offset", |c| c.fields[0].offset = -4), - ("offset past step", |c| c.fields[2].offset = i32::MAX), + ("oversized height", |c| c.height = u32::MAX), + ("oversized step", |c| c.point_step = u32::MAX), + ("oversized offset", |c| c.fields[0].offset = u32::MAX), + ("offset past step", |c| c.fields[2].offset = u32::MAX), ("width*height overflows", |c| { - c.width = i32::MAX; - c.height = i32::MAX; + c.width = u32::MAX; + c.height = u32::MAX; }), ("width*height*step overflows", |c| { - c.width = i32::MAX; - c.height = i32::MAX; - c.point_step = i32::MAX; + c.width = u32::MAX; + c.height = u32::MAX; + c.point_step = u32::MAX; }), ] { let mut cloud = cloud_of(&[[0.0; 3]]); diff --git a/native/rust/dimos-module/src/tf.rs b/native/rust/dimos-module/src/tf.rs index 50bf0b680ca..8cf74ffb505 100644 --- a/native/rust/dimos-module/src/tf.rs +++ b/native/rust/dimos-module/src/tf.rs @@ -25,7 +25,7 @@ use std::sync::atomic::AtomicU64; use std::sync::{Arc, Mutex, RwLock}; use std::time::Duration; -use dimos_generated_messages::codec::Message; +use crate::cdr; use nalgebra::{Isometry3, Quaternion, Translation3, UnitQuaternion, Vector3}; use tokio::sync::{mpsc, Notify}; use tracing::warn; @@ -385,10 +385,10 @@ impl Tf { buffer.receive(&t.parent, &t.child, t.ts, t.iso); } }); - let msg = dimos_generated_messages::tf2_msgs::msg::TFMessage { + let msg = dimos_generated_messages::tf2_msgs::msg::tf_message::TFMessage { transforms: transforms.iter().map(to_stamped).collect(), }; - let bytes = msg.encode()?; + let bytes = cdr::encode(&msg)?; crate::module::publish_encoded(&self.sender, bytes).await } } @@ -482,7 +482,9 @@ impl Lookup<'_> { } } -fn to_stamped(t: &Transform) -> dimos_generated_messages::geometry_msgs::msg::TransformStamped { +fn to_stamped( + t: &Transform, +) -> dimos_generated_messages::geometry_msgs::msg::transform_stamped::TransformStamped { let mut sec = t.ts.floor(); let mut nsec = ((t.ts - sec) * 1e9).round(); if nsec >= 1e9 { @@ -491,22 +493,22 @@ fn to_stamped(t: &Transform) -> dimos_generated_messages::geometry_msgs::msg::Tr } let p = t.iso.translation.vector; let q = t.iso.rotation; - dimos_generated_messages::geometry_msgs::msg::TransformStamped { - header: dimos_generated_messages::std_msgs::msg::Header { - stamp: dimos_generated_messages::builtin_interfaces::msg::Time { + dimos_generated_messages::geometry_msgs::msg::transform_stamped::TransformStamped { + header: dimos_generated_messages::std_msgs::msg::header::Header { + stamp: dimos_generated_messages::builtin_interfaces::msg::time::Time { sec: sec as i32, nanosec: nsec as u32, }, frame_id: t.parent.clone(), }, child_frame_id: t.child.clone(), - transform: dimos_generated_messages::geometry_msgs::msg::Transform { - translation: dimos_generated_messages::geometry_msgs::msg::Vector3 { + transform: dimos_generated_messages::geometry_msgs::msg::transform::Transform { + translation: dimos_generated_messages::geometry_msgs::msg::vector3::Vector3 { x: p.x, y: p.y, z: p.z, }, - rotation: dimos_generated_messages::geometry_msgs::msg::Quaternion { + rotation: dimos_generated_messages::geometry_msgs::msg::quaternion::Quaternion { x: q.i, y: q.j, z: q.k, @@ -525,7 +527,9 @@ struct TfRoute { impl Route for TfRoute { fn try_dispatch(&self, data: &[u8]) { - let msg = match dimos_generated_messages::tf2_msgs::msg::TFMessage::decode(data) { + let msg = match cdr::decode::( + data, + ) { Ok(msg) => msg, Err(e) => { crate::error_throttled!( @@ -870,34 +874,36 @@ mod tests { x: f64, quat: (f64, f64, f64, f64), ) -> Vec { - use dimos_generated_messages::builtin_interfaces::msg::Time; + use dimos_generated_messages::builtin_interfaces::msg::time::Time; use dimos_generated_messages::geometry_msgs::msg::{ - Quaternion as LQuat, Transform as LTransform, TransformStamped, Vector3 as LVec3, + quaternion::Quaternion as LQuat, transform::Transform as LTransform, + transform_stamped::TransformStamped, vector3::Vector3 as LVec3, }; - use dimos_generated_messages::std_msgs::msg::Header; + use dimos_generated_messages::std_msgs::msg::header::Header; let (x_q, y_q, z_q, w_q) = quat; - dimos_generated_messages::tf2_msgs::msg::TFMessage { - transforms: vec![TransformStamped { - header: Header { - stamp: Time { - sec: ts as i32, - nanosec: 0, + cdr::encode( + &dimos_generated_messages::tf2_msgs::msg::tf_message::TFMessage { + transforms: vec![TransformStamped { + header: Header { + stamp: Time { + sec: ts as i32, + nanosec: 0, + }, + frame_id: parent.to_string(), }, - frame_id: parent.to_string(), - }, - child_frame_id: child.to_string(), - transform: LTransform { - translation: LVec3 { x, y: 0.0, z: 0.0 }, - rotation: LQuat { - x: x_q, - y: y_q, - z: z_q, - w: w_q, + child_frame_id: child.to_string(), + transform: LTransform { + translation: LVec3 { x, y: 0.0, z: 0.0 }, + rotation: LQuat { + x: x_q, + y: y_q, + z: z_q, + w: w_q, + }, }, - }, - }], - } - .encode() + }], + }, + ) .unwrap() } @@ -1022,18 +1028,19 @@ mod tests { #[test] fn tf_route_decodes_into_graph() { - use dimos_generated_messages::builtin_interfaces::msg::Time; + use dimos_generated_messages::builtin_interfaces::msg::time::Time; use dimos_generated_messages::geometry_msgs::msg::{ - Quaternion as LQuat, Transform as LTransform, Vector3 as LVec3, + quaternion::Quaternion as LQuat, transform::Transform as LTransform, + vector3::Vector3 as LVec3, }; - use dimos_generated_messages::std_msgs::msg::Header; - use dimos_generated_messages::tf2_msgs::msg::TFMessage; + use dimos_generated_messages::std_msgs::msg::header::Header; + use dimos_generated_messages::tf2_msgs::msg::tf_message::TFMessage; let (tx, _rx) = mpsc::channel(8); let (tf, route) = tf_subscription("/tf".to_string(), DEFAULT_TF_WINDOW_SECS, tx); let msg = TFMessage { transforms: vec![ - dimos_generated_messages::geometry_msgs::msg::TransformStamped { + dimos_generated_messages::geometry_msgs::msg::transform_stamped::TransformStamped { header: Header { stamp: Time { sec: 5, @@ -1058,7 +1065,7 @@ mod tests { }, ], }; - route.try_dispatch(&msg.encode().unwrap()); + route.try_dispatch(&cdr::encode(&msg).unwrap()); let t = tf.get_latest("base_link", "mid360_link").unwrap(); assert!((t.translation().x - 0.1).abs() < 1e-9); diff --git a/native/rust/dimos-module/src/time.rs b/native/rust/dimos-module/src/time.rs index 45d0f9235a5..50075424b13 100644 --- a/native/rust/dimos-module/src/time.rs +++ b/native/rust/dimos-module/src/time.rs @@ -16,7 +16,7 @@ use std::time::{SystemTime, UNIX_EPOCH}; -use lcm_msgs::std_msgs::Time; +use dimos_generated_messages::builtin_interfaces::msg::time::Time; /// The wall clock as a message stamp. pub fn now() -> Time { @@ -25,7 +25,7 @@ pub fn now() -> Time { .unwrap_or_default(); Time { sec: dur.as_secs().min(i32::MAX as u64) as i32, - nsec: dur.subsec_nanos() as i32, + nanosec: dur.subsec_nanos(), } } diff --git a/packages/dimos-generated/catalog.json b/packages/dimos-generated/catalog.json index 8994a6eaedd..4cd5bf01262 100644 --- a/packages/dimos-generated/catalog.json +++ b/packages/dimos-generated/catalog.json @@ -74,6 +74,8 @@ "dimos_msgs/msg/Contacts": "DimOS runtime message contract", "dimos_msgs/msg/RegionBounds": "DimOS runtime message contract", "dimos_msgs/msg/RegionLineSegments3D": "DimOS runtime message contract", - "dimos_msgs/msg/RegionPointCloud2": "DimOS runtime message contract" + "dimos_msgs/msg/RegionPointCloud2": "DimOS runtime message contract", + "dimos_msgs/msg/TriangleMesh": "Published runtime consumer: dimos/mapping/experimental/mesh.py", + "dimos_msgs/msg/ControlValues": "Published runtime consumer: dimos/control/contract/validate.py" } } diff --git a/packages/dimos-generated/custom/dimos_msgs/msg/ControlValues.msg b/packages/dimos-generated/custom/dimos_msgs/msg/ControlValues.msg new file mode 100644 index 00000000000..3db412d903e --- /dev/null +++ b/packages/dimos-generated/custom/dimos_msgs/msg/ControlValues.msg @@ -0,0 +1,13 @@ +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +# One publisher's sparse interface values; lifecycle transitions use RPC. +string source +# Publisher Unix timestamp in seconds, not receiver-side freshness. +float64 source_ts +# State generation or command session; sequence ordering applies within an epoch. +int64 epoch +int64 sequence +# Parallel arrays; the control contract validates shape and finite values. +string[] interface_names +float64[] values diff --git a/packages/dimos-generated/custom/dimos_msgs/msg/TriangleMesh.msg b/packages/dimos-generated/custom/dimos_msgs/msg/TriangleMesh.msg new file mode 100644 index 00000000000..3d602fe36a5 --- /dev/null +++ b/packages/dimos-generated/custom/dimos_msgs/msg/TriangleMesh.msg @@ -0,0 +1,6 @@ +# A keyed mesh chunk. Flat arrays contain consecutive xyz vertices, normals, and triangles. +std_msgs/Header header +int32[3] key +float32[] vertices +float32[] normals +uint32[] faces diff --git a/scripts/install_native_messages.sh b/scripts/install_native_messages.sh index ae04d66c271..e350466567a 100644 --- a/scripts/install_native_messages.sh +++ b/scripts/install_native_messages.sh @@ -1,12 +1,27 @@ #!/usr/bin/env bash set -euo pipefail -# Generate/install the built-in C++ message package after setup_message_codegen.sh. -# Uses Python's standard library only; no ROS installation or Python extension build. +# Build the built-in C++ package with the same pinned support builder as dimos build. +# Requires dimos-cdrgen[native]; no ROS installation or runtime compilation. cd "$(dirname "$0")/.." -python3 -m dimos.message_codegen.generate --output build/message-codegen/native -cmake -S build/message-codegen/native/cpp -B build/message-codegen/native/cmake \ - -DDIMOS_BUILD_PYTHON=OFF \ - -DCMAKE_PREFIX_PATH="$PWD/build/message-codegen/install${CMAKE_PREFIX_PATH:+;$CMAKE_PREFIX_PATH}" \ - -DCMAKE_INSTALL_PREFIX="$PWD/build/message-codegen/install" -cmake --install build/message-codegen/native/cmake +"${DIMOS_CODEGEN_PYTHON:-python3}" - <<'BUILD' +import json +import os +from pathlib import Path + +from dimos_cdrgen.generate import generate +from dimos_cdrgen.native_build import installed_prefixes, prepare_cpp +from dimos_cdrgen.registry import schema_roots +from dimos.message_codegen.native_build import write_cmake_toolchain + +output = Path("build/message-codegen/native").resolve() +generate(list(schema_roots()), output, shared=True) +sources = {name: Path(path) for name, path in json.loads(os.environ.get("DIMOS_NATIVE_SOURCE_DIRS", "{}")).items()} +cache = os.environ.get("DIMOS_NATIVE_CACHE") +prefix = prepare_cpp( + output, cache=Path(cache) if cache else None, offline=os.environ.get("DIMOS_OFFLINE") == "1", + source_dirs=sources or None +) +write_cmake_toolchain(prefix, output / "toolchain.cmake") +(output / "prefixes.txt").write_text(";".join(installed_prefixes(prefix))) +BUILD diff --git a/scripts/package_native_sdk.py b/scripts/package_native_sdk.py index dd270af9a89..4818278b283 100644 --- a/scripts/package_native_sdk.py +++ b/scripts/package_native_sdk.py @@ -32,7 +32,10 @@ def main() -> None: shutil.copytree(source, target, ignore=shutil.ignore_patterns("target")) manifest = target / "Cargo.toml" manifest.write_text( - manifest.read_text().replace(', path = "../../../dimos/message_codegen"', "") + manifest.read_text().replace( + ', path = "../../../packages/dimos-generated/build/message-source/rust"', + "", + ) ) (args.output / "rust" / "Cargo.toml").write_text( '[workspace]\nmembers = ["dimos-module", "dimos-module-macros", "dimos-lcm-transport"]\nresolver = "2"\n' diff --git a/scripts/test_external_native_messages.sh b/scripts/test_external_native_messages.sh index 9648cbbfd1e..e033384ab7d 100644 --- a/scripts/test_external_native_messages.sh +++ b/scripts/test_external_native_messages.sh @@ -25,6 +25,6 @@ cmake -S examples/message-codegen/external-native -B "$native_demo/cpp" \ -DFETCHCONTENT_SOURCE_DIR_PFR="$PWD/build/native-cpp/_deps/pfr-src" cmake --build "$native_demo/cpp" -j 2 application_site=$("$application/venv/bin/python" -c 'import sysconfig; print(sysconfig.get_path("platlib"))') -PYTHONPATH="$PWD:$PWD/build/message-codegen/demo/cpp/build:$application_site${PYTHONPATH:+:$PYTHONPATH}" \ +PYTHONPATH="$PWD:$PWD/build/message-codegen/demo/python:$application_site${PYTHONPATH:+:$PYTHONPATH}" \ .venv/bin/python examples/message-codegen/demo_external_native.py --build "$native_demo" \ | tee "$native_demo/evidence/exchange.txt" diff --git a/scripts/test_message_user_story.sh b/scripts/test_message_user_story.sh index 9e7e9f8e685..d0958934916 100644 --- a/scripts/test_message_user_story.sh +++ b/scripts/test_message_user_story.sh @@ -4,20 +4,16 @@ cd "$(dirname "$0")/.." story_build="$PWD/build/message-codegen/user-story" story_python="${DIMOS_CODEGEN_PYTHON:-$PWD/.venv/bin/python}" story_fastcdr="${DIMOS_FASTCDR_PREFIX:-$PWD/build/message-codegen/install}" -"$story_python" -m dimos.message_codegen.generate \ +"$story_python" -m dimos_cdrgen.generate \ --package-root examples/message-codegen/user-story \ --type story_msgs/msg/DeviceReading --python-module story_messages \ --output "$story_build" -cmake -S "$story_build/cpp" -B "$story_build/cpp/build" \ - -DCMAKE_BUILD_TYPE=Release -DCMAKE_PREFIX_PATH="$story_fastcdr" \ - -DPython_EXECUTABLE="$story_python" -cmake --build "$story_build/cpp/build" -j 2 c++ -std=c++17 -I"$story_fastcdr/include" -I"$story_build/cpp" \ examples/message-codegen/user-story/consumer.cpp "$story_fastcdr/lib/libfastcdr.a" \ -o "$story_build/cpp-consumer" mkdir -p "$story_build/rust/src/bin" "$story_build/evidence" cp examples/message-codegen/user-story/consumer.rs "$story_build/rust/src/bin/story-consumer.rs" cargo build --offline --manifest-path "$story_build/rust/Cargo.toml" --bin story-consumer -PYTHONPATH="$story_build/cpp/build:$PWD${PYTHONPATH:+:$PYTHONPATH}" \ +PYTHONPATH="$story_build/python:$PWD${PYTHONPATH:+:$PYTHONPATH}" \ "$story_python" examples/message-codegen/user-story/demo_module.py --build "$story_build" \ | tee "$story_build/evidence/terminal.txt" diff --git a/web/cockpit/src/subscriptions.ts b/web/cockpit/src/subscriptions.ts index 5ff52e5b0ab..a1f309100e8 100644 --- a/web/cockpit/src/subscriptions.ts +++ b/web/cockpit/src/subscriptions.ts @@ -5,7 +5,7 @@ // manifest and view change. import type { ChannelSpec, DecoderRegistry, Session } from "@dimos/sdk"; -import { createDecoderRegistry, isLcmSchema, TRACK_ENCODING } from "@dimos/sdk"; +import { cdrHasSequences, createDecoderRegistry, isCdrSchema, TRACK_ENCODING } from "@dimos/sdk"; import type { PanelSpec } from "@dimos/shared"; import { getPanel } from "./panels/registry.tsx"; @@ -23,18 +23,15 @@ const PANEL_ONLY_ENCODINGS = new Set([ TRACK_ENCODING, ]); -/** Channels only a panel may subscribe: the encodings above, and a *.lcm.v1 +/** Channels only a panel may subscribe: the encodings above, and a *.cdr.v1 * channel whose schema has a variable-length array (a point cloud, a scan, a * path): every frame costs the message's full size, far more than the channel * table's preview is worth. A bounded message (a pose, an odometry) * subscribes like json.v1 so the table can show it. */ function panelOnly(spec: ChannelSpec): boolean { if (PANEL_ONLY_ENCODINGS.has(spec.encoding)) return true; - const lcm = spec.params.lcm; - return isLcmSchema(lcm) && - Object.values(lcm.structs).some((rows) => - rows.some(([, , dims]) => dims !== null && dims.some((d) => typeof d === "string")) - ); + const cdr = spec.params.cdr; + return isCdrSchema(cdr) && cdrHasSequences(cdr); } /** True when this build can put the channel to use: rx only, it has a diff --git a/web/deno.lock b/web/deno.lock index 0fe06125699..f87e68507f3 100644 --- a/web/deno.lock +++ b/web/deno.lock @@ -8,6 +8,8 @@ "jsr:@std/internal@^1.0.12": "1.0.12", "npm:@fontsource-variable/inter@^5.3.0": "5.3.0", "npm:@fontsource-variable/jetbrains-mono@^5.3.0": "5.3.0", + "npm:@foxglove/rosmsg2-serialization@3.1.2": "3.1.2", + "npm:@foxglove/rosmsg@5.0.5": "5.0.5", "npm:@peculiar/x509@2": "2.0.0", "npm:@types/node@^24.10.1": "24.12.2", "npm:@types/react-dom@^19.2.0": "19.2.3_@types+react@19.2.17", @@ -327,6 +329,30 @@ "@fontsource-variable/jetbrains-mono@5.3.0": { "integrity": "sha512-F32xpS2NsGYoQi2ADSkKTgpJj7ozajsGgDJ8woTnqjmIB+dxDIqImjl4pXZVEExu8UFZ2ndhmX18EBS/hdz3Lw==" }, + "@foxglove/cdr@3.5.1": { + "integrity": "sha512-Fubm04/xC1SkeNaviwUhbFk8nAiiNvIqlmP5MW7r6IVGFa35rsA0E3/Rg/pGdE77ZXf4XNwShZcQoEmje62ddQ==" + }, + "@foxglove/message-definition@0.5.0": { + "integrity": "sha512-B/2GvWn4wRSsPDywcmy6uNXgAfZRhiOAnmy3qAaS6MHVqiHQqs/RoDsoW6Gw9sXSweJpeH5qAqAihvm1pR9Thg==" + }, + "@foxglove/rosmsg2-serialization@3.1.2": { + "integrity": "sha512-KKtmHuuIQy6I0JLUi01LFhV7hrZb5xd+LegtuGz2+9L9pftBxtVVKhoxO6B/kQREGocIUOLP8494kxLF58LMWw==", + "dependencies": [ + "@foxglove/cdr", + "@foxglove/message-definition", + "@foxglove/rostime" + ] + }, + "@foxglove/rosmsg@5.0.5": { + "integrity": "sha512-+qY6xW+4HdbPI3ItAlFFMjXgAiplXyRsv+kqdBk8OcWGH38BmEd7s6v/VnCkbYovAoB5rfR5giuKfmPRNzwbrA==", + "dependencies": [ + "@foxglove/message-definition", + "md5-typescript" + ] + }, + "@foxglove/rostime@1.1.3": { + "integrity": "sha512-Gr3Xy4sNbo2NcXg0tBHb/jVUqIcDyZ1kgHj9dB7K+Y1A7AQ0oSCmSXIuf8oqeteSIs0pl8BJOWEwAtLNC99Izg==" + }, "@jridgewell/gen-mapping@0.3.13": { "integrity": "sha512-2kkt/7niJ6MgEPxF0bYdQ6etZaA+fQvDcLKckhy1yIQOzaoKjBBjSj63/aLVjYE3qhRt5dvM+uUyfCg6UKCBbA==", "dependencies": [ @@ -939,6 +965,9 @@ "@jridgewell/sourcemap-codec" ] }, + "md5-typescript@1.0.5": { + "integrity": "sha512-ovAc4EtiNt2dY8JPhPr/wkC9h4U5k/nuClNVcG0Ga3V1rMlYpAY24ZaaymFXJlz+ccJ6UMPo3FSaVKe7czBsXw==" + }, "meshoptimizer@1.1.1": { "integrity": "sha512-oRFNWJRDA/WTrVj7NWvqa5HqE1t9MYDj2VaWirQCzCCrAd2GHrqR/sQezCxiWATPNlKTcRaPRHPJwIRoPBAp5g==" }, @@ -1229,6 +1258,8 @@ "sdk": { "packageJson": { "dependencies": [ + "npm:@foxglove/rosmsg2-serialization@3.1.2", + "npm:@foxglove/rosmsg@5.0.5", "npm:@types/node@^24.10.1", "npm:@types/react-dom@^19.2.0", "npm:@types/react@^19.2.0", diff --git a/web/sdk/src/decoders/decoders.test.ts b/web/sdk/src/decoders/decoders.test.ts index d192f4b0b1c..d3d3d6ab3ad 100644 --- a/web/sdk/src/decoders/decoders.test.ts +++ b/web/sdk/src/decoders/decoders.test.ts @@ -1,7 +1,7 @@ import { describe, expect, it } from "vitest"; import type { FrameHeader } from "@dimos/shared"; import costmapFrames from "../../../shared/fixtures/costmap_frames.json"; -import lcmFrames from "../../../shared/fixtures/lcm_frames.json"; +import cdrFrames from "../../../shared/fixtures/cdr_frames.json"; import voxelFrames from "../../../shared/fixtures/voxel_frames.json"; import { spec } from "../testing/fakeRelay.ts"; import { @@ -38,14 +38,14 @@ async function deflate(bytes: Uint8Array): Promise { return new Uint8Array(await new Response(stream).arrayBuffer()); } -// The pose_stamped golden vector (Python-generated, see lcm.test.ts). -const POSE = (lcmFrames as { vectors: { name: string; schema: unknown; payload_b64: string }[] }) +// The pose_stamped golden vector (Python-generated, see cdr.test.ts). +const POSE = (cdrFrames as { vectors: { name: string; schema: unknown; payload_b64: string }[] }) .vectors[0]; const poseSpec = () => spec({ - ch: "lcm_pose", - encoding: "geometry_msgs.PoseStamped.lcm.v1", - params: { lcm: POSE.schema }, + ch: "cdr_pose", + encoding: "geometry_msgs/msg/PoseStamped.cdr.v1", + params: { cdr: POSE.schema }, }); /** Minimal scannable JPEG: SOI + SOF0 declaring w x h (no scan data). */ @@ -74,7 +74,7 @@ describe("decoder registry", () => { expect(registry.get(undefined)).toBeUndefined(); }); - it("resolve() compiles *.lcm.v1 from params.lcm and caches per manifest record", () => { + it("resolve() compiles *.cdr.v1 from params.cdr and caches per manifest record", () => { expect(POSE.name).toBe("pose_stamped"); const s = poseSpec(); const decode = registry.resolve(s); @@ -93,10 +93,10 @@ describe("decoder registry", () => { expect(registry.get(s.encoding)).toBeUndefined(); }); - it("resolve() lets an exact registration beat the lcm family rule", () => { + it("resolve() lets an exact registration beat the cdr family rule", () => { const own = createDecoderRegistry(); const mine: Decoder = () => ({ value: "mine" }); - own.register("geometry_msgs.PoseStamped.lcm.v1", mine); + own.register("geometry_msgs/msg/PoseStamped.cdr.v1", mine); expect(own.resolve(poseSpec())).toBe(mine); }); @@ -104,12 +104,12 @@ describe("decoder registry", () => { for ( const params of [ {}, - { lcm: null }, - { lcm: { type: "t.P", fp: "zz", structs: {} } }, - { lcm: { type: "t.P", fp: "0011223344556677", structs: {} } }, + { cdr: null }, + { cdr: { type: "t.P", fp: "zz", structs: {} } }, + { cdr: { type: "t.P", fp: "0011223344556677", structs: {} } }, ] ) { - const s = spec({ ch: "x", encoding: "t.P.lcm.v1", params }); + const s = spec({ ch: "x", encoding: "t/msg/P.cdr.v1", params }); expect(registry.resolve(s)).toBeUndefined(); expect(registry.resolve(s)).toBeUndefined(); } diff --git a/web/sdk/src/decoders/index.ts b/web/sdk/src/decoders/index.ts index dd8ce12e0d3..959cb267da6 100644 --- a/web/sdk/src/decoders/index.ts +++ b/web/sdk/src/decoders/index.ts @@ -1,5 +1,5 @@ // Payload decoder registry. resolve(spec) finds a manifest channel's decoder: -// a registered encoding id, the *.json.vN convention, or a *.lcm.v1 decoder +// a registered encoding id, the *.json.vN convention, or a *.cdr.v1 decoder // compiled from the schema the robot put in the channel's params. An encoding // without a decoder is not an error: the channel renders as "unsupported" // (forward compatibility with newer bridges). Binary decoders (h264.v1, ...) @@ -12,7 +12,7 @@ import { TRACK_ENCODING } from "@dimos/shared/manifest"; import { costmapDecoder } from "./costmap.ts"; import { jpegDecoder } from "./jpeg.ts"; import { jsonDecoder } from "./json.ts"; -import { LCM_ENCODING_RE, lcmDecoderFor } from "./lcm.ts"; +import { CDR_ENCODING_RE, cdrDecoderFor } from "./cdr.ts"; import { voxelsDecoder } from "./voxels.ts"; export interface Decoded { @@ -25,9 +25,9 @@ export type Decoder = (payload: Uint8Array, header: FrameHeader) => Decoded; export class DecoderRegistry { #decoders = new Map(); - // Compiled *.lcm.v1 decoders per adopted manifest record (null: params.lcm + // Compiled *.cdr.v1 decoders per adopted manifest record (null: params.cdr // unusable); entries die with the manifest object that carried them. - #lcm = new WeakMap(); + #cdr = new WeakMap(); /** Duplicate registration is an error unless `replace` is set. */ register(encoding: string, decoder: Decoder, opts: { replace?: boolean } = {}): void { @@ -47,16 +47,16 @@ export class DecoderRegistry { } /** The decoder for a manifest channel: get(spec.encoding) first (registered - * ids and the JSON convention win), else a *.lcm.v1 decoder compiled once - * from the schema in spec.params.lcm. */ + * ids and the JSON convention win), else a *.cdr.v1 decoder compiled once + * from the schema in spec.params.cdr. */ resolve(spec: ChannelSpec): Decoder | undefined { const known = this.get(spec.encoding); if (known !== undefined) return known; - if (!LCM_ENCODING_RE.test(spec.encoding)) return undefined; - let compiled = this.#lcm.get(spec); + if (!CDR_ENCODING_RE.test(spec.encoding)) return undefined; + let compiled = this.#cdr.get(spec); if (compiled === undefined) { - compiled = lcmDecoderFor(spec.params.lcm); - this.#lcm.set(spec, compiled); + compiled = cdrDecoderFor(spec.params.cdr); + this.#cdr.set(spec, compiled); } return compiled ?? undefined; } diff --git a/web/sdk/src/index.ts b/web/sdk/src/index.ts index f7efa3c7846..fc7a363b10c 100644 --- a/web/sdk/src/index.ts +++ b/web/sdk/src/index.ts @@ -14,15 +14,8 @@ export type { ChannelSnapshot, ChannelStats, SessionStatus, Slot } from "./store export { createDecoderRegistry, DecoderRegistry } from "./decoders/index.ts"; export type { Decoded, Decoder } from "./decoders/index.ts"; export { type CostmapValue, inflateCostmap } from "./decoders/costmap.ts"; +export { cdrDecoder, cdrHasSequences, type CdrSchema, isCdrSchema } from "./decoders/cdr.ts"; export { inflateVoxels, type VoxelsValue } from "./decoders/voxels.ts"; -export { - isLcmSchema, - lcmDecoder, - type LcmDecoderOptions, - type LcmSchema, - type LcmValue, - MAX_LCM_ARRAY_ELEMENTS, -} from "./decoders/lcm.ts"; export { PublishError, type PublishOutcome, From 4dcc03a3f01c61f6de41d6c3de7033e541dd3c69 Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 15 Oct 2026 00:20:21 -0700 Subject: [PATCH 10/11] fix(cdr): include native recorder stamped geometry schemas --- dimos/message_codegen/test_builtin_catalog.py | 2 +- dimos/message_codegen/test_builtin_install.py | 2 +- dimos/message_codegen/test_builtin_preparation.py | 4 ++-- packages/dimos-generated/catalog.json | 5 ++++- packages/dimos-generated/schema-sources.json | 12 ++++++++++++ 5 files changed, 20 insertions(+), 5 deletions(-) diff --git a/dimos/message_codegen/test_builtin_catalog.py b/dimos/message_codegen/test_builtin_catalog.py index 5c17805a86f..fbc37ed8c5f 100644 --- a/dimos/message_codegen/test_builtin_catalog.py +++ b/dimos/message_codegen/test_builtin_catalog.py @@ -36,7 +36,7 @@ def test_unlisted_sources_do_not_expand_the_public_catalog(tmp_path): manifest = json.loads( (source / "dimos_generated_schemas/package/message-package.json").read_text() ) - assert len(manifest["owned"]) == 80 + assert len(manifest["owned"]) == 83 assert "extra_msgs/msg/NotExported" not in manifest["owned"] assert "builtin_interfaces/msg/Time" in manifest["owned"] diff --git a/dimos/message_codegen/test_builtin_install.py b/dimos/message_codegen/test_builtin_install.py index 97c532478c0..d5c862d7271 100644 --- a/dimos/message_codegen/test_builtin_install.py +++ b/dimos/message_codegen/test_builtin_install.py @@ -38,7 +38,7 @@ assert type(r.header) is Header and r == m assert type(pickle.loads(pickle.dumps(m)).header) is Header assert 'MSG: std_msgs/Header' in schema(m.__msgtype__) -assert len(message_types()) == 80 +assert len(message_types()) == 83 assert 'visualization_msgs/msg/Marker' not in message_types() root = Path(package_root()) assert (root / 'schema-sources.json').is_file() diff --git a/dimos/message_codegen/test_builtin_preparation.py b/dimos/message_codegen/test_builtin_preparation.py index d4c4c7a068b..20c948ad673 100644 --- a/dimos/message_codegen/test_builtin_preparation.py +++ b/dimos/message_codegen/test_builtin_preparation.py @@ -48,11 +48,11 @@ def no_network(*args, **kwargs): monkeypatch.setattr("dimos_cdrgen.source_inputs.requests.get", no_network) inputs = prepare(cache, offline=True, project=project) - assert len(list(inputs.rglob("*.msg"))) == 80 + assert len(list(inputs.rglob("*.msg"))) == 83 output = generate(project) resources = output / "dimos_generated_schemas/package" manifest = json.loads((resources / "message-package.json").read_text()) - assert len(manifest["owned"]) == 80 + assert len(manifest["owned"]) == 83 assert len(list(project.glob("custom/dimos_msgs/msg/*.msg"))) == 21 lock = json.loads((project / "schema-sources.json").read_text()) assert all(len(item["sha256"]) == 64 for item in lock["sources"]) diff --git a/packages/dimos-generated/catalog.json b/packages/dimos-generated/catalog.json index 4cd5bf01262..af5d96e666f 100644 --- a/packages/dimos-generated/catalog.json +++ b/packages/dimos-generated/catalog.json @@ -76,6 +76,9 @@ "dimos_msgs/msg/RegionLineSegments3D": "DimOS runtime message contract", "dimos_msgs/msg/RegionPointCloud2": "DimOS runtime message contract", "dimos_msgs/msg/TriangleMesh": "Published runtime consumer: dimos/mapping/experimental/mesh.py", - "dimos_msgs/msg/ControlValues": "Published runtime consumer: dimos/control/contract/validate.py" + "dimos_msgs/msg/ControlValues": "Published runtime consumer: dimos/control/contract/validate.py", + "geometry_msgs/msg/PoseWithCovarianceStamped": "Published native runtime consumer: dimos/experimental/memory/rust/src/decoding.rs", + "geometry_msgs/msg/TwistWithCovarianceStamped": "Published native runtime consumer: dimos/experimental/memory/rust/src/decoding.rs", + "geometry_msgs/msg/WrenchStamped": "Published native runtime consumer: dimos/experimental/memory/rust/src/decoding.rs" } } diff --git a/packages/dimos-generated/schema-sources.json b/packages/dimos-generated/schema-sources.json index 65e2d266b09..9bc56261473 100644 --- a/packages/dimos-generated/schema-sources.json +++ b/packages/dimos-generated/schema-sources.json @@ -215,6 +215,18 @@ "licenses/ros2_common_interfaces_visualization_msgs_LICENSE": { "path": "visualization_msgs/LICENSE", "sha256": "cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30" + }, + "geometry_msgs/msg/PoseWithCovarianceStamped.msg": { + "path": "geometry_msgs/msg/PoseWithCovarianceStamped.msg", + "sha256": "87b6851480902effb5232f87bf68e76f19bf7b26affa11001be9e515542c21fb" + }, + "geometry_msgs/msg/TwistWithCovarianceStamped.msg": { + "path": "geometry_msgs/msg/TwistWithCovarianceStamped.msg", + "sha256": "af369aea38a14e57ff350a3f2c63c2c66335a367822d134fecdef60789e49682" + }, + "geometry_msgs/msg/WrenchStamped.msg": { + "path": "geometry_msgs/msg/WrenchStamped.msg", + "sha256": "853a8534230b795b6c1afcb8483002f5d1cf14e967ffa0da2cba8da5453b2162" } } }, From 8c24da76ee5f2fb676b9bf9ddb7c7060ddcbed1f Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Mon, 19 Oct 2026 18:14:40 +0000 Subject: [PATCH 11/11] [autofix.ci] apply automated fixes --- .../sensors/lidar/fastlio2/cpp/flake.lock | 24 ++------------ .../trajectory_follower/fancy/rust/flake.lock | 32 ------------------- dimos/simulation/habitat/nix/flake.lock | 6 ++-- examples/native-modules/cpp/flake.lock | 24 ++------------ 4 files changed, 9 insertions(+), 77 deletions(-) diff --git a/dimos/hardware/sensors/lidar/fastlio2/cpp/flake.lock b/dimos/hardware/sensors/lidar/fastlio2/cpp/flake.lock index 196d160769a..9bb331092dd 100644 --- a/dimos/hardware/sensors/lidar/fastlio2/cpp/flake.lock +++ b/dimos/hardware/sensors/lidar/fastlio2/cpp/flake.lock @@ -1,22 +1,5 @@ { "nodes": { - "dimos-lcm": { - "flake": false, - "locked": { - "lastModified": 1769774949, - "narHash": "sha256-icRK7jerqNlwK1WZBrnIP04I2WozzFqTD7qsmnPxQuo=", - "owner": "dimensionalOS", - "repo": "dimos-lcm", - "rev": "0aa72b7b1bd3a65f50f5c03485ee9b728df56afe", - "type": "github" - }, - "original": { - "owner": "dimensionalOS", - "ref": "main", - "repo": "dimos-lcm", - "type": "github" - } - }, "dimos-native-cpp": { "inputs": { "flake-utils": "flake-utils", @@ -28,11 +11,11 @@ }, "locked": { "dir": "native/cpp", - "lastModified": 1792432073, - "narHash": "sha256-5BgFoVp3/X5bKasI5G+f41sqVGAnZ4VPi1w3eC6+DIM=", + "lastModified": 1792433568, + "narHash": "sha256-3O7ECiGlq/HThKj65pi3rzjg0ny6iMGAuKXXEwt4GRc=", "owner": "dimensionalOS", "repo": "dimos", - "rev": "ea1d296b163d325d2dea782c4057dbaa3b697fc4", + "rev": "4c0afa2cc34247e537bbb57ffb4a5bc6d7007555", "type": "github" }, "original": { @@ -207,7 +190,6 @@ }, "root": { "inputs": { - "dimos-lcm": "dimos-lcm", "dimos-native-cpp": "dimos-native-cpp", "fast-lio": "fast-lio", "flake-utils": [ diff --git a/dimos/navigation/trajectory_follower/fancy/rust/flake.lock b/dimos/navigation/trajectory_follower/fancy/rust/flake.lock index d3f170199e1..12d4afa66f6 100644 --- a/dimos/navigation/trajectory_follower/fancy/rust/flake.lock +++ b/dimos/navigation/trajectory_follower/fancy/rust/flake.lock @@ -73,37 +73,6 @@ "type": "github" } }, - "dimos-local-planner": { - "inputs": { - "crate2nix": [ - "crate2nix" - ], - "flake-utils": [ - "flake-utils" - ], - "nix-filter": [ - "nix-filter" - ], - "nixpkgs": [ - "nixpkgs" - ] - }, - "locked": { - "dir": "dimos/navigation/local_planner/rust", - "lastModified": 1792134367, - "narHash": "sha256-MFGVh66kb5TXulN5CAtnNbUTDBDxjdSogd7gqhVSt14=", - "owner": "dimensionalOS", - "repo": "dimos", - "rev": "b938c8a601dae4a0de3af093f1af2e36950bfdbb", - "type": "github" - }, - "original": { - "dir": "dimos/navigation/local_planner/rust", - "owner": "dimensionalOS", - "repo": "dimos", - "type": "github" - } - }, "flake-compat": { "locked": { "lastModified": 1733328505, @@ -322,7 +291,6 @@ "root": { "inputs": { "crate2nix": "crate2nix", - "dimos-local-planner": "dimos-local-planner", "flake-utils": "flake-utils", "nix-filter": "nix-filter", "nixpkgs": "nixpkgs_2" diff --git a/dimos/simulation/habitat/nix/flake.lock b/dimos/simulation/habitat/nix/flake.lock index 5ef7dc7f483..cbc1b1b64bb 100644 --- a/dimos/simulation/habitat/nix/flake.lock +++ b/dimos/simulation/habitat/nix/flake.lock @@ -11,11 +11,11 @@ }, "locked": { "dir": "native/cpp", - "lastModified": 1792432073, - "narHash": "sha256-5BgFoVp3/X5bKasI5G+f41sqVGAnZ4VPi1w3eC6+DIM=", + "lastModified": 1792433568, + "narHash": "sha256-3O7ECiGlq/HThKj65pi3rzjg0ny6iMGAuKXXEwt4GRc=", "owner": "dimensionalOS", "repo": "dimos", - "rev": "ea1d296b163d325d2dea782c4057dbaa3b697fc4", + "rev": "4c0afa2cc34247e537bbb57ffb4a5bc6d7007555", "type": "github" }, "original": { diff --git a/examples/native-modules/cpp/flake.lock b/examples/native-modules/cpp/flake.lock index adbcf88e1f5..1fc58728c68 100644 --- a/examples/native-modules/cpp/flake.lock +++ b/examples/native-modules/cpp/flake.lock @@ -1,22 +1,5 @@ { "nodes": { - "dimos-lcm": { - "flake": false, - "locked": { - "lastModified": 1782118008, - "narHash": "sha256-z88X38ZM3vr1eneG1CkJ4GAGWBW3JYWPEX/2x7aRVJI=", - "owner": "dimensionalOS", - "repo": "dimos-lcm", - "rev": "0a1c24710ce2f7a569e1673617102cc5254a75e2", - "type": "github" - }, - "original": { - "owner": "dimensionalOS", - "ref": "main", - "repo": "dimos-lcm", - "type": "github" - } - }, "dimos-native-cpp": { "inputs": { "flake-utils": "flake-utils", @@ -28,11 +11,11 @@ }, "locked": { "dir": "native/cpp", - "lastModified": 1792432073, - "narHash": "sha256-5BgFoVp3/X5bKasI5G+f41sqVGAnZ4VPi1w3eC6+DIM=", + "lastModified": 1792433568, + "narHash": "sha256-3O7ECiGlq/HThKj65pi3rzjg0ny6iMGAuKXXEwt4GRc=", "owner": "dimensionalOS", "repo": "dimos", - "rev": "ea1d296b163d325d2dea782c4057dbaa3b697fc4", + "rev": "4c0afa2cc34247e537bbb57ffb4a5bc6d7007555", "type": "github" }, "original": { @@ -190,7 +173,6 @@ }, "root": { "inputs": { - "dimos-lcm": "dimos-lcm", "dimos-native-cpp": "dimos-native-cpp", "flake-utils": [ "dimos-native-cpp",