simulate_rl_c.py

from __future__ import annotations

import argparse
import ctypes
import os
import re
import subprocess
import time
from pathlib import Path
from typing import Any

os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")

import numpy as np

from quadcopter_animation import animation
from quadcopter_envs import Quadcopter3DGates, RANDOMIZATION


PROJECT_ROOT = Path(__file__).resolve().parent


def read_define(header: Path, name: str) -> int:
    pattern = re.compile(rf"^\s*#define\s+{re.escape(name)}\s+(\d+)\s*$")
    for line in header.read_text().splitlines():
        match = pattern.match(line)
        if match:
            return int(match.group(1))
    raise ValueError(f"Could not find #define {name} in {header}")


class CController:
    """Small ctypes wrapper around one generated C policy folder."""

    def __init__(self, model_dir: Path):
        self.model_dir = model_dir.resolve()
        header = self.model_dir / "nn_parameters.h"
        self.obs_size = read_define(header, "OBS_SIZE")
        self.action_size = read_define(header, "NUM_CONTROLS")
        self.library_path = self.model_dir / "libcontroller.so"
        self._build_shared_library()

        self.lib = ctypes.CDLL(str(self.library_path))
        float_ptr = ctypes.POINTER(ctypes.c_float)
        self.lib.nn_reset.argtypes = []
        self.lib.nn_reset.restype = None
        self.lib.nn_control.argtypes = [float_ptr, float_ptr]
        self.lib.nn_control.restype = None

        self._obs = np.zeros(self.obs_size, dtype=np.float32)
        self._action = np.zeros(self.action_size, dtype=np.float32)
        self.reset()

    def _build_shared_library(self) -> None:
        sources = [
            self.model_dir / "nn_operations.c",
            self.model_dir / "nn_parameters.c",
        ]
        command = [
            "gcc",
            "-std=c99",
            "-O2",
            "-Wall",
            "-Wextra",
            "-Wno-unused-function",
            "-fPIC",
            "-shared",
            *[str(source) for source in sources],
            "-lm",
            "-o",
            str(self.library_path),
        ]
        subprocess.run(command, check=True)

    def reset(self) -> None:
        self.lib.nn_reset()

    def predict(self, obs: np.ndarray) -> np.ndarray:
        flat_obs = np.asarray(obs, dtype=np.float32).reshape(-1)
        if flat_obs.size != self.obs_size:
            raise ValueError(
                f"C controller expects obs size {self.obs_size}, got {flat_obs.size}. "
                "Use observation flags matching the exported checkpoint."
            )
        self._obs[:] = flat_obs
        self.lib.nn_control(
            self._obs.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
            self._action.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
        )
        return self._action.copy()


def infer_observation_flags(model_dir: Path, obs_size: int) -> dict[str, bool]:
    """Infer env observation flags only when the exported obs size requires it."""
    name = model_dir.name.lower()
    flags = {"low_obs": False, "no_vel": False, "no_ang_vel": False}
    if obs_size == 17:
        if "no_ang_vel" in name:
            flags["no_ang_vel"] = True
        elif "no_vel" in name:
            flags["no_vel"] = True
    return flags


def make_env(args: argparse.Namespace, controller: CController) -> Quadcopter3DGates:
    flags = infer_observation_flags(controller.model_dir, controller.obs_size)
    if args.low_obs:
        flags = {"low_obs": True, "no_vel": False, "no_ang_vel": False}
    if args.no_vel:
        flags = {"low_obs": False, "no_vel": True, "no_ang_vel": False}
    if args.no_ang_vel:
        flags = {"low_obs": False, "no_vel": False, "no_ang_vel": True}

    env = Quadcopter3DGates(
        num_envs=1,
        randomization=RANDOMIZATION,
        gates_ahead=1,
        initialize_uniform=False,
        initialize_at_random_gates=True,
        seed=args.seed,
        low_obs=flags["low_obs"],
        no_vel=flags["no_vel"],
        no_ang_vel=flags["no_ang_vel"],
        param_input=args.param_input,
        param_input_noise=args.param_input_noise,
    )
    if env.observation_space.shape[0] != controller.obs_size:
        raise ValueError(
            f"Environment obs size {env.observation_space.shape[0]} does not match "
            f"C controller obs size {controller.obs_size}. "
            f"Flags used: {flags}."
        )
    return env


def append_if_present(values: list[float], value: Any) -> None:
    if value is not None:
        values.append(float(value))


def mean_or_nan(values: list[float]) -> float:
    return sum(values) / len(values) if values else float("nan")


