DQN 学不稳？未冻结 target、忘记 detach、终止状态仍 bootstrap 的三连坑

# 触发三连坑
python dqn_pitfalls.py --bug on  --env CartPole-v1

# 修复版本
python dqn_pitfalls.py --bug off --env CartPole-v1
# dqn_pitfalls.py
import argparse, random, collections, numpy as np, torch, torch.nn as nn, torch.nn.functional as F
import gymnasium as gym
torch.manual_seed(0); np.random.seed(0); random.seed(0)

Transition = collections.namedtuple("T", "s a r s2 terminated truncated")

class QNet(nn.Module):
    def __init__(self, s_dim, a_dim, h=128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(s_dim, h), nn.ReLU(),
            nn.Linear(h, h), nn.ReLU(),
            nn.Linear(h, a_dim)
        )
    def forward(self, x): return self.net(x)

class Replay:
    def __init__(self, cap=50000): self.cap, self.buf, self.pos = cap, [None]*cap, 0
    def push(self, *args):
        self.buf[self.pos] = Transition(*args); self.pos = (self.pos+1) % self.cap
    def sample(self, bs):
        batch = random.sample([b for b in self.buf if b is not None], bs)
        t = Transition(*zip(*batch))
        s  = torch.tensor(np.stack(t.s),  dtype=torch.float32)
        a  = torch.tensor(t.a,           dtype=torch.int64).unsqueeze(-1)
        r  = torch.tensor(t.r,           dtype=torch.float32).unsqueeze(-1)
        s2 = torch.tensor(np.stack(t.s2), dtype=torch.float32)
        term = torch.tensor(t.terminated, dtype=torch.float32).unsqueeze(-1)
        trunc = torch.tensor(t.truncated, dtype=torch.float32).unsqueeze(-1)
        return s, a, r, s2, term, trunc
    def __len__(self): return sum(b is not None for b in self.buf)

def epsilon_greedy(q, eps, a_dim):
    if random.random() < eps:
        return random.randrange(a_dim), None
    with torch.no_grad():
        return int(q.argmax().item()), None

def train(args):
    env = gym.make(args.env)
    s_dim, a_dim = env.observation_space.shape[0], env.action_space.n
    online = QNet(s_dim, a_dim)
    target = online if args.bug else QNet(s_dim, a_dim)
    if not args.bug:
        target.load_state_dict(online.state_dict())
        target.eval()

    opt = torch.optim.AdamW(online.parameters(), lr=1e-3)
    buf = Replay(50000)
    gamma = 0.99
    huber = nn.SmoothL1Loss()  # 稳定版损失

    eps_start, eps_end, eps_steps = 1.0, 0.05, 20000
    def eps_at(t): 
        return max(eps_end, eps_start - (eps_start-eps_end)*t/eps_steps)

    s, _ = env.reset(seed=np.random.randint(10000))
    ret, epi = 0.0, 0
    metrics = {"ret":[], "tdm":[], "tds":[], "qabs":[], "eps":[]}

    total_steps = 40000
    warmup, batch = 1000, 64
    target_sync = 1000

    for t in range(1, total_steps+1):
        eps = eps_at(t)
        a, _ = epsilon_greedy(online(torch.tensor(s, dtype=torch.float32)), eps, a_dim)
        s2, r, terminated, truncated, _ = env.step(a)
        buf.push(s, a, np.clip(r, -1, 1), s2, terminated, truncated)  # 奖励裁剪便于稳态
        s, ret = (env.reset()[0] if (terminated or truncated) else s2), (0.0 if (terminated or truncated) else ret+r)
        if terminated or truncated: 
            metrics["ret"].append(ret); epi += 1

        if len(buf) >= warmup:
            s_b, a_b, r_b, s2_b, term_b, trunc_b = buf.sample(batch)
            q = online(s_b).gather(1, a_b)  # [B,1]

            if args.bug:
                # 错误：1) 用在线网络做 target；2) 未 detach；3) 终止也 bootstrap
                q_next = online(s2_b).max(1, keepdim=True).values
                y = r_b + gamma * q_next  # ❌ 三连坑
                loss = F.mse_loss(q, y)
            else:
                with torch.no_grad():                       # 修复：冻结 target 分支
                    # time-limit 截断不算真正终止，只有 terminated=1 才不 bootstrap
                    nonterminal = (1.0 - term_b)            # [B,1]
                    q_next = target(s2_b).max(1, keepdim=True).values
                    y = r_b + gamma * nonterminal * q_next  # ✅ 正确 bootstrap
                loss = huber(q, y)                          # ✅ Huber 更稳

            opt.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(online.parameters(), 10.0)  # ✅ 梯度裁剪避免爆炸
            opt.step()

            if (not args.bug) and (t % target_sync == 0):
                target.load_state_dict(online.state_dict())

            with torch.no_grad():
                td = (q - y).detach()
                metrics["tdm"].append(float(td.mean()))
                metrics["tds"].append(float(td.std()))
                metrics["qabs"].append(float(online(s_b).abs().mean()))
                metrics["eps"].append(eps)

        if t % 1000 == 0 and len(metrics["ret"])>0:
            print(f"[{'BUG' if args.bug else 'FIX'}] step={t:05d} "
                  f"ep={epi:04d} avg_ret={np.mean(metrics['ret'][-10:]):6.2f} "
                  f"td(mean/std)={np.mean(metrics['tdm'][-100:]):+.3f}/{np.mean(metrics['tds'][-100:]):.3f} "
                  f"|Q|≈{np.mean(metrics['qabs'][-100:]):.2f} eps={eps:.3f}")

if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--env", default="CartPole-v1")
    ap.add_argument("--bug", choices=["on","off"], default="on")
    args = ap.parse_args()
    train(args)

target = QNet(s_dim, a_dim).eval()
target.load_state_dict(online.state_dict())
if step % target_sync == 0:
    target.load_state_dict(online.state_dict())

with torch.no_grad():
    nonterminal = (1.0 - terminated.float().unsqueeze(-1))
    q_next = target(s2).max(1, keepdim=True).values
    y = r + gamma * nonterminal * q_next
loss = F.smooth_l1_loss(q_selected, y)

torch.nn.utils.clip_grad_norm_(online.parameters(), max_norm=10.0)
r = np.clip(r, -1, 1)  # reward clip（经典 DQN 习惯）

def assert_target_frozen(y):
    assert (not y.requires_grad), "TD 目标 y 可导，疑似未 detach 或 target 未冻结"

def log_health(step, td, q_abs):
    print(f"[dbg] step={step} td_mean={float(td.mean()):+.3f} td_std={float(td.std()):.3f} |Q|={float(q_abs):.2f}")