A U-Net model that predicts directional viability maps for non-holonomic robot navigation: given a 2D occupancy grid and robot dimensions, the model outputs four binary masks (N/S/E/W) indicating which pixels are topological traps - cells the robot can enter but cannot escape by translating in that heading without first rotating.

Models in this repo

File	Description	Val IoU
`models/viability_cardinal_best_iou.pth`	Cardinal model - 3-ch input, 4-ch binary output (N/S/E/W)	0.9793
`models/viability_continuous_angle_best_iou.pth`	Angle model - 5-ch input (sin/cos theta), 1-ch arbitrary heading	0.984
`configs/training_config.yaml`	Training configuration	-

Architecture

Backbone: U-Net with ResNet34 encoder (ImageNet pre-trained), ~24.5M parameters
Input: 3-channel tensor - binary occupancy grid + normalised robot L/512, W/512
Output: 4-channel sigmoid map, one per cardinal direction
Training: AdamW (lr=1e-4, wd=1e-4), CosineAnnealingLR, BCE+Dice loss, FP16, batch 16, 50 epochs

Results

Metric	Value
IoU (seen sizes)	0.978
IoU (unseen size 25x18)	0.953
Generalisation gap	0.025
Inference time	8.8 ms/map (GPU)
Oracle speedup	21.5x
PRM trap reduction	81%

Quick start

python

import torch, numpy as np
from huggingface_hub import hf_hub_download
import segmentation_models_pytorch as smp

ckpt_path = hf_hub_download(
    repo_id="DanielDDDs/topological-traps",
    filename="models/viability_cardinal_best_iou.pth"
)
model = smp.Unet(encoder_name="resnet34", encoder_weights=None, in_channels=3, classes=4)
state = torch.load(ckpt_path, map_location="cpu")
model.load_state_dict(state["model_state_dict"])
model.eval()

def predict(occupancy, robot_L=30, robot_W=20):
    x = np.zeros((1, 3, 512, 512), dtype=np.float32)
    x[0, 0] = occupancy
    x[0, 1] = robot_L / 512.0
    x[0, 2] = robot_W / 512.0
    with torch.no_grad():
        out = torch.sigmoid(model(torch.from_numpy(x)))
    return out[0].numpy()  # (4, 512, 512) N/S/E/W viability

viability = predict(occupancy_grid)
north_viable = viability[0] > 0.5

Dataset

Full processed maps and Oracle labels: DanielDDDs/topological-traps-dataset

Code

github.com/danielsddd/topological-traps-project

Compute

TAU CS SLURM cluster:

NVIDIA TITAN Xp (12 GB)
NVIDIA GeForce RTX 2080 (8 GB)

If you use this model in your research, please cite it as follows:

bibtex

@misc{simanovsky2026neural,
  author = {Daniel Simanovsky},
  title  = {Neural Prediction of Heading-Dependent Topological Traps},
  year   = {2026},
  school = {Tel Aviv University},
}

DanielDDDS/topological-traps

Neural Prediction of Heading-Dependent Topological Traps