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]