p43_planner_node.py

#!/usr/bin/env python3
"""
ROS2 Planner Node - Runs AVLite local planner and publishes trajectory.

Worker process: bootstraps planner from profile, subscribes localization/perception,
runs step+replan on a ROS timer, publishes trajectory.
"""
import argparse
import json
import logging
import os
import sys
from typing import Optional

import rclpy
from rclpy.node import Node
from std_msgs.msg import String, Header, Bool, Empty

from avlite import ExecutionSettings, LocalPlanningStrategy, PerceptionStrategy, TrajectoryTracker
from avlite.c10_perception.c11_perception_model import EgoState
from avlite.c50_common.c54_trajectory_tracker import trajectory_path_fingerprint
from avlite.c50_common.c56_fps_tracker import FpsTracker

from .p46_autoware_converters import (
    AUTOWARE_AVAILABLE,
    ego_state_from_kinematic_state,
    trajectory_to_autoware,
    agents_from_bounding_boxes,
)
from .spawn_commands import agents_from_perception_json
from .global_plan_commands import parse_global_plan, GLOBAL_PLAN_QOS
from .settings import PluginSettings

log = logging.getLogger(__name__)

if AUTOWARE_AVAILABLE:
    from autoware_auto_msgs.msg import VehicleKinematicState
    from autoware_auto_msgs.msg import Trajectory as AutowareTrajectory
    from autoware_auto_msgs.msg import BoundingBoxArray


