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()