Upload folder using huggingface_hub

tests/conftest.py

@@ -0,0 +1,47 @@
+import os
+import tempfile
+
+import pytest
+import torch
+
+from mapgs.config import load_config
+
+CUDA = torch.cuda.is_available()
+requires_cuda = pytest.mark.skipif(not CUDA, reason="needs CUDA (gsplat rasterizer)")
+
+
+def tiny_overrides(root):
+    return [
+        "model.embed_dim=128", "model.enc_depth=1", "model.dec_depth=2", "model.n_heads=4",
+        "model.tokens.n_map=128", "model.tokens.n_free=128",
+        "model.tokens.gaussians_per_token=4", "model.tokens.n_dyn_per_instance=16",
+        "model.feature_dim=8",
+        "data.height=48", "data.width=64", "data.num_frames=8", "data.synth_dynamic_actors=1",
+        f"data.root={root}",
+        "train.batch_size=2", "train.amp=false", "train.num_workers=0",
+        "train.warmup=2", "train.extrap_ramp_iter=2", "train.log_every=1000", "train.ckpt_every=0",
+        "loss.lambda_lpips=0.0", "tt.steps=3",
+    ]
+
+
+@pytest.fixture(scope="session")
+def device():
+    return "cuda" if CUDA else "cpu"
+
+
+@pytest.fixture(scope="session")
+def data_root(tmp_path_factory):
+    return str(tmp_path_factory.mktemp("synthetic"))
+
+
+@pytest.fixture(scope="session")
+def cfg(data_root):
+    return load_config(overrides=tiny_overrides(data_root))
+
+
+@pytest.fixture(scope="session")
+def dataset(cfg):
+    if not CUDA:
+        pytest.skip("synthetic GT rendering needs CUDA")
+    from mapgs.data import SyntheticDataset
+    return SyntheticDataset(cfg, "train", n_scenes=4, n_sup_views=4, device="cuda")

tests/test_data_train_eval.py

