Belief-State RWKV for Reinforcement Learning under Partial Observability

TL;DR

This paper introduces a belief-state formulation for RL with RWKV models, explicitly modeling uncertainty to improve decision-making under partial observability.

Contribution

It proposes a belief-state recurrent model that explicitly encodes uncertainty, enhancing RL performance in partially observed environments.

Findings

Belief-state policies nearly match the best recurrent baseline.

They slightly improve return on the hardest in-distribution regime.

The simple belief readout outperforms more structured extensions.

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 = (\mu_t, \Sigma_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.