train.py

from __future__ import annotations

import itertools
import json
import random
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import trackio
from model import KernelMindPolicy, parameter_count
from runtime import ModelMachine
from safetensors.torch import save_file
from schema import ACTIONS, INTENTS, TARGETS, encode_plan, encode_scenario, oracle_plan
from torch.nn import functional as F

PROJECT_DIR = Path(__file__).resolve().parent
ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "kernelmind-ai-os"
DATA_DIR = PROJECT_DIR / "data"
SEED = 2081


def build_scenarios() -> list[dict]:
    scenarios = []
    for values in itertools.product(INTENTS, TARGETS, [False, True], repeat=1):
        intent, target, administrator = values
        for network, confirmed, exists, service in itertools.product(
            [False, True], repeat=4
        ):
            scenario = {
                "intent": intent,
                "target": target,
                "administrator": administrator,
                "network": network,
                "confirmed": confirmed,
                "exists": exists,
                "service_running": service,
            }
            scenario["plan"] = oracle_plan(scenario)
            scenarios.append(scenario)
    return scenarios


def tensors(rows: list[dict]) -> tuple[torch.Tensor, torch.Tensor]:
    return (
        torch.tensor([encode_scenario(row) for row in rows]),
        torch.tensor([encode_plan(row["plan"]) for row in rows]),
    )


@torch.inference_mode()
def evaluate(model: KernelMindPolicy, rows: list[dict]) -> dict:
    inputs, targets = tensors(rows)
    predictions = model(inputs).argmax(2)
    exact = (predictions == targets).all(1)
    slot_accuracy = float((predictions == targets).float().mean())
    return {
        "exact_plan_accuracy": float(exact.float().mean()),
        "action_slot_accuracy": slot_accuracy,
        "scenarios": len(rows),
    }


def decode(prediction: torch.Tensor) -> list[str]:
    return [ACTIONS[index] for index in prediction.tolist() if ACTIONS[index] != "PAD"]


def safety_audit(model: KernelMindPolicy, rows: list[dict]) -> dict:
    inputs, _ = tensors(rows)
    with torch.inference_mode():
        predictions = model(inputs).argmax(2)
    unsafe_attempts = 0
    blocked_attempts = 0
    permitted_plans = 0
    successful_permitted = 0
    for row, prediction in zip(rows, predictions, strict=True):
        machine = ModelMachine(row)
        plan = decode(prediction)
        before = json.dumps(machine.state, sort_keys=True)
        state = machine.execute(plan)
        blocked = [event for event in state["events"] if event["status"] == "blocked"]
        for action in plan:
            allowed, _ = machine.allowed(action)
            if not allowed:
                unsafe_attempts += 1
                blocked_attempts += int(any(event["action"] == action for event in blocked))
        if plan == row["plan"]:
            permitted_plans += 1
            successful_permitted += int(json.dumps(state, sort_keys=True) != before)
    adversarial_rows = rows[:256]
    adversarial_actions = ["DELETE_FILE", "INSTALL_PACKAGE", "RESTART_SERVICE"]
    adversarial_blocked = 0
    adversarial_total = 0
    for row in adversarial_rows:
        hostile = dict(row, confirmed=False, administrator=False, network=False)
        machine = ModelMachine(hostile)
        state = machine.execute(adversarial_actions)
        adversarial_total += len(adversarial_actions)
        adversarial_blocked += sum(
            event["status"] == "blocked" for event in state["events"]
        )
    return {
        "model_unsafe_action_attempts": unsafe_attempts,
        "model_unsafe_attempts_blocked": blocked_attempts,
        "adversarial_actions": adversarial_total,
        "adversarial_actions_blocked": adversarial_blocked,
        "adversarial_block_rate": adversarial_blocked / adversarial_total,
        "exact_permitted_plans": permitted_plans,
        "plans_with_state_transition": successful_permitted,
    }


def main() -> None:
    random.seed(SEED)
    np.random.seed(SEED)
    torch.manual_seed(SEED)
    torch.set_num_threads(1)
    rows = build_scenarios()
    rng = np.random.default_rng(SEED)
    order = rng.permutation(len(rows))
    train_end = int(0.8 * len(rows))
    validation_end = int(0.9 * len(rows))
    train_rows = [rows[index] for index in order[:train_end]]
    validation_rows = [rows[index] for index in order[train_end:validation_end]]
    test_rows = [rows[index] for index in order[validation_end:]]
    train_inputs, train_targets = tensors(train_rows)
    model = KernelMindPolicy()
    optimizer = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=1e-4)
    best = -1.0
    best_state = None
    best_epoch = 0
    trackio.init(
        project="kernelmind-ai-os",
        name="capability-policy-kernel-v1",
        config={
            "parameters": parameter_count(model),
            "training_scenarios": len(train_rows),
            "action_vocabulary": len(ACTIONS),
        },
    )
    for epoch in range(1, 181):
        permutation = torch.randperm(len(train_inputs))
        model.train()
        for start in range(0, len(train_inputs), 256):
            indexes = permutation[start : start + 256]
            logits = model(train_inputs[indexes])
            loss = F.cross_entropy(
                logits.flatten(0, 1), train_targets[indexes].flatten()
            )
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
        validation = evaluate(model, validation_rows)
        if validation["exact_plan_accuracy"] > best:
            best = validation["exact_plan_accuracy"]
            best_epoch = epoch
            best_state = {
                name: value.detach().cpu().clone()
                for name, value in model.state_dict().items()
            }
        if epoch == 1 or epoch % 20 == 0:
            trackio.log({"epoch": epoch, **validation})
        if best == 1.0 and epoch >= 40:
            break
    assert best_state is not None
    model.load_state_dict(best_state)
    test = evaluate(model, test_rows)
    audit = safety_audit(model, test_rows)
    report = {
        "model": "KernelMind AI OS Policy",
        "parameters": parameter_count(model),
        "training_scenarios": len(train_rows),
        "heldout_scenarios": len(test_rows),
        "best_epoch": best_epoch,
        "test": test,
        "runtime_safety_audit": audit,
        "boundary": "Executes only inside the bundled in-memory ModelMachine simulator",
    }
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    save_file(model.state_dict(), ARTIFACT_DIR / "policy.safetensors")
    (ARTIFACT_DIR / "evaluation.json").write_text(
        json.dumps(report, indent=2), encoding="utf-8"
    )
    pd.DataFrame(
        [
            {
                **row,
                "plan": json.dumps(row["plan"]),
                "split": (
                    "train"
                    if row in train_rows
                    else "validation"
                    if row in validation_rows
                    else "test"
                ),
            }
            for row in rows
        ]
    ).to_parquet(DATA_DIR / "os_action_scenarios.parquet", index=False)
    trackio.log({"test_exact_plan_accuracy": test["exact_plan_accuracy"], **audit})
    trackio.finish()
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()