使用DeepSeek的GRPO算法训练迷宫推理大模型AlphaMaze

AlphaMaze 是一个训练大语言模型解决迷宫问题的开源项目，这个项目灵感来自于 DeepSeek-R1 拥有空间推理能力，但是 R1 不具备视觉能力，并且蒸馏出来的小模型却并没有空间推理能力。

由此 AlphaMaze 把视觉的迷宫问题转换成文字描述，然后通过 SFT 加 GPRO训练得到 AlphaMaze-v0.2-1.5B，使这个小模型拥有了空间推理能力，训练后的迷宫成功率93%。

0 背景

之前微软的多模态思维可视化（MVoT）论文，表明模型可以通过生成其推理过程的视觉表示来解决基于图像的迷宫。但是这些模型在创建准确视觉表示的能力上仍有局限性，并且通常使用多项选择题（MCQ）进行评估，这可能无法完全验证它们的推理能力。

AlphaMaze 采用不同的方法，这是基于 DeepSeek-R1-Distill-Qwen-1.5B 蒸馏模型构建的，核心假设在于：如果一个模型能够成功地从文本描述中迷宫走出来，并输出 Token 来规划下一步行动，那么它就展示了有效的视觉推理能力。

经过验证，当前的非推理的大语言模型都无法走出文本描述的迷宫，包括 GPT-4o 和 Qwen-Max，在迷宫导航的任务中表现不好，这说明他们的空间推理能力确实存在差距。

1 迷宫数据集

使用了 Understanding Search 的多样化的迷宫数据集，都是利用深度优先搜索（DFS）算法生成 5x5 迷宫，接着要把这些数据转换成语言模型能理解的格式，使用一种基于标记的表示系统，将迷宫的视觉结构翻译成模型能理解的语言。

1. 迷宫结构：从迷宫的邻接表开始，它定义了单元格之间的连接方式。

2. 起始和结束：在 5x5 网格中指定起始位置和目标位置。

3. 解决方案路径：提供一系列坐标，代表通过迷宫的正确路径。

关键创新在于 Token 系统，分别使用特殊的 Token 来表示不同的组件：

坐标：用于识别网格中的每个单元格。

墙壁：使用、、等标记来表示每个方向上是否存在墙壁。这允许模型“看到”迷宫的结构。

动作：，，，表示模型可以采取的可能动作。

特殊标记：和表示起始点和结束点。

最终通过这个系统的编译生成了一个包含 10 万个迷宫示例的数据集，并且已经开源：huggingface.co/datasets/jan-hq/Maze-Reasoning

2 训练

训练过程包括两个关键阶段：使用创建好的迷宫数据集，进行监督微调（SFT）和强化学习（RL）的组合训练。

训练使用的数据包含了逐步生成的导航序列，每一步都囊括了整个迷宫的状态和单一方向的指示器，这为训练模型进行运动和空间推理打下了很好的基础。

3 阶段一：监督微调

开始之前明确目标：首先为了确定 SFT 能否让模型拥有基本的视觉思维能力；在此基础上再建立于后续强化学习的比较基准。使用到两种方法：

1. 基线模型：训练后可直接预测完整的解决方案路径，无需中间推理步骤。

2. 视觉思考模型：训练以逐步预测解决方案，在推理过程中结合运动标记。这鼓励模型“思考”其移动。

使用 llama-factory 进行 SFT 训练，10 万个样本，过滤掉其中超过 4096 Token 的样本。

结果，添加新的迷宫标记会导致性能下降，但在没有额外标记的情况下，模型在相同的训练步骤中表现出了令人惊讶的良好性能。

4 阶段二：GRPO 强化学习

强化学习中的一个关键方面是奖励函数，它引导模型朝着期望的行为发展。设计了一个包含多个组件的复合奖励函数，每个组件都关注良好迷宫解决能力的不同方面：

1. 正确性：

目的：检查模型的最终答案是否与正确解决方案匹配。

机制：提取模型预测的路径（使用作为分隔符）并将其与真实答案进行比较。

奖励：正确解决方案得 +2.0，否则得 0.0。这是权重最高的组件，强调准确性。

2. 有效性：

目的：确保模型仅使用有效的移动标记。

机制：验证推理步骤是否仅包含

奖励：有效序列得 +0.5，否则得 0.0。这强制使用正确的导航“语言”。

3. 严格格式化：

目的：强制模型输出严格的 XML 类似格式。

机制：使用正则表达式（^<think>\n.*?\n</think>\n\n.*?\n$）检查精确的换行符放置以及和标签的存在。

奖励：完美格式化得 +0.5，否则得 0.0。这鼓励结构化输出。

4. 宽松格式化：

目的：检查基本的 XML 结构，但对格式的要求更宽松。

机制：使用更宽松的正则表达式（<think>.*?</think>\s*.*?），允许在空白符和换行符上有变化。

奖励：正确使用和标签得 +0.5，否则得 0.0

5. 标签格式化：

目的：鼓励正确使用 XML 标签。

机制：检查是否恰好有一个打开的<think>\n（+0.125）和一个关闭的\n</think>\n（+0.125）。

最大奖励：+0.25。这为 XML 结构的正确部分提供部分奖励。

最终奖励是所有这些组件的总和，最大可能奖励：2.0 + 0.5 + 0.5 + 0.5 + 0.25 =3.75

在 GRPO 训练中，精选 10,000 个迷宫示例，训练时长为 1,000 步，上下文长度设为 4096 个标记；硬件方面使用单个 NVIDIA A6000 GPU 并采用 Low-Rank Adaptation（LoRA）进行高效训练；基础模型选用 DeepSeek R1 Distill Qwen 1.5B，SFT 模型为 AlphaMaze - SFT 且使用迷宫数据集进行监督式微调预训练，同时使用 Unsloth 框架并进行 4 位量化。

在训练的三个关键时间点（100 步、300 步和 600 步）监控训练性能。以下表格总结了每个模型在这些时间点的观察行为和结果。

5 关键收获

没有冷启动的 GRPO 具有挑战性：从基础模型（没有 SFT）开始 GRPO 训练被证明是困难的。模型难以学习任务，并表现出“过度思考”的行为，超过了上下文长度。

SFT 提供了强大的基础：使用 SFT 模型初始化显著加速了学习并提高了 GRPO 的效果。模型迅速收敛并产生了可用的结果。

LoRA 是有效的：使用 LoRA 在这项任务上取得了良好的结果，展示了其在资源受限情况下的潜力。

奖励函数设计至关重要：精心设计的奖励函数在塑造模型行为方面发挥了关键作用，引导其朝着正确的解决方案和适当的格式发展。

对大模型相关深度内容以及开源项目感兴趣，欢迎关注 极客开源。相关资料放在下面：