| import pytest |
| import numpy as np |
| from pathlib import Path |
| from types import SimpleNamespace |
|
|
| torch = pytest.importorskip("torch") |
|
|
| from cil.models import CTTConfig, CausalTangentTransport, ChartEncoder, TangentNormalizer |
| from cil.models.ctt import ( |
| chamfer_to_target_set, |
| diversity_loss, |
| entropic_ot_alignment_loss, |
| negative_boundary_loss, |
| ) |
| from scripts.eval_ctt_generated_rollout import ( |
| ChartItem, |
| Proposal, |
| _action_bound_diagnostics_4d, |
| _clip_to_action_space_4d, |
| _measured_row_from_rollout, |
| _parse_execution_action_scale_vector, |
| _scale_actions_4d, |
| _source_pool_for_target, |
| _transform_actions_4d, |
| ) |
|
|
|
|
| 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 torch.isfinite(entropic_ot_alignment_loss(predicted, positives, epsilon=0.1, iterations=5)) |
| 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"] |
|
|
|
|
| def test_action_bound_diagnostics_measure_preclip_violation(): |
| env = SimpleNamespace( |
| single_action_space=SimpleNamespace( |
| low=np.asarray([-1.0, -0.5], dtype=np.float32), |
| high=np.asarray([1.0, 0.5], dtype=np.float32), |
| ) |
| ) |
| actions = np.asarray( |
| [[[[0.0, 0.0], [1.2, -0.75]], [[0.1, 0.2], [0.3, 0.4]]]], |
| dtype=np.float32, |
| ) |
|
|
| diagnostics = _action_bound_diagnostics_4d(actions, env) |
| clipped = _clip_to_action_space_4d(actions, env) |
|
|
| assert diagnostics is not None |
| assert diagnostics["violation"].tolist() == [[True, False]] |
| assert diagnostics["max_abs"][0, 0] == pytest.approx(0.25) |
| assert clipped[0, 0, 1, 0] == pytest.approx(1.0) |
| assert clipped[0, 0, 1, 1] == pytest.approx(-0.5) |
|
|
|
|
| def test_execution_action_scale_applies_before_bound_diagnostics(): |
| env = SimpleNamespace( |
| single_action_space=SimpleNamespace( |
| low=np.asarray([-1.0, -1.0], dtype=np.float32), |
| high=np.asarray([1.0, 1.0], dtype=np.float32), |
| ) |
| ) |
| actions = np.asarray([[[[2.0, -4.0]]]], dtype=np.float32) |
|
|
| scaled = _scale_actions_4d(actions, 0.25) |
| diagnostics = _action_bound_diagnostics_4d(scaled, env) |
|
|
| assert scaled.tolist() == [[[[0.5, -1.0]]]] |
| assert diagnostics is not None |
| assert diagnostics["violation"].tolist() == [[False]] |
|
|
|
|
| def test_execution_action_scale_vector_is_dimension_wise(): |
| actions = np.asarray([[[[2.0, -4.0, 3.0]]]], dtype=np.float32) |
| vector = _parse_execution_action_scale_vector("0.5,0.25,2.0") |
|
|
| scaled = _scale_actions_4d(actions, 0.5, scale_vector=vector) |
|
|
| assert vector is not None |
| assert scaled.tolist() == [[[[0.5, -0.5, 3.0]]]] |
|
|
|
|
| def test_execution_action_scale_vector_pads_after_adapted_dims(): |
| actions = np.asarray([[[[2.0, -4.0, 3.0, 8.0]]]], dtype=np.float32) |
| vector = _parse_execution_action_scale_vector("0.5,0.25,2.0") |
|
|
| scaled = _scale_actions_4d(actions, 0.5, scale_vector=vector) |
|
|
| assert vector is not None |
| assert scaled.tolist() == [[[[0.5, -0.5, 3.0, 4.0]]]] |
|
|
|
|
| def test_tanh_execution_transform_maps_to_env_bounds(): |
| env = SimpleNamespace( |
| single_action_space=SimpleNamespace( |
| low=np.asarray([-1.0, -2.0], dtype=np.float32), |
| high=np.asarray([1.0, 2.0], dtype=np.float32), |
| ) |
| ) |
| actions = np.asarray([[[[10.0, -10.0]]]], dtype=np.float32) |
|
|
| transformed = _transform_actions_4d(actions, env, "tanh") |
| diagnostics = _action_bound_diagnostics_4d(transformed, env) |
|
|
| assert transformed[0, 0, 0, 0] <= 1.0 |
| assert transformed[0, 0, 0, 1] >= -2.0 |
| assert diagnostics is not None |
| assert diagnostics["violation"].tolist() == [[False]] |
|
|
|
|
| def test_env_clip_execution_transform_is_logged_bounded_convention(): |
| env = SimpleNamespace( |
| single_action_space=SimpleNamespace( |
| low=np.asarray([-1.0, -2.0], dtype=np.float32), |
| high=np.asarray([1.0, 2.0], dtype=np.float32), |
| ) |
| ) |
| actions = np.asarray([[[[10.0, -3.0], [0.25, 1.5]]]], dtype=np.float32) |
|
|
| transformed = _transform_actions_4d(actions, env, "env_clip") |
| diagnostics = _action_bound_diagnostics_4d(transformed, env) |
|
|
| assert transformed.tolist() == [[[[1.0, -2.0], [0.25, 1.5]]]] |
| assert diagnostics is not None |
| assert diagnostics["violation"].tolist() == [[False]] |
| assert diagnostics["max_abs"].tolist() == [[0.0]] |
|
|
|
|
| def test_measured_rollout_row_records_action_bound_safety_source(): |
| target = ChartItem( |
| chart_id="chart_a", |
| task_id="task", |
| seed="0", |
| state_hash="state_a", |
| instruction="", |
| source_dataset=Path("."), |
| base_action=np.zeros((2, 2), 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=[1.0], |
| hidden_candidate_types=["expert"], |
| stored_base_utility=0.3, |
| ) |
| proposal = Proposal( |
| tangent=np.ones(21, dtype=np.float32), |
| action=np.ones((2, 2), dtype=np.float32), |
| score=0.7, |
| source_chart_id="source", |
| source_task_id="task", |
| source_rank=1, |
| ) |
|
|
| row = _measured_row_from_rollout( |
| target, |
| [proposal], |
| progress=[0.3, 0.4], |
| success=[False, False], |
| utilities=[0.3, 0.4], |
| safety_violations=[False, True], |
| action_clip_max_abs=[0.0, 0.2], |
| restore_error=0.0, |
| execution_action_scale=0.25, |
| execution_action_scale_vector=np.asarray([0.5, 2.0], dtype=np.float32), |
| execution_action_transform="tanh", |
| ) |
|
|
| assert row["execution_action_scale"] == pytest.approx(0.25) |
| assert row["execution_action_scale_vector"] == [0.5, 2.0] |
| assert row["execution_action_transform"] == "tanh" |
| assert row["base_outcome"]["safety_violation"] is False |
| assert row["base_outcome"]["safety_violation_source"] == "action_bounds" |
| assert row["candidate_outcomes"][0]["safety_violation"] is True |
| assert row["candidate_outcomes"][0]["action_bound_violation"] is True |
| assert row["candidate_action_clip_max_abs"] == [0.2] |
|
|
