draft integration of sdtagnet sd map encoder

Author: immel-f

.gitignore

@@ -2,3 +2,12 @@ __pycache__/
 *.py[cod]
 ._*
 benchmark_results.json
+# Downloaded SDTagNet NLP encoder checkpoint (default download dir)
+checkpoints/
+# Generated OSM SD-map artifacts: caches, raw Overpass fetches, tokenized
+# shards, built road graphs, and visualization output.
+cache/
+*.osm.gz
+osm_bbox_*.xml.gz
+graph_*.pkl
+osm_map_vis.png

Model/README.md

@@ -34,4 +34,20 @@ trajectory, ego_hidden, future_visual_features = model(
 **To learn the driving policy:**
 - Imitation Learning is used to penalize trajectory prediction as well as World Model Simulation based Reinforcement Learning
 
+## Map priors
+
+Two interchangeable map-branch flavours feed road-network priors into the BEV:
+
+- **`map_type="rasterized"`** (default): a rendered nav-map image is encoded to a
+  spatial map BEV and fused with the camera BEV (`map_fusion_mode="residual"` or
+  `"cross_attn"`). See [data_parsing/map_rendering](./data_parsing/map_rendering).
+- **`map_type="osm_vector"`**: OpenStreetMap SD-map elements are encoded as
+  **text-annotated vector tokens** (ported from
+  [SDTagNet](https://arxiv.org/abs/2506.08997)) and the camera BEV
+  cross-attends to them (`map_fusion_mode="osm_cross_attn"`). This keeps arbitrary
+  OSM semantics (lane counts, surface, turn restrictions, ...) without a fixed
+  class taxonomy. See [data_parsing/osm_sd_map](./data_parsing/osm_sd_map) for the
+  offline pipeline and [model_components/map_encoder/osm_vector](./model_components/map_encoder/osm_vector)
+  for the encoder. Both branches share one on-disk OSM cache.
+
 

Model/data_parsing/l2d/dataset.py

@@ -68,10 +68,17 @@ def __init__(
         episodes: list[int] | None = None,
         backbone_name: str = "swinv2_tiny_window8_256",
         local_files_only: bool = False,
+        osm_cache_dir: str | None = None,
     ) -> None:
         from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
 
         self.repo_id = repo_id
+        # When set, per-episode tokenised OSM shards (built offline via
+        # data_parsing.osm_sd_map.cache.build_episode_osm_cache) are merged into
+        # each sample as osm_map_* keys for the OSM vector encoder. When None,
+        # samples are unchanged (existing behaviour).
+        self.osm_cache_dir = osm_cache_dir
+        self._osm_shard_cache: tuple[int | None, dict | None] = (None, None)
 
         self.lerobot_dataset = LeRobotDataset(
             repo_id=repo_id,
@@ -162,11 +169,49 @@ def __getitem__(self, idx: int) -> L2DSample:
 
         visual_history = torch.zeros(_VISUAL_HISTORY_DIM, dtype=torch.float32)
 
-        return L2DSample(
+        sample = L2DSample(
             visual_tiles=visual_tiles,
             egomotion_history=egomotion_history,
             visual_history=visual_history,
             trajectory_target=trajectory_target,
             episode_index=ep_idx,
             frame_index=local_frame_idx,
         )
+
+        if self.osm_cache_dir is not None:
+            osm_sample = self._load_osm_sample(ep_idx, local_frame_idx)
+            if osm_sample is not None:
+                # Merge osm_map_* keys; collate_osm_batch groups them per batch.
+                sample.update(osm_sample)
+
+        return sample
+
+    def _load_osm_sample(self, ep_idx: int, local_frame_idx: int):
+        """Load one frame's tokenised OSM data from the per-episode shard cache."""
+        from data_parsing.osm_sd_map.cache import load_episode_osm_cache
+
+        cached_ep, shard = self._osm_shard_cache
+        if cached_ep != ep_idx:
+            shard = load_episode_osm_cache(self.osm_cache_dir, ep_idx)
+            self._osm_shard_cache = (ep_idx, shard)
+        if shard is None:
+            return None
+        return shard.get(local_frame_idx)
+
+    def episode_ego_poses(self, ep_idx: int) -> list[tuple[int, float, float, float]]:
+        """Per-valid-frame ``(local_frame_idx, lat, lon, heading)`` for OSM caching.
+
+        Feed the result to ``osm_sd_map.cache.build_episode_osm_cache``. Uses
+        the vehicle state columns ``[1]=heading, [3]=lat, [4]=lon``.
+        """
+        episode_data_index = self.lerobot_dataset.episode_data_index
+        ep_start = episode_data_index["from"][ep_idx].item()
+
+        poses = []
+        for sample_ep_idx, global_frame_idx in self._samples:
+            if sample_ep_idx != ep_idx:
+                continue
+            local_frame_idx = global_frame_idx - ep_start
+            state = self.lerobot_dataset[global_frame_idx]["observation.state.vehicle"].numpy()
+            poses.append((local_frame_idx, float(state[3]), float(state[4]), float(state[1])))
+        return poses

Model/data_parsing/map_rendering/cache.py

@@ -63,8 +63,9 @@ def render_and_cache_tiles(
         output_dir: where rendered PNG tiles are written (`{clip_id}.png`).
         radius_m: render radius around each clip's centroid.
         image_size: output `(W, H)`.
-        network_cache_dir: if given, fetched graphs are persisted here keyed by
-            centroid so neighbouring clips share a cached download.
+        network_cache_dir: if given, fetched graphs and the shared raw-OSM XML
+            are persisted here (keyed by the clip's trajectory bbox) so
+            neighbouring clips and the OSM vector branch reuse downloads.
         skip_existing: do not re-render clips whose PNG already exists.
 
     Returns:
@@ -91,8 +92,11 @@ def render_and_cache_tiles(
         ego_lon = float(lons[-1])
         ego_heading = _heading_from_trace(lats, lons, ego_lat)
 
+        # Fetch sized to the whole trajectory + the render radius, so the route
+        # and the ±radius draw window are always covered (a centroid radius can
+        # miss long clips).
         graph = _load_or_fetch_network(
-            ego_lat, ego_lon, radius_m, net_cache
+            list(lats), list(lons), radius_m, net_cache
         )
         if graph is None:
             logger.warning("clip %s: failed to obtain road network; skipping", clip_id)
@@ -146,35 +150,108 @@ def _heading_from_trace(
     return math.atan2(dx, dy)
 
 
+def _graph_from_shared_osm(
+    lats: Sequence[float],
+    lons: Sequence[float],
+    margin_m: int,
+    cache_dir: Path | None,
+) -> nx.MultiDiGraph | None:
+    """Build the road graph from the shared raw-OSM cache (trace bbox + margin).
+
+    The SD-map (OSM vector) branch and this rendered-tile branch share a single
+    on-disk artifact: the raw Overpass OSM XML (see
+    ``data_parsing.osm_sd_map.overpass``). Both fetch the bounding box of the
+    clip's trajectory expanded by their own margin (here the render radius); the
+    shared cache reuses any already-fetched XML that covers the request, so one
+    download serves both. Returns ``None`` (caller falls back to a direct fetch)
+    if the shared path is unavailable for any reason.
+    """
+    try:
+        import os
+        import tempfile
+
+        import osmnx as ox
+
+        from ..osm_sd_map.overpass import load_or_fetch_osm_for_trace
+    except Exception as exc:  # noqa: BLE001 — optional path; never break rendering
+        logger.debug("shared OSM path unavailable (%s); using direct fetch", exc)
+        return None
+
+    xml = load_or_fetch_osm_for_trace(lats, lons, margin_m, cache_dir)
+    if xml is None:
+        return None
+
+    tmp = tempfile.NamedTemporaryFile("w", suffix=".osm", delete=False, encoding="utf-8")
+    try:
+        tmp.write(xml)
+        tmp.close()
+        return ox.graph_from_xml(tmp.name, bidirectional=False, simplify=True, retain_all=True)
+    except Exception as exc:  # noqa: BLE001 — osmnx version/parse differences
+        logger.warning("graph_from_xml failed (%s); using direct fetch", exc)
+        return None
+    finally:
+        os.unlink(tmp.name)
+
+
+def _trace_centroid_radius(
+    lats: Sequence[float], lons: Sequence[float], margin_m: int
+) -> tuple[float, float, int]:
+    """Centroid + a radius that covers the whole trace plus ``margin_m``."""
+    clat = sum(lats) / len(lats)
+    clon = sum(lons) / len(lons)
+    deg_to_m = EARTH_RADIUS_M * math.pi / 180.0
+    cos_lat = math.cos(math.radians(clat))
+    max_d = 0.0
+    for la, lo in zip(lats, lons):
+        dx = (lo - clon) * cos_lat * deg_to_m
+        dy = (la - clat) * deg_to_m
+        max_d = max(max_d, math.hypot(dx, dy))
+    return clat, clon, int(max_d + margin_m)
+
+
+def _graph_cache_path(
+    lats: Sequence[float], lons: Sequence[float], margin_m: int, cache_dir: Path
+) -> Path | None:
+    """Pickle path for the built graph, keyed by the trace bbox."""
+    try:
+        from ..osm_sd_map.overpass import bbox_from_points
+    except Exception:  # noqa: BLE001
+        return None
+    s, w, n, e = bbox_from_points(lats, lons, margin_m)
+    return cache_dir / f"graph_{s:.4f}_{w:.4f}_{n:.4f}_{e:.4f}.pkl"
+
+
 def _load_or_fetch_network(
-    center_lat: float,
-    center_lon: float,
-    radius_m: int,
+    lats: Sequence[float],
+    lons: Sequence[float],
+    margin_m: int,
     cache_dir: Path | None,
 ) -> nx.MultiDiGraph | None:
-    """Return a cached graph if available, otherwise fetch and cache it.
+    """Return a cached graph if available, otherwise build/fetch and cache it.
 
-    Centroids are quantized to ~100 m so nearby clips reuse the same download.
+    The fetch is sized to the trajectory bounding box plus ``margin_m`` (the
+    render radius), so long clips are fully covered (a fixed centroid radius can
+    miss them). The graph is built from the shared raw-OSM cache when possible
+    (reused by the OSM vector branch via bbox containment); otherwise it falls
+    back to a direct osmnx fetch sized to cover the trace. The built graph is
+    memoized as a pickle keyed by the trace bbox.
     """
-    if cache_dir is not None:
-        key = f"{round(center_lat, 3)}_{round(center_lon, 3)}_{radius_m}.pkl"
-        cache_path = cache_dir / key
+    cache_path = _graph_cache_path(lats, lons, margin_m, cache_dir) if cache_dir is not None else None
+    if cache_path is not None:
         cached = load_cached_network(cache_path)
         if cached is not None:
             return cached
-    else:
-        cache_path = None
 
-    try:
-        graph = fetch_road_network(center_lat, center_lon, radius_m=radius_m)
-    except Exception as exc:  # noqa: BLE001 — network/Overpass failures
-        logger.warning(
-            "fetch_road_network(%.4f, %.4f) failed: %s",
-            center_lat,
-            center_lon,
-            exc,
-        )
-        return None
+    graph = _graph_from_shared_osm(lats, lons, margin_m, cache_dir)
+    if graph is None:
+        clat, clon, radius = _trace_centroid_radius(lats, lons, margin_m)
+        try:
+            graph = fetch_road_network(clat, clon, radius_m=radius)
+        except Exception as exc:  # noqa: BLE001 — network/Overpass failures
+            logger.warning(
+                "fetch_road_network(%.4f, %.4f) failed: %s", clat, clon, exc
+            )
+            return None
 
     if cache_path is not None:
         cache_network(graph, cache_path)