在浏览器中实现深度强化学习：基于 TensorFlow.js 的 DQN 完整实战

从零构建可交互的 DQN 智能体，包含可视化训练、Double DQN 优化与完整可运行代码

一、为什么选择 TensorFlow.js 做强化学习

传统强化学习工程通常基于 Python + PyTorch/TensorFlow 后端栈，依赖服务器算力和复杂的环境配置。而 TensorFlow.js（TF.js）将深度强化学习直接搬进浏览器，带来几个独特优势：

• 零服务端依赖：纯前端运行，打开网页即可训练，无需部署后端、配置环境
• 数据隐私：训练全程在本地完成，敏感数据不上传服务器
• 天然交互：直接对接 Canvas、DOM、网页游戏，实时控制与可视化
• 前端友好：前端开发者无需切换技术栈，用 JavaScript 即可完整实现 RL 算法

二、DQN 核心原理快速回顾

DQN（Deep Q-Network）是深度强化学习的入门基石，由 DeepMind 于 2013 年提出，核心思想是用神经网络拟合 Q 值函数。

2.1 基本概念

2.2 两个关键创新

经验回放（Experience Replay）：存储历史交互样本 (s, a, r, s')，训练时随机采样小批量，打破样本时序相关性，提升训练稳定性。

目标网络（Target Network）：维护一个参数滞后的副本网络计算目标 Q 值，避免"追逐移动目标"导致的训练震荡。本文实现的 Double DQN 进一步用当前网络选动作、目标网络估值，解决 Q 值过估计问题。

三、环境实现：GridWorld

先实现纯 JavaScript 的网格环境，不依赖 TF.js，负责状态流转和奖励计算。环境与算法解耦是良好的工程实践。

class GridWorld {
  constructor(size = 5) {
    this.size = size;
    this.agentPos = { x: 0, y: 0 };      // 智能体起点：左上角
    this.goalPos = { x: size - 1, y: size - 1 }; // 目标点：右下角
  }

  // 重置环境，返回初始状态
  reset() {
    this.agentPos = { x: 0, y: 0 };
    return [this.agentPos.x, this.agentPos.y];
  }

  // 执行动作，返回 [新状态, 奖励, 是否结束]
  step(action) {
    const moves = [
      { x: 0, y: -1 }, // 0: 上
      { x: 0, y: 1 },  // 1: 下
      { x: -1, y: 0 }, // 2: 左
      { x: 1, y: 0 }   // 3: 右
    ];

    const next = {
      x: this.agentPos.x + moves[action].x,
      y: this.agentPos.y + moves[action].y
    };

    // 撞墙：位置不变，负奖励
    if (next.x < 0 || next.x >= this.size ||
        next.y < 0 || next.y >= this.size) {
      return [[this.agentPos.x, this.agentPos.y], -1, false];
    }

    this.agentPos = next;

    // 到达目标
    if (next.x === this.goalPos.x && next.y === this.goalPos.y) {
      return [[next.x, next.y], 10, true];
    }

    // 普通移动：小负奖励鼓励走最短路径
    return [[next.x, next.y], -0.1, false];
  }
}

四、用 TF.js 构建 Double DQN 模型

我们用两层全连接网络拟合 Q 函数，输入 2 维状态坐标，输出 4 个动作的 Q 值。同时维护当前网络和目标网络两个副本。

function buildDQN(stateSize = 2, actionSize = 4) {
  const model = tf.sequential();

  // 隐藏层 1
  model.add(tf.layers.dense({
    units: 64,
    activation: 'relu',
    inputShape: [stateSize]
  }));

  // 隐藏层 2
  model.add(tf.layers.dense({
    units: 32,
    activation: 'relu'
  }));

  // 输出层：每个动作对应一个 Q 值，线性激活
  model.add(tf.layers.dense({
    units: actionSize,
    activation: 'linear'
  }));

  model.compile({
    optimizer: tf.train.adam(0.001),
    loss: 'meanSquaredError'
  });

  return model;
}

// 复制网络权重（用于目标网络更新）
async function copyWeights(source, target) {
  const srcWeights = source.getWeights();
  const cloned = srcWeights.map(w => w.clone());
  target.setWeights(cloned);
  cloned.forEach(w => w.dispose());
}

五、经验回放池与动作选择

class ReplayBuffer {
  constructor(capacity = 5000) {
    this.buffer = [];
    this.capacity = capacity;
  }

  add(state, action, reward, nextState, done) {
    this.buffer.push({ state, action, reward, nextState, done });
    if (this.buffer.length > this.capacity) {
      this.buffer.shift();
    }
  }

  sample(batchSize) {
    const samples = [];
    for (let i = 0; i < batchSize; i++) {
      samples.push(this.buffer[Math.floor(Math.random() * this.buffer.length)]);
    }
    return samples;
  }
}

// ε-greedy 动作选择
function chooseAction(model, state, epsilon, actionSize = 4) {
  if (Math.random() < epsilon) {
    return Math.floor(Math.random() * actionSize); // 随机探索
  }

  return tf.tidy(() => {
    const stateTensor = tf.tensor2d([state]);
    const qValues = model.predict(stateTensor);
    const qArr = qValues.dataSync();
    return qArr.indexOf(Math.max(...qArr)); // 贪心选择
  });
}

六、Double DQN 训练核心

Double DQN 的关键在于：用当前网络选择下一状态的最优动作，用目标网络评估该动作的 Q 值，从而解耦选择与估值。

async function trainStep(onlineNet, targetNet, replayBuffer, batchSize = 64, gamma = 0.95) {
  if (replayBuffer.buffer.length < batchSize) return 0;

  const samples = replayBuffer.sample(batchSize);
  const states = samples.map(s => s.state);
  const nextStates = samples.map(s => s.nextState);

  let stateTensor, targetTensor;

  tf.tidy(() => {
    const stateT = tf.tensor2d(states);
    const nextStateT = tf.tensor2d(nextStates);

    // 当前网络：计算当前 Q 值，并选择下一状态的最优动作
    const currentQ = onlineNet.predict(stateT);
    const nextQOnline = onlineNet.predict(nextStateT);
    const currentQArr = currentQ.arraySync();
    const nextQOnlineArr = nextQOnline.arraySync();

    // 目标网络：评估当前网络选出的动作的 Q 值（Double DQN 核心）
    const nextQTarget = targetNet.predict(nextStateT);
    const nextQTargetArr = nextQTarget.arraySync();

    for (let i = 0; i < batchSize; i++) {
      const { action, reward, done } = samples[i];

      if (done) {
        currentQArr[i][action] = reward;
      } else {
        // Double DQN：online 选动作，target 估值
        const bestAction = nextQOnlineArr[i].indexOf(Math.max(...nextQOnlineArr[i]));
        currentQArr[i][action] = reward + gamma * nextQTargetArr[i][bestAction];
      }
    }

    stateTensor = tf.keep(stateT);
    targetTensor = tf.keep(tf.tensor2d(currentQArr));
  });

  const loss = await onlineNet.trainOnBatch(stateTensor, targetTensor);
  stateTensor.dispose();
  targetTensor.dispose();
  return loss;
}

七、进阶优化方向

7.1 Dueling DQN

将 Q 值拆分为状态价值 V(s) 和动作优势 A(s,a)：Q(s,a) = V(s) + (A(s,a) - mean(A))。这种结构让网络可以独立评估状态的好坏，在动作差异不大的环境中提升学习效率。

7.2 优先经验回放（PER）

按 TD 误差 |δ| 对样本加权采样，让模型更多学习"预测偏差大"的样本，显著提高样本利用率。实现上需要维护一个按误差排序的优先级队列。

7.3 更复杂的环境

本文的 GridWorld 过于简单，你可以扩展到：

• 带障碍物的迷宫（增加状态复杂度）
• CartPole 平衡杆（连续状态空间）
• Flappy Bird 类游戏（像素输入 + CNN 特征提取）

7.4 浏览器端性能优化

八、常见问题与排错

Q1: 训练不收敛，奖励一直很低？

检查：① ε 衰减是否过快（建议 0.995-0.999）；② 学习率是否过高（建议 0.0005-0.001）；③ 奖励设计是否合理（目标奖励应远大于单步惩罚）；④ 经验回放池是否足够大（建议 ≥2000）。

Q2: 浏览器内存持续增长？

这是 TF.js 最常见的问题。确保：① 所有推理代码用 tf.tidy() 包裹；② 异步训练的张量用 tf.keep() + 手动 dispose()；③ 定期调用 tf.nextFrame() 释放 GPU 内存。可以用 tf.memory().numTensors 监控张量数量。

Q3: 不同浏览器训练结果差异大？

WebGL 实现在不同浏览器中存在浮点精度差异。建议：① 统一使用 Chrome/Edge；② 如需更高精度可切换到 WASM 后端（tf.setBackend('wasm')），但速度会慢很多。