Belief-State RWKV for Reinforcement Learning under Partial Observability
For researchers working on reinforcement learning under partial observability, this work offers a simple yet effective method to incorporate uncertainty into recurrent policies, though the gains are incremental and limited to specific regimes.
The paper introduces a belief-state formulation for RWKV-style recurrent models in reinforcement learning under partial observability, where the recurrent state is explicitly interpreted as a belief state with uncertainty. In pilot experiments with hidden observation noise, the belief-state policy nearly matches the best recurrent baseline overall and slightly improves return on the hardest in-distribution regime and under a held-out noise shift.
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.