mapgs/config.py

"""Typed configuration for MapGS.

Defaults below mirror the values in the design doc (§2-§3). Override via YAML
(``configs/*.yaml``, with optional ``_base_`` inheritance) or dotted CLI args
(e.g. ``train.iters=1000 model.embed_dim=512``).
"""

from __future__ import annotations

import dataclasses
from dataclasses import dataclass, field
from typing import Any, get_type_hints
import os

import yaml


# --------------------------------------------------------------------------- #
# Map / anchors (§2.1, §2.2)
# --------------------------------------------------------------------------- #
@dataclass
class MapConfig:
    crop_lateral: float = 100.0       # m, lateral extent of the local map crop
    crop_longitudinal: float = 150.0  # m, longitudinal extent
    ground_spacing: float = 0.5       # m, ground anchor grid spacing
    lane_spacing: float = 0.2         # m, lane polyline anchor spacing
    boundary_spacing: float = 0.5     # m, road-edge anchor spacing
    ratio_ground: float = 0.6         # ~60% of anchors on drivable area
    ratio_lane: float = 0.3           # ~30% on lanes
    ratio_boundary: float = 0.1       # ~10% on boundaries
    n_anchors: int = 2048             # N_a == number of map-anchored tokens
    junction_weight: float = 3.0      # FPS density multiplier at junctions
    curvature_weight: float = 1.0     # FPS density multiplier with curvature
    ground_z_below: float = 0.15      # m, free-space delta below ground (§2.6 second term)


# --------------------------------------------------------------------------- #
# Tokens (§2.3)
# --------------------------------------------------------------------------- #
@dataclass
class TokenConfig:
    n_map: int = 2048                 # T^M map-anchored tokens
    n_free: int = 2048                # T^F free static tokens
    n_dyn_per_instance: int = 256     # T^D tokens per dynamic instance
    gaussians_per_token: int = 64     # N_G gaussians decoded per token
    s0: float = 0.5                   # m, initial residual bound for anchored means
    s_max: float = 2.0                # m, max residual bound (tempered up)


# --------------------------------------------------------------------------- #
# Model (§2.3, §2.4, §2.5)
# --------------------------------------------------------------------------- #
@dataclass
class ModelConfig:
    embed_dim: int = 1024             # C
    patch_size: int = 8               # p
    enc_depth: int = 6                # N_enc (base); finetune uses 3
    dec_depth: int = 24               # N_dec (base); finetune uses 12
    n_heads: int = 16
    mlp_ratio: float = 4.0
    qk_norm: bool = True              # QK-norm (TokenGS)
    layerscale_init: float = 1e-5     # LayerScale init (TokenGS)
    shared_decoder_kv: bool = True    # share K/V projections across decoder layers
    drop_final_ln: bool = True        # omit final LayerNorm before the head (TokenGS)

    # Appearance (§2.4)
    feature_color: bool = True        # render N_c feature channels then UNet-decode to RGB
    feature_dim: int = 128            # N_c
    use_unet: bool = True

    # Dynamics (§2.5)
    causal_dynamic_mask: bool = True  # dynamic tokens attend causally to static only
    st_fusion: bool = False           # optional PointForward-style ST-fusion (ablation)

    # Gaussian activation ranges
    scale_min: float = 1e-3
    scale_max: float = 3.0            # m (clipped exp)
    scale_init: float = 0.05          # m, nominal gaussian size at init
    opacity_bias: float = -2.0        # logit bias -> ~0.12 opacity at init

    tokens: TokenConfig = field(default_factory=TokenConfig)


# --------------------------------------------------------------------------- #
# Losses + tempering (§2.6, §2.7)
# --------------------------------------------------------------------------- #
@dataclass
class LossConfig:
    lambda_ssim: float = 0.2
    lambda_lpips: float = 0.2
    lambda_vis: float = 1.0           # free tokens only
    lambda_md: float = 0.5            # map-guided depth
    lambda_fs: float = 0.1            # free-space / sub-surface term (inside L_mapdepth)
    lambda_ext: float = 0.3           # extrapolation
    lambda_warp: float = 0.3          # within L_extrap
    lambda_lane: float = 0.1          # within L_extrap
    lambda_vert: float = 0.1          # vertical-structure weak depth
    lambda_ground_coupling: float = 0.05  # dynamic gaussians pulled to map ground z

    huber_delta: float = 0.5          # m, Huber transition for depth losses
    warp_tau: float = 0.5             # m, z-buffer occlusion threshold for warping
    lateral_shift_range: tuple = (1.0, 3.0)  # m, training deviated-view lateral sampling
    yaw_jitter_deg: float = 3.0
    pitch_jitter_deg: float = 1.5

    # Uncertainty tempering curriculum (§2.7)
    eps0: float = 0.1                 # m
    eps_max: float = 1.0              # m
    temper_gamma: float = 1.0005
    md_late_decay: float = 0.5        # multiply lambda_md by this in the last 20% of training


# --------------------------------------------------------------------------- #
# Data (§2.1, §3.1)
# --------------------------------------------------------------------------- #
@dataclass
class DataConfig:
    name: str = "synthetic"           # synthetic | waymo | nuscenes | argoverse
    root: str = ""                    # dataset root for real adapters
    clip_seconds: float = 2.0
    num_frames: int = 20
    fps: int = 10
    context_times: tuple = (0.0, 0.5, 1.0, 1.5)  # s, input frames
    cameras: tuple = ("front_left", "front", "front_right")
    height: int = 256
    width: int = 448
    depth_scale_target: float = 1.0   # rescale translations so mean depth ~= 1 (TokenGS)
    use_lidar: bool = False
    use_mono_depth: bool = True       # for L_vert and optional free-token init
    use_boxes: bool = True            # dynamic instances
    max_instances: int = 8
    # synthetic-only
    synth_num_scenes: int = 64
    synth_dynamic_actors: int = 3


