feat(actpred): HailoPredictor 适配 v4 onehot HEF——1864 logits 解码 + 固定列覆盖

v4 onehot 模型 HEF 输出 (1,1,1864)，原回归假设（输出=action_dim）导致
输入分类与维度校验失败。_classify_inputs 改为 pred_token 先判（s[1]==1）、
其余 3 维归 history；_derive_dims 识别 TOTAL_BINS 输出；predict 经 bins.decode
期望值解码并固定权重 0 列（与 predict.py 同语义），新增 column_weights 参数。
_npu_test.py 指向 v4 并增加 onehot 参考分支，NPU 实测 4 角 PASS、
HEF vs torch max_diff 0.5°。

actpred/_npu_test.py

@@ -4,7 +4,7 @@
 
 4 角测试（真实/全零帧 × 真实/全零历史）+ 与 PyTorch 参考对比。
 HEF 尺寸/布局从 vstream info 自动推导，checkpoint 用
-checkpoints/pretrain_best_inference.pt。
+checkpoints/best_inference.pt（v4 onehot，fixed_range）。
 """
 import sys
 import time
@@ -20,8 +20,8 @@ from config import parse_checkpoint_config
 from model import ActionPredictor
 from hailo_infer import HailoPredictor
 
-HEF_PATH = "checkpoints/actpred_pretrain.hef"
-CKPT = "checkpoints/pretrain_best_inference.pt"
+HEF_PATH = "checkpoints/actpred_v4.hef"
+CKPT = "checkpoints/best_inference.pt"
 IMG_DIR = "data/episode_0000/images"
 
 hef = HEF(HEF_PATH)
@@ -40,11 +40,11 @@ imgs_real = [np.array(Image.open(f"{IMG_DIR}/000{i}.jpg").convert("RGB"))
 imgs_zero = [np.zeros_like(imgs_real[0]) for _ in range(3)]
 
 hist_real = np.array([
-    [45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 10, 200, 0, 0, 0, 0, 0, 0],
-    [45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 20, 0, 0, 0, 0, 0, 0],
-    [45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 0, 0, 0, 0, 0, 0, 0],
-    [45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 0, 0, 0, 0, 0, 0, 0],
-    [45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 0, 0, 0, 0, 0, 0, 0],
+    [45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 10, 0, 0, 0, 0, 0, 0],
+    [45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 10, 0, 0, 0, 0, 0, 0],
+    [45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 10, 0, 0, 0, 0, 0, 0],
+    [45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 10, 0, 0, 0, 0, 0, 0],
+    [45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 10, 0, 0, 0, 0, 0, 0],
 ], dtype=np.float32)
 hist_zero = np.zeros((5, 18), dtype=np.float32)
 
@@ -99,12 +99,24 @@ def torch_pred(imgs, hist, use_learned_pred=True, imu=None):
         ii = torch.from_numpy(normalize_imu(imu, imu_dim)).unsqueeze(0)
     if use_learned_pred:
         with torch.no_grad():
-            return model(x, h, ii).squeeze(0).numpy()
-    saved = model.pred_token.data.clone()
-    model.pred_token.data.zero_()
-    with torch.no_grad():
-        out = model(x, h, ii).squeeze(0).numpy()
-    model.pred_token.data.copy_(saved)
+            out = model(x, h, ii).squeeze(0).numpy()
+    else:
+        saved = model.pred_token.data.clone()
+        model.pred_token.data.zero_()
+        with torch.no_grad():
+            out = model(x, h, ii).squeeze(0).numpy()
+        model.pred_token.data.copy_(saved)
+    if cfg.data.onehot:
+        # 多头 logits → 期望值解码 → 原始量纲（与 predict.py 同语义，
+        # 含权重 0 列固定常量），对比前无需反归一化
+        from bins import decode as bin_decode, fixed_zero_columns
+        from action_norm import default_column_weights
+        out = bin_decode(out, cfg.data.action_dim).astype(np.float32)
+        w = (np.asarray(cfg.data.column_weights, dtype=np.float32)
+             if cfg.data.column_weights is not None
+             else default_column_weights(cfg.data.action_dim))
+        fixed = fixed_zero_columns(w, cfg.data.action_dim)
+        out[w <= 0] = fixed[w <= 0]
     return out
 
 
@@ -141,15 +153,18 @@ with HailoPredictor(HEF_PATH, stats_path=stats_path,
     ref = torch_pred(imgs_real, hist_real, True, imu_real)
     ref_zero_pred = torch_pred(imgs_real, hist_real, False, imu_real)
     # torch 参考是模型原始归一化空间输出；HailoPredictor 已反归一化
-    # 到原始动作空间，比较前须同空间对齐，否则归一化模型下 max_diff 无意义
-    if norm_mode == "fixed_range":
-        from action_norm import denormalize_action
-        ref = denormalize_action(ref, cfg.data.action_dim)
-        ref_zero_pred = denormalize_action(ref_zero_pred, cfg.data.action_dim)
-    elif stats_path is not None:
-        s = np.load(stats_path)
-        ref = ref * s[1] + s[0]
-        ref_zero_pred = ref_zero_pred * s[1] + s[0]
+    # 到原始动作空间，比较前须同空间对齐，否则归一化模型下 max_diff
+    # 无意义（onehot 模型 torch_pred 已 decode 为原始量纲，无需再处理）
+    if not cfg.data.onehot:
+        if norm_mode == "fixed_range":
+            from action_norm import denormalize_action
+            ref = denormalize_action(ref, cfg.data.action_dim)
+            ref_zero_pred = denormalize_action(ref_zero_pred,
+                                               cfg.data.action_dim)
+        elif stats_path is not None:
+            s = np.load(stats_path)
+            ref = ref * s[1] + s[0]
+            ref_zero_pred = ref_zero_pred * s[1] + s[0]
     print(f"torch learned_pred  : {np.round(ref[:6], 3)}...")
     print(f"torch zero_pred     : {np.round(ref_zero_pred[:6], 3)}...")
     print(f"HEF vs torch(learned) max_diff={np.abs(a - ref).max():.4f}"

actpred/hailo_infer.py

@@ -39,6 +39,7 @@ from imu_norm import normalize_imu
 from action_norm import (normalize_action, denormalize_action,
                       is_range_mode)
 from utils import parse_action_history
+from bins import TOTAL_BINS, decode as bin_decode, fixed_zero_columns
 
 IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
 IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
@@ -62,7 +63,7 @@ class HailoPredictor:
 
     def __init__(self, hef_path: str, stats_path: str = None, vdevice=None,
                  predict_delta: bool = False, pred_token=None,
-                 normalize_mode: str = "zscore"):
+                 normalize_mode: str = "zscore", column_weights=None):
         """HEF 推理器。
 
         predict_delta: 模型用 --predict-delta 训练时须置 True，推理输出
