dovla_cil/experiments/baselines.py

from __future__ import annotations

from dataclasses import asdict, dataclass
from dataclasses import replace as dataclass_replace
from pathlib import Path
from typing import Any

try:
    from pydantic import BaseModel, ConfigDict, Field

    PYDANTIC_AVAILABLE = True
except ImportError:  # pragma: no cover - keeps bare smoke environments usable
    BaseModel = object
    ConfigDict = None
    Field = None
    PYDANTIC_AVAILABLE = False

from dovla_cil.data.datasets import CILDataset
from dovla_cil.data.schema import ActionChunk, CILRecord, RewardInfo, compute_regret_and_ranks
from dovla_cil.data.sharding import write_cil_shards
from dovla_cil.eval.causalstress import CausalStressBenchmark, CausalStressConfig
from dovla_cil.training.losses import InterventionalLossWeights
from dovla_cil.training.trainer import DoVLATrainer, TrainerConfig
from dovla_cil.utils.io import ensure_dir, write_json

BASELINES = (
    "expert_only_bc",
    "more_independent_demos",
    "random_negatives",
    "cross_state_negatives",
    "label_only_counterfactual",
    "world_model_auxiliary",
    "no_effect_head",
    "no_rank_regret",
)

ALIASES = {
    "more_independent_demonstrations": "more_independent_demos",
    "label_only_counterfactuals": "label_only_counterfactual",
}


if PYDANTIC_AVAILABLE:

    class BaselineConfig(BaseModel):
        model_config = ConfigDict(arbitrary_types_allowed=True)

        baseline: str
        dataset: str | Path
        out: str | Path
        backend: str = "toy"
        epochs: int = 1
        batch_groups: int = 4
        records_per_group: int | None = 8
        hidden_dim: int = 128
        lr: float = 1e-3
        device: str = "auto"
        seed: int = 0
        shard_size: int = 1024
        eval_num_tasks: int = 6
        eval_k: int = 4
        success_loss_weight: float = 1.0
        metadata: dict[str, Any] = Field(default_factory=dict)

        def normalized_baseline(self) -> str:
            return normalize_baseline_name(self.baseline)

else:

    @dataclass
    class BaselineConfig:
        baseline: str
        dataset: str | Path
        out: str | Path
        backend: str = "toy"
        epochs: int = 1
        batch_groups: int = 4
        records_per_group: int | None = 8
        hidden_dim: int = 128
        lr: float = 1e-3
        device: str = "auto"
        seed: int = 0
        shard_size: int = 1024
        eval_num_tasks: int = 6
        eval_k: int = 4
        success_loss_weight: float = 1.0
        metadata: dict[str, Any] | None = None

        @classmethod
        def model_validate(cls, payload: Any) -> BaselineConfig:
            if isinstance(payload, cls):
                return payload
            if isinstance(payload, dict):
                return cls(**payload)
            raise TypeError(f"Cannot validate {type(payload).__name__} as BaselineConfig")

        def model_dump(self, **_: Any) -> dict[str, Any]:
            payload = asdict(self)
            payload["metadata"] = dict(self.metadata or {})
            return payload

        def normalized_baseline(self) -> str:
            return normalize_baseline_name(self.baseline)


def list_baselines() -> tuple[str, ...]:
    return BASELINES


def normalize_baseline_name(name: str) -> str:
    normalized = ALIASES.get(name, name)
    if normalized not in BASELINES:
        raise ValueError(f"Unknown baseline {name!r}. Available: {', '.join(BASELINES)}")
    return normalized


def prepare_dataset_for_baseline(
    dataset_dir: str | Path,
    baseline_name: str,
    out_dir: str | Path,
    *,
    shard_size: int = 1024,
    seed: int = 0,
) -> Path:
    baseline = normalize_baseline_name(baseline_name)
    dataset = CILDataset(dataset_dir)
    output_path = ensure_dir(out_dir)
    groups = list(dataset.iter_groups())
    transformed_groups: list[list[CILRecord]] = []
    for group in groups:
        transformed_groups.append(_transform_group(group, baseline=baseline))

    records: list[CILRecord] = []
    for group in transformed_groups:
        records.extend(compute_regret_and_ranks(group))

    metadata = dataset.index.metadata
    manifest = write_cil_shards(
        records,
        output_dir=output_path,
        max_records_per_shard=shard_size,
        dataset_name=f"{metadata.get('dataset_name', 'dovla_cil')}_{baseline}",
        backend=str(metadata.get("backend", "unknown")),
        k=max((len(group) for group in transformed_groups), default=0),
        task_count=int(
            metadata.get("task_count", 0) or len({record.task_id for record in records})
        ),
        seed=seed,
    )
    baseline_metadata = {
        "baseline": baseline,
        "source_dataset": str(dataset_dir),
        "prepared_dataset": str(output_path),
        "approximate": baseline == "label_only_counterfactual",
        "num_groups": manifest.get("num_groups", manifest.get("group_count", 0)),
        "num_records": manifest.get("num_records", manifest.get("record_count", 0)),
        "notes": _baseline_notes(baseline),
    }
    write_json(baseline_metadata, output_path / "baseline_metadata.json")
    return output_path


