GateNet — Vision-Based Gate Segmentation for Drone Racing
A compact U-Net that replaces privileged simulator segmentation with RGB-only gate prediction, closing the sim-to-real perception gap for DreamerV3 drone racing. Achieves 0.93 mean IoU and 1.0 median IoU on held-out validation.
View on GitHub ↗Motivation
The DreamerV3 world model trained on the drone-racer environment originally consumed the simulator’s privileged semantic_segmentation channel — a per-pixel class label that is impossible to obtain on a real drone with only an RGB camera.
GateNet closes the sim-to-real gap on the perception side. It is a compact U-Net that takes the same 64×64 RGB observation the policy already sees and produces a single-channel binary gate mask (SkyDreamer convention, see arXiv:2510.14783 Appendix A). The mask replaces the privileged segmentation in the DreamerV3 training loop, making the full pipeline hardware-transferable.
Architecture
Five-level U-Net scaled for 64×64 images, with deep supervision at every scale:
inc : in_channels → 64
down1 : 64 → 128 (MaxPool 2×2)
down2 : 128 → 256
down3 : 256 → 512
down4 : 512 → 512 (bottleneck, 4×4)
up4 : skip + ConvTranspose → 256
up3 : → 128
up2 : → 64
up1 : → 64
out_i (×5): 1×1 Conv → 1-channel logit at each scale
Training loss with deep-supervision weighting:
total = 4·L0 + 2·L1 + L2 + L3 + L4
L_i = Dice + 2·BCE
Optional extensions: with_pose=True adds smooth-L1 pose-regression heads (body-frame position + quaternion); multi_gate=True emits one mask per gate on the track; frame_stack=K stacks K consecutive RGB frames for temporal parallax.
Results
Evaluated on a 100k-sample dataset (10% held-out, seed 42):
| Metric | Value |
|---|---|
| Mean IoU | 0.9255 |
| Median IoU | 1.0000 |
| IoU (gate visible) | 0.8125 |
| IoU (frame empty) | 0.9947 |
| Precision | 0.9567 |
| Recall | 0.8788 |
| Pixel accuracy | 0.9969 |
| Position error (visible, median) | 1.62 m |
| Quaternion error (visible, median) | 18.5° |
| Visibility classifier accuracy | 0.8934 |
Median IoU of 1.0 means most frames with a fully-visible or fully-absent gate are predicted exactly. The mean-IoU drag comes from edge cases: gates entering/leaving the frame, heavy yaw rotations compressing the gate to a few pixels, and motion blur during fast turns.
Integration with DreamerV3
Passing --gatenet_ckpt <path> to the training script causes the env wrapper to swap the oracle segmentation for GateNet predictions transparently:
# dreamer/env_wrapper.py
if gatenet_ckpt is not None:
mask_u8 = _gatenet_predict_mask_u8(self._gatenet, rgb_u8)
else:
mask_u8 = _extract_gate_mask_u8(seg, ...) # oracle fallback
The downstream world model sees an identical (image, state) dict either way — no changes required to the policy or training loop.
Key Files
| File | Purpose |
|---|---|
dreamer/gatenet.py | U-Net with multi-scale supervision and optional pose/multi-gate heads |
scripts/data/collect_gatenet_data.py | Random-action data collector with pose labels and frame-stack support |
scripts/train_gatenet.py | Supervised training with Dice + BCE deep supervision |
scripts/eval_gatenet.py | IoU / pose metrics, worst-case PNG grids |
scripts/render_gatenet_video.py | `[RGB |
dreamer/env_wrapper.py | Drop-in GateNet integration for DreamerV3 training |