· 31 min read

Teaching an AI to Play Mini Metro: A Custom Go Engine and a 4,087-Action Graph Neural Network

For a very long time, I wanted to build an AI that could play a game noticeably better than a human. Something like Stockfish for chess, but scaled to a project that one person could actually design from scratch.

A while back, I played a game called Mini Metro. At first glance, it looks convincingly clean and simple: a white canvas, a few colored lines, geometric stations, and tiny rectangular trains picking up passengers. But once you play for ten minutes, the city expands, passenger demand explodes, transfer stations get overwhelmed, and whatever tidy transit network you thought you had designed collapses into chaos. I struggled to make real progress in it, and eventually I stepped away.

A few months later, the idea hit me: Mini Metro would be the perfect testbed for an autonomous planning agent. It isn’t just about fast reflexes; it’s a dynamic network-optimization problem under strict resource scarcity.

Before getting into how I built it, here is Mini Metro in simple terms:

Stations continuously spawn across an expanding urban map. Each station is assigned a geometric shape—predominantly circles, triangles, and squares, with rarer unique shapes like stars or pentagons appearing later. Passengers appear at these stations, and each passenger wants to travel to a station matching their target shape. You are given a limited pool of colored transit lines, locomotives, carriages, river tunnels, and high-capacity interchange station upgrades. Using these resources, you draw lines between stations, assign trains, and dynamically reroute tracks. If any station accumulates more passengers than its waiting capacity for too long (a continuous 45-second overcrowding countdown), the station overflows, the transit system collapses, and the game is over.

When I decided to tackle this with reinforcement learning, I knew the very first step was to turn the game state into something an AI could observe and act on. But Mini Metro is a proprietary commercial game with no public source code or API. I initially considered training a computer vision model on raw video capture from the actual game, but that would have added unnecessary complexity before even touching the core game mechanics.

Instead, I decided to build my own simulation engine from scratch in Go.

Before getting into the physics and neural network details, here is the high-level architecture of how the entire system connects together:

