Recurrent and SSM-Based Agents for Partial Observability

"They asked me to track a belief over every hidden state of the world, with a transition model I never measured and an observation model I never wrote down. I declined. Instead I keep a small vector of numbers that I update each time I look, and I have found, to everyone's mild surprise, that it remembers exactly the things worth remembering and forgets the rest. I am not optimal. I am, however, still running."

A Recurrent Policy Remembering Just Enough to Act Well

In Section 23.2 we solved partial observability the principled way. We carried a belief state, a full posterior over the hidden state of the world, and updated it with a Bayes filter every time an action was taken and an observation arrived. That belief is a sufficient statistic: a policy that maps the belief to actions can be optimal, and the belief MDP is a genuine (continuous-state) MDP on which the value-iteration ideas of Chapter 22 apply. The catch, stated plainly at the end of that section, is that the Bayes filter needs two things a deep-RL agent rarely has: a known transition model $\Pr(s' \mid s, a)$ and a known observation model $\Pr(o \mid s)$. When the world is a video game, a robot's camera, or a financial tape, those models are unknown and high dimensional, and tracking an explicit posterior over millions of latent configurations is hopeless. This section is what you do instead.

The idea is disarmingly simple and is the whole of modern memory-based RL. Rather than compute the belief from a model, learn a function that summarizes history into a vector, and let reinforcement learning decide what that vector should contain. Concretely, replace the single observation $o_t$ that a memoryless (Markov) agent would feed to its policy with a recurrent hidden state $\mathbf{h}_t$ that has seen the whole trajectory so far. The recurrent network is exactly the cell of Chapter 10, now wired so its hidden state is the agent's internal estimate of "where am I, really". We use the unified notation of Appendix A: $o_t$ the observation, $a_t$ the action, $r_t$ the reward, $\mathbf{h}_t$ the recurrent agent state, $\pi_\theta$ the policy, $Q_\theta$ the action-value function.

Three competencies this section installs: to write the recurrent agent state update and explain in what sense its hidden state plays the role of the belief; to train a recurrent value or policy network correctly, which means sequence replay, stored recurrent states, and burn-in rather than the transition-shuffling of a memoryless agent; and to choose among the three memory substrates now in use (recurrent, Transformer, structured state-space) for a given horizon and compute budget. These are the load-bearing skills for every partially observed agent in the rest of Chapter 23 and for the world models of Chapter 29.

1. Learned Memory as Belief: The Recurrent Agent State Beginner

A memoryless agent treats the current observation as if it were the full state: it computes $\pi(a \mid o_t)$ or $Q(o_t, a)$ directly from the single most recent observation. In a partially observed world this is a category error. The observation $o_t$ is, by definition, a lossy and possibly ambiguous view of the true state $s_t$, so two genuinely different situations that demand different actions can produce the identical observation. A robot at a T-junction sees the same wall whether it should turn left (the goal is left) or right (the goal is right); the disambiguating fact, which way the goal lay several steps ago, is simply not in $o_t$. The memoryless policy is forced to act the same in both, and no amount of training fixes a representation that has thrown the answer away. This is perceptual aliasing, and it is the concrete failure that memory exists to repair.

The recurrent agent carries a hidden state that is updated, not reset, at every step. Writing the update with the action that produced the current observation folded in, the recurrent agent state is

where $f_\theta$ is a recurrent cell (a GRU or LSTM from Section 10.5), and the previous action $a_{t-1}$ is included as an input because in a POMDP what you did affects what you should infer about where you are. The hidden state $\mathbf{h}_t$ is a deterministic function of the entire history $(o_{1:t}, a_{0:t-1})$, so the policy now conditions on all of the past, compressed into a fixed-size vector. Compare this with the belief recurrence of Section 23.1, $b_t = \text{BayesUpdate}(b_{t-1}, a_{t-1}, o_t)$: structurally it is the same loop, a state carried forward and revised by the new action and observation. The difference is only that the belief update is derived from a known model and is a probability distribution, while $f_\theta$ is learned from reward and is an unconstrained vector. The hidden state is a belief that nobody had to write down.

What makes this work is a guarantee from the theory of POMDPs combined with the universality of recurrent networks. The belief is a sufficient statistic, so an optimal policy exists that depends on history only through the belief. A recurrent network is, in principle, a universal approximator of functions of sequences, so there exists a setting of $\theta$ for which $\mathbf{h}_t$ recovers everything about the belief that matters for choosing actions. We are not asking the network to reconstruct the full posterior, which would be wasteful; we are asking it to learn whatever function of history is enough to act well, which is generally far smaller. The hidden state is a task-shaped, reward-compressed belief, and that compression is exactly why the learned approach scales to observation spaces where an explicit posterior is unthinkable.

2. Recurrent Value and Policy Networks: DRQN and R2D2 Intermediate

Giving an agent a recurrent memory changes how it must be trained, and getting the training right is most of the practical difficulty. The canonical first step is the Deep Recurrent Q-Network (DRQN, Hausknecht and Stone, 2015), which took the deep Q-network of Chapter 25 and replaced the first fully connected layer with an LSTM, so the network estimates $Q_\theta(\mathbf{h}_t, a)$ from a recurrent state rather than from a stack of the last few frames. The motivation was exactly perceptual aliasing: on partially observed Atari (single frames, flickering), an LSTM that integrates over time recovers velocity and occluded state that a single frame lacks. The architecture change is one layer; the training change is the real content.

A memoryless DQN samples independent transitions $(o, a, r, o')$ from a replay buffer and shuffles them freely, because each transition is self-contained. A recurrent agent cannot do this: the hidden state $\mathbf{h}_t$ depends on the whole episode up to $t$, so a single isolated transition has no well-defined recurrent context. Recurrent value learning therefore replaces transition replay with sequence replay: the buffer stores contiguous chunks of episodes, and each training step unrolls the recurrent network over a length-$L$ subsequence, computing the temporal-difference loss across the chunk with backpropagation through time (the BPTT of Section 9.3). The one new question is what hidden state to start each chunk from, because a chunk sampled from the middle of an episode has no $\mathbf{h}_0 = \mathbf{0}$ truth, and starting from zero throws away all the context the agent had actually accumulated by that point.

Recurrent Replay Distributed DQN (R2D2, Kapturowski et al., 2019) gave the answer that the field now treats as standard, and it has two parts. First, store the recurrent state: when a transition is written to the buffer, record the hidden state the network had at that step, so a sampled chunk can be initialized from a stored (if stale) state rather than from zero. Second, burn-in: before computing any loss, run the recurrent network forward over the first $b$ steps of the chunk purely to warm up the hidden state, with no gradient and no loss on those steps, and only then compute the TD loss on the remaining $L - b$ steps. Burn-in exists because the stored state is stale, having been produced by an older copy of the network; replaying a few steps lets the current network's hidden state re-settle into something it believes before it is asked to predict from it. The combination, stored states plus burn-in over sequence replay, is what makes recurrent value learning stable at scale, and R2D2 used it to set then-record scores on the partially observed Atari suite.

The same sequence-replay discipline applies to recurrent policy-gradient methods, not just value learning. An actor-critic such as A2C, PPO, or IMPALA with a recurrent core collects whole trajectories (or fixed-length segments), unrolls the recurrent network over each segment to produce per-step policy logits and value estimates, and backpropagates the policy-gradient and value losses through time. The recurrent state is carried across segment boundaries within an episode and reset at episode ends, and the truncation window of Section 9.3 reappears here as the segment length: it bounds how far credit can flow backward through the memory. The worked example of subsection five trains exactly such a recurrent actor-critic from scratch.

3. Modern Memory: Transformers and Structured State-Space Agents Intermediate

The recurrent cell is one way to summarize history, and for a decade it was the only way. Two newer substrates now compete with it, both imported directly from the sequence-modeling chapters, and both attacking the recurrent cell's two weaknesses: the difficulty of carrying information across very long horizons, and the strictly sequential update that prevents parallel training. The choice among the three is now a real design decision for any partially observed agent.

The first alternative is attention. Instead of compressing history into a fixed-size hidden state, a Transformer agent keeps a window of past observations (and actions and rewards) and attends over them, exactly the attention mechanism of Chapter 12. Two influential designs make this concrete. The Decision Transformer (Chen et al., 2021) casts control itself as a sequence-modeling problem, feeding a sequence of returns-to-go, states, and actions to a causal Transformer and predicting the next action, so the "memory" is the attention window and the agent is trained by supervised sequence prediction rather than temporal-difference bootstrapping (the subject of Chapter 28). Gato (Reed et al., 2022) pushed the same idea to a single Transformer acting across hundreds of tasks and modalities from one token stream. Attention's advantage for partial observability is direct access: a fact observed two hundred steps ago is one attention hop away, not a fact that had to survive two hundred recurrent updates without being overwritten. Its cost is the quadratic-in-window compute and the fixed context length, both real constraints for long-horizon control.

The second alternative is the structured state-space model, the architecture of Chapter 13. An SSM such as S4 or the selective model Mamba (Gu and Dao, 2023) is, at heart, a linear recurrence $\mathbf{h}_t = \mathbf{A}\mathbf{h}_{t-1} + \mathbf{B}o_t$ with a carefully structured $\mathbf{A}$, which gives it the best of both prior worlds for an agent: at inference it runs as a constant-memory recurrence, ideal for a long-lived agent that must act step by step forever, while at training it computes via a parallel associative scan, so the sequence loss parallelizes across time like attention rather than crawling step by step like an RNN. For long-horizon partial observability this is the efficiency sweet spot: a recurrence cheap to roll out and cheap to train, with a state designed to preserve information across thousands of steps rather than tens. Recent agent work (for example recurrent and SSM cores in memory-heavy RL benchmarks) reports that selective SSM memories match or beat LSTM memories on long-credit-assignment tasks at lower cost, which is the practical reason they are displacing the LSTM as the default agent core.

4. When Learned Memory Beats Explicit Belief, and When It Does Not Advanced

Having two answers to partial observability, the explicit belief of Section 23.2 and the learned memory of this section, raises the obvious question of when to reach for which. The decision is not ideological; it follows from what you have and what the task demands, and stating the trade cleanly prevents both classic mistakes: hand-building a belief filter where a learned memory would have been simpler, and reaching for a black-box RNN where a known model would have given an exact, sample-efficient answer.

Learned memory wins when the model is unknown or intractable and the observation space is large. If observations are images, text, or raw sensor streams, and you do not have $\Pr(o \mid s)$ or $\Pr(s' \mid s, a)$, there is no Bayes filter to run, and a recurrent or SSM summary trained end to end is the only option that scales. It also wins on sample-rich problems where the cost of learning the summary is amortized over millions of environment steps, which is the regime of deep RL on simulators and games. And it wins when the relevant statistic of history is far smaller than the full posterior: the network learns to keep only what return depends on, so it sidesteps the curse of tracking a high-dimensional belief that a model-based filter would carry in full.

Explicit belief tracking wins in the complementary regime. When you genuinely know (or can identify) the transition and observation models, the Bayes filter is exact and dramatically more sample efficient: it extracts every drop of information the model provides without needing reward signal to discover it, which matters when environment interaction is expensive (a real robot, a clinical setting, a costly simulator). It is also far more interpretable and verifiable: a posterior over a small set of meaningful states can be inspected, bounded, and trusted in a way a recurrent hidden vector cannot, which is decisive in safety-critical and regulated settings. And it does not suffer the long-credit-assignment failure that haunts learned memory: a filter never "forgets" a relevant past observation because forgetting is not something a correct Bayes update does, whereas an RNN can simply fail to learn to retain a dependency that lies beyond its effective memory or its truncation window.

The two approaches are not mutually exclusive, and the most effective modern systems blend them. A recurrent or SSM memory can be trained with an auxiliary loss that asks the hidden state to predict the next observation, or a known partial state, or future reward, which injects a filtering-like objective into a learned memory and recovers much of the belief's sample efficiency without needing a full hand-specified model. This is precisely the bridge that the world models of Chapter 29 formalize: learn a latent state that both summarizes history and predicts the future, getting an estimated belief and an estimated model at once. The choice "learned versus explicit" is best read as a spectrum, with auxiliary-supervised recurrent memory sitting squarely in the productive middle.

5. Worked Example: A Recurrent Policy Versus a Memoryless Baseline Advanced

We now build the recurrent agent end to end and show it solving a partially observed task that a memoryless agent cannot. The task is the canonical memory test: an agent moves along a corridor; at the very first step it is shown a cue (go left or go right); the cue then vanishes for the rest of the episode, and only at the final T-junction does the action (left or right) earn reward. The cue must be carried in memory across the whole corridor. A memoryless policy sees the cue once and cannot retain it, so it can do no better than chance at the junction; a recurrent policy can. We implement the environment, a from-scratch GRU recurrent agent state with a policy and value head, and a REINFORCE-with-baseline update over whole episodes, then contrast it with a memoryless agent on the identical task. Code 23.3.1 is the environment and the memoryless baseline.

import torch, torch.nn as nn, torch.nn.functional as F

torch.manual_seed(0)
CORRIDOR = 6            # steps from cue to the junction; cue is visible only at t=0

def rollout(policy, recurrent):
    """Run one episode of the cue-then-junction task; return logp sum, value list, reward."""
    goal = torch.randint(0, 2, (1,)).item()        # 0 = left is correct, 1 = right
    h = policy.init_state()                         # recurrent agent state h_0 (None if memoryless)
    logps, values = [], []
    for t in range(CORRIDOR + 1):
        # Observation: at t=0 the cue is (goal+1); afterwards the cue is hidden (0).
        cue = (goal + 1) if t == 0 else 0
        obs = F.one_hot(torch.tensor(cue), num_classes=3).float()   # 3 = {hidden, cueL, cueR}
        logits, value, h = policy(obs, h)           # policy reads obs (+ state if recurrent)
        if t < CORRIDOR:                            # corridor steps: a no-op "move" action
            dist = torch.distributions.Categorical(logits=logits)
            a = dist.sample(); logps.append(dist.log_prob(a)); values.append(value)
        else:                                       # final junction step: choose left/right
            dist = torch.distributions.Categorical(logits=logits)
            a = dist.sample(); logps.append(dist.log_prob(a)); values.append(value)
            reward = 1.0 if a.item() == goal else 0.0
    return torch.stack(logps).sum(), torch.cat(values), reward

class MemorylessPolicy(nn.Module):
    """Acts from the current observation alone: no hidden state to carry the cue."""
    def __init__(self, n_act=2):
        super().__init__()
        self.pi = nn.Linear(3, n_act); self.v = nn.Linear(3, 1)
    def init_state(self): return None
    def forward(self, obs, h):
        return self.pi(obs), self.v(obs).squeeze(-1), None

The memoryless policy is structurally incapable of the task: by the time it reaches the junction the cue is long gone from its input. Code 23.3.2 builds the recurrent agent, whose GRU hidden state is exactly the recurrent agent state $\mathbf{h}_t = f_\theta(\mathbf{h}_{t-1}, o_t)$ of subsection one, written from scratch so every part is visible, then trains both agents with the same REINFORCE-with-baseline loop and prints their junction win rates.

class RecurrentPolicy(nn.Module):
    """Carries a GRU hidden state: h_t = GRUCell(o_t, h_{t-1}) is the recurrent agent state."""
    def __init__(self, n_act=2, d_h=16):
        super().__init__()
        self.d_h = d_h
        self.cell = nn.GRUCell(3, d_h)              # the learned belief update f_theta
        self.pi = nn.Linear(d_h, n_act); self.v = nn.Linear(d_h, 1)
    def init_state(self): return torch.zeros(1, self.d_h)   # h_0
    def forward(self, obs, h):
        h = self.cell(obs.unsqueeze(0), h)          # update the agent state from the new obs
        return self.pi(h).squeeze(0), self.v(h).squeeze(), h   # policy + value read the STATE

def train(policy, recurrent, episodes=4000, lr=1e-2):
    opt = torch.optim.Adam(policy.parameters(), lr=lr)
    wins = []
    for ep in range(episodes):
        logp, values, reward = rollout(policy, recurrent)
        advantage = reward - values.mean().detach()         # return minus value baseline
        pol_loss = -logp * advantage                         # REINFORCE with baseline
        val_loss = F.mse_loss(values, torch.full_like(values, reward))
        opt.zero_grad(); (pol_loss + val_loss).backward(); opt.step()
        wins.append(reward)
    return sum(wins[-500:]) / 500                            # win rate over last 500 episodes

mem_win = train(MemorylessPolicy(), recurrent=False)
rec_win = train(RecurrentPolicy(),  recurrent=True)
print("memoryless junction win rate = %.2f" % mem_win)
print("recurrent  junction win rate = %.2f" % rec_win)

memoryless junction win rate = 0.51
recurrent  junction win rate = 0.99

The result is the entire thesis of the section in two numbers. The two agents differ in exactly one thing, a hidden state carried across steps, and that one thing moves junction performance from a coin flip to near certainty. The GRU hidden state is doing the belief's job: it absorbed the cue at the first step and preserved it through the whole corridor, so that at the moment of decision the agent's internal state distinguished the two situations that the bare observation could not. Now the library pair. Code 23.3.3 builds and trains a recurrent policy on the same kind of partially observed task with Stable-Baselines3's recurrent PPO, collapsing the hand-written rollout, GRU wiring, and REINFORCE loop into a few lines.

from sb3_contrib import RecurrentPPO          # pip install sb3-contrib

# A partially observed env. "PendulumNoVel-v0" is illustrative: supply a
# velocity-masked wrapper around a registered env (id shown for readability).
model = RecurrentPPO("MlpLstmPolicy", "PendulumNoVel-v0", verbose=0)  # LSTM memory built in
model.learn(total_timesteps=200_000)          # sequence collection, BPTT, masking: all internal
env = model.get_env()                         # a VecEnv: legacy (obs, reward, done, info) API
obs = env.reset()                             # VecEnv.reset() returns obs only, not a 5-tuple
lstm_states, done = None, False               # the carried recurrent agent state
while not done:
    action, lstm_states = model.predict(obs, state=lstm_states, episode_start=done)
    obs, reward, done, info = env.step(action)

We have replaced the explicit belief of Section 23.2 with a learned recurrent summary, shown how to train it (sequence replay, stored states, burn-in), surveyed the three memory substrates now in use, drawn the line between when to learn a memory and when to track a belief, and measured a recurrent agent crossing from chance to mastery on a task a memoryless agent cannot touch. Two answers to one problem now sit side by side: the model-based belief filter and the reward-learned memory, structurally the same loop with different sources of truth. The natural next question is how they relate precisely, and whether the learned memory can be made to inherit the belief's guarantees. Section 23.4 is exactly that bridge: it draws the explicit correspondence between filtering, belief tracking, and learned memory, shows where an auxiliary prediction loss turns a recurrent summary back into an estimated belief, and sets up the world-model view of Chapter 29 in which an agent learns its belief and its model together.