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.namename self._state_typestate_type self.initial_stateinitial_state self._is_goal_predis_goal_pred self.action_enumaction_enum self._transform_functransform_func self._inverse_action_funcinverse_action_func self._render_funcrender_func self.encoded_shapeencoded_shape self._encode_funcencode_func魔方环境与 Gymnasium API 并不兼容本节使用这个示例来说明如何突破Gym的框架。(2)CubeEnv API中的部分方法仅是对构造函数所传函数的封装。这种设计允许新环境在独立模块中实现在环境注册表中完成注册并为代码其他部分提供统一接口def__repr__(self):returnCubeEnv(%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_actionNone):whileTrue:resself.action_enum(random.randrange(len(self.action_enum)))ifprev_actionisNoneorself.inverse_action(res)!prev_action:returnres(4)scramble()方法将动作列表(作为参数传入)应用于魔方初始状态返回最终状态defscramble(self,actions):sself.initial_stateforactioninactions:sself.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_inverseFalse,include_initialFalse):assertisinstance(scrambles_count,int)assertscrambles_count0stateself.initial_state result[]ifinclude_initial:assertnotreturn_inverse result.append((1,state))prev_actionNonefordepthinrange(scrambles_count):actionself.sample_action(prev_actionprev_action)stateself.transform(state,action)prev_actionactionifreturn_inverse:inv_actionself.inverse_action(action)res(depth1,state,inv_action)else:res(depth1,state)result.append(res)returnresult(6)explore_states()方法实现了自主迭代所需的功能对给定魔方状态应用所有可能动作。返回结果为元组其中第一个列表包含扩展后的状态第二个列表标记这些状态是否为目标状态defexplore_state(self,state):res_states,res_flags[],[]foractioninself.action_enum:new_stateself.transform(state,action)is_initself.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(namecube2x2,state_typeState,initial_stateinitial_state,is_goal_predis_initial,action_enumAction,transform_functransform,inverse_action_funcinverse_action,render_funcrender,encoded_shapeencoded_shape,encode_funcencode_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_typecube2x2 run_namepaper [train] cudaTrue lr1e-5 batch_size5000 scramble_depth200 report_batches10 checkpoint_batches100 lr_decayTrue lr_decay_gamma0.95 lr_decay_batches1000启动训练时需要将.ini文件传递给train.py工具。例如使用上述.ini文件训练模型的命令如下$ python ./train.py-iini/cube2x2-paper-d200.ini-nt1参数-n指定运行名称该名称将与.ini文件中的名称组合作为TensorBoard序列的名称。4. 搜索过程训练结果是包含网络权重的模型文件。该文件可通过MCTS方法用于求解魔方相关功能实现在solver.py工具和mcts.py模块中。solver工具非常灵活支持多种模式通过-p选项传递一个以逗号分隔的动作索引列表来解决一个单一的打乱魔方。例如-p 1,6,1表示依次执行第二个、第七个、第二个动作打乱的魔方。动作具体含义由-e选项指定的环境决定可在对应魔方环境模块中查看动作索引映射(如2×2魔方的动作1,6,1对应L→R′→L变换)从文本文件中读取排列(每行一个魔方)并求解。通过-i选项指定文件名(每行一个魔方状态)。cubes_tests文件夹提供若干测试问题集也可使用gen_cubes.py工具生成自定义随机问题集(支持设置随机种子、打乱深度等参数)生成指定深度的随机打乱并求解运行复杂度递增(打乱深度递增)的测试序列求解并将结果写入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)))yviamax(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)))0if 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 LearningRL详解PyTorch强化学习实战2——强化学习环境库GymnasiumPyTorch强化学习实战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步DQNN-step DQNPyTorch强化学习实战12——Double DQNDDQNPyTorch强化学习实战13——噪声网络NoisyNet-DQNPyTorch强化学习实战14——优先经验回放机制PyTorch强化学习实战15——Dueling DQNPyTorch强化学习实战16——Categorical DQNPyTorch强化学习实战17——强化学习训练加速PyTorch强化学习实战18——基于DQN处理股票交易问题PyTorch强化学习实战19——策略梯度法PyTorch强化学习实战20——优势演员-评论家Advantage Actor-Critic, A2CPyTorch强化学习实战21——异步优势演员-评论家Asynchronous Advantage Actor-Critic, A3CPyTorch强化学习实战22——将强化学习应用于TextWorld互动小说游戏PyTorch强化学习实战23——强化学习在网页导航中的应用PyTorch强化学习实战24——连续动作空间中的强化学习PyTorch强化学习实战25——深度确定性策略梯度DDPGPyTorch强化学习实战26——提升随机策略梯度稳定性PyTorch强化学习实战27——进化策略在强化学习中的应用PyTorch强化学习实战28——遗传算法在强化学习中的应用PyTorch强化学习实战29——深入解析强化学习中的探索策略PyTorch强化学习实战30——基于人类反馈的强化学习RLHFPyTorch强化学习实战31——AlphaGo Zero详解与实现PyTorch强化学习实战32——MuZero详解与实现PyTorch强化学习实战33——深度强化学习在离散优化中的应用
