Belief-State RWKV for Reinforcement Learning under Partial Observability

Resource information

Authors

Liu Xiao

Research area

Machine learning

Application domain

General AI research

AI methods

Reinforcement learning

Research topics

cs.LG

arXiv record date

2026-04-01

Publication

arXiv

Upstream last modified

Apr 14, 2026

License

CC BY 4.0

Last record update

Nov 4, 2026

Abstract

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 h_t, we maintain a compact uncertainty-aware state b_t = (μ_t, Σ_t) 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.