ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

PyTorch强化学习实战——从零开始实现魔方求解器

PyTorch强化学习实战——从零开始实现魔方求解器

PyTorch强化学习实战——从零开始实现魔方求解器

    • 0. 前言
    • 1. 通用魔方求解框架
    • 2. 魔方环境
    • 3. 模型训练
    • 4. 搜索过程
    • 5. 实验结果
      • 5.1 2×2 魔方
      • 5.2 3×3 魔方
    • 6. 模型改进方向
    • 相关链接

0. 前言

我们已经探索了强化学习 (Reinforcement Learning, RL) 在离散优化(数学中研究离散结构优化问题的分支)领域的应用,并以著名的魔方谜题作为示范。在本节中,我们将详细介绍并复现McAleer等人在论文《Solving the Rubik’s cube without human knowledge》中提出的方法,同时探讨该方法的改进方向。

1. 通用魔方求解框架

接下来,将简要介绍我们的实现框架和关键设计决策,但首先必须强调几点以设定合理预期:

  • 本节仅复现论文方法,但论文未详述具体超参数,因此经过大量实验后结果仍与论文存在差异
  • 同时尝试以通用方式实现所有模块以便后续实验。例如将魔方状态和变换的具体细节抽象化,允许通过添加新模块实现类似3×3魔方的其他谜题。代码已实现2×2和3×3两种魔方,任何具完全可观测性和固定动作集的环境均可扩展实现
  • 代码优先考虑清晰度与简洁性而非性能。当然在不过度增加复杂度的情况下会优化性能,例如通过分离打乱魔方生成和网络前向传播将训练速度提升五倍。但若需重构为多GPU多线程模式,则倾向于保持简洁。典型案例如蒙特卡洛树搜索 (Monte Carlo Tree Search, MCTS) 过程通常实现为共享树的多线程代码,虽能数倍提速但需复杂进程同步,因此本节采用串行实现,仅对批量搜索进行基础优化

总体来说,代码由以下几个部分组成:

  • 魔方环境:定义观测空间、可用动作及网络状态表示,位于 cubes.py 模块
  • 神经网络部分:包含待训练模型、训练样本生成及训练循环,包括训练工具train.py和模块model.py
  • 魔方求解器:包含搜索过程,含solver.py工具和实现MCTS的mcts.py模块
  • 辅助工具:包括超参数配置文件、生成魔方问题集的工具等整合组件

2. 魔方环境

离散优化问题规模庞大且类型多样。仅魔方类谜题就包含数十种变体,最流行的包括2×2、3×3、4×4魔方、Square-1和金字塔魔方。而论文提出的方法具有通用性,不依赖先验领域知识、动作数量或状态空间规模。该方法对问题的关键假设包括:

  • 环境状态需完全可观测,且观测值能区分不同状态。魔方满足此条件(可观测所有面状态),但扑克游戏多数变体不满足(如无法看到对手牌面)
  • 动作数量需离散且有限。魔方动作数量有限,但若动作空间为"旋转方向盘角度α∈[-120°…120°]",则属于不同问题领域(这属于连续控制问题)
  • 需具备可靠的环境模型,即能回答"对状态sⱼ执行动作aᵢ会产生什么结果?"若无此模型,自主迭代和MCTS均无法适用。该要求严格,多数问题缺乏此类模型或其输出噪声较大。但国际象棋、围棋等游戏具备此类模型(即游戏规则)
  • 同时,如 MuZero 方法所示,可用神经网络近似模型,但需以性能下降为代价
  • 此外该领域需具有确定性,相同动作作用于相同状态总是产生相同结果。反例为双陆棋,玩家每回合掷骰子决定移动步数,但该方法很可能也能推广至此类场景

为将该方法应用于3×3魔方以外的领域,所有具体环境细节均封装至独立模块,通过抽象接口CubeEnv与代码其他部分通信(定义于_env.py)。接下来,我们解析该接口。

