From 28665bf663eb4ef31f003bc9e02105d5552a1b62 Mon Sep 17 00:00:00 2001
From: KoalaTrapLord <265133236+KoalaTrapLord@users.noreply.github.com>
Date: Fri, 3 Apr 2026 15:16:49 -0400
Subject: [PATCH] fix: apply _unwrap_logits guard in _infer_from_vector
 (routes/policy.py)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The crash `unsupported operand type(s) for /: 'dict' and 'float'` was
firing from _infer_from_vector(), which is the code path taken when
has_precomputed_vector() is True (GameSession 1.1.0+ snapshots with
a precomputed state_vector float[29]).

Unlike PolicyInferenceService.predict() — which was fixed in the
previous commit via _unwrap_logits — this function calls
policy_service.model(state_tensor) directly and immediately uses the
return value in arithmetic (logits / max(temperature, 0.01)), so any
tuple/dict return (DataParallel, torch.compile, PolicyValueNetwork
swap) causes the TypeError.

Fix: import _unwrap_logits from ml.serving.policy_server and apply it
to the raw model output in _infer_from_vector before the temperature
division and argmax calls. All three return shapes are now handled:
  - plain tensor  → passthrough (normal case, zero overhead)
  - tuple/list    → raw[0]  (PolicyValueNetwork (logits, value))
  - dict          → first recognised key or first tensor value
---
 routes/policy.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/routes/policy.py b/routes/policy.py
index ba68911..95f1f84 100644
--- a/routes/policy.py
+++ b/routes/policy.py
@@ -518,16 +518,24 @@ def _infer_from_vector(state_vec: np.ndarray, policy_service, temperature: float
 
     Mirrors the core of PolicyInferenceService.predict() but accepts an
     ndarray instead of a snapshot dict, avoiding a second encode() call.
+
+    Uses _unwrap_logits (same helper as predict()) to guard against
+    tuple/dict returns from wrapped or swapped model variants before any
+    arithmetic on the logits tensor.
     """
     import torch
     from ml.config.scope import NUM_ACTIONS, IDX_TO_ACTION
+    from ml.serving.policy_server import _unwrap_logits
 
     t_start = time.time()
     try:
         state_tensor = torch.from_numpy(state_vec).unsqueeze(0).to(policy_service.device)
 
         with torch.no_grad():
-            logits = policy_service.model(state_tensor)
+            raw_output = policy_service.model(state_tensor)
+            # Guard: unwrap tuple/dict returns (DataParallel, torch.compile,
+            # PolicyValueNetwork swap) before temperature division / argmax.
+            logits = _unwrap_logits(raw_output)
 
             if greedy:
                 action_idx = logits.argmax(dim=-1).item()