08_cortex_agent.py

"""Cortex agent with vision, visual Q&A, speech, and a custom action.

This example demonstrates a Cortex-based agent that can:
- Take pictures using the Vision component's ``take_picture`` action
- Describe what it sees using the VLM component's ``describe`` action
- Speak using the TextToSpeech component's ``say`` action,
  or automatically speak VLM output routed via the text_in topic
- Toggle an LED via a custom action defined in this script

The Cortex component acts as the system monitor and task planner. Send it a
goal like "describe what you see" and it will inspect the available
components, plan the steps, and execute them in sequence. The VLM's text
output is routed to the TTS component so descriptions are spoken aloud.

Usage:
    python3 examples/cortex_agent.py

    # In another terminal, send a goal:
    ros2 action send_goal /cortex_<process_id>/vision_language_action automatika_embodied_agents/action/VisionLanguageAction "{task: 'describe what you see'}"
"""
import os
import numpy as np

from agents.components import Vision, VLM, TextToSpeech, Cortex
from agents.config import VisionConfig, TextToSpeechConfig, CortexConfig
from agents.models import OllamaModel
from agents.clients import OllamaClient
from agents.ros import Topic, Action, Launcher

# from agents.models import GenericLLM, GenericMLLM
# from agents.clients import GenericHTTPClient

from endpoints import OLLAMA_HOST
# from endpoints import VLM_BASE_URL, VLM_API_KEY, VLM_CHECKPOINT
from kompass.robot import (
    AngularCtrlLimits,
    LinearCtrlLimits,
    RobotGeometry,
    RobotType,
    RobotConfig,
    RobotFrames,
)
from kompass.control import ControllersID, MapConfig

from kompass.components import (
    Controller,
    DriveManager,
    DriveManagerConfig,
    Planner,
    PlannerConfig,
    LocalMapper,
    LocalMapperConfig,
    MapServer,
    MapServerConfig,
    TopicsKeys,
)

# -- Model clients --
# Planner (Cortex) and VLM both served by a local Ollama (use a multimodal,
# tool-calling model so Cortex can plan and the VLM can see).
planner_model = OllamaModel(name="planner", checkpoint="qwen3.5:latest")
planner_client = OllamaClient(planner_model, host=OLLAMA_HOST)

vlm_model = OllamaModel(name="vlm", checkpoint="qwen3.5:latest")
vlm_client = OllamaClient(vlm_model, host=OLLAMA_HOST)

# Workshop endpoint (offline). The Modal vLLM clients these replaced:
# planner_model = GenericLLM(name="planner", checkpoint=VLM_CHECKPOINT)
# planner_client = GenericHTTPClient(
#     planner_model, host=VLM_BASE_URL, api_key=VLM_API_KEY, inference_timeout=300
# )
# vlm_model = GenericMLLM(name="vlm", checkpoint=VLM_CHECKPOINT)
# vlm_client = GenericHTTPClient(vlm_model, host=VLM_BASE_URL, api_key=VLM_API_KEY)

# -- Vision component (RoboML client, with tracking output) --
image_in = Topic(name="/image_raw", msg_type="Image")
# Camera info is required for controller vision following
image_in_info = Topic(name="/camera_info", msg_type="CameraInfo")
detections_out = Topic(name="detections", msg_type="Detections")
trackings_out = Topic(name="trackings", msg_type="Trackings")

vision = Vision(
    inputs=[image_in],
    outputs=[detections_out, trackings_out],
    config=VisionConfig(threshold=0.5, enable_local_classifier=True),
    trigger=0.5,
    component_name="vision",
)

# -- VLM component (describe action, output routed to TTS) --
vlm_query = Topic(name="vlm_query", msg_type="String")
vlm_output = Topic(name="text_in", msg_type="String")  # same topic as TTS input

vlm = VLM(
    inputs=[vlm_query, image_in],
    outputs=[vlm_output],
    model_client=vlm_client,
    trigger=vlm_query,
    component_name="vlm",
)

# -- Text-to-Speech component (local model, triggered by VLM output) --
audio_out = Topic(name="audio_out", msg_type="Audio")
tts = TextToSpeech(
    inputs=[vlm_output],
    outputs=[audio_out],
    config=TextToSpeechConfig(enable_local_model=True, play_on_device=False),
    trigger=vlm_output,
    component_name="tts",
)


