from __future__ import annotations import argparse import json import sys import time from dataclasses import asdict from pathlib import Path from types import SimpleNamespace from typing import Any import numpy as np import pandas as pd import torch SCRIPT_DIR = Path(__file__).resolve().parent ROOT_DIR = SCRIPT_DIR.parents[1] V4P4_SCRIPT_DIR = ROOT_DIR / "v4p4_world_model" / "scripts" V3P5_SCRIPT_DIR = ROOT_DIR / "v3p5_static" / "scripts" for path in (SCRIPT_DIR, V4P4_SCRIPT_DIR, V3P5_SCRIPT_DIR): if str(path) not in sys.path: sys.path.insert(0, str(path)) from action_ontology_v5 import build_action_ontology, load_json # noqa: E402 from build_v5_action_tensors import ( # noqa: E402 build_action_arrays, build_target_trial_labels, medication_flags_from_stage0, resolve_split_path, validate_medication_flags, ) from config_v5_train import DEFAULT_STAGE0_DIR, DEFAULT_TENSOR_DIR, V5_LOSS_WEIGHTS # noqa: E402 from loss_v5 import sctm_v5_loss # noqa: E402 from model_v5 import SCTMv5, SCTMv5Config # noqa: E402 from smoke_v3p4_architecture import build_field_value_mask, load_service_prior # noqa: E402 from train_v4p4_cloud import TrainIndexSampler, autocast_context, batch_from_indices, load_npz_to_memory, metadata_config # noqa: E402 def tensorize_actions(action_arrays: dict[str, np.ndarray], idx: np.ndarray, device: torch.device) -> dict[str, torch.Tensor]: return { "action_type_ids": torch.as_tensor(action_arrays["action_type_ids"][idx], dtype=torch.long, device=device), "action_value_ids": torch.as_tensor(action_arrays["action_value_ids"][idx], dtype=torch.long, device=device), "action_mask": torch.as_tensor(action_arrays["action_mask"][idx], dtype=torch.bool, device=device), "action_available_at": torch.as_tensor(action_arrays["action_available_at"][idx], dtype=torch.long, device=device), } def validate_action_arrays(action_arrays: dict[str, np.ndarray], arrays: dict[str, np.ndarray], ontology: dict[str, Any], *, source: Path | str) -> None: required = ("action_type_ids", "action_value_ids", "action_mask", "action_available_at") missing = [key for key in required if key not in action_arrays] if missing: raise ValueError(f"Action tensor source {source} is missing required arrays: {missing}") n_windows, seq_len = arrays["valid_mask"].shape n_slots = len(ontology["slots"]) for key in required: shape = tuple(action_arrays[key].shape) expected = (n_windows, seq_len, n_slots) if shape != expected: raise ValueError(f"Action tensor source {source} has {key} shape {shape}, expected {expected}") if not np.array_equal(action_arrays["action_mask"].astype(bool), action_arrays["action_mask"]): raise ValueError(f"Action tensor source {source} has non-boolean-compatible action_mask values") def metrics_to_float(metrics: dict[str, torch.Tensor]) -> dict[str, float]: return {key: float(value.detach().cpu()) for key, value in metrics.items()} def aggregate_eval_metrics(rows: list[dict[str, float]]) -> dict[str, float]: if not rows: return {} keys = sorted(set().union(*(row.keys() for row in rows))) weights = np.asarray([max(0.0, float(row.get("valid_next_positions", 0.0))) for row in rows], dtype=np.float64) weight_sum = float(weights.sum()) out: dict[str, float] = {} for key in keys: values = np.asarray([float(row[key]) for row in rows if key in row], dtype=np.float64) if values.size == 0: continue if key in {"valid_next_positions", "next_contact_positions"}: out[key] = float(values.sum()) continue if key.startswith("loss_") and weight_sum > 0 and values.size == len(rows): out[key] = float(np.sum(values * weights) / weight_sum) else: out[key] = float(values.mean()) return out def load_split_bundle( tensor_dir: Path, stage0_dir: Path, split: str, max_windows: int = 0, action_dir: Path | None = None, ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray], dict[str, Any], dict[str, Any]]: meta = load_json(tensor_dir / "tensor_metadata.json") vocab = load_json(tensor_dir / "cat_value_vocab.json") ordinal_direction = load_json(stage0_dir / "ordinal_direction_table.json") shard = resolve_split_path(tensor_dir, split) arrays = load_npz_to_memory(shard) if max_windows > 0: arrays = {k: v[:max_windows] for k, v in arrays.items()} ontology = build_action_ontology(meta, vocab) medication_flags = medication_flags_from_stage0(arrays, stage0_dir) action_path = action_dir / f"v5_{split}_action_tensors.npz" if action_dir is not None else None if action_path is not None and action_path.exists(): with np.load(action_path, allow_pickle=False) as z: action_keys = {"action_type_ids", "action_value_ids", "action_mask", "action_available_at", "lai_present", "oral_present"} action_arrays = {k: z[k][:max_windows] if max_windows > 0 else z[k] for k in z.files if k in action_keys} validate_action_arrays(action_arrays, arrays, ontology, source=action_path) validate_medication_flags(action_arrays, medication_flags, source=action_path) else: action_arrays = build_action_arrays(arrays, meta, ordinal_direction, vocab, ontology, medication_flags=medication_flags) _ = build_target_trial_labels(arrays, (30.0, 60.0, 90.0), medication_flags=medication_flags, stage0_dir=stage0_dir) validate_action_arrays(action_arrays, arrays, ontology, source="built_in_memory") validate_medication_flags(action_arrays, medication_flags, source="built_in_memory") return arrays, action_arrays, meta, ontology def make_adamw( parameters: Any, *, lr: float, weight_decay: float, fused: bool, device: torch.device, ) -> torch.optim.Optimizer: kwargs: dict[str, Any] = {"lr": lr, "weight_decay": weight_decay} if fused and device.type == "cuda": try: return torch.optim.AdamW(parameters, fused=True, **kwargs) except TypeError: pass return torch.optim.AdamW(parameters, **kwargs) def _extract_model_state_dict(checkpoint: Any) -> tuple[dict[str, torch.Tensor], str]: if isinstance(checkpoint, dict): for key in ("model", "model_state_dict", "state_dict"): state = checkpoint.get(key) if isinstance(state, dict): return state, key if isinstance(checkpoint, dict) and all(isinstance(v, torch.Tensor) for v in checkpoint.values()): return checkpoint, "root" raise ValueError("Checkpoint does not contain a model state dict under model/model_state_dict/state_dict or root tensor keys.") def _normalize_state_dict_keys(state: dict[str, torch.Tensor], model_keys: set[str]) -> tuple[dict[str, torch.Tensor], str]: if any(key in model_keys for key in state): return state, "as_is" for prefix in ("module.", "_orig_mod."): stripped = {key.removeprefix(prefix): value for key, value in state.items() if key.startswith(prefix)} if stripped and any(key in model_keys for key in stripped): return stripped, f"strip_{prefix}" return state, "as_is_no_direct_match" def load_checkpoint_weights(model: torch.nn.Module, checkpoint_path: Path, *, strict: bool) -> dict[str, Any]: checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) raw_state, source_key = _extract_model_state_dict(checkpoint) model_keys = set(model.state_dict().keys()) state, key_transform = _normalize_state_dict_keys(raw_state, model_keys) loadable_keys = sorted(key for key in state if key in model_keys and tuple(state[key].shape) == tuple(model.state_dict()[key].shape)) if not strict: state = {key: value for key, value in state.items() if key in loadable_keys} if not state: raise RuntimeError(f"Checkpoint {checkpoint_path} did not match any model parameters; refusing to continue.") result = model.load_state_dict(state, strict=strict) return { "path": str(checkpoint_path), "strict": bool(strict), "source_key": source_key, "key_transform": key_transform, "checkpoint_parameter_keys": int(len(raw_state)), "loadable_parameter_keys": int(len(loadable_keys)), "missing_keys": list(result.missing_keys), "unexpected_keys": list(result.unexpected_keys), } def set_base_trainable(model: torch.nn.Module, trainable: bool) -> dict[str, int]: action_prefixes = ( "action_encoder.", "action_pwe_delta_head.", "action_active_delta_head.", "action_missing_delta_head.", "action_field_delta_head.", "action_numeric_delta_head.", "action_ordinal_delta_head.", "action_event_delta_head.", "behavior_policy_head.", ) frozen = 0 trainable_count = 0 for name, param in model.named_parameters(): if trainable or name.startswith(action_prefixes): param.requires_grad_(True) trainable_count += param.numel() else: param.requires_grad_(False) frozen += param.numel() return {"trainable_parameters": int(trainable_count), "frozen_parameters": int(frozen)} def checkpoint_state( model: torch.nn.Module, config: SCTMv5Config, ontology: dict[str, Any], args: argparse.Namespace, prior_audit: dict[str, Any], *, step: int | None = None, best_step: int | None = None, best_val_loss: float | None = None, current_val_loss: float | None = None, ) -> dict[str, Any]: raw_model = getattr(model, "_orig_mod", model) payload = { "model": raw_model.state_dict(), "config": asdict(config), "ontology": ontology, "args": vars(args), "service_prior_audit": prior_audit, } if step is not None: payload["step"] = int(step) if best_step is not None: payload["best_step"] = int(best_step) if best_val_loss is not None: payload["best_val_loss"] = float(best_val_loss) if current_val_loss is not None: payload["current_val_loss"] = float(current_val_loss) return payload @torch.inference_mode() def evaluate(model: SCTMv5, arrays: dict[str, np.ndarray], action_arrays: dict[str, np.ndarray], config: SCTMv5Config, args: argparse.Namespace, device: torch.device) -> dict[str, float]: model.eval() n = arrays["valid_mask"].shape[0] n_eval = min(n, args.eval_windows) if args.eval_windows > 0 else n losses = [] for start in range(0, n_eval, args.eval_batch_size): idx = np.arange(start, min(start + args.eval_batch_size, n_eval)) batch = batch_from_indices(arrays, idx, device) batch.update(tensorize_actions(action_arrays, idx, device)) with autocast_context(device, args.precision): out = model(batch, rollout_steps=1, compute_pwe_diagnostics=False) _, metrics = sctm_v5_loss(out, batch, config, return_metrics=True) losses.append(metrics_to_float(metrics)) return aggregate_eval_metrics(losses) def main() -> None: parser = argparse.ArgumentParser(description="Train SCTM-v5 action-conditioned heads on observed-action likelihood.") parser.add_argument("--tensor-dir", type=Path, default=DEFAULT_TENSOR_DIR) parser.add_argument("--stage0-dir", type=Path, default=DEFAULT_STAGE0_DIR) parser.add_argument("--out-dir", type=Path, required=True) parser.add_argument("--train-split", default="train") parser.add_argument("--val-split", default="val") parser.add_argument("--action-dir", type=Path, default=None, help="Optional directory containing precomputed v5_{split}_action_tensors.npz files.") parser.add_argument("--max-windows", type=int, default=0) parser.add_argument("--max-steps", type=int, default=100) parser.add_argument("--batch-size", type=int, default=8) parser.add_argument("--eval-batch-size", type=int, default=16) parser.add_argument("--log-every", type=int, default=25) parser.add_argument("--eval-every", type=int, default=500) parser.add_argument("--eval-windows", type=int, default=64) parser.add_argument("--save-every", type=int, default=1000) parser.add_argument("--save-eval-checkpoints", action="store_true", help="Save a checkpoint at every validation point for post-hoc model selection audits.") parser.add_argument("--top-k-checkpoints", type=int, default=0, help="If >0, keep only the top-k eval checkpoints by validation loss in eval_checkpoints/.") parser.add_argument("--early-stopping-patience", type=int, default=0, help="If >0, stop after this many consecutive validation checks without improvement.") parser.add_argument("--early-stopping-min-delta", type=float, default=0.0, help="Minimum val loss improvement required to reset early-stopping patience.") parser.add_argument("--min-steps-before-stopping", type=int, default=0, help="Do not trigger early stopping before this step.") parser.add_argument("--learning-rate", type=float, default=1.0e-4) parser.add_argument("--weight-decay", type=float, default=1.0e-4) parser.add_argument("--grad-clip", type=float, default=1.0) parser.add_argument("--precision", choices=("fp32", "bf16", "fp16"), default="fp32") parser.add_argument("--device", default="cpu") parser.add_argument("--fused-adamw", action="store_true") parser.add_argument("--compile", action="store_true") parser.add_argument("--matmul-precision", choices=("highest", "high", "medium"), default=None) parser.add_argument("--d-model", type=int, default=64) parser.add_argument("--n-heads", type=int, default=4) parser.add_argument("--n-layers", type=int, default=2) parser.add_argument("--dropout", type=float, default=0.1) parser.add_argument("--seed", type=int, default=20260525) parser.add_argument("--init-checkpoint", type=Path, default=None, help="Optional v4/v5 checkpoint used to initialize matching model weights with strict=False.") parser.add_argument("--resume-checkpoint", type=Path, default=None, help="Optional v5 checkpoint used to resume with strict=True.") parser.add_argument("--freeze-base-steps", type=int, default=0, help="If >0, train only v5 action/behavior heads for this many initial steps, then unfreeze all parameters.") for key, value in V5_LOSS_WEIGHTS.items(): parser.add_argument(f"--{key.replace('_', '-')}", type=float, default=float(value)) args = parser.parse_args() args.out_dir.mkdir(parents=True, exist_ok=True) device = torch.device(args.device) if args.matmul_precision: torch.set_float32_matmul_precision(args.matmul_precision) rng = np.random.default_rng(args.seed) torch.manual_seed(args.seed) train_arrays, train_actions, meta, ontology = load_split_bundle(args.tensor_dir, args.stage0_dir, args.train_split, args.max_windows, args.action_dir) val_arrays, val_actions, _, _ = load_split_bundle(args.tensor_dir, args.stage0_dir, args.val_split, args.max_windows, args.action_dir) vocab = load_json(args.tensor_dir / "cat_value_vocab.json") base_cfg = metadata_config(meta, vocab, SimpleNamespace(d_model=args.d_model, n_heads=args.n_heads, n_layers=args.n_layers, dropout=args.dropout, n_missing=5)) cfg_dict = asdict(base_cfg) cfg_dict.update( { "n_action_slots": len(ontology["slots"]), "action_value_vocab_size": len(ontology["action_value_vocab"]), "n_action_availability": len(ontology["action_availability_vocab"]), } ) config = SCTMv5Config(**cfg_dict) field_value_mask, _ = build_field_value_mask(meta, vocab, {k: train_arrays[k] for k in ["cat_value_ids", "missing_ids"]}) prior, prior_audit = load_service_prior(args.stage0_dir, int(meta.get("n_service_states", 8)), "service_state_transitions_train.json") model = SCTMv5(config, field_value_mask=field_value_mask, service_prior_bias=prior).to(device) checkpoint_audit: dict[str, Any] | None = None if args.resume_checkpoint is not None: checkpoint_audit = load_checkpoint_weights(model, args.resume_checkpoint, strict=True) elif args.init_checkpoint is not None: checkpoint_audit = load_checkpoint_weights(model, args.init_checkpoint, strict=False) freeze_audit = set_base_trainable(model, trainable=args.freeze_base_steps <= 0) if args.compile: model = torch.compile(model) optimizer = make_adamw(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay, fused=args.fused_adamw, device=device) sampler = TrainIndexSampler(rng, train_arrays["valid_mask"].shape[0], sample_with_replacement=False) loss_kwargs = {key: getattr(args, key) for key in V5_LOSS_WEIGHTS} log_rows: list[dict[str, Any]] = [] best_val = float("inf") best_step = 0 completed_steps = 0 no_improve_evals = 0 eval_checkpoint_records: list[tuple[float, int, Path]] = [] start_time = time.time() for step in range(1, args.max_steps + 1): completed_steps = step if args.freeze_base_steps > 0 and step == args.freeze_base_steps + 1: raw_model = getattr(model, "_orig_mod", model) freeze_audit = set_base_trainable(raw_model, trainable=True) optimizer = make_adamw(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay, fused=args.fused_adamw, device=device) model.train() idx = sampler.next(args.batch_size) batch = batch_from_indices(train_arrays, idx, device) batch.update(tensorize_actions(train_actions, idx, device)) optimizer.zero_grad(set_to_none=True) need_log = step == 1 or step % args.log_every == 0 or step % args.eval_every == 0 or step == args.max_steps need_eval = step == 1 or step % args.eval_every == 0 or step == args.max_steps with autocast_context(device, args.precision): out = model(batch, rollout_steps=1, compute_pwe_diagnostics=False) loss, metrics = sctm_v5_loss(out, batch, config, **loss_kwargs, return_metrics=need_log) loss.backward() grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) optimizer.step() row = { "step": step, "elapsed_sec": round(time.time() - start_time, 3), "loss": float(loss.detach().cpu()) if need_log else float("nan"), "grad_norm": float(grad_norm.detach().cpu()) if need_log else float("nan"), } if need_log: row.update({f"train_{k}": v for k, v in metrics_to_float(metrics).items()}) if need_eval: val = evaluate(model, val_arrays, val_actions, config, args, device) row.update({f"val_{k}": v for k, v in val.items()}) row.update({f"freeze_{k}": v for k, v in freeze_audit.items()}) val_loss = float(val.get("loss_total", float("inf"))) improved = val_loss < (best_val - float(args.early_stopping_min_delta)) if improved: best_val = val_loss best_step = step no_improve_evals = 0 torch.save( checkpoint_state( model, config, ontology, args, prior_audit, step=step, best_step=best_step, best_val_loss=best_val, current_val_loss=val_loss, ), args.out_dir / "best.pt", ) else: no_improve_evals += 1 if args.save_eval_checkpoints or args.top_k_checkpoints > 0: eval_dir = args.out_dir / "eval_checkpoints" eval_dir.mkdir(parents=True, exist_ok=True) eval_path = eval_dir / f"step_{step:06d}_val_{val_loss:.6f}.pt" torch.save( checkpoint_state( model, config, ontology, args, prior_audit, step=step, best_step=best_step, best_val_loss=best_val, current_val_loss=val_loss, ), eval_path, ) eval_checkpoint_records.append((val_loss, step, eval_path)) if args.top_k_checkpoints > 0: eval_checkpoint_records = sorted(eval_checkpoint_records, key=lambda x: (x[0], x[1])) for _, _, stale_path in eval_checkpoint_records[args.top_k_checkpoints :]: stale_path.unlink(missing_ok=True) eval_checkpoint_records = eval_checkpoint_records[: args.top_k_checkpoints] row["best_step"] = best_step row["best_val_loss"] = best_val row["early_stopping_no_improve_evals"] = no_improve_evals should_stop = ( args.early_stopping_patience > 0 and step >= args.min_steps_before_stopping and no_improve_evals >= args.early_stopping_patience ) row["early_stopped"] = bool(should_stop) print(json.dumps(row, ensure_ascii=False), flush=True) log_rows.append(row) if need_eval and row.get("early_stopped"): break if args.save_every > 0 and step % args.save_every == 0: torch.save( checkpoint_state( model, config, ontology, args, prior_audit, step=step, best_step=best_step, best_val_loss=best_val, ), args.out_dir / "latest.pt", ) pd.DataFrame(log_rows).to_csv(args.out_dir / "train_log.csv", index=False) torch.save( checkpoint_state( model, config, ontology, args, prior_audit, step=completed_steps, best_step=best_step, best_val_loss=best_val, ), args.out_dir / "latest.pt", ) summary = { "status": "complete", "out_dir": str(args.out_dir), "steps": completed_steps, "planned_steps": args.max_steps, "train_rows": len(log_rows), "best_step": best_step, "best_val_loss": best_val, "early_stopped": bool(completed_steps < args.max_steps), "checkpoint_audit": checkpoint_audit, "freeze_audit": freeze_audit, "eval_checkpoint_count": len(eval_checkpoint_records), } (args.out_dir / "train_result.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8") print(json.dumps(summary, ensure_ascii=False), flush=True) if __name__ == "__main__": main()