(1)如以下代码片段所示,该类的构造函数接收一系列参数:

  • 环境名称
  • 环境状态类型
  • 魔方初始(已还原)状态实例
  • 检查特定状态是否为还原状态的判定函数。对3×3魔方而言,直接与initial_state参数比较即可,但2×2、4×4等魔方可能存在多个最终状态,因此需要独立判定函数覆盖此类情况
  • 可应用于状态的动作枚举
  • 变换函数:接收状态和动作,返回结果状态
  • 逆函数:将每个动作映射为其逆动作
  • 渲染函数:将状态转换为人类可读形式
  • 编码状态张量的形状
  • 将紧凑状态表示编码为神经网络友好格式的函数
classCubeEnv:def__init__(self,name,state_type,initial_state,is_goal_pred,action_enum,transform_func,inverse_action_func,render_func,encoded_shape,encode_func):self.name=name self._state_type=state_type self.initial_state=initial_state self._is_goal_pred=is_goal_pred self.action_enum=action_enum self._transform_func=transform_func self._inverse_action_func=inverse_action_func self._render_func=render_func self.encoded_shape=encoded_shape self._encode_func=encode_func

魔方环境与 Gymnasium API 并不兼容,本节使用这个示例来说明如何突破Gym的框架。

(2)CubeEnv API中的部分方法仅是对构造函数所传函数的封装。这种设计允许新环境在独立模块中实现,在环境注册表中完成注册,并为代码其他部分提供统一接口:

def__repr__(self):return"CubeEnv(%r)"%self.name# wrapper functionsdefis_goal(self,state):assertisinstance(state,self._state_type)returnself._is_goal_pred(state)deftransform(self,state,action):assertisinstance(state,self._state_type)assertisinstance(action,self.action_enum)returnself._transform_func(state,action)definverse_action(self,action):returnself._inverse_action_func(action)defrender(self,state):assertisinstance(state,self._state_type)returnself._render_func(state)defencode_inplace(self,target,state):assertisinstance(state,self._state_type)returnself._encode_func(target,state)

类中的所有其他方法都基于这些基本操作提供扩展的统一功能。

(3)sample_action()方法实现随机采样动作的功能。若传入prev_action参数,则从可能结果中排除反向动作,这可有效避免生成短循环(如R→r或L→l):

defsample_action(self,prev_action=None):whileTrue:res=self.action_enum(random.randrange(len(self.action_enum)))ifprev_actionisNoneorself.inverse_action(res)!=prev_action:returnres

(4)scramble()方法将动作列表(作为参数传入)应用于魔方初始状态,返回最终状态:

defscramble(self,actions):s=self.initial_stateforactioninactions:s=self.transform(s,action)returns

(5)scramble_cube()方法提供随机打乱魔方的功能,返回所有中间状态。当return_inverse参数为False时,函数返回包含打乱过程每一步(depth, state)的元组列表;若为True,则返回包含三个值的元组(depth, state, inv_action):

defis_state(self,state):returnisinstance(state,self._state_type)defscramble_cube(self,scrambles_count,return_inverse=False,include_initial=False):assertisinstance(scrambles_count,int)assertscrambles_count>0state=self.initial_state result=[]ifinclude_initial:assertnotreturn_inverse result.append((1,state))prev_action=Nonefordepthinrange(scrambles_count):action=self.sample_action(prev_action=prev_action)state=self.transform(state,action)prev_action=actionifreturn_inverse:inv_action=self.inverse_action(action)res=(depth+1,state,inv_action)else:res=(depth+1,state)result.append(res)returnresult

(6)explore_states()方法实现了自主迭代所需的功能,对给定魔方状态应用所有可能动作。返回结果为元组,其中第一个列表包含扩展后的状态,第二个列表标记这些状态是否为目标状态:

defexplore_state(self,state):res_states,res_flags=[],[]foractioninself.action_enum:new_state=self.transform(state,action)is_init=self.is_goal(new_state)res_states.append(new_state)res_flags.append(is_init)returnres_states,res_flags

