diff --git a/bluebird-gymnasium/bluebird_gymnasium/__init__.py b/bluebird-gymnasium/bluebird_gymnasium/__init__.py index e05dc49..b547975 100644 --- a/bluebird-gymnasium/bluebird_gymnasium/__init__.py +++ b/bluebird-gymnasium/bluebird_gymnasium/__init__.py @@ -18,3 +18,5 @@ register(id="SectorYEnv-v0", entry_point="bluebird_gymnasium.envs:SectorYEnv") register(id="SpringfieldEnv-v0", entry_point="bluebird_gymnasium.envs:SpringfieldEnv") + +register(id="FlightSchoolEnv-v0", entry_point="bluebird_gymnasium.envs:FlightSchoolEnv") diff --git a/bluebird-gymnasium/bluebird_gymnasium/envs/__init__.py b/bluebird-gymnasium/bluebird_gymnasium/envs/__init__.py index 0ce78ca..bace6b7 100644 --- a/bluebird-gymnasium/bluebird_gymnasium/envs/__init__.py +++ b/bluebird-gymnasium/bluebird_gymnasium/envs/__init__.py @@ -152,6 +152,7 @@ class EnvConfig: # aircraft scenario generator class(es) from bluebird_dt.scenario_manager import ( # noqa: E402 + Infinite, Regular, Tactical, TwoAircraft, @@ -159,6 +160,7 @@ class EnvConfig: from bluebird_dt.scenario_manager.scenario_manager import ScenarioManager # noqa: E402 SCENARIO_CLS: dict[str, ScenarioManager] = { + "infinite": Infinite, "regular": Regular, "tactical": Tactical, "twoaircraft": TwoAircraft, @@ -167,6 +169,7 @@ class EnvConfig: # now envs module imports from bluebird_gymnasium.envs.base import BaseEnv # noqa: E402 +from bluebird_gymnasium.envs.flight_school import FlightSchoolEnv # noqa: E402 from bluebird_gymnasium.envs.infinite import CustomInfiniteEnv, InfiniteEnv # noqa: E402 from bluebird_gymnasium.envs.sector_i import SectorIEnv # noqa: E402 from bluebird_gymnasium.envs.sector_x import SectorXEnv # noqa: E402 @@ -183,6 +186,7 @@ class EnvConfig: "SectorXPlusEnv-v0": SectorXPlusEnv, "SectorYEnv-v0": SectorYEnv, "SpringfieldEnv-v0": SpringfieldEnv, + "FlightSchoolEnv-v0": FlightSchoolEnv, } name_to_gym_key: dict[str, str] = { @@ -193,6 +197,7 @@ class EnvConfig: "SectorXPlusEnv-v0": "sector_xplus", "SectorYEnv-v0": "sector_y", "SpringfieldEnv-v0": "springfield", + "FlightSchoolEnv-v0": "flight_school", } available_names = ", ".join(name_to_gym_key.values()) @@ -238,6 +243,7 @@ def get_env_cls_and_config(env_name: str) -> tuple[type[BaseEnv], EnvConfig]: __all__ = [ "BaseEnv", "CustomInfiniteEnv", + "FlightSchoolEnv", "InfiniteEnv", "SectorIEnv", "SectorXEnv", diff --git a/bluebird-gymnasium/bluebird_gymnasium/envs/flight_school.py b/bluebird-gymnasium/bluebird_gymnasium/envs/flight_school.py new file mode 100644 index 0000000..4846798 --- /dev/null +++ b/bluebird-gymnasium/bluebird_gymnasium/envs/flight_school.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +import datetime + +from bluebird_dt.simulator import Simulator + +from bluebird_gymnasium.envs import EnvConfig, ViewType +from bluebird_gymnasium.envs.sector_xplus import SectorXPlusEnv + + +class FlightSchoolEnv(SectorXPlusEnv): + """Gymnasium environment for the Flight School scenario.""" + + def _generate_scenario(self) -> Simulator: + category = "Flight School" + scenario = "Xplus-Sector" + timestamp = datetime.datetime.now().strftime("%Y_%m_%d__%H_%M_%S") + + suffix = self.config.simulation_log_config.get("log_suffix", None) + suffix = "" if suffix is None or suffix == "" else f"__{suffix}" + log_filename = f"{category}_{scenario}_{timestamp}{suffix}" + + return self.scenario_manager.to_simulator( + category=category, + scenario_name=scenario, + save_log_to_file=False, + log_filename=log_filename, + predictor=None, + ) + + @classmethod + def get_default_env_config(cls, view_type: ViewType | str = ViewType.CENTRALIZED) -> EnvConfig: + config = super().get_default_env_config(view_type) + config.scenario_config = { + "cls": "infinite", + "args": { + "random_seed": None, + "num_starter_aircraft": 2, + "initial_spawn_rate": 0.002, + "spawn_rate_increment": 0.002, + "spawn_rate_increase_interval": 60, + "max_spawn_rate": 0.03, + "total_time_seconds": 3600.0, + }, + } + config.reward_config = { + "fns": [ + "position_status_const", + "lateral_centreline_distance_shaped", + "safety_simple_avoidance_exp", + ], + "coeffs": [1.0, 1.0, 1.2], + } + config.scenario_duration = 10 * 60 + return config diff --git a/bluebird-gymnasium/bluebird_gymnasium/rewards/__init__.py b/bluebird-gymnasium/bluebird_gymnasium/rewards/__init__.py index 9942991..c2558f9 100644 --- a/bluebird-gymnasium/bluebird_gymnasium/rewards/__init__.py +++ b/bluebird-gymnasium/bluebird_gymnasium/rewards/__init__.py @@ -98,3 +98,20 @@ mod_name = f"{base_pkg}.custom.reward_drlan" registry_reward_fn.register("reward_drlan", f"{mod_name}:reward_drlan") registry_reward_fn.register("custom_reward", f"{mod_name}:custom_reward_fn") +mod_name = f"{base_pkg}.custom.custom_reward" +registry_reward_fn.register( + "route_progress_terminal_reward", + f"{mod_name}:route_progress_terminal_reward", +) +registry_reward_fn.register( + "lateral_termination_check_sac_env", + f"{mod_name}:lateral_termination_check_sac_env", +) +registry_reward_fn.register( + "lateral_termination_check_mac_env", + f"{mod_name}:lateral_termination_check_mac_env", +) +registry_reward_fn.register( + "anti_loiter_route_rejoin_reward", + f"{mod_name}:anti_loiter_route_rejoin_reward", +) diff --git a/bluebird-gymnasium/bluebird_gymnasium/rewards/custom/custom_reward.py b/bluebird-gymnasium/bluebird_gymnasium/rewards/custom/custom_reward.py index a1179d4..c95cf8b 100644 --- a/bluebird-gymnasium/bluebird_gymnasium/rewards/custom/custom_reward.py +++ b/bluebird-gymnasium/bluebird_gymnasium/rewards/custom/custom_reward.py @@ -1,4 +1,11 @@ +from typing import Any + from bluebird_gymnasium.envs.base import BaseEnv +from bluebird_gymnasium.rewards.lateral_termination_check import ( + lateral_termination_check_mac, + lateral_termination_check_sac, +) +from bluebird_gymnasium.utils.types import PositionStatus def custom_reward_fn(gym_env: BaseEnv, callsign: str, action: int, **kwargs) -> float: # noqa: ARG001, ANN003 @@ -16,3 +23,151 @@ def custom_reward_fn(gym_env: BaseEnv, callsign: str, action: int, **kwargs) -> # implement reward function here. return 0.0 + + +def route_progress_terminal_reward(gym_env: BaseEnv, callsign: str, action: int, **kwargs) -> float: # noqa: ARG001, ANN003 + """Reward progress toward exit and penalize stalling at episode end. + + This reward is intended for the PPO curriculum notebooks where the agent + was observed to loiter near its entry fix. It combines three signals: + + - a positive terminal bonus for correctly reaching the exit window + - a per-step progress reward for reducing route-track distance to exit + - a timeout penalty if the episode ends with the aircraft still far away + + Returns: + float: shaped reward that encourages forward progress through sector. + """ + + reward = 0.0 + + ac_tracked_state = gym_env.get_tracked_aircraft_data(callsign) + prev_ac_tracked_state = gym_env.get_tracked_aircraft_data_previous(callsign) + + if ac_tracked_state is None: + return reward + + if ac_tracked_state.pos_status == PositionStatus.EXIT_REACHED: + reward += 20.0 + + if ( + prev_ac_tracked_state is not None + and ac_tracked_state.track_dist_to_exit_cr is not None + and prev_ac_tracked_state.track_dist_to_exit_cr is not None + ): + curr_dist = float(ac_tracked_state.track_dist_to_exit_cr) + prev_dist = float(prev_ac_tracked_state.track_dist_to_exit_cr) + distance_improvement = prev_dist - curr_dist + + # Positive when the aircraft gets closer to its exit, negative when it + # drifts away or loiters. Clip to keep this component bounded. + reward += max(min(distance_improvement / 5.0, 2.0), -2.0) + + if ( + gym_env.timestep >= gym_env.maxstep + and ac_tracked_state.track_dist_to_exit_cr is not None + and ac_tracked_state.pos_status != PositionStatus.EXIT_REACHED + ): + timeout_distance = float(ac_tracked_state.track_dist_to_exit_cr) + reward -= min(timeout_distance / 20.0, 10.0) + + return float(reward) + + +def lateral_termination_check_sac_env( + gym_env: BaseEnv, callsign: str, action: int, **kwargs: dict[str, dict[str, Any]] +) -> float: + """Notebook-safe wrapper for the SAC lateral termination reward. + + The base reward registry invokes reward functions with only + `(gym_env, callsign, action)`. This wrapper adapts the original helper, + which expects explicit `timestep` and `maxstep` arguments. + """ + + return float( + lateral_termination_check_sac( + gym_env=gym_env, + callsign=callsign, + action=action, + timestep=gym_env.timestep, + maxstep=gym_env.maxstep, + **kwargs, + ) + ) + + +def lateral_termination_check_mac_env( + gym_env: BaseEnv, callsign: str, action: int, **kwargs: dict[str, dict[str, Any]] +) -> float: + """Notebook-safe wrapper for the MAC lateral termination reward.""" + + ac_tracked_state = gym_env.get_tracked_aircraft_data(callsign) + transferred = ac_tracked_state is not None and ac_tracked_state.pos_status == PositionStatus.EXIT_REACHED + + return float( + lateral_termination_check_mac( + gym_env=gym_env, + callsign=callsign, + action=action, + timestep=gym_env.timestep, + maxstep=gym_env.maxstep, + transferred=transferred, + **kwargs, + ) + ) + + +def anti_loiter_route_rejoin_reward(gym_env: BaseEnv, callsign: str, action: int) -> float: + """Reward progress and route rejoin, penalize stalling near the entry. + + This is used by the PPO curriculum notebook to address a specific failure + mode seen in multi-aircraft stages: orbiting near the start fix to avoid + future conflict instead of making progress, taking an avoidance action, + then returning to the route. + """ + + simulator_env = gym_env.get_simulator_env() + aircraft = simulator_env.aircraft[callsign] + + ac_tracked_state = gym_env.get_tracked_aircraft_data(callsign) + prev_ac_tracked_state = gym_env.get_tracked_aircraft_data_previous(callsign) + + if ac_tracked_state is None or prev_ac_tracked_state is None: + return 0.0 + + curr_exit_dist = ac_tracked_state.track_dist_to_exit_cr + prev_exit_dist = prev_ac_tracked_state.track_dist_to_exit_cr + + if curr_exit_dist is None or prev_exit_dist is None: + return 0.0 + + reward = 0.0 + + progress_nm = float(prev_exit_dist - curr_exit_dist) + + curr_centre_dist = None + prev_centre_dist = None + if ac_tracked_state.centreline_info_fr is not None: + curr_centre_dist = float(ac_tracked_state.centreline_info_fr[0]) + if prev_ac_tracked_state.centreline_info_fr is not None: + prev_centre_dist = float(prev_ac_tracked_state.centreline_info_fr[0]) + + if progress_nm > 0.15: + reward += min(progress_nm / 2.0, 1.0) + elif float(curr_exit_dist) > 20.0: + reward -= 0.35 + + if curr_centre_dist is not None and prev_centre_dist is not None: + centreline_improvement = prev_centre_dist - curr_centre_dist + if centreline_improvement > 0.1: + reward += min(centreline_improvement / 2.0, 0.75) + elif curr_centre_dist > 8.0 and action == 0: + reward -= 0.25 + + if curr_centre_dist < 3.0 and progress_nm > 0.15: + reward += 0.4 + + if aircraft.on_route and progress_nm > 0.15: + reward += 0.25 + + return float(reward) diff --git a/bluebird-gymnasium/bluebird_gymnasium/rewards/expeditious.py b/bluebird-gymnasium/bluebird_gymnasium/rewards/expeditious.py index 9287c84..8f75717 100644 --- a/bluebird-gymnasium/bluebird_gymnasium/rewards/expeditious.py +++ b/bluebird-gymnasium/bluebird_gymnasium/rewards/expeditious.py @@ -75,9 +75,9 @@ def expeditious_linear(gym_env: BaseEnv, callsign: str, action: int, **kwargs) - else: curr_dist = ac_tracked_state.track_dist_to_exit_cr prev_dist = prev_ac_tracked_state.track_dist_to_exit_cr - dist_diff = curr_dist - prev_dist + distance_improvement = prev_dist - curr_dist - reward = np.clip(dist_diff, -DIFF_THRESHOLD, DIFF_THRESHOLD) + reward = np.clip(distance_improvement, -DIFF_THRESHOLD, DIFF_THRESHOLD) # scale the reward based on the aircraft's speed. (i.e., it is # more impressive for an aircraft with a slower speed to travel @@ -136,9 +136,9 @@ def expeditious_quad(gym_env: BaseEnv, callsign: str, action: int, **kwargs) -> else: curr_dist = ac_tracked_state.track_dist_to_exit_cr prev_dist = prev_ac_tracked_state.track_dist_to_exit_cr - dist_diff = curr_dist - prev_dist + distance_improvement = prev_dist - curr_dist - reward = np.clip(dist_diff, -DIFF_THRESHOLD, DIFF_THRESHOLD) + reward = np.clip(distance_improvement, -DIFF_THRESHOLD, DIFF_THRESHOLD) sign = np.sign(reward) reward = 1.5 * (reward**2) @@ -199,14 +199,14 @@ def expeditious_exp(gym_env: BaseEnv, callsign: str, action: int, **kwargs) -> f else: curr_dist = ac_tracked_state.track_dist_to_exit_cr prev_dist = prev_ac_tracked_state.track_dist_to_exit_cr - dist_diff = curr_dist - prev_dist + distance_improvement = prev_dist - curr_dist - if dist_diff < 0: + if distance_improvement > 0: # clip difference to a minimum of -1.0 - dist_diff = max(dist_diff, -DIFF_THRESHOLD) + distance_improvement = min(distance_improvement, DIFF_THRESHOLD) - # the closer `dist_diff` is to -1.0, the higher the reward - reward = np.exp(dist_diff + DIFF_THRESHOLD) + # the closer `distance_improvement` is to 1.0, the higher the reward + reward = np.exp(distance_improvement - DIFF_THRESHOLD) # scale the reward based on the aircraft's speed. (i.e., it is # more impressive for an aircraft with a slower speed to travel diff --git a/bluebird-gymnasium/bluebird_gymnasium/rewards/lateral_next_fix_proximity.py b/bluebird-gymnasium/bluebird_gymnasium/rewards/lateral_next_fix_proximity.py index 524fe77..620c2da 100644 --- a/bluebird-gymnasium/bluebird_gymnasium/rewards/lateral_next_fix_proximity.py +++ b/bluebird-gymnasium/bluebird_gymnasium/rewards/lateral_next_fix_proximity.py @@ -123,4 +123,10 @@ def lateral_next_fix_proximity_dist_exp(gym_env: BaseEnv, callsign: str, action: pf_nf_dist = prev_fix.distance(next_fix) ac_nf_dist = aircraft.distance(next_fix) + # Some scenarios can yield a degenerate route segment where the previous + # and next fixes coincide. Fall back to a direct exponential in that case + # instead of dividing by zero. + if pf_nf_dist <= 1e-6: + return np.exp(-ac_nf_dist) + return np.exp(-ac_nf_dist / pf_nf_dist) diff --git a/bluebird-gymnasium/examples/curriculum_scenarios.py b/bluebird-gymnasium/examples/curriculum_scenarios.py new file mode 100644 index 0000000..5f5b91a --- /dev/null +++ b/bluebird-gymnasium/examples/curriculum_scenarios.py @@ -0,0 +1,141 @@ +from __future__ import annotations + +from datetime import timedelta +from typing import Literal + +import pandas as pd +from bluebird_dt.core import Aircraft, Coordination, FlightPlan, Route +from bluebird_dt.events.event_handler import EventHandler +from bluebird_dt.scenario_manager.tactical import Tactical +from pydantic import BaseModel, Field +from typing_extensions import override + +from bluebird_gymnasium.envs import SCENARIO_CLS + + +class FixedSequenceScenarioManagerConfig(BaseModel): + """Configuration for the fixed-sequence curriculum scenario manager.""" + + scenario_manager: Literal["fixed_sequence"] = Field(default="fixed_sequence") + + +class FixedSequenceTactical(Tactical): + """Tactical scenario manager with explicit per-aircraft route/time/speed specs. + + This is intended for curriculum stages where the notebook needs deterministic + traffic geometry instead of the stock random route sampling used by Tactical. + """ + + def __init__( + self, + airspace: object, + routes: list[Route], + aircraft_specs: list[dict], + start_time: int = 0, + vertical_buffer_distance: float | int = 500, + lateral_buffer_distance: float | int = 20, + initialise_with_event_handler: bool = True, + ) -> None: + super().__init__( + num_aircraft=len(aircraft_specs), + airspace=airspace, + routes=routes, + balance=[0.0, 0.0, 1.0], + speed_range=[400.0, 400.0], + time_entry_gap=0.0, + lateral_offset=None, + env_manager_class=None, + start_time=start_time, + vertical_buffer_distance=vertical_buffer_distance, + lateral_buffer_distance=lateral_buffer_distance, + initialise_with_event_handler=initialise_with_event_handler, + ) + self.aircraft_specs = aircraft_specs + + def _resolve_route(self, route_spec: dict) -> Route: + if "route_filed" in route_spec: + return Route(route_spec["route_filed"]) + + route_index = route_spec.get("route_index") + if route_index is None: + raise ValueError("Each aircraft spec must define either route_filed or route_index.") + + base_route = self.routes[route_index] + if route_spec.get("reverse", False): + return Route(base_route.filed[::-1]) + + return base_route + + @override + def create_event_handler(self) -> EventHandler: + sector_name = next(iter(self.airspace.sectors.keys())) + event_handler = EventHandler(ignore=self.event_handler_ignore_flags) + + for index, spec in enumerate(self.aircraft_specs): + route = self._resolve_route(spec) + callsign = spec.get("callsign", f"AIR{index}") + entry_fl = float(spec["entry_fl"]) + exit_fl = float(spec.get("exit_fl", entry_fl)) + speed_tas = float(spec["speed_tas"]) + start_time_seconds = float(spec.get("start_time_seconds", self.start_time)) + + start_fix = self.airspace.fixes.places[route.filed[0]] + next_fix = self.airspace.fixes.places[route.filed[1]] + heading = start_fix.bearing_to(next_fix) + + flight_plan = FlightPlan(route) + pos = start_fix.pos3d(entry_fl) + + aircraft = Aircraft( + pos.lat, + pos.lon, + pos.fl, + heading, + flight_plan, + callsign, + selected_fl=pos.fl, + current_sector=None, + ) + aircraft.speed_tas = speed_tas + aircraft.simulated = True + aircraft.selected_instructions.cas = speed_tas + + coordination_entry = Coordination( + callsign=callsign, + from_sector="background", + to_sector=sector_name, + fl=entry_fl, + fix=route.filed[0], + direction="Horizontal", + ) + coordination_exit = Coordination( + callsign=callsign, + from_sector=sector_name, + to_sector="background", + fl=exit_fl, + fix=route.filed[-1], + direction="Horizontal", + ) + + event_start_time = pd.to_datetime(start_time_seconds, unit="s") + event_handler.add_coordination( + event_start_time - timedelta(seconds=1), + coordination_exit, + ) + event_handler.add_coordination( + event_start_time - timedelta(seconds=1), + coordination_entry, + ) + event_handler.add_aircraft(event_start_time, aircraft) + + return event_handler + + @override + def config(self) -> FixedSequenceScenarioManagerConfig: + return FixedSequenceScenarioManagerConfig() + + +def register_curriculum_scenario_managers() -> None: + """Register notebook-specific scenario managers with the gymnasium env lookup.""" + + SCENARIO_CLS["fixed_sequence"] = FixedSequenceTactical diff --git a/bluebird-gymnasium/examples/flight_school_demo.ipynb b/bluebird-gymnasium/examples/flight_school_demo.ipynb new file mode 100644 index 0000000..d33217c --- /dev/null +++ b/bluebird-gymnasium/examples/flight_school_demo.ipynb @@ -0,0 +1,262 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "flight-school-intro", + "metadata": {}, + "source": [ + "# Flight School Demo\n", + "\n", + "This notebook demonstrates the Bluebird Gymnasium `FlightSchoolEnv-v0` environment. Flight School uses the X-plus sector with an infinite traffic generator, so aircraft continue to enter the sector during the episode." + ] + }, + { + "cell_type": "markdown", + "id": "imports-heading", + "metadata": {}, + "source": [ + "## Imports" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "imports", + "metadata": {}, + "outputs": [], + "source": [ + "from enum import Enum\n", + "from pprint import pprint\n", + "\n", + "import gymnasium as gym\n", + "import IPython.display\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "import bluebird_gymnasium\n", + "from bluebird_gymnasium.envs import BaseEnv, FlightSchoolEnv" + ] + }, + { + "cell_type": "markdown", + "id": "helpers-heading", + "metadata": {}, + "source": [ + "## Helpers" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "helpers", + "metadata": {}, + "outputs": [], + "source": [ + "def actions_to_enum(env: BaseEnv) -> Enum:\n", + " actions_map = env.get_action_parser().action_formatter_map\n", + " enum_items = {}\n", + "\n", + " for action_int, action_name in actions_map.items():\n", + " parts = action_name.split(\"__\")\n", + " base_name = parts[0]\n", + " magnitude = parts[1] if len(parts) == 2 else None\n", + "\n", + " if base_name == \"action_noop\":\n", + " enum_name = \"NOOP\"\n", + " elif base_name == \"simple_heading_left\":\n", + " enum_name = \"LEFT\"\n", + " elif base_name == \"simple_heading_right\":\n", + " enum_name = \"RIGHT\"\n", + " elif base_name == \"simple_heading_route_parallel\":\n", + " enum_name = \"ROUTE_PARALLEL\"\n", + " else:\n", + " enum_name = base_name.upper()\n", + "\n", + " if magnitude is not None:\n", + " enum_name = f\"{enum_name}_{magnitude}\"\n", + "\n", + " enum_items[enum_name] = action_int\n", + "\n", + " return Enum(\"Actions\", list(enum_items.items()))\n", + "\n", + "\n", + "def active_aircraft_count(env: BaseEnv) -> int:\n", + " return len(env.get_manager().environment.aircraft)" + ] + }, + { + "cell_type": "markdown", + "id": "scenario-config-heading", + "metadata": {}, + "source": [ + "## Scenario Config\n", + "\n", + "The default Flight School config uses the X-plus sector, two starter aircraft, and a gradual increasing spawn rate. This notebook runs a 10-minute Gymnasium episode from that generator." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "scenario-config", + "metadata": {}, + "outputs": [], + "source": [ + "seed = 7\n", + "\n", + "config = FlightSchoolEnv.get_default_env_config()\n", + "config.scenario_config[\"args\"][\"random_seed\"] = seed\n", + "config.scenario_duration = 10 * 60\n", + "config.view_config[\"type\"] = \"centralized\"\n", + "\n", + "pprint(config.scenario_config)" + ] + }, + { + "cell_type": "markdown", + "id": "instantiate-heading", + "metadata": {}, + "source": [ + "## Instantiate The Environment" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "instantiate", + "metadata": {}, + "outputs": [], + "source": [ + "env = gym.make(\"FlightSchoolEnv-v0\", config=config).unwrapped\n", + "env.set_render_mode(\"human\")\n", + "\n", + "obs, info = env.reset(seed=seed)\n", + "Actions = actions_to_enum(env)\n", + "\n", + "print(f\"Observation shape: {obs.shape}\")\n", + "print(f\"Action count: {env.action_space.n}\")\n", + "print(f\"Active aircraft: {active_aircraft_count(env)}\")\n", + "print(Actions.__members__)" + ] + }, + { + "cell_type": "markdown", + "id": "baseline-heading", + "metadata": {}, + "source": [ + "## Baseline Run\n", + "\n", + "This baseline uses `NOOP` at each step. It is useful as a smoke test and as a reference trace before adding an agent." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "baseline-run", + "metadata": {}, + "outputs": [], + "source": [ + "num_steps = config.scenario_duration // config.scenario_sec_per_step\n", + "noop = Actions.NOOP.value\n", + "\n", + "reward_history = []\n", + "traffic_history = []\n", + "time_history = []\n", + "\n", + "obs, info = env.reset(seed=seed)\n", + "\n", + "for step in range(num_steps):\n", + " obs, reward, done, truncated, info = env.step(noop)\n", + " reward_history.append(reward)\n", + " traffic_history.append(active_aircraft_count(env))\n", + " time_history.append(env.get_manager().environment.time)\n", + "\n", + " if step % 10 == 0 or done or truncated:\n", + " env.render()\n", + " IPython.display.display(env.radar.figure)\n", + " IPython.display.clear_output(wait=True)\n", + "\n", + " if done or truncated:\n", + " break\n", + "\n", + "env.render()\n", + "IPython.display.display(env.radar.figure)\n", + "print(f\"Completed {len(reward_history)} steps\")" + ] + }, + { + "cell_type": "markdown", + "id": "trace-heading", + "metadata": {}, + "source": [ + "## Reward And Traffic Trace" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "trace-plot", + "metadata": {}, + "outputs": [], + "source": [ + "fig, axes = plt.subplots(2, 1, figsize=(10, 6), sharex=True)\n", + "\n", + "axes[0].plot(time_history, reward_history, marker=\"o\", linewidth=1)\n", + "axes[0].set_ylabel(\"Reward\")\n", + "axes[0].grid(True, alpha=0.3)\n", + "\n", + "axes[1].step(time_history, traffic_history, where=\"post\")\n", + "axes[1].set_xlabel(\"Simulation time (s)\")\n", + "axes[1].set_ylabel(\"Active aircraft\")\n", + "axes[1].grid(True, alpha=0.3)\n", + "\n", + "fig.tight_layout()" + ] + }, + { + "cell_type": "markdown", + "id": "direct-construction-heading", + "metadata": {}, + "source": [ + "## Appendix: Direct Class Construction\n", + "\n", + "`gym.make` is convenient when using Gymnasium wrappers. Direct construction is useful when you want the concrete Bluebird environment instance immediately." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "direct-construction", + "metadata": {}, + "outputs": [], + "source": [ + "direct_env = FlightSchoolEnv(config=config)\n", + "direct_obs, direct_info = direct_env.reset(seed=seed)\n", + "\n", + "print(type(direct_env).__name__)\n", + "print(direct_obs.shape)\n", + "print(type(direct_env.scenario_manager).__name__)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/bluebird-gymnasium/examples/minimal_ppo_agent.ipynb b/bluebird-gymnasium/examples/minimal_ppo_agent.ipynb new file mode 100644 index 0000000..59f63e3 --- /dev/null +++ b/bluebird-gymnasium/examples/minimal_ppo_agent.ipynb @@ -0,0 +1,1094 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Minimal PPO Agent for Bluebird Gymnasium\n", + "\n", + "This notebook mirrors the REINFORCE example, but swaps in a minimal actor-critic PPO training loop.\n", + "\n", + "It keeps the same broad workflow:\n", + "\n", + "- small `SectorIEnv` configuration\n", + "- decentralized control with one aircraft\n", + "- periodic evaluation during training\n", + "- comparison against a random baseline\n", + "- checkpoint saving and best-model restore\n", + "- final GIF rendering of the best policy\n", + "\n", + "This is still educational rather than production-grade PPO, but it is much closer to what you would use in practice than plain REINFORCE." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Environment note\n", + "\n", + "This notebook assumes you are already running the **correct Python/Jupyter kernel**:\n", + "one with `gymnasium`, `torch`, and the Bluebird project dependencies installed.\n", + "\n", + "Do the package setup in your shell or project virtual environment first,\n", + "then open this notebook with that kernel. The notebook does **not** try to install\n", + "packages itself." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Imports and path setup\n", + "\n", + "This cell makes the notebook runnable from either the `bluebird-gymnasium` directory\n", + "or the repo root by adding the local package paths to `sys.path`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from __future__ import annotations\n", + "\n", + "import random\n", + "import shutil\n", + "import sys\n", + "from pathlib import Path\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import torch\n", + "import torch.nn as nn\n", + "import torch.optim as optim\n", + "from IPython.display import Image, display\n", + "\n", + "search_roots = [Path.cwd().resolve(), *Path.cwd().resolve().parents]\n", + "gym_root = None\n", + "dt_root = None\n", + "\n", + "for candidate in search_roots:\n", + " if (candidate / 'bluebird_gymnasium').exists():\n", + " gym_root = candidate\n", + " sibling_dt = candidate.parent / 'bluebird-dt'\n", + " if sibling_dt.exists():\n", + " dt_root = sibling_dt\n", + " break\n", + " if (candidate / 'bluebird-gymnasium').exists() and (candidate / 'bluebird-dt').exists():\n", + " gym_root = candidate / 'bluebird-gymnasium'\n", + " dt_root = candidate / 'bluebird-dt'\n", + " break\n", + "\n", + "if gym_root is None or dt_root is None:\n", + " raise RuntimeError('Could not locate local bluebird-gymnasium and bluebird-dt package roots.')\n", + "\n", + "sys.path.insert(0, str(gym_root))\n", + "sys.path.insert(0, str(dt_root))\n", + "\n", + "from bluebird_gymnasium.envs import EnvConfig, ViewType\n", + "from bluebird_gymnasium.envs.sector_i import SectorIEnv\n", + "from bluebird_gymnasium.utils.video import generate_video\n", + "\n", + "print(f'Using bluebird-gymnasium from: {gym_root}')\n", + "print(f'Using bluebird-dt from: {dt_root}')\n", + "print(f'Torch version: {torch.__version__}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## PPO actor-critic model and agents\n", + "\n", + "The actor-critic network has:\n", + "\n", + "- a shared feature trunk\n", + "- a **policy head** that outputs action logits\n", + "- a **value head** that predicts the expected return from the current observation\n", + "\n", + "The PPO update then uses:\n", + "\n", + "- old action log-probabilities from rollout collection\n", + "- new action log-probabilities from the current policy\n", + "- a clipped objective to avoid excessively large policy updates\n", + "- a value loss for the critic\n", + "- a small entropy bonus to keep exploration from collapsing too early" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "class ActorCriticNetwork(nn.Module):\n", + " \"\"\"Shared-trunk actor-critic network for PPO.\"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " observation_dimension: int,\n", + " number_of_actions: int,\n", + " hidden_units: int = 128,\n", + " ) -> None:\n", + " super().__init__()\n", + " self.trunk = nn.Sequential(\n", + " nn.Linear(observation_dimension, hidden_units),\n", + " nn.ReLU(),\n", + " nn.Linear(hidden_units, hidden_units),\n", + " nn.ReLU(),\n", + " )\n", + " self.policy_head = nn.Linear(hidden_units, number_of_actions)\n", + " self.value_head = nn.Linear(hidden_units, 1)\n", + "\n", + " def forward(self, observation_batch: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:\n", + " features = self.trunk(observation_batch)\n", + " action_logits = self.policy_head(features)\n", + " state_value = self.value_head(features).squeeze(-1)\n", + " return action_logits, state_value\n", + "\n", + "\n", + "class PPOAgent:\n", + " \"\"\"Minimal PPO agent with actor-critic network and checkpoint helpers.\"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " observation_dimension: int,\n", + " number_of_actions: int,\n", + " learning_rate: float = 3e-4,\n", + " hidden_units: int = 128,\n", + " clip_epsilon: float = 0.2,\n", + " value_loss_coefficient: float = 0.5,\n", + " entropy_coefficient: float = 0.01,\n", + " ppo_epochs: int = 4,\n", + " ) -> None:\n", + " self.actor_critic = ActorCriticNetwork(\n", + " observation_dimension=observation_dimension,\n", + " number_of_actions=number_of_actions,\n", + " hidden_units=hidden_units,\n", + " )\n", + " self.optimizer = optim.Adam(self.actor_critic.parameters(), lr=learning_rate)\n", + " self.clip_epsilon = clip_epsilon\n", + " self.value_loss_coefficient = value_loss_coefficient\n", + " self.entropy_coefficient = entropy_coefficient\n", + " self.ppo_epochs = ppo_epochs\n", + "\n", + " def choose_training_action(\n", + " self,\n", + " observation_vector: np.ndarray,\n", + " ) -> tuple[int, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:\n", + " observation_tensor = torch.tensor(\n", + " observation_vector,\n", + " dtype=torch.float32,\n", + " ).unsqueeze(0)\n", + " action_logits, state_value = self.actor_critic(observation_tensor)\n", + " action_distribution = torch.distributions.Categorical(logits=action_logits)\n", + " sampled_action = action_distribution.sample()\n", + " action_log_probability = action_distribution.log_prob(sampled_action)\n", + " action_entropy = action_distribution.entropy()\n", + "\n", + " return (\n", + " sampled_action.item(),\n", + " observation_tensor.squeeze(0),\n", + " action_log_probability.squeeze(0),\n", + " state_value.squeeze(0),\n", + " action_entropy.squeeze(0),\n", + " )\n", + "\n", + " def choose_evaluation_actions(\n", + " self,\n", + " observation_by_callsign: dict[str, np.ndarray],\n", + " ) -> dict[str, int]:\n", + " callsign, observation_vector = unpack_single_aircraft_observation(observation_by_callsign)\n", + " with torch.no_grad():\n", + " observation_tensor = torch.tensor(\n", + " observation_vector,\n", + " dtype=torch.float32,\n", + " ).unsqueeze(0)\n", + " action_logits, _state_value = self.actor_critic(observation_tensor)\n", + " chosen_action = torch.argmax(action_logits, dim=-1).item()\n", + " return {callsign: chosen_action}\n", + "\n", + " def update_from_trajectory(\n", + " self,\n", + " observations: list[torch.Tensor],\n", + " actions: list[torch.Tensor],\n", + " old_log_probabilities: list[torch.Tensor],\n", + " returns: torch.Tensor,\n", + " advantages: torch.Tensor,\n", + " ) -> dict[str, float] | None:\n", + " if not observations:\n", + " return None\n", + "\n", + " observation_tensor = torch.stack(observations)\n", + " action_tensor = torch.stack(actions).long()\n", + " old_log_probability_tensor = torch.stack(old_log_probabilities).detach()\n", + " returns_tensor = returns.detach()\n", + " advantages_tensor = advantages.detach()\n", + "\n", + " if advantages_tensor.numel() > 1:\n", + " advantages_std = advantages_tensor.std(unbiased=False)\n", + " if advantages_std > 1e-8:\n", + " advantages_tensor = (\n", + " (advantages_tensor - advantages_tensor.mean())\n", + " / (advantages_std + 1e-8)\n", + " )\n", + "\n", + " mean_policy_loss = 0.0\n", + " mean_value_loss = 0.0\n", + " mean_entropy = 0.0\n", + " mean_total_loss = 0.0\n", + "\n", + " for _epoch in range(self.ppo_epochs):\n", + " new_action_logits, new_state_values = self.actor_critic(observation_tensor)\n", + " action_distribution = torch.distributions.Categorical(logits=new_action_logits)\n", + " new_log_probabilities = action_distribution.log_prob(action_tensor)\n", + " entropy = action_distribution.entropy().mean()\n", + "\n", + " probability_ratio = torch.exp(new_log_probabilities - old_log_probability_tensor)\n", + " unclipped_objective = probability_ratio * advantages_tensor\n", + " clipped_objective = torch.clamp(\n", + " probability_ratio,\n", + " 1.0 - self.clip_epsilon,\n", + " 1.0 + self.clip_epsilon,\n", + " ) * advantages_tensor\n", + "\n", + " policy_loss = -torch.min(unclipped_objective, clipped_objective).mean()\n", + " value_loss = torch.nn.functional.mse_loss(new_state_values, returns_tensor)\n", + " total_loss = (\n", + " policy_loss\n", + " + self.value_loss_coefficient * value_loss\n", + " - self.entropy_coefficient * entropy\n", + " )\n", + "\n", + " self.optimizer.zero_grad()\n", + " total_loss.backward()\n", + " self.optimizer.step()\n", + "\n", + " mean_policy_loss += float(policy_loss.item())\n", + " mean_value_loss += float(value_loss.item())\n", + " mean_entropy += float(entropy.item())\n", + " mean_total_loss += float(total_loss.item())\n", + "\n", + " epoch_divisor = float(self.ppo_epochs)\n", + " return {\n", + " 'policy_loss': mean_policy_loss / epoch_divisor,\n", + " 'value_loss': mean_value_loss / epoch_divisor,\n", + " 'entropy': mean_entropy / epoch_divisor,\n", + " 'total_loss': mean_total_loss / epoch_divisor,\n", + " }\n", + "\n", + " def save_checkpoint(self, checkpoint_path: Path, metadata: dict | None = None) -> None:\n", + " checkpoint_path.parent.mkdir(parents=True, exist_ok=True)\n", + " payload = {\n", + " 'model_state_dict': self.actor_critic.state_dict(),\n", + " 'optimizer_state_dict': self.optimizer.state_dict(),\n", + " 'metadata': _to_python_types(metadata or {}),\n", + " }\n", + " torch.save(payload, checkpoint_path)\n", + "\n", + " def load_checkpoint(self, checkpoint_path: Path, map_location: str = 'cpu') -> dict:\n", + " payload = torch.load(checkpoint_path, map_location=map_location, weights_only=False)\n", + " self.actor_critic.load_state_dict(payload['model_state_dict'])\n", + " if 'optimizer_state_dict' in payload:\n", + " self.optimizer.load_state_dict(payload['optimizer_state_dict'])\n", + " return payload.get('metadata', {})\n", + "\n", + "\n", + "\n", + "def _to_python_types(value):\n", + " if isinstance(value, dict):\n", + " return {key: _to_python_types(val) for key, val in value.items()}\n", + " if isinstance(value, (list, tuple)):\n", + " return [_to_python_types(item) for item in value]\n", + " if isinstance(value, Path):\n", + " return str(value)\n", + " if isinstance(value, np.generic):\n", + " return value.item()\n", + " return value\n", + "\n", + "\n", + "class RandomAgent:\n", + " \"\"\"Simple random baseline for comparison.\"\"\"\n", + "\n", + " def __init__(self, number_of_actions: int) -> None:\n", + " self.number_of_actions = number_of_actions\n", + "\n", + " def choose_evaluation_actions(\n", + " self,\n", + " observation_by_callsign: dict[str, np.ndarray],\n", + " ) -> dict[str, int]:\n", + " return {\n", + " callsign: random.randrange(self.number_of_actions)\n", + " for callsign in observation_by_callsign.keys()\n", + " }" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Environment configuration and rollout helpers\n", + "\n", + "This notebook assumes the same deliberately small setup as the REINFORCE version:\n", + "\n", + "- `SectorIEnv`\n", + "- decentralized control\n", + "- one aircraft\n", + "- `extra_minimal` state encoding\n", + "- lateral actions only\n", + "\n", + "That keeps the PPO implementation easier to follow." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def unpack_single_aircraft_observation(\n", + " observation_by_callsign: dict[str, np.ndarray],\n", + ") -> tuple[str, np.ndarray]:\n", + " \"\"\"Extract the single controllable aircraft observation expected by this notebook.\"\"\"\n", + " if len(observation_by_callsign) != 1:\n", + " raise ValueError(\n", + " 'This PPO notebook assumes exactly one controllable aircraft in the observation dict. '\n", + " f'Got {len(observation_by_callsign)} entries instead.'\n", + " )\n", + " return next(iter(observation_by_callsign.items()))\n", + "\n", + "\n", + "def make_sector_i_training_config() -> EnvConfig:\n", + " config = SectorIEnv.get_default_env_config(ViewType.DECENTRALIZED)\n", + "\n", + " config.state_repr_config = {\n", + " 'encoder_cls': 'extra_minimal',\n", + " 'k_nearest_aircraft': 1,\n", + " }\n", + "\n", + " config.action_config = {\n", + " 'simple_heading_left': True,\n", + " 'simple_heading_right': True,\n", + " 'simple_fl_climb': False,\n", + " 'simple_fl_descent': False,\n", + " 'simple_fl_exit': False,\n", + " }\n", + "\n", + " config.reward_config = {\n", + " 'fns': [\n", + " 'position_status_const',\n", + " 'lateral_centreline_distance_shaped',\n", + " 'safety_simple_avoidance_exp',\n", + " ],\n", + " 'coeffs': [1.0, 1.0, 1.2],\n", + " }\n", + "\n", + " config.view_config = {\n", + " 'type': ViewType.DECENTRALIZED.value,\n", + " 'decentralized_params': {},\n", + " }\n", + "\n", + " config.scenario_config = {\n", + " 'cls': 'tactical',\n", + " 'args': {\n", + " 'num_aircraft': 1,\n", + " 'balance': [0.0, 0.0, 1.0],\n", + " },\n", + " }\n", + "\n", + " return config\n", + "\n", + "\n", + "def compute_returns_and_advantages(\n", + " rewards: list[float],\n", + " values: list[torch.Tensor],\n", + " dones: list[bool],\n", + " discount_factor_gamma: float,\n", + " gae_lambda: float,\n", + ") -> tuple[torch.Tensor, torch.Tensor]:\n", + " rewards_tensor = torch.tensor(rewards, dtype=torch.float32)\n", + " values_tensor = torch.stack(values).detach().float()\n", + " dones_tensor = torch.tensor(dones, dtype=torch.float32)\n", + "\n", + " advantages = torch.zeros_like(rewards_tensor)\n", + " last_gae = torch.tensor(0.0)\n", + "\n", + " for timestep in reversed(range(len(rewards))):\n", + " if timestep == len(rewards) - 1:\n", + " next_value = torch.tensor(0.0)\n", + " else:\n", + " next_value = values_tensor[timestep + 1]\n", + "\n", + " next_nonterminal = 1.0 - dones_tensor[timestep]\n", + " delta = (\n", + " rewards_tensor[timestep]\n", + " + discount_factor_gamma * next_value * next_nonterminal\n", + " - values_tensor[timestep]\n", + " )\n", + " last_gae = delta + discount_factor_gamma * gae_lambda * next_nonterminal * last_gae\n", + " advantages[timestep] = last_gae\n", + "\n", + " returns = advantages + values_tensor\n", + " return returns, advantages\n", + "\n", + "\n", + "def run_one_training_episode(\n", + " environment: SectorIEnv,\n", + " agent: PPOAgent,\n", + " random_seed: int,\n", + " discount_factor_gamma: float,\n", + " gae_lambda: float,\n", + ") -> tuple[float, int, dict[str, float] | None]:\n", + " random.seed(random_seed)\n", + " np.random.seed(random_seed)\n", + " torch.manual_seed(random_seed)\n", + "\n", + " observation_by_callsign, _info = environment.reset(seed=random_seed)\n", + "\n", + " episode_is_done = False\n", + " episode_step_count = 0\n", + " episode_total_reward = 0.0\n", + "\n", + " observations: list[torch.Tensor] = []\n", + " actions: list[torch.Tensor] = []\n", + " old_log_probabilities: list[torch.Tensor] = []\n", + " values: list[torch.Tensor] = []\n", + " rewards: list[float] = []\n", + " dones: list[bool] = []\n", + "\n", + " while not episode_is_done:\n", + " callsign, observation_vector = unpack_single_aircraft_observation(observation_by_callsign)\n", + " (\n", + " action_int,\n", + " observation_tensor,\n", + " action_log_probability,\n", + " state_value,\n", + " _action_entropy,\n", + " ) = agent.choose_training_action(observation_vector)\n", + "\n", + " action_by_callsign = {callsign: action_int}\n", + " (\n", + " next_observation_by_callsign,\n", + " reward_by_callsign,\n", + " done_by_callsign,\n", + " truncated_by_callsign,\n", + " _info,\n", + " ) = environment.step(action_by_callsign)\n", + "\n", + " timestep_reward = float(sum(reward_by_callsign.values())) if reward_by_callsign else 0.0\n", + " timestep_done = all(done_by_callsign.values()) if done_by_callsign else True\n", + "\n", + " observations.append(observation_tensor)\n", + " actions.append(torch.tensor(action_int))\n", + " old_log_probabilities.append(action_log_probability.detach())\n", + " values.append(state_value.detach())\n", + " rewards.append(timestep_reward)\n", + " dones.append(bool(timestep_done))\n", + "\n", + " _ = truncated_by_callsign\n", + " episode_total_reward += timestep_reward\n", + " episode_is_done = timestep_done\n", + " observation_by_callsign = next_observation_by_callsign\n", + " episode_step_count += 1\n", + "\n", + " returns, advantages = compute_returns_and_advantages(\n", + " rewards=rewards,\n", + " values=values,\n", + " dones=dones,\n", + " discount_factor_gamma=discount_factor_gamma,\n", + " gae_lambda=gae_lambda,\n", + " )\n", + "\n", + " update_metrics = agent.update_from_trajectory(\n", + " observations=observations,\n", + " actions=actions,\n", + " old_log_probabilities=old_log_probabilities,\n", + " returns=returns,\n", + " advantages=advantages,\n", + " )\n", + "\n", + " return episode_total_reward, episode_step_count, update_metrics\n", + "\n", + "\n", + "def run_one_evaluation_episode(\n", + " environment: SectorIEnv,\n", + " evaluation_agent,\n", + " random_seed: int,\n", + ") -> tuple[float, int]:\n", + " random.seed(random_seed)\n", + " np.random.seed(random_seed)\n", + " torch.manual_seed(random_seed)\n", + "\n", + " observation_by_callsign, _info = environment.reset(seed=random_seed)\n", + "\n", + " episode_is_done = False\n", + " episode_step_count = 0\n", + " episode_total_reward = 0.0\n", + "\n", + " while not episode_is_done:\n", + " action_by_callsign = evaluation_agent.choose_evaluation_actions(\n", + " observation_by_callsign,\n", + " )\n", + "\n", + " (\n", + " next_observation_by_callsign,\n", + " reward_by_callsign,\n", + " done_by_callsign,\n", + " truncated_by_callsign,\n", + " _info,\n", + " ) = environment.step(action_by_callsign)\n", + "\n", + " _ = truncated_by_callsign\n", + " episode_total_reward += (\n", + " float(sum(reward_by_callsign.values())) if reward_by_callsign else 0.0\n", + " )\n", + " episode_is_done = all(done_by_callsign.values()) if done_by_callsign else True\n", + " observation_by_callsign = next_observation_by_callsign\n", + " episode_step_count += 1\n", + "\n", + " return episode_total_reward, episode_step_count\n", + "\n", + "\n", + "def evaluate_agent_over_seeds(\n", + " environment: SectorIEnv,\n", + " evaluation_agent,\n", + " evaluation_seeds: list[int],\n", + ") -> dict:\n", + " rewards: list[float] = []\n", + " steps: list[int] = []\n", + "\n", + " for random_seed in evaluation_seeds:\n", + " total_reward, step_count = run_one_evaluation_episode(\n", + " environment=environment,\n", + " evaluation_agent=evaluation_agent,\n", + " random_seed=random_seed,\n", + " )\n", + " rewards.append(total_reward)\n", + " steps.append(step_count)\n", + "\n", + " return {\n", + " 'seeds': evaluation_seeds,\n", + " 'rewards': rewards,\n", + " 'steps': steps,\n", + " 'mean_reward': float(np.mean(rewards)),\n", + " 'std_reward': float(np.std(rewards)),\n", + " 'mean_steps': float(np.mean(steps)),\n", + " }\n", + "\n", + "\n", + "def render_evaluation_rollout_to_gif(\n", + " agent: PPOAgent,\n", + " random_seed: int,\n", + " render_dir: Path,\n", + " gif_name: str = 'trained_policy_eval',\n", + ") -> Path:\n", + " render_config = make_sector_i_training_config()\n", + " render_config.radar_config['display_actions'] = True\n", + " render_config.radar_config['render_dir'] = str(render_dir)\n", + " render_config.radar_config['prefix'] = 'frame'\n", + "\n", + " if render_dir.exists():\n", + " shutil.rmtree(render_dir)\n", + " render_dir.mkdir(parents=True, exist_ok=True)\n", + "\n", + " render_environment = SectorIEnv(config=render_config)\n", + " render_environment.set_render_mode('file')\n", + "\n", + " random.seed(random_seed)\n", + " np.random.seed(random_seed)\n", + " torch.manual_seed(random_seed)\n", + "\n", + " observation_by_callsign, _info = render_environment.reset(seed=random_seed)\n", + " render_environment.render()\n", + "\n", + " episode_is_done = False\n", + "\n", + " while not episode_is_done:\n", + " action_by_callsign = agent.choose_evaluation_actions(observation_by_callsign)\n", + " (\n", + " next_observation_by_callsign,\n", + " reward_by_callsign,\n", + " done_by_callsign,\n", + " truncated_by_callsign,\n", + " _info,\n", + " ) = render_environment.step(action_by_callsign)\n", + " _ = reward_by_callsign, truncated_by_callsign\n", + " render_environment.render()\n", + " episode_is_done = all(done_by_callsign.values()) if done_by_callsign else True\n", + " observation_by_callsign = next_observation_by_callsign\n", + "\n", + " png_frames = sorted(render_dir.glob(f\"{render_config.radar_config['prefix']}_*.png\"))\n", + " if not png_frames:\n", + " raise RuntimeError(\n", + " f'No rendered PNG frames were written to {render_dir}. '\n", + " 'Expected at least one frame before GIF generation.'\n", + " )\n", + "\n", + " generate_video(\n", + " render_dir=str(render_dir),\n", + " frame_prefix=render_config.radar_config['prefix'],\n", + " video_filename=gif_name,\n", + " clean_up=False,\n", + " )\n", + " render_environment.close()\n", + " return render_dir / f'{gif_name}.gif'" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Set up the environment and inspect the shapes\n", + "\n", + "For this notebook, the most important values are:\n", + "\n", + "- `observation_dimension`: how many numbers are in one aircraft observation vector\n", + "- `number_of_actions`: how many discrete actions the policy can choose from" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = make_sector_i_training_config()\n", + "environment = SectorIEnv(config=config)\n", + "\n", + "observation_dimension = environment.observation_space.shape[0]\n", + "number_of_actions = environment.action_space.n\n", + "\n", + "print(\n", + " 'environment shapes:',\n", + " f'observation_dimension={observation_dimension}',\n", + " f'number_of_actions={number_of_actions}',\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Hyperparameters and experiment settings\n", + "\n", + "This PPO version uses a richer experiment loop with:\n", + "\n", + "- generalized advantage estimation (GAE)\n", + "- periodic evaluation during training\n", + "- larger held-out evaluation set\n", + "- checkpoint saving\n", + "- best-checkpoint restore\n", + "- random-policy baseline comparison" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "learning_rate = 3e-4\n", + "hidden_units = 128\n", + "discount_factor_gamma = 0.99\n", + "gae_lambda = 0.95\n", + "clip_epsilon = 0.2\n", + "value_loss_coefficient = 0.5\n", + "entropy_coefficient = 0.01\n", + "ppo_epochs = 4\n", + "number_of_training_episodes = 100\n", + "training_seed_start = 100\n", + "periodic_eval_interval = 10\n", + "heldout_evaluation_seeds = list(range(200, 220))\n", + "checkpoint_dir = Path.cwd() / 'checkpoints' / 'minimal_ppo_agent'\n", + "latest_checkpoint_path = checkpoint_dir / 'latest.pt'\n", + "best_checkpoint_path = checkpoint_dir / 'best.pt'\n", + "\n", + "agent = PPOAgent(\n", + " observation_dimension=observation_dimension,\n", + " number_of_actions=number_of_actions,\n", + " learning_rate=learning_rate,\n", + " hidden_units=hidden_units,\n", + " clip_epsilon=clip_epsilon,\n", + " value_loss_coefficient=value_loss_coefficient,\n", + " entropy_coefficient=entropy_coefficient,\n", + " ppo_epochs=ppo_epochs,\n", + ")\n", + "random_agent = RandomAgent(number_of_actions=number_of_actions)\n", + "\n", + "training_rewards: list[float] = []\n", + "training_steps: list[int] = []\n", + "training_policy_losses: list[float] = []\n", + "training_value_losses: list[float] = []\n", + "training_entropies: list[float] = []\n", + "training_total_losses: list[float] = []\n", + "\n", + "periodic_eval_episodes: list[int] = []\n", + "periodic_eval_learned_mean_rewards: list[float] = []\n", + "periodic_eval_learned_std_rewards: list[float] = []\n", + "periodic_eval_random_mean_rewards: list[float] = []\n", + "periodic_eval_random_std_rewards: list[float] = []\n", + "periodic_eval_learned_mean_steps: list[float] = []\n", + "periodic_eval_random_mean_steps: list[float] = []\n", + "\n", + "best_mean_evaluation_reward = float('-inf')\n", + "best_checkpoint_metadata: dict = {}" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## PPO training loop with periodic evaluation and checkpointing\n", + "\n", + "Every `periodic_eval_interval` episodes, the notebook:\n", + "\n", + "- evaluates the current learned policy on the held-out evaluation seeds\n", + "- evaluates a random baseline on the same seeds\n", + "- saves a `latest.pt` checkpoint\n", + "- overwrites `best.pt` if the learned policy achieves a new best mean evaluation reward" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "for episode_index in range(number_of_training_episodes):\n", + " random_seed = training_seed_start + episode_index\n", + " total_reward, step_count, update_metrics = run_one_training_episode(\n", + " environment=environment,\n", + " agent=agent,\n", + " random_seed=random_seed,\n", + " discount_factor_gamma=discount_factor_gamma,\n", + " gae_lambda=gae_lambda,\n", + " )\n", + "\n", + " training_rewards.append(total_reward)\n", + " training_steps.append(step_count)\n", + " training_policy_losses.append(float('nan') if update_metrics is None else update_metrics['policy_loss'])\n", + " training_value_losses.append(float('nan') if update_metrics is None else update_metrics['value_loss'])\n", + " training_entropies.append(float('nan') if update_metrics is None else update_metrics['entropy'])\n", + " training_total_losses.append(float('nan') if update_metrics is None else update_metrics['total_loss'])\n", + "\n", + " print(\n", + " '[train]',\n", + " f'episode={episode_index:03d}',\n", + " f'seed={random_seed}',\n", + " f'reward={total_reward:.3f}',\n", + " f'steps={step_count}',\n", + " f'policy_loss={None if update_metrics is None else update_metrics[\"policy_loss\"]:.6f}' if update_metrics is not None else 'policy_loss=None',\n", + " f'value_loss={None if update_metrics is None else update_metrics[\"value_loss\"]:.6f}' if update_metrics is not None else 'value_loss=None',\n", + " )\n", + "\n", + " should_run_periodic_eval = (\n", + " (episode_index + 1) % periodic_eval_interval == 0\n", + " or episode_index == number_of_training_episodes - 1\n", + " )\n", + "\n", + " if should_run_periodic_eval:\n", + " learned_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + " )\n", + " random_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=random_agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + " )\n", + "\n", + " periodic_eval_episodes.append(episode_index + 1)\n", + " periodic_eval_learned_mean_rewards.append(learned_eval['mean_reward'])\n", + " periodic_eval_learned_std_rewards.append(learned_eval['std_reward'])\n", + " periodic_eval_random_mean_rewards.append(random_eval['mean_reward'])\n", + " periodic_eval_random_std_rewards.append(random_eval['std_reward'])\n", + " periodic_eval_learned_mean_steps.append(learned_eval['mean_steps'])\n", + " periodic_eval_random_mean_steps.append(random_eval['mean_steps'])\n", + "\n", + " metadata = {\n", + " 'episode': episode_index + 1,\n", + " 'train_seed': random_seed,\n", + " 'learned_mean_reward': learned_eval['mean_reward'],\n", + " 'learned_std_reward': learned_eval['std_reward'],\n", + " 'random_mean_reward': random_eval['mean_reward'],\n", + " 'random_std_reward': random_eval['std_reward'],\n", + " 'evaluation_seeds': heldout_evaluation_seeds,\n", + " }\n", + " agent.save_checkpoint(latest_checkpoint_path, metadata=metadata)\n", + "\n", + " if learned_eval['mean_reward'] > best_mean_evaluation_reward:\n", + " best_mean_evaluation_reward = learned_eval['mean_reward']\n", + " best_checkpoint_metadata = metadata\n", + " agent.save_checkpoint(best_checkpoint_path, metadata=metadata)\n", + " checkpoint_note = 'new best checkpoint'\n", + " else:\n", + " checkpoint_note = 'latest checkpoint only'\n", + "\n", + " print(\n", + " '[periodic-eval]',\n", + " f'episode={episode_index + 1:03d}',\n", + " f'learned_mean_reward={learned_eval[\"mean_reward\"]:.3f}',\n", + " f'learned_std_reward={learned_eval[\"std_reward\"]:.3f}',\n", + " f'random_mean_reward={random_eval[\"mean_reward\"]:.3f}',\n", + " f'random_std_reward={random_eval[\"std_reward\"]:.3f}',\n", + " checkpoint_note,\n", + " )" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Restore the best evaluated model\n", + "\n", + "The training loop may end on a policy that is not the best one seen so far.\n", + "This cell reloads the checkpoint with the highest held-out mean evaluation reward." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "if not best_checkpoint_path.exists():\n", + " raise FileNotFoundError(f'Best checkpoint not found: {best_checkpoint_path}')\n", + "\n", + "loaded_metadata = agent.load_checkpoint(best_checkpoint_path)\n", + "print('Reloaded best checkpoint from:', best_checkpoint_path)\n", + "print('Best checkpoint metadata:')\n", + "loaded_metadata" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Final evaluation of the best checkpoint vs random baseline\n", + "\n", + "This uses the larger held-out evaluation set and compares:\n", + "\n", + "- the best learned PPO checkpoint\n", + "- a random baseline on the same seeds" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "best_policy_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + ")\n", + "random_policy_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=random_agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + ")\n", + "\n", + "print('Best PPO policy evaluation mean reward:', best_policy_eval['mean_reward'])\n", + "print('Best PPO policy evaluation std reward:', best_policy_eval['std_reward'])\n", + "print('Random baseline mean reward:', random_policy_eval['mean_reward'])\n", + "print('Random baseline std reward:', random_policy_eval['std_reward'])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Useful plots\n", + "\n", + "These are the most useful quick-look plots for this PPO setup:\n", + "\n", + "- **Training reward**: raw reward and moving average during training\n", + "- **Episode length**: how long training episodes run\n", + "- **Periodic evaluation**: learned policy vs random baseline over training\n", + "- **Final evaluation by seed**: best learned checkpoint vs random on the held-out seeds\n", + "- **PPO losses**: policy loss, value loss, and entropy over training" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def moving_average(values: list[float], window: int) -> np.ndarray:\n", + " if len(values) < window:\n", + " return np.array([])\n", + " kernel = np.ones(window) / window\n", + " return np.convolve(np.asarray(values, dtype=float), kernel, mode='valid')\n", + "\n", + "plot_window = min(10, len(training_rewards))\n", + "smoothed_rewards = moving_average(training_rewards, plot_window)\n", + "training_episode_indices = np.arange(1, len(training_rewards) + 1)\n", + "heldout_seed_indices = np.arange(len(heldout_evaluation_seeds))\n", + "periodic_eval_episodes_arr = np.asarray(periodic_eval_episodes)\n", + "learned_mean_arr = np.asarray(periodic_eval_learned_mean_rewards)\n", + "learned_std_arr = np.asarray(periodic_eval_learned_std_rewards)\n", + "random_mean_arr = np.asarray(periodic_eval_random_mean_rewards)\n", + "random_std_arr = np.asarray(periodic_eval_random_std_rewards)\n", + "\n", + "fig, axes = plt.subplots(3, 2, figsize=(15, 14))\n", + "\n", + "axes[0, 0].plot(training_episode_indices, training_rewards, marker='o', alpha=0.25, label='raw reward')\n", + "if len(smoothed_rewards) > 0:\n", + " axes[0, 0].plot(\n", + " np.arange(plot_window, len(training_rewards) + 1),\n", + " smoothed_rewards,\n", + " linewidth=2.5,\n", + " color='tab:blue',\n", + " label=f'moving average (window={plot_window})',\n", + " )\n", + "axes[0, 0].set_title('Training Reward per Episode')\n", + "axes[0, 0].set_xlabel('Episode')\n", + "axes[0, 0].set_ylabel('Total Reward')\n", + "axes[0, 0].legend()\n", + "axes[0, 0].grid(alpha=0.3)\n", + "\n", + "axes[0, 1].plot(training_episode_indices, training_steps, marker='o', color='tab:orange')\n", + "axes[0, 1].set_title('Training Episode Length')\n", + "axes[0, 1].set_xlabel('Episode')\n", + "axes[0, 1].set_ylabel('Steps')\n", + "axes[0, 1].grid(alpha=0.3)\n", + "\n", + "axes[1, 0].plot(periodic_eval_episodes_arr, learned_mean_arr, marker='o', label='learned policy')\n", + "axes[1, 0].fill_between(\n", + " periodic_eval_episodes_arr,\n", + " learned_mean_arr - learned_std_arr,\n", + " learned_mean_arr + learned_std_arr,\n", + " alpha=0.2,\n", + ")\n", + "axes[1, 0].plot(periodic_eval_episodes_arr, random_mean_arr, marker='s', label='random baseline')\n", + "axes[1, 0].fill_between(\n", + " periodic_eval_episodes_arr,\n", + " random_mean_arr - random_std_arr,\n", + " random_mean_arr + random_std_arr,\n", + " alpha=0.2,\n", + ")\n", + "axes[1, 0].set_title('Periodic Evaluation: Learned vs Random')\n", + "axes[1, 0].set_xlabel('Training Episode')\n", + "axes[1, 0].set_ylabel('Mean Evaluation Reward')\n", + "axes[1, 0].legend()\n", + "axes[1, 0].grid(alpha=0.3)\n", + "\n", + "axes[1, 1].plot(\n", + " heldout_seed_indices,\n", + " best_policy_eval['rewards'],\n", + " marker='o',\n", + " linewidth=2,\n", + " label='best PPO checkpoint',\n", + ")\n", + "axes[1, 1].plot(\n", + " heldout_seed_indices,\n", + " random_policy_eval['rewards'],\n", + " marker='s',\n", + " linewidth=2,\n", + " label='random baseline',\n", + ")\n", + "axes[1, 1].set_xticks(heldout_seed_indices)\n", + "axes[1, 1].set_xticklabels([str(seed) for seed in heldout_evaluation_seeds], rotation=45)\n", + "axes[1, 1].set_title('Final Evaluation Reward by Seed')\n", + "axes[1, 1].set_xlabel('Held-out Evaluation Seed')\n", + "axes[1, 1].set_ylabel('Total Reward')\n", + "axes[1, 1].legend()\n", + "axes[1, 1].grid(alpha=0.3)\n", + "\n", + "axes[2, 0].plot(training_episode_indices, training_policy_losses, label='policy loss')\n", + "axes[2, 0].plot(training_episode_indices, training_value_losses, label='value loss')\n", + "axes[2, 0].plot(training_episode_indices, training_total_losses, label='total loss')\n", + "axes[2, 0].set_title('PPO Loss Terms')\n", + "axes[2, 0].set_xlabel('Episode')\n", + "axes[2, 0].set_ylabel('Loss')\n", + "axes[2, 0].legend()\n", + "axes[2, 0].grid(alpha=0.3)\n", + "\n", + "axes[2, 1].plot(training_episode_indices, training_entropies, label='entropy', color='tab:green')\n", + "axes[2, 1].set_title('Policy Entropy During Training')\n", + "axes[2, 1].set_xlabel('Episode')\n", + "axes[2, 1].set_ylabel('Entropy')\n", + "axes[2, 1].grid(alpha=0.3)\n", + "\n", + "fig.suptitle('Minimal PPO Training Summary with Baseline and Checkpoints', fontsize=16)\n", + "fig.tight_layout()\n", + "plt.show()\n", + "\n", + "print(f'Best PPO checkpoint mean evaluation reward: {best_policy_eval[\"mean_reward\"]:.3f}')\n", + "print(f'Random baseline mean evaluation reward: {random_policy_eval[\"mean_reward\"]:.3f}')\n", + "print(f'Latest checkpoint path: {latest_checkpoint_path}')\n", + "print(f'Best checkpoint path: {best_checkpoint_path}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Render the best checkpoint and save a GIF\n", + "\n", + "This section runs the **best restored PPO checkpoint** in evaluation mode with Bluebird radar rendering enabled.\n", + "It saves individual frames to disk and then combines them into a GIF.\n", + "\n", + "Notes:\n", + "\n", + "- `display_actions=True` overlays actions on the radar frames\n", + "- the GIF path is printed and displayed in the notebook\n", + "- by default this uses the first held-out evaluation seed" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "gif_seed = heldout_evaluation_seeds[0]\n", + "render_dir = Path.cwd() / 'renders' / 'minimal_ppo_agent_eval'\n", + "gif_path = render_evaluation_rollout_to_gif(\n", + " agent=agent,\n", + " random_seed=gif_seed,\n", + " render_dir=render_dir,\n", + " gif_name=f'best_ppo_checkpoint_eval_seed_{gif_seed}',\n", + ")\n", + "\n", + "print(f'Saved GIF to: {gif_path}')\n", + "display(Image(filename=str(gif_path)))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Optional cleanup\n", + "\n", + "Close the main environment if you are done with the notebook session." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "environment.close()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/bluebird-gymnasium/examples/minimal_ppo_flight_school.ipynb b/bluebird-gymnasium/examples/minimal_ppo_flight_school.ipynb new file mode 100644 index 0000000..1d98120 --- /dev/null +++ b/bluebird-gymnasium/examples/minimal_ppo_flight_school.ipynb @@ -0,0 +1,1035 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Minimal PPO Agent for Flight School\n", + "\n", + "This notebook mirrors the existing PPO example, but swaps the environment to **Flight School**.\n", + "\n", + "It keeps the same broad workflow:\n", + "\n", + "- small `FlightSchoolEnv` configuration\n", + "- centralized control with one observation vector and one discrete action per step\n", + "- periodic evaluation during training\n", + "- comparison against a random baseline\n", + "- checkpoint saving and best-model restore\n", + "- final GIF rendering of the best policy\n", + "\n", + "This is still educational rather than production-grade PPO, but it is much closer to what you would use in practice than plain REINFORCE.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Environment note\n", + "\n", + "This notebook assumes you are already running the **correct Python/Jupyter kernel**:\n", + "one with `gymnasium`, `torch`, and the Bluebird project dependencies installed.\n", + "\n", + "Do the package setup in your shell or project virtual environment first,\n", + "then open this notebook with that kernel. The notebook does **not** try to install\n", + "packages itself." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Imports and path setup\n", + "\n", + "This cell makes the notebook runnable from either the `bluebird-gymnasium` directory\n", + "or the repo root by adding the local package paths to `sys.path`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from __future__ import annotations\n", + "\n", + "import random\n", + "import shutil\n", + "import sys\n", + "from pathlib import Path\n", + "\n", + "import imageio\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import torch\n", + "import torch.nn as nn\n", + "import torch.optim as optim\n", + "from IPython.display import Image, display\n", + "\n", + "search_roots = [Path.cwd().resolve(), *Path.cwd().resolve().parents]\n", + "gym_root = None\n", + "dt_root = None\n", + "\n", + "for candidate in search_roots:\n", + " if (candidate / 'bluebird_gymnasium').exists():\n", + " gym_root = candidate\n", + " sibling_dt = candidate.parent / 'bluebird-dt'\n", + " if sibling_dt.exists():\n", + " dt_root = sibling_dt\n", + " break\n", + " if (candidate / 'bluebird-gymnasium').exists() and (candidate / 'bluebird-dt').exists():\n", + " gym_root = candidate / 'bluebird-gymnasium'\n", + " dt_root = candidate / 'bluebird-dt'\n", + " break\n", + "\n", + "if gym_root is None or dt_root is None:\n", + " raise RuntimeError('Could not locate local bluebird-gymnasium and bluebird-dt package roots.')\n", + "\n", + "sys.path.insert(0, str(gym_root))\n", + "sys.path.insert(0, str(dt_root))\n", + "\n", + "from bluebird_gymnasium.envs import EnvConfig, ViewType\n", + "from bluebird_gymnasium.envs.flight_school import FlightSchoolEnv\n", + "\n", + "print(f'Using bluebird-gymnasium from: {gym_root}')\n", + "print(f'Using bluebird-dt from: {dt_root}')\n", + "print(f'Torch version: {torch.__version__}')\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## PPO actor-critic model and agents\n", + "\n", + "The actor-critic network has:\n", + "\n", + "- a shared feature trunk\n", + "- a **policy head** that outputs action logits\n", + "- a **value head** that predicts the expected return from the current observation\n", + "\n", + "The PPO update then uses:\n", + "\n", + "- old action log-probabilities from rollout collection\n", + "- new action log-probabilities from the current policy\n", + "- a clipped objective to avoid excessively large policy updates\n", + "- a value loss for the critic\n", + "- a small entropy bonus to keep exploration from collapsing too early\n", + "\n", + "For Flight School, the policy is **centralized**:\n", + "\n", + "- input: one observation vector\n", + "- output: one discrete action\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "class ActorCriticNetwork(nn.Module):\n", + " \"\"\"Shared-trunk actor-critic network for PPO.\"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " observation_dimension: int,\n", + " number_of_actions: int,\n", + " hidden_units: int = 128,\n", + " ) -> None:\n", + " super().__init__()\n", + " self.trunk = nn.Sequential(\n", + " nn.Linear(observation_dimension, hidden_units),\n", + " nn.ReLU(),\n", + " nn.Linear(hidden_units, hidden_units),\n", + " nn.ReLU(),\n", + " )\n", + " self.policy_head = nn.Linear(hidden_units, number_of_actions)\n", + " self.value_head = nn.Linear(hidden_units, 1)\n", + "\n", + " def forward(self, observation_batch: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:\n", + " features = self.trunk(observation_batch)\n", + " action_logits = self.policy_head(features)\n", + " state_value = self.value_head(features).squeeze(-1)\n", + " return action_logits, state_value\n", + "\n", + "\n", + "class PPOAgent:\n", + " \"\"\"Minimal PPO agent with actor-critic network and checkpoint helpers.\"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " observation_dimension: int,\n", + " number_of_actions: int,\n", + " learning_rate: float = 3e-4,\n", + " hidden_units: int = 128,\n", + " clip_epsilon: float = 0.2,\n", + " value_loss_coefficient: float = 0.5,\n", + " entropy_coefficient: float = 0.01,\n", + " ppo_epochs: int = 4,\n", + " ) -> None:\n", + " self.actor_critic = ActorCriticNetwork(\n", + " observation_dimension=observation_dimension,\n", + " number_of_actions=number_of_actions,\n", + " hidden_units=hidden_units,\n", + " )\n", + " self.optimizer = optim.Adam(self.actor_critic.parameters(), lr=learning_rate)\n", + " self.clip_epsilon = clip_epsilon\n", + " self.value_loss_coefficient = value_loss_coefficient\n", + " self.entropy_coefficient = entropy_coefficient\n", + " self.ppo_epochs = ppo_epochs\n", + "\n", + " def choose_training_action(\n", + " self,\n", + " observation_vector: np.ndarray,\n", + " ) -> tuple[int, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:\n", + " observation_tensor = torch.tensor(\n", + " observation_vector,\n", + " dtype=torch.float32,\n", + " ).unsqueeze(0)\n", + " action_logits, state_value = self.actor_critic(observation_tensor)\n", + " action_distribution = torch.distributions.Categorical(logits=action_logits)\n", + " sampled_action = action_distribution.sample()\n", + " action_log_probability = action_distribution.log_prob(sampled_action)\n", + " action_entropy = action_distribution.entropy()\n", + "\n", + " return (\n", + " sampled_action.item(),\n", + " observation_tensor.squeeze(0),\n", + " action_log_probability.squeeze(0),\n", + " state_value.squeeze(0),\n", + " action_entropy.squeeze(0),\n", + " )\n", + "\n", + " def choose_evaluation_action(\n", + " self,\n", + " observation_vector: np.ndarray,\n", + " ) -> int:\n", + " with torch.no_grad():\n", + " observation_tensor = torch.tensor(\n", + " observation_vector,\n", + " dtype=torch.float32,\n", + " ).unsqueeze(0)\n", + " action_logits, _state_value = self.actor_critic(observation_tensor)\n", + " chosen_action = torch.argmax(action_logits, dim=-1).item()\n", + " return chosen_action\n", + "\n", + " def update_from_trajectory(\n", + " self,\n", + " observations: list[torch.Tensor],\n", + " actions: list[torch.Tensor],\n", + " old_log_probabilities: list[torch.Tensor],\n", + " returns: torch.Tensor,\n", + " advantages: torch.Tensor,\n", + " ) -> dict[str, float] | None:\n", + " if not observations:\n", + " return None\n", + "\n", + " observation_tensor = torch.stack(observations)\n", + " action_tensor = torch.stack(actions).long()\n", + " old_log_probability_tensor = torch.stack(old_log_probabilities).detach()\n", + " returns_tensor = returns.detach()\n", + " advantages_tensor = advantages.detach()\n", + "\n", + " if advantages_tensor.numel() > 1:\n", + " advantages_std = advantages_tensor.std(unbiased=False)\n", + " if advantages_std > 1e-8:\n", + " advantages_tensor = (\n", + " (advantages_tensor - advantages_tensor.mean())\n", + " / (advantages_std + 1e-8)\n", + " )\n", + "\n", + " mean_policy_loss = 0.0\n", + " mean_value_loss = 0.0\n", + " mean_entropy = 0.0\n", + " mean_total_loss = 0.0\n", + "\n", + " for _epoch in range(self.ppo_epochs):\n", + " new_action_logits, new_state_values = self.actor_critic(observation_tensor)\n", + " action_distribution = torch.distributions.Categorical(logits=new_action_logits)\n", + " new_log_probabilities = action_distribution.log_prob(action_tensor)\n", + " entropy = action_distribution.entropy().mean()\n", + "\n", + " probability_ratio = torch.exp(new_log_probabilities - old_log_probability_tensor)\n", + " unclipped_objective = probability_ratio * advantages_tensor\n", + " clipped_objective = torch.clamp(\n", + " probability_ratio,\n", + " 1.0 - self.clip_epsilon,\n", + " 1.0 + self.clip_epsilon,\n", + " ) * advantages_tensor\n", + "\n", + " policy_loss = -torch.min(unclipped_objective, clipped_objective).mean()\n", + " value_loss = torch.nn.functional.mse_loss(new_state_values, returns_tensor)\n", + " total_loss = (\n", + " policy_loss\n", + " + self.value_loss_coefficient * value_loss\n", + " - self.entropy_coefficient * entropy\n", + " )\n", + "\n", + " self.optimizer.zero_grad()\n", + " total_loss.backward()\n", + " self.optimizer.step()\n", + "\n", + " mean_policy_loss += float(policy_loss.item())\n", + " mean_value_loss += float(value_loss.item())\n", + " mean_entropy += float(entropy.item())\n", + " mean_total_loss += float(total_loss.item())\n", + "\n", + " epoch_divisor = float(self.ppo_epochs)\n", + " return {\n", + " 'policy_loss': mean_policy_loss / epoch_divisor,\n", + " 'value_loss': mean_value_loss / epoch_divisor,\n", + " 'entropy': mean_entropy / epoch_divisor,\n", + " 'total_loss': mean_total_loss / epoch_divisor,\n", + " }\n", + "\n", + " def save_checkpoint(self, checkpoint_path: Path, metadata: dict | None = None) -> None:\n", + " checkpoint_path.parent.mkdir(parents=True, exist_ok=True)\n", + " payload = {\n", + " 'model_state_dict': self.actor_critic.state_dict(),\n", + " 'optimizer_state_dict': self.optimizer.state_dict(),\n", + " 'metadata': metadata or {},\n", + " }\n", + " torch.save(payload, checkpoint_path)\n", + "\n", + " def load_checkpoint(self, checkpoint_path: Path, map_location: str = 'cpu') -> dict:\n", + " payload = torch.load(checkpoint_path, map_location=map_location)\n", + " self.actor_critic.load_state_dict(payload['model_state_dict'])\n", + " if 'optimizer_state_dict' in payload:\n", + " self.optimizer.load_state_dict(payload['optimizer_state_dict'])\n", + " return payload.get('metadata', {})\n", + "\n", + "\n", + "class RandomAgent:\n", + " \"\"\"Simple random baseline for comparison.\"\"\"\n", + "\n", + " def __init__(self, number_of_actions: int) -> None:\n", + " self.number_of_actions = number_of_actions\n", + "\n", + " def choose_evaluation_action(self, observation_vector: np.ndarray) -> int:\n", + " _ = observation_vector\n", + " return random.randrange(self.number_of_actions)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Environment configuration and rollout helpers\n", + "\n", + "This notebook uses the same kind of compact setup as the Flight School REINFORCE version:\n", + "\n", + "- `FlightSchoolEnv`\n", + "- centralized control\n", + "- one observation vector per step\n", + "- one discrete action per step\n", + "\n", + "That keeps the PPO implementation easier to follow while still showing the core mechanics.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def make_flight_school_training_config(random_seed: int | None = None) -> EnvConfig:\n", + " config = FlightSchoolEnv.get_default_env_config(ViewType.CENTRALIZED)\n", + " config.scenario_config['args']['random_seed'] = random_seed\n", + " config.scenario_duration = 10 * 60\n", + " config.view_config['type'] = ViewType.CENTRALIZED.value\n", + " return config\n", + "\n", + "\n", + "def compute_returns_and_advantages(\n", + " rewards: list[float],\n", + " values: list[torch.Tensor],\n", + " dones: list[bool],\n", + " discount_factor_gamma: float,\n", + " gae_lambda: float,\n", + ") -> tuple[torch.Tensor, torch.Tensor]:\n", + " rewards_tensor = torch.tensor(rewards, dtype=torch.float32)\n", + " values_tensor = torch.stack(values).detach().float()\n", + " dones_tensor = torch.tensor(dones, dtype=torch.float32)\n", + "\n", + " advantages = torch.zeros_like(rewards_tensor)\n", + " last_gae = torch.tensor(0.0)\n", + "\n", + " for timestep in reversed(range(len(rewards))):\n", + " if timestep == len(rewards) - 1:\n", + " next_value = torch.tensor(0.0)\n", + " else:\n", + " next_value = values_tensor[timestep + 1]\n", + "\n", + " next_nonterminal = 1.0 - dones_tensor[timestep]\n", + " delta = (\n", + " rewards_tensor[timestep]\n", + " + discount_factor_gamma * next_value * next_nonterminal\n", + " - values_tensor[timestep]\n", + " )\n", + " last_gae = delta + discount_factor_gamma * gae_lambda * next_nonterminal * last_gae\n", + " advantages[timestep] = last_gae\n", + "\n", + " returns = advantages + values_tensor\n", + " return returns, advantages\n", + "\n", + "\n", + "def run_one_training_episode(\n", + " environment: FlightSchoolEnv,\n", + " agent: PPOAgent,\n", + " random_seed: int,\n", + " discount_factor_gamma: float,\n", + " gae_lambda: float,\n", + ") -> tuple[float, int, dict[str, float] | None]:\n", + " random.seed(random_seed)\n", + " np.random.seed(random_seed)\n", + " torch.manual_seed(random_seed)\n", + " environment.config.scenario_config['args']['random_seed'] = random_seed\n", + "\n", + " observation_vector, _info = environment.reset(seed=random_seed)\n", + "\n", + " episode_is_done = False\n", + " episode_step_count = 0\n", + " episode_total_reward = 0.0\n", + "\n", + " observations: list[torch.Tensor] = []\n", + " actions: list[torch.Tensor] = []\n", + " old_log_probabilities: list[torch.Tensor] = []\n", + " values: list[torch.Tensor] = []\n", + " rewards: list[float] = []\n", + " dones: list[bool] = []\n", + "\n", + " while not episode_is_done:\n", + " (\n", + " action_int,\n", + " observation_tensor,\n", + " action_log_probability,\n", + " state_value,\n", + " _action_entropy,\n", + " ) = agent.choose_training_action(observation_vector)\n", + "\n", + " (\n", + " next_observation_vector,\n", + " reward,\n", + " done,\n", + " truncated,\n", + " _info,\n", + " ) = environment.step(action_int)\n", + "\n", + " timestep_reward = float(reward)\n", + " timestep_done = bool(done or truncated)\n", + "\n", + " observations.append(observation_tensor)\n", + " actions.append(torch.tensor(action_int))\n", + " old_log_probabilities.append(action_log_probability.detach())\n", + " values.append(state_value.detach())\n", + " rewards.append(timestep_reward)\n", + " dones.append(timestep_done)\n", + "\n", + " episode_total_reward += timestep_reward\n", + " episode_is_done = timestep_done\n", + " observation_vector = next_observation_vector\n", + " episode_step_count += 1\n", + "\n", + " returns, advantages = compute_returns_and_advantages(\n", + " rewards=rewards,\n", + " values=values,\n", + " dones=dones,\n", + " discount_factor_gamma=discount_factor_gamma,\n", + " gae_lambda=gae_lambda,\n", + " )\n", + "\n", + " update_metrics = agent.update_from_trajectory(\n", + " observations=observations,\n", + " actions=actions,\n", + " old_log_probabilities=old_log_probabilities,\n", + " returns=returns,\n", + " advantages=advantages,\n", + " )\n", + "\n", + " return episode_total_reward, episode_step_count, update_metrics\n", + "\n", + "\n", + "def run_one_evaluation_episode(\n", + " environment: FlightSchoolEnv,\n", + " evaluation_agent,\n", + " random_seed: int,\n", + ") -> tuple[float, int]:\n", + " random.seed(random_seed)\n", + " np.random.seed(random_seed)\n", + " torch.manual_seed(random_seed)\n", + " environment.config.scenario_config['args']['random_seed'] = random_seed\n", + "\n", + " observation_vector, _info = environment.reset(seed=random_seed)\n", + "\n", + " episode_is_done = False\n", + " episode_step_count = 0\n", + " episode_total_reward = 0.0\n", + "\n", + " while not episode_is_done:\n", + " action_int = evaluation_agent.choose_evaluation_action(observation_vector)\n", + " (\n", + " next_observation_vector,\n", + " reward,\n", + " done,\n", + " truncated,\n", + " _info,\n", + " ) = environment.step(action_int)\n", + "\n", + " episode_total_reward += float(reward)\n", + " episode_is_done = bool(done or truncated)\n", + " observation_vector = next_observation_vector\n", + " episode_step_count += 1\n", + "\n", + " return episode_total_reward, episode_step_count\n", + "\n", + "\n", + "def evaluate_agent_over_seeds(\n", + " environment: FlightSchoolEnv,\n", + " evaluation_agent,\n", + " evaluation_seeds: list[int],\n", + ") -> dict:\n", + " rewards: list[float] = []\n", + " steps: list[int] = []\n", + "\n", + " for random_seed in evaluation_seeds:\n", + " total_reward, step_count = run_one_evaluation_episode(\n", + " environment=environment,\n", + " evaluation_agent=evaluation_agent,\n", + " random_seed=random_seed,\n", + " )\n", + " rewards.append(total_reward)\n", + " steps.append(step_count)\n", + "\n", + " return {\n", + " 'seeds': evaluation_seeds,\n", + " 'rewards': rewards,\n", + " 'steps': steps,\n", + " 'mean_reward': float(np.mean(rewards)),\n", + " 'std_reward': float(np.std(rewards)),\n", + " 'mean_steps': float(np.mean(steps)),\n", + " }\n", + "\n", + "\n", + "def render_evaluation_rollout_to_gif(\n", + " agent: PPOAgent,\n", + " random_seed: int,\n", + " render_dir: Path,\n", + " gif_name: str = 'trained_policy_eval',\n", + " render_every_n_steps: int = 5,\n", + " gif_frame_duration_seconds: float = 0.2,\n", + ") -> Path:\n", + " render_config = make_flight_school_training_config(random_seed=random_seed)\n", + " render_config.radar_config['display_actions'] = True\n", + " render_config.radar_config['render_dir'] = str(render_dir)\n", + " render_config.radar_config['prefix'] = 'frame'\n", + "\n", + " if render_dir.exists():\n", + " shutil.rmtree(render_dir)\n", + " render_dir.mkdir(parents=True, exist_ok=True)\n", + "\n", + " render_environment = FlightSchoolEnv(config=render_config)\n", + " render_environment.set_render_mode('file')\n", + "\n", + " random.seed(random_seed)\n", + " np.random.seed(random_seed)\n", + " torch.manual_seed(random_seed)\n", + "\n", + " observation_vector, _info = render_environment.reset(seed=random_seed)\n", + " render_environment.render()\n", + "\n", + " episode_is_done = False\n", + " step_index = 0\n", + "\n", + " while not episode_is_done:\n", + " action_int = agent.choose_evaluation_action(observation_vector)\n", + " (\n", + " next_observation_vector,\n", + " reward,\n", + " done,\n", + " truncated,\n", + " _info,\n", + " ) = render_environment.step(action_int)\n", + " _ = reward\n", + " step_index += 1\n", + " episode_is_done = bool(done or truncated)\n", + " if step_index % render_every_n_steps == 0 or episode_is_done:\n", + " render_environment.render()\n", + " observation_vector = next_observation_vector\n", + "\n", + " png_frames = sorted(render_dir.glob(f\"{render_config.radar_config['prefix']}_*.png\"))\n", + " if not png_frames:\n", + " raise RuntimeError(\n", + " f'No rendered PNG frames were written to {render_dir}. '\n", + " 'Expected at least one frame before GIF generation.'\n", + " )\n", + "\n", + " gif_path = render_dir / f'{gif_name}.gif'\n", + " images = [imageio.v3.imread(frame_path) for frame_path in png_frames]\n", + " imageio.mimsave(gif_path, images, loop=0, duration=gif_frame_duration_seconds)\n", + " render_environment.close()\n", + " return gif_path\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Set up the environment and inspect the shapes\n", + "\n", + "For this notebook, the most important values are:\n", + "\n", + "- `observation_dimension`: how many numbers are in the centralized observation vector\n", + "- `number_of_actions`: how many discrete actions the PPO policy can choose from\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "initial_seed = 7\n", + "config = make_flight_school_training_config(random_seed=initial_seed)\n", + "environment = FlightSchoolEnv(config=config)\n", + "\n", + "observation_dimension = environment.observation_space.shape[0]\n", + "number_of_actions = environment.action_space.n\n", + "\n", + "print(\n", + " 'environment shapes:',\n", + " f'observation_dimension={observation_dimension}',\n", + " f'number_of_actions={number_of_actions}',\n", + ")\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Hyperparameters and experiment settings\n", + "\n", + "This version keeps the same richer experiment loop:\n", + "\n", + "- periodic evaluation during training\n", + "- larger held-out evaluation set\n", + "- checkpoint saving\n", + "- automatic tracking of the best evaluated model\n", + "- random-policy baseline comparison\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "learning_rate = 3e-4\n", + "hidden_units = 128\n", + "discount_factor_gamma = 0.99\n", + "gae_lambda = 0.95\n", + "clip_epsilon = 0.2\n", + "value_loss_coefficient = 0.5\n", + "entropy_coefficient = 0.01\n", + "ppo_epochs = 4\n", + "number_of_training_episodes = 100\n", + "training_seed_start = 100\n", + "periodic_eval_interval = 10\n", + "heldout_evaluation_seeds = list(range(200, 220))\n", + "checkpoint_dir = Path.cwd() / 'checkpoints' / 'minimal_ppo_flight_school'\n", + "latest_checkpoint_path = checkpoint_dir / 'latest.pt'\n", + "best_checkpoint_path = checkpoint_dir / 'best.pt'\n", + "\n", + "agent = PPOAgent(\n", + " observation_dimension=observation_dimension,\n", + " number_of_actions=number_of_actions,\n", + " learning_rate=learning_rate,\n", + " hidden_units=hidden_units,\n", + " clip_epsilon=clip_epsilon,\n", + " value_loss_coefficient=value_loss_coefficient,\n", + " entropy_coefficient=entropy_coefficient,\n", + " ppo_epochs=ppo_epochs,\n", + ")\n", + "random_agent = RandomAgent(number_of_actions=number_of_actions)\n", + "\n", + "training_rewards: list[float] = []\n", + "training_steps: list[int] = []\n", + "training_policy_losses: list[float] = []\n", + "training_value_losses: list[float] = []\n", + "training_entropies: list[float] = []\n", + "training_total_losses: list[float] = []\n", + "\n", + "periodic_eval_episodes: list[int] = []\n", + "periodic_eval_learned_mean_rewards: list[float] = []\n", + "periodic_eval_learned_std_rewards: list[float] = []\n", + "periodic_eval_random_mean_rewards: list[float] = []\n", + "periodic_eval_random_std_rewards: list[float] = []\n", + "periodic_eval_learned_mean_steps: list[float] = []\n", + "periodic_eval_random_mean_steps: list[float] = []\n", + "\n", + "best_mean_evaluation_reward = float('-inf')\n", + "best_checkpoint_metadata: dict = {}\n", + "render_every_n_steps = 5\n", + "gif_frame_duration_seconds = 0.2\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## PPO training loop with periodic evaluation and checkpointing\n", + "\n", + "Every `periodic_eval_interval` episodes, the notebook:\n", + "\n", + "- evaluates the current learned policy on the held-out evaluation seeds\n", + "- evaluates a random baseline on the same seeds\n", + "- saves a `latest.pt` checkpoint\n", + "- overwrites `best.pt` if the learned policy achieves a new best mean evaluation reward" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "for episode_index in range(number_of_training_episodes):\n", + " random_seed = training_seed_start + episode_index\n", + " total_reward, step_count, update_metrics = run_one_training_episode(\n", + " environment=environment,\n", + " agent=agent,\n", + " random_seed=random_seed,\n", + " discount_factor_gamma=discount_factor_gamma,\n", + " gae_lambda=gae_lambda,\n", + " )\n", + "\n", + " training_rewards.append(total_reward)\n", + " training_steps.append(step_count)\n", + " training_policy_losses.append(float('nan') if update_metrics is None else update_metrics['policy_loss'])\n", + " training_value_losses.append(float('nan') if update_metrics is None else update_metrics['value_loss'])\n", + " training_entropies.append(float('nan') if update_metrics is None else update_metrics['entropy'])\n", + " training_total_losses.append(float('nan') if update_metrics is None else update_metrics['total_loss'])\n", + "\n", + " print(\n", + " '[train]',\n", + " f'episode={episode_index:03d}',\n", + " f'seed={random_seed}',\n", + " f'reward={total_reward:.3f}',\n", + " f'steps={step_count}',\n", + " f'policy_loss={None if update_metrics is None else update_metrics[\"policy_loss\"]:.6f}' if update_metrics is not None else 'policy_loss=None',\n", + " f'value_loss={None if update_metrics is None else update_metrics[\"value_loss\"]:.6f}' if update_metrics is not None else 'value_loss=None',\n", + " )\n", + "\n", + " should_run_periodic_eval = (\n", + " (episode_index + 1) % periodic_eval_interval == 0\n", + " or episode_index == number_of_training_episodes - 1\n", + " )\n", + "\n", + " if should_run_periodic_eval:\n", + " learned_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + " )\n", + " random_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=random_agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + " )\n", + "\n", + " periodic_eval_episodes.append(episode_index + 1)\n", + " periodic_eval_learned_mean_rewards.append(learned_eval['mean_reward'])\n", + " periodic_eval_learned_std_rewards.append(learned_eval['std_reward'])\n", + " periodic_eval_random_mean_rewards.append(random_eval['mean_reward'])\n", + " periodic_eval_random_std_rewards.append(random_eval['std_reward'])\n", + " periodic_eval_learned_mean_steps.append(learned_eval['mean_steps'])\n", + " periodic_eval_random_mean_steps.append(random_eval['mean_steps'])\n", + "\n", + " metadata = {\n", + " 'episode': episode_index + 1,\n", + " 'train_seed': random_seed,\n", + " 'learned_mean_reward': learned_eval['mean_reward'],\n", + " 'learned_std_reward': learned_eval['std_reward'],\n", + " 'random_mean_reward': random_eval['mean_reward'],\n", + " 'random_std_reward': random_eval['std_reward'],\n", + " 'evaluation_seeds': heldout_evaluation_seeds,\n", + " }\n", + " agent.save_checkpoint(latest_checkpoint_path, metadata=metadata)\n", + "\n", + " if learned_eval['mean_reward'] > best_mean_evaluation_reward:\n", + " best_mean_evaluation_reward = learned_eval['mean_reward']\n", + " best_checkpoint_metadata = metadata\n", + " agent.save_checkpoint(best_checkpoint_path, metadata=metadata)\n", + " checkpoint_note = 'new best checkpoint'\n", + " else:\n", + " checkpoint_note = 'latest checkpoint only'\n", + "\n", + " print(\n", + " '[periodic-eval]',\n", + " f'episode={episode_index + 1:03d}',\n", + " f'learned_mean_reward={learned_eval[\"mean_reward\"]:.3f}',\n", + " f'learned_std_reward={learned_eval[\"std_reward\"]:.3f}',\n", + " f'random_mean_reward={random_eval[\"mean_reward\"]:.3f}',\n", + " f'random_std_reward={random_eval[\"std_reward\"]:.3f}',\n", + " checkpoint_note,\n", + " )" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Restore the best evaluated model\n", + "\n", + "The training loop may end on a policy that is not the best one seen so far.\n", + "This cell reloads the checkpoint with the highest held-out mean evaluation reward." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "if not best_checkpoint_path.exists():\n", + " raise FileNotFoundError(f'Best checkpoint not found: {best_checkpoint_path}')\n", + "\n", + "loaded_metadata = agent.load_checkpoint(best_checkpoint_path)\n", + "print('Reloaded best checkpoint from:', best_checkpoint_path)\n", + "print('Best checkpoint metadata:')\n", + "loaded_metadata" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Final evaluation of the best checkpoint vs random baseline\n", + "\n", + "This uses the larger held-out evaluation set and compares:\n", + "\n", + "- the best learned PPO checkpoint\n", + "- a random baseline on the same seeds\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "best_policy_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + ")\n", + "random_policy_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=random_agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + ")\n", + "\n", + "print('Best PPO policy evaluation mean reward:', best_policy_eval['mean_reward'])\n", + "print('Best PPO policy evaluation std reward:', best_policy_eval['std_reward'])\n", + "print('Random baseline mean reward:', random_policy_eval['mean_reward'])\n", + "print('Random baseline std reward:', random_policy_eval['std_reward'])\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Useful plots\n", + "\n", + "These are the most useful quick-look plots for this Flight School PPO setup:\n", + "\n", + "- **Training reward**: raw reward and moving average during training\n", + "- **Episode length**: how long training episodes run\n", + "- **Periodic evaluation**: learned policy vs random baseline over training\n", + "- **Final evaluation by seed**: best learned checkpoint vs random on the held-out seeds\n", + "- **PPO diagnostics**: policy loss, value loss, total loss, and entropy\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def moving_average(values: list[float], window: int) -> np.ndarray:\n", + " if len(values) < window:\n", + " return np.array([])\n", + " kernel = np.ones(window) / window\n", + " return np.convolve(np.asarray(values, dtype=float), kernel, mode='valid')\n", + "\n", + "plot_window = min(10, len(training_rewards))\n", + "smoothed_rewards = moving_average(training_rewards, plot_window)\n", + "training_episode_indices = np.arange(1, len(training_rewards) + 1)\n", + "heldout_seed_indices = np.arange(len(heldout_evaluation_seeds))\n", + "periodic_eval_episodes_arr = np.asarray(periodic_eval_episodes)\n", + "learned_mean_arr = np.asarray(periodic_eval_learned_mean_rewards)\n", + "learned_std_arr = np.asarray(periodic_eval_learned_std_rewards)\n", + "random_mean_arr = np.asarray(periodic_eval_random_mean_rewards)\n", + "random_std_arr = np.asarray(periodic_eval_random_std_rewards)\n", + "\n", + "fig, axes = plt.subplots(3, 2, figsize=(15, 14))\n", + "\n", + "axes[0, 0].plot(training_episode_indices, training_rewards, marker='o', alpha=0.25, label='raw reward')\n", + "if len(smoothed_rewards) > 0:\n", + " axes[0, 0].plot(\n", + " np.arange(plot_window, len(training_rewards) + 1),\n", + " smoothed_rewards,\n", + " linewidth=2.5,\n", + " color='tab:blue',\n", + " label=f'moving average (window={plot_window})',\n", + " )\n", + "axes[0, 0].set_title('Training Reward per Episode')\n", + "axes[0, 0].set_xlabel('Episode')\n", + "axes[0, 0].set_ylabel('Total Reward')\n", + "axes[0, 0].legend()\n", + "axes[0, 0].grid(alpha=0.3)\n", + "\n", + "axes[0, 1].plot(training_episode_indices, training_steps, marker='o', color='tab:orange')\n", + "axes[0, 1].set_title('Training Episode Length')\n", + "axes[0, 1].set_xlabel('Episode')\n", + "axes[0, 1].set_ylabel('Steps')\n", + "axes[0, 1].grid(alpha=0.3)\n", + "\n", + "axes[1, 0].plot(periodic_eval_episodes_arr, learned_mean_arr, marker='o', label='learned policy')\n", + "axes[1, 0].fill_between(\n", + " periodic_eval_episodes_arr,\n", + " learned_mean_arr - learned_std_arr,\n", + " learned_mean_arr + learned_std_arr,\n", + " alpha=0.2,\n", + ")\n", + "axes[1, 0].plot(periodic_eval_episodes_arr, random_mean_arr, marker='s', label='random baseline')\n", + "axes[1, 0].fill_between(\n", + " periodic_eval_episodes_arr,\n", + " random_mean_arr - random_std_arr,\n", + " random_mean_arr + random_std_arr,\n", + " alpha=0.2,\n", + ")\n", + "axes[1, 0].set_title('Periodic Evaluation: Learned vs Random')\n", + "axes[1, 0].set_xlabel('Training Episode')\n", + "axes[1, 0].set_ylabel('Mean Evaluation Reward')\n", + "axes[1, 0].legend()\n", + "axes[1, 0].grid(alpha=0.3)\n", + "\n", + "axes[1, 1].plot(\n", + " heldout_seed_indices,\n", + " best_policy_eval['rewards'],\n", + " marker='o',\n", + " linewidth=2,\n", + " label='best PPO checkpoint',\n", + ")\n", + "axes[1, 1].plot(\n", + " heldout_seed_indices,\n", + " random_policy_eval['rewards'],\n", + " marker='s',\n", + " linewidth=2,\n", + " label='random baseline',\n", + ")\n", + "axes[1, 1].set_xticks(heldout_seed_indices)\n", + "axes[1, 1].set_xticklabels([str(seed) for seed in heldout_evaluation_seeds], rotation=45)\n", + "axes[1, 1].set_title('Final Evaluation Reward by Seed')\n", + "axes[1, 1].set_xlabel('Held-out Evaluation Seed')\n", + "axes[1, 1].set_ylabel('Total Reward')\n", + "axes[1, 1].legend()\n", + "axes[1, 1].grid(alpha=0.3)\n", + "\n", + "axes[2, 0].plot(training_episode_indices, training_policy_losses, label='policy loss')\n", + "axes[2, 0].plot(training_episode_indices, training_value_losses, label='value loss')\n", + "axes[2, 0].plot(training_episode_indices, training_total_losses, label='total loss')\n", + "axes[2, 0].set_title('PPO Loss Terms')\n", + "axes[2, 0].set_xlabel('Episode')\n", + "axes[2, 0].set_ylabel('Loss')\n", + "axes[2, 0].legend()\n", + "axes[2, 0].grid(alpha=0.3)\n", + "\n", + "axes[2, 1].plot(training_episode_indices, training_entropies, label='entropy', color='tab:green')\n", + "axes[2, 1].set_title('Policy Entropy During Training')\n", + "axes[2, 1].set_xlabel('Episode')\n", + "axes[2, 1].set_ylabel('Entropy')\n", + "axes[2, 1].grid(alpha=0.3)\n", + "\n", + "fig.suptitle('Minimal PPO Training Summary for Flight School', fontsize=16)\n", + "fig.tight_layout()\n", + "plt.show()\n", + "\n", + "print(f'Best PPO checkpoint mean evaluation reward: {best_policy_eval[\"mean_reward\"]:.3f}')\n", + "print(f'Random baseline mean evaluation reward: {random_policy_eval[\"mean_reward\"]:.3f}')\n", + "print(f'Latest checkpoint path: {latest_checkpoint_path}')\n", + "print(f'Best checkpoint path: {best_checkpoint_path}')\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Render the best checkpoint and save a GIF\n", + "\n", + "This section runs the **best restored checkpoint** in evaluation mode with Bluebird radar rendering enabled.\n", + "It saves individual frames to disk and then combines them into a GIF.\n", + "\n", + "Notes:\n", + "\n", + "- `display_actions=True` overlays actions on the radar frames\n", + "- the GIF path is printed and displayed in the notebook\n", + "- by default this uses the first held-out evaluation seed\n", + "- `render_every_n_steps` controls how many simulator steps to skip between saved frames\n", + "- `gif_frame_duration_seconds` controls how long each GIF frame is shown\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "gif_seed = heldout_evaluation_seeds[0]\n", + "render_dir = Path.cwd() / 'renders' / 'minimal_ppo_flight_school_eval'\n", + "gif_path = render_evaluation_rollout_to_gif(\n", + " agent=agent,\n", + " random_seed=gif_seed,\n", + " render_dir=render_dir,\n", + " gif_name=f'best_ppo_flight_school_checkpoint_eval_seed_{gif_seed}',\n", + " render_every_n_steps=render_every_n_steps,\n", + " gif_frame_duration_seconds=gif_frame_duration_seconds,\n", + ")\n", + "\n", + "print(f'Saved GIF to: {gif_path}')\n", + "display(Image(filename=str(gif_path)))\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Optional cleanup\n", + "\n", + "Close the main environment if you are done with the notebook session." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "environment.close()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/bluebird-gymnasium/examples/minimal_reinforce_agent.ipynb b/bluebird-gymnasium/examples/minimal_reinforce_agent.ipynb new file mode 100644 index 0000000..ea2d1b3 --- /dev/null +++ b/bluebird-gymnasium/examples/minimal_reinforce_agent.ipynb @@ -0,0 +1,966 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Minimal REINFORCE Agent for Bluebird Gymnasium\n", + "\n", + "This notebook is the notebook companion to `examples/minimal_reinforce_agent.py`.\n", + "\n", + "It is intentionally educational rather than production-grade. It shows:\n", + "\n", + "- how to configure a small Bluebird environment\n", + "- how a policy network maps observations to logits\n", + "- how stochastic action sampling works during training\n", + "- how episode rewards become discounted returns\n", + "- how loss, backpropagation, and optimizer updates fit together\n", + "- how to visualize learning progress\n", + "- how to compare learned performance against a random baseline\n", + "- how to run periodic evaluation, save checkpoints, restore the best model, and render a GIF" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Environment note\n", + "\n", + "This notebook assumes you are already running the **correct Python/Jupyter kernel**:\n", + "one with `gymnasium`, `torch`, and the Bluebird project dependencies installed.\n", + "\n", + "Do the package setup in your shell or project virtual environment first,\n", + "then open this notebook with that kernel. The notebook does **not** try to install\n", + "packages itself." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Imports and path setup\n", + "\n", + "This cell makes the notebook runnable from either the `bluebird-gymnasium` directory\n", + "or the repo root by adding the local package paths to `sys.path`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from __future__ import annotations\n", + "\n", + "import random\n", + "import shutil\n", + "import sys\n", + "from pathlib import Path\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import torch\n", + "import torch.nn as nn\n", + "import torch.optim as optim\n", + "from IPython.display import Image, display\n", + "\n", + "search_roots = [Path.cwd().resolve(), *Path.cwd().resolve().parents]\n", + "gym_root = None\n", + "dt_root = None\n", + "\n", + "for candidate in search_roots:\n", + " if (candidate / 'bluebird_gymnasium').exists():\n", + " gym_root = candidate\n", + " sibling_dt = candidate.parent / 'bluebird-dt'\n", + " if sibling_dt.exists():\n", + " dt_root = sibling_dt\n", + " break\n", + " if (candidate / 'bluebird-gymnasium').exists() and (candidate / 'bluebird-dt').exists():\n", + " gym_root = candidate / 'bluebird-gymnasium'\n", + " dt_root = candidate / 'bluebird-dt'\n", + " break\n", + "\n", + "if gym_root is None or dt_root is None:\n", + " raise RuntimeError('Could not locate local bluebird-gymnasium and bluebird-dt package roots.')\n", + "\n", + "sys.path.insert(0, str(gym_root))\n", + "sys.path.insert(0, str(dt_root))\n", + "\n", + "from bluebird_gymnasium.envs import EnvConfig, ViewType\n", + "from bluebird_gymnasium.envs.sector_i import SectorIEnv\n", + "from bluebird_gymnasium.utils.video import generate_video\n", + "\n", + "print(f'Using bluebird-gymnasium from: {gym_root}')\n", + "print(f'Using bluebird-dt from: {dt_root}')\n", + "print(f'Torch version: {torch.__version__}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Policy networks and agents\n", + "\n", + "This notebook uses two simple agents:\n", + "\n", + "- `SharedPolicyAgent`: a trainable neural-network policy\n", + "- `RandomAgent`: a fixed baseline for comparison\n", + "\n", + "The policy network outputs one **logit** per action. During training we sample\n", + "from the categorical distribution defined by those logits. During evaluation we\n", + "use deterministic `argmax` action selection." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "class PolicyNetwork(nn.Module):\n", + " \"\"\"Neural network that maps one aircraft observation to action logits.\"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " observation_dimension: int,\n", + " number_of_actions: int,\n", + " hidden_units: int = 128,\n", + " ) -> None:\n", + " super().__init__()\n", + " self.layers = nn.Sequential(\n", + " nn.Linear(observation_dimension, hidden_units),\n", + " nn.ReLU(),\n", + " nn.Linear(hidden_units, hidden_units),\n", + " nn.ReLU(),\n", + " nn.Linear(hidden_units, number_of_actions),\n", + " )\n", + "\n", + " def forward(self, observation_batch: torch.Tensor) -> torch.Tensor:\n", + " return self.layers(observation_batch)\n", + "\n", + "\n", + "class SharedPolicyAgent:\n", + " \"\"\"One shared policy network reused for every aircraft.\"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " observation_dimension: int,\n", + " number_of_actions: int,\n", + " learning_rate: float = 1e-3,\n", + " hidden_units: int = 128,\n", + " ) -> None:\n", + " self.policy_network = PolicyNetwork(\n", + " observation_dimension=observation_dimension,\n", + " number_of_actions=number_of_actions,\n", + " hidden_units=hidden_units,\n", + " )\n", + " self.optimizer = optim.Adam(\n", + " self.policy_network.parameters(),\n", + " lr=learning_rate,\n", + " )\n", + "\n", + " def sample_training_actions(\n", + " self,\n", + " observation_by_callsign: dict[str, np.ndarray],\n", + " ) -> tuple[dict[str, int], dict[str, torch.Tensor]]:\n", + " action_by_callsign: dict[str, int] = {}\n", + " log_probability_by_callsign: dict[str, torch.Tensor] = {}\n", + "\n", + " for callsign, observation_vector in observation_by_callsign.items():\n", + " observation_tensor = torch.tensor(\n", + " observation_vector,\n", + " dtype=torch.float32,\n", + " ).unsqueeze(0)\n", + " action_logits = self.policy_network(observation_tensor)\n", + " action_distribution = torch.distributions.Categorical(\n", + " logits=action_logits,\n", + " )\n", + " sampled_action = action_distribution.sample()\n", + " sampled_action_log_probability = action_distribution.log_prob(\n", + " sampled_action,\n", + " )\n", + "\n", + " action_by_callsign[callsign] = sampled_action.item()\n", + " log_probability_by_callsign[callsign] = (\n", + " sampled_action_log_probability.squeeze()\n", + " )\n", + "\n", + " return action_by_callsign, log_probability_by_callsign\n", + "\n", + " def choose_evaluation_actions(\n", + " self,\n", + " observation_by_callsign: dict[str, np.ndarray],\n", + " ) -> dict[str, int]:\n", + " action_by_callsign: dict[str, int] = {}\n", + "\n", + " with torch.no_grad():\n", + " for callsign, observation_vector in observation_by_callsign.items():\n", + " observation_tensor = torch.tensor(\n", + " observation_vector,\n", + " dtype=torch.float32,\n", + " ).unsqueeze(0)\n", + " action_logits = self.policy_network(observation_tensor)\n", + " chosen_action = torch.argmax(action_logits, dim=-1).item()\n", + " action_by_callsign[callsign] = chosen_action\n", + "\n", + " return action_by_callsign\n", + "\n", + " def update_policy_from_episode(\n", + " self,\n", + " log_probability_per_step: list[torch.Tensor],\n", + " reward_per_step: list[float],\n", + " discount_factor_gamma: float = 0.99,\n", + " ) -> float | None:\n", + " if not log_probability_per_step:\n", + " return None\n", + "\n", + " discounted_return_per_step: list[float] = []\n", + " running_discounted_return = 0.0\n", + "\n", + " for reward in reversed(reward_per_step):\n", + " running_discounted_return = (\n", + " reward + discount_factor_gamma * running_discounted_return\n", + " )\n", + " discounted_return_per_step.append(running_discounted_return)\n", + "\n", + " discounted_return_per_step.reverse()\n", + " returns_tensor = torch.tensor(\n", + " discounted_return_per_step,\n", + " dtype=torch.float32,\n", + " )\n", + "\n", + " if returns_tensor.numel() > 1:\n", + " returns_std = returns_tensor.std(unbiased=False)\n", + " if returns_std > 1e-8:\n", + " returns_tensor = (\n", + " (returns_tensor - returns_tensor.mean())\n", + " / (returns_std + 1e-8)\n", + " )\n", + "\n", + " policy_loss_terms: list[torch.Tensor] = []\n", + " for action_log_probability, discounted_return in zip(\n", + " log_probability_per_step,\n", + " returns_tensor,\n", + " ):\n", + " policy_loss_terms.append(-action_log_probability * discounted_return)\n", + "\n", + " loss = torch.stack(policy_loss_terms).sum()\n", + " self.optimizer.zero_grad()\n", + " loss.backward()\n", + " self.optimizer.step()\n", + " return loss.item()\n", + "\n", + " def save_checkpoint(self, checkpoint_path: Path, metadata: dict | None = None) -> None:\n", + " checkpoint_path.parent.mkdir(parents=True, exist_ok=True)\n", + " payload = {\n", + " 'policy_state_dict': self.policy_network.state_dict(),\n", + " 'optimizer_state_dict': self.optimizer.state_dict(),\n", + " 'metadata': metadata or {},\n", + " }\n", + " torch.save(payload, checkpoint_path)\n", + "\n", + " def load_checkpoint(self, checkpoint_path: Path, map_location: str = 'cpu') -> dict:\n", + " payload = torch.load(checkpoint_path, map_location=map_location)\n", + " self.policy_network.load_state_dict(payload['policy_state_dict'])\n", + " if 'optimizer_state_dict' in payload:\n", + " self.optimizer.load_state_dict(payload['optimizer_state_dict'])\n", + " return payload.get('metadata', {})\n", + "\n", + "\n", + "class RandomAgent:\n", + " \"\"\"Simple random baseline for comparison.\"\"\"\n", + "\n", + " def __init__(self, number_of_actions: int) -> None:\n", + " self.number_of_actions = number_of_actions\n", + "\n", + " def choose_evaluation_actions(\n", + " self,\n", + " observation_by_callsign: dict[str, np.ndarray],\n", + " ) -> dict[str, int]:\n", + " return {\n", + " callsign: random.randrange(self.number_of_actions)\n", + " for callsign in observation_by_callsign.keys()\n", + " }" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Environment configuration and rollout helpers\n", + "\n", + "This configuration keeps the task deliberately small:\n", + "\n", + "- `SectorIEnv`\n", + "- decentralized control\n", + "- one aircraft\n", + "- `extra_minimal` state encoding\n", + "- lateral actions only\n", + "\n", + "The helper functions below support:\n", + "\n", + "- one training episode\n", + "- one evaluation episode\n", + "- evaluation over many seeds\n", + "- GIF rendering for a final evaluation rollout" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def make_sector_i_training_config() -> EnvConfig:\n", + " config = SectorIEnv.get_default_env_config(ViewType.DECENTRALIZED)\n", + "\n", + " config.state_repr_config = {\n", + " 'encoder_cls': 'extra_minimal',\n", + " 'k_nearest_aircraft': 1,\n", + " }\n", + "\n", + " config.action_config = {\n", + " 'simple_heading_left': True,\n", + " 'simple_heading_right': True,\n", + " 'simple_fl_climb': False,\n", + " 'simple_fl_descent': False,\n", + " 'simple_fl_exit': False,\n", + " }\n", + "\n", + " config.reward_config = {\n", + " 'fns': [\n", + " 'position_status_const',\n", + " 'lateral_centreline_distance_shaped',\n", + " 'safety_simple_avoidance_exp',\n", + " ],\n", + " 'coeffs': [1.0, 1.0, 1.2],\n", + " }\n", + "\n", + " config.view_config = {\n", + " 'type': ViewType.DECENTRALIZED.value,\n", + " 'decentralized_params': {},\n", + " }\n", + "\n", + " config.scenario_config = {\n", + " 'cls': 'tactical',\n", + " 'args': {\n", + " 'num_aircraft': 1,\n", + " 'balance': [0.0, 0.0, 1.0],\n", + " },\n", + " }\n", + "\n", + " return config\n", + "\n", + "\n", + "def run_one_training_episode(\n", + " environment: SectorIEnv,\n", + " agent: SharedPolicyAgent,\n", + " random_seed: int,\n", + " discount_factor_gamma: float,\n", + ") -> tuple[float, int, float | None]:\n", + " random.seed(random_seed)\n", + " np.random.seed(random_seed)\n", + " torch.manual_seed(random_seed)\n", + "\n", + " observation_by_callsign, _info = environment.reset(seed=random_seed)\n", + "\n", + " episode_is_done = False\n", + " episode_step_count = 0\n", + " episode_total_reward = 0.0\n", + "\n", + " log_probability_per_step: list[torch.Tensor] = []\n", + " reward_per_step: list[float] = []\n", + "\n", + " while not episode_is_done:\n", + " action_by_callsign, log_probability_by_callsign = (\n", + " agent.sample_training_actions(observation_by_callsign)\n", + " )\n", + "\n", + " (\n", + " next_observation_by_callsign,\n", + " reward_by_callsign,\n", + " done_by_callsign,\n", + " truncated_by_callsign,\n", + " _info,\n", + " ) = environment.step(action_by_callsign)\n", + "\n", + " timestep_reward = (\n", + " float(sum(reward_by_callsign.values())) if reward_by_callsign else 0.0\n", + " )\n", + "\n", + " if log_probability_by_callsign:\n", + " timestep_log_probability = torch.stack(\n", + " list(log_probability_by_callsign.values())\n", + " ).sum()\n", + " log_probability_per_step.append(timestep_log_probability)\n", + " reward_per_step.append(timestep_reward)\n", + "\n", + " episode_total_reward += timestep_reward\n", + " _ = truncated_by_callsign\n", + " episode_is_done = all(done_by_callsign.values()) if done_by_callsign else True\n", + " observation_by_callsign = next_observation_by_callsign\n", + " episode_step_count += 1\n", + "\n", + " loss_value = agent.update_policy_from_episode(\n", + " log_probability_per_step=log_probability_per_step,\n", + " reward_per_step=reward_per_step,\n", + " discount_factor_gamma=discount_factor_gamma,\n", + " )\n", + "\n", + " return episode_total_reward, episode_step_count, loss_value\n", + "\n", + "\n", + "def run_one_evaluation_episode(\n", + " environment: SectorIEnv,\n", + " evaluation_agent,\n", + " random_seed: int,\n", + ") -> tuple[float, int]:\n", + " random.seed(random_seed)\n", + " np.random.seed(random_seed)\n", + " torch.manual_seed(random_seed)\n", + "\n", + " observation_by_callsign, _info = environment.reset(seed=random_seed)\n", + "\n", + " episode_is_done = False\n", + " episode_step_count = 0\n", + " episode_total_reward = 0.0\n", + "\n", + " while not episode_is_done:\n", + " action_by_callsign = evaluation_agent.choose_evaluation_actions(\n", + " observation_by_callsign,\n", + " )\n", + "\n", + " (\n", + " next_observation_by_callsign,\n", + " reward_by_callsign,\n", + " done_by_callsign,\n", + " truncated_by_callsign,\n", + " _info,\n", + " ) = environment.step(action_by_callsign)\n", + "\n", + " _ = truncated_by_callsign\n", + " episode_total_reward += (\n", + " float(sum(reward_by_callsign.values())) if reward_by_callsign else 0.0\n", + " )\n", + " episode_is_done = all(done_by_callsign.values()) if done_by_callsign else True\n", + " observation_by_callsign = next_observation_by_callsign\n", + " episode_step_count += 1\n", + "\n", + " return episode_total_reward, episode_step_count\n", + "\n", + "\n", + "def evaluate_agent_over_seeds(\n", + " environment: SectorIEnv,\n", + " evaluation_agent,\n", + " evaluation_seeds: list[int],\n", + ") -> dict:\n", + " rewards: list[float] = []\n", + " steps: list[int] = []\n", + "\n", + " for random_seed in evaluation_seeds:\n", + " total_reward, step_count = run_one_evaluation_episode(\n", + " environment=environment,\n", + " evaluation_agent=evaluation_agent,\n", + " random_seed=random_seed,\n", + " )\n", + " rewards.append(total_reward)\n", + " steps.append(step_count)\n", + "\n", + " return {\n", + " 'seeds': evaluation_seeds,\n", + " 'rewards': rewards,\n", + " 'steps': steps,\n", + " 'mean_reward': float(np.mean(rewards)),\n", + " 'std_reward': float(np.std(rewards)),\n", + " 'mean_steps': float(np.mean(steps)),\n", + " }\n", + "\n", + "\n", + "def render_evaluation_rollout_to_gif(\n", + " agent: SharedPolicyAgent,\n", + " random_seed: int,\n", + " render_dir: Path,\n", + " gif_name: str = 'trained_policy_eval',\n", + ") -> Path:\n", + " render_config = make_sector_i_training_config()\n", + " render_config.radar_config['display_actions'] = True\n", + " render_config.radar_config['render_dir'] = str(render_dir)\n", + " render_config.radar_config['prefix'] = 'frame'\n", + "\n", + " if render_dir.exists():\n", + " shutil.rmtree(render_dir)\n", + " render_dir.mkdir(parents=True, exist_ok=True)\n", + "\n", + " render_environment = SectorIEnv(config=render_config)\n", + " render_environment.set_render_mode('file')\n", + "\n", + " random.seed(random_seed)\n", + " np.random.seed(random_seed)\n", + " torch.manual_seed(random_seed)\n", + "\n", + " observation_by_callsign, _info = render_environment.reset(seed=random_seed)\n", + "\n", + " # In this Bluebird branch, file-mode rendering is not triggered\n", + " # automatically by reset() or step(). Call render() explicitly.\n", + " render_environment.render()\n", + "\n", + " episode_is_done = False\n", + "\n", + " while not episode_is_done:\n", + " action_by_callsign = agent.choose_evaluation_actions(observation_by_callsign)\n", + " (\n", + " next_observation_by_callsign,\n", + " reward_by_callsign,\n", + " done_by_callsign,\n", + " truncated_by_callsign,\n", + " _info,\n", + " ) = render_environment.step(action_by_callsign)\n", + " _ = reward_by_callsign, truncated_by_callsign\n", + "\n", + " # Save one frame after every evaluation step.\n", + " render_environment.render()\n", + "\n", + " episode_is_done = all(done_by_callsign.values()) if done_by_callsign else True\n", + " observation_by_callsign = next_observation_by_callsign\n", + "\n", + " png_frames = sorted(render_dir.glob(f\"{render_config.radar_config['prefix']}_*.png\"))\n", + " if not png_frames:\n", + " raise RuntimeError(\n", + " f'No rendered PNG frames were written to {render_dir}. '\n", + " 'Expected at least one frame before GIF generation.'\n", + " )\n", + "\n", + " generate_video(\n", + " render_dir=str(render_dir),\n", + " frame_prefix=render_config.radar_config['prefix'],\n", + " video_filename=gif_name,\n", + " clean_up=False,\n", + " )\n", + " render_environment.close()\n", + " return render_dir / f'{gif_name}.gif'\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Set up the environment and inspect the shapes\n", + "\n", + "For this notebook, the most important values are:\n", + "\n", + "- `observation_dimension`: how many numbers are in one aircraft observation vector\n", + "- `number_of_actions`: how many discrete actions the policy can choose from" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = make_sector_i_training_config()\n", + "environment = SectorIEnv(config=config)\n", + "\n", + "observation_dimension = environment.observation_space.shape[0]\n", + "number_of_actions = environment.action_space.n\n", + "\n", + "print(\n", + " 'environment shapes:',\n", + " f'observation_dimension={observation_dimension}',\n", + " f'number_of_actions={number_of_actions}',\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Hyperparameters and experiment settings\n", + "\n", + "This version of the notebook adds:\n", + "\n", + "- periodic evaluation during training\n", + "- a larger held-out evaluation set\n", + "- checkpoint saving\n", + "- automatic tracking of the best evaluated model\n", + "- a random-policy baseline" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "learning_rate = 1e-3\n", + "hidden_units = 128\n", + "discount_factor_gamma = 0.99\n", + "number_of_training_episodes = 100\n", + "training_seed_start = 100\n", + "periodic_eval_interval = 10\n", + "heldout_evaluation_seeds = list(range(200, 220))\n", + "checkpoint_dir = Path.cwd() / 'checkpoints' / 'minimal_reinforce_agent'\n", + "latest_checkpoint_path = checkpoint_dir / 'latest.pt'\n", + "best_checkpoint_path = checkpoint_dir / 'best.pt'\n", + "\n", + "agent = SharedPolicyAgent(\n", + " observation_dimension=observation_dimension,\n", + " number_of_actions=number_of_actions,\n", + " learning_rate=learning_rate,\n", + " hidden_units=hidden_units,\n", + ")\n", + "random_agent = RandomAgent(number_of_actions=number_of_actions)\n", + "\n", + "training_rewards: list[float] = []\n", + "training_steps: list[int] = []\n", + "training_losses: list[float] = []\n", + "\n", + "periodic_eval_episodes: list[int] = []\n", + "periodic_eval_learned_mean_rewards: list[float] = []\n", + "periodic_eval_learned_std_rewards: list[float] = []\n", + "periodic_eval_random_mean_rewards: list[float] = []\n", + "periodic_eval_random_std_rewards: list[float] = []\n", + "periodic_eval_learned_mean_steps: list[float] = []\n", + "periodic_eval_random_mean_steps: list[float] = []\n", + "\n", + "best_mean_evaluation_reward = float('-inf')\n", + "best_checkpoint_metadata: dict = {}" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Training loop with periodic evaluation and checkpointing\n", + "\n", + "Every `periodic_eval_interval` episodes, the notebook:\n", + "\n", + "- evaluates the current learned policy on the held-out evaluation seeds\n", + "- evaluates a random baseline on the same seeds\n", + "- saves a `latest.pt` checkpoint\n", + "- overwrites `best.pt` if the learned policy achieves a new best mean evaluation reward" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "for episode_index in range(number_of_training_episodes):\n", + " random_seed = training_seed_start + episode_index\n", + " total_reward, step_count, loss_value = run_one_training_episode(\n", + " environment=environment,\n", + " agent=agent,\n", + " random_seed=random_seed,\n", + " discount_factor_gamma=discount_factor_gamma,\n", + " )\n", + "\n", + " training_rewards.append(total_reward)\n", + " training_steps.append(step_count)\n", + " training_losses.append(float('nan') if loss_value is None else loss_value)\n", + "\n", + " print(\n", + " '[train]',\n", + " f'episode={episode_index:03d}',\n", + " f'seed={random_seed}',\n", + " f'reward={total_reward:.3f}',\n", + " f'steps={step_count}',\n", + " f'loss={loss_value}',\n", + " )\n", + "\n", + " should_run_periodic_eval = (\n", + " (episode_index + 1) % periodic_eval_interval == 0\n", + " or episode_index == number_of_training_episodes - 1\n", + " )\n", + "\n", + " if should_run_periodic_eval:\n", + " learned_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + " )\n", + " random_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=random_agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + " )\n", + "\n", + " periodic_eval_episodes.append(episode_index + 1)\n", + " periodic_eval_learned_mean_rewards.append(learned_eval['mean_reward'])\n", + " periodic_eval_learned_std_rewards.append(learned_eval['std_reward'])\n", + " periodic_eval_random_mean_rewards.append(random_eval['mean_reward'])\n", + " periodic_eval_random_std_rewards.append(random_eval['std_reward'])\n", + " periodic_eval_learned_mean_steps.append(learned_eval['mean_steps'])\n", + " periodic_eval_random_mean_steps.append(random_eval['mean_steps'])\n", + "\n", + " metadata = {\n", + " 'episode': episode_index + 1,\n", + " 'train_seed': random_seed,\n", + " 'learned_mean_reward': learned_eval['mean_reward'],\n", + " 'learned_std_reward': learned_eval['std_reward'],\n", + " 'random_mean_reward': random_eval['mean_reward'],\n", + " 'random_std_reward': random_eval['std_reward'],\n", + " 'evaluation_seeds': heldout_evaluation_seeds,\n", + " }\n", + " agent.save_checkpoint(latest_checkpoint_path, metadata=metadata)\n", + "\n", + " if learned_eval['mean_reward'] > best_mean_evaluation_reward:\n", + " best_mean_evaluation_reward = learned_eval['mean_reward']\n", + " best_checkpoint_metadata = metadata\n", + " agent.save_checkpoint(best_checkpoint_path, metadata=metadata)\n", + " checkpoint_note = 'new best checkpoint'\n", + " else:\n", + " checkpoint_note = 'latest checkpoint only'\n", + "\n", + " print(\n", + " '[periodic-eval]',\n", + " f'episode={episode_index + 1:03d}',\n", + " f'learned_mean_reward={learned_eval[\"mean_reward\"]:.3f}',\n", + " f'learned_std_reward={learned_eval[\"std_reward\"]:.3f}',\n", + " f'random_mean_reward={random_eval[\"mean_reward\"]:.3f}',\n", + " f'random_std_reward={random_eval[\"std_reward\"]:.3f}',\n", + " checkpoint_note,\n", + " )" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Restore the best evaluated model\n", + "\n", + "The training loop may end on a policy that is not the best one seen so far.\n", + "This cell reloads the checkpoint with the highest held-out mean evaluation reward." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "if not best_checkpoint_path.exists():\n", + " raise FileNotFoundError(f'Best checkpoint not found: {best_checkpoint_path}')\n", + "\n", + "loaded_metadata = agent.load_checkpoint(best_checkpoint_path)\n", + "print('Reloaded best checkpoint from:', best_checkpoint_path)\n", + "print('Best checkpoint metadata:')\n", + "loaded_metadata" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Final evaluation of the best checkpoint vs random baseline\n", + "\n", + "This uses the larger held-out evaluation set and compares:\n", + "\n", + "- the best learned policy checkpoint\n", + "- a random baseline on the same seeds" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "best_policy_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + ")\n", + "random_policy_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=random_agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + ")\n", + "\n", + "print('Best policy evaluation mean reward:', best_policy_eval['mean_reward'])\n", + "print('Best policy evaluation std reward:', best_policy_eval['std_reward'])\n", + "print('Random baseline mean reward:', random_policy_eval['mean_reward'])\n", + "print('Random baseline std reward:', random_policy_eval['std_reward'])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Useful plots\n", + "\n", + "These are the most useful quick-look plots for this richer experiment setup:\n", + "\n", + "- **Training reward**: raw reward and moving average during training\n", + "- **Episode length**: how long training episodes run\n", + "- **Periodic evaluation**: learned policy vs random baseline over training\n", + "- **Final evaluation by seed**: best learned checkpoint vs random on the held-out seeds" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def moving_average(values: list[float], window: int) -> np.ndarray:\n", + " if len(values) < window:\n", + " return np.array([])\n", + " kernel = np.ones(window) / window\n", + " return np.convolve(np.asarray(values, dtype=float), kernel, mode='valid')\n", + "\n", + "plot_window = min(10, len(training_rewards))\n", + "smoothed_rewards = moving_average(training_rewards, plot_window)\n", + "training_episode_indices = np.arange(1, len(training_rewards) + 1)\n", + "heldout_seed_indices = np.arange(len(heldout_evaluation_seeds))\n", + "\n", + "fig, axes = plt.subplots(2, 2, figsize=(15, 11))\n", + "\n", + "axes[0, 0].plot(training_episode_indices, training_rewards, marker='o', alpha=0.25, label='raw reward')\n", + "if len(smoothed_rewards) > 0:\n", + " axes[0, 0].plot(\n", + " np.arange(plot_window, len(training_rewards) + 1),\n", + " smoothed_rewards,\n", + " linewidth=2.5,\n", + " color='tab:blue',\n", + " label=f'moving average (window={plot_window})',\n", + " )\n", + "axes[0, 0].set_title('Training Reward per Episode')\n", + "axes[0, 0].set_xlabel('Episode')\n", + "axes[0, 0].set_ylabel('Total Reward')\n", + "axes[0, 0].legend()\n", + "axes[0, 0].grid(alpha=0.3)\n", + "\n", + "axes[0, 1].plot(training_episode_indices, training_steps, marker='o', color='tab:orange')\n", + "axes[0, 1].set_title('Training Episode Length')\n", + "axes[0, 1].set_xlabel('Episode')\n", + "axes[0, 1].set_ylabel('Steps')\n", + "axes[0, 1].grid(alpha=0.3)\n", + "\n", + "periodic_eval_episodes_arr = np.asarray(periodic_eval_episodes)\n", + "learned_mean_arr = np.asarray(periodic_eval_learned_mean_rewards)\n", + "learned_std_arr = np.asarray(periodic_eval_learned_std_rewards)\n", + "random_mean_arr = np.asarray(periodic_eval_random_mean_rewards)\n", + "random_std_arr = np.asarray(periodic_eval_random_std_rewards)\n", + "\n", + "axes[1, 0].plot(periodic_eval_episodes_arr, learned_mean_arr, marker='o', label='learned policy')\n", + "axes[1, 0].fill_between(\n", + " periodic_eval_episodes_arr,\n", + " learned_mean_arr - learned_std_arr,\n", + " learned_mean_arr + learned_std_arr,\n", + " alpha=0.2,\n", + ")\n", + "axes[1, 0].plot(periodic_eval_episodes_arr, random_mean_arr, marker='s', label='random baseline')\n", + "axes[1, 0].fill_between(\n", + " periodic_eval_episodes_arr,\n", + " random_mean_arr - random_std_arr,\n", + " random_mean_arr + random_std_arr,\n", + " alpha=0.2,\n", + ")\n", + "axes[1, 0].set_title('Periodic Evaluation: Learned vs Random')\n", + "axes[1, 0].set_xlabel('Training Episode')\n", + "axes[1, 0].set_ylabel('Mean Evaluation Reward')\n", + "axes[1, 0].legend()\n", + "axes[1, 0].grid(alpha=0.3)\n", + "\n", + "axes[1, 1].plot(\n", + " heldout_seed_indices,\n", + " best_policy_eval['rewards'],\n", + " marker='o',\n", + " linewidth=2,\n", + " label='best learned checkpoint',\n", + ")\n", + "axes[1, 1].plot(\n", + " heldout_seed_indices,\n", + " random_policy_eval['rewards'],\n", + " marker='s',\n", + " linewidth=2,\n", + " label='random baseline',\n", + ")\n", + "axes[1, 1].set_xticks(heldout_seed_indices)\n", + "axes[1, 1].set_xticklabels([str(seed) for seed in heldout_evaluation_seeds], rotation=45)\n", + "axes[1, 1].set_title('Final Evaluation Reward by Seed')\n", + "axes[1, 1].set_xlabel('Held-out Evaluation Seed')\n", + "axes[1, 1].set_ylabel('Total Reward')\n", + "axes[1, 1].legend()\n", + "axes[1, 1].grid(alpha=0.3)\n", + "\n", + "fig.suptitle('Minimal REINFORCE Training Summary with Baseline and Checkpoints', fontsize=16)\n", + "fig.tight_layout()\n", + "plt.show()\n", + "\n", + "print(f'Best checkpoint mean evaluation reward: {best_policy_eval[\"mean_reward\"]:.3f}')\n", + "print(f'Random baseline mean evaluation reward: {random_policy_eval[\"mean_reward\"]:.3f}')\n", + "print(f'Latest checkpoint path: {latest_checkpoint_path}')\n", + "print(f'Best checkpoint path: {best_checkpoint_path}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Render the best checkpoint and save a GIF\n", + "\n", + "This section runs the **best restored checkpoint** in evaluation mode with Bluebird radar rendering enabled.\n", + "It saves individual frames to disk and then combines them into a GIF.\n", + "\n", + "Notes:\n", + "\n", + "- `display_actions=True` overlays actions on the radar frames\n", + "- the GIF path is printed and displayed in the notebook\n", + "- by default this uses the first held-out evaluation seed" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "gif_seed = heldout_evaluation_seeds[0]\n", + "render_dir = Path.cwd() / 'renders' / 'minimal_reinforce_agent_eval'\n", + "gif_path = render_evaluation_rollout_to_gif(\n", + " agent=agent,\n", + " random_seed=gif_seed,\n", + " render_dir=render_dir,\n", + " gif_name=f'best_checkpoint_eval_seed_{gif_seed}',\n", + ")\n", + "\n", + "print(f'Saved GIF to: {gif_path}')\n", + "display(Image(filename=str(gif_path)))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Optional cleanup\n", + "\n", + "Close the main environment if you are done with the notebook session." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "environment.close()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/bluebird-gymnasium/examples/minimal_reinforce_agent.py b/bluebird-gymnasium/examples/minimal_reinforce_agent.py new file mode 100644 index 0000000..9e070ba --- /dev/null +++ b/bluebird-gymnasium/examples/minimal_reinforce_agent.py @@ -0,0 +1,434 @@ +"""Minimal policy-gradient agent for Bluebird Gymnasium. + +This file is an educational example that connects the core RL concepts: + +- observation vectors +- a neural-network policy +- logits and stochastic action sampling +- rewards and discounted returns +- loss calculation +- backpropagation +- gradient-descent-based parameter updates + +It implements a very small REINFORCE-style training loop. This is useful for +understanding the mechanics, but it is not intended to be a strong baseline. +For serious experiments, use a more stable algorithm such as PPO. +""" + +from __future__ import annotations + +import random +import sys +from pathlib import Path + +import numpy as np +import torch +import torch.nn as nn +import torch.optim as optim + +# Allow this example to be run from the monorepo without installing packages. +THIS_FILE = Path(__file__).resolve() +BLUEBIRD_GYMNASIUM_ROOT = THIS_FILE.parents[1] +BLUEBIRD_DT_ROOT = THIS_FILE.parents[2] / "bluebird-dt" +sys.path.insert(0, str(BLUEBIRD_GYMNASIUM_ROOT)) +sys.path.insert(0, str(BLUEBIRD_DT_ROOT)) + +from bluebird_gymnasium.envs import EnvConfig, ViewType # noqa: E402 +from bluebird_gymnasium.envs.sector_i import SectorIEnv # noqa: E402 + + +class PolicyNetwork(nn.Module): + """Neural network that maps one aircraft observation to action logits.""" + + def __init__( + self, + observation_dimension: int, + number_of_actions: int, + hidden_units: int = 128, + ) -> None: + super().__init__() + + # A unit is one neuron in a layer. This network has two hidden layers, + # each with hidden_units neurons. + self.layers = nn.Sequential( + nn.Linear(observation_dimension, hidden_units), + nn.ReLU(), + nn.Linear(hidden_units, hidden_units), + nn.ReLU(), + nn.Linear(hidden_units, number_of_actions), + ) + + def forward(self, observation_batch: torch.Tensor) -> torch.Tensor: + """Return raw action scores, also called logits.""" + return self.layers(observation_batch) + + +class SharedPolicyAgent: + """One shared policy network reused for every aircraft.""" + + def __init__( + self, + observation_dimension: int, + number_of_actions: int, + learning_rate: float = 1e-3, + hidden_units: int = 128, + ) -> None: + self.policy_network = PolicyNetwork( + observation_dimension=observation_dimension, + number_of_actions=number_of_actions, + hidden_units=hidden_units, + ) + + # Adam is a gradient-descent-based optimizer. It updates the network's + # weights and biases using gradients computed by backpropagation. + self.optimizer = optim.Adam( + self.policy_network.parameters(), + lr=learning_rate, + ) + + def sample_training_actions( + self, + observation_by_callsign: dict[str, np.ndarray], + ) -> tuple[dict[str, int], dict[str, torch.Tensor]]: + """Choose stochastic actions and keep log-probabilities for learning. + + During training, we sample from the action distribution instead of + always choosing the largest logit. This is stochastic training: the + agent explores actions that are not currently its top choice. + """ + action_by_callsign: dict[str, int] = {} + log_probability_by_callsign: dict[str, torch.Tensor] = {} + + for callsign, observation_vector in observation_by_callsign.items(): + # Convert the NumPy observation vector into a PyTorch tensor. + # Shape changes from (observation_dimension,) to + # (1, observation_dimension). The leading 1 is the batch dimension. + observation_tensor = torch.tensor( + observation_vector, + dtype=torch.float32, + ).unsqueeze(0) + + # The network outputs logits: raw, unconstrained action scores. + # If there are three actions, this might look like: + # [[0.2, 1.1, -0.5]] + action_logits = self.policy_network(observation_tensor) + + # Categorical(logits=...) applies softmax internally to create a + # probability distribution over discrete actions. + action_distribution = torch.distributions.Categorical( + logits=action_logits, + ) + + # Sample one action from the distribution. This is what enables + # exploration during training. + sampled_action = action_distribution.sample() + + # The log-probability is needed for the policy-gradient loss. + # If an action later receives high return, the loss will push the + # network to increase this log-probability next time. + sampled_action_log_probability = action_distribution.log_prob( + sampled_action, + ) + + action_by_callsign[callsign] = sampled_action.item() + log_probability_by_callsign[callsign] = sampled_action_log_probability.squeeze() + + return action_by_callsign, log_probability_by_callsign + + def choose_evaluation_actions( + self, + observation_by_callsign: dict[str, np.ndarray], + ) -> dict[str, int]: + """Choose deterministic actions for evaluation or deployment.""" + action_by_callsign: dict[str, int] = {} + + # No gradients are needed when evaluating a trained policy. + with torch.no_grad(): + for callsign, observation_vector in observation_by_callsign.items(): + observation_tensor = torch.tensor( + observation_vector, + dtype=torch.float32, + ).unsqueeze(0) + + action_logits = self.policy_network(observation_tensor) + + # argmax chooses the highest-scoring action. + chosen_action = torch.argmax(action_logits, dim=-1).item() + action_by_callsign[callsign] = chosen_action + + return action_by_callsign + + def update_policy_from_episode( + self, + log_probability_per_step: list[torch.Tensor], + reward_per_step: list[float], + discount_factor_gamma: float = 0.99, + ) -> float | None: + """Turn episode rewards into a loss, then update network parameters.""" + if not log_probability_per_step: + return None + + discounted_return_per_step: list[float] = [] + running_discounted_return = 0.0 + + # Compute discounted returns in reverse: + # G_t = r_t + gamma*r_{t+1} + gamma^2*r_{t+2} + ... + for reward in reversed(reward_per_step): + running_discounted_return = reward + discount_factor_gamma * running_discounted_return + discounted_return_per_step.append(running_discounted_return) + + discounted_return_per_step.reverse() + + returns_tensor = torch.tensor( + discounted_return_per_step, + dtype=torch.float32, + ) + + # Normalizing returns often makes simple policy-gradient training less + # numerically erratic. This changes scale, not the episode ordering. + if returns_tensor.numel() > 1: + returns_std = returns_tensor.std(unbiased=False) + if returns_std > 1e-8: + returns_tensor = (returns_tensor - returns_tensor.mean()) / (returns_std + 1e-8) + + # REINFORCE objective: + # - If return is high, increase probability of the sampled action. + # - If return is low/negative, decrease probability of that action. + policy_loss_terms: list[torch.Tensor] = [] + for action_log_probability, discounted_return in zip( + log_probability_per_step, + returns_tensor, + strict=False, + ): + policy_loss_terms.append( + -action_log_probability * discounted_return, + ) + + loss = torch.stack(policy_loss_terms).sum() + + # Clear gradients from the previous update. + self.optimizer.zero_grad() + + # Backpropagation: compute gradients of loss with respect to every + # trainable weight and bias in the policy network. + loss.backward() + + # Gradient descent update: Adam uses the gradients to modify the + # policy network's parameters. + self.optimizer.step() + + return loss.item() + + +def make_sector_i_training_config() -> EnvConfig: + """Create a small, easy Bluebird environment configuration.""" + config = SectorIEnv.get_default_env_config(ViewType.DECENTRALIZED) + + # State encoding: convert raw simulator state into compact numeric vectors. + config.state_repr_config = { + "encoder_cls": "extra_minimal", + "k_nearest_aircraft": 1, + } + + # Keep the first action space small: no-op, left turn, right turn. + config.action_config = { + "simple_heading_left": True, + "simple_heading_right": True, + "simple_fl_climb": False, + "simple_fl_descent": False, + "simple_fl_exit": False, + } + + # Reward components define what behavior the agent is encouraged to learn. + config.reward_config = { + "fns": [ + "position_status_const", + "lateral_centreline_distance_shaped", + "safety_simple_avoidance_exp", + ], + "coeffs": [1.0, 1.0, 1.2], + } + + # Decentralized mode means obs/action/reward/done are dicts keyed by + # aircraft callsign. + config.view_config = { + "type": ViewType.DECENTRALIZED.value, + "decentralized_params": {}, + } + + # Start with one aircraft. Increase difficulty only after this works. + config.scenario_config = { + "cls": "tactical", + "args": { + "num_aircraft": 1, + "balance": [0.0, 0.0, 1.0], + }, + } + + return config + + +def run_one_training_episode( + environment: SectorIEnv, + agent: SharedPolicyAgent, + random_seed: int, + discount_factor_gamma: float, +) -> tuple[float, int, float | None]: + """Run one episode, then update the policy from the collected trajectory.""" + random.seed(random_seed) + np.random.seed(random_seed) + torch.manual_seed(random_seed) + + observation_by_callsign, _info = environment.reset(seed=random_seed) + + episode_is_done = False + episode_step_count = 0 + episode_total_reward = 0.0 + + log_probability_per_step: list[torch.Tensor] = [] + reward_per_step: list[float] = [] + + while not episode_is_done: + action_by_callsign, log_probability_by_callsign = agent.sample_training_actions(observation_by_callsign) + + ( + next_observation_by_callsign, + reward_by_callsign, + done_by_callsign, + truncated_by_callsign, + _info, + ) = environment.step(action_by_callsign) + + # For this simple example, aggregate per-aircraft rewards into one + # scalar reward for the timestep. With one aircraft, this is just that + # aircraft's reward. + timestep_reward = float(sum(reward_by_callsign.values())) if reward_by_callsign else 0.0 + + # Aggregate log-probabilities for all aircraft active this step. With + # one aircraft, this is just that aircraft's sampled action log-prob. + if log_probability_by_callsign: + timestep_log_probability = torch.stack( + list(log_probability_by_callsign.values()), + ).sum() + log_probability_per_step.append(timestep_log_probability) + reward_per_step.append(timestep_reward) + + episode_total_reward += timestep_reward + + # In decentralized mode, done is also per aircraft. + # truncated_by_callsign is not used separately here because Bluebird + # sets done when the episode is time-truncated in this path. + _ = truncated_by_callsign + episode_is_done = all(done_by_callsign.values()) if done_by_callsign else True + + observation_by_callsign = next_observation_by_callsign + episode_step_count += 1 + + loss_value = agent.update_policy_from_episode( + log_probability_per_step=log_probability_per_step, + reward_per_step=reward_per_step, + discount_factor_gamma=discount_factor_gamma, + ) + + return episode_total_reward, episode_step_count, loss_value + + +def run_one_evaluation_episode( + environment: SectorIEnv, + agent: SharedPolicyAgent, + random_seed: int, +) -> tuple[float, int]: + """Run one episode without updating network weights.""" + random.seed(random_seed) + np.random.seed(random_seed) + torch.manual_seed(random_seed) + + observation_by_callsign, _info = environment.reset(seed=random_seed) + + episode_is_done = False + episode_step_count = 0 + episode_total_reward = 0.0 + + while not episode_is_done: + action_by_callsign = agent.choose_evaluation_actions( + observation_by_callsign, + ) + + ( + next_observation_by_callsign, + reward_by_callsign, + done_by_callsign, + truncated_by_callsign, + _info, + ) = environment.step(action_by_callsign) + + _ = truncated_by_callsign + episode_total_reward += float(sum(reward_by_callsign.values())) if reward_by_callsign else 0.0 + episode_is_done = all(done_by_callsign.values()) if done_by_callsign else True + observation_by_callsign = next_observation_by_callsign + episode_step_count += 1 + + return episode_total_reward, episode_step_count + + +def main() -> None: + config = make_sector_i_training_config() + environment = SectorIEnv(config=config) + + # In decentralized mode, this is the length of one aircraft's observation + # vector. For extra_minimal with k_nearest_aircraft=1, it is typically 4. + observation_dimension = environment.observation_space.shape[0] + + # This is the number of discrete actions available to each aircraft. + number_of_actions = environment.action_space.n + + agent = SharedPolicyAgent( + observation_dimension=observation_dimension, + number_of_actions=number_of_actions, + learning_rate=1e-3, + hidden_units=128, + ) + + print( + "environment shapes:", + f"observation_dimension={observation_dimension}", + f"number_of_actions={number_of_actions}", + ) + + discount_factor_gamma = 0.99 + + for episode_index in range(20): + random_seed = 100 + episode_index + total_reward, step_count, loss_value = run_one_training_episode( + environment=environment, + agent=agent, + random_seed=random_seed, + discount_factor_gamma=discount_factor_gamma, + ) + + print( + "[train]", + f"episode={episode_index:02d}", + f"seed={random_seed}", + f"reward={total_reward:.3f}", + f"steps={step_count}", + f"loss={loss_value}", + ) + + for random_seed in [200, 201, 202]: + total_reward, step_count = run_one_evaluation_episode( + environment=environment, + agent=agent, + random_seed=random_seed, + ) + + print( + "[eval]", + f"seed={random_seed}", + f"reward={total_reward:.3f}", + f"steps={step_count}", + ) + + +if __name__ == "__main__": + main() diff --git a/bluebird-gymnasium/examples/minimal_reinforce_flight_school.ipynb b/bluebird-gymnasium/examples/minimal_reinforce_flight_school.ipynb new file mode 100644 index 0000000..c42988f --- /dev/null +++ b/bluebird-gymnasium/examples/minimal_reinforce_flight_school.ipynb @@ -0,0 +1,892 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Minimal REINFORCE Agent for Flight School\n", + "\n", + "This notebook adapts the original REINFORCE example to use **Bluebird Flight School** instead of `SectorIEnv`.\n", + "\n", + "Flight School differs from the earlier example in an important way:\n", + "\n", + "- it uses an **infinite traffic generator**\n", + "- the notebook uses the **centralized** control view\n", + "- the observation is one combined state vector\n", + "- the action is one discrete action integer each step\n", + "\n", + "The notebook still includes:\n", + "\n", + "- a small policy network\n", + "- REINFORCE training\n", + "- periodic evaluation during training\n", + "- comparison against a random baseline\n", + "- checkpoint saving and best-model restore\n", + "- final GIF rendering of the best policy" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Environment note\n", + "\n", + "This notebook assumes you are already running the **correct Python/Jupyter kernel**:\n", + "one with `gymnasium`, `torch`, and the Bluebird project dependencies installed.\n", + "\n", + "Do the package setup in your shell or project virtual environment first,\n", + "then open this notebook with that kernel. The notebook does **not** try to install\n", + "packages itself." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Imports and path setup\n", + "\n", + "This cell makes the notebook runnable from either the `bluebird-gymnasium` directory\n", + "or the repo root by adding the local package paths to `sys.path`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from __future__ import annotations\n", + "\n", + "import random\n", + "import shutil\n", + "import sys\n", + "from pathlib import Path\n", + "\n", + "import imageio\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import torch\n", + "import torch.nn as nn\n", + "import torch.optim as optim\n", + "from IPython.display import Image, display\n", + "\n", + "search_roots = [Path.cwd().resolve(), *Path.cwd().resolve().parents]\n", + "gym_root = None\n", + "dt_root = None\n", + "\n", + "for candidate in search_roots:\n", + " if (candidate / 'bluebird_gymnasium').exists():\n", + " gym_root = candidate\n", + " sibling_dt = candidate.parent / 'bluebird-dt'\n", + " if sibling_dt.exists():\n", + " dt_root = sibling_dt\n", + " break\n", + " if (candidate / 'bluebird-gymnasium').exists() and (candidate / 'bluebird-dt').exists():\n", + " gym_root = candidate / 'bluebird-gymnasium'\n", + " dt_root = candidate / 'bluebird-dt'\n", + " break\n", + "\n", + "if gym_root is None or dt_root is None:\n", + " raise RuntimeError('Could not locate local bluebird-gymnasium and bluebird-dt package roots.')\n", + "\n", + "sys.path.insert(0, str(gym_root))\n", + "sys.path.insert(0, str(dt_root))\n", + "\n", + "from bluebird_gymnasium.envs import EnvConfig, ViewType\n", + "from bluebird_gymnasium.envs.flight_school import FlightSchoolEnv\n", + "\n", + "print(f'Using bluebird-gymnasium from: {gym_root}')\n", + "print(f'Using bluebird-dt from: {dt_root}')\n", + "print(f'Torch version: {torch.__version__}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Policy network and agents\n", + "\n", + "Because Flight School is used here in **centralized** mode:\n", + "\n", + "- the input is one observation vector `obs`\n", + "- the policy outputs one set of action logits\n", + "- the chosen action is one integer\n", + "\n", + "This makes the REINFORCE logic a little simpler than the decentralized notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "class PolicyNetwork(nn.Module):\n", + " \"\"\"Neural network that maps one observation vector to action logits.\"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " observation_dimension: int,\n", + " number_of_actions: int,\n", + " hidden_units: int = 128,\n", + " ) -> None:\n", + " super().__init__()\n", + " self.layers = nn.Sequential(\n", + " nn.Linear(observation_dimension, hidden_units),\n", + " nn.ReLU(),\n", + " nn.Linear(hidden_units, hidden_units),\n", + " nn.ReLU(),\n", + " nn.Linear(hidden_units, number_of_actions),\n", + " )\n", + "\n", + " def forward(self, observation_batch: torch.Tensor) -> torch.Tensor:\n", + " return self.layers(observation_batch)\n", + "\n", + "\n", + "class SharedPolicyAgent:\n", + " \"\"\"Single policy network for centralized Flight School control.\"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " observation_dimension: int,\n", + " number_of_actions: int,\n", + " learning_rate: float = 1e-3,\n", + " hidden_units: int = 128,\n", + " ) -> None:\n", + " self.policy_network = PolicyNetwork(\n", + " observation_dimension=observation_dimension,\n", + " number_of_actions=number_of_actions,\n", + " hidden_units=hidden_units,\n", + " )\n", + " self.optimizer = optim.Adam(\n", + " self.policy_network.parameters(),\n", + " lr=learning_rate,\n", + " )\n", + "\n", + " def sample_training_action(\n", + " self,\n", + " observation_vector: np.ndarray,\n", + " ) -> tuple[int, torch.Tensor]:\n", + " observation_tensor = torch.tensor(\n", + " observation_vector,\n", + " dtype=torch.float32,\n", + " ).unsqueeze(0)\n", + " action_logits = self.policy_network(observation_tensor)\n", + " action_distribution = torch.distributions.Categorical(logits=action_logits)\n", + " sampled_action = action_distribution.sample()\n", + " action_log_probability = action_distribution.log_prob(sampled_action)\n", + " return sampled_action.item(), action_log_probability.squeeze(0)\n", + "\n", + " def choose_evaluation_action(\n", + " self,\n", + " observation_vector: np.ndarray,\n", + " ) -> int:\n", + " with torch.no_grad():\n", + " observation_tensor = torch.tensor(\n", + " observation_vector,\n", + " dtype=torch.float32,\n", + " ).unsqueeze(0)\n", + " action_logits = self.policy_network(observation_tensor)\n", + " return torch.argmax(action_logits, dim=-1).item()\n", + "\n", + " def update_policy_from_episode(\n", + " self,\n", + " log_probability_per_step: list[torch.Tensor],\n", + " reward_per_step: list[float],\n", + " discount_factor_gamma: float = 0.99,\n", + " ) -> float | None:\n", + " if not log_probability_per_step:\n", + " return None\n", + "\n", + " discounted_return_per_step: list[float] = []\n", + " running_discounted_return = 0.0\n", + "\n", + " for reward in reversed(reward_per_step):\n", + " running_discounted_return = reward + discount_factor_gamma * running_discounted_return\n", + " discounted_return_per_step.append(running_discounted_return)\n", + "\n", + " discounted_return_per_step.reverse()\n", + " returns_tensor = torch.tensor(discounted_return_per_step, dtype=torch.float32)\n", + "\n", + " if returns_tensor.numel() > 1:\n", + " returns_std = returns_tensor.std(unbiased=False)\n", + " if returns_std > 1e-8:\n", + " returns_tensor = (\n", + " (returns_tensor - returns_tensor.mean())\n", + " / (returns_std + 1e-8)\n", + " )\n", + "\n", + " policy_loss_terms: list[torch.Tensor] = []\n", + " for action_log_probability, discounted_return in zip(\n", + " log_probability_per_step,\n", + " returns_tensor,\n", + " ):\n", + " policy_loss_terms.append(-action_log_probability * discounted_return)\n", + "\n", + " loss = torch.stack(policy_loss_terms).sum()\n", + " self.optimizer.zero_grad()\n", + " loss.backward()\n", + " self.optimizer.step()\n", + " return loss.item()\n", + "\n", + " def save_checkpoint(self, checkpoint_path: Path, metadata: dict | None = None) -> None:\n", + " checkpoint_path.parent.mkdir(parents=True, exist_ok=True)\n", + " payload = {\n", + " 'policy_state_dict': self.policy_network.state_dict(),\n", + " 'optimizer_state_dict': self.optimizer.state_dict(),\n", + " 'metadata': metadata or {},\n", + " }\n", + " torch.save(payload, checkpoint_path)\n", + "\n", + " def load_checkpoint(self, checkpoint_path: Path, map_location: str = 'cpu') -> dict:\n", + " payload = torch.load(checkpoint_path, map_location=map_location)\n", + " self.policy_network.load_state_dict(payload['policy_state_dict'])\n", + " if 'optimizer_state_dict' in payload:\n", + " self.optimizer.load_state_dict(payload['optimizer_state_dict'])\n", + " return payload.get('metadata', {})\n", + "\n", + "\n", + "class RandomAgent:\n", + " \"\"\"Simple random baseline for centralized control.\"\"\"\n", + "\n", + " def __init__(self, number_of_actions: int) -> None:\n", + " self.number_of_actions = number_of_actions\n", + "\n", + " def choose_evaluation_action(self, observation_vector: np.ndarray) -> int:\n", + " _ = observation_vector\n", + " return random.randrange(self.number_of_actions)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Flight School configuration and rollout helpers\n", + "\n", + "This notebook uses `FlightSchoolEnv` in centralized mode.\n", + "\n", + "The default Flight School config already includes:\n", + "\n", + "- infinite traffic generation\n", + "- a shaped reward with safety terms\n", + "- a 10-minute Gymnasium episode horizon\n", + "\n", + "The helper functions below support:\n", + "\n", + "- one training episode\n", + "- one evaluation episode\n", + "- evaluation over many seeds\n", + "- GIF rendering for a final evaluation rollout" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def make_flight_school_training_config(random_seed: int | None = None) -> EnvConfig:\n", + " config = FlightSchoolEnv.get_default_env_config(ViewType.CENTRALIZED)\n", + " config.scenario_config['args']['random_seed'] = random_seed\n", + " config.scenario_duration = 10 * 60\n", + " config.view_config['type'] = ViewType.CENTRALIZED.value\n", + " return config\n", + "\n", + "\n", + "def run_one_training_episode(\n", + " environment: FlightSchoolEnv,\n", + " agent: SharedPolicyAgent,\n", + " random_seed: int,\n", + " discount_factor_gamma: float,\n", + ") -> tuple[float, int, float | None]:\n", + " random.seed(random_seed)\n", + " np.random.seed(random_seed)\n", + " torch.manual_seed(random_seed)\n", + " environment.config.scenario_config['args']['random_seed'] = random_seed\n", + "\n", + " observation_vector, _info = environment.reset(seed=random_seed)\n", + "\n", + " episode_is_done = False\n", + " episode_step_count = 0\n", + " episode_total_reward = 0.0\n", + "\n", + " log_probability_per_step: list[torch.Tensor] = []\n", + " reward_per_step: list[float] = []\n", + "\n", + " while not episode_is_done:\n", + " action_int, action_log_probability = agent.sample_training_action(observation_vector)\n", + " (\n", + " next_observation_vector,\n", + " reward,\n", + " done,\n", + " truncated,\n", + " _info,\n", + " ) = environment.step(action_int)\n", + "\n", + " log_probability_per_step.append(action_log_probability)\n", + " reward_per_step.append(float(reward))\n", + " episode_total_reward += float(reward)\n", + " episode_is_done = bool(done or truncated)\n", + " observation_vector = next_observation_vector\n", + " episode_step_count += 1\n", + "\n", + " loss_value = agent.update_policy_from_episode(\n", + " log_probability_per_step=log_probability_per_step,\n", + " reward_per_step=reward_per_step,\n", + " discount_factor_gamma=discount_factor_gamma,\n", + " )\n", + "\n", + " return episode_total_reward, episode_step_count, loss_value\n", + "\n", + "\n", + "def run_one_evaluation_episode(\n", + " environment: FlightSchoolEnv,\n", + " evaluation_agent,\n", + " random_seed: int,\n", + ") -> tuple[float, int]:\n", + " random.seed(random_seed)\n", + " np.random.seed(random_seed)\n", + " torch.manual_seed(random_seed)\n", + " environment.config.scenario_config['args']['random_seed'] = random_seed\n", + "\n", + " observation_vector, _info = environment.reset(seed=random_seed)\n", + "\n", + " episode_is_done = False\n", + " episode_step_count = 0\n", + " episode_total_reward = 0.0\n", + "\n", + " while not episode_is_done:\n", + " action_int = evaluation_agent.choose_evaluation_action(observation_vector)\n", + " (\n", + " next_observation_vector,\n", + " reward,\n", + " done,\n", + " truncated,\n", + " _info,\n", + " ) = environment.step(action_int)\n", + "\n", + " episode_total_reward += float(reward)\n", + " episode_is_done = bool(done or truncated)\n", + " observation_vector = next_observation_vector\n", + " episode_step_count += 1\n", + "\n", + " return episode_total_reward, episode_step_count\n", + "\n", + "\n", + "def evaluate_agent_over_seeds(\n", + " environment: FlightSchoolEnv,\n", + " evaluation_agent,\n", + " evaluation_seeds: list[int],\n", + ") -> dict:\n", + " rewards: list[float] = []\n", + " steps: list[int] = []\n", + "\n", + " for random_seed in evaluation_seeds:\n", + " total_reward, step_count = run_one_evaluation_episode(\n", + " environment=environment,\n", + " evaluation_agent=evaluation_agent,\n", + " random_seed=random_seed,\n", + " )\n", + " rewards.append(total_reward)\n", + " steps.append(step_count)\n", + "\n", + " return {\n", + " 'seeds': evaluation_seeds,\n", + " 'rewards': rewards,\n", + " 'steps': steps,\n", + " 'mean_reward': float(np.mean(rewards)),\n", + " 'std_reward': float(np.std(rewards)),\n", + " 'mean_steps': float(np.mean(steps)),\n", + " }\n", + "\n", + "\n", + "def render_evaluation_rollout_to_gif(\n", + " agent: SharedPolicyAgent,\n", + " random_seed: int,\n", + " render_dir: Path,\n", + " gif_name: str = 'flight_school_reinforce_eval',\n", + " render_every_n_steps: int = 5,\n", + " gif_frame_duration_seconds: float = 0.2,\n", + ") -> Path:\n", + " render_config = make_flight_school_training_config(random_seed=random_seed)\n", + " render_config.radar_config['display_actions'] = True\n", + " render_config.radar_config['render_dir'] = str(render_dir)\n", + " render_config.radar_config['prefix'] = 'frame'\n", + "\n", + " if render_dir.exists():\n", + " shutil.rmtree(render_dir)\n", + " render_dir.mkdir(parents=True, exist_ok=True)\n", + "\n", + " render_environment = FlightSchoolEnv(config=render_config)\n", + " render_environment.set_render_mode('file')\n", + "\n", + " random.seed(random_seed)\n", + " np.random.seed(random_seed)\n", + " torch.manual_seed(random_seed)\n", + "\n", + " observation_vector, _info = render_environment.reset(seed=random_seed)\n", + " render_environment.render()\n", + "\n", + " episode_is_done = False\n", + " step_index = 0\n", + "\n", + " while not episode_is_done:\n", + " action_int = agent.choose_evaluation_action(observation_vector)\n", + " (\n", + " next_observation_vector,\n", + " reward,\n", + " done,\n", + " truncated,\n", + " _info,\n", + " ) = render_environment.step(action_int)\n", + " _ = reward\n", + " step_index += 1\n", + " episode_is_done = bool(done or truncated)\n", + " if step_index % render_every_n_steps == 0 or episode_is_done:\n", + " render_environment.render()\n", + " observation_vector = next_observation_vector\n", + "\n", + " png_frames = sorted(render_dir.glob(f\"{render_config.radar_config['prefix']}_*.png\"))\n", + " if not png_frames:\n", + " raise RuntimeError(\n", + " f'No rendered PNG frames were written to {render_dir}. '\n", + " 'Expected at least one frame before GIF generation.'\n", + " )\n", + "\n", + " gif_path = render_dir / f'{gif_name}.gif'\n", + " images = [imageio.v3.imread(frame_path) for frame_path in png_frames]\n", + " imageio.mimsave(gif_path, images, loop=0, duration=gif_frame_duration_seconds)\n", + " render_environment.close()\n", + " return gif_path\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Set up the environment and inspect the shapes\n", + "\n", + "For this notebook, the most important values are:\n", + "\n", + "- `observation_dimension`: how many numbers are in the centralized observation vector\n", + "- `number_of_actions`: how many discrete actions the policy can choose from" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "initial_seed = 7\n", + "config = make_flight_school_training_config(random_seed=initial_seed)\n", + "environment = FlightSchoolEnv(config=config)\n", + "\n", + "observation_dimension = environment.observation_space.shape[0]\n", + "number_of_actions = environment.action_space.n\n", + "\n", + "print(\n", + " 'environment shapes:',\n", + " f'observation_dimension={observation_dimension}',\n", + " f'number_of_actions={number_of_actions}',\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Hyperparameters and experiment settings\n", + "\n", + "This version keeps the same richer experiment loop:\n", + "\n", + "- periodic evaluation during training\n", + "- larger held-out evaluation set\n", + "- checkpoint saving\n", + "- automatic tracking of the best evaluated model\n", + "- random-policy baseline comparison" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "learning_rate = 1e-3\n", + "hidden_units = 128\n", + "discount_factor_gamma = 0.99\n", + "number_of_training_episodes = 100\n", + "training_seed_start = 100\n", + "periodic_eval_interval = 10\n", + "heldout_evaluation_seeds = list(range(200, 220))\n", + "checkpoint_dir = Path.cwd() / 'checkpoints' / 'minimal_reinforce_flight_school'\n", + "latest_checkpoint_path = checkpoint_dir / 'latest.pt'\n", + "best_checkpoint_path = checkpoint_dir / 'best.pt'\n", + "\n", + "agent = SharedPolicyAgent(\n", + " observation_dimension=observation_dimension,\n", + " number_of_actions=number_of_actions,\n", + " learning_rate=learning_rate,\n", + " hidden_units=hidden_units,\n", + ")\n", + "random_agent = RandomAgent(number_of_actions=number_of_actions)\n", + "\n", + "training_rewards: list[float] = []\n", + "training_steps: list[int] = []\n", + "training_losses: list[float] = []\n", + "\n", + "periodic_eval_episodes: list[int] = []\n", + "periodic_eval_learned_mean_rewards: list[float] = []\n", + "periodic_eval_learned_std_rewards: list[float] = []\n", + "periodic_eval_random_mean_rewards: list[float] = []\n", + "periodic_eval_random_std_rewards: list[float] = []\n", + "periodic_eval_learned_mean_steps: list[float] = []\n", + "periodic_eval_random_mean_steps: list[float] = []\n", + "\n", + "best_mean_evaluation_reward = float('-inf')\n", + "best_checkpoint_metadata: dict = {}\n", + "render_every_n_steps = 5\n", + "gif_frame_duration_seconds = 0.2\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Training loop with periodic evaluation and checkpointing\n", + "\n", + "Every `periodic_eval_interval` episodes, the notebook:\n", + "\n", + "- evaluates the current learned policy on the held-out evaluation seeds\n", + "- evaluates a random baseline on the same seeds\n", + "- saves a `latest.pt` checkpoint\n", + "- overwrites `best.pt` if the learned policy achieves a new best mean evaluation reward" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "for episode_index in range(number_of_training_episodes):\n", + " random_seed = training_seed_start + episode_index\n", + " total_reward, step_count, loss_value = run_one_training_episode(\n", + " environment=environment,\n", + " agent=agent,\n", + " random_seed=random_seed,\n", + " discount_factor_gamma=discount_factor_gamma,\n", + " )\n", + "\n", + " training_rewards.append(total_reward)\n", + " training_steps.append(step_count)\n", + " training_losses.append(float('nan') if loss_value is None else loss_value)\n", + "\n", + " print(\n", + " '[train]',\n", + " f'episode={episode_index:03d}',\n", + " f'seed={random_seed}',\n", + " f'reward={total_reward:.3f}',\n", + " f'steps={step_count}',\n", + " f'loss={loss_value}',\n", + " )\n", + "\n", + " should_run_periodic_eval = (\n", + " (episode_index + 1) % periodic_eval_interval == 0\n", + " or episode_index == number_of_training_episodes - 1\n", + " )\n", + "\n", + " if should_run_periodic_eval:\n", + " learned_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + " )\n", + " random_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=random_agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + " )\n", + "\n", + " periodic_eval_episodes.append(episode_index + 1)\n", + " periodic_eval_learned_mean_rewards.append(learned_eval['mean_reward'])\n", + " periodic_eval_learned_std_rewards.append(learned_eval['std_reward'])\n", + " periodic_eval_random_mean_rewards.append(random_eval['mean_reward'])\n", + " periodic_eval_random_std_rewards.append(random_eval['std_reward'])\n", + " periodic_eval_learned_mean_steps.append(learned_eval['mean_steps'])\n", + " periodic_eval_random_mean_steps.append(random_eval['mean_steps'])\n", + "\n", + " metadata = {\n", + " 'episode': episode_index + 1,\n", + " 'train_seed': random_seed,\n", + " 'learned_mean_reward': learned_eval['mean_reward'],\n", + " 'learned_std_reward': learned_eval['std_reward'],\n", + " 'random_mean_reward': random_eval['mean_reward'],\n", + " 'random_std_reward': random_eval['std_reward'],\n", + " 'evaluation_seeds': heldout_evaluation_seeds,\n", + " }\n", + " agent.save_checkpoint(latest_checkpoint_path, metadata=metadata)\n", + "\n", + " if learned_eval['mean_reward'] > best_mean_evaluation_reward:\n", + " best_mean_evaluation_reward = learned_eval['mean_reward']\n", + " best_checkpoint_metadata = metadata\n", + " agent.save_checkpoint(best_checkpoint_path, metadata=metadata)\n", + " checkpoint_note = 'new best checkpoint'\n", + " else:\n", + " checkpoint_note = 'latest checkpoint only'\n", + "\n", + " print(\n", + " '[periodic-eval]',\n", + " f'episode={episode_index + 1:03d}',\n", + " f'learned_mean_reward={learned_eval[\"mean_reward\"]:.3f}',\n", + " f'learned_std_reward={learned_eval[\"std_reward\"]:.3f}',\n", + " f'random_mean_reward={random_eval[\"mean_reward\"]:.3f}',\n", + " f'random_std_reward={random_eval[\"std_reward\"]:.3f}',\n", + " checkpoint_note,\n", + " )" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Restore the best evaluated model\n", + "\n", + "The training loop may end on a policy that is not the best one seen so far.\n", + "This cell reloads the checkpoint with the highest held-out mean evaluation reward." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "if not best_checkpoint_path.exists():\n", + " raise FileNotFoundError(f'Best checkpoint not found: {best_checkpoint_path}')\n", + "\n", + "loaded_metadata = agent.load_checkpoint(best_checkpoint_path)\n", + "print('Reloaded best checkpoint from:', best_checkpoint_path)\n", + "print('Best checkpoint metadata:')\n", + "loaded_metadata" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Final evaluation of the best checkpoint vs random baseline\n", + "\n", + "This uses the larger held-out evaluation set and compares:\n", + "\n", + "- the best learned policy checkpoint\n", + "- a random baseline on the same seeds" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "best_policy_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + ")\n", + "random_policy_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=random_agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + ")\n", + "\n", + "print('Best policy evaluation mean reward:', best_policy_eval['mean_reward'])\n", + "print('Best policy evaluation std reward:', best_policy_eval['std_reward'])\n", + "print('Random baseline mean reward:', random_policy_eval['mean_reward'])\n", + "print('Random baseline std reward:', random_policy_eval['std_reward'])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Useful plots\n", + "\n", + "These are the most useful quick-look plots for this Flight School REINFORCE setup:\n", + "\n", + "- **Training reward**: raw reward and moving average during training\n", + "- **Episode length**: how long training episodes run\n", + "- **Periodic evaluation**: learned policy vs random baseline over training\n", + "- **Final evaluation by seed**: best learned checkpoint vs random on the held-out seeds" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def moving_average(values: list[float], window: int) -> np.ndarray:\n", + " if len(values) < window:\n", + " return np.array([])\n", + " kernel = np.ones(window) / window\n", + " return np.convolve(np.asarray(values, dtype=float), kernel, mode='valid')\n", + "\n", + "plot_window = min(10, len(training_rewards))\n", + "smoothed_rewards = moving_average(training_rewards, plot_window)\n", + "training_episode_indices = np.arange(1, len(training_rewards) + 1)\n", + "heldout_seed_indices = np.arange(len(heldout_evaluation_seeds))\n", + "periodic_eval_episodes_arr = np.asarray(periodic_eval_episodes)\n", + "learned_mean_arr = np.asarray(periodic_eval_learned_mean_rewards)\n", + "learned_std_arr = np.asarray(periodic_eval_learned_std_rewards)\n", + "random_mean_arr = np.asarray(periodic_eval_random_mean_rewards)\n", + "random_std_arr = np.asarray(periodic_eval_random_std_rewards)\n", + "\n", + "fig, axes = plt.subplots(2, 2, figsize=(15, 11))\n", + "\n", + "axes[0, 0].plot(training_episode_indices, training_rewards, marker='o', alpha=0.25, label='raw reward')\n", + "if len(smoothed_rewards) > 0:\n", + " axes[0, 0].plot(\n", + " np.arange(plot_window, len(training_rewards) + 1),\n", + " smoothed_rewards,\n", + " linewidth=2.5,\n", + " color='tab:blue',\n", + " label=f'moving average (window={plot_window})',\n", + " )\n", + "axes[0, 0].set_title('Training Reward per Episode')\n", + "axes[0, 0].set_xlabel('Episode')\n", + "axes[0, 0].set_ylabel('Total Reward')\n", + "axes[0, 0].legend()\n", + "axes[0, 0].grid(alpha=0.3)\n", + "\n", + "axes[0, 1].plot(training_episode_indices, training_steps, marker='o', color='tab:orange')\n", + "axes[0, 1].set_title('Training Episode Length')\n", + "axes[0, 1].set_xlabel('Episode')\n", + "axes[0, 1].set_ylabel('Steps')\n", + "axes[0, 1].grid(alpha=0.3)\n", + "\n", + "axes[1, 0].plot(periodic_eval_episodes_arr, learned_mean_arr, marker='o', label='learned policy')\n", + "axes[1, 0].fill_between(\n", + " periodic_eval_episodes_arr,\n", + " learned_mean_arr - learned_std_arr,\n", + " learned_mean_arr + learned_std_arr,\n", + " alpha=0.2,\n", + ")\n", + "axes[1, 0].plot(periodic_eval_episodes_arr, random_mean_arr, marker='s', label='random baseline')\n", + "axes[1, 0].fill_between(\n", + " periodic_eval_episodes_arr,\n", + " random_mean_arr - random_std_arr,\n", + " random_mean_arr + random_std_arr,\n", + " alpha=0.2,\n", + ")\n", + "axes[1, 0].set_title('Periodic Evaluation: Learned vs Random')\n", + "axes[1, 0].set_xlabel('Training Episode')\n", + "axes[1, 0].set_ylabel('Mean Evaluation Reward')\n", + "axes[1, 0].legend()\n", + "axes[1, 0].grid(alpha=0.3)\n", + "\n", + "axes[1, 1].plot(\n", + " heldout_seed_indices,\n", + " best_policy_eval['rewards'],\n", + " marker='o',\n", + " linewidth=2,\n", + " label='best learned checkpoint',\n", + ")\n", + "axes[1, 1].plot(\n", + " heldout_seed_indices,\n", + " random_policy_eval['rewards'],\n", + " marker='s',\n", + " linewidth=2,\n", + " label='random baseline',\n", + ")\n", + "axes[1, 1].set_xticks(heldout_seed_indices)\n", + "axes[1, 1].set_xticklabels([str(seed) for seed in heldout_evaluation_seeds], rotation=45)\n", + "axes[1, 1].set_title('Final Evaluation Reward by Seed')\n", + "axes[1, 1].set_xlabel('Held-out Evaluation Seed')\n", + "axes[1, 1].set_ylabel('Total Reward')\n", + "axes[1, 1].legend()\n", + "axes[1, 1].grid(alpha=0.3)\n", + "\n", + "fig.suptitle('Flight School REINFORCE Summary with Baseline and Checkpoints', fontsize=16)\n", + "fig.tight_layout()\n", + "plt.show()\n", + "\n", + "print(f'Best checkpoint mean evaluation reward: {best_policy_eval[\"mean_reward\"]:.3f}')\n", + "print(f'Random baseline mean evaluation reward: {random_policy_eval[\"mean_reward\"]:.3f}')\n", + "print(f'Latest checkpoint path: {latest_checkpoint_path}')\n", + "print(f'Best checkpoint path: {best_checkpoint_path}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Render the best checkpoint and save a GIF\n", + "\n", + "This section runs the **best restored checkpoint** in evaluation mode with Bluebird radar rendering enabled.\n", + "It saves individual frames to disk and then combines them into a GIF.\n", + "\n", + "Notes:\n", + "\n", + "- `display_actions=True` overlays actions on the radar frames\n", + "- the GIF path is printed and displayed in the notebook\n", + "- by default this uses the first held-out evaluation seed\n", + "- `render_every_n_steps` controls how many simulator steps to skip between saved frames\n", + "- `gif_frame_duration_seconds` controls how long each GIF frame is shown\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "gif_seed = heldout_evaluation_seeds[0]\n", + "render_dir = Path.cwd() / 'renders' / 'minimal_reinforce_flight_school_eval'\n", + "gif_path = render_evaluation_rollout_to_gif(\n", + " agent=agent,\n", + " random_seed=gif_seed,\n", + " render_dir=render_dir,\n", + " gif_name=f'best_checkpoint_eval_seed_{gif_seed}',\n", + " render_every_n_steps=render_every_n_steps,\n", + " gif_frame_duration_seconds=gif_frame_duration_seconds,\n", + ")\n", + "\n", + "print(f'Saved GIF to: {gif_path}')\n", + "display(Image(filename=str(gif_path)))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Optional cleanup\n", + "\n", + "Close the main environment if you are done with the notebook session." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "environment.close()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/bluebird-gymnasium/examples/ppo_curriculum_agent.ipynb b/bluebird-gymnasium/examples/ppo_curriculum_agent.ipynb new file mode 100644 index 0000000..ca7b20a --- /dev/null +++ b/bluebird-gymnasium/examples/ppo_curriculum_agent.ipynb @@ -0,0 +1,1456 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "d04355c0", + "metadata": {}, + "source": [ + "# PPO Curriculum Agent for Bluebird Gymnasium\n", + "\n", + "This notebook is a **follow-on** to [minimal_ppo_agent.ipynb](/home/sprite/BluebirdATC/bluebird-gymnasium/examples/minimal_ppo_agent.ipynb).\n", + "\n", + "It keeps the same overall notebook style and PPO structure, but extends the training setup to support:\n", + "\n", + "- staged curriculum learning\n", + "- checkpoint reuse across stages and notebook sessions\n", + "- gradual aircraft-count increases\n", + "- a larger-sector follow-on stage after Sector I\n", + "\n", + "The important design choice is that the curriculum stays **decentralized** and **lateral-only**.\n", + "That keeps the observation dimension and action count stable, so checkpoints remain reusable.\n" + ] + }, + { + "cell_type": "markdown", + "id": "ac0a96e1", + "metadata": {}, + "source": [ + "## Environment note\n", + "\n", + "This notebook can bootstrap `torch` into the current Python/Jupyter kernel if it\n", + "is missing. It prefers `uv pip install --python ` when `uv` is\n", + "available, and falls back to `python -m pip install` otherwise. Bluebird project\n", + "dependencies still need to be available in the kernel or local workspace.\n", + "\n", + "It is intentionally based on the older `minimal_ppo_agent.ipynb` line rather than replacing it.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "f6ba8f0e", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Torch available in kernel: 2.12.0+cu130\n" + ] + } + ], + "source": [ + "from __future__ import annotations\n", + "\n", + "import importlib\n", + "import shutil\n", + "import subprocess\n", + "import sys\n", + "\n", + "\n", + "def ensure_python_package(module_name: str, install_name: str | None = None) -> None:\n", + " if importlib.util.find_spec(module_name) is not None:\n", + " return\n", + "\n", + " package_name = install_name or module_name\n", + " uv_executable = shutil.which('uv')\n", + "\n", + " if uv_executable is not None:\n", + " print(f'Installing {package_name} with uv into the current kernel environment...')\n", + " subprocess.check_call([\n", + " uv_executable,\n", + " 'pip',\n", + " 'install',\n", + " '--python',\n", + " sys.executable,\n", + " package_name,\n", + " ])\n", + " return\n", + "\n", + " try:\n", + " import pip # noqa: F401\n", + " except ImportError:\n", + " import ensurepip\n", + " ensurepip.bootstrap(upgrade=True)\n", + "\n", + " print(f'Installing {package_name} with pip into the current kernel environment...')\n", + " subprocess.check_call([sys.executable, '-m', 'pip', 'install', package_name])\n", + "\n", + "\n", + "ensure_python_package('torch')\n", + "\n", + "import torch\n", + "\n", + "print(f'Torch available in kernel: {torch.__version__}')\n" + ] + }, + { + "cell_type": "markdown", + "id": "ab8e3539", + "metadata": {}, + "source": [ + "## Imports and path setup\n", + "\n", + "This cell makes the notebook runnable from either the repo root or the\n", + "`bluebird-gymnasium` directory by adding the local package paths to `sys.path`.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "f95bed66", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Using bluebird-gymnasium from: /mnt/74F4CA39F4C9FCFC/Giles/Projects/project-bluebird/BluebirdATC/bluebird-gymnasium\n", + "Using bluebird-dt from: /mnt/74F4CA39F4C9FCFC/Giles/Projects/project-bluebird/BluebirdATC/bluebird-dt\n", + "Using notebook helpers from: /mnt/74F4CA39F4C9FCFC/Giles/Projects/project-bluebird/BluebirdATC/bluebird-gymnasium/examples\n", + "Torch version: 2.12.0+cu130\n" + ] + } + ], + "source": [ + "from __future__ import annotations\n", + "\n", + "import json\n", + "import random\n", + "import shutil\n", + "import sys\n", + "from pathlib import Path\n", + "\n", + "import imageio\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import torch\n", + "import torch.nn as nn\n", + "import torch.optim as optim\n", + "from IPython.display import Image, display\n", + "\n", + "search_roots = [Path.cwd().resolve(), *Path.cwd().resolve().parents]\n", + "gym_root = None\n", + "dt_root = None\n", + "\n", + "for candidate in search_roots:\n", + " if (candidate / 'bluebird_gymnasium').exists():\n", + " gym_root = candidate\n", + " sibling_dt = candidate.parent / 'bluebird-dt'\n", + " if sibling_dt.exists():\n", + " dt_root = sibling_dt\n", + " break\n", + " if (candidate / 'bluebird-gymnasium').exists() and (candidate / 'bluebird-dt').exists():\n", + " gym_root = candidate / 'bluebird-gymnasium'\n", + " dt_root = candidate / 'bluebird-dt'\n", + " break\n", + "\n", + "if gym_root is None or dt_root is None:\n", + " raise RuntimeError('Could not locate local bluebird-gymnasium and bluebird-dt package roots.')\n", + "\n", + "examples_root = gym_root / 'examples'\n", + "sys.path.insert(0, str(gym_root))\n", + "sys.path.insert(0, str(dt_root))\n", + "sys.path.insert(0, str(examples_root))\n", + "\n", + "from bluebird_gymnasium.envs import EnvConfig, ViewType\n", + "from bluebird_gymnasium.envs.sector_i import SectorIEnv\n", + "from bluebird_gymnasium.envs.sector_xplus import SectorXPlusEnv\n", + "from curriculum_scenarios import register_curriculum_scenario_managers\n", + "\n", + "register_curriculum_scenario_managers()\n", + "\n", + "print(f'Using bluebird-gymnasium from: {gym_root}')\n", + "print(f'Using bluebird-dt from: {dt_root}')\n", + "print(f'Using notebook helpers from: {examples_root}')\n", + "print(f'Torch version: {torch.__version__}')\n" + ] + }, + { + "cell_type": "markdown", + "id": "5bdf453e", + "metadata": {}, + "source": [ + "## PPO actor-critic model and agents\n", + "\n", + "This is the same style of PPO agent used in the older minimal PPO notebook.\n", + "The main additions are shape-aware checkpoint loading and explicit metadata storage.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "ed228f7b", + "metadata": {}, + "outputs": [], + "source": [ + "class ActorCriticNetwork(nn.Module):\n", + " \"\"\"Shared-trunk actor-critic network for PPO.\"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " observation_dimension: int,\n", + " number_of_actions: int,\n", + " hidden_units: int = 128,\n", + " ) -> None:\n", + " super().__init__()\n", + " self.trunk = nn.Sequential(\n", + " nn.Linear(observation_dimension, hidden_units),\n", + " nn.ReLU(),\n", + " nn.Linear(hidden_units, hidden_units),\n", + " nn.ReLU(),\n", + " )\n", + " self.policy_head = nn.Linear(hidden_units, number_of_actions)\n", + " self.value_head = nn.Linear(hidden_units, 1)\n", + "\n", + " def forward(self, observation_batch: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:\n", + " features = self.trunk(observation_batch)\n", + " action_logits = self.policy_head(features)\n", + " state_value = self.value_head(features).squeeze(-1)\n", + " return action_logits, state_value\n", + "\n", + "\n", + "class PPOAgent:\n", + " \"\"\"Minimal PPO agent with actor-critic network and checkpoint helpers.\"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " observation_dimension: int,\n", + " number_of_actions: int,\n", + " learning_rate: float = 3e-4,\n", + " hidden_units: int = 128,\n", + " clip_epsilon: float = 0.2,\n", + " value_loss_coefficient: float = 0.5,\n", + " entropy_coefficient: float = 0.01,\n", + " ppo_epochs: int = 4,\n", + " ) -> None:\n", + " self.observation_dimension = observation_dimension\n", + " self.number_of_actions = number_of_actions\n", + " self.actor_critic = ActorCriticNetwork(\n", + " observation_dimension=observation_dimension,\n", + " number_of_actions=number_of_actions,\n", + " hidden_units=hidden_units,\n", + " )\n", + " self.optimizer = optim.Adam(self.actor_critic.parameters(), lr=learning_rate)\n", + " self.clip_epsilon = clip_epsilon\n", + " self.value_loss_coefficient = value_loss_coefficient\n", + " self.entropy_coefficient = entropy_coefficient\n", + " self.ppo_epochs = ppo_epochs\n", + "\n", + " def choose_training_action(\n", + " self,\n", + " observation_vector: np.ndarray,\n", + " ) -> tuple[int, torch.Tensor, torch.Tensor, torch.Tensor]:\n", + " observation_tensor = torch.tensor(\n", + " observation_vector,\n", + " dtype=torch.float32,\n", + " ).unsqueeze(0)\n", + " action_logits, state_value = self.actor_critic(observation_tensor)\n", + " action_distribution = torch.distributions.Categorical(logits=action_logits)\n", + " sampled_action = action_distribution.sample()\n", + " action_log_probability = action_distribution.log_prob(sampled_action)\n", + "\n", + " return (\n", + " sampled_action.item(),\n", + " observation_tensor.squeeze(0),\n", + " action_log_probability.squeeze(0),\n", + " state_value.squeeze(0),\n", + " )\n", + "\n", + " def choose_evaluation_actions(\n", + " self,\n", + " observation_by_callsign: dict[str, np.ndarray],\n", + " ) -> dict[str, int]:\n", + " chosen_actions: dict[str, int] = {}\n", + " with torch.no_grad():\n", + " for callsign, observation_vector in observation_by_callsign.items():\n", + " observation_tensor = torch.tensor(\n", + " observation_vector,\n", + " dtype=torch.float32,\n", + " ).unsqueeze(0)\n", + " action_logits, _state_value = self.actor_critic(observation_tensor)\n", + " chosen_actions[callsign] = torch.argmax(action_logits, dim=-1).item()\n", + " return chosen_actions\n", + "\n", + " def update_from_trajectory(\n", + " self,\n", + " observations: list[torch.Tensor],\n", + " actions: list[torch.Tensor],\n", + " old_log_probabilities: list[torch.Tensor],\n", + " returns: torch.Tensor,\n", + " advantages: torch.Tensor,\n", + " ) -> dict[str, float] | None:\n", + " if not observations:\n", + " return None\n", + "\n", + " observation_tensor = torch.stack(observations)\n", + " action_tensor = torch.stack(actions).long()\n", + " old_log_probability_tensor = torch.stack(old_log_probabilities).detach()\n", + " returns_tensor = returns.detach()\n", + " advantages_tensor = advantages.detach()\n", + "\n", + " if advantages_tensor.numel() > 1:\n", + " advantages_std = advantages_tensor.std(unbiased=False)\n", + " if advantages_std > 1e-8:\n", + " advantages_tensor = (\n", + " (advantages_tensor - advantages_tensor.mean())\n", + " / (advantages_std + 1e-8)\n", + " )\n", + "\n", + " mean_policy_loss = 0.0\n", + " mean_value_loss = 0.0\n", + " mean_entropy = 0.0\n", + " mean_total_loss = 0.0\n", + "\n", + " for _epoch in range(self.ppo_epochs):\n", + " new_action_logits, new_state_values = self.actor_critic(observation_tensor)\n", + " action_distribution = torch.distributions.Categorical(logits=new_action_logits)\n", + " new_log_probabilities = action_distribution.log_prob(action_tensor)\n", + " entropy = action_distribution.entropy().mean()\n", + "\n", + " probability_ratio = torch.exp(new_log_probabilities - old_log_probability_tensor)\n", + " unclipped_objective = probability_ratio * advantages_tensor\n", + " clipped_objective = torch.clamp(\n", + " probability_ratio,\n", + " 1.0 - self.clip_epsilon,\n", + " 1.0 + self.clip_epsilon,\n", + " ) * advantages_tensor\n", + "\n", + " policy_loss = -torch.min(unclipped_objective, clipped_objective).mean()\n", + " value_loss = torch.nn.functional.mse_loss(new_state_values, returns_tensor)\n", + " total_loss = (\n", + " policy_loss\n", + " + self.value_loss_coefficient * value_loss\n", + " - self.entropy_coefficient * entropy\n", + " )\n", + "\n", + " self.optimizer.zero_grad()\n", + " total_loss.backward()\n", + " self.optimizer.step()\n", + "\n", + " mean_policy_loss += float(policy_loss.item())\n", + " mean_value_loss += float(value_loss.item())\n", + " mean_entropy += float(entropy.item())\n", + " mean_total_loss += float(total_loss.item())\n", + "\n", + " epoch_divisor = float(self.ppo_epochs)\n", + " return {\n", + " 'policy_loss': mean_policy_loss / epoch_divisor,\n", + " 'value_loss': mean_value_loss / epoch_divisor,\n", + " 'entropy': mean_entropy / epoch_divisor,\n", + " 'total_loss': mean_total_loss / epoch_divisor,\n", + " }\n", + "\n", + " def save_checkpoint(self, checkpoint_path: Path, metadata: dict | None = None) -> None:\n", + " checkpoint_path.parent.mkdir(parents=True, exist_ok=True)\n", + " payload = {\n", + " 'model_state_dict': self.actor_critic.state_dict(),\n", + " 'optimizer_state_dict': self.optimizer.state_dict(),\n", + " 'metadata': _to_python_types(metadata or {}),\n", + " }\n", + " torch.save(payload, checkpoint_path)\n", + "\n", + " def load_checkpoint(\n", + " self,\n", + " checkpoint_path: Path,\n", + " map_location: str = 'cpu',\n", + " load_optimizer_state: bool = True,\n", + " strict_shape_check: bool = True,\n", + " ) -> dict:\n", + " payload = torch.load(checkpoint_path, map_location=map_location, weights_only=False)\n", + " metadata = payload.get('metadata', {})\n", + " if strict_shape_check:\n", + " saved_obs_dim = metadata.get('observation_dimension')\n", + " saved_num_actions = metadata.get('number_of_actions')\n", + " if saved_obs_dim is not None and saved_obs_dim != self.observation_dimension:\n", + " raise ValueError(\n", + " f'Checkpoint observation dimension {saved_obs_dim} does not match current {self.observation_dimension}. '\n", + " 'Use a different checkpoint lineage.'\n", + " )\n", + " if saved_num_actions is not None and saved_num_actions != self.number_of_actions:\n", + " raise ValueError(\n", + " f'Checkpoint action count {saved_num_actions} does not match current {self.number_of_actions}. '\n", + " 'Use a different checkpoint lineage.'\n", + " )\n", + " self.actor_critic.load_state_dict(payload['model_state_dict'])\n", + " if load_optimizer_state and 'optimizer_state_dict' in payload:\n", + " self.optimizer.load_state_dict(payload['optimizer_state_dict'])\n", + " return metadata\n", + "\n", + "\n", + "\n", + "def _to_python_types(value):\n", + " if isinstance(value, dict):\n", + " return {key: _to_python_types(val) for key, val in value.items()}\n", + " if isinstance(value, (list, tuple)):\n", + " return [_to_python_types(item) for item in value]\n", + " if isinstance(value, Path):\n", + " return str(value)\n", + " if isinstance(value, np.generic):\n", + " return value.item()\n", + " return value\n", + "\n", + "\n", + "class RandomAgent:\n", + " \"\"\"Simple random baseline for comparison.\"\"\"\n", + "\n", + " def __init__(self, number_of_actions: int) -> None:\n", + " self.number_of_actions = number_of_actions\n", + "\n", + " def choose_evaluation_actions(\n", + " self,\n", + " observation_by_callsign: dict[str, np.ndarray],\n", + " ) -> dict[str, int]:\n", + " return {\n", + " callsign: random.randrange(self.number_of_actions)\n", + " for callsign in observation_by_callsign.keys()\n", + " }\n" + ] + }, + { + "cell_type": "markdown", + "id": "b8322b19", + "metadata": {}, + "source": [ + "## Curriculum configuration and rollout helpers\n", + "\n", + "The curriculum deliberately keeps these fixed across stages:\n", + "\n", + "- decentralized control\n", + "- `extra_minimal` state encoder\n", + "- `k_nearest_aircraft = 1`\n", + "- lateral-only action set\n", + "- same shared policy network shape\n", + "\n", + "It now uses a more progress-oriented reward mix:\n", + "\n", + "- a much lighter centreline shaping term\n", + "- `lateral_next_fix_proximity_dist_exp` to reward getting closer to the next fix\n", + "- the existing safety penalty\n", + "- `expeditious_linear` to reward reducing distance-to-exit\n", + "- a stronger custom `route_progress_terminal_reward` that gives:\n", + " - a positive bonus for correct exit\n", + " - a per-step progress signal\n", + " - a timeout penalty for loitering\n", + "\n", + "That is intended to make circling near the entry fix much less attractive.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "14ffddc3", + "metadata": {}, + "outputs": [], + "source": [ + "def make_training_config(\n", + " env_cls,\n", + " num_aircraft: int,\n", + " k_nearest_aircraft: int = 1,\n", + " enable_vertical_actions: bool = False,\n", + " scenario_duration_seconds: int = 1800,\n", + " scenario_cls: str = 'tactical',\n", + " scenario_args: dict | None = None,\n", + " exit_window_width_nmi: float = 5.0,\n", + " reward_coeff_overrides: list[float] | None = None,\n", + ") -> EnvConfig:\n", + " config = env_cls.get_default_env_config(ViewType.DECENTRALIZED)\n", + " config.airspace_config['exit_window_width'] = exit_window_width_nmi\n", + "\n", + " config.state_repr_config = {\n", + " 'encoder_cls': 'extra_minimal',\n", + " 'k_nearest_aircraft': k_nearest_aircraft,\n", + " }\n", + "\n", + " config.action_config = {\n", + " 'simple_heading_left': True,\n", + " 'simple_heading_right': True,\n", + " 'simple_heading_route_parallel': True,\n", + " 'simple_fl_climb': enable_vertical_actions,\n", + " 'simple_fl_descent': enable_vertical_actions,\n", + " 'simple_fl_exit': False,\n", + " }\n", + "\n", + " reward_fns = [\n", + " 'position_status_const',\n", + " 'lateral_centreline_distance_shaped',\n", + " 'lateral_next_fix_proximity_dist_exp',\n", + " 'safety_simple_avoidance_exp',\n", + " 'expeditious_linear',\n", + " 'route_progress_terminal_reward',\n", + " 'anti_loiter_route_rejoin_reward',\n", + " 'route_parallel_exp',\n", + " 'action_penalty_thresh',\n", + " ]\n", + " reward_coeffs = [1.0, 0.06, 0.9, 1.6, 2.2, 3.5, 1.6, 0.8, 0.02]\n", + " if reward_coeff_overrides is not None:\n", + " reward_coeffs = list(reward_coeff_overrides)\n", + "\n", + " if num_aircraft <= 1:\n", + " reward_fns.append('lateral_termination_check_sac_env')\n", + " else:\n", + " reward_fns.append('lateral_termination_check_mac_env')\n", + "\n", + " if len(reward_coeffs) == len(reward_fns) - 1:\n", + " reward_coeffs.append(0.05)\n", + " elif len(reward_coeffs) != len(reward_fns):\n", + " raise ValueError(\n", + " 'reward_coeff_overrides must match either the base reward count '\n", + " 'or the full reward count including termination shaping.'\n", + " )\n", + "\n", + " config.reward_config = {\n", + " 'fns': reward_fns,\n", + " 'coeffs': reward_coeffs,\n", + " }\n", + "\n", + " config.view_config = {\n", + " 'type': ViewType.DECENTRALIZED.value,\n", + " 'decentralized_params': {},\n", + " }\n", + "\n", + " if scenario_args is None:\n", + " scenario_args = {\n", + " 'num_aircraft': num_aircraft,\n", + " 'balance': [0.0, 0.0, 1.0],\n", + " }\n", + "\n", + " config.scenario_config = {\n", + " 'cls': scenario_cls,\n", + " 'args': scenario_args,\n", + " }\n", + "\n", + " config.scenario_duration = scenario_duration_seconds\n", + " return config\n", + "\n", + "\n", + "def compute_returns_and_advantages(\n", + " rewards: list[float],\n", + " values: list[torch.Tensor],\n", + " dones: list[bool],\n", + " discount_factor_gamma: float,\n", + " gae_lambda: float,\n", + ") -> tuple[torch.Tensor, torch.Tensor]:\n", + " rewards_tensor = torch.tensor(rewards, dtype=torch.float32)\n", + " values_tensor = torch.stack(values).detach().float()\n", + " dones_tensor = torch.tensor(dones, dtype=torch.float32)\n", + "\n", + " advantages = torch.zeros_like(rewards_tensor)\n", + " last_gae = torch.tensor(0.0)\n", + "\n", + " for timestep in reversed(range(len(rewards))):\n", + " if timestep == len(rewards) - 1:\n", + " next_value = torch.tensor(0.0)\n", + " else:\n", + " next_value = values_tensor[timestep + 1]\n", + "\n", + " next_nonterminal = 1.0 - dones_tensor[timestep]\n", + " delta = (\n", + " rewards_tensor[timestep]\n", + " + discount_factor_gamma * next_value * next_nonterminal\n", + " - values_tensor[timestep]\n", + " )\n", + " last_gae = delta + discount_factor_gamma * gae_lambda * next_nonterminal * last_gae\n", + " advantages[timestep] = last_gae\n", + "\n", + " returns = advantages + values_tensor\n", + " return returns, advantages\n", + "\n", + "\n", + "def run_one_training_episode(\n", + " environment,\n", + " agent: PPOAgent,\n", + " random_seed: int,\n", + " discount_factor_gamma: float,\n", + " gae_lambda: float,\n", + ") -> tuple[float, int, dict[str, float] | None]:\n", + " random.seed(random_seed)\n", + " np.random.seed(random_seed)\n", + " torch.manual_seed(random_seed)\n", + "\n", + " observation_by_callsign, _info = environment.reset(seed=random_seed)\n", + "\n", + " episode_is_done = False\n", + " episode_step_count = 0\n", + " episode_total_reward = 0.0\n", + "\n", + " observations: list[torch.Tensor] = []\n", + " actions: list[torch.Tensor] = []\n", + " old_log_probabilities: list[torch.Tensor] = []\n", + " values: list[torch.Tensor] = []\n", + " rewards: list[float] = []\n", + " dones: list[bool] = []\n", + "\n", + " while not episode_is_done:\n", + " action_by_callsign: dict[str, int] = {}\n", + " step_rollout_rows: list[tuple[str, torch.Tensor, int, torch.Tensor, torch.Tensor]] = []\n", + "\n", + " for callsign, observation_vector in observation_by_callsign.items():\n", + " action_int, observation_tensor, action_log_probability, state_value = agent.choose_training_action(observation_vector)\n", + " action_by_callsign[callsign] = action_int\n", + " step_rollout_rows.append(\n", + " (\n", + " callsign,\n", + " observation_tensor,\n", + " action_int,\n", + " action_log_probability.detach(),\n", + " state_value.detach(),\n", + " )\n", + " )\n", + "\n", + " (\n", + " next_observation_by_callsign,\n", + " reward_by_callsign,\n", + " done_by_callsign,\n", + " truncated_by_callsign,\n", + " _info,\n", + " ) = environment.step(action_by_callsign)\n", + "\n", + " timestep_reward = float(sum(reward_by_callsign.values())) if reward_by_callsign else 0.0\n", + " timestep_done = all(done_by_callsign.values()) if done_by_callsign else True\n", + "\n", + " for callsign, observation_tensor, action_int, action_log_probability, state_value in step_rollout_rows:\n", + " observations.append(observation_tensor)\n", + " actions.append(torch.tensor(action_int))\n", + " old_log_probabilities.append(action_log_probability)\n", + " values.append(state_value)\n", + " rewards.append(float(reward_by_callsign.get(callsign, 0.0)))\n", + " dones.append(\n", + " bool(done_by_callsign.get(callsign, False))\n", + " or bool(truncated_by_callsign.get(callsign, False))\n", + " )\n", + "\n", + " _ = truncated_by_callsign\n", + " episode_total_reward += timestep_reward\n", + " episode_is_done = timestep_done\n", + " observation_by_callsign = next_observation_by_callsign\n", + " episode_step_count += 1\n", + "\n", + " returns, advantages = compute_returns_and_advantages(\n", + " rewards=rewards,\n", + " values=values,\n", + " dones=dones,\n", + " discount_factor_gamma=discount_factor_gamma,\n", + " gae_lambda=gae_lambda,\n", + " )\n", + "\n", + " update_metrics = agent.update_from_trajectory(\n", + " observations=observations,\n", + " actions=actions,\n", + " old_log_probabilities=old_log_probabilities,\n", + " returns=returns,\n", + " advantages=advantages,\n", + " )\n", + "\n", + " return episode_total_reward, episode_step_count, update_metrics\n", + "\n", + "\n", + "def run_one_evaluation_episode(environment, evaluation_agent, random_seed: int) -> tuple[float, int]:\n", + " random.seed(random_seed)\n", + " np.random.seed(random_seed)\n", + " torch.manual_seed(random_seed)\n", + "\n", + " observation_by_callsign, _info = environment.reset(seed=random_seed)\n", + "\n", + " episode_is_done = False\n", + " episode_step_count = 0\n", + " episode_total_reward = 0.0\n", + "\n", + " while not episode_is_done:\n", + " action_by_callsign = evaluation_agent.choose_evaluation_actions(observation_by_callsign)\n", + " (\n", + " next_observation_by_callsign,\n", + " reward_by_callsign,\n", + " done_by_callsign,\n", + " truncated_by_callsign,\n", + " _info,\n", + " ) = environment.step(action_by_callsign)\n", + "\n", + " _ = truncated_by_callsign\n", + " episode_total_reward += float(sum(reward_by_callsign.values())) if reward_by_callsign else 0.0\n", + " episode_is_done = all(done_by_callsign.values()) if done_by_callsign else True\n", + " observation_by_callsign = next_observation_by_callsign\n", + " episode_step_count += 1\n", + "\n", + " return episode_total_reward, episode_step_count\n", + "\n", + "\n", + "def evaluate_agent_over_seeds(environment, evaluation_agent, evaluation_seeds: list[int]) -> dict:\n", + " rewards: list[float] = []\n", + " steps: list[int] = []\n", + "\n", + " for random_seed in evaluation_seeds:\n", + " total_reward, step_count = run_one_evaluation_episode(environment, evaluation_agent, random_seed)\n", + " rewards.append(total_reward)\n", + " steps.append(step_count)\n", + "\n", + " return {\n", + " 'seeds': evaluation_seeds,\n", + " 'rewards': rewards,\n", + " 'steps': steps,\n", + " 'mean_reward': float(np.mean(rewards)),\n", + " 'std_reward': float(np.std(rewards)),\n", + " 'mean_steps': float(np.mean(steps)),\n", + " }\n", + "\n", + "\n", + "def render_evaluation_rollout_to_gif(\n", + " env_cls,\n", + " stage_num_aircraft: int,\n", + " enable_vertical_actions: bool,\n", + " scenario_duration_seconds: int,\n", + " agent: PPOAgent,\n", + " random_seed: int,\n", + " render_dir: Path,\n", + " gif_name: str,\n", + " render_every_n_steps: int = 5,\n", + " gif_frame_duration_seconds: float = 0.2,\n", + " scenario_cls: str = 'tactical',\n", + " scenario_args: dict | None = None,\n", + " exit_window_width_nmi: float = 5.0,\n", + " reward_coeff_overrides: list[float] | None = None,\n", + ") -> dict[str, Path]:\n", + " render_config = make_training_config(\n", + " env_cls=env_cls,\n", + " num_aircraft=stage_num_aircraft,\n", + " enable_vertical_actions=enable_vertical_actions,\n", + " scenario_duration_seconds=scenario_duration_seconds,\n", + " scenario_cls=scenario_cls,\n", + " scenario_args=scenario_args,\n", + " exit_window_width_nmi=exit_window_width_nmi,\n", + " reward_coeff_overrides=reward_coeff_overrides,\n", + " )\n", + " render_config.radar_config['display_actions'] = True\n", + " render_config.radar_config['render_dir'] = str(render_dir)\n", + " render_config.radar_config['prefix'] = 'frame'\n", + "\n", + " if render_dir.exists():\n", + " shutil.rmtree(render_dir)\n", + " render_dir.mkdir(parents=True, exist_ok=True)\n", + "\n", + " render_environment = env_cls(config=render_config)\n", + " render_environment.set_render_mode('file')\n", + "\n", + " random.seed(random_seed)\n", + " np.random.seed(random_seed)\n", + " torch.manual_seed(random_seed)\n", + "\n", + " observation_by_callsign, _info = render_environment.reset(seed=random_seed)\n", + " action_formatter_map = render_environment.get_action_parser().action_formatter_map\n", + " action_trace: list[dict] = []\n", + " render_environment.render()\n", + "\n", + " episode_is_done = False\n", + " step_index = 0\n", + "\n", + " while not episode_is_done:\n", + " action_by_callsign = agent.choose_evaluation_actions(observation_by_callsign)\n", + " action_trace.append({\n", + " 'step': step_index,\n", + " 'actions': {\n", + " callsign: action_formatter_map.get(action_int, str(action_int))\n", + " for callsign, action_int in action_by_callsign.items()\n", + " },\n", + " })\n", + " (\n", + " next_observation_by_callsign,\n", + " reward_by_callsign,\n", + " done_by_callsign,\n", + " truncated_by_callsign,\n", + " _info,\n", + " ) = render_environment.step(action_by_callsign)\n", + " _ = reward_by_callsign, truncated_by_callsign\n", + " step_index += 1\n", + " episode_is_done = all(done_by_callsign.values()) if done_by_callsign else True\n", + " if step_index % render_every_n_steps == 0 or episode_is_done:\n", + " render_environment.render()\n", + " observation_by_callsign = next_observation_by_callsign\n", + "\n", + " png_frames = sorted(render_dir.glob(f\"{render_config.radar_config['prefix']}_*.png\"))\n", + " if not png_frames:\n", + " raise RuntimeError(\n", + " f'No rendered PNG frames were written to {render_dir}. ' \n", + " 'Expected at least one frame before GIF generation.'\n", + " )\n", + "\n", + " gif_path = render_dir / f'{gif_name}.gif'\n", + " frames = [imageio.v3.imread(frame_path) for frame_path in png_frames]\n", + " imageio.mimsave(gif_path, frames, duration=gif_frame_duration_seconds, loop=0)\n", + "\n", + " action_trace_path = render_dir / f'{gif_name}_actions.json'\n", + " action_trace_path.write_text(json.dumps(action_trace, indent=2))\n", + "\n", + " render_environment.close()\n", + " return {\n", + " 'gif_path': gif_path,\n", + " 'action_trace_path': action_trace_path,\n", + " }\n" + ] + }, + { + "cell_type": "markdown", + "id": "a5d61556", + "metadata": {}, + "source": [ + "## Curriculum stages and CPU-friendly defaults\n", + "\n", + "These defaults are intended for a standard desktop CPU:\n", + "\n", + "- `hidden_units = 128`\n", + "- `ppo_epochs = 3`\n", + "- modest per-stage episode counts\n", + "- evaluation on 10 held-out seeds\n", + "- shorter early-stage horizons so loitering is more expensive\n", + "\n", + "The stage ordering reflects the recommendation from the earlier discussion:\n", + "\n", + "1. `SectorIEnv`, 1 aircraft\n", + "2. `SectorIEnv`, 2 aircraft\n", + "3. `SectorIEnv`, 3 aircraft\n", + "4. `SectorXPlusEnv`, 2 aircraft\n" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "2553f942", + "metadata": {}, + "outputs": [], + "source": [ + "global_seed = 7\n", + "learning_rate = 3e-4\n", + "hidden_units = 128\n", + "discount_factor_gamma = 0.99\n", + "gae_lambda = 0.95\n", + "clip_epsilon = 0.2\n", + "value_loss_coefficient = 0.5\n", + "entropy_coefficient = 0.01\n", + "ppo_epochs = 3\n", + "periodic_eval_interval = 10\n", + "heldout_evaluation_seeds = list(range(300, 310))\n", + "render_every_n_steps = 5\n", + "gif_frame_duration_seconds = 0.2\n", + "\n", + "sector_i_exit_window_width_nmi = 5.0\n", + "sector_i_forward_route = ['FIRE', 'EARTH', 'WATER', 'AIR', 'SPIRIT']\n", + "sector_i_reverse_route = list(reversed(sector_i_forward_route))\n", + "\n", + "curriculum_stages = [\n", + " {\n", + " 'stage_name': 'sector_i_1ac_lateral',\n", + " 'env_cls': SectorIEnv,\n", + " 'num_aircraft': 1,\n", + " 'enable_vertical_actions': False,\n", + " 'scenario_duration_seconds': 900,\n", + " 'training_episodes': 60,\n", + " 'training_seed_start': 100,\n", + " 'scenario_cls': 'tactical',\n", + " 'reward_coeff_overrides': [1.0, 0.06, 0.9, 1.6, 2.2, 3.5, 1.6, 0.8, 0.02, 0.05],\n", + " 'scenario_args': {\n", + " 'num_aircraft': 1,\n", + " 'balance': [0.0, 0.0, 1.0],\n", + " },\n", + " 'exit_window_width_nmi': sector_i_exit_window_width_nmi,\n", + " },\n", + " {\n", + " 'stage_name': 'sector_i_2ac_head_on_wide_gap_lateral',\n", + " 'env_cls': SectorIEnv,\n", + " 'num_aircraft': 2,\n", + " 'enable_vertical_actions': False,\n", + " 'scenario_duration_seconds': 1500,\n", + " 'training_episodes': 140,\n", + " 'training_seed_start': 2000,\n", + " 'scenario_cls': 'fixed_sequence',\n", + " 'reward_coeff_overrides': [1.0, 0.04, 0.75, 4.2, 1.4, 2.4, 1.8, 1.0, 0.02, 0.05],\n", + " 'scenario_args': {\n", + " 'aircraft_specs': [\n", + " {\n", + " 'callsign': 'AIR0',\n", + " 'route_filed': sector_i_forward_route,\n", + " 'start_time_seconds': 0,\n", + " 'speed_tas': 380.0,\n", + " 'entry_fl': 200.0,\n", + " 'exit_fl': 200.0,\n", + " },\n", + " {\n", + " 'callsign': 'AIR1',\n", + " 'route_filed': sector_i_reverse_route,\n", + " 'start_time_seconds': 150,\n", + " 'speed_tas': 380.0,\n", + " 'entry_fl': 200.0,\n", + " 'exit_fl': 200.0,\n", + " },\n", + " ],\n", + " },\n", + " 'exit_window_width_nmi': sector_i_exit_window_width_nmi,\n", + " },\n", + " {\n", + " 'stage_name': 'sector_i_2ac_head_on_bridge_gap_lateral',\n", + " 'env_cls': SectorIEnv,\n", + " 'num_aircraft': 2,\n", + " 'enable_vertical_actions': False,\n", + " 'scenario_duration_seconds': 1500,\n", + " 'training_episodes': 140,\n", + " 'training_seed_start': 3000,\n", + " 'scenario_cls': 'fixed_sequence',\n", + " 'reward_coeff_overrides': [1.0, 0.04, 0.75, 4.4, 1.35, 2.3, 1.9, 1.05, 0.02, 0.05],\n", + " 'scenario_args': {\n", + " 'aircraft_specs': [\n", + " {\n", + " 'callsign': 'AIR0',\n", + " 'route_filed': sector_i_forward_route,\n", + " 'start_time_seconds': 0,\n", + " 'speed_tas': 380.0,\n", + " 'entry_fl': 200.0,\n", + " 'exit_fl': 200.0,\n", + " },\n", + " {\n", + " 'callsign': 'AIR1',\n", + " 'route_filed': sector_i_reverse_route,\n", + " 'start_time_seconds': 120,\n", + " 'speed_tas': 380.0,\n", + " 'entry_fl': 200.0,\n", + " 'exit_fl': 200.0,\n", + " },\n", + " ],\n", + " },\n", + " 'exit_window_width_nmi': sector_i_exit_window_width_nmi,\n", + " },\n", + " {\n", + " 'stage_name': 'sector_i_2ac_head_on_medium_gap_lateral',\n", + " 'env_cls': SectorIEnv,\n", + " 'num_aircraft': 2,\n", + " 'enable_vertical_actions': False,\n", + " 'scenario_duration_seconds': 1550,\n", + " 'training_episodes': 150,\n", + " 'training_seed_start': 4000,\n", + " 'scenario_cls': 'fixed_sequence',\n", + " 'reward_coeff_overrides': [1.0, 0.04, 0.75, 5.0, 1.2, 2.1, 2.0, 1.15, 0.02, 0.05],\n", + " 'scenario_args': {\n", + " 'aircraft_specs': [\n", + " {\n", + " 'callsign': 'AIR0',\n", + " 'route_filed': sector_i_forward_route,\n", + " 'start_time_seconds': 0,\n", + " 'speed_tas': 380.0,\n", + " 'entry_fl': 200.0,\n", + " 'exit_fl': 200.0,\n", + " },\n", + " {\n", + " 'callsign': 'AIR1',\n", + " 'route_filed': sector_i_reverse_route,\n", + " 'start_time_seconds': 90,\n", + " 'speed_tas': 380.0,\n", + " 'entry_fl': 200.0,\n", + " 'exit_fl': 200.0,\n", + " },\n", + " ],\n", + " },\n", + " 'exit_window_width_nmi': sector_i_exit_window_width_nmi,\n", + " },\n", + "]\n", + "\n", + "# Same-direction bridge stages are intentionally disabled in this notebook.\n", + "# The current focus is to stabilize pure head-on avoidance and route rejoin.\n", + "# 3-aircraft Sector I and SectorXPlus remain disabled until that behavior is stable.\n", + "\n", + "curriculum_checkpoint_dir = Path.cwd() / 'checkpoints' / 'ppo_curriculum_agent'\n", + "resume_from_checkpoint: Path | None = None\n", + "resume_from_stage_index = 0\n", + "load_optimizer_state = True\n", + "history_json_path = curriculum_checkpoint_dir / 'curriculum_history.json'\n" + ] + }, + { + "cell_type": "markdown", + "id": "e7ca1570", + "metadata": {}, + "source": [ + "## Initialize the first stage and optionally resume a checkpoint\n", + "\n", + "The checkpoint loader will reject incompatible observation/action shapes.\n", + "That is deliberate: if you change the action set, you should start a new checkpoint lineage.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "301e139f", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Global seed: 7\n", + "Shared observation dimension: 4\n", + "Shared action count: 4\n" + ] + } + ], + "source": [ + "random.seed(global_seed)\n", + "np.random.seed(global_seed)\n", + "torch.manual_seed(global_seed)\n", + "\n", + "first_stage = curriculum_stages[resume_from_stage_index]\n", + "first_config = make_training_config(\n", + " env_cls=first_stage['env_cls'],\n", + " num_aircraft=first_stage['num_aircraft'],\n", + " enable_vertical_actions=first_stage['enable_vertical_actions'],\n", + " scenario_duration_seconds=first_stage['scenario_duration_seconds'],\n", + " scenario_cls=first_stage['scenario_cls'],\n", + " scenario_args=first_stage['scenario_args'],\n", + " exit_window_width_nmi=first_stage['exit_window_width_nmi'],\n", + " reward_coeff_overrides=first_stage.get('reward_coeff_overrides'),\n", + ")\n", + "first_environment = first_stage['env_cls'](config=first_config)\n", + "observation_dimension = first_environment.observation_space.shape[0]\n", + "number_of_actions = first_environment.action_space.n\n", + "first_environment.close()\n", + "\n", + "agent = PPOAgent(\n", + " observation_dimension=observation_dimension,\n", + " number_of_actions=number_of_actions,\n", + " learning_rate=learning_rate,\n", + " hidden_units=hidden_units,\n", + " clip_epsilon=clip_epsilon,\n", + " value_loss_coefficient=value_loss_coefficient,\n", + " entropy_coefficient=entropy_coefficient,\n", + " ppo_epochs=ppo_epochs,\n", + ")\n", + "random_agent = RandomAgent(number_of_actions=number_of_actions)\n", + "\n", + "if resume_from_checkpoint is not None:\n", + " loaded_metadata = agent.load_checkpoint(\n", + " resume_from_checkpoint,\n", + " load_optimizer_state=load_optimizer_state,\n", + " strict_shape_check=True,\n", + " )\n", + " print('Resumed from checkpoint:', resume_from_checkpoint)\n", + " print('Loaded metadata:', loaded_metadata)\n", + "\n", + "print('Global seed:', global_seed)\n", + "print('Shared observation dimension:', observation_dimension)\n", + "print('Shared action count:', number_of_actions)\n" + ] + }, + { + "cell_type": "markdown", + "id": "ec9ce9f5", + "metadata": {}, + "source": [ + "## Run the curriculum\n", + "\n", + "Each stage:\n", + "\n", + "- builds a fresh environment for that stage\n", + "- keeps the same PPO weights in memory\n", + "- saves `latest.pt` and `best.pt` inside a stage-specific checkpoint directory\n", + "- writes stage history to `curriculum_history.json`\n", + "- warms the next stage from the best checkpoint of the current stage\n" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "6f315830", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "=== Starting stage: sector_i_1ac_lateral ===\n" + ] + }, + { + "ename": "ValueError", + "evalue": "zip() argument 2 is longer than argument 1", + "output_type": "error", + "traceback": [ + "\u001b[31m---------------------------------------------------------------------------\u001b[39m", + "\u001b[31mValueError\u001b[39m Traceback (most recent call last)", + "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[7]\u001b[39m\u001b[32m, line 53\u001b[39m\n\u001b[32m 49\u001b[39m print(\u001b[33m'=== Starting stage:'\u001b[39m, stage_name, \u001b[33m'==='\u001b[39m)\n\u001b[32m 50\u001b[39m \n\u001b[32m 51\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m episode_index \u001b[38;5;28;01min\u001b[39;00m range(stage[\u001b[33m'training_episodes'\u001b[39m]):\n\u001b[32m 52\u001b[39m random_seed = stage[\u001b[33m'training_seed_start'\u001b[39m] + episode_index\n\u001b[32m---> \u001b[39m\u001b[32m53\u001b[39m total_reward, step_count, update_metrics = run_one_training_episode(\n\u001b[32m 54\u001b[39m environment=environment,\n\u001b[32m 55\u001b[39m agent=agent,\n\u001b[32m 56\u001b[39m random_seed=random_seed,\n", + "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[4]\u001b[39m\u001b[32m, line 156\u001b[39m, in \u001b[36mrun_one_training_episode\u001b[39m\u001b[34m(environment, agent, random_seed, discount_factor_gamma, gae_lambda)\u001b[39m\n\u001b[32m 152\u001b[39m reward_by_callsign,\n\u001b[32m 153\u001b[39m done_by_callsign,\n\u001b[32m 154\u001b[39m truncated_by_callsign,\n\u001b[32m 155\u001b[39m _info,\n\u001b[32m--> \u001b[39m\u001b[32m156\u001b[39m ) = environment.step(action_by_callsign)\n\u001b[32m 157\u001b[39m \n\u001b[32m 158\u001b[39m timestep_reward = float(sum(reward_by_callsign.values())) \u001b[38;5;28;01mif\u001b[39;00m reward_by_callsign \u001b[38;5;28;01melse\u001b[39;00m \u001b[32m0.0\u001b[39m\n\u001b[32m 159\u001b[39m timestep_done = all(done_by_callsign.values()) \u001b[38;5;28;01mif\u001b[39;00m done_by_callsign \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m\n", + "\u001b[36mFile \u001b[39m\u001b[32m/mnt/74F4CA39F4C9FCFC/Giles/Projects/project-bluebird/BluebirdATC/bluebird-gymnasium/bluebird_gymnasium/envs/base.py:1518\u001b[39m, in \u001b[36mBaseEnv.step\u001b[39m\u001b[34m(self, action)\u001b[39m\n\u001b[32m 1514\u001b[39m \u001b[38;5;28mself\u001b[39m.timestep += \u001b[32m1\u001b[39m\n\u001b[32m 1516\u001b[39m \u001b[38;5;66;03m# now call step function specific to centralised or decentralised\u001b[39;00m\n\u001b[32m 1517\u001b[39m \u001b[38;5;66;03m# update aircraft tracker, compute reward and generate the next state\u001b[39;00m\n\u001b[32m-> \u001b[39m\u001b[32m1518\u001b[39m ret_values = \u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mstep_fn\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mactions_rf_int\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mactions_st\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43maction\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 1520\u001b[39m \u001b[38;5;66;03m# reset the current time step outcomm buffer\u001b[39;00m\n\u001b[32m 1521\u001b[39m \u001b[38;5;28mself\u001b[39m._current_time_step_outcomm_buffer.clear()\n", + "\u001b[36mFile \u001b[39m\u001b[32m/mnt/74F4CA39F4C9FCFC/Giles/Projects/project-bluebird/BluebirdATC/bluebird-gymnasium/bluebird_gymnasium/envs/base.py:2207\u001b[39m, in \u001b[36mBaseEnv._step_decentralized\u001b[39m\u001b[34m(self, actions_rf_int, actions_st, original_actions)\u001b[39m\n\u001b[32m 2203\u001b[39m \u001b[38;5;66;03m# rewards\u001b[39;00m\n\u001b[32m 2204\u001b[39m \u001b[38;5;66;03m# if this is a new aircraft, there is no previous\u001b[39;00m\n\u001b[32m 2205\u001b[39m \u001b[38;5;66;03m# action. hence, assume a NOOP action.\u001b[39;00m\n\u001b[32m 2206\u001b[39m act = actions.get(callsign, ACTION_NOOP)\n\u001b[32m-> \u001b[39m\u001b[32m2207\u001b[39m _rewards[callsign] = \u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mcompute_reward\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mcallsign\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mact\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 2209\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m callsign \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m.ac_tracker_prev_step:\n\u001b[32m 2210\u001b[39m \u001b[38;5;66;03m# previous step (p_) outcomm and position status\u001b[39;00m\n\u001b[32m 2211\u001b[39m p_outcomm = \u001b[38;5;28mself\u001b[39m.ac_tracker_prev_step[callsign].outcomm_status\n", + "\u001b[36mFile \u001b[39m\u001b[32m/mnt/74F4CA39F4C9FCFC/Giles/Projects/project-bluebird/BluebirdATC/bluebird-gymnasium/bluebird_gymnasium/envs/base.py:2813\u001b[39m, in \u001b[36mBaseEnv.compute_reward\u001b[39m\u001b[34m(self, callsign, action)\u001b[39m\n\u001b[32m 2811\u001b[39m ac = simulator_env.aircraft[callsign]\n\u001b[32m 2812\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m ac.controllable:\n\u001b[32m-> \u001b[39m\u001b[32m2813\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43;01mfor\u001b[39;49;00m\u001b[30;43m \u001b[39;49m\u001b[30;43mstr_fn\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mcoeff\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43;01min\u001b[39;49;00m\u001b[30;43m \u001b[39;49m\u001b[30;43mzip\u001b[39;49m\u001b[30;43m(\u001b[39;49m\n\u001b[32m 2814\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mconfig\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mreward_config\u001b[39;49m\u001b[30;43m[\u001b[39;49m\u001b[30;43m\"\u001b[39;49m\u001b[30;43mfns\u001b[39;49m\u001b[30;43m\"\u001b[39;49m\u001b[30;43m]\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mconfig\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mreward_config\u001b[39;49m\u001b[30;43m[\u001b[39;49m\u001b[30;43m\"\u001b[39;49m\u001b[30;43mcoeffs\u001b[39;49m\u001b[30;43m\"\u001b[39;49m\u001b[30;43m]\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mstrict\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43;01mTrue\u001b[39;49;00m\n\u001b[32m 2815\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43m)\u001b[39;49m\u001b[30;43m:\u001b[39;49m\n\u001b[32m 2816\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mfn\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mregistry_reward_fn\u001b[39;49m\u001b[30;43m[\u001b[39;49m\u001b[30;43mstr_fn\u001b[39;49m\u001b[30;43m]\u001b[39;49m\n\u001b[32m 2817\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mrewards\u001b[39;49m\u001b[30;43m[\u001b[39;49m\u001b[30;43mstr_fn\u001b[39;49m\u001b[30;43m]\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mcoeff\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mfn\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mself\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mcallsign\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43maction\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n", + "\u001b[31mValueError\u001b[39m: zip() argument 2 is longer than argument 1" + ] + } + ], + "source": [ + "curriculum_history: list[dict] = []\n", + "\n", + "for stage_index in range(resume_from_stage_index, len(curriculum_stages)):\n", + " stage = curriculum_stages[stage_index]\n", + " stage_name = stage['stage_name']\n", + " stage_dir = curriculum_checkpoint_dir / stage_name\n", + " latest_checkpoint_path = stage_dir / 'latest.pt'\n", + " best_checkpoint_path = stage_dir / 'best.pt'\n", + "\n", + " config = make_training_config(\n", + " env_cls=stage['env_cls'],\n", + " num_aircraft=stage['num_aircraft'],\n", + " enable_vertical_actions=stage['enable_vertical_actions'],\n", + " scenario_duration_seconds=stage['scenario_duration_seconds'],\n", + " scenario_cls=stage['scenario_cls'],\n", + " scenario_args=stage['scenario_args'],\n", + " exit_window_width_nmi=stage['exit_window_width_nmi'],\n", + " reward_coeff_overrides=stage.get('reward_coeff_overrides'),\n", + " )\n", + " environment = stage['env_cls'](config=config)\n", + "\n", + " stage_observation_dimension = environment.observation_space.shape[0]\n", + " stage_number_of_actions = environment.action_space.n\n", + " if stage_observation_dimension != agent.observation_dimension:\n", + " raise ValueError(\n", + " f'Stage {stage_name} observation dimension {stage_observation_dimension} does not match current agent {agent.observation_dimension}.'\n", + " )\n", + " if stage_number_of_actions != agent.number_of_actions:\n", + " raise ValueError(\n", + " f'Stage {stage_name} action count {stage_number_of_actions} does not match current agent {agent.number_of_actions}.'\n", + " )\n", + "\n", + " training_rewards: list[float] = []\n", + " training_steps: list[int] = []\n", + " training_policy_losses: list[float] = []\n", + " training_value_losses: list[float] = []\n", + " training_entropies: list[float] = []\n", + " training_total_losses: list[float] = []\n", + " periodic_eval_episodes: list[int] = []\n", + " periodic_eval_learned_mean_rewards: list[float] = []\n", + " periodic_eval_learned_std_rewards: list[float] = []\n", + " periodic_eval_random_mean_rewards: list[float] = []\n", + " periodic_eval_random_std_rewards: list[float] = []\n", + "\n", + " best_mean_evaluation_reward = float('-inf')\n", + " best_stage_metadata: dict = {}\n", + "\n", + " print()\n", + " print('=== Starting stage:', stage_name, '===')\n", + "\n", + " for episode_index in range(stage['training_episodes']):\n", + " random_seed = stage['training_seed_start'] + episode_index\n", + " total_reward, step_count, update_metrics = run_one_training_episode(\n", + " environment=environment,\n", + " agent=agent,\n", + " random_seed=random_seed,\n", + " discount_factor_gamma=discount_factor_gamma,\n", + " gae_lambda=gae_lambda,\n", + " )\n", + "\n", + " training_rewards.append(total_reward)\n", + " training_steps.append(step_count)\n", + " training_policy_losses.append(float('nan') if update_metrics is None else update_metrics['policy_loss'])\n", + " training_value_losses.append(float('nan') if update_metrics is None else update_metrics['value_loss'])\n", + " training_entropies.append(float('nan') if update_metrics is None else update_metrics['entropy'])\n", + " training_total_losses.append(float('nan') if update_metrics is None else update_metrics['total_loss'])\n", + "\n", + " print(\n", + " '[train]',\n", + " f'stage={stage_name}',\n", + " f'episode={episode_index:03d}',\n", + " f'seed={random_seed}',\n", + " f'reward={total_reward:.3f}',\n", + " f'steps={step_count}',\n", + " )\n", + "\n", + " should_run_periodic_eval = (\n", + " (episode_index + 1) % periodic_eval_interval == 0\n", + " or episode_index == stage['training_episodes'] - 1\n", + " )\n", + "\n", + " if should_run_periodic_eval:\n", + " learned_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + " )\n", + " random_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=random_agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + " )\n", + "\n", + " periodic_eval_episodes.append(episode_index + 1)\n", + " periodic_eval_learned_mean_rewards.append(learned_eval['mean_reward'])\n", + " periodic_eval_learned_std_rewards.append(learned_eval['std_reward'])\n", + " periodic_eval_random_mean_rewards.append(random_eval['mean_reward'])\n", + " periodic_eval_random_std_rewards.append(random_eval['std_reward'])\n", + "\n", + " metadata = {\n", + " 'stage_name': stage_name,\n", + " 'stage_index': stage_index,\n", + " 'environment': stage['env_cls'].__name__,\n", + " 'view_type': 'decentralized',\n", + " 'num_aircraft': stage['num_aircraft'],\n", + " 'scenario_cls': stage['scenario_cls'],\n", + " 'scenario_args': stage['scenario_args'],\n", + " 'exit_window_width_nmi': stage['exit_window_width_nmi'],\n", + " 'enable_vertical_actions': stage['enable_vertical_actions'],\n", + " 'reward_coeff_overrides': stage.get('reward_coeff_overrides'),\n", + " 'observation_dimension': stage_observation_dimension,\n", + " 'number_of_actions': stage_number_of_actions,\n", + " 'episode': episode_index + 1,\n", + " 'learned_mean_reward': learned_eval['mean_reward'],\n", + " 'learned_std_reward': learned_eval['std_reward'],\n", + " 'random_mean_reward': random_eval['mean_reward'],\n", + " 'random_std_reward': random_eval['std_reward'],\n", + " 'evaluation_seeds': heldout_evaluation_seeds,\n", + " }\n", + " agent.save_checkpoint(latest_checkpoint_path, metadata=metadata)\n", + "\n", + " if learned_eval['mean_reward'] > best_mean_evaluation_reward:\n", + " best_mean_evaluation_reward = learned_eval['mean_reward']\n", + " best_stage_metadata = metadata\n", + " agent.save_checkpoint(best_checkpoint_path, metadata=metadata)\n", + " checkpoint_note = 'new best checkpoint'\n", + " else:\n", + " checkpoint_note = 'latest checkpoint only'\n", + "\n", + " print(\n", + " '[periodic-eval]',\n", + " f'stage={stage_name}',\n", + " f'episode={episode_index + 1:03d}',\n", + " f'learned_mean_reward={learned_eval[\"mean_reward\"]:.3f}',\n", + " f'random_mean_reward={random_eval[\"mean_reward\"]:.3f}',\n", + " checkpoint_note,\n", + " )\n", + "\n", + " stage_history = {\n", + " 'stage_name': stage_name,\n", + " 'environment': stage['env_cls'].__name__,\n", + " 'num_aircraft': stage['num_aircraft'],\n", + " 'scenario_cls': stage['scenario_cls'],\n", + " 'scenario_args': stage['scenario_args'],\n", + " 'exit_window_width_nmi': stage['exit_window_width_nmi'],\n", + " 'reward_coeff_overrides': stage.get('reward_coeff_overrides'),\n", + " 'best_mean_evaluation_reward': best_mean_evaluation_reward,\n", + " 'best_checkpoint_path': str(best_checkpoint_path),\n", + " 'latest_checkpoint_path': str(latest_checkpoint_path),\n", + " 'training_rewards': training_rewards,\n", + " 'training_steps': training_steps,\n", + " 'training_policy_losses': training_policy_losses,\n", + " 'training_value_losses': training_value_losses,\n", + " 'training_entropies': training_entropies,\n", + " 'training_total_losses': training_total_losses,\n", + " 'periodic_eval_episodes': periodic_eval_episodes,\n", + " 'periodic_eval_learned_mean_rewards': periodic_eval_learned_mean_rewards,\n", + " 'periodic_eval_random_mean_rewards': periodic_eval_random_mean_rewards,\n", + " 'best_stage_metadata': best_stage_metadata,\n", + " }\n", + " curriculum_history.append(stage_history)\n", + "\n", + " curriculum_checkpoint_dir.mkdir(parents=True, exist_ok=True)\n", + " with history_json_path.open('w', encoding='utf-8') as fp:\n", + " json.dump(\n", + " curriculum_history,\n", + " fp,\n", + " indent=2,\n", + " default=lambda obj: obj.item() if isinstance(obj, np.generic) else str(obj),\n", + " )\n", + "\n", + " if best_checkpoint_path.exists():\n", + " reloaded_metadata = agent.load_checkpoint(\n", + " best_checkpoint_path,\n", + " load_optimizer_state=False,\n", + " strict_shape_check=True,\n", + " )\n", + " print('Loaded best checkpoint for next stage warm start:', reloaded_metadata)\n", + "\n", + " environment.close()\n" + ] + }, + { + "cell_type": "markdown", + "id": "0fbaf36c", + "metadata": {}, + "source": [ + "## Stage summary plots\n", + "\n", + "This plot keeps the same “quick-look” style as the older PPO notebooks,\n", + "but compares progress stage by stage.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "82379466", + "metadata": {}, + "outputs": [], + "source": [ + "if not curriculum_history:\n", + " raise ValueError('No curriculum history recorded yet. Run the curriculum loop first.')\n", + "\n", + "num_stages = len(curriculum_history)\n", + "fig, axes = plt.subplots(num_stages, 2, figsize=(15, 5 * num_stages))\n", + "if num_stages == 1:\n", + " axes = np.asarray([axes])\n", + "\n", + "for row_index, stage_history in enumerate(curriculum_history):\n", + " training_rewards = stage_history['training_rewards']\n", + " plot_window = min(10, len(training_rewards))\n", + " if len(training_rewards) >= plot_window and plot_window > 0:\n", + " smoothed_rewards = np.convolve(\n", + " np.asarray(training_rewards, dtype=float),\n", + " np.ones(plot_window) / plot_window,\n", + " mode='valid',\n", + " )\n", + " else:\n", + " smoothed_rewards = np.array([])\n", + "\n", + " training_episode_indices = np.arange(1, len(training_rewards) + 1)\n", + " eval_episodes = np.asarray(stage_history['periodic_eval_episodes'])\n", + " learned_mean = np.asarray(stage_history['periodic_eval_learned_mean_rewards'])\n", + " random_mean = np.asarray(stage_history['periodic_eval_random_mean_rewards'])\n", + "\n", + " ax_left = axes[row_index, 0]\n", + " ax_left.plot(training_episode_indices, training_rewards, marker='o', alpha=0.25, label='raw reward')\n", + " if len(smoothed_rewards) > 0:\n", + " ax_left.plot(\n", + " np.arange(plot_window, len(training_rewards) + 1),\n", + " smoothed_rewards,\n", + " linewidth=2.5,\n", + " label=f'moving average (window={plot_window})',\n", + " )\n", + " ax_left.set_title(f\"{stage_history['stage_name']} training reward\")\n", + " ax_left.set_xlabel('Episode')\n", + " ax_left.set_ylabel('Total reward')\n", + " ax_left.grid(alpha=0.3)\n", + " ax_left.legend()\n", + "\n", + " ax_right = axes[row_index, 1]\n", + " ax_right.plot(eval_episodes, learned_mean, marker='o', label='learned policy')\n", + " ax_right.plot(eval_episodes, random_mean, marker='s', label='random baseline')\n", + " ax_right.set_title(f\"{stage_history['stage_name']} periodic evaluation\")\n", + " ax_right.set_xlabel('Training episode')\n", + " ax_right.set_ylabel('Mean evaluation reward')\n", + " ax_right.grid(alpha=0.3)\n", + " ax_right.legend()\n", + "\n", + "fig.suptitle('PPO Curriculum Training Summary', fontsize=16)\n", + "fig.tight_layout()\n", + "plt.show()\n" + ] + }, + { + "cell_type": "markdown", + "id": "a62b18e8", + "metadata": {}, + "source": [ + "## Render best-checkpoint GIFs for each curriculum stage\n", + "\n", + "This renders one evaluation GIF for the best checkpoint from each stage.\n", + "The final stage is included automatically.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2f59dcd1", + "metadata": {}, + "outputs": [], + "source": [ + "if not curriculum_history:\n", + " raise ValueError('No curriculum history recorded yet. Run the curriculum loop first.')\n", + "\n", + "stage_render_artifacts: dict[str, dict[str, Path]] = {}\n", + "gif_seed = heldout_evaluation_seeds[0]\n", + "\n", + "for stage_history in curriculum_history:\n", + " stage_name_to_render = stage_history['stage_name']\n", + " stage_config = next(stage for stage in curriculum_stages if stage['stage_name'] == stage_name_to_render)\n", + " best_checkpoint_path = Path(stage_history['best_checkpoint_path'])\n", + "\n", + " if not best_checkpoint_path.exists():\n", + " print(f'Skipping {stage_name_to_render}: checkpoint not found at {best_checkpoint_path}')\n", + " continue\n", + "\n", + " loaded_metadata = agent.load_checkpoint(\n", + " best_checkpoint_path,\n", + " load_optimizer_state=False,\n", + " strict_shape_check=True,\n", + " )\n", + " print()\n", + " print(f'Rendering stage {stage_name_to_render} checkpoint metadata: {loaded_metadata}')\n", + "\n", + " render_dir = Path.cwd() / 'renders' / f'{stage_name_to_render}_eval'\n", + " render_artifacts = render_evaluation_rollout_to_gif(\n", + " env_cls=stage_config['env_cls'],\n", + " stage_num_aircraft=stage_config['num_aircraft'],\n", + " enable_vertical_actions=stage_config['enable_vertical_actions'],\n", + " scenario_duration_seconds=stage_config['scenario_duration_seconds'],\n", + " agent=agent,\n", + " random_seed=gif_seed,\n", + " render_dir=render_dir,\n", + " gif_name=f'{stage_name_to_render}_best_eval_seed_{gif_seed}',\n", + " render_every_n_steps=render_every_n_steps,\n", + " gif_frame_duration_seconds=gif_frame_duration_seconds,\n", + " scenario_cls=stage_config['scenario_cls'],\n", + " scenario_args=stage_config['scenario_args'],\n", + " exit_window_width_nmi=stage_config['exit_window_width_nmi'],\n", + " reward_coeff_overrides=stage_config.get('reward_coeff_overrides'),\n", + " )\n", + "\n", + " stage_render_artifacts[stage_name_to_render] = render_artifacts\n", + " print(f\"Saved GIF to: {render_artifacts['gif_path']}\")\n", + " print(f\"Saved action trace to: {render_artifacts['action_trace_path']}\")\n", + " display(Image(filename=str(render_artifacts['gif_path'])))\n", + "\n", + "stage_render_artifacts\n" + ] + }, + { + "cell_type": "markdown", + "id": "fce82f9b", + "metadata": {}, + "source": [ + "## Practical notes\n", + "\n", + "This notebook is the recommended place to try:\n", + "\n", + "- more episodes per stage\n", + "- different stage orderings\n", + "- different aircraft counts\n", + "- checkpoint resume across notebook sessions\n", + "\n", + "Do not add climb/descent into this notebook’s checkpoint lineage unless you are willing\n", + "to start a new branch of experiments.\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.14.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/bluebird-gymnasium/examples/ppo_vertical_actions_agent.ipynb b/bluebird-gymnasium/examples/ppo_vertical_actions_agent.ipynb new file mode 100644 index 0000000..f28514d --- /dev/null +++ b/bluebird-gymnasium/examples/ppo_vertical_actions_agent.ipynb @@ -0,0 +1,1059 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# PPO Vertical-Actions Branch for Bluebird Gymnasium\n", + "\n", + "This notebook is a **separate branch experiment** based on [minimal_ppo_agent.ipynb](/home/sprite/BluebirdATC/bluebird-gymnasium/examples/minimal_ppo_agent.ipynb).\n", + "\n", + "Its purpose is to test the user idea of adding climb/descent actions, without contaminating\n", + "the lateral-only checkpoint lineage used by the baseline and curriculum notebooks.\n", + "\n", + "Key difference from the lateral-only PPO notebooks:\n", + "\n", + "- `simple_fl_climb = True`\n", + "- `simple_fl_descent = True`\n", + "- default scenario uses 2 aircraft so vertical actions have some chance of mattering\n", + "\n", + "Because the action count changes, checkpoints from the lateral-only notebooks are not directly compatible.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Environment note\n", + "\n", + "Run this notebook from the same working kernel/environment as the older PPO notebooks.\n", + "It uses the same overall structure and plotting style, but starts a separate checkpoint lineage.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Imports and path setup\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "from __future__ import annotations\n", + "\n", + "import json\n", + "import random\n", + "import shutil\n", + "import sys\n", + "from pathlib import Path\n", + "\n", + "import imageio\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import torch\n", + "import torch.nn as nn\n", + "import torch.optim as optim\n", + "from IPython.display import Image, display\n", + "\n", + "search_roots = [Path.cwd().resolve(), *Path.cwd().resolve().parents]\n", + "gym_root = None\n", + "dt_root = None\n", + "\n", + "for candidate in search_roots:\n", + " if (candidate / 'bluebird_gymnasium').exists():\n", + " gym_root = candidate\n", + " sibling_dt = candidate.parent / 'bluebird-dt'\n", + " if sibling_dt.exists():\n", + " dt_root = sibling_dt\n", + " break\n", + " if (candidate / 'bluebird-gymnasium').exists() and (candidate / 'bluebird-dt').exists():\n", + " gym_root = candidate / 'bluebird-gymnasium'\n", + " dt_root = candidate / 'bluebird-dt'\n", + " break\n", + "\n", + "if gym_root is None or dt_root is None:\n", + " raise RuntimeError('Could not locate local bluebird-gymnasium and bluebird-dt package roots.')\n", + "\n", + "sys.path.insert(0, str(gym_root))\n", + "sys.path.insert(0, str(dt_root))\n", + "\n", + "from bluebird_gymnasium.envs import EnvConfig, ViewType\n", + "from bluebird_gymnasium.envs.sector_i import SectorIEnv\n", + "from bluebird_gymnasium.envs.sector_xplus import SectorXPlusEnv\n", + "\n", + "print(f'Using bluebird-gymnasium from: {gym_root}')\n", + "print(f'Using bluebird-dt from: {dt_root}')\n", + "print(f'Torch version: {torch.__version__}')\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## PPO actor-critic model and agents\n", + "\n", + "This is the same PPO implementation line as the older minimal PPO notebook.\n", + "The experimental change is in the environment action configuration, not the model architecture.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "class ActorCriticNetwork(nn.Module):\n", + " \"\"\"Shared-trunk actor-critic network for PPO.\"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " observation_dimension: int,\n", + " number_of_actions: int,\n", + " hidden_units: int = 128,\n", + " ) -> None:\n", + " super().__init__()\n", + " self.trunk = nn.Sequential(\n", + " nn.Linear(observation_dimension, hidden_units),\n", + " nn.ReLU(),\n", + " nn.Linear(hidden_units, hidden_units),\n", + " nn.ReLU(),\n", + " )\n", + " self.policy_head = nn.Linear(hidden_units, number_of_actions)\n", + " self.value_head = nn.Linear(hidden_units, 1)\n", + "\n", + " def forward(self, observation_batch: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:\n", + " features = self.trunk(observation_batch)\n", + " action_logits = self.policy_head(features)\n", + " state_value = self.value_head(features).squeeze(-1)\n", + " return action_logits, state_value\n", + "\n", + "\n", + "class PPOAgent:\n", + " \"\"\"Minimal PPO agent with actor-critic network and checkpoint helpers.\"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " observation_dimension: int,\n", + " number_of_actions: int,\n", + " learning_rate: float = 3e-4,\n", + " hidden_units: int = 128,\n", + " clip_epsilon: float = 0.2,\n", + " value_loss_coefficient: float = 0.5,\n", + " entropy_coefficient: float = 0.01,\n", + " ppo_epochs: int = 4,\n", + " ) -> None:\n", + " self.observation_dimension = observation_dimension\n", + " self.number_of_actions = number_of_actions\n", + " self.actor_critic = ActorCriticNetwork(\n", + " observation_dimension=observation_dimension,\n", + " number_of_actions=number_of_actions,\n", + " hidden_units=hidden_units,\n", + " )\n", + " self.optimizer = optim.Adam(self.actor_critic.parameters(), lr=learning_rate)\n", + " self.clip_epsilon = clip_epsilon\n", + " self.value_loss_coefficient = value_loss_coefficient\n", + " self.entropy_coefficient = entropy_coefficient\n", + " self.ppo_epochs = ppo_epochs\n", + "\n", + " def choose_training_action(\n", + " self,\n", + " observation_vector: np.ndarray,\n", + " ) -> tuple[int, torch.Tensor, torch.Tensor, torch.Tensor]:\n", + " observation_tensor = torch.tensor(\n", + " observation_vector,\n", + " dtype=torch.float32,\n", + " ).unsqueeze(0)\n", + " action_logits, state_value = self.actor_critic(observation_tensor)\n", + " action_distribution = torch.distributions.Categorical(logits=action_logits)\n", + " sampled_action = action_distribution.sample()\n", + " action_log_probability = action_distribution.log_prob(sampled_action)\n", + "\n", + " return (\n", + " sampled_action.item(),\n", + " observation_tensor.squeeze(0),\n", + " action_log_probability.squeeze(0),\n", + " state_value.squeeze(0),\n", + " )\n", + "\n", + " def choose_evaluation_actions(\n", + " self,\n", + " observation_by_callsign: dict[str, np.ndarray],\n", + " ) -> dict[str, int]:\n", + " chosen_actions: dict[str, int] = {}\n", + " with torch.no_grad():\n", + " for callsign, observation_vector in observation_by_callsign.items():\n", + " observation_tensor = torch.tensor(\n", + " observation_vector,\n", + " dtype=torch.float32,\n", + " ).unsqueeze(0)\n", + " action_logits, _state_value = self.actor_critic(observation_tensor)\n", + " chosen_actions[callsign] = torch.argmax(action_logits, dim=-1).item()\n", + " return chosen_actions\n", + "\n", + " def update_from_trajectory(\n", + " self,\n", + " observations: list[torch.Tensor],\n", + " actions: list[torch.Tensor],\n", + " old_log_probabilities: list[torch.Tensor],\n", + " returns: torch.Tensor,\n", + " advantages: torch.Tensor,\n", + " ) -> dict[str, float] | None:\n", + " if not observations:\n", + " return None\n", + "\n", + " observation_tensor = torch.stack(observations)\n", + " action_tensor = torch.stack(actions).long()\n", + " old_log_probability_tensor = torch.stack(old_log_probabilities).detach()\n", + " returns_tensor = returns.detach()\n", + " advantages_tensor = advantages.detach()\n", + "\n", + " if advantages_tensor.numel() > 1:\n", + " advantages_std = advantages_tensor.std(unbiased=False)\n", + " if advantages_std > 1e-8:\n", + " advantages_tensor = (\n", + " (advantages_tensor - advantages_tensor.mean())\n", + " / (advantages_std + 1e-8)\n", + " )\n", + "\n", + " mean_policy_loss = 0.0\n", + " mean_value_loss = 0.0\n", + " mean_entropy = 0.0\n", + " mean_total_loss = 0.0\n", + "\n", + " for _epoch in range(self.ppo_epochs):\n", + " new_action_logits, new_state_values = self.actor_critic(observation_tensor)\n", + " action_distribution = torch.distributions.Categorical(logits=new_action_logits)\n", + " new_log_probabilities = action_distribution.log_prob(action_tensor)\n", + " entropy = action_distribution.entropy().mean()\n", + "\n", + " probability_ratio = torch.exp(new_log_probabilities - old_log_probability_tensor)\n", + " unclipped_objective = probability_ratio * advantages_tensor\n", + " clipped_objective = torch.clamp(\n", + " probability_ratio,\n", + " 1.0 - self.clip_epsilon,\n", + " 1.0 + self.clip_epsilon,\n", + " ) * advantages_tensor\n", + "\n", + " policy_loss = -torch.min(unclipped_objective, clipped_objective).mean()\n", + " value_loss = torch.nn.functional.mse_loss(new_state_values, returns_tensor)\n", + " total_loss = (\n", + " policy_loss\n", + " + self.value_loss_coefficient * value_loss\n", + " - self.entropy_coefficient * entropy\n", + " )\n", + "\n", + " self.optimizer.zero_grad()\n", + " total_loss.backward()\n", + " self.optimizer.step()\n", + "\n", + " mean_policy_loss += float(policy_loss.item())\n", + " mean_value_loss += float(value_loss.item())\n", + " mean_entropy += float(entropy.item())\n", + " mean_total_loss += float(total_loss.item())\n", + "\n", + " epoch_divisor = float(self.ppo_epochs)\n", + " return {\n", + " 'policy_loss': mean_policy_loss / epoch_divisor,\n", + " 'value_loss': mean_value_loss / epoch_divisor,\n", + " 'entropy': mean_entropy / epoch_divisor,\n", + " 'total_loss': mean_total_loss / epoch_divisor,\n", + " }\n", + "\n", + " def save_checkpoint(self, checkpoint_path: Path, metadata: dict | None = None) -> None:\n", + " checkpoint_path.parent.mkdir(parents=True, exist_ok=True)\n", + " payload = {\n", + " 'model_state_dict': self.actor_critic.state_dict(),\n", + " 'optimizer_state_dict': self.optimizer.state_dict(),\n", + " 'metadata': _to_python_types(metadata or {}),\n", + " }\n", + " torch.save(payload, checkpoint_path)\n", + "\n", + " def load_checkpoint(\n", + " self,\n", + " checkpoint_path: Path,\n", + " map_location: str = 'cpu',\n", + " load_optimizer_state: bool = True,\n", + " strict_shape_check: bool = True,\n", + " ) -> dict:\n", + " payload = torch.load(checkpoint_path, map_location=map_location, weights_only=False)\n", + " metadata = payload.get('metadata', {})\n", + " if strict_shape_check:\n", + " saved_obs_dim = metadata.get('observation_dimension')\n", + " saved_num_actions = metadata.get('number_of_actions')\n", + " if saved_obs_dim is not None and saved_obs_dim != self.observation_dimension:\n", + " raise ValueError(\n", + " f'Checkpoint observation dimension {saved_obs_dim} does not match current {self.observation_dimension}. '\n", + " 'Use a different checkpoint lineage.'\n", + " )\n", + " if saved_num_actions is not None and saved_num_actions != self.number_of_actions:\n", + " raise ValueError(\n", + " f'Checkpoint action count {saved_num_actions} does not match current {self.number_of_actions}. '\n", + " 'Use a different checkpoint lineage.'\n", + " )\n", + " self.actor_critic.load_state_dict(payload['model_state_dict'])\n", + " if load_optimizer_state and 'optimizer_state_dict' in payload:\n", + " self.optimizer.load_state_dict(payload['optimizer_state_dict'])\n", + " return metadata\n", + "\n", + "\n", + "\n", + "def _to_python_types(value):\n", + " if isinstance(value, dict):\n", + " return {key: _to_python_types(val) for key, val in value.items()}\n", + " if isinstance(value, (list, tuple)):\n", + " return [_to_python_types(item) for item in value]\n", + " if isinstance(value, Path):\n", + " return str(value)\n", + " if isinstance(value, np.generic):\n", + " return value.item()\n", + " return value\n", + "\n", + "\n", + "class RandomAgent:\n", + " \"\"\"Simple random baseline for comparison.\"\"\"\n", + "\n", + " def __init__(self, number_of_actions: int) -> None:\n", + " self.number_of_actions = number_of_actions\n", + "\n", + " def choose_evaluation_actions(\n", + " self,\n", + " observation_by_callsign: dict[str, np.ndarray],\n", + " ) -> dict[str, int]:\n", + " return {\n", + " callsign: random.randrange(self.number_of_actions)\n", + " for callsign in observation_by_callsign.keys()\n", + " }\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Environment configuration and rollout helpers\n", + "\n", + "This branch keeps the decentralized PPO structure, but turns on vertical actions.\n", + "It also defaults to 2 aircraft so the extra actions are not completely idle.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "def make_training_config(\n", + " env_cls,\n", + " num_aircraft: int,\n", + " k_nearest_aircraft: int = 1,\n", + " enable_vertical_actions: bool = False,\n", + ") -> EnvConfig:\n", + " config = env_cls.get_default_env_config(ViewType.DECENTRALIZED)\n", + "\n", + " config.state_repr_config = {\n", + " 'encoder_cls': 'extra_minimal',\n", + " 'k_nearest_aircraft': k_nearest_aircraft,\n", + " }\n", + "\n", + " config.action_config = {\n", + " 'simple_heading_left': True,\n", + " 'simple_heading_right': True,\n", + " 'simple_fl_climb': enable_vertical_actions,\n", + " 'simple_fl_descent': enable_vertical_actions,\n", + " 'simple_fl_exit': False,\n", + " }\n", + "\n", + " config.reward_config = {\n", + " 'fns': [\n", + " 'position_status_const',\n", + " 'lateral_centreline_distance_shaped',\n", + " 'safety_simple_avoidance_exp',\n", + " ],\n", + " 'coeffs': [1.0, 1.0, 1.2],\n", + " }\n", + "\n", + " config.view_config = {\n", + " 'type': ViewType.DECENTRALIZED.value,\n", + " 'decentralized_params': {},\n", + " }\n", + "\n", + " config.scenario_config = {\n", + " 'cls': 'tactical',\n", + " 'args': {\n", + " 'num_aircraft': num_aircraft,\n", + " 'balance': [0.0, 0.0, 1.0],\n", + " },\n", + " }\n", + "\n", + " return config\n", + "\n", + "\n", + "def compute_returns_and_advantages(\n", + " rewards: list[float],\n", + " values: list[torch.Tensor],\n", + " dones: list[bool],\n", + " discount_factor_gamma: float,\n", + " gae_lambda: float,\n", + ") -> tuple[torch.Tensor, torch.Tensor]:\n", + " rewards_tensor = torch.tensor(rewards, dtype=torch.float32)\n", + " values_tensor = torch.stack(values).detach().float()\n", + " dones_tensor = torch.tensor(dones, dtype=torch.float32)\n", + "\n", + " advantages = torch.zeros_like(rewards_tensor)\n", + " last_gae = torch.tensor(0.0)\n", + "\n", + " for timestep in reversed(range(len(rewards))):\n", + " if timestep == len(rewards) - 1:\n", + " next_value = torch.tensor(0.0)\n", + " else:\n", + " next_value = values_tensor[timestep + 1]\n", + "\n", + " next_nonterminal = 1.0 - dones_tensor[timestep]\n", + " delta = (\n", + " rewards_tensor[timestep]\n", + " + discount_factor_gamma * next_value * next_nonterminal\n", + " - values_tensor[timestep]\n", + " )\n", + " last_gae = delta + discount_factor_gamma * gae_lambda * next_nonterminal * last_gae\n", + " advantages[timestep] = last_gae\n", + "\n", + " returns = advantages + values_tensor\n", + " return returns, advantages\n", + "\n", + "\n", + "def run_one_training_episode(\n", + " environment,\n", + " agent: PPOAgent,\n", + " random_seed: int,\n", + " discount_factor_gamma: float,\n", + " gae_lambda: float,\n", + ") -> tuple[float, int, dict[str, float] | None]:\n", + " random.seed(random_seed)\n", + " np.random.seed(random_seed)\n", + " torch.manual_seed(random_seed)\n", + "\n", + " observation_by_callsign, _info = environment.reset(seed=random_seed)\n", + "\n", + " episode_is_done = False\n", + " episode_step_count = 0\n", + " episode_total_reward = 0.0\n", + "\n", + " observations: list[torch.Tensor] = []\n", + " actions: list[torch.Tensor] = []\n", + " old_log_probabilities: list[torch.Tensor] = []\n", + " values: list[torch.Tensor] = []\n", + " rewards: list[float] = []\n", + " dones: list[bool] = []\n", + "\n", + " while not episode_is_done:\n", + " action_by_callsign: dict[str, int] = {}\n", + " actions_this_step = 0\n", + "\n", + " for callsign, observation_vector in observation_by_callsign.items():\n", + " action_int, observation_tensor, action_log_probability, state_value = agent.choose_training_action(observation_vector)\n", + " action_by_callsign[callsign] = action_int\n", + " observations.append(observation_tensor)\n", + " actions.append(torch.tensor(action_int))\n", + " old_log_probabilities.append(action_log_probability.detach())\n", + " values.append(state_value.detach())\n", + " actions_this_step += 1\n", + "\n", + " (\n", + " next_observation_by_callsign,\n", + " reward_by_callsign,\n", + " done_by_callsign,\n", + " truncated_by_callsign,\n", + " _info,\n", + " ) = environment.step(action_by_callsign)\n", + "\n", + " timestep_reward = float(sum(reward_by_callsign.values())) if reward_by_callsign else 0.0\n", + " timestep_done = all(done_by_callsign.values()) if done_by_callsign else True\n", + "\n", + " rewards.extend([timestep_reward] * max(actions_this_step, 1))\n", + " dones.extend([bool(timestep_done)] * max(actions_this_step, 1))\n", + "\n", + " _ = truncated_by_callsign\n", + " episode_total_reward += timestep_reward\n", + " episode_is_done = timestep_done\n", + " observation_by_callsign = next_observation_by_callsign\n", + " episode_step_count += 1\n", + "\n", + " returns, advantages = compute_returns_and_advantages(\n", + " rewards=rewards,\n", + " values=values,\n", + " dones=dones,\n", + " discount_factor_gamma=discount_factor_gamma,\n", + " gae_lambda=gae_lambda,\n", + " )\n", + "\n", + " update_metrics = agent.update_from_trajectory(\n", + " observations=observations,\n", + " actions=actions,\n", + " old_log_probabilities=old_log_probabilities,\n", + " returns=returns,\n", + " advantages=advantages,\n", + " )\n", + "\n", + " return episode_total_reward, episode_step_count, update_metrics\n", + "\n", + "\n", + "def run_one_evaluation_episode(environment, evaluation_agent, random_seed: int) -> tuple[float, int]:\n", + " random.seed(random_seed)\n", + " np.random.seed(random_seed)\n", + " torch.manual_seed(random_seed)\n", + "\n", + " observation_by_callsign, _info = environment.reset(seed=random_seed)\n", + "\n", + " episode_is_done = False\n", + " episode_step_count = 0\n", + " episode_total_reward = 0.0\n", + "\n", + " while not episode_is_done:\n", + " action_by_callsign = evaluation_agent.choose_evaluation_actions(observation_by_callsign)\n", + " (\n", + " next_observation_by_callsign,\n", + " reward_by_callsign,\n", + " done_by_callsign,\n", + " truncated_by_callsign,\n", + " _info,\n", + " ) = environment.step(action_by_callsign)\n", + "\n", + " _ = truncated_by_callsign\n", + " episode_total_reward += float(sum(reward_by_callsign.values())) if reward_by_callsign else 0.0\n", + " episode_is_done = all(done_by_callsign.values()) if done_by_callsign else True\n", + " observation_by_callsign = next_observation_by_callsign\n", + " episode_step_count += 1\n", + "\n", + " return episode_total_reward, episode_step_count\n", + "\n", + "\n", + "def evaluate_agent_over_seeds(environment, evaluation_agent, evaluation_seeds: list[int]) -> dict:\n", + " rewards: list[float] = []\n", + " steps: list[int] = []\n", + "\n", + " for random_seed in evaluation_seeds:\n", + " total_reward, step_count = run_one_evaluation_episode(environment, evaluation_agent, random_seed)\n", + " rewards.append(total_reward)\n", + " steps.append(step_count)\n", + "\n", + " return {\n", + " 'seeds': evaluation_seeds,\n", + " 'rewards': rewards,\n", + " 'steps': steps,\n", + " 'mean_reward': float(np.mean(rewards)),\n", + " 'std_reward': float(np.std(rewards)),\n", + " 'mean_steps': float(np.mean(steps)),\n", + " }\n", + "\n", + "\n", + "def render_evaluation_rollout_to_gif(\n", + " env_cls,\n", + " stage_num_aircraft: int,\n", + " enable_vertical_actions: bool,\n", + " agent: PPOAgent,\n", + " random_seed: int,\n", + " render_dir: Path,\n", + " gif_name: str,\n", + " render_every_n_steps: int = 5,\n", + " gif_frame_duration_seconds: float = 0.2,\n", + ") -> Path:\n", + " render_config = make_training_config(\n", + " env_cls=env_cls,\n", + " num_aircraft=stage_num_aircraft,\n", + " enable_vertical_actions=enable_vertical_actions,\n", + " )\n", + " render_config.radar_config['display_actions'] = True\n", + " render_config.radar_config['render_dir'] = str(render_dir)\n", + " render_config.radar_config['prefix'] = 'frame'\n", + "\n", + " if render_dir.exists():\n", + " shutil.rmtree(render_dir)\n", + " render_dir.mkdir(parents=True, exist_ok=True)\n", + "\n", + " render_environment = env_cls(config=render_config)\n", + " render_environment.set_render_mode('file')\n", + "\n", + " random.seed(random_seed)\n", + " np.random.seed(random_seed)\n", + " torch.manual_seed(random_seed)\n", + "\n", + " observation_by_callsign, _info = render_environment.reset(seed=random_seed)\n", + " render_environment.render()\n", + "\n", + " episode_is_done = False\n", + " step_index = 0\n", + "\n", + " while not episode_is_done:\n", + " action_by_callsign = agent.choose_evaluation_actions(observation_by_callsign)\n", + " (\n", + " next_observation_by_callsign,\n", + " reward_by_callsign,\n", + " done_by_callsign,\n", + " truncated_by_callsign,\n", + " _info,\n", + " ) = render_environment.step(action_by_callsign)\n", + " _ = reward_by_callsign, truncated_by_callsign\n", + " step_index += 1\n", + " episode_is_done = all(done_by_callsign.values()) if done_by_callsign else True\n", + " if step_index % render_every_n_steps == 0 or episode_is_done:\n", + " render_environment.render()\n", + " observation_by_callsign = next_observation_by_callsign\n", + "\n", + " png_frames = sorted(render_dir.glob(f\"{render_config.radar_config['prefix']}_*.png\"))\n", + " if not png_frames:\n", + " raise RuntimeError(\n", + " f'No rendered PNG frames were written to {render_dir}. '\n", + " 'Expected at least one frame before GIF generation.'\n", + " )\n", + "\n", + " gif_path = render_dir / f'{gif_name}.gif'\n", + " images = [imageio.v3.imread(frame_path) for frame_path in png_frames]\n", + " imageio.mimsave(gif_path, images, loop=0, duration=gif_frame_duration_seconds)\n", + " render_environment.close()\n", + " return gif_path\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Set up the environment and inspect the shapes\n", + "\n", + "This is where the branch diverges from the baseline.\n", + "You should expect a larger `number_of_actions` than the lateral-only notebooks.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "config = make_training_config(\n", + " env_cls=SectorIEnv,\n", + " num_aircraft=2,\n", + " enable_vertical_actions=True,\n", + ")\n", + "environment = SectorIEnv(config=config)\n", + "\n", + "observation_dimension = environment.observation_space.shape[0]\n", + "number_of_actions = environment.action_space.n\n", + "\n", + "print(\n", + " 'environment shapes:',\n", + " f'observation_dimension={observation_dimension}',\n", + " f'number_of_actions={number_of_actions}',\n", + ")\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Hyperparameters and experiment settings\n", + "\n", + "This branch is more exploratory, so the defaults stay conservative:\n", + "\n", + "- same network width as the baseline\n", + "- same PPO structure\n", + "- separate checkpoint directory\n", + "- no attempt to load checkpoints from the lateral-only notebooks\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "learning_rate = 3e-4\n", + "hidden_units = 128\n", + "discount_factor_gamma = 0.99\n", + "gae_lambda = 0.95\n", + "clip_epsilon = 0.2\n", + "value_loss_coefficient = 0.5\n", + "entropy_coefficient = 0.01\n", + "ppo_epochs = 4\n", + "number_of_training_episodes = 100\n", + "training_seed_start = 500\n", + "periodic_eval_interval = 10\n", + "heldout_evaluation_seeds = list(range(600, 610))\n", + "render_every_n_steps = 5\n", + "gif_frame_duration_seconds = 0.2\n", + "\n", + "checkpoint_dir = Path.cwd() / 'checkpoints' / 'ppo_vertical_actions_agent'\n", + "latest_checkpoint_path = checkpoint_dir / 'latest.pt'\n", + "best_checkpoint_path = checkpoint_dir / 'best.pt'\n", + "resume_from_checkpoint: Path | None = None\n", + "load_optimizer_state = True\n", + "\n", + "agent = PPOAgent(\n", + " observation_dimension=observation_dimension,\n", + " number_of_actions=number_of_actions,\n", + " learning_rate=learning_rate,\n", + " hidden_units=hidden_units,\n", + " clip_epsilon=clip_epsilon,\n", + " value_loss_coefficient=value_loss_coefficient,\n", + " entropy_coefficient=entropy_coefficient,\n", + " ppo_epochs=ppo_epochs,\n", + ")\n", + "random_agent = RandomAgent(number_of_actions=number_of_actions)\n", + "\n", + "if resume_from_checkpoint is not None:\n", + " loaded_metadata = agent.load_checkpoint(\n", + " resume_from_checkpoint,\n", + " load_optimizer_state=load_optimizer_state,\n", + " strict_shape_check=True,\n", + " )\n", + " print('Resumed from checkpoint:', resume_from_checkpoint)\n", + " print('Loaded metadata:', loaded_metadata)\n", + "\n", + "training_rewards: list[float] = []\n", + "training_steps: list[int] = []\n", + "training_policy_losses: list[float] = []\n", + "training_value_losses: list[float] = []\n", + "training_entropies: list[float] = []\n", + "training_total_losses: list[float] = []\n", + "\n", + "periodic_eval_episodes: list[int] = []\n", + "periodic_eval_learned_mean_rewards: list[float] = []\n", + "periodic_eval_learned_std_rewards: list[float] = []\n", + "periodic_eval_random_mean_rewards: list[float] = []\n", + "periodic_eval_random_std_rewards: list[float] = []\n", + "\n", + "best_mean_evaluation_reward = float('-inf')\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## PPO training loop with periodic evaluation and checkpointing\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "for episode_index in range(number_of_training_episodes):\n", + " random_seed = training_seed_start + episode_index\n", + " total_reward, step_count, update_metrics = run_one_training_episode(\n", + " environment=environment,\n", + " agent=agent,\n", + " random_seed=random_seed,\n", + " discount_factor_gamma=discount_factor_gamma,\n", + " gae_lambda=gae_lambda,\n", + " )\n", + "\n", + " training_rewards.append(total_reward)\n", + " training_steps.append(step_count)\n", + " training_policy_losses.append(float('nan') if update_metrics is None else update_metrics['policy_loss'])\n", + " training_value_losses.append(float('nan') if update_metrics is None else update_metrics['value_loss'])\n", + " training_entropies.append(float('nan') if update_metrics is None else update_metrics['entropy'])\n", + " training_total_losses.append(float('nan') if update_metrics is None else update_metrics['total_loss'])\n", + "\n", + " print(\n", + " '[train]',\n", + " f'episode={episode_index:03d}',\n", + " f'seed={random_seed}',\n", + " f'reward={total_reward:.3f}',\n", + " f'steps={step_count}',\n", + " f'policy_loss={None if update_metrics is None else update_metrics[\"policy_loss\"]:.6f}' if update_metrics is not None else 'policy_loss=None',\n", + " f'value_loss={None if update_metrics is None else update_metrics[\"value_loss\"]:.6f}' if update_metrics is not None else 'value_loss=None',\n", + " )\n", + "\n", + " should_run_periodic_eval = (\n", + " (episode_index + 1) % periodic_eval_interval == 0\n", + " or episode_index == number_of_training_episodes - 1\n", + " )\n", + "\n", + " if should_run_periodic_eval:\n", + " learned_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + " )\n", + " random_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=random_agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + " )\n", + "\n", + " periodic_eval_episodes.append(episode_index + 1)\n", + " periodic_eval_learned_mean_rewards.append(learned_eval['mean_reward'])\n", + " periodic_eval_learned_std_rewards.append(learned_eval['std_reward'])\n", + " periodic_eval_random_mean_rewards.append(random_eval['mean_reward'])\n", + " periodic_eval_random_std_rewards.append(random_eval['std_reward'])\n", + "\n", + " metadata = {\n", + " 'episode': episode_index + 1,\n", + " 'train_seed': random_seed,\n", + " 'environment': 'SectorIEnv',\n", + " 'view_type': 'decentralized',\n", + " 'num_aircraft': 2,\n", + " 'observation_dimension': observation_dimension,\n", + " 'number_of_actions': number_of_actions,\n", + " 'enable_vertical_actions': True,\n", + " 'learned_mean_reward': learned_eval['mean_reward'],\n", + " 'learned_std_reward': learned_eval['std_reward'],\n", + " 'random_mean_reward': random_eval['mean_reward'],\n", + " 'random_std_reward': random_eval['std_reward'],\n", + " 'evaluation_seeds': heldout_evaluation_seeds,\n", + " }\n", + " agent.save_checkpoint(latest_checkpoint_path, metadata=metadata)\n", + "\n", + " if learned_eval['mean_reward'] > best_mean_evaluation_reward:\n", + " best_mean_evaluation_reward = learned_eval['mean_reward']\n", + " agent.save_checkpoint(best_checkpoint_path, metadata=metadata)\n", + " checkpoint_note = 'new best checkpoint'\n", + " else:\n", + " checkpoint_note = 'latest checkpoint only'\n", + "\n", + " print(\n", + " '[periodic-eval]',\n", + " f'episode={episode_index + 1:03d}',\n", + " f'learned_mean_reward={learned_eval[\"mean_reward\"]:.3f}',\n", + " f'random_mean_reward={random_eval[\"mean_reward\"]:.3f}',\n", + " checkpoint_note,\n", + " )\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Restore the best evaluated model\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "if not best_checkpoint_path.exists():\n", + " raise FileNotFoundError(f'Best checkpoint not found: {best_checkpoint_path}')\n", + "\n", + "loaded_metadata = agent.load_checkpoint(\n", + " best_checkpoint_path,\n", + " load_optimizer_state=False,\n", + " strict_shape_check=True,\n", + ")\n", + "print('Reloaded best checkpoint from:', best_checkpoint_path)\n", + "print('Best checkpoint metadata:')\n", + "loaded_metadata\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Final evaluation of the best checkpoint vs random baseline\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "best_policy_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + ")\n", + "random_policy_eval = evaluate_agent_over_seeds(\n", + " environment=environment,\n", + " evaluation_agent=random_agent,\n", + " evaluation_seeds=heldout_evaluation_seeds,\n", + ")\n", + "\n", + "print('Best PPO policy evaluation mean reward:', best_policy_eval['mean_reward'])\n", + "print('Best PPO policy evaluation std reward:', best_policy_eval['std_reward'])\n", + "print('Random baseline mean reward:', random_policy_eval['mean_reward'])\n", + "print('Random baseline std reward:', random_policy_eval['std_reward'])\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Useful plots\n", + "\n", + "This keeps the same plot style as the older PPO line so comparisons stay easy.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "def moving_average(values: list[float], window: int) -> np.ndarray:\n", + " if len(values) < window:\n", + " return np.array([])\n", + " kernel = np.ones(window) / window\n", + " return np.convolve(np.asarray(values, dtype=float), kernel, mode='valid')\n", + "\n", + "plot_window = min(10, len(training_rewards))\n", + "smoothed_rewards = moving_average(training_rewards, plot_window)\n", + "training_episode_indices = np.arange(1, len(training_rewards) + 1)\n", + "heldout_seed_indices = np.arange(len(heldout_evaluation_seeds))\n", + "periodic_eval_episodes_arr = np.asarray(periodic_eval_episodes)\n", + "learned_mean_arr = np.asarray(periodic_eval_learned_mean_rewards)\n", + "learned_std_arr = np.asarray(periodic_eval_learned_std_rewards)\n", + "random_mean_arr = np.asarray(periodic_eval_random_mean_rewards)\n", + "random_std_arr = np.asarray(periodic_eval_random_std_rewards)\n", + "\n", + "fig, axes = plt.subplots(3, 2, figsize=(15, 14))\n", + "\n", + "axes[0, 0].plot(training_episode_indices, training_rewards, marker='o', alpha=0.25, label='raw reward')\n", + "if len(smoothed_rewards) > 0:\n", + " axes[0, 0].plot(\n", + " np.arange(plot_window, len(training_rewards) + 1),\n", + " smoothed_rewards,\n", + " linewidth=2.5,\n", + " color='tab:blue',\n", + " label=f'moving average (window={plot_window})',\n", + " )\n", + "axes[0, 0].set_title('Training Reward per Episode')\n", + "axes[0, 0].set_xlabel('Episode')\n", + "axes[0, 0].set_ylabel('Total Reward')\n", + "axes[0, 0].legend()\n", + "axes[0, 0].grid(alpha=0.3)\n", + "\n", + "axes[0, 1].plot(training_episode_indices, training_steps, marker='o', color='tab:orange')\n", + "axes[0, 1].set_title('Training Episode Length')\n", + "axes[0, 1].set_xlabel('Episode')\n", + "axes[0, 1].set_ylabel('Steps')\n", + "axes[0, 1].grid(alpha=0.3)\n", + "\n", + "axes[1, 0].plot(periodic_eval_episodes_arr, learned_mean_arr, marker='o', label='learned policy')\n", + "axes[1, 0].fill_between(\n", + " periodic_eval_episodes_arr,\n", + " learned_mean_arr - learned_std_arr,\n", + " learned_mean_arr + learned_std_arr,\n", + " alpha=0.2,\n", + ")\n", + "axes[1, 0].plot(periodic_eval_episodes_arr, random_mean_arr, marker='s', label='random baseline')\n", + "axes[1, 0].fill_between(\n", + " periodic_eval_episodes_arr,\n", + " random_mean_arr - random_std_arr,\n", + " random_mean_arr + random_std_arr,\n", + " alpha=0.2,\n", + ")\n", + "axes[1, 0].set_title('Periodic Evaluation: Learned vs Random')\n", + "axes[1, 0].set_xlabel('Training Episode')\n", + "axes[1, 0].set_ylabel('Mean Evaluation Reward')\n", + "axes[1, 0].legend()\n", + "axes[1, 0].grid(alpha=0.3)\n", + "\n", + "axes[1, 1].plot(\n", + " heldout_seed_indices,\n", + " best_policy_eval['rewards'],\n", + " marker='o',\n", + " linewidth=2,\n", + " label='best PPO checkpoint',\n", + ")\n", + "axes[1, 1].plot(\n", + " heldout_seed_indices,\n", + " random_policy_eval['rewards'],\n", + " marker='s',\n", + " linewidth=2,\n", + " label='random baseline',\n", + ")\n", + "axes[1, 1].set_xticks(heldout_seed_indices)\n", + "axes[1, 1].set_xticklabels([str(seed) for seed in heldout_evaluation_seeds], rotation=45)\n", + "axes[1, 1].set_title('Final Evaluation Reward by Seed')\n", + "axes[1, 1].set_xlabel('Held-out Evaluation Seed')\n", + "axes[1, 1].set_ylabel('Total Reward')\n", + "axes[1, 1].legend()\n", + "axes[1, 1].grid(alpha=0.3)\n", + "\n", + "axes[2, 0].plot(training_episode_indices, training_policy_losses, label='policy loss')\n", + "axes[2, 0].plot(training_episode_indices, training_value_losses, label='value loss')\n", + "axes[2, 0].plot(training_episode_indices, training_total_losses, label='total loss')\n", + "axes[2, 0].set_title('PPO Loss Terms')\n", + "axes[2, 0].set_xlabel('Episode')\n", + "axes[2, 0].set_ylabel('Loss')\n", + "axes[2, 0].legend()\n", + "axes[2, 0].grid(alpha=0.3)\n", + "\n", + "axes[2, 1].plot(training_episode_indices, training_entropies, label='entropy', color='tab:green')\n", + "axes[2, 1].set_title('Policy Entropy During Training')\n", + "axes[2, 1].set_xlabel('Episode')\n", + "axes[2, 1].set_ylabel('Entropy')\n", + "axes[2, 1].grid(alpha=0.3)\n", + "\n", + "fig.suptitle('PPO Vertical-Actions Branch Summary', fontsize=16)\n", + "fig.tight_layout()\n", + "plt.show()\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Render the best checkpoint and save a GIF\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "gif_seed = heldout_evaluation_seeds[0]\n", + "render_dir = Path.cwd() / 'renders' / 'ppo_vertical_actions_agent_eval'\n", + "gif_path = render_evaluation_rollout_to_gif(\n", + " env_cls=SectorIEnv,\n", + " stage_num_aircraft=2,\n", + " enable_vertical_actions=True,\n", + " agent=agent,\n", + " random_seed=gif_seed,\n", + " render_dir=render_dir,\n", + " gif_name=f'best_vertical_branch_eval_seed_{gif_seed}',\n", + " render_every_n_steps=render_every_n_steps,\n", + " gif_frame_duration_seconds=gif_frame_duration_seconds,\n", + ")\n", + "\n", + "print(f'Saved GIF to: {gif_path}')\n", + "display(Image(filename=str(gif_path)))\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Optional cleanup\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "environment.close()\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/bluebird-gymnasium/tests/envs/test_flight_school.py b/bluebird-gymnasium/tests/envs/test_flight_school.py new file mode 100644 index 0000000..3d2f6ff --- /dev/null +++ b/bluebird-gymnasium/tests/envs/test_flight_school.py @@ -0,0 +1,39 @@ +import gymnasium as gym +import numpy as np + +import bluebird_gymnasium # noqa: F401 +from bluebird_dt.scenario_manager import Infinite +from bluebird_gymnasium.envs import SCENARIO_CLS, ViewType +from bluebird_gymnasium.envs.flight_school import FlightSchoolEnv + + +def test_flight_school_uses_infinite_scenario_manager(): + assert SCENARIO_CLS["infinite"] is Infinite + + +def test_flight_school_reset_step(): + config = FlightSchoolEnv.get_default_env_config(ViewType.CENTRALIZED) + config.scenario_config["args"]["random_seed"] = 7 + config.scenario_duration = 60 + gym_env = FlightSchoolEnv(config=config) + + obs, info = gym_env.reset() + assert isinstance(obs, np.ndarray) + assert isinstance(info, dict) + assert obs.shape == gym_env.observation_space.shape + assert isinstance(gym_env.scenario_manager, Infinite) + + obs, reward, done, truncated, info = gym_env.step(0) + assert isinstance(obs, np.ndarray) + assert isinstance(reward, float) + assert isinstance(done, bool) + assert isinstance(truncated, bool) + assert isinstance(info, dict) + + +def test_flight_school_gym_make(): + gym_env = gym.make("FlightSchoolEnv-v0") + obs, info = gym_env.reset() + + assert isinstance(obs, np.ndarray) + assert isinstance(info, dict)