flowchart TB
    subgraph SIM["Go Simulation Engine (30 Hz Runtime)"]
        direction TB
        KINE["Kinematics Engine
30 Hz Sub-Ticks & Headways"]
        ROUTER["A* Pathfinding
Passenger Routing & Fair Boarding"]
        MASKER["Action Mask Generator
Validates 4,087 Action Slots"]
        STATE["Atomic State Pointer
atomic.Pointer[Simulator]"]

        KINE --> ROUTER --> MASKER --> STATE
    end

    subgraph CGO["Zero-Copy CGO Bridge (Shared Memory)"]
        direction TB
        BUFF_GRAPH["Graph Feature Tensors
30x32 Nodes, 200x2 Edges, 23 Global"]
        BUFF_MASK["Legal Action Mask
4,087 Slot Preallocated Bitmask"]
    end

    subgraph RL["Python ML Stack (PyTorch)"]
        direction TB
        GAT["3-Layer Dynamic GATv2
Spatial Attention on Nodes & Edges"]
        POOL["DiffPool Spatial Clustering
Soft-Assignment to 4 Regional Clusters"]
        LSTM["Recurrent LSTM Core
Preserves Temporal Congestion History"]
        HEADS["12-Way Hierarchical Policy
Action Intent + Spatial Cross-Attention"]

        GAT --> POOL --> LSTM --> HEADS
    end

    subgraph UI["Visual Client & Telemetry"]
        direction TB
        WS["WebSocket Server (Port 6969)
Broadcasts StateSnapshot at 30 Hz"]
        CANVAS["TypeScript Canvas 2D Client
Real-Time Track & Overcrowd Gauges"]

        WS --> CANVAS
    end

    %% End-to-End Pipeline Connections
    STATE -- "Zero-Copy Slice Fill (27 µs)" --> BUFF_GRAPH
    STATE -- "Bitmask Extraction" --> BUFF_MASK

    BUFF_GRAPH -- "C-ABI NumPy Tensors" --> GAT
    BUFF_MASK -- "Mask Ingestion" --> HEADS

    HEADS -- "Dispatched Action (Macro-Step)" --> KINE

    STATE -- "Thread-Safe Load" --> WS

    %% Excalidraw Color Palette
    style SIM fill:#FFE4E6,stroke:#F43F5E,stroke-width:2px,color:#881337
    style CGO fill:#FEF3C7,stroke:#D97706,stroke-width:2px,color:#78350F
    style RL fill:#E0E7FF,stroke:#6366F1,stroke-width:2px,color:#312E81
    style UI fill:#DCFCE7,stroke:#16A34A,stroke-width:2px,color:#14532D

    style KINE fill:#FFF1F2,stroke:#FB7185,color:#881337
    style ROUTER fill:#FFF1F2,stroke:#FB7185,color:#881337
    style MASKER fill:#FFF1F2,stroke:#FB7185,color:#881337
    style STATE fill:#FFF1F2,stroke:#FB7185,color:#881337

    style BUFF_GRAPH fill:#FFFBEB,stroke:#FBBF24,color:#78350F
    style BUFF_MASK fill:#FFFBEB,stroke:#FBBF24,color:#78350F

    style GAT fill:#EEF2FF,stroke:#818CF8,color:#312E81
    style POOL fill:#EEF2FF,stroke:#818CF8,color:#312E81
    style LSTM fill:#EEF2FF,stroke:#818CF8,color:#312E81
    style HEADS fill:#EEF2FF,stroke:#818CF8,color:#312E81

    style WS fill:#F0FDF4,stroke:#34D399,color:#14532D
    style CANVAS fill:#F0FDF4,stroke:#34D399,color:#14532D

1. Rebuilding the Simulation from Observable Behavior

Building a game engine without access to the original source code was hard. I had to deduce internal mechanics and model how different interdependent systems interact based on observation, community notes, and trial and error.

In Mini Metro, nothing lives in isolation:

  • How fast do stations appear, and how far apart must they spawn?
  • How does passenger arrival rate accelerate over time?
  • How do trains accelerate, cruise, brake into platforms, dwell while boarding, and turn around at line endpoints?
  • How do passengers choose their routes when multiple lines connect across the city?
  • How do lines cross rivers, and when are tunnel tokens refunded?

Early on, I spent a lot of time trying to match the exact hidden parameters of the original game. I experimented with different combinations of station and passenger spawn rates, but it was frustrating: some settings made the game progress too quickly and collapse in two minutes, while others felt too slow and boring. Eventually, I realized that my goal wasn’t to produce a 1:1 clone of Mini Metro for commercial release. My goal was to capture the authentic physical constraints, spatial trade-offs, and failure modes accurately enough to train an intelligent agent.

Once I adopted that mindset, development moved much faster.

The Physics: Continuous Kinematics and Headways

A train cannot simply jump discretely between stations like a piece on a chess board. In my engine, physics updates at a fixed sub-tick rate of 30 Hz ($\Delta t \approx 0.0333$ seconds).

Lines are ordered sequences of stations, and rail tracks are piecewise polylines. Each locomotive tracks its fractional progress $p \in [0.0, 1.0]$ along its current track segment, moving at a base cruising velocity ($v_{\text{base}} = 30.0$). To simulate realistic train momentum, trains smoothly ease in and out over the first and last 25% of each segment using a sinusoidal velocity envelope (dropping to a floor of $v_{\min} = 0.30$ at platform approach). Trains spend a minimum dwell time (0.4 seconds plus boarding time) picking up passengers, and slow down around sharp track bends based on the angle between incoming and outgoing track segments.

Continuous kinematics matter because transit timing directly dictates headway—the time interval between consecutive train arrivals at a station. A standard station holds up to 6 passengers before an overcrowding countdown begins. Once that countdown starts, you have roughly 45 seconds before the station overflows and triggers a game over. If a single train running an 8-station line takes 80 seconds to complete a round trip, no amount of clever routing can save that station during a passenger surge. The physics and the geometry determine whether survival is even physically possible.

Passenger Pathfinding: Why Real Commuters Hate Transfers

When passengers spawn, they need to know which train to board. I initially wrote an $A^*$ pathfinding algorithm that tried to be clever by factoring in live train locations and estimated waiting times.

That turned out to be a disaster. If a train was slightly delayed or dwelling at an adjacent station, passengers would suddenly decide that a complex, three-line detour across the city was two seconds faster. Platform queues would shuffle frantically, passengers would board the wrong trains, and when the delayed train finally pulled in, its intended passengers had already scattered.

To fix this, I changed the $A^*$ cost function to reflect how human transit riders actually think. Commuters care far more about avoiding transfers than shaving off hypothetical seconds:

$$\text{Cost} = \text{RideTime} + \text{DwellTime} + \text{TransferPenalty} + \text{ReversalPenalty}$$

By slapping a massive 1,000-second penalty on switching lines and a 500-second penalty on reversing direction, passengers consistently chose direct routes or clean single-transfer paths. This stabilized passenger flow throughout the network.

Boarding Equality at Transfer Hubs

Another subtle issue showed up at busy interchange stations where multiple lines intersect. In my first implementation, whichever train pulled up to the platform first would greedily board as many passengers as possible until its 6-passenger capacity was full.

If Line 1 arrived while a crowd was waiting, it would grab passengers who actually needed Line 2 simply because they were at the front of the queue. Line 1 would then depart carrying passengers it had no way of delivering, while Line 2 arrived moments later to find its passengers displaced. To stop transfer hubs from cannibalizing themselves, I implemented fair platform boarding: passengers only board trains that directly advance them along their $A^*$ itinerary toward their destination shape.

flowchart TB
    subgraph QUEUE["Platform Queue: Interchange Station"]
        direction LR
        P1["Passenger 1
Target: Square (Line 1)"]
        P2["Passenger 2
Target: Triangle (Line 2)"]
        P3["Passenger 3
Target: Square (Line 1)"]
        P4["Passenger 4
Target: Circle (Line 2)"]
        P5["Passenger 5
Target: Triangle (Line 1)"]
    end

    subgraph TRAINS["Fair Platform Boarding Dispatch"]
        direction LR
        T1["Line 1 Locomotive
Boards Sq-1 & Tri-1 Only"]
        T2["Line 2 Locomotive
Boards Tri-2 & Cir-2 Only"]
    end

    P1 & P3 & P5 --> T1
    P2 & P4 --> T2

    style QUEUE fill:#EFF6FF,stroke:#3B82F6,stroke-width:2px,color:#1E3A8A
    style TRAINS fill:#FEF3C7,stroke:#D97706,stroke-width:2px,color:#78350F
    style P1 fill:#DBEAFE,stroke:#2563EB,color:#1E40AF
    style P2 fill:#FCE7F3,stroke:#DB2777,color:#831843
    style P3 fill:#DBEAFE,stroke:#2563EB,color:#1E40AF
    style P4 fill:#FCE7F3,stroke:#DB2777,color:#831843
    style P5 fill:#DBEAFE,stroke:#2563EB,color:#1E40AF
    style T1 fill:#FFFBEB,stroke:#F59E0B,color:#78350F
    style T2 fill:#FFF1F2,stroke:#F43F5E,color:#881337

2. Bridging Go and Python: The Zero-Copy CGO Bridge

Once the simulation was working, I had to connect it to the machine learning stack.

Python is the undisputed for deep learning because of PyTorch, automatic differentiation, and GPU libraries. But Python’s interpreter overhead and Global Interpreter Lock (GIL) make it poorly suited for running dozens of high-speed simulation environments in parallel. Go, on the other hand, is fast, statically typed, compiles in seconds, and handles concurrency with goroutines effortlessly.

The question was: how could I let a Go simulation and a Python RL pipeline talk to each other without killing throughput?

In past projects, I had connected different languages using JSON over standard OS pipes (stdout/stdin) or local HTTP sockets. But in reinforcement learning, the training loop steps environments millions of times. Serializing graph objects with dozens of stations and edges into JSON strings, piping them across an OS boundary, and parsing them back into Python objects creates a massive bottleneck. The GPU spends half its time sitting idle waiting for the CPU to finish serializing text.

Instead, I compiled the Go simulation engine into a C-shared library:

go build -buildmode=c-shared -o libminimetro.so ./c_api

On the Python side, I used ctypes to load libminimetro.so directly into Python’s process memory.

flowchart TB
    subgraph PY["Python RL Runtime (PyTorch / ctypes)"]
        direction TB
        ALLOC["Preallocated NumPy Memory Buffers:
• out_nodes: 30 × 32 floats
• out_edges: 200 × 2 ints
• out_globals: 23 floats
• out_mask: 4,087 bytes"]
    end

    subgraph GO["Go Simulation Engine (CGO C-ABI)"]
        direction TB
        RECV["Receives Raw Memory Pointers
(c_api/main.go)"]
        SLICE["Direct Memory Fill via unsafe.Slice (27 µs):
nodesSlice := unsafe.Slice((*float32)(outNodes), 30*32)
obs.WriteVectorizedObservation(...)"]
        RECV --> SLICE
    end

    ALLOC -- "Direct Memory Pointers
(Zero-Copy C-ABI FFI)" --> RECV

    style PY fill:#EFF6FF,stroke:#3B82F6,stroke-width:2px,color:#1E3A8A
    style GO fill:#FEF3C7,stroke:#D97706,stroke-width:2px,color:#78350F
    style ALLOC fill:#DBEAFE,stroke:#2563EB,color:#1E40AF
    style RECV fill:#FFFBEB,stroke:#F59E0B,color:#78350F
    style SLICE fill:#FFFBEB,stroke:#F59E0B,color:#78350F

Memory ownership is kept strictly on the Python side. During environment initialization in ml/env.py, Python allocates contiguous memory buffers for station node features, track edge lists, global metrics, and legal action masks. When the environment steps, Python passes raw memory pointers into Go via CGO. The Go engine uses unsafe.Slice to populate Python’s memory directly in place.

There are no network sockets, no pipes, and zero JSON serialization on the training path. When I benchmarked this with ml/bench.py, the results were clear:

  • Advancing a macro-step in Go took 77 microseconds (~13,000 steps/second).
  • Extracting the full graph observation into Python memory took 27 microseconds.

That gave me between 450 and 650 environment steps per second across 16 parallel CPU workers on a standard laptop (train_local.py), and 1,200 to 2,000 steps per second during multi-map GPU training (train.py).

3. The Visual Client: Rapid Prototyping with Canvas 2D and WebSockets

Early on, I realized that debugging transit policies purely through terminal loss numbers and reward curves is like flying blind. If an agent’s score suddenly drops from 100 to 12, did it run out of tunnel tokens? Did it get trapped in a line-deletion loop? Or did a station overflow because a train got stuck dwelling at a platform?

I needed to see what the agent was actually building in real time.

My primary focus on this project was the core systems engine—the Go simulator, CGO shared memory bindings, and the PyTorch RL architecture. Rather than building a heavyweight frontend application, I wanted a lean, responsive visualizer that could keep up with real-time simulation updates. I built a lightweight web client in ui/ using Vanilla TypeScript, HTML5 Canvas 2D, and Vite, keeping it deliberately dependency-light so I could pipe raw WebSocket data directly into a 2D rendering loop.

It’s a practical, focused debugging tool: it has some visual quirks like overlapping track curve offsets when multiple lines share a corridor and basic HUD controls, but for streaming game state at 30 Hz and verifying the agent’s spatial decisions visually, it got the job done.

flowchart TB
    subgraph UI_BOX["Frontend Client (TypeScript / Canvas 2D)"]
        direction TB
        CANV["Canvas 2D Renderer
Bezier Track Paths & Station Shapes"]
        HUD["HUD Telemetry Overlay
Radial Overcrowd Timers"]
        CANV <--> HUD
    end

    subgraph SERVER_BOX["Go Standalone Server (Port 6969)"]
        direction TB
        PHYSICS["Physics Loop (30 Hz Sub-Ticks)
atomic.Pointer[Simulator]"]
        WS_HUB["WebSocket Hub (/ws)
Broadcasts StateSnapshot JSON"]
        API["REST Endpoint (GET /api/obs)
Direct Inference Hook"]
        PHYSICS --> WS_HUB & API
    end

    SERVER_BOX <== "WebSocket Stream
(30 Hz StateSnapshot JSON)" ==> UI_BOX

    style UI_BOX fill:#DCFCE7,stroke:#16A34A,stroke-width:2px,color:#14532D
    style SERVER_BOX fill:#EFF6FF,stroke:#3B82F6,stroke-width:2px,color:#1E3A8A
    style CANV fill:#F0FDF4,stroke:#34D399,color:#14532D
    style HUD fill:#F0FDF4,stroke:#34D399,color:#14532D
    style PHYSICS fill:#DBEAFE,stroke:#2563EB,color:#1E40AF
    style WS_HUB fill:#DBEAFE,stroke:#2563EB,color:#1E40AF
    style API fill:#DBEAFE,stroke:#2563EB,color:#1E40AF

On the backend, hooking it up was clean and reliable. The Go server runs a standalone binary on port 6969 (simulator/cmd/server/main.go). It broadcasts StateSnapshot JSON objects over WebSockets at 30 Hz, carrying station coordinates, active line paths, moving train positions, passenger counts, and overcrowding timer percentages.

To keep the web server completely thread-safe without lock contention, I stored the running simulator behind an atomic pointer:

simPtr atomic.Pointer[engine.Simulator]

The 30 Hz broadcast loop reads the snapshot safely using s.simPtr.Load(), while player inputs or live AI actions arriving from WebSockets are pushed into an internal Go channel (chan actionMsg, capacity 512) and executed cleanly on the next tick.

Even with a few rough edges in track curve overlap, having this visual canvas made debugging so much easier. I could watch the trains move, see radial countdown gauges fill up above congested platforms, and test lines manually with mouse drags.

4. Representing the Transit Network as a Graph

With the simulation and visualization in place, the next challenge was figuring out what the neural network should observe.

In many arcade games, you can feed raw screen pixels into a Convolutional Neural Network (CNN). But in Mini Metro, that would be wasteful. Most of the screen is empty white space, stations appear at arbitrary coordinates, and what truly matters is connectivity—how stations link together into a transit network.

A naive alternative is to flatten all station coordinates and queues into a fixed-length vector. But that fails for three reasons:

  1. Variable Station Count: A game starts with 3 stations and grows to 20 or 30 over time.
  2. Permutation Invariance: A station’s physical role does not change just because its internal slice index in Go happens to be 2 instead of 7.
  3. Relational Structure: A station with 5 waiting passengers is completely safe if an empty train is 3 seconds away; that exact same station is in critical danger if no train is scheduled to arrive for the next 45 seconds.

Because transit networks are inherently graphs, I formatted the observation as an attributed spatial graph: stations are nodes, rail tracks are edges, and inventory metrics are global attributes.

flowchart TB
    subgraph NODES["Station Nodes (Up to 30 × 32 Floats)"]
        direction TB
        N_FEATS["Spatial Coordinates (x, y)
One-Hot Station Shape (10)
Passenger Queue per Shape
Queue Velocity (dq/dt)
Approaching Train Proximity
Overcrowd Timer Fraction"]
    end

    subgraph EDGES["Track Edges (Up to 200 × 10 Floats)"]
        direction TB
        E_FEATS["Connected Line ID (0..6)
Euclidean Segment Distance
Water / River Crossing Flag"]
    end

    subgraph GLOBALS["Global Context Vector (23 Floats)"]
        direction TB
        G_FEATS["Inventory: Lines, Locomotives, Carriages, Tunnels, Interchanges
Simulation Time, Delivered Passengers, Active Sunday Reward Cards"]
    end

    NODES <== "Network Adjacency Topology" ==> EDGES
    NODES --> GLOBALS
    EDGES --> GLOBALS

    style NODES fill:#EFF6FF,stroke:#3B82F6,stroke-width:2px,color:#1E3A8A
    style EDGES fill:#FEF3C7,stroke:#D97706,stroke-width:2px,color:#78350F
    style GLOBALS fill:#FDF2F8,stroke:#EC4899,stroke-width:2px,color:#831843
    style N_FEATS fill:#DBEAFE,stroke:#2563EB,color:#1E40AF
    style E_FEATS fill:#FFFBEB,stroke:#F59E0B,color:#78350F
    style G_FEATS fill:#FCE7F3,stroke:#DB2777,color:#831843

The feature set went through several iterations:

  • Station Node Features (32 dimensions): Spatial coordinates $(x, y)$, one-hot station shape, passenger queue counts categorized by target shape (10 shapes), active line membership flags, interchange status, local connection degree, overcrowding timer fraction, and two crucial dynamic features: queue growth rate ($\Delta q / \Delta t$) and approaching train proximity/capacity.
  • Track Edge Features (10 dimensions): One-hot line identity, Euclidean segment length, and water-crossing indicators.
  • Global Context Vector (23 dimensions): Remaining inventory tokens (lines, locomotives, carriages, river tunnels, interchange upgrades), elapsed time, total passengers delivered, and the active Sunday upgrade cards offered by the game.

Learning the Hard Way: Why Static Snapshots Fail

In my first version, node features were only 25 dimensions and did not include temporal derivatives or train proximity.

The policy struggled badly. When you inspect why, the problem is obvious: a static snapshot of a queue is fundamentally ambiguous. Five passengers waiting at a station with a high-capacity train pulling up to the platform is a solved problem. Five passengers at a station with an empty track and no train in sight is an emergency.

Adding explicit train proximity and queue velocity features gave the network the context it needed to gauge urgency before timers started expiring.

Similarly, my earliest global vector tracked resource counts but completely omitted the weekly upgrade options offered every in-game Sunday. The model was expected to choose between upgrade cards (like 1 Line vs. 1 Locomotive + 1 Tunnel) without being able to see what choices were on the table! Adding one-hot representations of the active reward cards into the global vector fixed that blindness immediately.

5. The Neural Network: GATv2, DiffPool, and LSTM

To process this graph observation, I designed an Actor-Critic architecture (ml/model.py) combining three key mechanisms:

flowchart TB
    subgraph INPUTS["Input Graph Representations"]
        direction LR
        N_IN["Station Features
[N, 32]"]
        E_IN["Track Features
[E, 10]"]
        G_IN["Global Context
[23]"]
    end

    subgraph GAT_LAYER["Spatial Attention Backbone"]
        direction TB
        GAT_CORE["3-Layer Dynamic GATv2
Multi-Head Attention (K=4 Heads, H=128)"]
        DIFF["DiffPool Spatial Clustering
Softmax Assignment to 4 Regional Clusters"]
        GAT_CORE --> DIFF
    end

    subgraph TEMPORAL["Temporal Integration Core"]
        direction TB
        CONCAT["Concatenation (DiffPool 4H + Global 1H)
Unified 5H Embedding Vector"]
        LSTM_CORE["Recurrent LSTM Core (5H → 5H)
Preserves Temporal Congestion History"]
        CONCAT --> LSTM_CORE
    end

    subgraph OUTPUTS["Hierarchical Action Heads"]
        direction LR
        H_INTENT["12-Way Action Type Head
Selects High-Level Intent"]
        H_PARAMS["Conditioned Parameter Heads
Relational Transformer Scorer"]
    end

    N_IN & E_IN --> GAT_CORE
    DIFF --> CONCAT
    G_IN --> CONCAT
    LSTM_CORE --> H_INTENT & H_PARAMS

    style INPUTS fill:#EFF6FF,stroke:#3B82F6,stroke-width:2px,color:#1E3A8A
    style GAT_LAYER fill:#FEF3C7,stroke:#D97706,stroke-width:2px,color:#78350F
    style TEMPORAL fill:#F5F3FF,stroke:#8B5CF6,stroke-width:2px,color:#4C1D95
    style OUTPUTS fill:#DCFCE7,stroke:#16A34A,stroke-width:2px,color:#14532D

    style N_IN fill:#DBEAFE,stroke:#2563EB,color:#1E40AF
    style E_IN fill:#DBEAFE,stroke:#2563EB,color:#1E40AF
    style G_IN fill:#DBEAFE,stroke:#2563EB,color:#1E40AF
    style GAT_CORE fill:#FFFBEB,stroke:#F59E0B,color:#78350F
    style DIFF fill:#FFFBEB,stroke:#F59E0B,color:#78350F
    style CONCAT fill:#EDE9FE,stroke:#7C3AED,color:#5B21B6
    style LSTM_CORE fill:#EDE9FE,stroke:#7C3AED,color:#5B21B6
    style H_INTENT fill:#F0FDF4,stroke:#34D399,color:#14532D
    style H_PARAMS fill:#F0FDF4,stroke:#34D399,color:#14532D

1. Dynamic Graph Attention (GATv2)

Standard Graph Convolutional Networks (GCNs) average neighborhood features using static, degree-based weights. But in transit planning, neighboring stations are not equally important: a heavily congested interchange hub demands far more attention from adjacent stations than an empty terminus.

I used 3 layers of GATv2 (Brody et al., 2021) with 4 attention heads. Unlike standard GATv1—where the ranking of neighbors cannot depend on the query node—GATv2 computes dynamic attention conditioned on the source station, destination station, and connecting track edge features.

I also added an edge-update MLP so that edge representations evolve across layers alongside node embeddings, allowing track segments to carry accumulated congestion context across the network.

2. Hierarchical Spatial Pooling (DiffPool)

After passing through the GATv2 layers, the graph representations need to be summarized into a vector for the policy and value heads.

A simple global mean or max pool over all station embeddings completely flattens regional structure—a passenger crisis in North London looks identical to a crisis in South London. To preserve regional context, I used Differentiable Pooling (DiffPool) (Ying et al., 2018). A lightweight assignment MLP maps the $N$ station embeddings into 4 learned regional clusters, soft-pooling the graph into a compact 4-region representation.

Flattening this pooled representation gives a $4H$-dimensional regional summary. Concatenating it with the updated global context ($1H$) produces a unified $5H$-dimensional system representation that captures both macro city-wide inventory and regional demand clusters.

3. Recurrent Memory (LSTM)

Because passenger queues accumulate dynamically and train trips take dozens of simulated seconds, a single observation cannot capture the full trajectory. The unified $5H$ representation feeds into an LSTM recurrent core ($\text{hidden\_dim} = 5H$), preserving temporal memory across consecutive macro-steps.

6. The Action Space Nightmare: 4,087 Actions and the Inaction Freeze

Defining the action space for Mini Metro was easily one of the most frustrating parts of the project.

A single decision in Mini Metro requires choosing an operation and specifying its parameters: which line to create, which endpoints to extend, which track segments to insert stations into, where to assign rolling stock, or which weekly reward to accept.

I initially mapped these operations into a single flat discrete action space containing 4,087 action slots:

  • NoOp (1 action)
  • AddLine across pairwise station combinations ($\binom{30}{2} = 435$)
  • ExtendLine from line endpoints (7 lines $\times$ 2 ends $\times$ 30 stations = 420)
  • InsertStation into existing track segments (7 lines $\times$ 15 segments $\times$ 30 stations = 3,150)
  • AddTrain and AddCarriage (7 lines each = 14)
  • UpgradeInterchange (30 stations)
  • ChooseReward (2 weekly card options)
  • CloseLoop and OpenLoop (7 lines each = 14)
  • RemoveLine and ShortenLine (21 actions)

The Inaction Freeze Trap

In any given state, the vast majority of these 4,087 actions are physically illegal. You cannot extend a line that doesn’t exist, build a bridge without tunnels, or add a locomotive when your inventory is zero.

The Go engine computes a boolean mask of legal actions on every step, and Python masks out illegal moves before the softmax:

# Mask illegal actions with a safe negative value before softmax
logits = torch.where(mask, logits, -1e4)

(Note: In 16-bit half precision, using -1e9 for masking causes immediate underflow to -inf, producing NaN gradients during backprop. Using -1e4 provides safe masking without numerical collapse).

Even with action masking, flat categorical softmax produced a bizarre failure mode: the agent completely froze during evaluation.

Here is why: in a typical mid-game state, there might be 1 valid NoOp action and 400 valid track-building combinations. If the policy places 80% of its total probability mass on building infrastructure, that 80% gets spread thinly across 400 options—giving just 0.002 per track action ($0.80 / 400$).

Meanwhile, the single NoOp action gets the remaining 20% ($P(\text{NoOp}) = 0.20$).

During training with stochastic sampling, track actions get selected frequently because their cumulative mass is 80%. But during deterministic evaluation (greedy argmax decoding to pick the single highest probability action): NoOp (at 0.20) beats every individual track action (at 0.002) every single time! The agent would sit completely motionless, watching stations overcrowd and collapse without laying a single piece of track, despite having placed four-fifths of its probability mass on taking action.

The Solution: 12-Way Hierarchical Factorization

To eliminate this dilution, I refactored the policy into a two-stage decision hierarchy:

$$P(a \mid s) = P(\text{action type} \mid s) \cdot P(\text{parameters} \mid \text{action type}, s)$$
  1. Stage 1 (12-Way Action Type Head): The network first predicts high-level intent across the 12 primary operation categories (NoOp, AddLine, ExtendLine, InsertStation, AddTrain, etc.). Here, AddLine competes directly against NoOp on equal terms.
  2. Stage 2 (Conditioned Parameter Heads): Once an action type is chosen, specialized sub-heads evaluate legal arguments:
    • AddLine / ExtendLine: Evaluated by a Spatial Cross-Attention Scorer (SpatialCrossAttentionScorer), which combines bilinear station embeddings with a geometric distance bias [distance, dx, dy, distance^2]. I initialized the distance bias weights to naturally penalize long Euclidean distances, teaching the network from the start to favor compact local connections over cross-map sprawl.
    • AddTrain / AddCarriage: Evaluated by candidate-conditioned dispatch MLPs that score each line’s passenger backlog and train density.

Decoupling high-level intent from combinatorial parameter selection immediately eliminated the inaction freeze.

7. The Forensic Debugging Journey: Bugs That Broke Training

When a deep reinforcement learning agent fails to learn, it almost never throws an exception or prints a clean stack trace. Instead, it quietly exploits physical loopholes, unhandled edge cases, and subtle state discrepancies in your simulation.

A massive portion of this project wasn’t tuning hyperparameter learning rates; it was tracking down simulation bugs that masqueraded as learning failures.

+-----------------------------------------------------------------------------+
|                            FIVE CRITICAL BUGS FOUND                         |
+-----------------------------------------------------------------------------+
| 1. INFINITE LINE DELETION LOOP                                              |
|    Symptom: Policy achieved 11 passengers; kept adding and deleting track.   |
|    Root Cause: Random uninitialized logits on RemoveLine dominated NoOp.    |
|    Fix: Negative initialization, deletion cooldowns, disruption penalties.  |
+-----------------------------------------------------------------------------+
| 2. KINEMATIC TELEPORTATION ON REVERSAL                                      |
|    Symptom: Trains jumped across the map whenever a line was reversed.      |
|    Root Cause: Segment index flipped, but progress p in [0, 1] was un-flipped.|
|    Fix: Updated progress to (1.0 - p) on line inversion.                    |
+-----------------------------------------------------------------------------+
| 3. PASSENGERS TRAPPED ON LINE SHORTENING                                    |
|    Symptom: Train capacity permanently degraded over time.                  |
|    Root Cause: Passengers bound for removed stations stayed on board forever.|
|    Fix: Disembark and re-queue passengers at nearest surviving platform.   |
+-----------------------------------------------------------------------------+
| 4. REWARD TIME-SCALE SKEW                                                   |
|    Symptom: Extreme value loss variance across macro-steps.                 |
|    Root Cause: 1-tick macro-steps penalized equally to 150-tick steps.      |
|    Fix: Scaled all continuous penalties by elapsed duration Delta t.        |
+-----------------------------------------------------------------------------+
| 5. MULTI-TUNNEL SLICE BOUNDS & MASK PANICS                                  |
|    Symptom: Engine crashed with slice index out of bounds panics at 4087.   |
|    Root Cause: Loop iterated past MaxLines; mask checked 1 tunnel for split. |
|    Fix: Added lID < MaxLines guard and exact geometric net-tunnel checks.    |
+-----------------------------------------------------------------------------+

1. The Infinite Add/Delete Churn Loop

When I first enabled actions to delete or shorten lines so the agent could reorganize routes, evaluation scores collapsed to single digits.

Inspecting rollout logs revealed an infinite deletion loop:

  • Step 40: AddLine(A -> B)
  • Step 41: RemoveLine(Line 0)
  • Step 42: AddLine(A -> B)
  • Step 43: RemoveLine(Line 0)

Because the line removal head was newly added, its weights were uninitialized and its logits were essentially random noise. If a removal action happened to have a logit slightly higher than NoOp, the agent immediately demolished the track it had laid one step prior, stranding passengers repeatedly.

I resolved this by adding an operational cooldown (lines could not be deleted within a fixed time window of creation), initializing the removal head with a negative bias, and applying a passenger-disruption penalty whenever active lines with waiting passengers were torn down.

2. Kinematic Spatial Teleportation on Line Reversal

In Mini Metro, a non-loop line operates as a bidirectional shuttle: trains travel from station $s_0$ to $s_{k-1}$, dwell, and reverse back toward $s_0$.

To support line reorganization, the simulator included a ReverseLine function. But during testing, I noticed trains occasionally appeared to jump across the canvas.

In the kinematic engine, a train’s location is parameterized by its segment index and its fractional progress $p \in [0.0, 1.0]$. When reversing a line of $n$ stations, my code updated the segment index (new_segment = (n - 1) - old_segment) and inverted the train’s direction flag. But I forgot to invert the fractional progress coordinate! A train that was 80% of the way along a segment ($p = 0.8$) suddenly became 80% along the reversed segment, instantly teleporting across the track.

I corrected the inversion to p = 1.0 - p and wrote an automated unit test (simulator/cmd/physics_test.go) validating spatial continuity within $10^{-4}$ tolerance.

3. Stranded Passengers on Line Shortening

When the agent shortened a line by removing a terminus station, any passengers currently on board trains whose destination was that removed station remained on the train indefinitely.

Because their destination stop no longer existed on the route, they never alighted. They rode back and forth forever, permanently consuming carriage capacity and silently reducing the train’s throughput to zero.

I updated shortenLine() in simulator/engine/simulator.go to iterate over all active trains on the line, disembark any passengers whose destination was cut off, and re-queue them at the nearest surviving platform with an updated $A^*$ routing request.

4. Macro-Step Reward Time-Scale Distortion

To speed up training, the environment runs macro-steps: the Go engine steps physics forward for up to 150 sub-ticks (5.0 seconds) between AI decisions, but interrupts early if an asynchronous event occurs (such as a station spawning or an overcrowding warning triggering).

In my early reward formulation, continuous crowding penalties were added as fixed per-step quantities. This introduced massive variance: a macro-step that lasted only 1 tick (0.033 seconds) received the exact same crowding penalty as a macro-step that ran for the full 150 ticks (5.0 seconds)!

All continuous rewards and penalties had to be formulated as rates and scaled by the actual elapsed duration ($\Delta t = \text{StepTicks} \times dt$), calculating $\text{Reward}_{\text{step}} = R_{\text{rate}} \cdot \Delta t$.

5. Multi-Tunnel Slice Bounds & Mask Panics

When lines multiplied, an off-by-one error in action_space.go caused the Go engine to crash with a slice index out of bounds panic. The loop checking for line modifications iterated for lID := 0; lID < len(s.State.Lines); lID++ without checking lID < MaxLines (7), computing an index offset equal to 4087 in a 4087-length slice.

At the same time, when the agent attempted to insert a station into an existing segment that crossed water, splitting that track in two actually required two tunnels instead of one. My action mask only checked if the player could afford one tunnel (CanSpend(RewardTunnel)), meaning the agent picked moves that were marked legal by the mask but immediately failed during execution. Adding explicit bounds guards (lID < MaxLines) and accurate geometric multi-tunnel requirements (s.State.Resources.Tunnels >= netTunnels) fixed the crashes and kept the action mask synchronized with the simulator.

8. Guiding the Agent: Reward Shaping and Curriculum Learning

With the environment hardened and bugs fixed, I turned to the reinforcement learning pipeline.

Potential-Based Reward Shaping (PBRS)

In Mini Metro, passenger deliveries are sparse, especially early in training before an agent has learned to link complementary shapes. If you only reward deliveries, the agent spends millions of steps receiving zero feedback.

However, adding ad-hoc heuristic bonuses (like rewarding total track laid) frequently causes reward exploitation. (In one early run, the agent learned to maximize its reward by repeatedly building redundant loop tracks between two adjacent circle stations, completely ignoring expanding demand elsewhere!).

To provide dense guidance without altering the optimal policy, I implemented Potential-Based Reward Shaping (PBRS) (Ng et al., 1999). Augmenting rewards with a potential difference:

$$R'(s, a, s') = R(s, a, s') + \gamma \Phi(s') - \Phi(s)$$

This provides continuous guidance during early exploration without altering the true optimal policy.

In ml/env.py, the state potential $\Phi(s)$ penalizes four systemic stress metrics:

# Dense potential: guide the agent with continuous penalties
potential = -(
    0.20 * queue_stress +      # quadratic penalty as stations approach capacity
    1.00 * overcrowd_stress +  # heavy penalty on active countdown timers
    0.50 * unconnected_count + # penalty for newly spawned stations left isolated
    0.30 * headway_stress      # penalty for stretching lines beyond 4 segments
)

4-Stage Progressive Expanding Pool Curriculum

Training an untrained model directly on complex maps like New York City (which features narrow Manhattan river bottlenecks and severe tunnel scarcity) causes immediate collapse: the agent burns its tunnels and overflows before learning basic track connectivity.

Conversely, sequentially training on one map and then switching to another (Berlin $\to$ London $\to$ Tokyo) causes catastrophic forgetting: the network overwrites its open-grid routing heuristics to adapt to river crossings.

To solve this, I built a Progressive Expanding Pool Curriculum (ml/curriculum.py):

Stage 1: Berlin Fundamentals (Map ID 3)
└── 100% Berlin (Open geographic grid, no early water chokepoints)
      │
      ▼ (Promoted at rolling score >= 100.0)
Stage 2: River Crossing (Maps 3, 0)
└── 50% Berlin + 50% London (Thames river crossings, tunnel budgeting)
      │
      ▼ (Promoted at rolling score >= 150.0)
Stage 3: Coastal Islands (Maps 3, 0, 2)
└── 33% Berlin + 33% London + 34% Tokyo (Dense passenger surges, island clusters)
      │
      ▼ (Promoted at rolling score >= 200.0)
Stage 4: All-Map Mastery (Maps 0, 1, 2, 3)
└── 25% London + 25% NYC + 25% Tokyo + 25% Berlin (Full multi-map surge mastery)

By expanding the active pool rather than replacing maps, the agent retains basic routing strategies while progressively learning to budget tunnels and manage island bottlenecks.

9. Results and What Actually Happened

Now for the honest engineering truth: what happened when the system was evaluated?

The Reality of Tabula-Rasa PPO: Conservative Idling Collapse

If you take the trained neural network checkpoint (ml/runs/minimetro_ppo/model_best.pt) and evaluate it using pure deterministic argmax decoding without any heuristic fallback, it scores 0.0 passengers over 50 steps on seed 1000.

Why? Because reinforcement learning in a 4,087-dimensional discrete action space with severe delayed credit assignment is brutally difficult. During early exploration, taking random structural actions—like connecting distant stations or burning tunnels—frequently led to rapid catastrophic failure. The PPO optimizer discovered a local optimum: selecting NoOp (Action 0) preserves system survival longer than taking erratic random actions.

The network fell into conservative idling collapse. Under stochastic sampling during training, it expanded lines and delivered passengers (historically reaching ~100–140 passengers on London and NYC in commit 75195ad). But under strict deterministic greedy evaluation, its highest individual logit was almost always NoOp.

The Pragmatic Engineering Solution: The Hybrid Live Agent

Rather than abandoning the trained network or relying purely on hand-written heuristics, I took a hybrid engineering approach in the live production agent (ml/agent.py and ModelLiveAgentPolicy in ml/eval.py).

The hybrid agent queries the neural actor-critic model on every macro-step, but wraps it in four deterministic safety invariants:

  1. Train Reservation Protection (Invariant P5-1): If the model attempts to deploy an extra locomotive (AddTrain) to an existing route, but doing so would leave fewer spare trains than unbuilt lines (unused_trains <= unused_lines), the action is suppressed. You must never strand an unbuilt line token without a train to run it.
  2. Proactive Line Extension: If the neural model outputs NoOp while new stations remain completely unconnected to the transit network, the agent overrides NoOp and connects the isolated station to the nearest compatible line.
  3. Emergency Interchange Upgrades: If a platform exceeds capacity and enters a critical countdown, an interchange token is deployed to triple holding capacity from 6 to 18.
  4. Headway-Balanced Rolling Stock: Locomotives and carriages are dispatched based on real-time passenger density and corridor length.

When evaluated on the standardized London benchmark (Map 0, Seed 1000, 50 Steps), here is how the policies compare:

Policy ModePassengers DeliveredSurvival StepsActive LinesTrains DeployedTunnels Used
Model (Deterministic Raw)0.050.00.00.00.0
Greedy Heuristic Baseline55.050.01.04.03.0
Model (Live Agent Hybrid)70.050.04.04.03.0

Historical multi-seed benchmarks across longer multi-hundred step rollouts:

  • Greedy / Grandmaster Heuristic: London: 220.2 ± 61.2, NYC: 200.0 ± 85.5, Tokyo: 205.2 ± 21.7.
  • Model Live Agent Hybrid: Consistently reaches 200–220+ passengers, matching or exceeding human-tier survival times while maintaining cleaner multi-line topologies than greedy heuristics.
Delivered Passengers over 50 Steps (London, Seed 1000)
Passengers
 80 ┼─────────────────────────────────────────────────────────────╭── Hybrid Agent (70)
 70 ┼───────────────────────────────────────────────────────╭─────╯
 60 ┼─────────────────────────────────────────────────╭─────╯
 50 ┼────────────────────────────────────╭────────────╯────────────── Greedy Heuristic (55)
 40 ┼──────────────────────────────╭─────╯
 30 ┼────────────────────────╭─────╯
 20 ┼─────────────────╭──────╯
 10 ┼──────────╭──────╯
  0 ┼──────────┴───────────────────────────────────────────────────── Raw Deterministic Model (0)
    └──────────┬───────────┬───────────┬───────────┬───────────┬──────►
             Step 10     Step 20     Step 30     Step 40     Step 50

Emergent Strategic Behaviors

When running the live agent, several genuine transit-planning strategies emerged from the combination of relational attention and safety invariants:

  1. Geometric Shape Alternation: The network naturally learned to avoid connecting stations of identical shapes (e.g., Circle $\to$ Circle). Tracks naturally alternate between Circles, Triangles, and Squares, ensuring that trains continuously pick up and drop off passengers at alternating stops.
  2. Circular Loop Closures: On open maps like Berlin and London, the agent actively utilizes CloseLoop to turn open lines into continuous rings, eliminating terminal turnaround dwell delays.
  3. Interchange Shock-Absorption: The agent reserves scarce interchange upgrades for central transfer nexuses serving 3 or more intersecting lines, creating natural passenger hubs.

Probing: Did the Model Actually Learn Demand Scarcity?

In standard Mini Metro, Squares are rare bottleneck destinations, while Circles are abundant origins. A neural network could score reasonably well simply by memorizing a static correlation: always connect tracks to Squares.

To test whether the model was genuinely reasoning about demand or just following a fixed rule of thumb, I set up a counterfactual experiment in ml/probing.py:

  • Standard Map: 5 Circles, 1 Triangle, 1 Square (Square is the rare bottleneck).
  • Inverted Map: 5 Squares, 1 Triangle, 1 Circle (Circle is now the rare bottleneck).

In my earliest models with static one-hot station features, the policy completely failed to adapt: it continued blindly routing tracks toward redundant Squares because the training data had strongly correlated “Square” with high reward. Only after incorporating explicit dynamic features—passenger queue arrival rates and approaching train capacity—did the attention layers begin shifting focus based on actual waiting passenger volume and congestion urgency, rather than relying on static shape memorization.

10. What I Learned and What’s Next

Taking a project like this from an empty directory to a functioning reinforcement learning system was a humbling engineering experience. Here are my biggest takeaways:

  1. Simulation Throughput Dictates Everything: If your simulator runs at 50 FPS, you cannot do deep reinforcement learning. Writing the engine in Go and bridging it directly to Python via zero-copy CGO shared memory was the single best architectural decision in this project.
  2. Flat Action Spaces Collapse Under Combinatorics: Attempting to sample from thousands of flat action slots with standard categorical softmax is a recipe for failure. Hierarchical intent factorization and domain-aware masking are mandatory.
  3. Pure RL from Scratch is the Hard Way: Training an agent entirely tabula-rasa in a complex combinatorial environment with delayed credit assignment wastes immense compute. The natural next step—which is already coded in the repository in ml/train_bc.py—is Behavioral Cloning (BC): pre-training the neural network on demonstration datasets from the heuristic controller to give it strong warm-start weights before running PPO.
  4. State Cloning Enables Tactical Lookahead: In simulator/engine/simulator.go, I implemented an in-memory deep copy method:
    func (s *Simulator) Clone() *Simulator
    Because Go memory is fast and pointers can be duplicated cleanly, ml/mcts.py can clone the entire simulation in a microsecond and simulate forward rollouts 4 to 8 steps into the future during crisis states. Combining tree search with neural evaluation is the most promising path toward grandmaster-level play.
  5. Refining the Frontend Visualizer: The current visualizer was built to prioritize quick debugging over UI polish. Next steps include fixing track overlap rendering on shared corridors, smoothing out mouse drag-and-drop controls, and adding an overlay that displays the neural network’s GATv2 attention weights directly across the canvas.

The full source code, test suites, and models are available under the MIT License on GitHub.