def loss_weights_for_baseline(
    baseline_name: str, *, success_loss_weight: float = 1.0
) -> InterventionalLossWeights:
    baseline = normalize_baseline_name(baseline_name)
    if baseline in {"expert_only_bc", "more_independent_demos"}:
        return InterventionalLossWeights(
            bc=1.0,
            effect=0.0,
            success=success_loss_weight,
            progress=0.0,
            rank=0.0,
            regret=0.0,
            contrast=0.0,
            lang_pair=0.0,
        )
    if baseline == "world_model_auxiliary":
        return InterventionalLossWeights(
            bc=1.0,
            effect=1.0,
            success=1.0,
            progress=1.0,
            rank=0.0,
            regret=0.0,
            contrast=0.0,
            lang_pair=0.0,
        )
    if baseline == "no_effect_head":
        return InterventionalLossWeights(effect=0.0)
    if baseline == "no_rank_regret":
        return InterventionalLossWeights(rank=0.0, regret=0.0)
    if baseline == "label_only_counterfactual":
        return InterventionalLossWeights(effect=0.0, rank=1.0, regret=0.5)
    if baseline == "cross_state_negatives":
        return InterventionalLossWeights(rank=1.0, regret=0.5)
    return InterventionalLossWeights()


def train_baseline(baseline_config: BaselineConfig | dict[str, Any]) -> dict[str, Any]:
    config = (
        BaselineConfig.model_validate(baseline_config)
        if hasattr(BaselineConfig, "model_validate")
        else BaselineConfig(**dict(baseline_config))
    )
    baseline = config.normalized_baseline()
    out_dir = ensure_dir(config.out)
    prepared_dataset = prepare_dataset_for_baseline(
        config.dataset,
        baseline,
        out_dir / "dataset",
        shard_size=config.shard_size,
        seed=config.seed,
    )
    weights = loss_weights_for_baseline(
        baseline, success_loss_weight=float(config.success_loss_weight)
    )
    train_dir = out_dir / "train"
    train_result = DoVLATrainer(
        TrainerConfig(
            dataset_dir=prepared_dataset,
            output_dir=train_dir,
            epochs=config.epochs,
            batch_groups=config.batch_groups,
            records_per_group=config.records_per_group,
            hidden_dim=config.hidden_dim,
            learning_rate=config.lr,
            device=config.device,
            seed=config.seed,
            losses=weights,
            objective="legacy",
            pair_scope=(
                "cross_state" if baseline == "cross_state_negatives" else "same_state"
            ),
        )
    ).train()
    eval_metrics = evaluate_baseline(
        train_dir / "best.pt",
        backend=config.backend,
        out_path=out_dir / "causalstress.json",
        num_tasks=config.eval_num_tasks,
        k=config.eval_k,
        seed=config.seed,
        device=config.device,
    )
    summary = {
        "baseline": baseline,
        "config": _config_to_dict(config),
        "prepared_dataset": str(prepared_dataset),
        "train_dir": str(train_dir),
        "checkpoint": str(train_dir / "best.pt"),
        "train": train_result,
        "eval": eval_metrics,
        "loss_weights": _loss_weights_to_dict(weights),
    }
    write_json(summary, out_dir / "metrics.json")
    write_json(_config_to_dict(config), out_dir / "baseline_config.json")
    return summary


def evaluate_baseline(
    checkpoint: str | Path,
    *,
    backend: str = "toy",
    out_path: str | Path | None = None,
    num_tasks: int = 6,
    k: int = 4,
    seed: int = 0,
    device: str = "auto",
) -> dict[str, Any]:
    metrics = CausalStressBenchmark(
        CausalStressConfig(backend=backend, num_tasks=num_tasks, k=k, seed=seed)
    ).evaluate(checkpoint, device=device)
    metrics["checkpoint"] = str(checkpoint)
    if out_path is not None:
        write_json(metrics, out_path)
    return metrics


