File size: 4,136 Bytes
20c251e | 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 | from __future__ import annotations
import pytest
from dovla_cil.transfercritic.eval import compare_selection_strategies
from dovla_cil.transfercritic.labeling import make_utility_labels, toy_utility_value
from dovla_cil.transfercritic.schema import DataAtom, TransferContext
from dovla_cil.transfercritic.selection import greedy_marginal_selection
def _atom(
atom_id: str,
*,
score: float,
success: bool = False,
candidate_type: str = "near_miss",
task_id: str = "task_a",
cost: float = 1.0,
) -> DataAtom:
return DataAtom(
record_id=atom_id,
embedding=[score, 1.0 if success else 0.0, cost, 0.5],
candidate_type=candidate_type,
task_metadata={"task_id": task_id, "family": "pick"},
reward_summary={
"progress": score,
"score": score + (1.0 if success else 0.0),
"success": 1.0 if success else 0.0,
"regret": max(0.0, 1.0 - score),
},
effect_summary={"moved_object_count": 1.0, "true_relation_count": score},
cost=cost,
)
def test_data_atom_and_context_schema_roundtrip() -> None:
atom = _atom("r1", score=0.5, success=True)
context = TransferContext(
benchmark_name="CausalStress",
task_family="pick",
target_objects=["red_mug"],
ood_factor="wrong_target",
validation_ref="val://small",
)
restored = DataAtom.from_dict(atom.to_dict())
context_restored = TransferContext.from_dict(context.to_dict())
assert restored == atom
assert context_restored == context
assert restored.atom_id == "r1"
def test_greedy_selection_uses_score_over_cost() -> None:
atoms = [
_atom("cheap_good", score=0.8, success=True, cost=1.0),
_atom("expensive_best", score=1.0, success=True, cost=3.0),
_atom("bad", score=0.1, success=False, cost=1.0),
]
context = TransferContext(benchmark_name="CausalStress", task_family="pick")
result = greedy_marginal_selection(
atoms,
context,
budget=2.0,
score_fn=lambda atom, _selected, _context: float(atom.reward_summary["score"]),
)
assert result.name == "transfercritic"
assert result.total_cost <= 2.0
assert result.atom_ids[0] == "cheap_good"
assert "expensive_best" not in result.atom_ids
def test_toy_utility_labels_prefer_successful_useful_atoms() -> None:
context = TransferContext(benchmark_name="CausalStress", task_family="pick")
good = _atom("good", score=0.9, success=True, candidate_type="near_miss")
poor = _atom("poor", score=0.1, success=False, candidate_type="noop")
labels = make_utility_labels([good, poor], context, method="toy_retraining_delta")
assert labels[0].utility == toy_utility_value(good, context)
assert labels[0].utility > labels[1].utility
assert labels[0].metadata["approximate"] is True
def test_selection_experiments_return_all_baselines() -> None:
atoms = [
_atom("a", score=0.6, task_id="task_a"),
_atom("b", score=0.9, success=True, task_id="task_b"),
_atom("c", score=0.2, task_id="task_a"),
]
context = TransferContext(benchmark_name="CausalStress", task_family="pick")
rows = compare_selection_strategies(atoms, context, budget=2.0, seed=0)
assert {row["name"] for row in rows} == {
"random_subset",
"top_reward_subset",
"task_balanced_subset",
"transfercritic",
}
assert all(float(row["total_cost"]) <= 2.0 for row in rows)
def test_transfercritic_model_shapes() -> None:
torch = pytest.importorskip("torch")
from dovla_cil.transfercritic.model import SetConditionedTransferCritic, TransferCriticConfig
config = TransferCriticConfig(atom_dim=4, set_dim=4, context_dim=4, hidden_dim=8)
model = SetConditionedTransferCritic(config)
atom = torch.randn(3, 4)
current_set = torch.zeros(3, 4)
context = torch.randn(3, 4)
score = model(atom, current_set, context)
assert score.shape == (3,)
score.mean().backward()
assert any(parameter.grad is not None for parameter in model.parameters())
|