通过这种通用功能,类似环境可以轻松实现并接入现有的训练和测试方法,只需极少的样板代码。作为示例,我们提供了实验使用的2×2魔方和3×3魔方环境,其具体实现位于 cube2x2.py 和 cube3x3.py 中,可作为实现自定义环境的基础。
每个环境需要通过创建CubeEnv类实例,并将实例传递给_env.py中定义的register()函数来完成注册。以下是cube2x2.py模块中的相关代码片段:

_env.register(_env.CubeEnv(name="cube2x2",state_type=State,initial_state=initial_state,is_goal_pred=is_initial,action_enum=Action,transform_func=transform,inverse_action_func=inverse_action,render_func=render,encoded_shape=encoded_shape,encode_func=encode_inplace))

完成注册后,可通过libcube.cubes.get()方法获取魔方环境(以环境名称为参数)。代码其余部分仅使用CubeEnv类的公共接口,这使得代码与具体魔方类型解耦,并简化了扩展性。

3. 模型训练

训练流程实现在 train.py 工具和 model.py 模块中,基本按照论文描述的方法实现,但有一处不同:代码支持两种计算网络价值头目标值的方法。一种是论文原方法,另一种是我们的改进版本。
为简化实验并确保结果可复现,所有训练参数均在单独的.ini文件中指定,包含以下配置选项:

  • 使用的环境名称(当前支持cube2x2和cube3x3)
  • 运行名称,用于TensorBoard命名和模型保存目录
  • 在ADI中使用的目标值计算方法。本节实现了两种:论文原方法和改进版——实验表明改进版收敛更稳定
  • 训练参数:批大小、CUDA使用、学习率、学习率衰减等

在ini文件夹中可找到实验配置示例。训练期间,参数指标将记录到runs文件夹的TensorBoard中,最佳损失值的模型将保存至saves目录。
以下展示配置文件示例 (cube2x2-paper-d200.ini),该配置定义了2×2魔方实验,使用论文中的价值计算方法,打乱深度为200:

[general] cube_type=cube2x2 run_name=paper [train] cuda=True lr=1e-5 batch_size=5000 scramble_depth=200 report_batches=10 checkpoint_batches=100 lr_decay=True lr_decay_gamma=0.95 lr_decay_batches=1000

启动训练时,需要将.ini文件传递给train.py工具。例如,使用上述.ini文件训练模型的命令如下:

$ python ./train.py-iini/cube2x2-paper-d200.ini-nt1

参数-n指定运行名称,该名称将与.ini文件中的名称组合,作为TensorBoard序列的名称。

4. 搜索过程

训练结果是包含网络权重的模型文件。该文件可通过MCTS方法用于求解魔方,相关功能实现在solver.py工具和mcts.py模块中。
solver工具非常灵活,支持多种模式:

  1. 通过-p选项传递一个以逗号分隔的动作索引列表来解决一个单一的打乱魔方。例如-p 1,6,1表示依次执行第二个、第七个、第二个动作打乱的魔方。动作具体含义由-e选项指定的环境决定,可在对应魔方环境模块中查看动作索引映射(如2×2魔方的动作1,6,1对应L→R′→L变换)
  2. 从文本文件中读取排列(每行一个魔方)并求解。通过-i选项指定文件名(每行一个魔方状态)。cubes_tests文件夹提供若干测试问题集,也可使用gen_cubes.py工具生成自定义随机问题集(支持设置随机种子、打乱深度等参数)
  3. 生成指定深度的随机打乱并求解
  4. 运行复杂度递增(打乱深度递增)的测试序列,求解并将结果写入CSV文件。通过-o选项启用此模式,对评估训练模型质量非常有用(但可能耗时较长),还可选择生成测试结果图

所有情况都需通过-e选项指定环境名称,-m选项指定模型权重文件。此外还有其他参数可用于调整MCTS选项和时间/搜索步数限制,具体选项名称可在 solver.py 代码中查看。

5. 实验结果

