Latent State RL: VAE-based Implicit World Model for Post-Inference Training Optimization

The core innovations of this project are: using a Variational Autoencoder (VAE) to learn a compact implicit Markov state representation from inference trajectories, replacing traditional token history; and introducing an uncertainty-driven exploration mechanism, providing new state modeling ideas for post-training reinforcement learning methods like GRPO, which is expected to solve problems such as high computational overhead and key information being buried in long inference chain training.

In post-training reinforcement learning for large language models, the traditional approach of using complete token history as state input has limitations: sequence length grows linearly with inference steps leading to huge computational overhead; key information in long sequences is easily buried; and it's difficult to capture high-level abstract patterns. With the success of inference models like DeepSeek-R1 and OpenAI o1, post-training methods based on GRPO have gained attention, but how to extract meaningful state signals from trajectories remains an open question.

The core solution of Latent State RL is to use a VAE to encode inference trajectories (including token sequences, hidden layer states, and final rewards) into low-dimensional continuous vectors z, capturing key trajectory features while discarding redundant details. This implicit state has Markov property—current state z_t contains all information needed for the next decision, without needing to backtrack the complete token history, similar to the high-level understanding of a problem's core structure and progress by human experts.

The project introduces a cognitive uncertainty metric based on the variance of the VAE's posterior distribution: when encountering unfamiliar inference scenarios, the VAE encoding uncertainty increases (posterior variance expands), and this signal is used as part of the exploration reward to encourage attempts in high-uncertainty regions. Compared to traditional exploration strategies, it has advantages of context sensitivity (only explores when truly uncertain), interpretability (variance explicitly reflects confidence), and efficiency (avoids unnecessary exploration).

The project uses a four-stage experiment:

1. Phase A: Train a standard GRPO baseline on the MATH-Beyond benchmark, collect trajectory data, establish a performance ceiling, and provide VAE training data;
2. Phase B: Train the VAE to verify the latent space structure (correct/incorrect trajectories are distinguishable, variance does not collapse);
3. Phase C: Integrate the VAE encoder into the GRPO loop, where the policy receives implicit state z instead of original tokens, to verify the stability of joint training;
4. Phase D: Design four groups of controlled experiments (standard GRPO, token Markov state, VAE implicit state, VAE + uncertainty reward) to ensure fair comparison.

The project uses a modular code structure (directories like configs, scripts, eval), and the training script supports multiple configuration options:

• --state-mode: Select state representation method (token history, Markov token, VAE implicit);
• --uncertainty-bonus: Enable uncertainty reward;
• --freeze-vae: Freeze VAE parameters during joint training;
• --beta: Weight coefficient for uncertainty reward. Each experiment generates a manifest.json file that records configurations, random seeds, Git hashes, etc., to ensure reproducibility of results.

The significance of this work lies in:

1. Challenging the assumption that 'complete token history must be retained', demonstrating the feasibility of compressed state representation—if successful, it will significantly reduce the training cost of long inference chains;
2. Uncertainty exploration provides a new perspective for the RL exploration-exploitation trade-off, especially suitable for sparse reward environments like mathematical reasoning;
3. It is expected to be extended to fields requiring long-range reasoning such as code generation, theorem proving, and scientific discovery—any task involving multi-step decision-making plus intermediate evaluation may benefit.

As an ongoing project, there are still issues to be resolved:

• How to choose the dimension of the implicit state?
• How much trajectory data is needed for VAE training?
• How to adapt the uncertainty reward weight to task difficulty?
• Is the method equally effective across different inference tasks (mathematics, logic, common sense)? The answers to these questions will become clearer as the project progresses.