from __future__ import annotations import numpy as np import pytest import torch from types import SimpleNamespace from src.data_contract import ( CosmosRBMCollator, adapt_official_trajectory, audit_group_overlap, deterministic_frame_indices, make_preference_contract, prepare_trajectory, validate_span_bounds, worker_invariant_sample_ids, ) from src.reward_heads import SharedRewardHeads def synthetic_frames(count: int, height: int = 12, width: int = 20) -> np.ndarray: frames = np.empty((count, height, width, 3), dtype=np.uint8) for index in range(count): frames[index].fill(index) return frames def prepared( *, strategy: str = "forward", quality: str = "successful", frame_count: int = 13, trajectory_id: str = "trajectory-a", ): return prepare_trajectory( trajectory_id=trajectory_id, frames=synthetic_frames(frame_count), task="insert cable", data_source="unit_test", quality_label=quality, strategy=strategy, canvas_size=16, ) def test_forward_progress_is_nondecreasing_and_sampling_is_exact() -> None: item = prepared() expected, _ = deterministic_frame_indices(13, 8) assert np.array_equal(item.source_frame_indices, expected) assert np.all(np.diff(item.progress) >= 0) assert item.frames_uint8.shape == (8, 16, 16, 3) def test_reverse_keeps_all_frame_fields_synchronized() -> None: forward = prepared() reverse = prepared(strategy="reverse") assert np.array_equal(reverse.source_frame_indices, forward.source_frame_indices[::-1]) assert np.array_equal(reverse.progress, forward.progress[::-1]) assert np.array_equal(reverse.frame_valid_mask, forward.frame_valid_mask[::-1]) assert np.array_equal(reverse.success_labels, forward.success_labels[::-1]) def test_rewind_has_a_real_progress_drop() -> None: rewind = prepared(strategy="rewind") assert np.any(np.diff(rewind.progress) < 0) def test_different_task_is_zero_with_supervision() -> None: item = prepared(strategy="different_task", quality="failure") assert np.all(item.progress == 0) assert np.array_equal(item.progress_loss_mask, item.frame_valid_mask) def test_preference_swap_flips_label_and_bradley_terry_logit() -> None: chosen = prepared(trajectory_id="chosen") rejected = prepared(strategy="rewind", trajectory_id="rejected") pair = make_preference_contract(chosen, rejected, strategy="rewind") swapped = pair.swapped() assert swapped.preference_label == 0.0 assert swapped.trajectory_a.trajectory_id == "rejected" torch.manual_seed(7) heads = SharedRewardHeads(hidden_size=8, dropout=0.0) features_a = torch.randn(1, 8, 8) features_b = torch.randn(1, 8, 8) mask = torch.ones(1, 8, dtype=torch.bool) ab = heads.preference_logit(features_a, mask, features_b, mask) ba = heads.preference_logit(features_b, mask, features_a, mask) assert torch.allclose(ab, -ba) def test_padding_never_enters_loss_or_pooling() -> None: item = prepared(frame_count=3) batch = CosmosRBMCollator()([item])["trajectory_inputs"] assert batch["frame_valid_mask"].tolist() == [ [True, True, True, False, False, False, False, False] ] assert not batch["progress_loss_mask"][0, 3:].any() assert not batch["success_loss_mask"][0, 3:].any() pool_features = torch.randn(1, 8, 4) heads = SharedRewardHeads(hidden_size=4, dropout=0.0) score_a = heads(pool_features, batch["frame_valid_mask"])["trajectory_scores"] pool_features[:, 3:] = 1000 score_b = heads(pool_features, batch["frame_valid_mask"])["trajectory_scores"] assert torch.allclose(score_a, score_b) def test_failure_does_not_receive_dense_progress() -> None: failure = prepared(quality="failure") assert not failure.progress_loss_mask.any() assert not failure.success_labels.any() def test_official_sampler_adapter_preserves_labels_and_indices() -> None: official_progress = np.linspace(0.0, 1.0, 8, dtype=np.float32) official_success = [0.0] * 7 + [1.0] trajectory = SimpleNamespace( id="official-a", task="insert cable", data_source="official_unit_test", quality_label="successful", frames=synthetic_frames(8), target_progress=official_progress.tolist(), success_label=official_success, predict_last_frame_mask=[1.0] * 8, metadata={ "selected_frame_indices": [0, 4, 9, 13, 18, 22, 27, 31], "frame_valid_mask": [True] * 8, }, ) adapted = adapt_official_trajectory( trajectory, strategy="forward", progress_supervised=True, canvas_size=16, ) assert np.allclose(adapted.progress, official_progress) assert adapted.success_labels.tolist() == official_success assert adapted.source_frame_indices.tolist() == trajectory.metadata[ "selected_frame_indices" ] assert adapted.metadata["adapter"] == "official_robometer_sampler_v1" def test_all_backbones_consume_the_same_frozen_indices() -> None: frozen, _ = deterministic_frame_indices(32, 8) for _backbone in ("qwen3-vl", "cosmos3-edge", "frozen-cosmos"): item = prepared(frame_count=32) assert np.array_equal(item.source_frame_indices, frozen) def test_visual_span_contract() -> None: validate_span_bounds([(1, 4), (8, 11), (12, 15)], 3, 16) with pytest.raises(ValueError, match="overlap"): validate_span_bounds([(1, 5), (4, 8)], 2, 10) with pytest.raises(ValueError, match="count"): validate_span_bounds([(1, 5)], 2, 10) def test_worker_count_does_not_change_seeded_sampling_list() -> None: ids = [f"trajectory-{index}" for index in range(100)] expected = worker_invariant_sample_ids(ids, seed=17, worker_count=1) assert worker_invariant_sample_ids(ids, seed=17, worker_count=2) == expected assert worker_invariant_sample_ids(ids, seed=17, worker_count=8) == expected def test_group_overlap_gate() -> None: assert audit_group_overlap(["a", "b", "c"], ["d", "e"]) == set() assert audit_group_overlap(["a", "b"], ["b", "c"]) == {"b"} def test_one_hundred_trajectories_cpu_collate_and_backward() -> None: item = prepared() batch = CosmosRBMCollator()([item] * 100)["trajectory_inputs"] assert batch["frames_uint8"].shape == (100, 8, 16, 16, 3) features = torch.randn(8, 8, 16, requires_grad=True) heads = SharedRewardHeads(hidden_size=16, dropout=0.0) output = heads(features, batch["frame_valid_mask"][:8]) loss = output["progress_logits"].float().mean() + output[ "success_logits" ].float().mean() loss.backward() assert features.grad is not None and torch.isfinite(features.grad).all()