@@ -0,0 +1,83 @@
+import torch
+
+from mapgs.data import collate_samples
+from mapgs.train import Trainer, compute_step_losses
+from mapgs.train.trainer import _move_batch
+from mapgs.eval import Evaluator, psnr, d_rmse
+from .conftest import requires_cuda
+
+
+@requires_cuda
+def test_sample_shapes(cfg, dataset):
+    from mapgs.data.synthetic import _context_frames
+    s = dataset[0]
+    Vc = len(_context_frames(cfg)) * 3  # context frames clamp+dedup within num_frames
+    assert s.ctx_images.shape == (Vc, 3, cfg.data.height, cfg.data.width)
+    assert s.ctx_tids.shape[0] == Vc
+    assert s.anchor_pos.shape == (cfg.model.tokens.n_map, 3)
+    assert s.sup_depth.shape[0] == s.sup_images.shape[0]
+
+
+@requires_cuda
+def test_collate_and_dynamic_padding(cfg, dataset):
+    batch = collate_samples([dataset[0], dataset[1]])
+    assert batch["ctx_images"].shape[0] == 2
+    if batch["dynamic"] is not None:
+        assert batch["dynamic"]["box_centers"].shape[0] == 2
+        assert batch["dynamic"]["valid"].shape[0] == 2
+
+
+@requires_cuda
+def test_step_losses_finite_and_backward(cfg, dataset):
+    trainer = Trainer(cfg)
+    batch = _move_batch(collate_samples([dataset[0], dataset[1]]), "cuda")
+    total, log = compute_step_losses(trainer.model, trainer.ras, batch, trainer.criterion, 100, cfg, "cuda")
+    assert torch.isfinite(total)
+    for k in ["rgb", "mapdepth", "vis", "vert", "free_space"]:
+        assert k in log
+    total.backward()
+    assert sum(p.grad.abs().sum() for p in trainer.model.parameters() if p.grad is not None) > 0
+
+
+@requires_cuda
+def test_amp_train_step_runs(cfg, dataset):
+    # guards the autocast (bf16) -> fp32 rendering boundary
+    prev = cfg.train.amp
+    cfg.train.amp = True
+    try:
+        trainer = Trainer(cfg)
+        log = trainer.train_step(collate_samples([dataset[0], dataset[1]]))
+        assert log["total"] == log["total"]  # finite (not NaN)
+        assert log["grad_norm"] >= 0
+    finally:
+        cfg.train.amp = prev
+
+
+@requires_cuda
+def test_training_decreases_loss(cfg, dataset):
+    trainer = Trainer(cfg)
+    losses = []
+    for _ in range(20):
+        batch = collate_samples([dataset[0], dataset[1]])
+        losses.append(trainer.train_step(batch)["total"])
+    assert min(losses[10:]) < losses[0]  # improved over the first step
+
+
+@requires_cuda
+def test_metrics_sanity():
+    img = torch.rand(3, 16, 16, device="cuda")
+    assert psnr(img, img) > 50
+    d = torch.rand(1, 8, 8, device="cuda") + 1
+    assert d_rmse(d, d.clone()) < 1e-5
+
+
+@requires_cuda
+def test_evaluator_runs(cfg, dataset):
+    trainer = Trainer(cfg)
+    ev = Evaluator(trainer.model, cfg, device="cuda")
+    interp = ev.interpolation(dataset, max_scenes=2)
+    assert "PSNR" in interp and "D-RMSE" in interp
+    sweep = ev.extrapolation_sweep(dataset, shifts=[2.0], max_scenes=2, frame=cfg.data.num_frames // 2)
+    assert 2.0 in sweep
+    lane = ev.lane_consistency(dataset, max_scenes=2, frame=cfg.data.num_frames // 2)
+    assert "lane_mIoU" in lane

tests/test_geometry.py

@@ -0,0 +1,56 @@
+import math
+
+import torch
+
+from mapgs.geometry import (
+    quat_to_rotmat, rotmat_to_quat, se3_inverse, project_points,
+    backproject_depth, plucker_embedding, look_at_pose,
+)
+
+
+def test_quat_rotmat_roundtrip():
+    q = torch.randn(20, 4)
+    q = q / q.norm(dim=-1, keepdim=True)
+    R = quat_to_rotmat(q)
+    # orthonormal, det 1
+    eye = R @ R.transpose(-1, -2)
+    assert torch.allclose(eye, torch.eye(3).expand_as(eye), atol=1e-5)
+    assert torch.allclose(torch.det(R), torch.ones(20), atol=1e-5)
+    q2 = rotmat_to_quat(R)
+    R2 = quat_to_rotmat(q2)
+    assert torch.allclose(R, R2, atol=1e-5)
+
+
+def test_se3_inverse():
+    T = torch.eye(4).repeat(5, 1, 1)
+    T[:, :3, :3] = quat_to_rotmat(torch.randn(5, 4))
+    T[:, :3, 3] = torch.randn(5, 3)
+    I = T @ se3_inverse(T)
+    assert torch.allclose(I, torch.eye(4).expand_as(I), atol=1e-5)
+
+
+def test_projection_roundtrip():
+    K = torch.tensor([[300., 0, 224], [0, 300, 128], [0, 0, 1]])[None]
+    c2w = look_at_pose(torch.tensor([0., 0, 1.5]), torch.tensor([0., 20, 0.5]))[None]
+    pts = torch.randn(1, 100, 3)
+    pts[..., 1] = pts[..., 1].abs() + 5  # in front
+    uv, z = project_points(pts, K, se3_inverse(c2w))
+    back = backproject_depth(uv, z, K, c2w)
+    assert torch.allclose(back, pts, atol=1e-3)
+    assert (z > 0).all()
+
+
+def test_look_at_forward_is_positive_depth():
+    # a point in front of a forward-looking camera must have positive cam-z
+    c2w = look_at_pose(torch.tensor([0., 0, 1.5]), torch.tensor([0., 30, 0.5]))[None]
+    K = torch.tensor([[300., 0, 224], [0, 300, 128], [0, 0, 1]])[None]
+    pt = torch.tensor([[[0., 10, 0.5]]])
+    _, z = project_points(pt, K, se3_inverse(c2w))
+    assert z.item() > 0
+
+
+def test_plucker_shape():
+    K = torch.tensor([[300., 0, 224], [0, 300, 128], [0, 0, 1]])[None].repeat(3, 1, 1)
+    c2w = torch.eye(4)[None].repeat(3, 1, 1)
+    pl = plucker_embedding(K, c2w, 32, 48)
+    assert pl.shape == (3, 6, 32, 48)

tests/test_hdmap.py

@@ -0,0 +1,69 @@
+import numpy as np
+import torch
+
+from mapgs.config import MapConfig
+from mapgs.hdmap import (
+    grid_field_from_points, HDMap, sample_anchors, rasterize_map_depth, project_polylines,
+)
+from mapgs.hdmap.anchors import resample_polyline, weighted_fps
+from .conftest import requires_cuda
+
+
+def _toy_map():
+    pts = np.random.RandomState(0).uniform([-20, -5], [20, 50], size=(2000, 2))
+    z = 0.02 * pts[:, 0]
+    gp = np.concatenate([pts, z[:, None]], 1)
+    gf = grid_field_from_points(gp, 0.5)
+    lanes = [torch.tensor([[x, y, 0.02 * x] for y in np.arange(0, 45, 1.0)], dtype=torch.float32)
+             for x in [-3.5, 0, 3.5]]
+    bnds = [torch.tensor([[x, y, 0.02 * x] for y in np.arange(0, 45, 2.0)], dtype=torch.float32)
+            for x in [-7, 7]]
+    return HDMap(ground=gf, lanes=lanes, boundaries=bnds)
+
+
+def test_ground_field_bilinear_and_grad():
+    hd = _toy_map()
+    xy = torch.tensor([[1.0, 2.0], [3.0, 4.0]], requires_grad=True)
+    z, valid = hd.height_at(xy)
+    assert valid.all()
+    assert torch.allclose(z, 0.02 * xy[:, 0].detach(), atol=0.05)
+    z.sum().backward()
+    assert xy.grad.abs().sum() > 0
+
+
+def test_anchor_ratios_and_count():
+    hd = _toy_map()
+    cfg = MapConfig(n_anchors=1000)
+    A = sample_anchors(hd, cfg, seed=0)
+    assert len(A) == 1000
+    counts = torch.bincount(A.types, minlength=3).float() / 1000
+    assert abs(counts[0] - 0.6) < 0.05
+    assert abs(counts[1] - 0.3) < 0.05
+    assert abs(counts[2] - 0.1) < 0.05
+
+
+def test_resample_polyline_spacing():
+    pl = torch.tensor([[0., 0, 0], [0, 10, 0]])
+    out = resample_polyline(pl, 1.0)
+    assert out.shape[0] >= 9
+
+
+def test_weighted_fps_prefers_high_weight():
+    pts = torch.randn(200, 3)
+    w = torch.ones(200)
+    w[0] = 100.0
+    sel = weighted_fps(pts, w, 10)
+    assert 0 in sel.tolist()  # highest-weight point is picked first
+
+
+@requires_cuda
+def test_map_depth_coverage():
+    from mapgs.geometry import look_at_pose
+    hd = _toy_map().to("cuda")
+    K = torch.tensor([[60., 0, 32], [0, 60, 24], [0, 0, 1]], device="cuda")[None]
+    c2w = look_at_pose(torch.tensor([0., 0, 1.5]), torch.tensor([0., 20, 0.5]))[None].cuda()
+    depth, mask = rasterize_map_depth(hd.ground, K, c2w, 48, 64)
+    assert mask.float().mean() > 0.2  # lower half sees ground
+    assert (depth[mask] > 0).all()
+    uv = project_polylines(hd.lanes, K, c2w, 48, 64)
+    assert uv[0].shape[0] > 0

tests/test_losses.py

@@ -0,0 +1,81 @@
+import torch
+
+from mapgs.config import LossConfig, TokenConfig
+from mapgs.losses import (
+    Tempering, ssim, chamfer_2d, mapdepth_loss, free_space_loss,
+    ground_coupling_loss, silog_loss, visibility_loss, huber,
+)
+from mapgs.losses.extrap import perturb_pose
+from mapgs.geometry import look_at_pose
+from mapgs.hdmap import grid_field_from_points
+from mapgs.render.gaussians import Gaussians, GROUP_DYNAMIC
+import numpy as np
+
+
+def test_tempering_monotone_and_capped():
+    t = Tempering(LossConfig(), TokenConfig(), total_iters=150000)
+    assert t.eps(0) < t.eps(1000) <= t.eps_max
+    assert t.s(0) < t.s(1000) <= t.s_max
+    assert t.eps(10_000_000) == t.eps_max  # capped
+    assert t.lambda_md_scale(0) == 1.0
+    assert t.lambda_md_scale(149000) == LossConfig().md_late_decay
+
+
+def test_ssim_identical_is_one():
+    img = torch.rand(1, 3, 32, 32)
+    assert ssim(img, img) > 0.999
+
+
+def test_mapdepth_zero_when_equal():
+    d = torch.rand(2, 16, 16) + 1
+    mask = torch.ones(2, 16, 16, dtype=torch.bool)
+    assert mapdepth_loss(d, d.clone(), mask, eps=0.3) < 1e-6
+
+
+def test_huber_matches_quadratic_for_small():
+    x = torch.tensor([0.1, -0.1])
+    assert torch.allclose(huber(x, 1.0), 0.5 * x ** 2)
+
+
+def test_visibility_inside_vs_outside():
+    K = torch.tensor([[60., 0, 32], [0, 60, 24], [0, 0, 1]])[None]
+    c2w = look_at_pose(torch.tensor([0., 0, 1.5]), torch.tensor([0., 20, 0.5]))[None]
+    inside = torch.tensor([[0., 12, 0.5]])      # in front, centered
+    outside = torch.tensor([[0., -30, 0.5]])    # behind camera
+    assert visibility_loss(inside, K, c2w, 48, 64) < 1e-4
+    assert visibility_loss(outside, K, c2w, 48, 64) > 0
+
+
+def test_free_space_and_ground_coupling():
+    gp = np.concatenate([np.random.RandomState(0).uniform([-10, -5], [10, 20], (500, 2)),
+                         np.zeros((500, 1))], 1)
+    gf = grid_field_from_points(gp, 0.5)
+    n = 50
+    means = torch.zeros(n, 3); means[:, 2] = -1.0  # all below ground
+    g = Gaussians(means, torch.full((n, 3), 0.1), _identq(n), torch.full((n,), 0.8), torch.rand(n, 3),
+                  group=torch.full((n,), GROUP_DYNAMIC))
+    assert free_space_loss(g, gf, delta=0.1) > 0          # penalizes below-ground opacity
+    assert ground_coupling_loss(g, gf, eps=1.0) > 0       # pulls dynamic z up
+
+
+def test_silog_zero_when_equal():
+    d = torch.rand(1, 16, 16) + 1
+    mask = torch.ones(1, 16, 16, dtype=torch.bool)
+    assert abs(float(silog_loss(d, d.clone(), mask))) < 1e-6
+
+
+def test_chamfer_zero_for_identical():
+    a = torch.randn(20, 2)
+    assert chamfer_2d(a, a.clone()) < 1e-6
+
+
+def test_perturb_pose_shifts_laterally():
+    c2w = look_at_pose(torch.tensor([0., 0, 1.5]), torch.tensor([0., 20, 0.5]))
+    dev = perturb_pose(c2w, lateral=2.0)
+    shift = (dev[:3, 3] - c2w[:3, 3]).norm()
+    assert abs(float(shift) - 2.0) < 1e-4
+
+
+def _identq(n):
+    q = torch.zeros(n, 4); q[:, 0] = 1
+    return q

tests/test_model_render.py

@@ -0,0 +1,82 @@
+import torch
+
+from mapgs.model import MapGS
+from mapgs.render import Gaussians, GaussianRasterizer, GROUP_MAP, GROUP_DYNAMIC
+from mapgs.geometry import plucker_embedding, look_at_pose
+from .conftest import requires_cuda
+
+
+def _model_inputs(cfg, device, B=1, dynamic=False):
+    V, H, W = 3, cfg.data.height, cfg.data.width
+    imgs = torch.rand(B, V, 3, H, W, device=device)
+    K = torch.tensor([[60., 0, W / 2], [0, 60, H / 2], [0, 0, 1]], device=device)[None, None].repeat(B, V, 1, 1)
+    c2w = torch.stack([look_at_pose(torch.tensor([dx, 0., 1.5]), torch.tensor([dx, 20., 0.5]))
+                       for dx in [-1., 0, 1]]).to(device)[None].repeat(B, 1, 1, 1)
+    pl = torch.stack([plucker_embedding(K[b], c2w[b], H, W) for b in range(B)])
+    tids = torch.zeros(B, V, dtype=torch.long, device=device)
+    nmap = cfg.model.tokens.n_map
+    apos = torch.randn(B, nmap, 3, device=device); apos[..., 2] *= 0.1
+    atype = torch.randint(0, 3, (B, nmap), device=device)
+    anorm = torch.zeros(B, nmap, 3, device=device); anorm[..., 2] = 1
+    dyn = None
+    if dynamic:
+        I, F = 2, cfg.data.num_frames
+        centers = torch.zeros(B, I, F, 3, device=device)
+        for f in range(F):
+            centers[:, 0, f] = torch.tensor([3., 5 + 0.5 * f, 0.5], device=device)
+            centers[:, 1, f] = torch.tensor([-3., 8., 0.5], device=device)
+        dyn = dict(box_centers=centers, box_rots=torch.eye(3, device=device).view(1, 1, 1, 3, 3).repeat(B, I, F, 1, 1),
+                   box_size=torch.ones(B, I, 3, device=device) * 2, valid=torch.ones(B, I, dtype=torch.bool, device=device),
+                   canon_idx=torch.zeros(B, I, dtype=torch.long, device=device))
+    return imgs, pl, tids, apos, atype, anorm, K, c2w, dyn
+
+
+@requires_cuda
+def test_forward_shapes_and_budget(cfg):
+    model = MapGS(cfg).cuda()
+    imgs, pl, tids, apos, atype, anorm, K, c2w, _ = _model_inputs(cfg, "cuda")
+    g = model(imgs, pl, tids, apos, atype, anorm, s_t=0.5)
+    M = (cfg.model.tokens.n_map + cfg.model.tokens.n_free) * cfg.model.tokens.gaussians_per_token
+    assert g.means.shape == (1, M, 3)
+    assert g.colors.shape[-1] == cfg.model.feature_dim
+
+
+@requires_cuda
+def test_bounded_residual_for_map_tokens(cfg):
+    model = MapGS(cfg).cuda()
+    imgs, pl, tids, apos, atype, anorm, K, c2w, _ = _model_inputs(cfg, "cuda")
+    s_t = 0.5
+    g = model(imgs, pl, tids, apos, atype, anorm, s_t=s_t)
+    is_map = g.group[0] == GROUP_MAP
+    # map gaussian means must lie within s_t of their anchor (sqrt(3)*s_t bound on the cube)
+    means = g.means[0][is_map]
+    ng = cfg.model.tokens.gaussians_per_token
+    anchors = apos[0].repeat_interleave(ng, 0)[: means.shape[0]]
+    dev = (means - anchors).abs().max()
+    assert dev <= s_t + 1e-4
+
+
+@requires_cuda
+def test_dynamic_placement_moves_only_dynamic(cfg):
+    model = MapGS(cfg).cuda()
+    imgs, pl, tids, apos, atype, anorm, K, c2w, dyn = _model_inputs(cfg, "cuda", dynamic=True)
+    g = model(imgs, pl, tids, apos, atype, anorm, s_t=0.5, dynamic=dyn)
+    g0 = model.place_dynamics(g, dyn, 0)
+    g1 = model.place_dynamics(g, dyn, cfg.data.num_frames - 1)
+    dynm = g.group[0] == GROUP_DYNAMIC
+    assert (g1.means[0][dynm] - g0.means[0][dynm]).norm(dim=-1).mean() > 0.1
+    assert torch.allclose(g1.means[0][~dynm], g0.means[0][~dynm])
+
+
+@requires_cuda
+def test_render_and_backward(cfg):
+    model = MapGS(cfg).cuda()
+    imgs, pl, tids, apos, atype, anorm, K, c2w, _ = _model_inputs(cfg, "cuda")
+    g = model(imgs, pl, tids, apos, atype, anorm, s_t=0.5)
+    ras = GaussianRasterizer()
+    out = ras.render(g.scene(0), K[0], c2w[0], cfg.data.height, cfg.data.width)
+    rgb = model.feature_to_rgb(out.color)
+    assert rgb.shape == (3, 3, cfg.data.height, cfg.data.width)
+    assert out.aux is not None  # lane channel rendered
+    (rgb.mean() + out.depth.mean()).backward()
+    assert sum(p.grad.abs().sum() for p in model.parameters() if p.grad is not None) > 0

tests/test_ttt.py

@@ -0,0 +1,31 @@
+import torch
+
+from mapgs.model import MapGS
+from mapgs.ttt import token_tuning, map_guided_densify
+from mapgs.ttt.densify import anchors_along_trajectory
+from .conftest import requires_cuda
+
+
+@requires_cuda
+def test_token_tuning_runs(cfg, dataset):
+    model = MapGS(cfg).cuda()
+    g = token_tuning(model, dataset[0], cfg, steps=3)
+    M = (cfg.model.tokens.n_map + cfg.model.tokens.n_free) * cfg.model.tokens.gaussians_per_token
+    assert g.means.shape == (1, M, 3)
+    # network weights must be unchanged by TT (only tokens are tuned)
+    before = {n: p.clone() for n, p in model.named_parameters()}
+    token_tuning(model, dataset[0], cfg, steps=2)
+    for n, p in model.named_parameters():
+        assert torch.allclose(before[n], p)
+
+
+@requires_cuda
+def test_densify_increases_gaussians(cfg, dataset):
+    model = MapGS(cfg).cuda()
+    s = dataset[0]
+    traj = torch.stack([torch.zeros(15, device="cuda"), torch.linspace(2, 25, 15, device="cuda")], -1)
+    npos, ntyp, nnrm = anchors_along_trajectory(s.ground.to("cuda"), traj, spacing=1.0)
+    n_new = min(64, npos.shape[0])
+    g = map_guided_densify(model, s, npos[:n_new], ntyp[:n_new], nnrm[:n_new], cfg)
+    base = (cfg.model.tokens.n_map + cfg.model.tokens.n_free) * cfg.model.tokens.gaussians_per_token
+    assert g.means.shape[1] == base + n_new * cfg.model.tokens.gaussians_per_token

tests/test_unified.py

@@ -0,0 +1,68 @@
+import os
+
+import torch
+
+from mapgs.config import load_config
+from mapgs.geometry import resize_with_intrinsics
+from mapgs.hdmap.ground_field import GridGroundField
+from .conftest import requires_cuda, tiny_overrides
+
+
+def test_ground_field_from_raster_and_scaled():
+    gf = GridGroundField.from_raster(torch.ones(10, 12) * 2.0, x0=-3.0, y0=-2.0, dx=0.5)
+    z, valid = gf.height_at(torch.tensor([[0.0, 0.0]]))
+    assert valid.all() and abs(float(z[0]) - 2.0) < 1e-4
+    g2 = gf.scaled(0.5)
+    assert abs(g2.x0 - (-1.5)) < 1e-6 and abs(g2.dx - 0.25) < 1e-6
+    # scaling halves heights too (spatial)
+    z2, _ = g2.height_at(torch.tensor([[0.0, 0.0]]))
+    assert abs(float(z2[0]) - 1.0) < 1e-4
+
+
+def test_resize_with_intrinsics():
+    img = torch.rand(3, 100, 200)
+    K = torch.tensor([[200.0, 0, 100], [0, 200, 50], [0, 0, 1]])
+    out, Ks = resize_with_intrinsics(img, K, 50, 50)
+    assert out.shape == (3, 50, 50)
+    assert abs(float(Ks[0, 0]) - 50.0) < 1e-3   # fx * (50/200)
+    assert abs(float(Ks[1, 1]) - 100.0) < 1e-3  # fy * (50/100)
+    assert abs(float(Ks[0, 2]) - 25.0) < 1e-3   # cx * (50/200)
+
+
+@requires_cuda
+def test_unified_convert_load_and_mixed_roots(tmp_path):
+    cfg = load_config(overrides=tiny_overrides(str(tmp_path / "src")) + ["data.name=unified"])
+    from mapgs.data.convert import convert_synthetic
+    from mapgs.data import UnifiedClipDataset, collate_samples
+
+    ra = str(tmp_path / "ua"); rb = str(tmp_path / "ub")
+    convert_synthetic(cfg, ra, "train", n_clips=2, device="cuda")
+    convert_synthetic(cfg, rb, "train", n_clips=2, device="cuda")
+
+    # on-disk layout
+    clip = sorted(os.listdir(os.path.join(ra, "train")))[0]
+    files = os.listdir(os.path.join(ra, "train", clip))
+    assert "meta.pt" in files and "images" in files
+
+    ds = UnifiedClipDataset(cfg, roots=[ra, rb], split="train", n_sup_views=3)
+    assert len(ds) == 4  # mixed roots concatenated
+    s = ds[0]
+    assert s.ctx_images.shape[-2:] == (cfg.data.height, cfg.data.width)
+    assert s.anchor_pos.shape[0] == cfg.model.tokens.n_map
+    batch = collate_samples([ds[0], ds[1]])
+    assert batch["ctx_images"].shape[0] == 2
+
+
+@requires_cuda
+def test_unified_train_step(tmp_path):
+    cfg = load_config(overrides=tiny_overrides(str(tmp_path / "src")) + ["data.name=unified"])
+    from mapgs.data.convert import convert_synthetic
+    from mapgs.data import UnifiedClipDataset, collate_samples
+    from mapgs.train import Trainer
+
+    root = str(tmp_path / "u")
+    convert_synthetic(cfg, root, "train", n_clips=2, device="cuda")
+    ds = UnifiedClipDataset(cfg, roots=root, split="train", n_sup_views=3)
+    tr = Trainer(cfg)
+    log = tr.train_step(collate_samples([ds[0], ds[1]]))
+    assert log["total"] == log["total"]  # finite