mlsanityData to 7 Oct 2026

Belief-State RWKV for Reinforcement Learning under Partial Observability

arXiv 2604.09671 ↗First posted 1 Apr 2026cs.LGcited by 0

Not on the map. Why?

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.

Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 19 of its 21 references that are in the map. First posted on arXiv 1 Apr 2026.

Builds on

19 of 21 references.

References · 19

  1. 2017Proximal Policy Optimization Algorithmscited by 14,441
  2. 2023Mastering Diverse Domains through World Modelscited by 1,193
  3. 2023RWKV: Reinventing RNNs for the Transformer Eracited by 814
  4. 2015Deep Recurrent Q-Learning for Partially Observable MDPscited by 457
  5. 2019Stabilizing Transformers for Reinforcement Learningcited by 262
  6. 2024Eagle and Finch: RWKV with Matrix-Valued States and Dynamic Recurrencecited by 155
  7. 2018Deep Variational Reinforcement Learning for POMDPscited by 122
  8. 2021Recurrent Model-Free RL Can Be a Strong Baseline for Many POMDPscited by 115
  9. 2021Why Generalization in RL is Difficult: Epistemic POMDPs and Implicit Partial Observabilitycited by 102
  10. 2023POPGym: Benchmarking Partially Observable Reinforcement Learningcited by 59
  11. 2024Decision Mamba: Reinforcement Learning via Sequence Modeling with Selective State Spacescited by 23
  12. 2025A Survey of State Representation Learning for Deep Reinforcement Learningcited by 17
7 more
  1. 2017Predictive State Recurrent Neural Networkscited by 14
  2. 2025Accelerating Model-Based Reinforcement Learning with State-Space World Modelscited by 10
  3. 2024KalMamba: Towards Efficient Probabilistic State Space Models for RL under Uncertaintycited by 9
  4. 2025PIGDreamer: Privileged Information Guided World Models for Safe Partially Observable Reinforcement Learningcited by 7
  5. 2018Recurrent Predictive State Policy Networkscited by 6
  6. 2025Benchmarking Partial Observability in Reinforcement Learning with a Suite of Memory-Improvable Domainscited by 6
  7. 2025Guided Policy Optimization under Partial Observabilitycited by 3

Built on by

Not cited yet.