feat(actpred): R1 残差预测+多步展开训练 与 R4 离线评估基准
R1(可选开关,默认关,向后兼容):
- --predict-delta:残差预测,target=Δ=a[t]-a[t-1](归一化空间),
推理端 predict 自动加回 a_{t-1}(惯性先验平滑动作)
- --rollout-len K + --ss-prob:多步展开 + scheduled sampling,
首步真历史、后续步按概率用真历史或自预测回传(detach 截断),
治 exposure bias;dataset 多步样本返回 {steps:[...]},自定义 collate
- 旧 checkpoint getattr 默认 False 兼容;pybot 闭环/部署链零改动
R4(新增 evaluate.py):
- 留出 episode 开环滚动预测,复刻闭环口径(图像真值滚动、动作自预测回写)
- 5 个闭环代理指标:overall_mse/error_growth(发散)/error_curve/
smoothness_ratio(平滑)/oob_rate(越界,数据经验界 min/max+margin)
- 支持 --compare A/B 对比与 --output JSON 报告
- 端到端验证:正确捕捉开环发散(单步 MSE 盲区)与残差平滑收益
新增 17 用例(R1×10 + R4×7),全量 100 passed 无回归
This commit is contained in:
1 parent
fe12dd7b33
commit
cc2f7bca79
7 files changed
+739
-42
No files matched your search
@@ -33,6 +33,13 @@ bash _run.sh train --data-dir ./data --epochs 100 --action-dim 18
|
||||
# 断点续训练(从 last.pt 恢复 model/optimizer/scheduler/epoch)
|
||||
bash _run.sh train --data-dir ./data --epochs 100 --resume checkpoints/last.pt
|
||||
|
||||
# R1 改进:残差预测 + 多步展开训练(可叠加)
|
||||
# --predict-delta 预测 Δ=actions[t]-actions[t-1](惯性先验,平滑动作;推理端自动加回 a_{t-1})
|
||||
# --rollout-len K 多步展开训练(K>1 时 scheduled sampling,治 exposure bias)
|
||||
# --ss-prob P scheduled sampling 真实历史概率(默认 1.0=全真;0.5 表示半数步用自预测回传)
|
||||
bash _run.sh train --data-dir ./data --epochs 100 --action-dim 18 \
|
||||
--predict-delta --rollout-len 4 --ss-prob 0.7
|
||||
|
||||
# 推理(--images 按时间顺序旧→新传入 frame_len 个图像路径,冷启动时可重复传同一帧;
|
||||
# --action-history 格式:冒号分隔各步,每步 action_dim 个值,共 history_len 步)
|
||||
bash _run.sh --model checkpoints/best_inference.pt \
|
||||
@@ -42,6 +49,12 @@ bash _run.sh --model checkpoints/best_inference.pt \
|
||||
# 导出 ONNX(如安装了 onnxruntime,导出后自动验证一致性)
|
||||
bash _run.sh export --model checkpoints/best_inference.pt --output model.onnx
|
||||
|
||||
# R4 离线评估:开环滚动预测,量化闭环代理指标(误差发散/平滑度/越界率)
|
||||
python evaluate.py --model checkpoints/best_inference.pt --data-dir ./data
|
||||
# A/B 对比两个 checkpoint(负值 = 主模型更优)
|
||||
python evaluate.py --model checkpoints/r1.pt --compare checkpoints/baseline.pt \
|
||||
--data-dir ./data --episodes 3 --horizon 50 --output report.json
|
||||
|
||||
# 运行单元测试
|
||||
python test_actpred.py
|
||||
```
|
||||
@@ -144,6 +157,26 @@ recorder.flush()
|
||||
| `--init-from` | 无 | 预训练 checkpoint 路径,作为微调起点(strict=False) |
|
||||
| `--resume` | 无 | 从 last.pt 断点续训练(恢复 model/optimizer/scheduler/epoch) |
|
||||
| `--save-config` | 无 | 保存配置到 JSON 文件 |
|
||||
| `--predict-delta` | False | R1-C 残差预测(target 改 Δ,推理端自动加回 a_{t-1}) |
|
||||
| `--rollout-len` | 1 | R1-B 多步展开步数 K(>1 时 scheduled sampling) |
|
||||
| `--ss-prob` | 1.0 | scheduled sampling 真实历史概率(仅 rollout-len>1 生效) |
|
||||
|
||||
## 离线评估(R4)
|
||||
|
||||
`evaluate.py` 在留出 episode 上**开环滚动预测**(复现 pybot 闭环口径:图像用真值滚动、
|
||||
动作用自预测回写),输出闭环代理指标。单步 MSE(teacher forcing)看不见闭环误差累积,
|
||||
本评估补上这一盲区:
|
||||
|
||||
| 指标 | 含义 | 判读 |
|
||||
|------|------|------|
|
||||
| `overall_mse` | 全部滚动步的平均 MSE | 越低越好 |
|
||||
| `error_growth` | 末段(后25%)误差 / 首段(前25%)误差 | >1 表示误差随步数发散(闭环不稳) |
|
||||
| `error_curve` | 误差随步数的增长曲线 | 看发散拐点 |
|
||||
| `smoothness_ratio` | 预测平滑度 / 真值平滑度 | <1 表示比真值更平滑(残差模式的预期收益) |
|
||||
| `oob_rate` | 预测落在训练数据经验范围外的比例 | 越高越危险(舵机越界风险) |
|
||||
|
||||
物理可行性用数据分布的经验界(min/max + 10% margin)而非硬编码范围——18 列混含
|
||||
舵机角度/频率/延时/预留,无统一硬限位。`--compare` 用于 A/B 对比(如 R1 开/关)。
|
||||
|
||||
## 部署注意
|
||||
|
||||
|
||||
@@ -41,6 +41,14 @@ class TrainConfig:
|
||||
seed: int = 42
|
||||
max_grad_norm: float = 1.0
|
||||
augment: bool = False
|
||||
# R1-C 残差预测:True 时 target 改为 Δ=actions[t]-actions[t-1],
|
||||
# 推理端 â_t = a_{t-1} + Δ̂(惯性先验,平滑动作)。默认 False 兼容旧 checkpoint。
|
||||
predict_delta: bool = False
|
||||
# R1-B 多步展开训练:rollout_len>1 时 dataset 返回 K 步连续序列,
|
||||
# 训练用 scheduled sampling(前 ss_prob 比例步用真历史,其余用自预测回传),
|
||||
# 治 exposure bias。rollout_len=1 时退化为现状单步。
|
||||
rollout_len: int = 1
|
||||
ss_prob: float = 1.0
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -90,6 +98,13 @@ def build_argparser(description: str = "") -> argparse.ArgumentParser:
|
||||
parser.add_argument("--max-grad-norm", type=float, default=None)
|
||||
parser.add_argument("--augment", action=argparse.BooleanOptionalAction, default=None,
|
||||
help="训练时启用安全图像增强(ColorJitter:亮度/对比度/饱和度/色调)")
|
||||
parser.add_argument("--predict-delta", action=argparse.BooleanOptionalAction, default=None,
|
||||
help="R1-C 残差预测:target 改为 Δ=actions[t]-actions[t-1],"
|
||||
"推理端 â_t=a_{t-1}+Δ̂(惯性先验平滑动作,默认关兼容旧 checkpoint)")
|
||||
parser.add_argument("--rollout-len", type=int, default=None,
|
||||
help="R1-B 多步展开训练步数 K(默认 1=单步;K>1 时 scheduled sampling)")
|
||||
parser.add_argument("--ss-prob", type=float, default=None,
|
||||
help="scheduled sampling 真实历史概率(仅 rollout-len>1 生效,默认 1.0=全真)")
|
||||
parser.add_argument("--init-from", type=str, default=None,
|
||||
help="预训练 checkpoint 路径,作为微调起点(strict=False 加载兼容权重)")
|
||||
parser.add_argument("--resume", type=str, default=None,
|
||||
@@ -142,6 +157,8 @@ def validate_config_values(data_cfg: DataConfig, model_cfg: ModelConfig,
|
||||
("patience", train_cfg.patience, 0, 100000, int),
|
||||
("num_workers", train_cfg.num_workers, 0, 256, int),
|
||||
("max_grad_norm", train_cfg.max_grad_norm, 0.0, 1e6, float),
|
||||
("rollout_len", train_cfg.rollout_len, 1, 1024, int),
|
||||
("ss_prob", train_cfg.ss_prob, 0.0, 1.0, float),
|
||||
]
|
||||
for name, val, lo, hi, typ in t_checks:
|
||||
if typ is int:
|
||||
@@ -236,6 +253,9 @@ def apply_overrides(cfg: Config, args: argparse.Namespace) -> Config:
|
||||
"seed": ("train", "seed"),
|
||||
"max_grad_norm": ("train", "max_grad_norm"),
|
||||
"augment": ("train", "augment"),
|
||||
"predict_delta": ("train", "predict_delta"),
|
||||
"rollout_len": ("train", "rollout_len"),
|
||||
"ss_prob": ("train", "ss_prob"),
|
||||
}
|
||||
args_dict = vars(args)
|
||||
for arg_name, (section, field) in mapping.items():
|
||||
|
||||
+63
-19
@@ -21,13 +21,18 @@ class ActionDataset(Dataset):
|
||||
image_width: int = 640, image_height: int = 480,
|
||||
normalize_actions: bool = False,
|
||||
action_stats_path: str = None, episode_indices: list = None,
|
||||
frame_len: int = 3, augment: bool = False):
|
||||
frame_len: int = 3, augment: bool = False,
|
||||
predict_delta: bool = False, rollout_len: int = 1):
|
||||
self.history_len = history_len
|
||||
self.action_dim = action_dim
|
||||
self.image_width = image_width
|
||||
self.image_height = image_height
|
||||
self.frame_len = frame_len
|
||||
self.augment = augment
|
||||
# R1-C 残差预测:target 改为 Δ=actions[t]-actions[t-1](归一化空间内求差)
|
||||
self.predict_delta = predict_delta
|
||||
# R1-B 多步展开:rollout_len>1 时每个样本返回 K 步连续 (images,history,target)
|
||||
self.rollout_len = max(1, rollout_len)
|
||||
|
||||
self.transform = transforms.Compose([
|
||||
transforms.Resize((image_height, image_width)),
|
||||
@@ -39,6 +44,8 @@ class ActionDataset(Dataset):
|
||||
if episode_indices is not None:
|
||||
episodes = [episodes[i] for i in episode_indices]
|
||||
|
||||
# 样本存 (image_files, actions, t_start):t_start 是 rollout 第一步的目标时刻,
|
||||
# __getitem__ 按 rollout_len 展开 K 步。存原始数组引用避免重复加载/拷贝。
|
||||
self.samples = []
|
||||
for ep_dir in episodes:
|
||||
actions_path = ep_dir / "actions.npy"
|
||||
@@ -58,11 +65,12 @@ class ActionDataset(Dataset):
|
||||
print(f"WARNING: {ep_dir} 动作帧数 {n_frames} 与图像数 {n_images} 不一致,"
|
||||
f"按较少者({min(n_frames, n_images)})使用,请检查录制完整性")
|
||||
n = min(n_frames, n_images)
|
||||
# 每个样本:帧 [t-frame_len+1, t](旧→新)+ 动作历史 [t-history_len, t) → 动作 t
|
||||
# 第 k 步(k=0..K-1)目标时刻 t=t_start+k:帧 [t-frame_len+1, t]、
|
||||
# 历史 [t-history_len, t)、目标 actions[t]。t_start 下界保证第 0 步窗口完整,
|
||||
# 上界保证第 K-1 步目标不越界。
|
||||
start = max(history_len, frame_len - 1)
|
||||
for t in range(start, n):
|
||||
self.samples.append((image_files[t - frame_len + 1: t + 1],
|
||||
actions[t - history_len:t], actions[t]))
|
||||
for t_start in range(start, n - (self.rollout_len - 1)):
|
||||
self.samples.append((image_files, actions, t_start))
|
||||
|
||||
self._action_mean = None
|
||||
self._action_std = None
|
||||
@@ -83,24 +91,34 @@ class ActionDataset(Dataset):
|
||||
raise ValueError(
|
||||
f"数据集为空(0 样本),无法计算归一化统计量。"
|
||||
f"请检查 data_dir={data_dir} 中的 episode 数据。")
|
||||
all_actions = np.stack(
|
||||
[np.concatenate([s[1], s[2][None]]) for s in self.samples]
|
||||
).reshape(-1, self.action_dim)
|
||||
# 归一化统计基于绝对动作(与 Δ 无关):覆盖所有样本 rollout 范围内
|
||||
# 出现的 history 与 target 动作,保持与旧版一致的统计口径。
|
||||
all_actions = np.stack([
|
||||
s[1][s[2] - self.history_len: s[2] + self.rollout_len]
|
||||
for s in self.samples
|
||||
]).reshape(-1, self.action_dim)
|
||||
self._action_mean = all_actions.mean(axis=0)
|
||||
self._action_std = all_actions.std(axis=0) + 1e-8
|
||||
|
||||
def __len__(self):
|
||||
return len(self.samples)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
image_paths, action_history, target = self.samples[idx]
|
||||
frames = []
|
||||
def _norm(self, a):
|
||||
"""z-score 归一化(numpy → torch 前在 numpy 空间做,保持 float32)"""
|
||||
if self._action_mean is None:
|
||||
return a
|
||||
return (a - self._action_mean) / self._action_std
|
||||
|
||||
def _load_frames(self, image_files, t):
|
||||
"""加载时刻 t 的 F 帧窗口 [t-frame_len+1, t](旧→新),augment 对整个窗口一致"""
|
||||
paths = image_files[t - self.frame_len + 1: t + 1]
|
||||
if self.augment:
|
||||
bf = random.uniform(0.8, 1.2)
|
||||
cf = random.uniform(0.8, 1.2)
|
||||
sf = random.uniform(0.8, 1.2)
|
||||
hf = random.uniform(-0.05, 0.05)
|
||||
for p in image_paths:
|
||||
frames = []
|
||||
for p in paths:
|
||||
image = Image.open(p).convert("RGB")
|
||||
if self.augment:
|
||||
image = TF.adjust_brightness(image, bf)
|
||||
@@ -108,13 +126,39 @@ class ActionDataset(Dataset):
|
||||
image = TF.adjust_saturation(image, sf)
|
||||
image = TF.adjust_hue(image, hf)
|
||||
frames.append(self.transform(image))
|
||||
images = torch.stack(frames) # (F, 3, S, S)
|
||||
action_history = torch.tensor(action_history, dtype=torch.float32)
|
||||
target = torch.tensor(target, dtype=torch.float32)
|
||||
if self._action_mean is not None:
|
||||
action_history = (action_history - self._action_mean) / self._action_std
|
||||
target = (target - self._action_mean) / self._action_std
|
||||
return {"images": images, "action_history": action_history, "target": target}
|
||||
return torch.stack(frames) # (F, 3, H, W)
|
||||
|
||||
def _step_target(self, actions, t):
|
||||
"""时刻 t 的监督目标(归一化空间);predict_delta 时为 Δ=actions[t]-actions[t-1]"""
|
||||
if self.predict_delta:
|
||||
return self._norm(actions[t]) - self._norm(actions[t - 1])
|
||||
return self._norm(actions[t])
|
||||
|
||||
def __getitem__(self, idx):
|
||||
image_files, actions, t_start = self.samples[idx]
|
||||
if self.rollout_len == 1:
|
||||
# 单步:与旧版输出结构完全一致
|
||||
images = self._load_frames(image_files, t_start)
|
||||
history = torch.tensor(self._norm(
|
||||
actions[t_start - self.history_len: t_start]), dtype=torch.float32)
|
||||
target = torch.tensor(self._step_target(actions, t_start),
|
||||
dtype=torch.float32)
|
||||
return {"images": images, "action_history": history, "target": target}
|
||||
|
||||
# 多步:返回 K 步序列,scheduled sampling 在 train 循环内按步取用
|
||||
steps = []
|
||||
for k in range(self.rollout_len):
|
||||
t = t_start + k
|
||||
steps.append({
|
||||
"images": self._load_frames(image_files, t),
|
||||
# 真实历史(scheduled sampling 的真值来源;首步必用)
|
||||
"action_history": torch.tensor(
|
||||
self._norm(actions[t - self.history_len: t]), dtype=torch.float32),
|
||||
"target": torch.tensor(self._step_target(actions, t), dtype=torch.float32),
|
||||
# 上一步真实动作(Δ 残差的基准 a_{t-1},归一化空间)
|
||||
"prev_action": torch.tensor(self._norm(actions[t - 1]), dtype=torch.float32),
|
||||
})
|
||||
return {"steps": steps}
|
||||
|
||||
def get_action_stats(self):
|
||||
if self._action_mean is not None:
|
||||
|
||||
@@ -0,0 +1,232 @@
|
||||
"""R4 离线评估基准:开环滚动预测,量化闭环代理指标
|
||||
|
||||
现状只有开环单步 MSE(teacher forcing),无法反映真实闭环表现。本脚本在留出
|
||||
episode 上复现 pybot 闭环的开环滚动(图像用真值滚动、动作用自预测回写),
|
||||
输出三类闭环代理指标:
|
||||
|
||||
累积误差 逐步 MSE 及随步数的增长曲线(衡量误差累积/发散)
|
||||
平滑度 预测动作相邻差分方差(越低越平滑,对物理舵机越友好),与真值对比
|
||||
物理可行性 预测动作落在训练数据经验范围外的比例(越界率,越高越危险)
|
||||
|
||||
用法:
|
||||
python evaluate.py --model checkpoints/best_inference.pt --data-dir ./data
|
||||
python evaluate.py --model A.pt --compare B.pt --data-dir ./data # A/B 对比
|
||||
python evaluate.py --model m.pt --data-dir ./data --episodes 3 --horizon 50 \
|
||||
--output report.json
|
||||
"""
|
||||
import argparse
|
||||
import json
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from dataset import ActionDataset
|
||||
from predict import ActionPredictorInference
|
||||
|
||||
|
||||
def _episode_ids(data_dir):
|
||||
return [d.name for d in ActionDataset._valid_episode_dirs(data_dir)]
|
||||
|
||||
|
||||
def _load_episode(data_dir, ep_name):
|
||||
"""加载单个 episode 的图像路径(自然序)与动作真值 (T, D)"""
|
||||
ep_dir = [d for d in ActionDataset._valid_episode_dirs(data_dir)
|
||||
if d.name == ep_name][0]
|
||||
from dataset import _natural_key
|
||||
images_dir = ep_dir / "images"
|
||||
image_files = sorted(list(images_dir.glob("*.jpg")) +
|
||||
list(images_dir.glob("*.png")), key=_natural_key)
|
||||
actions = np.load(ep_dir / "actions.npy").astype(np.float32)
|
||||
n = min(len(image_files), len(actions))
|
||||
return image_files[:n], actions[:n]
|
||||
|
||||
|
||||
def _action_bounds(data_dir, action_dim, margin_ratio=0.1):
|
||||
"""从全数据集统计动作经验范围(min/max 加 margin),作为物理可行性判据。
|
||||
|
||||
actpred 侧无硬限位定义(18 列混含舵机角度/频率/延时/预留),硬编码范围不可靠;
|
||||
用数据分布的经验界更通用。margin 默认 10% 容忍合理外推。
|
||||
"""
|
||||
mins = np.full(action_dim, np.inf, dtype=np.float32)
|
||||
maxs = np.full(action_dim, -np.inf, dtype=np.float32)
|
||||
for ep_dir in ActionDataset._valid_episode_dirs(data_dir):
|
||||
a = np.load(ep_dir / "actions.npy").astype(np.float32)
|
||||
mins = np.minimum(mins, a.min(axis=0))
|
||||
maxs = np.maximum(maxs, a.max(axis=0))
|
||||
span = np.maximum(maxs - mins, 1e-6)
|
||||
return mins - margin_ratio * span, maxs + margin_ratio * span
|
||||
|
||||
|
||||
@np.errstate(all="ignore")
|
||||
def rollout_episode(predictor, image_files, actions, horizon=None,
|
||||
lo=None, hi=None):
|
||||
"""单 episode 开环滚动,复现闭环口径(图像真值滚动、动作自预测回写)。
|
||||
|
||||
从 t0 = max(history_len, frame_len-1) 起步(窗口完整处),每步:
|
||||
帧窗 = 真值图像 [t-frame_len+1, t](旧→新)
|
||||
历史 = 前 H 步动作(首步用真值填充,之后用自预测回写,复刻 run_loop)
|
||||
预测 â_t → 与真值 a_t 比较 → 回写历史
|
||||
返回逐步记录 dict。horizon=None 表示滚到 episode 末尾。
|
||||
"""
|
||||
F, H = predictor.frame_len, predictor.history_len
|
||||
t0 = max(H, F - 1)
|
||||
n = len(actions)
|
||||
if horizon is None:
|
||||
horizon = n - t0
|
||||
horizon = min(horizon, n - t0)
|
||||
if horizon <= 0:
|
||||
return None
|
||||
|
||||
cache = {}
|
||||
|
||||
def load(t):
|
||||
if t not in cache:
|
||||
cache[t] = np.array(Image.open(image_files[t]).convert("RGB"))
|
||||
return cache[t]
|
||||
|
||||
history = [actions[t0 - H + k].copy() for k in range(H)] # 真值冷启动
|
||||
steps = []
|
||||
for k in range(horizon):
|
||||
t = t0 + k
|
||||
frames = [load(t - F + 1 + i) for i in range(F)]
|
||||
pred = np.asarray(
|
||||
predictor.predict(frames, np.array(history)),
|
||||
dtype=np.float32).reshape(-1)
|
||||
true = actions[t]
|
||||
err = pred - true
|
||||
rec = {
|
||||
"t": int(t),
|
||||
"mse": float(np.mean(err ** 2)),
|
||||
"pred": pred,
|
||||
"true": true,
|
||||
}
|
||||
if lo is not None:
|
||||
rec["oob"] = bool(np.any((pred < lo) | (pred > hi)))
|
||||
steps.append(rec)
|
||||
history = history[1:] + [pred.copy()] # 自预测回写(复刻 run_loop)
|
||||
return steps
|
||||
|
||||
|
||||
def _smoothness(seq):
|
||||
"""相邻差分方差(多维取均值);序列 <2 返回 0。越低越平滑。"""
|
||||
seq = np.asarray(seq, dtype=np.float32)
|
||||
if len(seq) < 2:
|
||||
return 0.0
|
||||
return float(np.mean(np.var(np.diff(seq, axis=0), axis=0)))
|
||||
|
||||
|
||||
def summarize(all_steps, lo, hi):
|
||||
"""汇总多个 episode 的滚动记录为报告指标"""
|
||||
if not all_steps:
|
||||
raise ValueError("无可评估的滚动步(episode 太短或 horizon 过小)")
|
||||
flat = [s for steps in all_steps for s in steps]
|
||||
mses = np.array([s["mse"] for s in flat], dtype=np.float32)
|
||||
|
||||
# 误差随步数增长曲线:按相对步位置对齐各 episode,取每位置的平均 MSE
|
||||
max_len = max(len(s) for s in all_steps)
|
||||
curve = []
|
||||
for k in range(max_len):
|
||||
vals = [s[k]["mse"] for s in all_steps if k < len(s)]
|
||||
if vals:
|
||||
curve.append(float(np.mean(vals)))
|
||||
|
||||
preds = np.array([s["pred"] for s in flat], dtype=np.float32)
|
||||
trues = np.array([s["true"] for s in flat], dtype=np.float32)
|
||||
oob = [s.get("oob", False) for s in flat]
|
||||
|
||||
return {
|
||||
"n_episodes": len(all_steps),
|
||||
"n_steps": len(flat),
|
||||
"overall_mse": float(np.mean(mses)),
|
||||
# 末段(后 25% 步)相对首段(前 25%)的误差比,>1 表示误差随步数发散
|
||||
"error_growth": float(np.mean(mses[int(len(mses) * 0.75):]) /
|
||||
max(np.mean(mses[:max(1, int(len(mses) * 0.25))]),
|
||||
1e-12)),
|
||||
"error_curve": curve,
|
||||
"pred_smoothness": _smoothness(preds),
|
||||
"true_smoothness": _smoothness(trues),
|
||||
# 平滑度比:<1 表示预测比真值还平滑(好),>1 表示预测更抖(坏)
|
||||
"smoothness_ratio": _smoothness(preds) / max(_smoothness(trues), 1e-12),
|
||||
"oob_rate": float(np.mean(oob)) if lo is not None else None,
|
||||
}
|
||||
|
||||
|
||||
def evaluate(checkpoint, data_dir, episodes=None, horizon=None, device=None):
|
||||
predictor = ActionPredictorInference(checkpoint, device)
|
||||
ep_names = _episode_ids(data_dir)
|
||||
if episodes is not None:
|
||||
ep_names = ep_names[-episodes:] # 取尾部 episode 作为留出集
|
||||
lo, hi = _action_bounds(data_dir, predictor.action_dim)
|
||||
all_steps = []
|
||||
for name in ep_names:
|
||||
image_files, actions = _load_episode(data_dir, name)
|
||||
steps = rollout_episode(predictor, image_files, actions, horizon, lo, hi)
|
||||
if steps:
|
||||
all_steps.append(steps)
|
||||
rep = summarize(all_steps, lo, hi)
|
||||
rep["checkpoint"] = checkpoint
|
||||
rep["frame_len"] = predictor.frame_len
|
||||
rep["history_len"] = predictor.history_len
|
||||
rep["predict_delta"] = getattr(predictor, "_predict_delta", False)
|
||||
return rep
|
||||
|
||||
|
||||
def _fmt(rep):
|
||||
lines = [
|
||||
f"checkpoint : {rep['checkpoint']}",
|
||||
f"predict_delta : {rep['predict_delta']}",
|
||||
f"episodes/steps : {rep['n_episodes']} / {rep['n_steps']}",
|
||||
f"overall_mse : {rep['overall_mse']:.6f}",
|
||||
f"error_growth : {rep['error_growth']:.3f} (>1 = 误差随步数发散)",
|
||||
f"smoothness_ratio: {rep['smoothness_ratio']:.3f} (<1 = 比真值更平滑)",
|
||||
f" pred/true : {rep['pred_smoothness']:.4f} / {rep['true_smoothness']:.4f}",
|
||||
]
|
||||
if rep["oob_rate"] is not None:
|
||||
lines.append(f"oob_rate : {rep['oob_rate']:.4f} (越界率, 越低越安全)")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def main():
|
||||
p = argparse.ArgumentParser(description="actpred 离线闭环代理评估(R4)")
|
||||
p.add_argument("--model", required=True, help="主 checkpoint 路径")
|
||||
p.add_argument("--data-dir", required=True, help="episode 数据目录")
|
||||
p.add_argument("--compare", default=None, help="可选对比 checkpoint(A/B)")
|
||||
p.add_argument("--episodes", type=int, default=None,
|
||||
help="评估用尾部 episode 数(默认全部)")
|
||||
p.add_argument("--horizon", type=int, default=None,
|
||||
help="每 episode 滚动步数上限(默认滚到末尾)")
|
||||
p.add_argument("--device", default=None)
|
||||
p.add_argument("--output", default=None, help="保存 JSON 报告路径")
|
||||
args = p.parse_args()
|
||||
|
||||
rep = evaluate(args.model, args.data_dir, args.episodes, args.horizon,
|
||||
args.device)
|
||||
print("== 主模型 ==")
|
||||
print(_fmt(rep))
|
||||
|
||||
result = {"main": rep}
|
||||
if args.compare:
|
||||
rep_b = evaluate(args.compare, args.data_dir, args.episodes,
|
||||
args.horizon, args.device)
|
||||
print("\n== 对比模型 ==")
|
||||
print(_fmt(rep_b))
|
||||
# 关键指标对比(负值=主模型更优)
|
||||
result["compare"] = rep_b
|
||||
result["delta"] = {
|
||||
"overall_mse": rep["overall_mse"] - rep_b["overall_mse"],
|
||||
"error_growth": rep["error_growth"] - rep_b["error_growth"],
|
||||
"smoothness_ratio": rep["smoothness_ratio"] - rep_b["smoothness_ratio"],
|
||||
}
|
||||
d = result["delta"]
|
||||
print("\n== 主 − 对比 (负值=主模型更优) ==")
|
||||
print(f"overall_mse : {d['overall_mse']:+.6f}")
|
||||
print(f"error_growth : {d['error_growth']:+.3f}")
|
||||
print(f"smoothness_ratio: {d['smoothness_ratio']:+.3f}")
|
||||
|
||||
if args.output:
|
||||
with open(args.output, "w", encoding="utf-8") as f:
|
||||
json.dump(result, f, ensure_ascii=False, indent=2)
|
||||
print(f"\n报告已保存: {args.output}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+8
-1
@@ -29,6 +29,10 @@ class ActionPredictorInference:
|
||||
self.action_dim = self._data_cfg.action_dim
|
||||
self.history_len = self._data_cfg.history_len
|
||||
self.frame_len = self._data_cfg.frame_len
|
||||
# R1-C 残差预测:checkpoint 记录 predict_delta=True 时,模型输出 Δ̂,
|
||||
# 需在归一化空间加回 a_{t-1}(history 末步)得到绝对动作。
|
||||
# 旧 checkpoint 无该字段,getattr 默认 False,行为与旧版一致。
|
||||
self._predict_delta = getattr(cfg.train, "predict_delta", False)
|
||||
|
||||
self._transform = transforms.Compose([
|
||||
transforms.Resize((self._data_cfg.image_height,
|
||||
@@ -84,7 +88,10 @@ class ActionPredictorInference:
|
||||
if self._action_mean is not None:
|
||||
history = (history - self._action_mean) / self._action_std
|
||||
history = history.unsqueeze(0).to(self._device)
|
||||
pred = self._model(img_tensor, history).squeeze(0)
|
||||
pred = self._model(img_tensor, history).squeeze(0) # 归一化空间
|
||||
if self._predict_delta:
|
||||
# 残差:Δ̂ 加回 a_{t-1}(history 末步,同为归一化空间)→ 绝对动作
|
||||
pred = history[0, -1] + pred
|
||||
if self._action_mean is not None:
|
||||
pred = pred * self._action_std + self._action_mean
|
||||
return pred.cpu().numpy()
|
||||
|
||||
@@ -1509,5 +1509,293 @@ class TestDfcCompile(unittest.TestCase):
|
||||
self.assertIn("n/input_layer2", str(ctx.exception))
|
||||
|
||||
|
||||
class TestR1PredictDelta(unittest.TestCase):
|
||||
"""R1-C 残差预测:dataset Δ target 构造与 predict 逆变换"""
|
||||
|
||||
@staticmethod
|
||||
def _create_episode(base_dir, actions, image_size=16):
|
||||
ep_dir = base_dir / "episode_0001"
|
||||
images_dir = ep_dir / "images"
|
||||
images_dir.mkdir(parents=True, exist_ok=True)
|
||||
n = len(actions)
|
||||
for i in range(n):
|
||||
Image.fromarray(np.random.randint(0, 255, (image_size, image_size, 3),
|
||||
dtype=np.uint8)).save(
|
||||
images_dir / f"{i:04d}.jpg")
|
||||
np.save(ep_dir / "actions.npy", np.array(actions, dtype=np.float32))
|
||||
|
||||
def test_delta_target_unnormalized(self):
|
||||
"""不归一化时 target 应为 Δ=actions[t]-actions[t-1]"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 线性递增动作,Δ 恒为 [1,1]
|
||||
actions = [[float(i), float(i)] for i in range(8)]
|
||||
self._create_episode(Path(tmpdir), actions)
|
||||
ds = ActionDataset(tmpdir, history_len=2, action_dim=2,
|
||||
image_width=16, image_height=16, frame_len=2,
|
||||
predict_delta=True)
|
||||
sample = ds[0]
|
||||
# 第 0 样本 t_start=2,target = a[2]-a[1] = [1,1]
|
||||
np.testing.assert_allclose(sample["target"].numpy(), [1.0, 1.0],
|
||||
atol=1e-6)
|
||||
|
||||
def test_absolute_target_when_disabled(self):
|
||||
"""predict_delta=False 时 target 应为绝对动作 actions[t]"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
actions = [[float(i), float(i)] for i in range(8)]
|
||||
self._create_episode(Path(tmpdir), actions)
|
||||
ds = ActionDataset(tmpdir, history_len=2, action_dim=2,
|
||||
image_width=16, image_height=16, frame_len=2,
|
||||
predict_delta=False)
|
||||
# t_start=2,target = a[2] = [2,2]
|
||||
np.testing.assert_allclose(ds[0]["target"].numpy(), [2.0, 2.0],
|
||||
atol=1e-6)
|
||||
|
||||
def test_delta_target_normalized_space(self):
|
||||
"""归一化开启时 Δ 应在归一化空间求差(先 z-score 再相减)"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
rng = np.random.RandomState(0)
|
||||
actions = rng.randn(12, 3).astype(np.float32)
|
||||
self._create_episode(Path(tmpdir), actions)
|
||||
ds = ActionDataset(tmpdir, history_len=2, action_dim=3,
|
||||
image_width=16, image_height=16, frame_len=2,
|
||||
normalize_actions=True, predict_delta=True)
|
||||
mean, std = ds._action_mean, ds._action_std
|
||||
sample = ds[0]
|
||||
t = 2 # 第 0 样本 t_start
|
||||
expected = (actions[t] - mean) / std - (actions[t - 1] - mean) / std
|
||||
np.testing.assert_allclose(sample["target"].numpy(), expected, atol=1e-5)
|
||||
|
||||
|
||||
class TestR1Rollout(unittest.TestCase):
|
||||
"""R1-B 多步展开:dataset K 步序列结构与窗口对齐"""
|
||||
|
||||
@staticmethod
|
||||
def _create_episode(base_dir, n, action_dim, image_size=16):
|
||||
ep_dir = base_dir / "episode_0001"
|
||||
images_dir = ep_dir / "images"
|
||||
images_dir.mkdir(parents=True, exist_ok=True)
|
||||
for i in range(n):
|
||||
Image.fromarray(np.random.randint(0, 255, (image_size, image_size, 3),
|
||||
dtype=np.uint8)).save(
|
||||
images_dir / f"{i:04d}.jpg")
|
||||
actions = np.arange(n * action_dim, dtype=np.float32).reshape(n, action_dim)
|
||||
np.save(ep_dir / "actions.npy", actions)
|
||||
return actions
|
||||
|
||||
def test_rollout_sample_count(self):
|
||||
"""rollout_len=K 时样本数 = n - start - (K-1)"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
self._create_episode(Path(tmpdir), n=10, action_dim=2)
|
||||
ds = ActionDataset(tmpdir, history_len=2, action_dim=2,
|
||||
image_width=16, image_height=16, frame_len=2,
|
||||
rollout_len=3)
|
||||
# start=max(2,1)=2,n-(K-1)=8,样本数=8-2=6
|
||||
self.assertEqual(len(ds), 6)
|
||||
|
||||
def test_rollout_structure_and_alignment(self):
|
||||
"""多步样本返回 K 步,每步 history/target/prev_action 与时刻对齐"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
actions = self._create_episode(Path(tmpdir), n=10, action_dim=2)
|
||||
ds = ActionDataset(tmpdir, history_len=2, action_dim=2,
|
||||
image_width=16, image_height=16, frame_len=2,
|
||||
rollout_len=3)
|
||||
sample = ds[0]
|
||||
self.assertIn("steps", sample)
|
||||
steps = sample["steps"]
|
||||
self.assertEqual(len(steps), 3)
|
||||
for k, step in enumerate(steps):
|
||||
t = 2 + k # t_start=2
|
||||
self.assertEqual(step["images"].shape, (2, 3, 16, 16))
|
||||
self.assertEqual(step["action_history"].shape, (2, 2))
|
||||
self.assertEqual(step["target"].shape, (2,))
|
||||
self.assertEqual(step["prev_action"].shape, (2,))
|
||||
# 未归一化:history=actions[t-2:t]、target=actions[t]、prev=actions[t-1]
|
||||
np.testing.assert_allclose(step["action_history"].numpy(),
|
||||
actions[t - 2:t], atol=1e-6)
|
||||
np.testing.assert_allclose(step["target"].numpy(),
|
||||
actions[t], atol=1e-6)
|
||||
np.testing.assert_allclose(step["prev_action"].numpy(),
|
||||
actions[t - 1], atol=1e-6)
|
||||
|
||||
def test_rollout_delta_prev_action_is_base(self):
|
||||
"""残差+多步:prev_action 应为 prev+target 的基准(target=Δ 时 prev+Δ=abs)"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
actions = self._create_episode(Path(tmpdir), n=10, action_dim=2)
|
||||
ds = ActionDataset(tmpdir, history_len=2, action_dim=2,
|
||||
image_width=16, image_height=16, frame_len=2,
|
||||
predict_delta=True, rollout_len=3)
|
||||
step = ds[0]["steps"][1] # t=3
|
||||
# prev + Δ 应还原绝对动作 actions[t]
|
||||
recon = step["prev_action"].numpy() + step["target"].numpy()
|
||||
np.testing.assert_allclose(recon, actions[3], atol=1e-6)
|
||||
|
||||
def test_single_step_has_no_steps_key(self):
|
||||
"""rollout_len=1 时保持旧版扁平结构(无 steps 键)"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
self._create_episode(Path(tmpdir), n=8, action_dim=2)
|
||||
ds = ActionDataset(tmpdir, history_len=2, action_dim=2,
|
||||
image_width=16, image_height=16, frame_len=2,
|
||||
rollout_len=1)
|
||||
sample = ds[0]
|
||||
self.assertNotIn("steps", sample)
|
||||
self.assertIn("target", sample)
|
||||
|
||||
|
||||
class TestR1TrainHelpers(unittest.TestCase):
|
||||
"""R1 训练辅助:_shift_history 与 _collate"""
|
||||
|
||||
def test_shift_history_appends_and_drops(self):
|
||||
from train import _shift_history
|
||||
history = torch.tensor([[[1.0], [2.0], [3.0]]]) # (1,3,1)
|
||||
new_action = torch.tensor([[4.0]]) # (1,1)
|
||||
out = _shift_history(history, new_action)
|
||||
self.assertEqual(out.shape, (1, 3, 1))
|
||||
np.testing.assert_allclose(out[0, :, 0].numpy(), [2.0, 3.0, 4.0])
|
||||
|
||||
def test_collate_single_step_passthrough(self):
|
||||
from train import _collate
|
||||
b = [{"images": torch.zeros(2, 3, 4, 4),
|
||||
"action_history": torch.zeros(2, 2),
|
||||
"target": torch.zeros(2)} for _ in range(3)]
|
||||
out = _collate(b)
|
||||
self.assertNotIn("steps", out)
|
||||
self.assertEqual(out["images"].shape, (3, 2, 3, 4, 4))
|
||||
|
||||
def test_collate_multi_step_stacks_per_k(self):
|
||||
from train import _collate
|
||||
def mk(k):
|
||||
return {"images": torch.full((2, 3, 4, 4), float(k)),
|
||||
"action_history": torch.zeros(2, 2),
|
||||
"target": torch.zeros(2),
|
||||
"prev_action": torch.zeros(2)}
|
||||
# 2 个样本,每个 K=3 步
|
||||
b = [{"steps": [mk(k) for k in range(3)]} for _ in range(2)]
|
||||
out = _collate(b)
|
||||
self.assertEqual(len(out["steps"]), 3)
|
||||
# 第 k 步 batch 的 images 值应全为 k
|
||||
for k in range(3):
|
||||
self.assertEqual(out["steps"][k]["images"].shape, (2, 2, 3, 4, 4))
|
||||
self.assertTrue(
|
||||
torch.allclose(out["steps"][k]["images"],
|
||||
torch.full((2, 2, 3, 4, 4), float(k))))
|
||||
|
||||
|
||||
class TestR4Evaluate(unittest.TestCase):
|
||||
"""R4 离线评估:滚动口径、平滑度、误差增长、越界、汇总"""
|
||||
|
||||
@staticmethod
|
||||
def _create_episode(base_dir, name, actions, image_size=16):
|
||||
ep_dir = base_dir / name
|
||||
images_dir = ep_dir / "images"
|
||||
images_dir.mkdir(parents=True, exist_ok=True)
|
||||
for i in range(len(actions)):
|
||||
Image.fromarray(np.random.randint(0, 255, (image_size, image_size, 3),
|
||||
dtype=np.uint8)).save(
|
||||
images_dir / f"{i:04d}.jpg")
|
||||
np.save(ep_dir / "actions.npy", np.array(actions, dtype=np.float32))
|
||||
|
||||
class _ConstPredictor:
|
||||
"""恒定输出 c 的假推理器(鸭子类型),用于可控验证滚动逻辑"""
|
||||
frame_len = 2
|
||||
history_len = 2
|
||||
action_dim = 2
|
||||
|
||||
def __init__(self, c):
|
||||
self._c = np.asarray(c, dtype=np.float32)
|
||||
self.calls = 0
|
||||
|
||||
def predict(self, images, history):
|
||||
self.calls += 1
|
||||
return self._c.copy()
|
||||
|
||||
def test_rollout_uses_self_prediction_feedback(self):
|
||||
"""滚动应把自预测回写进历史(复刻 run_loop),步数 = n - t0"""
|
||||
import evaluate
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
actions = np.arange(20, dtype=np.float32).reshape(10, 2)
|
||||
self._create_episode(Path(tmpdir), "episode_0001", actions)
|
||||
image_files, acts = evaluate._load_episode(tmpdir, "episode_0001")
|
||||
pred = self._ConstPredictor([100.0, 100.0])
|
||||
steps = evaluate.rollout_episode(pred, image_files, acts)
|
||||
# t0 = max(H=2, F-1=1) = 2,horizon = 10-2 = 8
|
||||
self.assertEqual(len(steps), 8)
|
||||
self.assertEqual(pred.calls, 8)
|
||||
# 恒预测 100,真值 a[t],逐步 mse = mean((100-a[t])^2)
|
||||
self.assertAlmostEqual(steps[0]["mse"],
|
||||
float(np.mean((100 - acts[2]) ** 2)))
|
||||
|
||||
def test_rollout_respects_horizon(self):
|
||||
import evaluate
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
actions = np.zeros((10, 2), dtype=np.float32)
|
||||
self._create_episode(Path(tmpdir), "episode_0001", actions)
|
||||
image_files, acts = evaluate._load_episode(tmpdir, "episode_0001")
|
||||
pred = self._ConstPredictor([0.0, 0.0])
|
||||
steps = evaluate.rollout_episode(pred, image_files, acts, horizon=3)
|
||||
self.assertEqual(len(steps), 3)
|
||||
|
||||
def test_rollout_short_episode_returns_none(self):
|
||||
import evaluate
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
actions = np.zeros((2, 2), dtype=np.float32) # n=2 <= t0=2
|
||||
self._create_episode(Path(tmpdir), "episode_0001", actions)
|
||||
image_files, acts = evaluate._load_episode(tmpdir, "episode_0001")
|
||||
pred = self._ConstPredictor([0.0, 0.0])
|
||||
self.assertIsNone(evaluate.rollout_episode(pred, image_files, acts))
|
||||
|
||||
def test_smoothness_constant_vs_noisy(self):
|
||||
import evaluate
|
||||
const = np.tile([1.0, 2.0], (10, 1))
|
||||
rng = np.random.RandomState(0)
|
||||
noisy = const + rng.randn(10, 2).astype(np.float32)
|
||||
self.assertEqual(evaluate._smoothness(const), 0.0)
|
||||
self.assertGreater(evaluate._smoothness(noisy), 0.0)
|
||||
self.assertEqual(evaluate._smoothness(np.zeros((1, 2))), 0.0)
|
||||
|
||||
def test_oob_detection(self):
|
||||
"""越界:预测超出经验范围应标记 oob=True"""
|
||||
import evaluate
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
actions = np.zeros((10, 2), dtype=np.float32) # 真值全 0
|
||||
self._create_episode(Path(tmpdir), "episode_0001", actions)
|
||||
image_files, acts = evaluate._load_episode(tmpdir, "episode_0001")
|
||||
pred = self._ConstPredictor([5.0, 5.0]) # 远超 [0,0] 范围
|
||||
lo = np.zeros(2, dtype=np.float32)
|
||||
hi = np.zeros(2, dtype=np.float32)
|
||||
steps = evaluate.rollout_episode(pred, image_files, acts,
|
||||
lo=lo, hi=hi)
|
||||
self.assertTrue(all(s["oob"] for s in steps))
|
||||
|
||||
def test_summarize_keys_and_growth(self):
|
||||
import evaluate
|
||||
# 构造误差随步数线性增长的滚动记录
|
||||
steps = [{"t": k, "mse": float(k + 1),
|
||||
"pred": np.zeros(2), "true": np.zeros(2), "oob": False}
|
||||
for k in range(40)]
|
||||
lo = np.zeros(2); hi = np.zeros(2)
|
||||
rep = evaluate.summarize([steps], lo, hi)
|
||||
for key in ("n_episodes", "n_steps", "overall_mse", "error_growth",
|
||||
"error_curve", "pred_smoothness", "true_smoothness",
|
||||
"smoothness_ratio", "oob_rate"):
|
||||
self.assertIn(key, rep)
|
||||
self.assertEqual(rep["n_episodes"], 1)
|
||||
self.assertEqual(rep["n_steps"], 40)
|
||||
# 误差递增 → error_growth > 1
|
||||
self.assertGreater(rep["error_growth"], 1.0)
|
||||
self.assertEqual(len(rep["error_curve"]), 40)
|
||||
|
||||
def test_action_bounds_margin(self):
|
||||
import evaluate
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
self._create_episode(Path(tmpdir), "episode_0001",
|
||||
np.zeros((5, 2), dtype=np.float32))
|
||||
self._create_episode(Path(tmpdir), "episode_0002",
|
||||
np.full((5, 2), 10.0, dtype=np.float32))
|
||||
lo, hi = evaluate._action_bounds(tmpdir, 2, margin_ratio=0.1)
|
||||
# min=0,max=10,span=10,margin=1 → lo=-1, hi=11
|
||||
np.testing.assert_allclose(lo, [-1.0, -1.0], atol=1e-6)
|
||||
np.testing.assert_allclose(hi, [11.0, 11.0], atol=1e-6)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=2)
|
||||
+95
-22
@@ -21,39 +21,104 @@ def split_episodes(data_dir: str, val_ratio: float, seed: int = 42):
|
||||
return indices[val_count:], indices[:val_count]
|
||||
|
||||
|
||||
def train_one_epoch(model, loader, optimizer, device, max_grad_norm: float = None):
|
||||
def _collate(batch):
|
||||
"""多步样本({"steps":[...]})逐 k 聚合成 batch 列表;单步样本走默认 collate。
|
||||
|
||||
rollout_len==1 时 __getitem__ 返回标准 dict,直接默认 collate;
|
||||
rollout_len>1 时返回 {"steps":[step_0..step_{K-1}]},把每个 step_k 跨样本堆叠,
|
||||
返回 {"steps":[batch_0..batch_{K-1}]},供 train_one_epoch 按步 scheduled sampling。
|
||||
"""
|
||||
from torch.utils.data._utils.collate import default_collate
|
||||
if batch and isinstance(batch[0], dict) and "steps" in batch[0]:
|
||||
k = len(batch[0]["steps"])
|
||||
return {"steps": [default_collate([b["steps"][i] for b in batch])
|
||||
for i in range(k)]}
|
||||
return default_collate(batch)
|
||||
|
||||
|
||||
def _shift_history(history: torch.Tensor, new_action: torch.Tensor) -> torch.Tensor:
|
||||
"""动作历史左移一步、追加 new_action 到末尾(scheduled sampling 回传自预测用)。
|
||||
|
||||
history: (B, H, D),new_action: (B, D) → (B, H, D)
|
||||
"""
|
||||
return torch.cat([history[:, 1:], new_action.unsqueeze(1)], dim=1)
|
||||
|
||||
|
||||
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):
|
||||
model.train()
|
||||
total_loss = 0.0
|
||||
count = 0
|
||||
rng = np.random.RandomState()
|
||||
for batch in tqdm(loader, desc=" train", leave=False):
|
||||
images = batch["images"].to(device)
|
||||
history = batch["action_history"].to(device)
|
||||
target = batch["target"].to(device)
|
||||
optimizer.zero_grad()
|
||||
pred = model(images, history)
|
||||
loss = torch.nn.functional.mse_loss(pred, target)
|
||||
if rollout_len == 1:
|
||||
images = batch["images"].to(device)
|
||||
history = batch["action_history"].to(device)
|
||||
target = batch["target"].to(device)
|
||||
pred = model(images, history)
|
||||
loss = torch.nn.functional.mse_loss(pred, target)
|
||||
bs = images.size(0)
|
||||
else:
|
||||
# 多步 scheduled sampling:首步必用真历史;后续步以 ss_prob 用真历史、
|
||||
# 否则用上一步模型预测(detach 截断,避免梯度跨步爆炸)回传构造历史。
|
||||
loss = 0.0
|
||||
bs = 0
|
||||
history = None
|
||||
for k, step in enumerate(batch["steps"]):
|
||||
images = step["images"].to(device)
|
||||
true_history = step["action_history"].to(device)
|
||||
target = step["target"].to(device)
|
||||
prev_action = step["prev_action"].to(device)
|
||||
if k == 0 or rng.random_sample() < ss_prob:
|
||||
history = true_history
|
||||
# 否则沿用上一轮回传的 history(自预测构造)
|
||||
pred = model(images, history)
|
||||
loss = loss + torch.nn.functional.mse_loss(pred, target)
|
||||
bs = images.size(0)
|
||||
# 回传:残差模式先把 Δ̂ 加回 a_{t-1} 得到绝对动作,再移入历史
|
||||
with torch.no_grad():
|
||||
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()
|
||||
if max_grad_norm:
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
|
||||
optimizer.step()
|
||||
total_loss += loss.item() * images.size(0)
|
||||
count += images.size(0)
|
||||
total_loss += loss.item() * bs
|
||||
count += bs
|
||||
return total_loss / max(count, 1)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def validate(model, loader, device):
|
||||
def validate(model, loader, device, rollout_len: int = 1, predict_delta: bool = False):
|
||||
model.eval()
|
||||
total_loss = 0.0
|
||||
count = 0
|
||||
for batch in tqdm(loader, desc=" val ", leave=False):
|
||||
images = batch["images"].to(device)
|
||||
history = batch["action_history"].to(device)
|
||||
target = batch["target"].to(device)
|
||||
pred = model(images, history)
|
||||
loss = torch.nn.functional.mse_loss(pred, target)
|
||||
total_loss += loss.item() * images.size(0)
|
||||
count += images.size(0)
|
||||
if rollout_len == 1:
|
||||
images = batch["images"].to(device)
|
||||
history = batch["action_history"].to(device)
|
||||
target = batch["target"].to(device)
|
||||
pred = model(images, history)
|
||||
loss = torch.nn.functional.mse_loss(pred, target)
|
||||
bs = images.size(0)
|
||||
else:
|
||||
# 验证固定全真历史(确定性,不带 scheduled sampling 随机性),
|
||||
# 对 K 步累加,保证 best 模型选择口径稳定可比。
|
||||
loss = 0.0
|
||||
bs = 0
|
||||
for step in batch["steps"]:
|
||||
images = step["images"].to(device)
|
||||
history = step["action_history"].to(device)
|
||||
target = step["target"].to(device)
|
||||
pred = model(images, history)
|
||||
loss = loss + torch.nn.functional.mse_loss(pred, target)
|
||||
bs = images.size(0)
|
||||
loss = loss / len(batch["steps"])
|
||||
total_loss += loss.item() * bs
|
||||
count += bs
|
||||
if count == 0:
|
||||
return float("inf") # 空验证集不参与比较,避免假 best
|
||||
return total_loss / count
|
||||
@@ -85,7 +150,8 @@ def main():
|
||||
cfg.data.data_dir, cfg.data.history_len, cfg.data.action_dim,
|
||||
cfg.data.image_width, cfg.data.image_height, cfg.data.normalize_actions,
|
||||
cfg.data.action_stats_path, train_indices, cfg.data.frame_len,
|
||||
augment=cfg.train.augment)
|
||||
augment=cfg.train.augment,
|
||||
predict_delta=cfg.train.predict_delta, rollout_len=cfg.train.rollout_len)
|
||||
if cfg.data.normalize_actions and not cfg.data.action_stats_path:
|
||||
# 归一化统计量必须只用训练集计算,再供验证集复用,否则 val loss 不可信
|
||||
stats_path = os.path.join(cfg.train.checkpoint_dir, "action_stats.npy")
|
||||
@@ -96,7 +162,8 @@ def main():
|
||||
val_ds = ActionDataset(
|
||||
cfg.data.data_dir, cfg.data.history_len, cfg.data.action_dim,
|
||||
cfg.data.image_width, cfg.data.image_height, cfg.data.normalize_actions,
|
||||
cfg.data.action_stats_path, val_indices, cfg.data.frame_len)
|
||||
cfg.data.action_stats_path, val_indices, cfg.data.frame_len,
|
||||
predict_delta=cfg.train.predict_delta, rollout_len=cfg.train.rollout_len)
|
||||
|
||||
print(f"Samples: {len(train_ds)} train, {len(val_ds)} val")
|
||||
if len(train_ds) == 0:
|
||||
@@ -111,9 +178,10 @@ def main():
|
||||
loader_gen.manual_seed(cfg.train.seed)
|
||||
train_loader = DataLoader(train_ds, batch_size=cfg.train.batch_size,
|
||||
shuffle=True, num_workers=cfg.train.num_workers,
|
||||
generator=loader_gen)
|
||||
generator=loader_gen, collate_fn=_collate)
|
||||
val_loader = DataLoader(val_ds, batch_size=cfg.train.batch_size,
|
||||
shuffle=False, num_workers=cfg.train.num_workers)
|
||||
shuffle=False, num_workers=cfg.train.num_workers,
|
||||
collate_fn=_collate)
|
||||
|
||||
model = ActionPredictor(cfg.model, cfg.data).to(device)
|
||||
n_params = sum(p.numel() for p in model.parameters())
|
||||
@@ -154,8 +222,13 @@ 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)
|
||||
val_loss = validate(model, val_loader, device)
|
||||
cfg.train.max_grad_norm,
|
||||
rollout_len=cfg.train.rollout_len,
|
||||
ss_prob=cfg.train.ss_prob,
|
||||
predict_delta=cfg.train.predict_delta)
|
||||
val_loss = validate(model, val_loader, device,
|
||||
rollout_len=cfg.train.rollout_len,
|
||||
predict_delta=cfg.train.predict_delta)
|
||||
scheduler.step()
|
||||
print(f" train_loss={train_loss:.6f} val_loss={val_loss:.6f}")
|
||||
|
||||
|
||||
Reference in new issue
Block a user