File size: 6,804 Bytes
e41bef1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 | 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()
|