AgentATC — Agentic Air Traffic Control

Research-grade system that integrates a BlueSky ATM simulator, Graph RAG retrieval, and LLM-based ATC decision making. Includes a full PPO + LoRA reinforcement learning pipeline for fine-tuning the LLM to minimize controller interventions while maintaining safety.

Research Hypothesis

Graph-structured retrieval enables more context-efficient air traffic control, reducing controller interventions while maintaining safety.

The RL layer optimizes for minimal control actions (PRIMARY metric) under a hard safety constraint — exactly how a real ATC controller operates: speak only when necessary.

Quick Start

# Install dependencies
uv sync

# Run simulation (LLM-based ATC, no RL)
uv run python main.py                     # default config
uv run python main.py config.toml         # custom config

# Run RL training
uv run python -m rl.train_loop             # default config
uv run python -m rl.train_loop config.toml # custom config

Requirements

Python >= 3.10
CUDA GPU (for LLM inference and RL training; CPU works but is very slow)
~16 GB VRAM for 7B model with LoRA

Architecture

┌─────────────────────────────────────────────────────────────┐
│                        main.py                              │
│  11-step simulation loop (per tick):                        │
│  step sim → departures → classify phases → approach seq     │
│  → sectors → build graph → predict conflicts → detect       │
│  landed → agent decisions → metrics → status                │
└──────────────────┬──────────────────────────────────────────┘
                   │
    ┌──────────────┼──────────────┬────────────────┐
    ▼              ▼              ▼                ▼
simulator/     state/         rag/            agent/
BlueSky ATM    Airspace       Graph RAG       LLM Agent
Interface      Graph          Retriever       (Mistral-7B)
Traffic Gen    Conflict       k-hop ego       Prompt Builder
Lifecycle      Predictor      subgraph        Action Parser
Runway Mgr     Sectors        extraction      Coordinator
Weather                                       Cmd History
                                              Emergency
    ┌─────────────────────────────────────────────┐
    │                   rl/                       │
    │  PPO + LoRA training layer                  │
    │  reward.py      — ATC-aligned reward fn     │
    │  rollout_buffer  — per-aircraft GAE         │
    │  ppo_trainer     — LoRA + value head + PPO  │
    │  train_loop      — episode runner + wandb   │
    └─────────────────────────────────────────────┘

Module Breakdown

Directory	Purpose
`config/`	TOML-based configuration with nested dataclass schema
`simulator/`	BlueSky ATM wrapper: sim interface, traffic generation, lifecycle management, runway sequencing, weather
`state/`	Data contracts (`AircraftState`, `FlightPlan`), NetworkX airspace graph, conflict prediction, sectors
`rag/`	Graph RAG retriever: k-hop ego subgraph extraction with conflict chains, shared resources, proximity
`agent/`	LLM decision pipeline: prompt builder, action parser (11 actions), multi-aircraft coordinator, command history, emergency handling
`metrics/`	KPI collector and console/CSV dashboard
`rl/`	Reinforcement learning: reward function, rollout buffer, PPO trainer with LoRA, training loop

Configuration

All parameters live in config/schema.py with defaults. Override via TOML:

uv run python main.py config/default.toml

See config/default.toml for the full reference with inline comments.

Key Config Sections

Section	Key Parameters
`[sim]`	`n_steps`, `agent_interval_s` (decision cycle frequency)
`[traffic]`	`max_aircraft`, `spawn_interval`, `seed`
`[agent]`	`model_name`, `temperature`, `k_hops`, `max_context_aircraft`
`[graph]`	`proximity_radius_nm`, `conflict_horiz_nm` (5nm), `conflict_vert_ft` (1000ft)
`[prediction]`	`lookahead_times_s` (120, 180, 300s), `enabled`
`[weather]`	`enabled`, `scenario` (calm/westerly/storm)
`[rl]`	All RL hyperparameters — see below

Reinforcement Learning

Overview

