File size: 3,329 Bytes
04f12e2 | 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 | import pytest
import numpy as np
from pathlib import Path
torch = pytest.importorskip("torch")
from cil.models import CTTConfig, CausalTangentTransport, ChartEncoder, TangentNormalizer
from cil.models.ctt import chamfer_to_target_set, diversity_loss, negative_boundary_loss
from scripts.eval_ctt_generated_rollout import ChartItem, _source_pool_for_target
def test_ctt_variants_preserve_tangent_shape():
z_source = torch.randn(3, 8)
z_target = torch.randn(3, 8)
xi_source = torch.randn(3, 5)
for variant in ("residual", "gated_residual"):
model = CausalTangentTransport(
CTTConfig(chart_feature_dim=10, chart_dim=8, tangent_dim=5, hidden_dim=16, variant=variant)
)
output = model(z_source, z_target, xi_source)
assert output.shape == xi_source.shape
assert torch.isfinite(output).all()
def test_chart_encoder_and_losses_are_finite():
encoder = ChartEncoder(input_dim=6, hidden_dim=12, output_dim=8)
z = encoder(torch.randn(4, 6))
assert z.shape == (4, 8)
predicted = torch.randn(4, 5)
positives = predicted + 0.1
negatives = predicted + 10.0
assert torch.isfinite(chamfer_to_target_set(predicted, positives))
assert negative_boundary_loss(predicted, negatives, margin=0.2).item() == pytest.approx(0.0)
assert diversity_loss(predicted).item() >= 0.0
def test_gated_residual_matches_documented_formula():
model = CausalTangentTransport(
CTTConfig(chart_feature_dim=10, chart_dim=2, tangent_dim=2, hidden_dim=4, variant="gated_residual")
)
for parameter in model.delta.parameters():
parameter.data.zero_()
for parameter in model.gate.parameters():
parameter.data.zero_()
model.delta[-1].bias.data[:] = torch.tensor([4.0, -2.0])
model.gate[2].bias.data[:] = torch.tensor([0.0, 0.0])
xi = torch.tensor([[2.0, 6.0]])
output = model(torch.zeros(1, 2), torch.zeros(1, 2), xi)
assert torch.allclose(output, torch.tensor([[3.0, 2.0]]), atol=1.0e-6)
def test_tangent_normalizer_round_trips():
values = torch.randn(8, 5)
normalizer = TangentNormalizer.fit(values)
restored = normalizer.inverse_transform(normalizer.transform(values))
assert torch.allclose(values, restored, atol=1.0e-5)
def test_rollout_source_pool_excludes_target_chart_and_state_hash():
def chart(chart_id: str, state_hash: str) -> ChartItem:
return ChartItem(
chart_id=chart_id,
task_id="task",
seed="0",
state_hash=state_hash,
instruction="",
source_dataset=Path("."),
base_action=np.zeros((2, 7), dtype=np.float32),
feature=np.zeros(4, dtype=np.float32),
positive_tangents=np.zeros((1, 21), dtype=np.float32),
negative_tangents=np.zeros((0, 21), dtype=np.float32),
hidden_utilities=[],
hidden_candidate_types=[],
stored_base_utility=None,
)
target = chart("chart_a", "state_a")
sources = [chart("chart_a", "state_a"), chart("chart_b", "state_a"), chart("chart_c", "state_c")]
pool = _source_pool_for_target(
target,
task_pool=sources,
source_charts=sources,
exclude_self_source=True,
)
assert [item.chart_id for item in pool] == ["chart_c"]
|