# --------------------------------------------------------------------------- #
# Training (§3)
# --------------------------------------------------------------------------- #
@dataclass
class TrainConfig:
    stage: str = "base"               # base | finetune
    iters: int = 150000
    lr: float = 4e-4
    warmup: int = 2000
    min_lr_ratio: float = 0.01
    batch_size: int = 16          # literal per-step batch (sized to fill GPU memory)
    grad_accum: int = 1           # accumulate this many micro-batches per optimizer step
    weight_decay: float = 0.05
    grad_clip: float = 1.0
    extrap_ramp_iter: int = 10000     # ramp in L_extrap after this many iters
    deviated_views_per_step: int = 1
    amp: bool = True
    amp_dtype: str = "bf16"
    grad_checkpoint: bool = False     # checkpoint encoder/decoder blocks (full-scale memory)
    num_workers: int = 4
    log_every: int = 50
    ckpt_every: int = 5000
    val_every: int = 5000
    out_dir: str = "runs/mapgs"
    resume: str = ""
    init_translation_scale: float = 0.1  # TokenGS: start mean-depth scale near 0.1


# --------------------------------------------------------------------------- #
# Test-time scaling (§2.8)
# --------------------------------------------------------------------------- #
@dataclass
class TTConfig:
    enabled: bool = False
    steps: int = 50
    lr: float = 1e-4
    densify: bool = True
    densify_perturb_std: float = 0.01


# --------------------------------------------------------------------------- #
# Eval (§4)
# --------------------------------------------------------------------------- #
@dataclass
class EvalConfig:
    lateral_shifts: tuple = (1.0, 2.0, 3.0, 4.0, 6.0)  # m
    held_out_lane_change_frames: int = 50
    fid_enabled: bool = True
    lane_miou: bool = True


@dataclass
class MapGSConfig:
    map: MapConfig = field(default_factory=MapConfig)
    model: ModelConfig = field(default_factory=ModelConfig)
    loss: LossConfig = field(default_factory=LossConfig)
    data: DataConfig = field(default_factory=DataConfig)
    train: TrainConfig = field(default_factory=TrainConfig)
    tt: TTConfig = field(default_factory=TTConfig)
    eval: EvalConfig = field(default_factory=EvalConfig)
    seed: int = 0
    device: str = "cuda"

    # convenience
    @property
    def n_static_tokens(self) -> int:
        return self.model.tokens.n_map + self.model.tokens.n_free

    @property
    def static_gaussian_budget(self) -> int:
        return self.n_static_tokens * self.model.tokens.gaussians_per_token

    def to_dict(self) -> dict:
        return dataclasses.asdict(self)


# --------------------------------------------------------------------------- #
# (de)serialization
# --------------------------------------------------------------------------- #
def _from_dict(cls, d: dict):
    """Recursively build a (possibly nested) dataclass from a dict."""
    if not dataclasses.is_dataclass(cls):
        return d
    hints = get_type_hints(cls)
    kwargs = {}
    for f in dataclasses.fields(cls):
        if f.name not in d:
            continue
        val = d[f.name]
        ftype = hints[f.name]
        if dataclasses.is_dataclass(ftype) and isinstance(val, dict):
            kwargs[f.name] = _from_dict(ftype, val)
        elif ftype is tuple and isinstance(val, list):
            kwargs[f.name] = tuple(val)
        elif ftype is float and isinstance(val, (str, int)):
            kwargs[f.name] = float(val)  # YAML 1.1 parses '1e-3' as str
        elif ftype is int and isinstance(val, str):
            kwargs[f.name] = int(float(val))
        elif ftype is bool and isinstance(val, str):
            kwargs[f.name] = val.lower() in ("1", "true", "yes")
        else:
            kwargs[f.name] = val
    return cls(**kwargs)


def _deep_merge(base: dict, override: dict) -> dict:
    out = dict(base)
    for k, v in override.items():
        if k in out and isinstance(out[k], dict) and isinstance(v, dict):
            out[k] = _deep_merge(out[k], v)
        else:
            out[k] = v
    return out


def _load_yaml_with_base(path: str) -> dict:
    with open(path, "r") as fh:
        raw = yaml.safe_load(fh) or {}
    base_ref = raw.pop("_base_", None)
    if base_ref is None:
        return raw
    base_path = base_ref if os.path.isabs(base_ref) else os.path.join(os.path.dirname(path), base_ref)
    base = _load_yaml_with_base(base_path)
    return _deep_merge(base, raw)


def _set_dotted(d: dict, dotted: str, value: Any) -> None:
    keys = dotted.split(".")
    cur = d
    for k in keys[:-1]:
        cur = cur.setdefault(k, {})
    cur[keys[-1]] = value


def _coerce(value: str) -> Any:
    try:
        return yaml.safe_load(value)
    except Exception:
        return value


def load_config(path: str | None = None, overrides: list[str] | None = None) -> MapGSConfig:
    """Load a :class:`MapGSConfig` from YAML + dotted overrides.

    Args:
        path: optional YAML path (supports ``_base_``).
        overrides: list of ``a.b.c=value`` strings.
    """
    merged: dict = {}
    if path:
        merged = _load_yaml_with_base(path)
    for ov in overrides or []:
        if "=" not in ov:
            raise ValueError(f"override must be key=value, got {ov!r}")
        key, val = ov.split("=", 1)
        _set_dotted(merged, key.strip(), _coerce(val.strip()))
    return _from_dict(MapGSConfig, merged)