def _transform_group(group: list[CILRecord], *, baseline: str) -> list[CILRecord]:
    if baseline in {"expert_only_bc", "more_independent_demos"}:
        return [_annotate_record(_best_or_expert_record(group), baseline)]
    if baseline == "random_negatives":
        return [_as_random_negative_baseline(record, group) for record in group]
    if baseline == "label_only_counterfactual":
        return [_as_label_only_record(record, group) for record in group]
    if baseline == "cross_state_negatives":
        return [_annotate_record(record, baseline) for record in group]
    return [_annotate_record(record, baseline) for record in group]


def _best_or_expert_record(group: list[CILRecord]) -> CILRecord:
    experts = [record for record in group if record.candidate_type == "expert"]
    candidates = experts or group
    return max(
        candidates,
        key=lambda record: (record.reward.score, -(record.rank_within_group or 0)),
    )


def _as_random_negative_baseline(record: CILRecord, group: list[CILRecord]) -> CILRecord:
    best = _best_or_expert_record(group)
    if record.record_id == best.record_id:
        return _annotate_record(record, "random_negatives")
    action = _replace_action_metadata(
        record.action_chunk,
        {
            "candidate_type": "random_negative",
            "baseline_original_candidate_type": record.candidate_type,
            "baseline_random_negative": True,
        },
    )
    return _annotate_record(
        dataclass_replace(record, action_chunk=action, candidate_type="random_negative"),
        "random_negatives",
    )


def _as_label_only_record(record: CILRecord, group: list[CILRecord]) -> CILRecord:
    best = _best_or_expert_record(group)
    candidate = record.candidate_type
    if record.record_id == best.record_id or candidate == "expert":
        progress, success = 1.0, True
    elif candidate == "near_miss":
        progress, success = 0.45, False
    elif candidate == "alternative_skill":
        progress, success = 0.35, False
    elif candidate in {"wrong_target", "wrong_relation"}:
        progress, success = 0.1, False
    else:
        progress, success = 0.0, False
    reward = RewardInfo(
        progress=progress,
        success=success,
        terminal_success=success,
        dense_components={
            "label_only_heuristic": progress,
            "measured_reward_ignored": record.reward.progress,
        },
    )
    return _annotate_record(
        dataclass_replace(record, reward=reward),
        "label_only_counterfactual",
        approximate=True,
    )


def _annotate_record(
    record: CILRecord, baseline: str, *, approximate: bool = False
) -> CILRecord:
    metadata = {
        **record.metadata,
        "baseline": baseline,
        "baseline_approximate": approximate,
    }
    return dataclass_replace(record, metadata=metadata)


def _replace_action_metadata(action: ActionChunk, metadata: dict[str, Any]) -> ActionChunk:
    return dataclass_replace(action, metadata={**action.metadata, **metadata})


def _baseline_notes(baseline: str) -> str:
    notes = {
        "expert_only_bc": "One best/expert action per group; ranking and regret losses disabled.",
        "more_independent_demos": (
            "K=1-style subset from the source dataset; use a larger source for full comparison."
        ),
        "random_negatives": (
            "Non-expert candidates are relabeled as random negatives with measured outcomes."
        ),
        "cross_state_negatives": (
            "Pairs come from different states of the same task with the same pair budget."
        ),
        "label_only_counterfactual": (
            "Approximate baseline using heuristic rather than measured outcome labels."
        ),
        "world_model_auxiliary": (
            "Effect/progress/success losses enabled; ranking and regret disabled."
        ),
        "no_effect_head": "Effect vector loss disabled.",
        "no_rank_regret": "Ranking and regret losses disabled.",
    }
    return notes[baseline]


def _config_to_dict(config: BaselineConfig) -> dict[str, Any]:
    if hasattr(config, "model_dump"):
        payload = config.model_dump()
    else:
        payload = asdict(config)
    payload["baseline"] = normalize_baseline_name(str(payload["baseline"]))
    payload["dataset"] = str(payload["dataset"])
    payload["out"] = str(payload["out"])
    payload["metadata"] = dict(payload.get("metadata") or {})
    return payload


def _loss_weights_to_dict(weights: InterventionalLossWeights) -> dict[str, float]:
    return {
        "bc": weights.weight("bc"),
        "effect": weights.weight("effect"),
        "success": weights.weight("success"),
        "progress": weights.weight("progress"),
        "rank": weights.weight("rank"),
        "regret": weights.weight("regret"),
        "contrast": weights.weight("contrast"),
        "lang_pair": weights.weight("lang_pair"),
    }