feat(actpred): 多头 one-hot 分类输出——18 组独立 softmax 头替代连续回归
分箱（bins.py）：舵机 1°/箱 ×10（180 箱，匹配 0.75° 微动作插值粒度）、 spk.F/spk.delay 对数 32 箱 ×2，共 1864 logits；保留位无头。软标签 （onehot_sigma=1.0 高斯扩散）缓解小样本数据稀疏；期望值解码恢复 连续分辨率（箱生活在原始空间，无反变换放大）。 训练损失为逐头加权 CE（列权重沿用），闭环反馈经解码→范围归一化 回写历史；predict 输出仍为 18 维原始量纲，pybot 契约零变化。 结构性消除尺度放大（每箱等距）与物理界问题（箱天然有界）， 为 A-E 归一化实验失败后的正解路线。要求 normalize_mode 为范围类， 与 predict_delta 互斥。测试 +10（含 CLI 映射回归）。
This commit is contained in:
qing committed
1 parent
99c1bc1739
commit
a18a7de1ca
7 files changed
+409
-20
No files matched your search
+114
@@ -0,0 +1,114 @@
|
||||
"""多头 one-hot 分箱 —— 18 维连续动作向量与多组 one-hot 分类目标互转
|
||||
|
||||
分箱设计（actpred_onehot_multihead_dev_plan_20260814.md）：
|
||||
0-9 舵机角：均匀 1°/箱 → 180 箱（匹配 0.75° 微动作插值数据粒度）
|
||||
10 spk.F：对数分箱 log10(1+x) ∈ [0, log10(5001)] → 32 箱
|
||||
11 spk.delay：对数分箱 log10(1+x) ∈ [0, log10(2001)] → 32 箱
|
||||
12-17 保留位：无头（恒 0，不参与分类）
|
||||
|
||||
编码/解码与 action_norm 的关系：箱生活在原始物理空间（舵机箱=角度），
|
||||
decode 直接产出原始量纲，无需反归一化；输入历史仍按 normalize_mode
|
||||
归一化（与现状一致），仅监督目标换为 one-hot。
|
||||
|
||||
软标签：encode(x, sigma) 把真值 one-hot 高斯扩散到邻近箱，缓解
|
||||
804→1864 分类面的数据稀疏（多数箱极少被访问）。
|
||||
|
||||
仅依赖 numpy，推理/部署文件可直接复用。
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
# 每组 one-hot 的箱数（按列索引；保留位无头）
|
||||
BINS_PER_COL = [180] * 10 + [32, 32]
|
||||
# 有头的列
|
||||
HEAD_COLS = tuple(range(12))
|
||||
# 各组起始位置（组 i 占 [START[i], START[i+1])）
|
||||
START = np.concatenate([[0], np.cumsum(BINS_PER_COL)])
|
||||
TOTAL_BINS = int(START[-1]) # 1864
|
||||
|
||||
LOG_COLS = (10, 11) # 对数分箱列
|
||||
LOG_HI = {10: 5000.0, 11: 2000.0}
|
||||
|
||||
|
||||
def bin_edges(col: int) -> np.ndarray:
|
||||
"""列 col 的箱边界数组（长度 bins+1），原始量纲递增"""
|
||||
n = BINS_PER_COL[col]
|
||||
if col in LOG_COLS:
|
||||
hi = LOG_HI[col]
|
||||
lo_log, hi_log = 0.0, np.log10(hi + 1.0)
|
||||
# 对数空间均匀边界 → 原始空间箱边界
|
||||
return np.power(10.0, np.linspace(lo_log, hi_log, n + 1)) - 1.0
|
||||
lo, hi = 0.0, 180.0
|
||||
return np.linspace(lo, hi, n + 1)
|
||||
|
||||
|
||||
def bin_centers(col: int) -> np.ndarray:
|
||||
"""列 col 的箱中心（decode 期望值的坐标），原始量纲"""
|
||||
e = bin_edges(col)
|
||||
return 0.5 * (e[:-1] + e[1:])
|
||||
|
||||
|
||||
def encode(x: np.ndarray, action_dim: int = 18, sigma: float = 0.0,
|
||||
out: np.ndarray = None) -> np.ndarray:
|
||||
"""连续动作向量 (..., D) → 多组 one-hot (..., TOTAL_BINS)
|
||||
|
||||
sigma > 0 时启用软标签：真值箱高斯扩散 exp(-d²/2σ²)，d=与真值箱的
|
||||
距离（按箱索引）；归一化后总和为 1。保留位列（12+）无头，恒不编码。
|
||||
"""
|
||||
a = np.asarray(x, dtype=np.float32)
|
||||
flat = a.reshape(-1, a.shape[-1])
|
||||
n_rows = flat.shape[0]
|
||||
if out is None:
|
||||
out = np.zeros(a.shape[:-1] + (TOTAL_BINS,), dtype=np.float32)
|
||||
out_flat = out.reshape(n_rows, TOTAL_BINS)
|
||||
rows = np.arange(n_rows)
|
||||
for col in HEAD_COLS:
|
||||
if col >= flat.shape[1]:
|
||||
break
|
||||
edges = bin_edges(col)
|
||||
idx = np.clip(np.searchsorted(edges, flat[:, col], side="right") - 1,
|
||||
0, BINS_PER_COL[col] - 1)
|
||||
s, e = START[col], START[col + 1]
|
||||
if sigma <= 0.0:
|
||||
out_flat[rows, s + idx] = 1.0
|
||||
else:
|
||||
n = BINS_PER_COL[col]
|
||||
d = np.abs(np.arange(n)[None, :] - idx[:, None])
|
||||
g = np.exp(-(d * d) / (2.0 * sigma * sigma))
|
||||
g = g / g.sum(axis=1, keepdims=True)
|
||||
out_flat[:, s:s + n] = g
|
||||
return out
|
||||
|
||||
|
||||
def decode(logits: np.ndarray, action_dim: int = 18,
|
||||
temperature: float = 1.0) -> np.ndarray:
|
||||
"""多头 logits (..., TOTAL_BINS) → 连续动作向量 (..., D)，原始量纲
|
||||
|
||||
逐头 softmax → 期望值 Σ pᵢ·箱中心（分辨率超越箱宽）；保留位列填 0。
|
||||
temperature 用于探索采样（=1 为正常推理）。
|
||||
"""
|
||||
l = np.asarray(logits, dtype=np.float32)
|
||||
out = np.zeros(l.shape[:-1] + (action_dim,), dtype=np.float32)
|
||||
for col in HEAD_COLS:
|
||||
if col >= action_dim:
|
||||
break
|
||||
s, e = START[col], START[col + 1]
|
||||
z = l[..., s:e] / max(temperature, 1e-6)
|
||||
z = z - z.max(axis=-1, keepdims=True) # 数值稳定
|
||||
p = np.exp(z)
|
||||
p = p / p.sum(axis=-1, keepdims=True)
|
||||
out[..., col] = p @ bin_centers(col)
|
||||
return out
|
||||
|
||||
|
||||
def argmax_decode(logits: np.ndarray, action_dim: int = 18) -> np.ndarray:
|
||||
"""argmax 解码：每头取最大概率箱中心（确定性，量化到箱宽）"""
|
||||
l = np.asarray(logits, dtype=np.float32)
|
||||
out = np.zeros(l.shape[:-1] + (action_dim,), dtype=np.float32)
|
||||
for col in HEAD_COLS:
|
||||
if col >= action_dim:
|
||||
break
|
||||
s, e = START[col], START[col + 1]
|
||||
idx = np.argmax(l[..., s:e], axis=-1)
|
||||
out[..., col] = bin_centers(col)[idx]
|
||||
return out
|
||||
@@ -24,6 +24,15 @@ class DataConfig:
|
||||
# .default_column_weights(action_dim)（舵机 1.0 / spk 0.2 / 保留位 0.0）。
|
||||
# 显式指定时长度须等于 action_dim。
|
||||
column_weights: Optional[List[float]] = None
|
||||
# 多头 one-hot 分类输出（bins.py 分箱）：True 时监督目标为多组 one-hot
|
||||
# （舵机 1°/箱 ×10、spk 对数 32 箱 ×2，共 1864 logits），损失用加权 CE，
|
||||
# 推理端期望值解码。要求 normalize_actions=True 且 normalize_mode 为
|
||||
# 范围类（fixed_range/log_range）——历史输入沿用范围归一化，免 stats。
|
||||
# 与 predict_delta 互斥（残差在连续空间定义）。
|
||||
onehot: bool = False
|
||||
# one-hot 软标签高斯宽度（箱）：0=硬标签；>0 真值扩散到邻近箱，
|
||||
# 缓解小样本下分类面数据稀疏（1°/箱 × 804 样本多数箱极少被访问）
|
||||
onehot_sigma: float = 1.0
|
||||
# R2 方案一：MPU6050 姿态作为独立 IMU token 注入（序列 F+H+2）。
|
||||
# 0 = 关闭（默认，序列保持 F+H+1，旧 checkpoint 完全兼容）；
|
||||
# 8×N = N 个传感器的 roll,pitch,ax,ay,az,gx,gy,gz 按 pybot 脚本
|
||||
@@ -100,6 +109,11 @@ def build_argparser(description: str = "") -> argparse.ArgumentParser:
|
||||
help="动作归一化方式：zscore=数据统计（依赖 action_stats.npy）；"
|
||||
"fixed_range=固定物理范围 min-max→[0,1]（免 stats 文件）；"
|
||||
"log_range=spk 两列对数尺度（免 stats 文件）")
|
||||
parser.add_argument("--onehot", action=argparse.BooleanOptionalAction,
|
||||
default=None,
|
||||
help="多头 one-hot 分类输出（bins.py 分箱，舵机 1°/箱）")
|
||||
parser.add_argument("--onehot-sigma", type=float, default=None,
|
||||
help="one-hot 软标签高斯宽度（箱，0=硬标签，默认 1.0）")
|
||||
parser.add_argument("--column-weights", type=str, default=None,
|
||||
help="逐列损失权重（逗号分隔 action_dim 个浮点），"
|
||||
"缺省用默认表（舵机1.0/spk0.2/保留位0.0）；仅影响训练损失")
|
||||
@@ -180,6 +194,21 @@ def validate_config_values(data_cfg: DataConfig, model_cfg: ModelConfig,
|
||||
f"normalize_mode={data_cfg.normalize_mode!r} 非法，"
|
||||
f"应为 'zscore'、'fixed_range' 或 'log_range'")
|
||||
|
||||
if data_cfg.onehot:
|
||||
if not data_cfg.normalize_actions:
|
||||
raise ValueError("onehot=True 要求 normalize_actions=True"
|
||||
"（历史输入须范围归一化）")
|
||||
if not (data_cfg.normalize_mode in ("fixed_range", "log_range")):
|
||||
raise ValueError(
|
||||
f"onehot=True 要求 normalize_mode 为范围类，"
|
||||
f"实际为 {data_cfg.normalize_mode!r}")
|
||||
if train_cfg.predict_delta:
|
||||
raise ValueError("onehot 与 predict_delta 互斥"
|
||||
"（残差定义在连续空间）")
|
||||
if not (0.0 <= data_cfg.onehot_sigma <= 10.0):
|
||||
raise ValueError(
|
||||
f"onehot_sigma={data_cfg.onehot_sigma} 应在 [0, 10] 内")
|
||||
|
||||
if data_cfg.column_weights is not None:
|
||||
w = data_cfg.column_weights
|
||||
if len(w) != data_cfg.action_dim:
|
||||
@@ -293,6 +322,8 @@ def apply_overrides(cfg: Config, args: argparse.Namespace) -> Config:
|
||||
"action_stats_path": ("data", "action_stats_path"),
|
||||
"normalize_mode": ("data", "normalize_mode"),
|
||||
"column_weights": ("data", "column_weights"),
|
||||
"onehot": ("data", "onehot"),
|
||||
"onehot_sigma": ("data", "onehot_sigma"),
|
||||
"imu_dim": ("data", "imu_dim"),
|
||||
"backbone": ("model", "cnn_backbone"),
|
||||
"cnn_feature_dim": ("model", "cnn_feature_dim"),
|
||||
|
||||
+12
-2
@@ -12,6 +12,7 @@ from torchvision.transforms import functional as TF
|
||||
|
||||
from imu_norm import imu_scales
|
||||
from action_norm import normalize_action, is_range_mode
|
||||
from bins import encode as bin_encode
|
||||
|
||||
|
||||
def _natural_key(p):
|
||||
@@ -26,7 +27,8 @@ class ActionDataset(Dataset):
|
||||
action_stats_path: str = None, episode_indices: list = None,
|
||||
frame_len: int = 3, augment: bool = False,
|
||||
predict_delta: bool = False, rollout_len: int = 1,
|
||||
imu_dim: int = 0, normalize_mode: str = "zscore"):
|
||||
imu_dim: int = 0, normalize_mode: str = "zscore",
|
||||
onehot: bool = False, onehot_sigma: float = 1.0):
|
||||
self.history_len = history_len
|
||||
self.action_dim = action_dim
|
||||
# 固定范围类归一化（fixed_range / log_range，免 stats 文件）；
|
||||
@@ -34,6 +36,9 @@ class ActionDataset(Dataset):
|
||||
self._range_mode = (normalize_mode
|
||||
if (normalize_actions and is_range_mode(normalize_mode))
|
||||
else None)
|
||||
# 多头 one-hot：监督目标为多组 one-hot（bins.py），历史仍走 _norm
|
||||
self.onehot = onehot
|
||||
self.onehot_sigma = onehot_sigma
|
||||
# R2 方案一：>0 时样本带 "imu" 键（归一化姿态向量，缺失 episode 补零）
|
||||
self.imu_dim = imu_dim
|
||||
self.image_width = image_width
|
||||
@@ -151,7 +156,12 @@ class ActionDataset(Dataset):
|
||||
return torch.stack(frames) # (F, 3, H, W)
|
||||
|
||||
def _step_target(self, actions, t):
|
||||
"""时刻 t 的监督目标（归一化空间）；predict_delta 时为 Δ=actions[t]-actions[t-1]"""
|
||||
"""时刻 t 的监督目标；onehot 时为多组 one-hot（原始空间分箱），
|
||||
predict_delta 时为 Δ=actions[t]-actions[t-1]（归一化空间），
|
||||
否则为归一化连续值"""
|
||||
if self.onehot:
|
||||
return bin_encode(actions[t], self.action_dim,
|
||||
sigma=self.onehot_sigma)
|
||||
if self.predict_delta:
|
||||
return self._norm(actions[t]) - self._norm(actions[t - 1])
|
||||
return self._norm(actions[t])
|
||||
|
||||
+3
-1
@@ -127,8 +127,10 @@ class ActionPredictor(nn.Module):
|
||||
encoder_layer, num_layers=model_cfg.attn_layers,
|
||||
enable_nested_tensor=False)
|
||||
|
||||
from bins import TOTAL_BINS
|
||||
head_out = TOTAL_BINS if data_cfg.onehot else data_cfg.action_dim
|
||||
self.head = PredictionHead(
|
||||
d, model_cfg.fused_hidden_dim, data_cfg.action_dim, model_cfg.dropout)
|
||||
d, model_cfg.fused_hidden_dim, head_out, model_cfg.dropout)
|
||||
|
||||
self._frame_len = data_cfg.frame_len
|
||||
self._history_len = data_cfg.history_len
|
||||
|
||||
@@ -10,6 +10,7 @@ from config import parse_checkpoint_config
|
||||
from imu_norm import normalize_imu
|
||||
from action_norm import (normalize_action, denormalize_action,
|
||||
is_range_mode)
|
||||
from bins import decode as bin_decode
|
||||
from model import ActionPredictor
|
||||
from utils import parse_action_history
|
||||
|
||||
@@ -37,6 +38,7 @@ class ActionPredictorInference:
|
||||
# pybot 闭环鸭子类型约定：此属性为 None 时按旧两参调用，
|
||||
# 故 imu_dim=0（旧 checkpoint）必须映射为 None。
|
||||
self.imu_dim = self._data_cfg.imu_dim or None
|
||||
self._onehot = self._data_cfg.onehot
|
||||
# R1-C 残差预测：checkpoint 记录 predict_delta=True 时，模型输出 Δ̂，
|
||||
# 需在归一化空间加回 a_{t-1}（history 末步）得到绝对动作。
|
||||
# 旧 checkpoint 无该字段，getattr 默认 False，行为与旧版一致。
|
||||
@@ -117,6 +119,12 @@ class ActionPredictorInference:
|
||||
normalize_imu(imu, self._data_cfg.imu_dim)
|
||||
).unsqueeze(0).to(self._device)
|
||||
pred = self._model(img_tensor, history, imu_tensor).squeeze(0) # 归一化空间
|
||||
if self._onehot:
|
||||
# 多头 logits → 期望值解码 → 原始量纲（箱生活在原始空间）
|
||||
result = bin_decode(pred.cpu().numpy(), self.action_dim)
|
||||
if not np.all(np.isfinite(result)):
|
||||
result = np.nan_to_num(result, nan=0.0, posinf=0.0, neginf=0.0)
|
||||
return result
|
||||
if self._predict_delta:
|
||||
# 残差：Δ̂ 加回 a_{t-1}（history 末步，同为归一化空间）→ 绝对动作
|
||||
pred = history[0, -1] + pred
|
||||
|
||||
@@ -2347,6 +2347,176 @@ class TestLogRangeNorm(unittest.TestCase):
|
||||
validate_config_values(cfg.data, cfg.model)
|
||||
|
||||
|
||||
class TestOneHot(unittest.TestCase):
|
||||
"""多头 one-hot 输出：bins.py 编解码、软标签、dataset 目标、加权 CE、
|
||||
config 校验"""
|
||||
|
||||
# ---------------- bins.py ----------------
|
||||
|
||||
def test_encode_hard_onehot(self):
|
||||
from bins import encode, START, BINS_PER_COL
|
||||
a = np.zeros((2, 18), dtype=np.float32)
|
||||
a[0, 0] = 45.0
|
||||
a[0, 10] = 10.0
|
||||
a[0, 11] = 200.0
|
||||
a[1, 0] = 46.5
|
||||
o = encode(a)
|
||||
self.assertEqual(o.shape, (2, 1864))
|
||||
# 每行 12 头各一个 1，其余 0
|
||||
self.assertEqual(float(o.sum()), 24.0)
|
||||
self.assertEqual(np.argmax(o[0, START[0]:START[1]]), 45)
|
||||
self.assertEqual(np.argmax(o[1, START[0]:START[1]]), 46)
|
||||
# 保留位列无头，恒 0
|
||||
np.testing.assert_allclose(o[:, START[12]:], 0.0)
|
||||
|
||||
def test_encode_log_bins_spread_low_end(self):
|
||||
"""对数箱：10Hz 与 1000Hz 落在不同箱，低端不再挤在一起"""
|
||||
from bins import encode, START
|
||||
a = np.zeros((2, 18), dtype=np.float32)
|
||||
a[0, 10] = 10.0
|
||||
a[1, 10] = 1000.0
|
||||
o = encode(a)
|
||||
b0 = np.argmax(o[0, START[10]:START[11]])
|
||||
b1 = np.argmax(o[1, START[10]:START[11]])
|
||||
self.assertNotEqual(b0, b1)
|
||||
self.assertGreater(b1, b0)
|
||||
|
||||
def test_encode_soft_label(self):
|
||||
"""软标签：真值附近多箱非零且组内概率和为 1"""
|
||||
from bins import encode, START
|
||||
a = np.zeros((1, 18), dtype=np.float32)
|
||||
a[0, 0] = 45.0
|
||||
o = encode(a, sigma=1.0)
|
||||
head = o[0, START[0]:START[1]]
|
||||
self.assertGreater(np.count_nonzero(head), 1)
|
||||
self.assertAlmostEqual(float(head.sum()), 1.0, places=6)
|
||||
# 峰值在真值箱
|
||||
self.assertEqual(np.argmax(head), 45)
|
||||
|
||||
def test_decode_expectation(self):
|
||||
"""期望值解码：与真值一致时还原箱中心；子箱分辨率由概率内插"""
|
||||
from bins import encode, decode, START, BINS_PER_COL
|
||||
a = np.zeros((1, 18), dtype=np.float32)
|
||||
a[0, 0] = 45.0
|
||||
a[0, 10] = 10.0
|
||||
logits = np.log(encode(a) + 1e-8)
|
||||
d = decode(logits)
|
||||
self.assertAlmostEqual(float(d[0, 0]), 45.5, places=3) # 箱中心
|
||||
self.assertGreater(float(d[0, 10]), 10.0) # 对数箱中心偏右
|
||||
# 混合概率内插：45 箱 0.6 / 46 箱 0.4 → 45.9
|
||||
# （背景须 -1e9 使尾部概率可忽略；-10 背景的尾部质量 × 高箱中心
|
||||
# 会贡献 ~0.7°，属测试构造假象）
|
||||
logits2 = np.full((1, 1864), -1e9, dtype=np.float32)
|
||||
logits2[0, START[0] + 45] = np.log(0.6)
|
||||
logits2[0, START[0] + 46] = np.log(0.4)
|
||||
d2 = decode(logits2)
|
||||
self.assertAlmostEqual(float(d2[0, 0]), 45.9, places=4)
|
||||
|
||||
def test_decode_preserves_reserved_cols_zero(self):
|
||||
from bins import decode
|
||||
logits = np.zeros((2, 1864), dtype=np.float32)
|
||||
d = decode(logits)
|
||||
np.testing.assert_allclose(d[:, 12:], 0.0)
|
||||
|
||||
# ---------------- dataset 链路 ----------------
|
||||
|
||||
def test_dataset_onehot_targets(self):
|
||||
"""onehot 模式：目标为多组 one-hot，历史仍按范围归一化"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
base = Path(tmpdir)
|
||||
ep_dir = base / "episode_0001"
|
||||
images_dir = ep_dir / "images"
|
||||
images_dir.mkdir(parents=True)
|
||||
for i in range(8):
|
||||
Image.fromarray(np.random.randint(
|
||||
0, 255, (16, 16, 3), dtype=np.uint8)).save(
|
||||
images_dir / f"{i:04d}.jpg")
|
||||
actions = np.zeros((8, 12), dtype=np.float32)
|
||||
actions[:, 0] = 90.0
|
||||
np.save(ep_dir / "actions.npy", actions)
|
||||
ds = ActionDataset(tmpdir, history_len=3, action_dim=12,
|
||||
image_width=16, image_height=16, frame_len=3,
|
||||
normalize_actions=True,
|
||||
normalize_mode="fixed_range",
|
||||
onehot=True, onehot_sigma=0.0)
|
||||
s = ds[3]
|
||||
tgt = s["target"]
|
||||
self.assertEqual(tgt.shape, (1864,))
|
||||
from bins import START
|
||||
self.assertEqual(np.argmax(tgt.numpy()[START[0]:START[1]]), 90)
|
||||
# 历史仍为连续归一化值（90° → 0.5）
|
||||
np.testing.assert_allclose(
|
||||
s["action_history"].numpy()[:, 0], 0.5, atol=1e-6)
|
||||
|
||||
# ---------------- 加权 CE ----------------
|
||||
|
||||
def test_onehot_ce(self):
|
||||
from train import _onehot_ce
|
||||
from bins import encode, START
|
||||
logits = torch.zeros(2, 1864)
|
||||
a = np.zeros((2, 18), dtype=np.float32)
|
||||
a[0, 0] = 45.0
|
||||
a[1, 0] = 90.0
|
||||
targets = torch.from_numpy(encode(a, sigma=0.0))
|
||||
w = torch.tensor([1.0] * 10 + [0.2, 0.0] + [0.0] * 6)
|
||||
# 全部 12 头都预测正确 → CE≈0
|
||||
logits = torch.log(targets + 1e-9)
|
||||
loss = _onehot_ce(logits, targets, w, 18)
|
||||
self.assertLess(float(loss), 1e-3)
|
||||
# 第 2 行 head0 错箱（预测 bin 0，真值 90）→ 加权 CE 上升
|
||||
bad = logits.clone()
|
||||
bad[1, START[0]] = 0.0
|
||||
bad[1, START[0] + 90] = float(np.log(1e-9))
|
||||
loss_bad = _onehot_ce(bad, targets, w, 18)
|
||||
self.assertGreater(float(loss_bad), 0.5)
|
||||
|
||||
# ---------------- config 校验 ----------------
|
||||
|
||||
def test_onehot_cli_override_applied(self):
|
||||
"""--onehot CLI 标志经 apply_overrides 落到 cfg.data.onehot（回归：
|
||||
曾漏映射导致静默退化为回归训练）"""
|
||||
from config import build_argparser, apply_overrides
|
||||
parser = build_argparser("t")
|
||||
args = parser.parse_args(["--onehot", "--onehot-sigma", "2.0"])
|
||||
cfg = apply_overrides(Config(), args)
|
||||
self.assertTrue(cfg.data.onehot)
|
||||
self.assertEqual(cfg.data.onehot_sigma, 2.0)
|
||||
# 缺省不开启
|
||||
args = parser.parse_args([])
|
||||
cfg = apply_overrides(Config(), args)
|
||||
self.assertFalse(cfg.data.onehot)
|
||||
|
||||
def test_onehot_config_validation(self):
|
||||
cfg = Config()
|
||||
cfg.data.onehot = True
|
||||
cfg.data.normalize_actions = True
|
||||
cfg.data.normalize_mode = "fixed_range"
|
||||
validate_config_values(cfg.data, cfg.model, cfg.train)
|
||||
# 与 predict_delta 互斥
|
||||
cfg.train.predict_delta = True
|
||||
with self.assertRaises(ValueError):
|
||||
validate_config_values(cfg.data, cfg.model, cfg.train)
|
||||
# 要求范围类归一化
|
||||
cfg.train.predict_delta = False
|
||||
cfg.data.normalize_mode = "zscore"
|
||||
with self.assertRaises(ValueError):
|
||||
validate_config_values(cfg.data, cfg.model, cfg.train)
|
||||
|
||||
def test_onehot_model_output_shape(self):
|
||||
"""onehot 模型前向输出 TOTAL_BINS logits"""
|
||||
cfg = Config()
|
||||
cfg.data.onehot = True
|
||||
cfg.data.normalize_mode = "fixed_range"
|
||||
cfg.data.image_width = 32
|
||||
cfg.data.image_height = 32
|
||||
m = ActionPredictor(cfg.model, cfg.data, pretrained=False)
|
||||
from bins import TOTAL_BINS
|
||||
images = torch.randn(2, cfg.data.frame_len, 3, 32, 32)
|
||||
history = torch.randn(2, cfg.data.history_len, cfg.data.action_dim)
|
||||
out = m(images, history)
|
||||
self.assertEqual(out.shape, (2, TOTAL_BINS))
|
||||
|
||||
|
||||
class TestR2ImuModel(unittest.TestCase):
|
||||
"""R2 方案一模型：独立 IMU token 注入（序列 F+H+2）"""
|
||||
|
||||
|
||||
+71
-17
@@ -8,6 +8,8 @@ from tqdm import tqdm
|
||||
|
||||
from config import (Config, build_argparser, apply_overrides,
|
||||
validate_config_values, warn_init_from_mismatch)
|
||||
from action_norm import normalize_action, is_range_mode
|
||||
from bins import encode as bin_encode, decode as bin_decode, START, HEAD_COLS
|
||||
from model import ActionPredictor
|
||||
from dataset import ActionDataset
|
||||
from action_norm import default_column_weights
|
||||
@@ -48,6 +50,28 @@ def _collate(batch):
|
||||
return default_collate(batch)
|
||||
|
||||
|
||||
def _onehot_ce(logits: torch.Tensor, targets: torch.Tensor,
|
||||
weights: torch.Tensor, action_dim: int) -> torch.Tensor:
|
||||
"""多头加权交叉熵：每组 softmax 与 one-hot（或软标签）求 CE，按列权重加权
|
||||
|
||||
logits/targets: (B, TOTAL_BINS)；weights: (action_dim,) 列权重向量。
|
||||
软标签（sigma>0）各组概率和恒 1，CE 定义一致。
|
||||
"""
|
||||
ls = None
|
||||
head_ce = []
|
||||
for c in HEAD_COLS:
|
||||
if c >= action_dim:
|
||||
break
|
||||
s, e = START[c], START[c + 1]
|
||||
# 逐头 log_softmax（不能整行：各头 bin 数不同，且整行 max 会
|
||||
# 把其他头压到 -1e9 量级，CE 爆炸）
|
||||
head_ce.append(-(targets[:, s:e]
|
||||
* torch.log_softmax(logits[:, s:e], dim=-1)).sum(-1))
|
||||
head_ce = torch.stack(head_ce, dim=1) # (B, n_heads)
|
||||
w = weights[:head_ce.size(1)]
|
||||
return (head_ce * w).sum() / (w.sum() * head_ce.size(0))
|
||||
|
||||
|
||||
def _shift_history(history: torch.Tensor, new_action: torch.Tensor) -> torch.Tensor:
|
||||
"""动作历史左移一步、追加 new_action 到末尾（scheduled sampling 回传自预测用）。
|
||||
|
||||
@@ -58,7 +82,9 @@ def _shift_history(history: torch.Tensor, new_action: torch.Tensor) -> torch.Ten
|
||||
|
||||
def train_one_epoch(model, loader, optimizer, device, max_grad_norm: float = None,
|
||||
rollout_len: int = 1, ss_prob: float = 1.0,
|
||||
predict_delta: bool = False, weights: torch.Tensor = None):
|
||||
predict_delta: bool = False, weights: torch.Tensor = None,
|
||||
onehot: bool = False, range_mode: str = None,
|
||||
action_dim: int = 18):
|
||||
model.train()
|
||||
total_loss = 0.0
|
||||
count = 0
|
||||
@@ -71,7 +97,8 @@ def train_one_epoch(model, loader, optimizer, device, max_grad_norm: float = Non
|
||||
target = batch["target"].to(device)
|
||||
imu = batch["imu"].to(device) if "imu" in batch else None
|
||||
pred = model(images, history, imu)
|
||||
loss = _weighted_mse(pred, target, weights)
|
||||
loss = (_onehot_ce(pred, target, weights, action_dim) if onehot
|
||||
else _weighted_mse(pred, target, weights))
|
||||
bs = images.size(0)
|
||||
else:
|
||||
# 多步 scheduled sampling：首步必用真历史；后续步以 ss_prob 用真历史、
|
||||
@@ -89,11 +116,21 @@ def train_one_epoch(model, loader, optimizer, device, max_grad_norm: float = Non
|
||||
history = true_history
|
||||
# 否则沿用上一轮回传的 history（自预测构造）
|
||||
pred = model(images, history, imu)
|
||||
loss = loss + _weighted_mse(pred, target, weights)
|
||||
loss = loss + (_onehot_ce(pred, target, weights, action_dim)
|
||||
if onehot
|
||||
else _weighted_mse(pred, target, weights))
|
||||
bs = images.size(0)
|
||||
# 回传：残差模式先把 Δ̂ 加回 a_{t-1} 得到绝对动作，再移入历史
|
||||
# 回传：onehot 先期望值解码（原始空间）再按范围归一化；
|
||||
# 残差模式把 Δ̂ 加回 a_{t-1} 得到绝对动作；再移入历史
|
||||
with torch.no_grad():
|
||||
abs_action = prev_action + pred if predict_delta else pred
|
||||
if onehot:
|
||||
raw = bin_decode(pred.detach().cpu().numpy(),
|
||||
action_dim)
|
||||
abs_action = torch.from_numpy(normalize_action(
|
||||
raw, action_dim, mode=range_mode)).to(device)
|
||||
else:
|
||||
abs_action = (prev_action + pred
|
||||
if predict_delta else pred)
|
||||
history = _shift_history(true_history, abs_action.detach())
|
||||
loss = loss / len(batch["steps"])
|
||||
loss.backward()
|
||||
@@ -107,7 +144,8 @@ def train_one_epoch(model, loader, optimizer, device, max_grad_norm: float = Non
|
||||
|
||||
@torch.no_grad()
|
||||
def validate(model, loader, device, rollout_len: int = 1, predict_delta: bool = False,
|
||||
weights: torch.Tensor = None):
|
||||
weights: torch.Tensor = None, onehot: bool = False,
|
||||
action_dim: int = 18):
|
||||
model.eval()
|
||||
total_loss = 0.0
|
||||
count = 0
|
||||
@@ -118,7 +156,8 @@ def validate(model, loader, device, rollout_len: int = 1, predict_delta: bool =
|
||||
target = batch["target"].to(device)
|
||||
imu = batch["imu"].to(device) if "imu" in batch else None
|
||||
pred = model(images, history, imu)
|
||||
loss = _weighted_mse(pred, target, weights)
|
||||
loss = (_onehot_ce(pred, target, weights, action_dim) if onehot
|
||||
else _weighted_mse(pred, target, weights))
|
||||
bs = images.size(0)
|
||||
else:
|
||||
# 验证固定全真历史（确定性，不带 scheduled sampling 随机性），
|
||||
@@ -131,7 +170,9 @@ def validate(model, loader, device, rollout_len: int = 1, predict_delta: bool =
|
||||
target = step["target"].to(device)
|
||||
imu = step["imu"].to(device) if "imu" in step else None
|
||||
pred = model(images, history, imu)
|
||||
loss = loss + _weighted_mse(pred, target, weights)
|
||||
loss = loss + (_onehot_ce(pred, target, weights, action_dim)
|
||||
if onehot
|
||||
else _weighted_mse(pred, target, weights))
|
||||
bs = images.size(0)
|
||||
loss = loss / len(batch["steps"])
|
||||
total_loss += loss.item() * bs
|
||||
@@ -169,7 +210,8 @@ def main():
|
||||
cfg.data.action_stats_path, train_indices, cfg.data.frame_len,
|
||||
augment=cfg.train.augment,
|
||||
predict_delta=cfg.train.predict_delta, rollout_len=cfg.train.rollout_len,
|
||||
imu_dim=cfg.data.imu_dim, normalize_mode=cfg.data.normalize_mode)
|
||||
imu_dim=cfg.data.imu_dim, normalize_mode=cfg.data.normalize_mode,
|
||||
onehot=cfg.data.onehot, onehot_sigma=cfg.data.onehot_sigma)
|
||||
auto_stats = False
|
||||
if (cfg.data.normalize_actions and cfg.data.normalize_mode == "zscore"
|
||||
and not cfg.data.action_stats_path):
|
||||
@@ -186,7 +228,8 @@ def main():
|
||||
cfg.data.image_width, cfg.data.image_height, cfg.data.normalize_actions,
|
||||
cfg.data.action_stats_path, val_indices, cfg.data.frame_len,
|
||||
predict_delta=cfg.train.predict_delta, rollout_len=cfg.train.rollout_len,
|
||||
imu_dim=cfg.data.imu_dim, normalize_mode=cfg.data.normalize_mode)
|
||||
imu_dim=cfg.data.imu_dim, normalize_mode=cfg.data.normalize_mode,
|
||||
onehot=cfg.data.onehot, onehot_sigma=cfg.data.onehot_sigma)
|
||||
if auto_stats:
|
||||
# checkpoint 配置只存裸文件名：推理端按 checkpoint 目录解析相对路径
|
||||
cfg.data.action_stats_path = "action_stats.npy"
|
||||
@@ -219,6 +262,14 @@ def main():
|
||||
col_w = default_column_weights(cfg.data.action_dim)
|
||||
weights_t = torch.tensor(np.asarray(col_w, dtype=np.float32), device=device)
|
||||
print(f"Column weights: {[round(float(v), 4) for v in weights_t.tolist()]}")
|
||||
range_mode = (cfg.data.normalize_mode
|
||||
if (cfg.data.normalize_actions
|
||||
and is_range_mode(cfg.data.normalize_mode))
|
||||
else None)
|
||||
if cfg.data.onehot:
|
||||
print(f"One-hot 输出: {sum(1 for c in HEAD_COLS if c < cfg.data.action_dim)} 头"
|
||||
f" / {sum(START[c + 1] - START[c] for c in HEAD_COLS if c < cfg.data.action_dim)} bins"
|
||||
f"（sigma={cfg.data.onehot_sigma}）")
|
||||
|
||||
if args.init_from:
|
||||
init_ckpt = torch.load(args.init_from, map_location="cpu", weights_only=True)
|
||||
@@ -257,16 +308,19 @@ def main():
|
||||
|
||||
for epoch in range(start_epoch, cfg.train.epochs + 1):
|
||||
print(f"\nEpoch {epoch}/{cfg.train.epochs} lr={optimizer.param_groups[0]['lr']:.6f}")
|
||||
train_loss = train_one_epoch(model, train_loader, optimizer, device,
|
||||
cfg.train.max_grad_norm,
|
||||
rollout_len=cfg.train.rollout_len,
|
||||
ss_prob=cfg.train.ss_prob,
|
||||
predict_delta=cfg.train.predict_delta,
|
||||
weights=weights_t)
|
||||
train_loss = train_one_epoch(
|
||||
model, train_loader, optimizer, device,
|
||||
cfg.train.max_grad_norm,
|
||||
rollout_len=cfg.train.rollout_len,
|
||||
ss_prob=cfg.train.ss_prob,
|
||||
predict_delta=cfg.train.predict_delta,
|
||||
weights=weights_t, onehot=cfg.data.onehot,
|
||||
range_mode=range_mode, action_dim=cfg.data.action_dim)
|
||||
val_loss = validate(model, val_loader, device,
|
||||
rollout_len=cfg.train.rollout_len,
|
||||
predict_delta=cfg.train.predict_delta,
|
||||
weights=weights_t)
|
||||
weights=weights_t, onehot=cfg.data.onehot,
|
||||
action_dim=cfg.data.action_dim)
|
||||
scheduler.step()
|
||||
print(f" train_loss={train_loss:.6f} val_loss={val_loss:.6f}")
|
||||
|
||||
|
||||
Reference in new issue
Block a user