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:
1 parent
83d05294c8
commit
28665bf663
1 file changed
+9
-1
+9
-1
@@ -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()
|
||||
|
||||
Reference in new issue
Block a user