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

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

- [0. 前言](#0. 前言)
- [1. 通用魔方求解框架](#1. 通用魔方求解框架)
- [2. 魔方环境](#2. 魔方环境)
- [3. 模型训练](#3. 模型训练)
- [4. 搜索过程](#4. 搜索过程)
- [5. 实验结果](#5. 实验结果)
- - [5.1 2×2 魔方](#5.1 2×2 魔方)
  - [5.2 3×3 魔方](#5.2 3×3 魔方)
- [6. 模型改进方向](#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 等魔方可能存在多个最终状态，因此需要独立判定函数覆盖此类情况
可应用于状态的动作枚举
变换函数：接收状态和动作，返回结果状态
逆函数：将每个动作映射为其逆动作
渲染函数：将状态转换为人类可读形式
编码状态张量的形状
将紧凑状态表示编码为神经网络友好格式的函数

python 复制代码

class CubeEnv:
    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 中的部分方法仅是对构造函数所传函数的封装。这种设计允许新环境在独立模块中实现，在环境注册表中完成注册，并为代码其他部分提供统一接口：

python 复制代码

    def __repr__(self):
        return "CubeEnv(%r)" % self.name

    # wrapper functions
    def is_goal(self, state):
        assert isinstance(state, self._state_type)
        return self._is_goal_pred(state)

    def transform(self, state, action):
        assert isinstance(state, self._state_type)
        assert isinstance(action, self.action_enum)
        return self._transform_func(state, action)

    def inverse_action(self, action):
        return self._inverse_action_func(action)

    def render(self, state):
        assert isinstance(state, self._state_type)
        return self._render_func(state)

    def encode_inplace(self, target, state):
        assert isinstance(state, self._state_type)
        return self._encode_func(target, state)

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

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

python 复制代码

    def sample_action(self, prev_action=None):
        while True:
            res = self.action_enum(random.randrange(len(self.action_enum)))
            if prev_action is None or self.inverse_action(res) != prev_action:
                return res

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

python 复制代码

    def scramble(self, actions):
        s = self.initial_state
        for action in actions:
            s = self.transform(s, action)
        return s

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

python 复制代码

    def is_state(self, state):
        return isinstance(state, self._state_type)

    def scramble_cube(self, scrambles_count, return_inverse=False, include_initial=False):
        assert isinstance(scrambles_count, int)
        assert scrambles_count > 0

        state = self.initial_state
        result = []
        if include_initial:
            assert not return_inverse
            result.append((1, state))
        prev_action = None
        for depth in range(scrambles_count):
            action = self.sample_action(prev_action=prev_action)
            state = self.transform(state, action)
            prev_action = action
            if return_inverse:
                inv_action = self.inverse_action(action)
                res = (depth+1, state, inv_action)
            else:
                res = (depth+1, state)
            result.append(res)
        return result

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

python 复制代码

    def explore_state(self, state):
        res_states, res_flags = [], []
        for action in self.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)
        return res_states, res_flags

通过这种通用功能，类似环境可以轻松实现并接入现有的训练和测试方法，只需极少的样板代码。作为示例，我们提供了实验使用的 2×2 魔方和 3×3 魔方环境，其具体实现位于 cube2x2.py 和 cube3x3.py中，可作为实现自定义环境的基础。

每个环境需要通过创建 CubeEnv 类实例，并将实例传递给 _env.py 中定义的 register() 函数来完成注册。以下是 cube2x2.py 模块中的相关代码片段：

python 复制代码

_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：

xml 复制代码

[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 文件训练模型的命令如下：

shell 复制代码

$ python ./train.py -i ini/cube2x2-paper-d200.ini -n t1

参数 -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 选项指定模型权重文件。此外还有其他参数可用于调整 MCT S选项和时间/搜索步数限制，具体选项名称可在 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)))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)。

shell 复制代码

$ python solver.py -m saves/cube2x2-paper-d200-t1/chpt_016700.dat -e cube2x2 -i cubes_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强化学习实战——从零开始实现魔方求解器

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

0. 前言

1. 通用魔方求解框架

2. 魔方环境

3. 模型训练

4. 搜索过程

5. 实验结果

5.1 2×2 魔方

5.2 3×3 魔方

6. 模型改进方向

相关链接