diff --git a/actpred/bins.py b/actpred/bins.py new file mode 100644 index 0000000..c7bce83 --- /dev/null +++ b/actpred/bins.py @@ -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 diff --git a/actpred/config.py b/actpred/config.py index f1dbc83..030f965 100644 --- a/actpred/config.py +++ b/actpred/config.py @@ -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"), diff --git a/actpred/dataset.py b/actpred/dataset.py index d66cf5e..1eb7c89 100644 --- a/actpred/dataset.py +++ b/actpred/dataset.py @@ -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]) diff --git a/actpred/model.py b/actpred/model.py index 704cff3..e5160c8 100644 --- a/actpred/model.py +++ b/actpred/model.py @@ -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 diff --git a/actpred/predict.py b/actpred/predict.py index 70e0bb3..9ed2480 100644 --- a/actpred/predict.py +++ b/actpred/predict.py @@ -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 diff --git a/actpred/test_actpred.py b/actpred/test_actpred.py index 7fc3df9..95c5827 100644 --- a/actpred/test_actpred.py +++ b/actpred/test_actpred.py @@ -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）""" diff --git a/actpred/train.py b/actpred/train.py index aae7efa..73254d5 100644 --- a/actpred/train.py +++ b/actpred/train.py @@ -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}")