The RL layer trains LoRA adapters on the base LLM via PPO, using the full simulation as the environment. It wraps (never modifies) the existing agent pipeline.

Reward Design

Rewards are hierarchical — safety dominates everything:

Component	Signal	Value	Priority
Safety	Separation violation	-100	Highest
Safety	Near-miss (predicted, <120s)	-10
Safety	Conflict resolved	+20
Workload	Any non-maintain action	-1.0	PRIMARY
Workload	Repeated same action	-0.5
Efficiency	Speed below 60% cruise	-0.5/min
Success	Landing	+50
Success	Takeoff	+20
Smoothness	Abrupt alt/hdg change	-0.01 * delta
Deadlock	Stuck aircraft / bad hold	-5	Lowest

Final reward is clamped to [-200, +100].

RL Config (`[rl]` section in TOML)

[rl]
enabled = false                # activate RL in main loop
dry_run = false                # random actions, no GPU — pipeline testing
colab = false                  # mirror checkpoints to Google Drive
use_finetuned = true           # false = base model only (no LoRA)

# PPO
ppo_epochs = 4
clip_epsilon = 0.2
gamma = 0.99
gae_lambda = 0.95
learning_rate = 1e-5
mini_batch_size = 4
max_grad_norm = 0.5
kl_target = 0.02

# LoRA
lora_r = 16
lora_alpha = 32
lora_target_modules = ["q_proj", "v_proj"]

# Training
num_episodes = 100
max_steps_per_episode = 4000
train_temperature = 0.6        # exploration
eval_temperature = 0.2         # greedy

# Checkpoints
checkpoint_dir = "checkpoints"
checkpoint_every = 10
eval_every = 10

Running RL Training

# Full training
uv run python -m rl.train_loop config/default.toml

# Dry run (no GPU, random actions — tests full pipeline)
# Set dry_run = true in your TOML, then:
uv run python -m rl.train_loop config/default.toml

# Resume from checkpoint (automatic — scans checkpoint_dir)
uv run python -m rl.train_loop config/default.toml

Weights & Biases

All training metrics are logged to wandb. Set up credentials before training:

# Option 1: environment variable
export WANDB_API_KEY=<your-key>

# Option 2: CLI login (one-time)
wandb login

Configure project/entity in TOML:

[rl]
wandb_project = "agentatc-rl"
wandb_entity = "your-team"       # leave empty for default

Tracked metrics per episode:

episode/total_reward — cumulative reward
episode/num_actions — non-maintain actions (primary research metric)
episode/num_violations — separation violations
episode/landings, episode/takeoffs — throughput
episode/action_rate — fraction of decisions that are interventions
train/policy_loss, train/value_loss, train/kl_div, train/entropy