def evaluate(args: argparse.Namespace) -> None:
    controller = CController(args.model_dir)
    env = make_env(args, controller)
    obs = env.reset()
    controller.reset()

    episode_reward = 0.0
    episode_gates = 0
    total_reward = 0.0
    total_gates = 0
    total_crashes = 0
    total_gate_crashes = 0
    total_time_limits = 0
    completed = 0

    time_sim_list: list[float] = []
    distance_list: list[float] = []
    max_speed_list: list[float] = []
    avg_speed_list: list[float] = []
    avg_bank_angle_list: list[float] = []
    action_std_list: list[float] = []
    avg_action_rpm_diff_list: list[float] = []
    inference_time_list: list[float] = []

    while completed < args.episodes:
        start = time.perf_counter()
        action = controller.predict(obs[0])
        inference_time_list.append(time.perf_counter() - start)

        obs, rewards, dones, infos = env.step(action.reshape(1, -1))
        episode_reward += float(rewards[0])
        if infos[0].get("gate_passed", False):
            episode_gates += 1

        if dones[0]:
            info = infos[0]
            crashed = bool(
                info.get("ground_collision", False)
                or info.get("out_of_bounds", False)
                or info.get("gate_collision", False)
            )
            gate_crashed = bool(info.get("gate_collision", False))
            append_if_present(time_sim_list, info.get("TimeLimit.time"))
            append_if_present(distance_list, info.get("distance_travelled"))
            append_if_present(max_speed_list, info.get("max_speed"))
            append_if_present(avg_speed_list, info.get("avg_speed"))
            append_if_present(avg_bank_angle_list, info.get("avg_bank_angle"))
            append_if_present(action_std_list, info.get("action_std"))
            append_if_present(avg_action_rpm_diff_list, info.get("avg_action_rpm_diff"))

            completed += 1
            total_reward += episode_reward
            total_gates += episode_gates
            total_crashes += int(crashed)
            total_gate_crashes += int(gate_crashed)
            total_time_limits += int(bool(info.get("TimeLimit.truncated", False)))

            if args.verbose:
                if gate_crashed:
                    outcome = "gate-crashed"
                elif crashed:
                    outcome = "crashed"
                elif info.get("TimeLimit.truncated", False):
                    outcome = "finished (time limit)"
                else:
                    outcome = "finished"
                print(
                    f"Episode {completed} | reward={episode_reward:.3f} | "
                    f"{outcome} | gates_passed={episode_gates}"
                )

            episode_reward = 0.0
            episode_gates = 0
            controller.reset()

    env.close()

    avg_reward = total_reward / completed
    avg_gates = total_gates / completed
    crash_rate = total_crashes / completed
    gate_crash_rate = total_gate_crashes / completed
    time_limit_rate = total_time_limits / completed
    finish_rate = 1.0 - crash_rate
    avg_inference_time_ms = 1000.0 * mean_or_nan(inference_time_list)

    print(
        f"\nC RL simulation over {completed} episodes | "
        f"model_dir={controller.model_dir.relative_to(PROJECT_ROOT)} | "
        f"avg_reward={avg_reward:.3f} | "
        f"avg_gates_passed={avg_gates:.3f} | "
        f"crash_rate={crash_rate:.3f} | "
        f"gate_crash_rate={gate_crash_rate:.3f} | "
        f"finish_rate={finish_rate:.3f} | "
        f"time_limit_rate={time_limit_rate:.3f} | "
        f"avg_time_sim={mean_or_nan(time_sim_list):.3f} | "
        f"avg_distance={mean_or_nan(distance_list):.3f} | "
        f"avg_max_speed={mean_or_nan(max_speed_list):.3f} | "
        f"avg_speed={mean_or_nan(avg_speed_list):.3f} | "
        f"avg_bank_angle={mean_or_nan(avg_bank_angle_list):.3f} | "
        f"avg_action_std={mean_or_nan(action_std_list):.3f} | "
        f"avg_action_rpm_diff={mean_or_nan(avg_action_rpm_diff_list):.3f} | "
        f"avg_c_inference_time_ms={avg_inference_time_ms:.6f}\n"
    )


def animate(args: argparse.Namespace) -> None:
    controller = CController(args.model_dir)
    env = make_env(args, controller)
    obs = env.reset()
    controller.reset()

    def run():
        nonlocal obs
        action = controller.predict(obs[0])
        obs, _, dones, _ = env.step(action.reshape(1, -1))
        if dones[0]:
            controller.reset()
        return env.render()

    animation.view(run, gate_pos=env.gate_pos, gate_yaw=env.gate_yaw, fps=1 / env.dt)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Simulate the quadcopter using an exported RL C controller.")
    parser.add_argument(
        "--model-dir",
        type=Path,
        default=Path("C_codes/recurrent_ppo_cfc_no_vel"),
        help="Folder containing nn_operations.c and nn_parameters.c.",
    )
    parser.add_argument("--episodes", type=int, default=20)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--animate", action="store_true")
    parser.add_argument("--verbose", action="store_true")
    parser.add_argument("--low-obs", action="store_true")
    parser.add_argument("--no-vel", action="store_true")
    parser.add_argument("--no-ang-vel", action="store_true")
    parser.add_argument("--param-input", action="store_true")
    parser.add_argument("--param-input-noise", type=float, default=0.0)
    return parser.parse_args()


if __name__ == "__main__":
    parsed_args = parse_args()
    if parsed_args.animate:
        animate(parsed_args)
    else:
        evaluate(parsed_args)