Initial release of STOK-SpatialNet-v3: model weights, config, and inference

.gitattributes

@@ -0,0 +1 @@
+*.safetensors filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,109 @@
+---
+language:
+- en
+license: mit
+library_name: generic
+tags:
+- reinforcement-learning
+- reward-free-rl
+- state-time-option-kernels
+- stok
+- game-ai
+- spatial-attention
+- theory-of-mind
+- empowerment
+- simultaneous-turn
+pipeline_tag: tabular-classification
+---
+
+# STOK-SpatialNet-v3: A Reward-Free Spatial-Attention Amortizer
+
+**STOK-SpatialNet-v3** is a compact, multi-task spatial neural network for **reward-free reinforcement learning and simultaneous-turn planning**. 
+
+Instead of regressing a single scalar reward or value function, it computes the quantities of the **State-Time Option Kernel (STOK)** framework ([Ringstrom & Schrater, 2025](https://arxiv.org/abs/2506.09499)):
+- **Feasibility ($\hat{\kappa}$)**: Probability of reaching a target goal without constraint violation.
+- **Time Urgency ($\hat{\eta}$)**: Arrival time distribution $1 / (1 + \mathbb{E}[t_f])$.
+- **Valence ($\Delta\hat{E}$)**: Signed change in action-space empowerment (preserving future degrees of freedom).
+- **Survivability ($\hat{\text{lose}}$)**: Predicted adversarial lose rate.
+- **Theory of Mind**: Softmax distributions predicting the opponent's active tactical concept and goal family.
+
+---
+
+## Architecture: `spatial-attn2`
+
+The network operates on a 2D spatial field with relational goal conditioning:
+
+```text
+Board Cells (T tiles x 39 ch) ──► Embed (39 → 32) + ReLU
+                                         │
+                                         ▼
+                         Tile ⟷ Tile Self-Attention (Residual)
+                                         │
+                                         ▼
+Relational Goal (8 dims) ───────► Goal-Query Cross-Attention
+                                         │
+                                         ▼
+Context Vector (32) ─────┐
+High-level Features (80) ┼──────► Dense Trunk (138 → 96) + ReLU
+Intent History (18) ─────┘               │
+                                         ▼
+                               37 Output Heads (96 → 37)
+```
+
+### Input Channels (39 per tile)
+1. **Geometry**: Normalized absolute coordinates, viewer- and opponent-relative offsets, Chebyshev distances.
+2. **Occupancy & Stats**: Self/ally/enemy presence, normalized HP.
+3. **Terrain & Content**: Substrate, walls, rocks, hazardous elements, blocking flags.
+4. **Path-Aware Reachability**: BFS arrival steps for both agents under dynamic obstacle rules.
+5. **Plan Grounding**: Footprints of candidate movement, strikes, and terrain alterations.
+
+### 37 Multi-Task Output Heads
+- **Heads 0–3 (The Core Quartet)**: Feasibility $\hat{\kappa}$, urgency $\hat{\eta}$, signed valence $\Delta\hat{E}$, survivability $\hat{\text{lose}}$.
+- **Heads 4–8 (Goal-Setting Feasibilities)**: Predicted feasibility for 5 goal families: `Strike`, `Engage`, `Evade`, `Turtle`, `Zone`.
+- **Heads 9–11 (Arrival ETA Histogram)**: Softmax distribution over $[\text{kill@t1}, \text{kill@t2}, \text{not in window}]$.
+- **Heads 12–16 (Opponent Goal Prior)**: Predicted distribution over which goal the opponent will execute this turn.
+- **Heads 17–18 (Decomposed Valence)**: Own option-capacity growth vs denial of opponent option-capacity.
+- **Heads 19–26 (Tactical Concept Tagging)**: Distribution over 8 human-annotated tactical concepts (`trap`, `spacing`, `punish`, `bait`, `coverage`, `conversion`, `denial`, `tempo`).
+- **Heads 27–34 (Theory of Mind)**: Predicted tactical concept the **adversary** is currently playing.
+- **Heads 35–36 (Trajectory Valence)**: Long-horizon realized empowerment flow ($V_{\text{traj}}$, $V_{\text{past}}$).
+
+---
+
+## Quickstart (Zero External Dependencies)
+
+Requires only standard `numpy`. No PyTorch or C++ runtime required.
+
+```python
+import numpy as np
+from modeling_stok import STOKSpatialAmortizer
+
+# Load model from local files or Hugging Face Hub
+model = STOKSpatialAmortizer.from_pretrained(".")
+
+# Prepare inputs for a 10x10 grid (100 tiles x 39 channels)
+spatial_field = np.zeros((100, 39), dtype=np.float32)
+goal_vector = np.zeros(8, dtype=np.float32)      # 5-way one-hot + (dx, dz, has_target)
+goal_vector[0] = 1.0                             # Strike goal
+flat_features = np.zeros(106, dtype=np.float32)  # High-level features + goal + history
+
+# Forward inference
+pred = model.forward(spatial_field, goal_vector, flat_features)
+
+print(f"Goal Feasibility (kappa):   {pred.kappa:.4f}")
+print(f"Empowerment Valence (dE):   {pred.valence:.4f}")
+print(f"Expected time-to-goal:      {pred.expected_time():.2f} turns")
+print(f"Opponent Intent (Argmax):   {max(pred.opp_concept.items(), key=lambda kv: kv[1])}")
+```
+
+---
+
+## Citation & References
+
+```bibtex
+@article{ringstrom2025stok,
+  title   = {A Unified Theory of Compositionality, Modularity, and Interpretability in Markov Decision Processes},
+  author  = {Ringstrom, Tyler and Schrater, Paul},
+  journal = {arXiv preprint arXiv:2506.09499},
+  year    = {2025}
+}
+```

config.json

@@ -0,0 +1,60 @@
+{
+  "architectures": [
+    "STOKSpatialAmortizer"
+  ],
+  "model_type": "stok_spatial_amortizer",
+  "spatial_channels": 39,
+  "attn_dim": 32,
+  "goal_dim": 8,
+  "flat_dim": 106,
+  "hidden_dim": 96,
+  "num_heads": 37,
+  "config_digest": "daf428ff935666b8",
+  "goal_families": [
+    "Strike",
+    "Engage",
+    "Evade",
+    "Turtle",
+    "Zone"
+  ],
+  "tactical_concepts": [
+    "trap",
+    "spacing",
+    "punish",
+    "bait",
+    "coverage",
+    "conversion",
+    "denial",
+    "tempo"
+  ],
+  "head_slices": {
+    "kappa": 0,
+    "urgency": 1,
+    "valence": 2,
+    "lose": 3,
+    "goal_kappa": [
+      4,
+      9
+    ],
+    "eta": [
+      9,
+      12
+    ],
+    "opp_prior": [
+      12,
+      17
+    ],
+    "valence_own": 17,
+    "valence_deny": 18,
+    "concept": [
+      19,
+      27
+    ],
+    "opp_concept": [
+      27,
+      35
+    ],
+    "v_traj": 35,
+    "v_past": 36
+  }
+}

modeling_stok.py

@@ -0,0 +1,197 @@
+"""
+STOK-SpatialNet-v3: State-Time Option Kernel Spatial-Attention Amortizer
+Reference: Ringstrom & Schrater (2025), arXiv:2506.09499
+
+Zero-dependency Python / NumPy implementation of the spatial amortizer.
+Compatible with PyTorch if available, but requires only standard NumPy.
+"""
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Union
+import json
+import os
+import struct
+import numpy as np
+
+
+@dataclass
+class AmortizedPrediction:
+    """Structured predictions from the 37 output heads of STOK-SpatialNet."""
+    kappa: float                      # Feasibility in [0, 1]
+    urgency: float                    # Time-to-kill urgency: 1 / (1 + E[t_f])
+    valence: float                    # Normalized signed empowerment change: delta_E / 32
+    lose: float                       # Predicted lose rate (survivability constraint)
+    goal_kappa: Dict[str, float]      # Feasibility per goal family (Strike, Engage, Evade, Turtle, Zone)
+    eta: List[float]                  # [kill@t1, kill@t2, not-in-window] probability mass
+    opp_prior: Dict[str, float]       # Predicted distribution over opponent goal family
+    valence_own: float                # Decomposed own option-capacity growth
+    valence_deny: float               # Decomposed denial of opponent option-capacity
+    concept: Dict[str, float]         # Tactical concept distribution (trap, spacing, punish, ...)
+    opp_concept: Dict[str, float]     # Theory of Mind: predicted opponent concept distribution
+    v_traj: float                     # Long-horizon realized trajectory empowerment flow
+    v_past: float                     # Inherited past trajectory flow
+
+    def expected_time(self) -> float:
+        """Expected time-to-kill under eta, with not-in-window priced at 8 turns."""
+        return self.eta[0] + 2.0 * self.eta[1] + 8.0 * self.eta[2]
+
+
+class STOKSpatialAmortizer:
+    """
+    Spatial-Attention Amortizer for State-Time Option Kernels (STOK).
+    
+    Architecture:
+      1. Tile embedding (channels=39 -> attn=32) + ReLU
+      2. Residual tile-to-tile single-head self-attention
+      3. Goal-query cross-attention read (query=8 -> keys/values=32)
+      4. Dense trunk on [flat (106) || context (32)] -> hidden (96) + ReLU
+      5. Multi-head projection (37 heads)
+    """
+
+    def __init__(self, weights: Dict[str, np.ndarray], config: dict):
+        self.weights = weights
+        self.config = config
+        self.channels = config.get("spatial_channels", 39)
+        self.attn = config.get("attn_dim", 32)
+        self.goal_dim = config.get("goal_dim", 8)
+        self.flat_dim = config.get("flat_dim", 106)
+        self.hidden = config.get("hidden_dim", 96)
+        self.heads = config.get("num_heads", 37)
+        self.goal_families = config.get(
+            "goal_families", ["Strike", "Engage", "Evade", "Turtle", "Zone"]
+        )
+        self.concepts = config.get(
+            "tactical_concepts",
+            ["trap", "spacing", "punish", "bait", "coverage", "conversion", "denial", "tempo"],
+        )
+
+    @classmethod
+    def from_pretrained(cls, model_dir_or_path: str):
+        """Load model from a directory containing config.json and model.safetensors."""
+        if os.path.isfile(model_dir_or_path):
+            model_dir = os.path.dirname(model_dir_or_path)
+            weights_path = model_dir_or_path
+        else:
+            model_dir = model_dir_or_path
+            weights_path = os.path.join(model_dir, "model.safetensors")
+
+        config_path = os.path.join(model_dir, "config.json")
+        with open(config_path, "r", encoding="utf-8") as f:
+            config = json.load(f)
+
+        weights = cls._load_safetensors(weights_path)
+        return cls(weights, config)
+
+    @staticmethod
+    def _load_safetensors(path: str) -> Dict[str, np.ndarray]:
+        with open(path, "rb") as f:
+            header_len = struct.unpack("<Q", f.read(8))[0]
+            header_json = f.read(header_len).decode("utf-8")
+            header = json.loads(header_json)
+            raw_data = f.read()
+
+        tensors = {}
+        for name, meta in header.items():
+            if name == "__metadata__":
+                continue
+            start, end = meta["data_offsets"]
+            shape = meta["shape"]
+            dtype = meta["dtype"]
+            assert dtype == "F32", f"Unsupported dtype: {dtype}"
+            arr = np.frombuffer(raw_data[start:end], dtype=np.float32).reshape(shape)
+            tensors[name] = arr
+        return tensors
+
+    @staticmethod
+    def _softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
+        e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
+        return e_x / np.sum(e_x, axis=axis, keepdims=True)
+
+    def forward(
+        self,
+        spatial_field: np.ndarray,
+        goal_vector: np.ndarray,
+        flat_features: np.ndarray,
+    ) -> AmortizedPrediction:
+        """
+        Run forward inference.
+        
+        Args:
+            spatial_field: Array of shape [num_tiles, 39] (e.g. 100 tiles for 10x10).
+            goal_vector: Array of shape [8] (5-way one-hot + [dx, dz, has_target]).
+            flat_features: Array of shape [106] (high-level features + goal + history).
+            
+        Returns:
+            AmortizedPrediction with all 37 heads decoded and labeled.
+        """
+        assert spatial_field.shape[-1] == self.channels
+        assert goal_vector.shape[-1] == self.goal_dim
+        assert flat_features.shape[-1] == self.flat_dim
+
+        tiles = spatial_field.astype(np.float32)
+        scale = np.sqrt(self.attn, dtype=np.float32)
+
+        # 1. Tile Embedding: [T, 39] @ W_embed.T + b -> [T, 32]
+        h0 = np.maximum(0.0, np.dot(tiles, self.weights["embed.weight"].T) + self.weights["embed.bias"])
+
+        # 2. Residual Tile-to-Tile Self-Attention
+        q = np.dot(h0, self.weights["sa_q.weight"].T) + self.weights["sa_q.bias"]
+        k = np.dot(h0, self.weights["sa_k.weight"].T) + self.weights["sa_k.bias"]
+        v = np.dot(h0, self.weights["sa_v.weight"].T) + self.weights["sa_v.bias"]
+
+        attn_scores = np.dot(q, k.T) / scale  # [T, T]
+        attn_weights = self._softmax(attn_scores, axis=-1)  # [T, T]
+        h = h0 + np.dot(attn_weights, v)  # [T, 32]
+
+        # 3. Goal-Query Cross-Attention Read
+        gq = np.dot(goal_vector, self.weights["query.weight"].T) + self.weights["query.bias"]  # [32]
+        k2 = np.dot(h, self.weights["key.weight"].T) + self.weights["key.bias"]  # [T, 32]
+        v2 = np.dot(h, self.weights["value.weight"].T) + self.weights["value.bias"]  # [T, 32]
+
+        read_scores = np.dot(k2, gq) / scale  # [T]
+        read_weights = self._softmax(read_scores, axis=-1)  # [T]
+        ctx = np.dot(read_weights, v2)  # [32]
+
+        # 4. Dense Trunk
+        x = np.concatenate([flat_features, ctx], axis=-1)  # [138]
+        hid = np.maximum(0.0, np.dot(x, self.weights["l1.weight"].T) + self.weights["l1.bias"])  # [96]
+        out = np.dot(hid, self.weights["l2.weight"].T) + self.weights["l2.bias"]  # [37]
+
+        # 5. Decode Heads
+        kappa = float(out[0])
+        urgency = float(out[1])
+        valence = float(out[2])
+        lose = float(out[3])
+
+        goal_kappa = {fam: float(out[4 + i]) for i, fam in enumerate(self.goal_families)}
+        eta = self._softmax(out[9:12]).tolist()
+        opp_prior_vals = self._softmax(out[12:17])
+        opp_prior = {fam: float(opp_prior_vals[i]) for i, fam in enumerate(self.goal_families)}
+
+        valence_own = float(out[17])
+        valence_deny = float(out[18])
+
+        concept_vals = self._softmax(out[19:27])
+        concept = {c: float(concept_vals[i]) for i, c in enumerate(self.concepts)}
+
+        opp_concept_vals = self._softmax(out[27:35])
+        opp_concept = {c: float(opp_concept_vals[i]) for i, c in enumerate(self.concepts)}
+
+        v_traj = float(out[35])
+        v_past = float(out[36])
+
+        return AmortizedPrediction(
+            kappa=kappa,
+            urgency=urgency,
+            valence=valence,
+            lose=lose,
+            goal_kappa=goal_kappa,
+            eta=eta,
+            opp_prior=opp_prior,
+            valence_own=valence_own,
+            valence_deny=valence_deny,
+            concept=concept,
+            opp_concept=opp_concept,
+            v_traj=v_traj,
+            v_past=v_past,
+        )