Every eval_every episodes, an A/B comparison runs (finetuned vs baseline) and logs both under eval/finetuned/* and eval/baseline/*.

Google Colab

For training on Colab with Google Drive checkpoint persistence:

[rl]
colab = true
colab_save_dir = "/content/drive/MyDrive/aatc_models/"

Mount your Drive in the Colab notebook before training:

from google.colab import drive
drive.mount('/content/drive')

Evaluation Only

The evaluation loop runs automatically during training (every eval_every episodes). It compares the finetuned agent against the base model on identical scenarios, logging:

Conflicts, control actions, delay, throughput for both modes
Enables direct measurement of the research hypothesis

Dry Run

Set dry_run = true to test the full pipeline without loading any model:

Random actions with dummy logprobs
PPO update is skipped
Rewards, buffer, wandb logging all run normally
Useful for verifying the training loop, reward function, and infrastructure

Action Space

The LLM outputs one of 11 structured ATC actions:

Action	Parameters	Description
`climb`	`altitude` (ft)	Climb to altitude
`descend`	`altitude` (ft)	Descend to altitude
`turn_left`	`heading` (deg)	Turn left by degrees
`turn_right`	`heading` (deg)	Turn right by degrees
`maintain`	—	Keep current state (no-op)
`hold`	—	Orbit at current position
`land`	`speed`, `altitude`	Begin approach
`takeoff`	—	Begin departure climb
`go_around`	`altitude`	Abort approach
`emergency_descent`	`altitude`	Rapid descent
`priority_landing`	`speed`, `altitude`	Emergency landing

Key Conventions

All altitudes in feet, speeds in knots CAS, distances in nautical miles (SI internally in BlueSky)
ICAO separation minima: 5nm horizontal, 1000ft vertical
Emergency aircraft are always processed first in the agent cycle
Coordination: first-come priority — earlier decisions kept, later contradictions overridden
Dataclasses in state/airspace_graph.py are frozen with slots=True
Config system: all parameters in config/schema.py with defaults matching original hardcoded values

Project Structure

AgentATC/
├── config/
│   ├── schema.py              # Dataclass config schema (all defaults)
│   └── default.toml           # Reference TOML config
├── simulator/
│   ├── bluesky_interface.py   # BlueSky ATM wrapper
│   ├── traffic_generator.py   # Dense W-European airspace (10 airports, 70+ routes)
│   ├── lifecycle_manager.py   # Flight phase: GROUND→DEPARTURE→CRUISE→APPROACH→LANDED
│   ├── runway_manager.py      # Runway occupancy, departure spacing, approach sequencing
│   └── weather.py             # Wind layers synced to BlueSky
├── state/
│   ├── airspace_graph.py      # AircraftState, FlightPlan, FlightPhase, PredictedConflict
│   ├── graph_builder.py       # NetworkX graph: aircraft/waypoint/runway nodes, composite edges
│   ├── conflict_predictor.py  # Dead-reckoning trajectory projection (2-5 min lookahead)
│   └── sectors.py             # Airspace sectors, handoff detection
├── rag/
│   └── graph_retriever.py     # k-hop ego subgraph with conflict chains, shared resources, proximity
├── agent/
│   ├── atc_agent.py           # Singleton LLM pipeline (Mistral-7B / LLaMA-3-8B)
│   ├── prompt_builder.py      # Multi-section structured prompt
│   ├── action_parser.py       # Triple-fallback JSON extraction → 11 ATC actions
│   ├── coordinator.py         # Multi-aircraft contradiction detection + override
│   ├── command_history.py     # Per-aircraft ring buffer, oscillation detection
│   ├── emergency.py           # Emergency declarations, priority ordering
│   └── decision_logger.py     # JSONL trajectory logging
├── rl/
│   ├── reward.py              # ATC-aligned reward (safety >> workload >> efficiency)
│   ├── rollout_buffer.py      # Per-aircraft trajectories with GAE
│   ├── ppo_trainer.py         # LoRA + value head + PPO clipped surrogate
│   └── train_loop.py          # Episode runner, wandb logging, checkpointing
├── metrics/
│   └── collector.py           # KPI accumulator + console/CSV dashboard
├── main.py                    # Simulation orchestration loop
└── pyproject.toml

anirudhamar04/A-gen-TC

README

AgentATC — Agentic Air Traffic Control

Research Hypothesis

Quick Start

Requirements

Architecture

Module Breakdown

Configuration

Key Config Sections

Reinforcement Learning

Overview

Reward Design

RL Config (`[rl]` section in TOML)

Running RL Training

Weights & Biases

Google Colab

Evaluation Only

Dry Run

Action Space

Key Conventions

Project Structure

Contributors

Issues

anirudhamar04/A-gen-TC

README

AgentATC — Agentic Air Traffic Control

Research Hypothesis

Quick Start

Requirements

Architecture

Module Breakdown

Configuration

Key Config Sections

Reinforcement Learning

Overview

Reward Design

RL Config ([rl] section in TOML)

Running RL Training

Weights & Biases

Google Colab

Evaluation Only

Dry Run

Action Space

Key Conventions

Project Structure

Contributors

Issues

RL Config (`[rl]` section in TOML)