p45_world_node.py

#!/usr/bin/env python3
"""
ROS2 World Node - Runs world simulation in a dedicated process.

Subscribes to control commands, steps the simulator, publishes localization and perception.
"""
import argparse
import json
import logging
import os
import sys
import time

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

from avlite.c10_perception.c11_perception_model import EgoState, AgentState
from avlite.c40_execution.c41_world_bridge import (
    WorldBridge,
    is_world_capability_enabled,
    is_world_stack_capability_enabled,
)
from avlite.c30_control.c31_control_model import ControlCommand
from avlite.c40_execution.c49_settings import ExecutionSettings
from avlite.c50_common.c51_capabilities import StackCapability, WorldCapability

from .p46_autoware_converters import (
    AUTOWARE_AVAILABLE,
    ego_state_to_kinematic_state,
    control_from_vehicle_cmd,
)
from .spawn_commands import parse_spawn_command, parse_teleport_command
from .lidar_commands import encode_lidar_cloud
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 VehicleControlCommand
    from autoware_auto_msgs.msg import BoundingBoxArray, BoundingBox


class WorldNode(Node):
    """ROS2 node that runs world simulation in a worker process."""

    def __init__(self, world: WorldBridge, sim_dt: float | None = None, global_plan=None):
        super().__init__('avlite_world')

        self.settings = PluginSettings
        self.world = world
        self.global_plan = global_plan
        self.last_cmd: ControlCommand | None = None
        self.use_autoware = AUTOWARE_AVAILABLE and self.settings.use_autoware_msgs

        sim_dt = sim_dt if sim_dt is not None else ExecutionSettings.c40_sim_dt
        self.declare_parameter('sim_dt', sim_dt)
        self.sim_dt = self.get_parameter('sim_dt').get_parameter_value().double_value
        self.pace_sim = bool(getattr(self.settings, "pace_sim", True))

        self._tick_count = 0
        self._fps_update_time = time.time()
        self._elapsed_sim_time = 0.0
        self._last_sim_wall_t = None
        self._shutdown = False
        self.lidar_pub = None

        self._setup_subscribers()
        self._setup_publishers()
        timer_period = self.sim_dt if self.pace_sim else 0.001
        self.timer = self.create_timer(timer_period, self._sim_tick)

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

    def _setup_subscribers(self):
        """Setup control command subscriber."""
        if self.use_autoware:
            self.ctrl_sub = self.create_subscription(
                VehicleControlCommand,
                self.settings.control_out_topic,
                self._on_control_autoware,
                10
            )
        else:
            self.ctrl_sub = self.create_subscription(
                String,
                self.settings.control_out_topic,
                self._on_control_json,
                10
            )
        self.create_subscription(
            String,
            self.settings.spawn_agent_topic,
            self._on_spawn_agent,
            10,
        )
        self.create_subscription(
            Bool,
            self.settings.reset_agents_topic,
            self._on_reset_agents,
            10,
        )
        self.create_subscription(
            String,
            self.settings.teleport_ego_topic,
            self._on_teleport_ego,
            10,
        )
        self.create_subscription(
            String,
            self.settings.global_plan_topic,
            self._on_global_plan,
            GLOBAL_PLAN_QOS,
        )

    def _setup_publishers(self):
        """Setup ego state and ground-truth agent publishers."""
        if self.use_autoware:
            self.ego_pub = self.create_publisher(
                VehicleKinematicState,
                self.settings.localization_topic,
                10
            )
            self.gt_pub = self.create_publisher(
                BoundingBoxArray,
                self.settings.world_gt_topic,
                10
            )
        else:
            self.ego_pub = self.create_publisher(
                String,
                self.settings.localization_topic,
                10
            )
            self.gt_pub = self.create_publisher(
                String,
                self.settings.world_gt_topic,
                10
            )

    def _ensure_lidar_publisher(self) -> None:
        if self.lidar_pub is not None:
            return
        self.lidar_pub = self.create_publisher(
            String,
            self.settings.lidar_topic,
            10,
        )

    def _on_control_autoware(self, msg: 'VehicleControlCommand'):
        """Handle Autoware control command."""
        self.last_cmd = control_from_vehicle_cmd(msg)

    def _on_control_json(self, msg: String):
        """Handle JSON control command."""
        try:
            data = json.loads(msg.data)
            self.last_cmd = ControlCommand(
                steer=data.get('steer', 0),
                acceleration=data.get('acceleration', 0)
            )
        except json.JSONDecodeError as e:
            self.get_logger().error(f"Invalid JSON control: {e}")

    def _on_spawn_agent(self, msg: String):
        if self.world is None:
            return
        try:
            agent = parse_spawn_command(msg.data)
            if agent is not None:
                self.world.spawn_agent(agent, global_plan=self.global_plan)
        except json.JSONDecodeError as e:
            self.get_logger().error(f"Invalid JSON spawn command: {e}")

    def _on_global_plan(self, msg: String):
        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
        self.global_plan = global_plan

    def _on_reset_agents(self, msg: Bool):
        if self.world is None or not msg.data:
            return
        self.world.reset()

    def _on_teleport_ego(self, msg: String):
        if self.world is None:
            return
        try:
            parsed = parse_teleport_command(msg.data)
            if parsed is None:
                return
            x, y, theta = parsed
            self.world.teleport_ego(x, y, theta)
        except json.JSONDecodeError as e:
            self.get_logger().error(f"Invalid JSON teleport command: {e}")

    def _sim_tick(self):
        """Run one simulation step and publish results."""
        # Skip if shutting down or ROS context invalid
        if self._shutdown or not rclpy.ok():
            return
            
        if self.world is None:
            return

        now = time.time()
        if self.pace_sim:
            dt = self.sim_dt
            self._last_sim_wall_t = now
        else:
            if self._last_sim_wall_t is None:
                self._last_sim_wall_t = now
                dt = None
            else:
                dt = max(1e-4, min(now - self._last_sim_wall_t, 1.0))
                self._last_sim_wall_t = now

        if dt is not None:
            if self.last_cmd:
                self.world.control_ego_state(self.last_cmd, dt=dt)
            self._elapsed_sim_time += dt
            self._tick_count += 1

        ego_state = self.world.get_ego_state()
        self._publish_ego_state(ego_state)
        if is_world_stack_capability_enabled(StackCapability.DETECTION):
            self._publish_ground_truth()
        if is_world_capability_enabled(WorldCapability.LIDAR_2D):
            self._ensure_lidar_publisher()
            self._publish_lidar()

    def _publish_ego_state(self, ego_state: EgoState):
        """Publish ego vehicle state."""
        if ego_state is None:
            return
        
        try:
            if self.use_autoware:
                header = Header()
                header.stamp = self.get_clock().now().to_msg()
                header.frame_id = self.settings.map_frame
                msg = ego_state_to_kinematic_state(ego_state, header)
            else:
                msg = String()
                msg.data = json.dumps({
                    'x': float(ego_state.x),
                    'y': float(ego_state.y),
                    'theta': float(ego_state.theta),
                    'velocity': float(ego_state.velocity),
                })
            
            self.ego_pub.publish(msg)
        except (rclpy.exceptions.InvalidHandle, RuntimeError):
            # Suppress errors during shutdown (invalid handle or runtime errors)
            pass
        except Exception as e:
            if not self._shutdown:
                self.get_logger().error(f"Failed to publish ego state: {e}")

    def _publish_ground_truth(self):
        """Publish ground-truth agents for the perception worker."""
        try:
            # Get ground truth perception from world
            pm = self.world.get_ground_truth_perception_model()
            if pm is None:
                return
            
            agents = getattr(pm, 'agent_vehicles', []) or []
            
            if self.use_autoware:
                msg = BoundingBoxArray()
                msg.header.stamp = self.get_clock().now().to_msg()
                msg.header.frame_id = self.settings.map_frame
                
                for agent in agents:
                    box = BoundingBox()
                    box.centroid.x = float(agent.x)
                    box.centroid.y = float(agent.y)
                    box.centroid.z = 0.0
                    box.size.x = float(agent.length)
                    box.size.y = float(agent.width)
                    box.size.z = 1.5
                    box.heading = float(agent.theta)
                    box.velocity = float(agent.velocity)
                    box.vehicle_label = 1
                    msg.boxes.append(box)
            else:
                msg = String()
                objects_list = []
                for agent in agents:
                    objects_list.append({
                        'id': getattr(agent, 'agent_id', 0),
                        'x': float(agent.x),
                        'y': float(agent.y),
                        'theta': float(agent.theta),
                        'velocity': float(agent.velocity),
                        'length': float(agent.length),
                        'width': float(agent.width),
                    })
                msg.data = json.dumps({'objects': objects_list})
            
            self.gt_pub.publish(msg)
        except (rclpy.exceptions.InvalidHandle, RuntimeError):
            # Suppress errors during shutdown (invalid handle or runtime errors)
            pass
        except Exception as e:
            if not self._shutdown:
                self.get_logger().error(f"Failed to publish perception: {e}")

    def _publish_lidar(self):
        """Publish simulated LiDAR scan for the main-process mirror bridge."""
        try:
            cloud = self.world.get_lidar_data()
            msg = String()
            msg.data = encode_lidar_cloud(cloud)
            self.lidar_pub.publish(msg)
        except (rclpy.exceptions.InvalidHandle, RuntimeError):
            pass
        except Exception as e:
            if not self._shutdown:
                self.get_logger().error(f"Failed to publish lidar: {e}")

    def set_world(self, world: WorldBridge):
        """Set or update the world bridge."""
        self.world = world
        self.get_logger().info(f"World set: {world.__class__.__name__}")

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


def _parse_args(argv=None):
    parser = argparse.ArgumentParser(description="AVLite world 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("world", profile=parsed.profile)
    if boot.world is None:
        log.error("World bootstrap failed")
        sys.exit(1)

    rclpy.init(args=None)
    node = WorldNode(world=boot.world, global_plan=boot.global_plan)
    attach_worker_logging(node, "avlite.c40_execution")
    try:
        spin_node(node)
    finally:
        node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()


if __name__ == '__main__':
    main()