Latent State RL：基于VAE的隐式世界模型用于推理后训练优化

本项目提出核心创新：使用变分自编码器（VAE）从推理轨迹中学习紧凑的隐式马尔可夫状态表示，替代传统token历史；同时引入不确定性驱动的探索机制，为GRPO等强化学习后训练方法提供新的状态建模思路，有望解决长推理链训练的计算开销大、关键信息淹没等问题。

在大语言模型的强化学习后训练中，传统以完整token历史为状态输入的方式存在局限：序列长度随推理步骤线性增长导致计算开销巨大；长序列关键信息易被淹没；难以捕捉高层抽象模式。随着DeepSeek-R1、OpenAI o1等推理模型的成功，基于GRPO的后训练方法受关注，但如何从轨迹中提取有意义的状态信号仍是开放问题。

Latent State RL的核心方案是用VAE将推理轨迹（含token序列、隐藏层状态、最终奖励）编码为低维连续向量z，捕获轨迹关键特征并丢弃冗余细节。该隐式状态具有马尔可夫性——当前状态z_t包含下一步决策所需全部信息，无需回溯完整token历史，类似人类专家对问题核心结构和进展的高层理解。

项目引入基于VAE后验分布方差的认知不确定性度量：当遇到陌生推理情境时，VAE编码不确定性增加（后验方差扩大），此信号作为探索奖励一部分，鼓励在高不确定性区域尝试。相比传统探索策略，它具有情境敏感（只在真不确定时探索）、可解释（方差显式反映置信度）、高效（避免无谓探索）的优势。

项目采用四阶段实验：

1. Phase A：在MATH-Beyond基准训练标准GRPO基线，收集轨迹数据，建立性能天花板并提供VAE训练数据；
2. Phase B：训练VAE验证隐空间结构（正确/错误轨迹可区分、方差未塌陷）；
3. Phase C：将VAE编码器接入GRPO循环，策略接收隐状态z而非原始token，验证联合训练稳定性；
4. Phase D：设计4组对照实验（标准GRPO、token马尔可夫状态、VAE隐状态、VAE+不确定性奖励），确保公平比较。

项目采用模块化代码结构（configs、scripts、eval等目录），训练脚本支持多种配置选项：

• --state-mode：选择状态表示方式（token历史、马尔可夫token、VAE隐式）；
• --uncertainty-bonus：启用不确定性奖励；
• --freeze-vae：联合训练中冻结VAE参数；
• --beta：不确定性奖励权重系数。 每个实验生成manifest.json记录配置、随机种子、Git哈希等，确保结果可复现。

本工作的意义在于：

1. 挑战“必须保留完整token历史”的假设，展示压缩状态表示的可行性，若成功将大幅降低长推理链训练成本；
2. 不确定性探索为RL探索-利用权衡提供新视角，尤其适用于数学推理等稀疏奖励环境；
3. 有望推广到代码生成、定理证明、科学发现等需长程推理的领域，任何多步决策+中间评估的任务都可能受益。

作为在研项目，仍有问题待解决：

• 隐式状态维度如何选择？
• VAE训练需要多少轨迹数据？
• 不确定性奖励权重如何与任务难度自适应？
• 方法在不同推理任务（数学、逻辑、常识）上是否同样有效？ 这些问题的答案将随项目进展逐步明朗。