Differentiable Tree Search Network

📄 arXiv: 2401.11660v2 📥 PDF

作者: Dixant Mittal, Wee Sun Lee

分类: cs.LG, cs.AI

发布日期: 2024-01-22 (更新: 2024-08-02)

💡 一句话要点

提出可微分树搜索网络以解决有限训练数据下的决策问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 决策问题 深度学习 可微分搜索 世界模型 强化学习 算法归纳偏置 在线搜索 模型优化

📋 核心要点

  1. 现有方法在有限训练数据下，策略函数的性能常常受到学习的世界模型不准确性影响，导致决策效果不佳。
  2. 本文提出D-TSN，通过将最佳优先搜索算法的结构嵌入神经网络，增强了算法的归纳偏置，并实现了可微分的在线搜索。
  3. 在Procgen游戏和网格导航任务的离线强化学习实验中，D-TSN的表现超越了多种无模型和有模型的基线方法。

📝 摘要（中文）

在有限训练数据的决策问题中，使用深度神经网络近似的策略函数往往表现不佳。本文提出可微分树搜索网络（D-TSN），通过嵌入最佳优先在线搜索算法的结构，显著增强了算法的归纳偏置。D-TSN利用学习的世界模型进行完全可微的在线搜索，并与搜索算法共同优化，从而学习到更稳健的世界模型，减轻预测不准确带来的影响。实验表明，D-TSN在Procgen游戏和网格导航任务中优于流行的无模型和有模型基线。

🔬 方法详解

问题定义：本文旨在解决在有限训练数据下，深度学习模型在决策问题中因学习的世界模型不准确而导致的性能下降。现有方法如TreeQN虽然引入了算法归纳偏置，但其效果往往不足以应对复杂决策任务。

核心思路：D-TSN通过将最佳优先在线搜索算法的结构嵌入到神经网络中，增强了模型的归纳偏置。该方法通过学习的世界模型进行完全可微的在线搜索，并与搜索算法共同优化，从而提高了模型的鲁棒性。

技术框架：D-TSN的整体架构包括学习的世界模型和搜索算法两个主要模块。首先，世界模型通过有限的数据进行训练；然后，搜索算法在此模型上进行可微分的在线搜索，最终通过联合优化提升模型性能。

关键创新：D-TSN的核心创新在于采用了随机树扩展策略，避免了简单最佳优先搜索带来的损失函数不连续问题。这一设计使得搜索树的扩展被视为另一个决策任务，从而提升了模型的稳定性和有效性。

关键设计：在D-TSN中，损失函数设计为联合优化世界模型和搜索算法的目标，网络结构采用了适应性树扩展策略，并引入了有效的方差减少技术以优化梯度计算。具体参数设置和网络结构细节在实验部分进行了详细描述。

🖼️ 关键图片

fig_0
img_1

📊 实验亮点

在离线强化学习实验中，D-TSN在Procgen游戏和网格导航任务上表现优异，超越了多种流行的无模型和有模型基线，展示了显著的性能提升，具体提升幅度未明确给出。

🎯 应用场景

该研究的潜在应用领域包括游戏智能体、机器人决策系统以及其他需要在有限数据下进行高效决策的场景。D-TSN的创新方法可以为复杂决策任务提供更稳健的解决方案，推动智能体在动态环境中的表现提升。

📄 摘要（原文）

In decision-making problems with limited training data, policy functions approximated using deep neural networks often exhibit suboptimal performance. An alternative approach involves learning a world model from the limited data and determining actions through online search. However, the performance is adversely affected by compounding errors arising from inaccuracies in the learned world model. While methods like TreeQN have attempted to address these inaccuracies by incorporating algorithmic inductive biases into the neural network architectures, the biases they introduce are often weak and insufficient for complex decision-making tasks. In this work, we introduce Differentiable Tree Search Network (D-TSN), a novel neural network architecture that significantly strengthens the inductive bias by embedding the algorithmic structure of a best-first online search algorithm. D-TSN employs a learned world model to conduct a fully differentiable online search. The world model is jointly optimized with the search algorithm, enabling the learning of a robust world model and mitigating the effect of prediction inaccuracies. Further, we note that a naive incorporation of best-first search could lead to a discontinuous loss function in the parameter space. We address this issue by adopting a stochastic tree expansion policy, formulating search tree expansion as another decision-making task, and introducing an effective variance reduction technique for the gradient computation. We evaluate D-TSN in an offline-RL setting with a limited training data scenario on Procgen games and grid navigation task, and demonstrate that D-TSN outperforms popular model-free and model-based baselines.