Belief-State RWKV for Reinforcement Learning under Partial Observability

Liu Xiao

cs.LGApr 1, 2026

Belief-State RWKV for Reinforcement Learning under Partial Observability

Authors: Liu Xiao

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.

Figures & tables

Figure 1: The proposed RWKV-first interface. Long observation history is compressed into a fixed-size belief state $(\mu_{t},\Sigma_{t})$ . Policy and value read from this uncertainty-aware state, while reward drives parameter updates rather than directly defining the instantaneous recurrent state.

Figure 2: RWKV-specific control interface. The belief readout branches from the temporal state produced by RWKV time-mix, then emits $(\mu_{t},\Sigma_{t})$ for policy and value heads. Channel-mix continues to refine the backbone feature stream, but the actor-critic reads belief rather than the raw hidden state.

Model	Mean return	Hard return	Very-hard return
MLP (memoryless)	0.909 $\pm$ 0.006	0.836 $\pm$ 0.017	0.809 $\pm$ 0.018
RWKV-style summary state	0.924 $\pm$ 0.008	0.861 $\pm$ 0.016	0.820 $\pm$ 0.021
Belief-state RWKV-style	0.919 $\pm$ 0.013	0.862 $\pm$ 0.025	0.822 $\pm$ 0.030

Table 1: Pilot results on the hidden-noise stop-or-guess task. The summary-state model performs best overall, but the belief-state model slightly improves return in the hardest noise regimes, suggesting that explicit uncertainty can help when partial observability is most severe.

Model	OOD mean return	OOD very-hard return
MLP (memoryless)	0.630 $\pm$ 0.008	0.630 $\pm$ 0.008
RWKV-style summary state	0.643 $\pm$ 0.011	0.643 $\pm$ 0.011
Belief-state RWKV-style	0.650 $\pm$ 0.015	0.650 $\pm$ 0.015

Table 2: Held-out noise-shift evaluation. The belief-state model performs best when trained on easier episodes and tested on a strictly harder observation noise regime, supporting the claim that explicit uncertainty is especially useful under distribution shift.

Model	ID return	OOD return	ID ECE	OOD ECE
RWKV-style summary state	0.924	0.643	0.024	0.164
Belief-state RWKV-style	0.919	0.650	0.025	0.160
Belief-state + gated memory	0.917	0.643	0.024	0.162
Belief-state + privileged targets	0.921	0.634	0.027	0.170

Table 3: Ablation and calibration summary. The plain summary-state RWKV baseline remains strongest on in-distribution return, but the plain belief-state RWKV variant is still the best simple out-of-distribution choice and has the best held-out calibration. Extra structure does not automatically help on this toy task: gated memory is competitive but not better, while privileged targets reduce robustness under shift.

Figure 3: Robustness sweep across fixed evaluation noise levels. The summary-state RWKV baseline remains strongest on the easy side of the training range, but the belief-state variant degrades more slowly once hidden noise approaches and exceeds the train-time ceiling. The right panel shows that this comes with a smoother increase in decision latency, consistent with uncertainty-aware waiting rather than indiscriminate hesitation.

Explore similar work

Why Linear Recurrent Memory Works in Partially Observable Reinforcement Learning
May 29, 2026Yike Zhao, Onno Eberhard, Malek Khammassi +2Partially Observable RL Markov Models
EPFL, Lausanne, Switzerland · Max Planck Institute for Intelligent Systems, Tübingen, Germany
Learning Suffers More Than the Policy Class Under Partial Observability: A Closed-Form Analysis
Aug 7, 2026Idil GözelRL Control Reinforcement Learning
University College London · Department of Physics and Astronomy University College London, London, UK
Does Learning Recurrent Networks Matter in Partially Observable RL?
Oct 9, 2026Zekang Wang, Zhe He, Gaojie JinPartially Observable RL Partially Observable Markov Decision Processes
University of Macau Macau SAR, China

May 29, 2026J/K move · Enter open · S save

Why Linear Recurrent Memory Works in Partially Observable Reinforcement Learning

Yike Zhao, Onno Eberhard, Malek Khammassi, Ali H. Sayed+1

EPFL, Lausanne, Switzerland · Max Planck Institute for Intelligent Systems, Tübingen, Germany

Partially Observable RL Markov Models Recurrent Neural Networks Partially Observable Markov Decision Processes State Representation Learning Linear RNNs+1

Open paper AskPDF arXiv

The family of linear recurrent neural networks has shown strong performance as recurrent memory units in partially observable reinforcement learning. We provide a theoretical justification for their empirical effectiveness by constructing and studying two linear filters: (i) the first exactly reproduces the pre-softmax logits of the belief vector in a hidden Markov model (HMM) under a deterministic transition matrix, thereby serving as a sufficient statistic for optimal policy learning, (ii) the second achieves vanishing state-decoding error under a nearly deterministic transition matrix, thus reducing state ambiguity to near zero. The results extend to action-controlled HMMs, where the corresponding linear filters become time-varying with action-dependent dynamics. We illustrate our main results through numerical experiments and further show that the constructed linear filter serves as a strong feature extractor in a small reinforcement learning game.

Abstract

Explore similar work

Why Linear Recurrent Memory Works in Partially Observable Reinforcement Learning

Learning Suffers More Than the Policy Class Under Partial Observability: A Closed-Form Analysis

Does Learning Recurrent Networks Matter in Partially Observable RL?