@@ -72,12 +73,19 @@ class HailoPredictor:
         normalize_mode: 与训练时 DataConfig.normalize_mode 一致。
             "zscore"（默认）需传 stats_path；"fixed_range" / "log_range"
             免 stats 文件，按固定物理范围归一化（action_norm.py）。
+        column_weights: 逐列损失权重（onehot HEF 的权重 0 列按
+            bins.FIXED_OUTPUT 固定，防闭环垃圾回写）；None 时用
+            action_norm.default_column_weights。
         """
         from hailo_platform import HEF, VDevice
 
         self._predict_delta = predict_delta
         self._range_mode = (normalize_mode
                             if is_range_mode(normalize_mode) else None)
+        self._column_weights = None
+        if column_weights is not None:
+            self._column_weights = np.asarray(column_weights,
+                                              dtype=np.float32)
         self._default_pred_token = None
         if pred_token is not None:
             if isinstance(pred_token, str):
@@ -132,7 +140,6 @@ class HailoPredictor:
         imu 为 R2 方案一新增输入（1,1,1,imu_dim），旧 HEF 无此输入。
         """
         ins = list(self._in_infos)
-        out_dim = tuple(self._out_info.shape)[-1]
         # 新版：每帧独立 (H,W,3) 输入
         frames = [i for i in ins if len(i.shape) == 3
                   and i.shape[2] == 3 and i.shape[0] > 3]
@@ -145,10 +152,10 @@ class HailoPredictor:
         history, pred, imu = None, None, None
         for i in rest:
             s = tuple(i.shape)
-            if len(s) == 3 and s[-1] == out_dim:
-                history = i
-            elif len(s) == 3 and s[1] == 1:
-                pred = i
+            if len(s) == 3 and s[1] == 1:
+                pred = i                     # pred_token (1,1,d)
+            elif len(s) == 3:
+                history = i                  # 动作历史 (1,H_len,D)
             elif len(s) == 4 and s[0] == 1 and s[1] == 1 and s[2] == 1:
                 imu = i                     # R2：imu (1,1,1,imu_dim)
             else:
@@ -159,7 +166,8 @@ class HailoPredictor:
             raise ValueError("HEF 输入中未找到图像输入")
         if history is None:
             raise ValueError(
-                f"HEF 输入中未找到动作历史输入（输出维度 {out_dim}），"
+                f"HEF 输入中未找到动作历史输入（输出维度 "
+                f"{tuple(self._out_info.shape)}），"
                 f"实际: {[(i.name, tuple(i.shape)) for i in ins]}")
         return frames, history, pred, imu
 
@@ -209,9 +217,16 @@ class HailoPredictor:
 
         out_shape = tuple(self._out_info.shape)
         out_dim = out_shape[-1]
-        if out_dim != self.action_dim:
-            raise ValueError(
-                f"输出动作维度 {out_dim} 与动作历史维度 {self.action_dim} 不一致")
+        if out_dim == TOTAL_BINS:
+            # onehot 模型：输出为多头 logits，predict 时经 bins.decode
+            # 解码为 action_dim 原始量纲（无需反归一化）
+            self._onehot = True
+        else:
+            self._onehot = False
+            if out_dim != self.action_dim:
+                raise ValueError(
+                    f"输出动作维度 {out_dim} 与动作历史维度 "
+                    f"{self.action_dim} 不一致")
 
     # ---------------- 生命周期 ----------------
 
@@ -310,6 +325,23 @@ class HailoPredictor:
             q = self._out_info.quant_info
             result = (result - q.qp_zp) * q.qp_scale
         result = result.reshape(-1)
+        if self._onehot:
+            # 多头 logits → 期望值解码 → 原始量纲（箱生活在原始空间，
+            # 无需反归一化）；权重 0 列固定为常量（与 predict.py 对齐），
+            # 防闭环垃圾回写
+            result = bin_decode(result, self.action_dim).reshape(-1)
+            if self._column_weights is None:
+                from action_norm import default_column_weights
+                self._column_weights = default_column_weights(
+                    self.action_dim)
+            w = np.asarray(self._column_weights, dtype=np.float32)
+            fixed = fixed_zero_columns(w, self.action_dim)
+            result[w[:self.action_dim] <= 0.0] = (
+                fixed[w[:self.action_dim] <= 0.0])
+            if not np.all(np.isfinite(result)):
+                print("WARNING: HEF output contains NaN/Inf, replacing with zeros")
+                result = np.nan_to_num(result, nan=0.0, posinf=0.0, neginf=0.0)
+            return result
         if self._predict_delta:
             # 残差模式：Δ̂ + a_{t-1}（归一化空间），与 predict.py 对齐
             result = result + action_history[-1]