class PlannerNode(Node):
    """ROS2 node that runs AVLite local planner and publishes trajectory."""

    def __init__(
        self,
        planner: LocalPlanningStrategy,
        ego_state: EgoState,
        replan_dt: Optional[float] = None,
        perception: Optional[PerceptionStrategy] = None,
    ):
        super().__init__('avlite_planner')

        self.settings = PluginSettings
        self.planner = planner
        self.ego_state = ego_state
        self.perception = perception
        self.pm = getattr(planner, "env", None) or getattr(planner, "pm", None)
        self.use_autoware = AUTOWARE_AVAILABLE and self.settings.use_autoware_msgs
        self.call_replan = True

        self._fps_tracker = FpsTracker()
        self._shutdown = False
        self._last_published_fingerprint: tuple | None = None

        replan_dt = replan_dt if replan_dt is not None else ExecutionSettings.c40_replan_dt
        self.declare_parameter('replan_dt', replan_dt)
        replan_dt = self.get_parameter('replan_dt').get_parameter_value().double_value
        self._node_period = replan_dt
        pace_replan = bool(getattr(self.settings, "pace_replan", True))

        self._setup_subscriptions()
        self.traj_pub = self._create_trajectory_publisher()
        self.replan_tick_pub = self.create_publisher(
            Empty,
            self.settings.replan_tick_topic,
            10,
        )
        timer_period = replan_dt if pace_replan else 0.001
        self.timer = self.create_timer(timer_period, self._plan_tick)

        self.get_logger().info(
            f"PlannerNode worker started (autoware={self.use_autoware}, "
            f"{1.0/timer_period:.1f} Hz, pace_replan={pace_replan})"
        )

    def _create_trajectory_publisher(self):
        if self.use_autoware:
            return self.create_publisher(
                AutowareTrajectory,
                self.settings.trajectory_out_topic,
                10,
            )
        return self.create_publisher(String, self.settings.trajectory_out_topic, 10)

    def _setup_subscriptions(self):
        if self.use_autoware:
            self.create_subscription(
                VehicleKinematicState,
                self.settings.localization_topic,
                self._on_localization,
                10,
            )
            self.create_subscription(
                BoundingBoxArray,
                self.settings.perception_topic,
                self._on_perception,
                10,
            )
        else:
            self.create_subscription(
                String,
                self.settings.localization_topic,
                self._on_localization_json,
                10,
            )
            self.create_subscription(
                String,
                self.settings.perception_topic,
                self._on_perception_json,
                10,
            )
        self.create_subscription(
            Bool,
            self.settings.call_replan_topic,
            self._on_call_replan,
            10,
        )
        self.create_subscription(
            String,
            self.settings.global_plan_topic,
            self._on_global_plan,
            GLOBAL_PLAN_QOS,
        )

    def _on_global_plan(self, msg: String) -> None:
        if self.planner is None:
            return
        try:
            global_plan, ego_xy = parse_global_plan(msg.data)
        except (json.JSONDecodeError, KeyError, IndexError, TypeError, ValueError) as e:
            self.get_logger().error(f"Invalid global plan JSON: {e}")
            return
        try:
            if ego_xy is None:
                ego_xy = (self.ego_state.x, self.ego_state.y)
            self.planner.set_global_plan(global_plan, ego_xy=ego_xy)
            if hasattr(self.planner, "last_plan"):
                self.planner.last_plan = None
            self._last_published_fingerprint = None
            self.get_logger().info(
                f"Global plan updated from main process ({len(global_plan.path)} path points)"
            )
        except Exception as e:
            self.get_logger().error(f"Failed to apply global plan: {e}")

    def _on_call_replan(self, msg: Bool):
        self.call_replan = bool(msg.data)

    def _on_localization(self, msg: 'VehicleKinematicState'):
        ego_state_from_kinematic_state(msg, self.ego_state)

    def _on_localization_json(self, msg: String):
        try:
            data = json.loads(msg.data)
            self.ego_state.x = data.get('x', self.ego_state.x)
            self.ego_state.y = data.get('y', self.ego_state.y)
            self.ego_state.theta = data.get('theta', self.ego_state.theta)
            self.ego_state.velocity = data.get('velocity', self.ego_state.velocity)
        except json.JSONDecodeError:
            pass

    def _on_perception(self, msg: 'BoundingBoxArray'):
        if self.pm is not None:
            self.pm.agent_vehicles = agents_from_bounding_boxes(msg)

    def _on_perception_json(self, msg: String):
        if self.pm is None:
            return
        try:
            data = json.loads(msg.data)
            self.pm.agent_vehicles = agents_from_perception_json(data)
        except json.JSONDecodeError:
            pass

    def _run_planning(self) -> Optional[object]:
        if self.planner is None:
            return None
        self.planner.step(self.ego_state)
        if self.call_replan:
            self.planner.replan()
        return self.planner.get_local_plan()

    def _publish_trajectory(self, trajectory) -> None:
        if not trajectory.path:
            log.debug("Skipping trajectory publish: empty path")
            return

        tj = trajectory.as_trajectory() if hasattr(trajectory, "as_trajectory") else trajectory
        if tj is not None and tj.path:
            tj.update_waypoint_by_xy_forward(self.ego_state.x, self.ego_state.y)
            trajectory = tj

        if self.use_autoware:
            header = Header()
            header.stamp = self.get_clock().now().to_msg()
            header.frame_id = self.settings.map_frame
            msg = trajectory_to_autoware(trajectory, header)
        else:
            msg = String()
            msg.data = json.dumps({
                'path': [(float(p[0]), float(p[1])) for p in trajectory.path],
                'velocity': [float(v) for v in trajectory.velocity] if trajectory.velocity else [],
            })
        self.traj_pub.publish(msg)

    def _plan_tick(self):
        if self._shutdown or not rclpy.ok() or self.planner is None:
            return
        try:
            local_plan = self._run_planning()
            if local_plan is None:
                return
            trajectory = local_plan.as_trajectory()
            if trajectory is None or not trajectory.path:
                log.debug("Skipping plan tick: no trajectory path")
                return

            fingerprint = trajectory_path_fingerprint(trajectory)
            if fingerprint != self._last_published_fingerprint:
                self._publish_trajectory(trajectory)
                self._last_published_fingerprint = fingerprint
            self.replan_tick_pub.publish(Empty())
            self._fps_tracker.tick()
        except (rclpy.exceptions.InvalidHandle, RuntimeError):
            pass
        except Exception as e:
            if not self._shutdown:
                planner_name = type(self.planner).__name__ if self.planner else "unknown"
                self.get_logger().error(f"Planning failed ({planner_name}): {e}")

    def destroy_node(self):
        self._shutdown = True
        if self.timer:
            self.timer.cancel()
        super().destroy_node()


def _parse_args(argv=None):
    parser = argparse.ArgumentParser(description="AVLite planner ROS worker")
    parser.add_argument(
        "--profile",
        default=os.environ.get("AVLITE_PROFILE", "default"),
        help="Settings profile name",
    )
    return parser.parse_args(argv)


def main(args=None):
    parsed = _parse_args(args)
    os.environ["AVLITE_PROFILE"] = parsed.profile

    from .p48_node_bootstrap import attach_worker_logging, bootstrap_role, spin_node

    logging.basicConfig(level=logging.WARNING)
    boot = bootstrap_role("planner", profile=parsed.profile)
    if boot.local_planner is None:
        log.error("Planner bootstrap failed")
        sys.exit(1)

    rclpy.init(args=None)
    node = PlannerNode(
        planner=boot.local_planner,
        ego_state=boot.ego_state,
        perception=boot.perception,
    )
    attach_worker_logging(node, "avlite.c20_planning")
    try:
        spin_node(node)
    finally:
        node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()


if __name__ == '__main__':
    main()