# Belief-State RWKV for Reinforcement Learning under Partial Observability

arXiv: https://arxiv.org/abs/2604.09671 · HTML: https://mlsanity.com/paper/2604.09671/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 19 of its 21 references that are in the map.
- First posted on arXiv 1 Apr 2026.

## Placed in

- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 19 of 19)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [RWKV: Reinventing RNNs for the Transformer Era](https://mlsanity.com/paper/2305.13048/) (2023)
- [Deep Recurrent Q-Learning for Partially Observable MDPs](https://mlsanity.com/paper/1507.06527/) (2015)
- [Stabilizing Transformers for Reinforcement Learning](https://mlsanity.com/paper/1910.06764/) (2019)
- [Eagle and Finch: RWKV with Matrix-Valued States and Dynamic Recurrence](https://mlsanity.com/paper/2404.05892/) (2024)
- [Deep Variational Reinforcement Learning for POMDPs](https://mlsanity.com/paper/1806.02426/) (2018)
- [Recurrent Model-Free RL Can Be a Strong Baseline for Many POMDPs](https://mlsanity.com/paper/2110.05038/) (2021)
- [Why Generalization in RL is Difficult: Epistemic POMDPs and Implicit Partial Observability](https://mlsanity.com/paper/2107.06277/) (2021)
- [POPGym: Benchmarking Partially Observable Reinforcement Learning](https://mlsanity.com/paper/2303.01859/) (2023)
- [Decision Mamba: Reinforcement Learning via Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2403.19925/) (2024)
- [A Survey of State Representation Learning for Deep Reinforcement Learning](https://mlsanity.com/paper/2506.17518/) (2025)
- [Predictive State Recurrent Neural Networks](https://mlsanity.com/paper/1705.09353/) (2017)
- [Accelerating Model-Based Reinforcement Learning with State-Space World Models](https://mlsanity.com/paper/2502.20168/) (2025)
- [KalMamba: Towards Efficient Probabilistic State Space Models for RL under Uncertainty](https://mlsanity.com/paper/2406.15131/) (2024)
- [Benchmarking Partial Observability in Reinforcement Learning with a Suite of Memory-Improvable Domains](https://mlsanity.com/paper/2508.00046/) (2025)
- [PIGDreamer: Privileged Information Guided World Models for Safe Partially Observable Reinforcement Learning](https://mlsanity.com/paper/2508.02159/) (2025)
- [Recurrent Predictive State Policy Networks](https://mlsanity.com/paper/1803.01489/) (2018)
- [Guided Policy Optimization under Partial Observability](https://mlsanity.com/paper/2505.15418/) (2025)

## Abstract (arXiv)

We propose a stronger formulation of RL on top of RWKV-style recurrent sequence models, in which the fixed-size recurrent state is explicitly interpreted as a belief state rather than an opaque hidden vector. Instead of conditioning policy and value on a single summary ht, we maintain a compact uncertainty-aware state bt = (μt, Sigmat) derived from RWKV-style recurrent statistics and let control depend on both memory and uncertainty. This design targets a key weakness of plain fixed-state policies in partially observed settings: they may store evidence, but not necessarily confidence. We present the method, a theoretical program, and a pilot RL experiment with hidden episode-level observation noise together with a test-time noise sweep. The pilot shows that belief-state policies nearly match the best recurrent baseline overall while slightly improving return on the hardest in-distribution regime and under a held-out noise shift. Additional ablations show that this simple belief readout is currently stronger than two more structured extensions, namely gated memory control and privileged belief targets, underscoring the need for richer benchmarks.