fix(il): restore random val split and shape assert in sampler (#776)

get_val_mask hardcoded `val_idxs = -1` (the real `rng.choice(...)` selection was
commented out), so the validation set was always exactly the last episode
regardless of val_ratio/seed — every reported IL validation metric was computed
on one fixed episode. Restore the seeded random selection of n_val episodes.

create_indices had `episode_mask.shape == episode_ends.shape` as a bare no-op
expression (missing the `assert` keyword), so the shape precondition never ran.
Restore the assert (numba-nopython compatible; the sole caller always builds
matching shapes, so it never trips in practice but fails fast on misuse).

Adds general regression tests: val mask selects round(n*ratio) episodes and is
seed-reproducible; create_indices asserts on a shape mismatch. Red on pre-fix.

Author: geng-haoran

roboverse_learn/il/utils/sampler.py

@@ -25,7 +25,7 @@ def create_indices(
     pad_after: int = 0,
     debug: bool = True,
 ) -> np.ndarray:
-    episode_mask.shape == episode_ends.shape
+    assert episode_mask.shape == episode_ends.shape
     pad_before = min(max(pad_before, 0), sequence_length - 1)
     pad_after = min(max(pad_after, 0), sequence_length - 1)
 
@@ -68,8 +68,7 @@ def get_val_mask(n_episodes, val_ratio, seed=0):
     # have at least 1 episode for validation, and at least 1 episode for train
     n_val = min(max(1, round(n_episodes * val_ratio)), n_episodes - 1)
     rng = np.random.default_rng(seed=seed)
-    # val_idxs = rng.choice(n_episodes, size=n_val, replace=False)
-    val_idxs = -1
+    val_idxs = rng.choice(n_episodes, size=n_val, replace=False)
     val_mask[val_idxs] = True
     return val_mask
 

tests/test_il_numba_optional.py

@@ -66,3 +66,39 @@ def test_create_indices_runs_without_numba():
     assert len(indices) > 0
     # buffer windows never exceed the data length.
     assert indices[:, 1].max() <= episode_ends[-1]
+
+
+@pytest.mark.general
+def test_get_val_mask_selects_n_val_episodes():
+    """Regression: the validation split must select ~val_ratio episodes, not just one.
+
+    The real selection was commented out and replaced with ``val_idxs = -1`` so
+    the val set was always exactly the last episode regardless of val_ratio/seed.
+    """
+    from roboverse_learn.il.utils.sampler import get_val_mask
+
+    n_episodes, val_ratio = 10, 0.2
+    mask = get_val_mask(n_episodes, val_ratio, seed=0)
+    assert mask.sum() == 2, f"expected round(10*0.2)=2 val episodes, got {int(mask.sum())}"
+    # Not degenerate-to-last: with seed=0 the chosen set must not be only the last index.
+    assert not (mask[-1] and mask.sum() == 1)
+
+    # Reproducible under a fixed seed.
+    assert np.array_equal(mask, get_val_mask(n_episodes, val_ratio, seed=0))
+    # val_ratio <= 0 -> empty val set (unchanged contract).
+    assert get_val_mask(n_episodes, 0.0).sum() == 0
+
+
+@pytest.mark.general
+def test_create_indices_asserts_shape_mismatch():
+    """Regression: the shape precondition must actually assert (it was a no-op expr).
+
+    With episode_mask longer than episode_ends, the pre-fix code silently ran to
+    completion; the restored ``assert`` makes it fail fast.
+    """
+    from roboverse_learn.il.utils.sampler import create_indices
+
+    episode_ends = np.array([5, 10], dtype=np.int64)
+    episode_mask = np.ones(3, dtype=bool)  # deliberately mismatched (len 3 vs 2)
+    with pytest.raises(AssertionError):
+        create_indices(episode_ends, sequence_length=3, episode_mask=episode_mask, debug=False)