train.py

from __future__ import annotations

import itertools
import json
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import trackio
from model import NeuralModelMachine, parameter_count
from safetensors.torch import save_file
from schema import ACTIONS, CAPABILITY_NAMES, STATE_NAMES, reference_transition
from torch.nn import functional as F

PROJECT_DIR = Path(__file__).resolve().parent
ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "kernelmind-model-machine"
DATA_DIR = PROJECT_DIR / "data"
SEED = 2087


def build_rows() -> list[dict]:
    rows = []
    for state in itertools.product([0, 1], repeat=8):
        for capabilities in itertools.product([0, 1], repeat=4):
            for action_index, action in enumerate(ACTIONS):
                next_state, blocked = reference_transition(
                    list(state), list(capabilities), action
                )
                rows.append(
                    {
                        "state": list(state),
                        "capabilities": list(capabilities),
                        "action": action,
                        "action_index": action_index,
                        "next_state": next_state,
                        "blocked": blocked,
                    }
                )
    return rows


def tensors(rows: list[dict]) -> tuple[torch.Tensor, ...]:
    return (
        torch.tensor([row["state"] for row in rows], dtype=torch.float32),
        torch.tensor([row["capabilities"] for row in rows], dtype=torch.float32),
        torch.tensor([row["action_index"] for row in rows]),
        torch.tensor(
            [row["next_state"] + [row["blocked"]] for row in rows],
            dtype=torch.float32,
        ),
    )


@torch.inference_mode()
def evaluate(model: NeuralModelMachine, rows: list[dict]) -> dict:
    state, capabilities, action, targets = tensors(rows)
    probabilities = torch.sigmoid(model(state, capabilities, action))
    predictions = (probabilities >= 0.5).float()
    exact = (predictions == targets).all(1)
    return {
        "exact_transition_accuracy": float(exact.float().mean()),
        "state_bit_accuracy": float(
            (predictions[:, :8] == targets[:, :8]).float().mean()
        ),
        "blocked_accuracy": float(
            (predictions[:, 8] == targets[:, 8]).float().mean()
        ),
        "transitions": len(rows),
    }


@torch.inference_mode()
def rollout_audit(model: NeuralModelMachine, samples: int = 2_000) -> dict:
    rng = np.random.default_rng(SEED + 1)
    exact_steps = 0
    exact_rollouts = 0
    for _ in range(samples):
        state = rng.integers(0, 2, size=8).tolist()
        capabilities = rng.integers(0, 2, size=4).tolist()
        learned = torch.tensor([state], dtype=torch.float32)
        reference = state
        rollout_exact = True
        for action_index in rng.integers(0, len(ACTIONS), size=6):
            action = ACTIONS[action_index]
            reference, reference_blocked = reference_transition(
                reference, capabilities, action
            )
            predicted, blocked_probability = model.transition(
                learned,
                torch.tensor([capabilities], dtype=torch.float32),
                torch.tensor([action_index]),
            )
            learned = predicted
            step_exact = learned[0].int().tolist() == reference
            step_exact = step_exact and (
                int(blocked_probability[0] >= 0.5) == reference_blocked
            )
            exact_steps += int(step_exact)
            rollout_exact = rollout_exact and step_exact
        exact_rollouts += int(rollout_exact)
    return {
        "rollouts": samples,
        "steps_per_rollout": 6,
        "exact_step_fraction": exact_steps / (samples * 6),
        "exact_six_step_rollout_fraction": exact_rollouts / samples,
    }


def main() -> None:
    torch.manual_seed(SEED)
    torch.set_num_threads(1)
    rows = build_rows()
    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:]]
    state, capabilities, action, targets = tensors(train_rows)
    model = NeuralModelMachine()
    optimizer = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=1e-5)
    best = -1.0
    best_state = None
    best_epoch = 0
    trackio.init(
        project="kernelmind-model-machine",
        name="neural-os-transition-runtime-v1",
        config={
            "parameters": parameter_count(model),
            "training_transitions": len(train_rows),
            "state_bits": len(STATE_NAMES),
            "capability_bits": len(CAPABILITY_NAMES),
        },
    )
    for epoch in range(1, 81):
        permutation = torch.randperm(len(state))
        model.train()
        for start in range(0, len(state), 512):
            indexes = permutation[start : start + 512]
            logits = model(state[indexes], capabilities[indexes], action[indexes])
            loss = F.binary_cross_entropy_with_logits(logits, targets[indexes])
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
        validation = evaluate(model, validation_rows)
        if validation["exact_transition_accuracy"] > best:
            best = validation["exact_transition_accuracy"]
            best_epoch = epoch
            best_state = {
                name: value.detach().cpu().clone()
                for name, value in model.state_dict().items()
            }
        if epoch == 1 or epoch % 10 == 0:
            trackio.log({"epoch": epoch, **validation})
        if best == 1.0 and epoch >= 20:
            break
    assert best_state is not None
    model.load_state_dict(best_state)
    test = evaluate(model, test_rows)
    rollouts = rollout_audit(model)
    report = {
        "model": "KernelMind Neural Model Machine",
        "parameters": parameter_count(model),
        "training_transitions": len(train_rows),
        "heldout_transitions": len(test_rows),
        "best_epoch": best_epoch,
        "test": test,
        "multi_step_audit": rollouts,
        "boundary": "Models an eight-bit virtual state; it does not control the host OS",
    }
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    save_file(model.state_dict(), ARTIFACT_DIR / "model_machine.safetensors")
    (ARTIFACT_DIR / "evaluation.json").write_text(
        json.dumps(report, indent=2), encoding="utf-8"
    )
    pd.DataFrame(
        {
            "state": [row["state"] for row in rows],
            "capabilities": [row["capabilities"] for row in rows],
            "action": [row["action"] for row in rows],
            "next_state": [row["next_state"] for row in rows],
            "blocked": [row["blocked"] for row in rows],
        }
    ).to_parquet(DATA_DIR / "transitions.parquet", index=False)
    trackio.log({**test, **rollouts})
    trackio.finish()
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()