Teaching a Drone Swarm to Catch Its Own Rogue Agent

Most multi-agent reinforcement learning tutorials stop at "agents learn to cooperate". They rarely ask what happens when one of the agents shouldn't be trusted. That question, not coordination, but coordination despite compromise, became the actual center of gravity for this project.

The brief was open-ended: model something with graph structure, learn something over it. I picked drone swarm coordination, because a swarm is one of the cleanest real-world cases where the graph isn't a modeling convenience, it's the actual shape of the problem. Drones only know about neighbors they can communicate with. Who's a neighbor changes every timestep as the formation moves. A fixed-size state vector can't represent that cleanly; a graph can.

Why a graph neural network fits here

In a typical MARL setup, each agent gets some flattened observation vector and learns a policy independently, maybe with a shared critic. That works when the number of agents is fixed and the interaction structure is static. Drone swarms break both assumptions, the communication graph is dynamic (drones drift in and out of range), and the thing you actually care about (a drone's anomaly score, its coordination priority) is fundamentally relational, not just a function of its own position and velocity.

A GNN handles this naturally: each drone is a node, each communication link is an edge. Message passing lets a drone's representation absorb information from its actual current neighbors, not a fixed neighbor list baked in at training time. I used a 2-layer AttentionEdgeConv (a GraphSAGE-style layer with attention weighting edges rather than treating every neighbor equally). This mattered specifically for the rogue-detection piece as a compromised drone shouldn't get to influence its neighbors' anomaly scores as strongly as a trusted one. Moreover the attention gives the model a mechanism to learn that weighting instead of assuming it.

Four heads, one graph

The model doesn't output a single thing. It has four prediction heads running off the same shared GNN backbone:

• Node head — velocity adjustment (ax, ay) for that drone
• Edge head — communication priority between two connected drones
• Graph head — a single swarm-wide health score
• Anomaly head — per-node rogue-detection score

Splitting these into separate heads instead of one monolithic output was a deliberate choice. I'd defend even though it made the loss function messier to write. Movement, communication topology, and security are genuinely different tasks that happen to share the same underlying graph representation. Forcing them through one output made early experiments unstable, the movement loss would dominate and the anomaly head would just learn to predict "everything's fine" as a local minimum. Separate heads, each with their own loss term, fixed that: the anomaly head can be wrong about coordination without being punished for it, and vice versa.

Training: QMIX, and why credit assignment is the actual hard part

The training loop uses a QMIX-inspired approach layered with a supervised anomaly detection loss. QMIX exists to solve a specific multi-agent problem: if you have a single team-level reward, how do you know which agent's action actually deserved credit for it? A naive shared-reward setup lets a bad agent hide behind good ones. The swarm's overall reward goes up, so every agent's policy gets reinforced identically, including the one that did nothing useful.

QMIX addresses this with a mixing network that combines individual agents' value estimates into a team value. It is done in a way that's constrained to be monotonic. If one agent's individual value goes up, the team value can't go down because of it. That constraint is what makes credit assignment tractable here, and it's the piece that made rogue-drone detection actually learnable instead of just noise.

The reward function itself has six components:

• Formation coherence
• Collision avoidance
• Communication maintenance
• Mission progress
• Rogue detection accuracy
• False-positive penalty

That last one mattered more than I expected going in. Early versions of the reward over-weighted raw detection accuracy, and the model found the laziest possible solution i.e. flag drones as rogue somewhat randomly, since a true positive was rewarded far more than a false positive was punished. Rebalancing that penalty was most of what got false positives down to zero — not a model architecture change, a reward shaping one. That's a pattern I keep running into in RL work generally: the model usually isn't broken, the incentive structure is.

What "100% detection, 0% false positives" actually means

Over 3000 episodes (600,000 environment steps), the trained model went from a 0% rogue-detection baseline to 100%, with zero false positives across every evaluation scenario, and a 193% improvement in episode reward (1.57 → 4.60) over the untrained baseline. Numbers like that are easy to state and easy to misread, so worth being specific about what's being measured: this is evaluated across three distinct scenarios in evaluate.py, comparing trained-vs-untrained behavior on identical initial conditions, with the rogue drone's identity held out from the model and only revealed for scoring. The false-positive number is the one I actually trust most, precisely because it was the hardest to get to zero, a model that flags nothing is trivially free of false positives too, so that number only means something in combination with 100% true-positive detection, which rules out the lazy solution.

Communication dropout resilience was tested the same way. Dropping edges from the graph mid-episode to simulate a drone losing contact and the mothership-child hierarchy turned out to matter more here than I originally designed it for. It wasn't just an organizational convenience; it gave the swarm a fallback communication path when the peer-to-peer mesh got disrupted, which kept anomaly scores stable even when the graph itself was actively degrading.

Building the dashboard wasn't optional

The project includes a full Flask + HTML5 Canvas dashboard, a live tactical map, training curves, a before/after evaluation view, and a side-by-side trained-vs-untrained comparison mode. I didn't plan to spend as much time on this as I did, but debugging a multi-agent system from loss curves alone turned out to be close to useless. A loss going down doesn't tell you which drone the model just decided was rogue, or whether the swarm's formation is actually coherent or just statistically "good enough" on paper. Watching the tactical map live, especially in compare mode against the untrained baseline, is what actually surfaced the early reward-shaping bug. I could see the trained model flagging drones almost at random long before the aggregate metrics made it obvious.

What I'd do differently

The anomaly detection loss is currently supervised — the rogue drone's identity is known during training, even though it's held out at evaluation. That's a reasonable simplification for a course project, but it's the first thing I'd change with more time: a fully unsupervised or self-supervised anomaly signal (deviation from expected neighbor behavior, rather than a labeled ground truth) would generalize to rogue behaviors the model never saw labeled examples of, which is the actual threat model in a real deployment.

I'd also want to push past 2D. The simulation environment is currently a 600×600 2D space, which was the right call for iterating quickly on the GNN and reward design, but a real swarm coordination problem is 3D with altitude-dependent communication range — a meaningfully harder graph construction problem, not just an extra coordinate.

Full code, the trained checkpoints, and the training data export pipeline are on GitHub if you want to dig into the reward function or the GNN architecture directly.