yiliao2026/src/planner/scripts/odom_map_tf.py

#!/usr/bin/env python3
"""Align map to wheel odom using a user-provided initial pose.

This node replaces AMCL for controlled short-range tests. It publishes a fixed
map->odom transform computed from:

  known current map pose * inverse(current odom pose)

After initialization, pose updates come only from /odom.
"""

import math

import rclpy
from geometry_msgs.msg import PoseWithCovarianceStamped, TransformStamped
from nav_msgs.msg import Odometry
from rclpy.node import Node
from std_msgs.msg import String
from tf2_ros import TransformBroadcaster


def yaw_from_quaternion(q):
    return math.atan2(
        2.0 * (q.w * q.z + q.x * q.y),
        1.0 - 2.0 * (q.y * q.y + q.z * q.z),
    )


def normalize_angle(angle):
    return math.atan2(math.sin(angle), math.cos(angle))


def fill_yaw_quaternion(q, yaw):
    q.x = 0.0
    q.y = 0.0
    q.z = math.sin(0.5 * yaw)
    q.w = math.cos(0.5 * yaw)


class OdomMapTf(Node):
    def __init__(self):
        super().__init__("planner_odom_map_tf")
        self.declare_parameter("map_frame", "map")
        self.declare_parameter("odom_frame", "odom")
        self.declare_parameter("odom_topic", "/odom")
        self.declare_parameter("initial_pose_topic", "/initialpose")
        self.declare_parameter("status_topic", "/odom_localization_status")
        self.declare_parameter("initial_x", 0.0)
        self.declare_parameter("initial_y", 0.0)
        self.declare_parameter("initial_yaw", 0.0)
        self.declare_parameter("initialize_on_first_odom", True)
        self.declare_parameter("publish_rate", 30.0)
        self.declare_parameter("state_log_period", 1.0)

        self.map_frame = str(self.get_parameter("map_frame").value)
        self.default_odom_frame = str(self.get_parameter("odom_frame").value)
        self.active_odom_frame = self.default_odom_frame
        self.initialize_on_first_odom = bool(
            self.get_parameter("initialize_on_first_odom").value
        )
        self.initial_pose = (
            float(self.get_parameter("initial_x").value),
            float(self.get_parameter("initial_y").value),
            float(self.get_parameter("initial_yaw").value),
        )
        self.state_log_period = max(
            0.0, float(self.get_parameter("state_log_period").value)
        )

        self.latest_odom_pose = None
        self.pending_initial_pose = None
        self.have_transform = False
        self.last_state_log_ns = 0
        self.map_to_odom_x = 0.0
        self.map_to_odom_y = 0.0
        self.map_to_odom_yaw = 0.0

        self.tf_broadcaster = TransformBroadcaster(self)
        self.status_pub = self.create_publisher(
            String, str(self.get_parameter("status_topic").value), 10
        )
        self.odom_sub = self.create_subscription(
            Odometry,
            str(self.get_parameter("odom_topic").value),
            self._odom_callback,
            20,
        )
        self.initial_pose_sub = self.create_subscription(
            PoseWithCovarianceStamped,
            str(self.get_parameter("initial_pose_topic").value),
            self._initial_pose_callback,
            10,
        )

        publish_rate = max(1.0, float(self.get_parameter("publish_rate").value))
        self.timer = self.create_timer(1.0 / publish_rate, self._publish_transform)
        self._publish_status(
            "waiting_for_odom initial_pose=(%.3f, %.3f, %.3f)"
            % self.initial_pose
        )

    def _odom_callback(self, msg):
        source_frame = msg.header.frame_id or self.default_odom_frame
        self.active_odom_frame = source_frame
        self.latest_odom_pose = (
            msg.pose.pose.position.x,
            msg.pose.pose.position.y,
            yaw_from_quaternion(msg.pose.pose.orientation),
        )

        if self.pending_initial_pose is not None:
            pending = self.pending_initial_pose
            self.pending_initial_pose = None
            self._initialize_from_pose(*pending, reason="initialpose")
        elif not self.have_transform and self.initialize_on_first_odom:
            self._initialize_from_pose(*self.initial_pose, reason="launch_initial")

    def _initial_pose_callback(self, msg):
        frame_id = msg.header.frame_id or self.map_frame
        if frame_id != self.map_frame:
            self._publish_status(
                "initialpose_rejected frame=%s expected=%s"
                % (frame_id, self.map_frame)
            )
            return

        pose = msg.pose.pose
        initial_pose = (
            pose.position.x,
            pose.position.y,
            yaw_from_quaternion(pose.orientation),
        )
        if self.latest_odom_pose is None:
            self.pending_initial_pose = initial_pose
            self._publish_status("initialpose_queued waiting_for_odom")
            return
        self._initialize_from_pose(*initial_pose, reason="initialpose")

    def _initialize_from_pose(self, map_x, map_y, map_yaw, reason):
        if self.latest_odom_pose is None:
            self.pending_initial_pose = (map_x, map_y, map_yaw)
            return

        odom_x, odom_y, odom_yaw = self.latest_odom_pose
        delta_yaw = normalize_angle(map_yaw - odom_yaw)
        cos_yaw = math.cos(delta_yaw)
        sin_yaw = math.sin(delta_yaw)
        self.map_to_odom_x = map_x - (cos_yaw * odom_x - sin_yaw * odom_y)
        self.map_to_odom_y = map_y - (sin_yaw * odom_x + cos_yaw * odom_y)
        self.map_to_odom_yaw = delta_yaw
        self.have_transform = True
        self._publish_status(
            "%s_aligned %s->%s x=%.3f y=%.3f yaw=%.3f"
            % (
                reason,
                self.map_frame,
                self.active_odom_frame,
                self.map_to_odom_x,
                self.map_to_odom_y,
                self.map_to_odom_yaw,
            )
        )

    def _publish_transform(self):
        if not self.have_transform:
            self._maybe_log_state()
            return
        msg = TransformStamped()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.header.frame_id = self.map_frame
        msg.child_frame_id = self.active_odom_frame
        msg.transform.translation.x = self.map_to_odom_x
        msg.transform.translation.y = self.map_to_odom_y
        msg.transform.translation.z = 0.0
        fill_yaw_quaternion(msg.transform.rotation, self.map_to_odom_yaw)
        self.tf_broadcaster.sendTransform(msg)
        self._maybe_log_state()

    def _maybe_log_state(self):
        if self.state_log_period <= 0.0:
            return
        now_ns = self.get_clock().now().nanoseconds
        min_period_ns = int(self.state_log_period * 1_000_000_000)
        if now_ns - self.last_state_log_ns < min_period_ns:
            return
        self.last_state_log_ns = now_ns

        if self.latest_odom_pose is None:
            self.get_logger().info(
                "odom_tf_state waiting_for_odom "
                "initial_pose=(%.3f,%.3f,%.3f)"
                % self.initial_pose
            )
            return

        odom_x, odom_y, odom_yaw = self.latest_odom_pose
        if not self.have_transform:
            self.get_logger().info(
                "odom_tf_state waiting_for_alignment "
                "odom=(%.3f,%.3f,%.3f) pending_initialpose=%s"
                % (odom_x, odom_y, odom_yaw, self.pending_initial_pose is not None)
            )
            return

        cos_yaw = math.cos(self.map_to_odom_yaw)
        sin_yaw = math.sin(self.map_to_odom_yaw)
        map_x = self.map_to_odom_x + cos_yaw * odom_x - sin_yaw * odom_y
        map_y = self.map_to_odom_y + sin_yaw * odom_x + cos_yaw * odom_y
        map_yaw = normalize_angle(self.map_to_odom_yaw + odom_yaw)
        self.get_logger().info(
            "odom_tf_state aligned=%s odom=(%.3f,%.3f,%.3f) "
            "map_pose=(%.3f,%.3f,%.3f) map_to_odom=(%.3f,%.3f,%.3f)"
            % (
                self.have_transform,
                odom_x,
                odom_y,
                odom_yaw,
                map_x,
                map_y,
                map_yaw,
                self.map_to_odom_x,
                self.map_to_odom_y,
                self.map_to_odom_yaw,
            )
        )

    def _publish_status(self, text):
        msg = String()
        msg.data = text
        self.status_pub.publish(msg)
        self.get_logger().info(text)


def main(args=None):
    rclpy.init(args=args)
    node = None
    try:
        node = OdomMapTf()
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        if node is not None:
            node.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()