但原论文未提供该方法关键细节,如训练超参数、训练时的魔方打乱深度以及收敛情况。为填补这些空白,我们尝试了多种超参数组合,但结果仍与论文公布数据存在较大差异。实验表明原方法的训练收敛极不稳定:即使使用较小学习率和较大批次,训练最终仍会发散,价值损失分量呈指数级增长。下图展示了2×2环境中的典型发散情况:

经过多次实验,该问题可能源于方法中提出的价值目标函数存在缺陷。具体而言,在公式y v i = m a x a ( v s ( a ) + R ( A ( s , a ) ) ) y_{v_i} = \underset {a}{max}( v_s(a) + R(A(s,a)))yvi​​=amax​(vs​(a)+R(A(s,a)))中,即使对于目标状态,网络返回值v s ( a ) v_s(a)vs​(a)也会与实际奖励R ( s ) R(s)R(s)相加。这导致网络可能输出任意值(-100、10⁶或3.1415),对于使用均方误差目标的神经网络训练极为不利。为验证此问题,我们修改目标值计算方法,对目标状态直接赋予0值:
y v i = { m a x a ( v s ( a ) + R ( A ( s , a ) ) ) if s is not the goal 0 if s is the goal state y_{v_i} = \begin{cases} \underset {a}{max}( v_s(a) + R(A(s,a))) & \text{if s is not the goal} \\ 0 & \text{if s is the goal state} \end{cases}yvi​​={amax​(vs​(a)+R(A(s,a)))0​if s is not the goalif s is the goal state​
通过在.ini文件中将value_targets_method参数设为zero_goal_value(默认值为paper)即可启用此修改。
经过这一简单调整,训练过程更快收敛到稳定的网络价值头输出。下图展示了改进后的收敛情况:

5.1 2×2 魔方

在论文中,在配备三块Titan Xp GPU的机器上训练了44小时,期间模型处理了80亿个魔方状态,对应每秒5万个状态的训练速度。本节的实现在单块GPU上达到每秒1.5万个状态,速度相当。若在单GPU上复现训练过程需近六天时间,这对实验和超参数调优并不现实。
为了克服这个问题,本节实现了一个更简单的2×2魔方环境,训练只需要一到两个小时。我们提供了两个.ini文件用于复现训练:

  • ini/cube2x2-paper-d200.ini:使用论文原价值目标方法
  • ini/cube2x2-zero-goal-d200.ini:对目标状态设价值目标为0

两种配置均使用1万状态的批次和200的打乱深度,训练参数相同。训练后得到两个模型::

  • 论文方法:损失值0.032572
  • 零目标方法:损失值0.012226

为公平比较,我们生成了深度1-50的20组测试打乱(共1000个测试魔方),存放在cubes_test目录,并用solver.py工具对每种方法的最佳模型进行测试。每个测试打乱的搜索上限设为3万次。该工具生成包含每个测试结果的CSV文件(存于csvs)。

$ python solver.py-msaves/cube2x2-paper-d200-t1/chpt_016700.dat-ecube2x2-icubes_tests/2x2-d1-50.txt--cuda

实验表明:论文所述模型能解决55%的测试魔方,而零目标改进模型达到100%解决率。下图展示两种模型随打乱深度的结果:左侧部分显示解决比例,右侧部分显示平均MCTS搜索步数。可见改进版本需要显著减少(3-5倍)的MCTS搜索步骤,说明学习到的策略更优。

最后检查解的长度。下图显示朴素解与BFS解的长度对比:朴素解比BFS解长10倍左右,这种差异可能源于未优化的MCTS参数。在朴素解中,零目标方法找到的解更短(再次表明其策略更优)。

5.2 3×3 魔方

3×3魔方模型的训练负担重得多,我们仅进行了初步探索。但有限实验表明,训练方法的零目标修改能显著提升训练稳定性和最终模型质量。
本节结果并不如论文中的优异:获得的最佳模型能解决打乱深度12-15的魔方,但对更复杂问题始终失败。增加CPU核心数和并行MCTS可能提升性能。测试时搜索步数限制为10万步,每个打乱深度生成5个随机打乱。改进版本再次展现优势:论文方法训练的模型仅能解决深度9的问题,而改进版本能达到深度13。
下图对比了论文方法与零目标改进版本的解决率(左侧图像)和平均MCTS搜索次数(右侧图像)。

下图展示找到的最优解长度。如前所述,朴素搜索的结果远长于BFS优化解。BFS解长度与打乱深度几乎完全吻合:

理论上,深度超过20后解长度应趋于饱和(因为"上帝之数"为20),但本节版本无法解决打乱深度超过13的魔方,因此难以验证。

6. 模型改进方向

以下是模型有待探索改进的方向:

  • 输入与网络结构优化:魔方具有复杂结构,简单前馈网络可能非最佳选择。卷积结构可能带来显著提升
  • 训练过程中的振荡与不稳定可能源于步间关联性这一RL常见问题。常规解决方案是采用目标网络,即使用旧版本网络获取引导值
  • 优先经验回放缓冲区可能加速训练
  • 实验表明按打乱深度反比加权的样本权重有助于学习解决轻度打乱魔方的策略,但可能减缓对深层状态的学习。可设计自适应权重机制,在训练后期降低加权强度
  • 可引入熵损失函数作为策略正则化手段
  • 2×2魔方模型未考虑无中心块导致的整体旋转对称性。这对小状态空间影响不大,但对4×4魔方将至关重要
  • 需更多实验优化训练和MCTS参数

相关链接

PyTorch强化学习实战(1)——强化学习(Reinforcement Learning,RL)详解
PyTorch强化学习实战(2)——强化学习环境库Gymnasium
PyTorch强化学习实战(3)——Gymnasium API扩展功能
PyTorch强化学习实战(4)——PyTorch基础
PyTorch强化学习实战(5)——PyTorch Ignite 事件驱动机制与实践
PyTorch强化学习实战(6)——交叉熵方法详解与实现
PyTorch强化学习实战(7)——表格学习与贝尔曼方程
PyTorch强化学习实战(8)——Q学习详解与实现
PyTorch强化学习实战(9)——深度Q学习
PyTorch强化学习实战(10)——强化学习高级组件
PyTorch强化学习实战(11)——N步DQN(N-step DQN)
PyTorch强化学习实战(12)——Double DQN(DDQN)
PyTorch强化学习实战(13)——噪声网络(NoisyNet-DQN)
PyTorch强化学习实战(14)——优先经验回放机制
PyTorch强化学习实战(15)——Dueling DQN
PyTorch强化学习实战(16)——Categorical DQN
PyTorch强化学习实战(17)——强化学习训练加速
PyTorch强化学习实战(18)——基于DQN处理股票交易问题
PyTorch强化学习实战(19)——策略梯度法
PyTorch强化学习实战(20)——优势演员-评论家(Advantage Actor-Critic, A2C)
PyTorch强化学习实战(21)——异步优势演员-评论家(Asynchronous Advantage Actor-Critic, A3C)
PyTorch强化学习实战(22)——将强化学习应用于TextWorld互动小说游戏
PyTorch强化学习实战(23)——强化学习在网页导航中的应用
PyTorch强化学习实战(24)——连续动作空间中的强化学习
PyTorch强化学习实战(25)——深度确定性策略梯度(DDPG)
PyTorch强化学习实战(26)——提升随机策略梯度稳定性
PyTorch强化学习实战(27)——进化策略在强化学习中的应用
PyTorch强化学习实战(28)——遗传算法在强化学习中的应用
PyTorch强化学习实战(29)——深入解析强化学习中的探索策略
PyTorch强化学习实战(30)——基于人类反馈的强化学习(RLHF)
PyTorch强化学习实战(31)——AlphaGo Zero详解与实现
PyTorch强化学习实战(32)——MuZero详解与实现
PyTorch强化学习实战(33)——深度强化学习在离散优化中的应用

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进