fix: apply _unwrap_logits guard in _infer_from_vector (routes/policy.py)

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
This commit is contained in:
KoalaTrapLord committed 2026-04-03 15:16:49 -04:00
1 parent 83d05294c8
commit 28665bf663
1 file changed
+9 -1
+9 -1
View File
@@ -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()