# -- Custom action: toggle an LED --
led_on = False


def toggle_led():
    """Toggle an LED on the robot."""
    global led_on
    led_on = not led_on
    state = "ON" if led_on else "OFF"
    print(f"LED toggled {state}")


# -- Cortex: the planner / monitor --
cortex = Cortex(
    actions=[
        Action(method=toggle_led, description="Toggle the robot's LED on or off."),
    ],
    model_client=planner_client,
    config=CortexConfig(max_planning_steps=5, max_execution_steps=10),
    component_name="cortex",
)

# --- Kompass ---
# Setup your robot configuration
my_robot = RobotConfig(
    model_type=RobotType.DIFFERENTIAL_DRIVE,
    geometry_type=RobotGeometry.Type.CYLINDER,
    geometry_params=np.array([0.07, 0.3]),
    ctrl_vx_limits=LinearCtrlLimits(max_vel=0.4, max_acc=1.5, max_decel=2.5),
    ctrl_omega_limits=AngularCtrlLimits(
        max_vel=0.4, max_acc=2.0, max_decel=2.0, max_steer=np.pi / 3
    ),
)

# Configure the Global Planner
planner_config = PlannerConfig(loop_rate=1.0)
planner = Planner(component_name="planner", config=planner_config)
planner.run_type = "ActionServer"


# Configure the motion controller
controller = Controller(component_name="controller")
controller.algorithm = ControllersID.DWA
controller.direct_sensor = (
    False  # Get local perception from a "map" instead (from the local mapper)
)
controller.inputs(vision_detections=trackings_out, depth_camera_info=image_in_info)

# Configure the Drive Manager (Direct commands sending to robot)
driver_config = DriveManagerConfig(
    critical_zone_distance=0.05,
    critical_zone_angle=90.0,
    slowdown_zone_distance=0.3,
)
driver = DriveManager(component_name="drive_manager", config=driver_config)
# Publish Twist or TwistStamped from the DriveManager based on the distribution
if "ROS_DISTRO" in os.environ and (
    os.environ["ROS_DISTRO"] in ["rolling", "jazzy", "kilted"]
):
    cmd_msg_type: str = "TwistStamped"
else:
    cmd_msg_type = "Twist"

driver.outputs(robot_command=Topic(name="/cmd_vel", msg_type=cmd_msg_type))
driver.config.disable_safety_stop = True

# Configure a Local Mapper
local_mapper_config = LocalMapperConfig(
    map_params=MapConfig(width=3.0, height=3.0, resolution=0.05)
)
local_mapper = LocalMapper(component_name="mapper", config=local_mapper_config)

# Configure the global Map Server
map_file = "turtlebot3_webots.yaml"
map_server_config = MapServerConfig(
    loop_rate=1.0,
    map_file_path=map_file,  # Path to a 2D map yaml file or a point cloud file
    grid_resolution=0.5,
    pc_publish_row=False,
)
map_server = MapServer(component_name="global_map_server", config=map_server_config)


# -- Launch --
launcher = Launcher()

# Add Kompass components
launcher.add_pkg(
    components=[map_server, controller, planner, driver, local_mapper],
    multiprocessing=True,
    package_name="kompass"
)

launcher.add_pkg(
    components=[vision, vlm, cortex],
    multiprocessing=True,
    package_name="automatika_embodied_agents",
)

# Get odom from localizer filtered odom for all components
odom_topic = Topic(name="/odometry/filtered", msg_type="Odometry")
launcher.inputs(location=odom_topic)
launcher.robot = my_robot
launcher.frames = RobotFrames(
    world="map", odom="map", scan="LDS-01", rgb="camera_link", depth="camera_link"
)

launcher.enable_ui(
    inputs=[cortex.ui_main_action_input],
    outputs=[
        audio_out,
        vlm_output,
        map_server.get_out_topic(TopicsKeys.GLOBAL_MAP),
        odom_topic,
        planner.get_out_topic(TopicsKeys.GLOBAL_PLAN),
    ],
)


launcher.bringup()