tests/test_causalstress.py

from __future__ import annotations

import subprocess
import sys
from pathlib import Path

from dovla_cil.eval.causalstress import (
    CAUSALSTRESS_CATEGORIES,
    CausalStressConfig,
    compute_causalstress_metrics,
    generate_causalstress_groups,
)
from dovla_cil.generation.pipeline import generate_cil_dataset
from dovla_cil.tasks.library import built_in_toy_tasks
from dovla_cil.training.trainer import DoVLATrainer, TrainerConfig
from dovla_cil.utils.io import read_json


def test_causalstress_generation_works() -> None:
    groups = generate_causalstress_groups(
        CausalStressConfig(num_tasks=len(CAUSALSTRESS_CATEGORIES), k=4, seed=1)
    )
    assert len(groups) == len(CAUSALSTRESS_CATEGORIES)
    assert {group.category for group in groups} == set(CAUSALSTRESS_CATEGORIES)
    assert all(len(group.records) == 4 for group in groups)
    assert all(record.group_id == group.group_id for group in groups for record in group.records)


def test_each_causalstress_category_generates_one_group() -> None:
    for category in CAUSALSTRESS_CATEGORIES:
        groups = generate_causalstress_groups(
            CausalStressConfig(num_tasks=1, k=3, seed=4, categories=(category,))
        )
        assert len(groups) == 1
        assert groups[0].category == category
        assert groups[0].records
        assert groups[0].task.success_predicates


def test_hard_causalstress_categories_cycle_named_variants() -> None:
    expected_counts = {
        "similar_distractors": 4,
        "spatial_relation_minimal_pairs": 3,
        "negation_and_avoidance": 2,
        "sequential_tasks": 3,
        "irreversible_failure": 2,
        "physics_perturbation_placeholders": 3,
    }
    for category, count in expected_counts.items():
        groups = generate_causalstress_groups(
            CausalStressConfig(num_tasks=count, k=2, seed=5, categories=(category,))
        )
        assert len({group.task.task_id for group in groups}) == count


def test_causalstress_metrics_on_synthetic_predictions() -> None:
    groups = generate_causalstress_groups(CausalStressConfig(num_tasks=3, k=4, seed=2))
    predictions = {}
    for group in groups:
        predictions[group.group_id] = {
            "scores": [record.reward.score for record in group.records],
            "success": [1.0 if record.reward.terminal_success else 0.0 for record in group.records],
            "progress": [record.reward.progress for record in group.records],
            "regret": [float(record.regret or 0.0) for record in group.records],
            "effects": [
                [0.0] * 32 for _record in group.records
            ],
        }
    # Use target effect vectors as predictions to make the MAE exact zero.
    from dovla_cil.eval.causalstress import _effect_vector

    for group in groups:
        predictions[group.group_id]["effects"] = [
            _effect_vector(record, dim=32) for record in group.records
        ]

    metrics = compute_causalstress_metrics(groups, predictions)
    assert metrics["pairwise_ranking_accuracy"] == 1.0
    assert metrics["top1_action_selection"] == 1.0
    assert metrics["success_prediction_accuracy"] == 1.0
    assert metrics["effect_prediction_mae"] == 0.0
    assert metrics["regret_calibration_error"] == 0.0
    assert "per_category" in metrics
    assert "target_confusion_matrix" in metrics


def test_eval_causalstress_script_runs_on_smoke_checkpoint(tmp_path: Path) -> None:
    dataset_dir = tmp_path / "cil"
    run_dir = tmp_path / "run"
    out_path = tmp_path / "causalstress.json"
    generate_cil_dataset(
        backend="toy",
        tasks=built_in_toy_tasks()[:2],
        out_dir=dataset_dir,
        num_states_per_task=1,
        k=4,
        seed=3,
        shard_size=8,
        inline_observations=True,
    )
    DoVLATrainer(
        TrainerConfig(
            dataset_dir=dataset_dir,
            output_dir=run_dir,
            epochs=1,
            batch_groups=1,
            records_per_group=4,
            hidden_dim=32,
            seed=3,
            device="cpu",
        )
    ).train()

    subprocess.run(
        [
            sys.executable,
            "scripts/eval_causalstress.py",
            "--checkpoint",
            str(run_dir / "best.pt"),
            "--backend",
            "toy",
            "--out",
            str(out_path),
            "--num-tasks",
            "3",
            "--k",
            "4",
            "--seed",
            "3",
        ],
        check=True,
        capture_output=True,
        text=True,
    )
    metrics = read_json(out_path)
    assert metrics["num_groups"] == 3
    assert "pairwise_ranking_accuracy" in metrics
    assert "task_success_rate" in metrics