p44_controller_node.py

#!/usr/bin/env python3
"""
ROS2 Controller Node - Runs AVLite controller and publishes control commands.

Worker process: bootstraps controller from profile, subscribes localization/trajectory,
runs control on a ROS timer, publishes commands.
"""
import argparse
import json
import logging
import os
import sys

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

from avlite.c10_perception.c11_perception_model import EgoState
from avlite import TrajectoryTracker
from avlite.c50_common.c54_trajectory_tracker import trajectory_path_fingerprint
from avlite.c30_control.c32_control_strategy import ControlStrategy
from avlite.c40_execution.c49_settings import ExecutionSettings

from .p46_autoware_converters import (
    AUTOWARE_AVAILABLE,
    ego_state_from_kinematic_state,
    trajectory_from_autoware,
    control_to_vehicle_cmd,
)
from avlite.c50_common.c56_fps_tracker import FpsTracker
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 VehicleControlCommand


class ControllerNode(Node):
    """ROS2 node that runs AVLite controller and publishes control commands."""

    def __init__(
        self,
        controller: ControlStrategy,
        ego_state: EgoState,
        control_dt: float | None = None,
    ):
        super().__init__('avlite_controller')

        self.settings = PluginSettings
        self.controller = controller
        self.ego_state = ego_state
        self.current_trajectory: TrajectoryTracker | None = None
        self.use_autoware = AUTOWARE_AVAILABLE and self.settings.use_autoware_msgs
        self.call_control = True

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

        control_dt = control_dt if control_dt is not None else ExecutionSettings.c40_control_dt
        self.declare_parameter('control_dt', control_dt)
        control_dt = self.get_parameter('control_dt').get_parameter_value().double_value
        self._node_period = control_dt
        pace_control = bool(getattr(self.settings, "pace_control", True))

        self._setup_subscriptions()
        self.ctrl_pub = self._create_control_publisher()
        timer_period = control_dt if pace_control else 0.001
        self.timer = self.create_timer(timer_period, self._control_tick)

        self.get_logger().info(
            f"ControllerNode worker started ({1.0/timer_period:.1f} Hz, pace_control={pace_control})"
        )

    def _create_control_publisher(self):
        if self.use_autoware:
            return self.create_publisher(
                VehicleControlCommand,
                self.settings.control_out_topic,
                10,
            )
        return self.create_publisher(String, self.settings.control_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(
                AutowareTrajectory,
                self.settings.trajectory_out_topic,
                self._on_trajectory,
                10,
            )
        else:
            self.create_subscription(
                String,
                self.settings.localization_topic,
                self._on_localization_json,
                10,
            )
            self.create_subscription(
                String,
                self.settings.trajectory_out_topic,
                self._on_trajectory_json,
                10,
            )
        self.create_subscription(
            Bool,
            self.settings.call_control_topic,
            self._on_call_control,
            10,
        )

    def _on_call_control(self, msg: Bool):
        self.call_control = 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 _align_trajectory_to_ego(self, traj: TrajectoryTracker | None) -> TrajectoryTracker | None:
        if traj is not None and traj.path:
            min_wp = traj.current_wp
            traj.update_waypoint_by_xy_forward(self.ego_state.x, self.ego_state.y, min_wp=min_wp)
        return traj

    def _apply_trajectory_if_changed(self, traj: TrajectoryTracker | None) -> None:
        if traj is None or not traj.path:
            return
        fingerprint = trajectory_path_fingerprint(traj)
        if fingerprint == self._last_trajectory_fingerprint:
            return
        self._last_trajectory_fingerprint = fingerprint
        self.current_trajectory = traj
        if self.controller:
            self.controller.set_trajectory_tracker(traj)

    def _on_trajectory(self, msg: 'AutowareTrajectory'):
        aligned = self._align_trajectory_to_ego(trajectory_from_autoware(msg))
        self._apply_trajectory_if_changed(aligned)

    def _on_trajectory_json(self, msg: String):
        try:
            data = json.loads(msg.data)
            path = [(p[0], p[1]) for p in data.get('path', [])]
            velocity = data.get('velocity', [])
            aligned = self._align_trajectory_to_ego(
                TrajectoryTracker(path=path, velocity=velocity)
            )
            self._apply_trajectory_if_changed(aligned)
        except json.JSONDecodeError:
            pass

    def _control_tick(self):
        if self._shutdown or not rclpy.ok() or not self.call_control or self.controller is None:
            return

        traj = self.controller.tj if hasattr(self.controller, 'tj') else self.current_trajectory
        if traj is None:
            return
        traj = self._align_trajectory_to_ego(traj)
        self.controller.set_trajectory_tracker(traj)

        try:
            cmd = self.controller.control(self.ego_state)
            if not cmd:
                return
            if self.use_autoware:
                msg = control_to_vehicle_cmd(cmd)
                msg.stamp = self.get_clock().now().to_msg()
            else:
                msg = String()
                msg.data = json.dumps({'steer': cmd.steer, 'acceleration': cmd.acceleration})
            self.ctrl_pub.publish(msg)
            self._fps_tracker.tick()
        except (rclpy.exceptions.InvalidHandle, RuntimeError):
            pass
        except Exception as e:
            if not self._shutdown:
                self.get_logger().error(f"Control failed: {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 controller ROS worker")
    parser.add_argument(
        "--profile",
        default=os.environ.get("AVLITE_PROFILE", "default"),
    )
    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("controller", profile=parsed.profile)
    if boot.controller is None:
        log.error("Controller bootstrap failed")
        sys.exit(1)

    rclpy.init(args=None)
    node = ControllerNode(controller=boot.controller, ego_state=boot.ego_state)
    attach_worker_logging(node, "avlite.c30_control")
    try:
        spin_node(node)
    finally:
        node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()


if __name__ == '__main__':
    main()