Mitigating Partial Observability in Sequential Decision Processes via the Lambda Discrepancy
Cameron Allen, Aaron Kirtland, Ruo Yu Tao, Sam Lobel, Daniel Scott, Nicholas Petrocelli, Omer Gottesman, Ronald Parr, Michael L. Littman, George Konidaris
cs.LG, cs.AI, stat.ML
2024-07-10
Comparing TD value estimates at two different lambdas detects partial observability; minimizing it as an auxiliary loss improves recurrent PPO on all four hard POMDPs, never worse.
Most reinforcement learning quietly assumes the world is an MDP: given the current state, history stops mattering. Real environments often hand the agent only partial observations. In the T-maze, every corridor square looks identical, and the agent must remember the color seen at the start to pick the rewarding arm at the junction. The standard fixes are manual feature engineering (stacking past observations into the input) or bolting an RNN on and letting end-to-end training sort it out. Both share a blind spot: nothing tells you when the representation is good enough. Missing features silently cripple learning, and an RNN trained only on reward never receives a gradient that says its state is still missing information. This NeurIPS 2024 paper from Brown, Berkeley, Duke and Georgia Tech fills that gap.
The construction sits inside TD(λ), the classic interpolation between one-step bootstrapping and Monte Carlo, because the two extremes treat history differently.
If observations are Markovian, both converge to the same fixed point. If not, they can disagree, and that gap is the λ-discrepancy. Two theorems back the metric: in a block MDP (each observation generated by exactly one hidden state) it is exactly zero for every policy; and if any policy produces a nonzero discrepancy, almost all policies do, so probing with random policies almost surely reveals partial observability.
The T-maze makes the mechanism concrete: when seeing blue at the start guarantees +4 for going up at the junction, MC values the blue observation at +4 while TD(0) values it at 1.95, having averaged +4 and −0.1 into a coin flip.
Detection then turns into repair. The λ-discrepancy is differentiable with respect to memory-function parameters, so minimizing it by gradient descent pushes an RNN to encode the missing information. The deep version is deliberately plain: recurrent PPO with a GRU memory runs two critics with different λ values, and an auxiliary loss minimizes the squared difference between them, trained end to end with the usual PPO losses.
The small-scale check runs on seven classic POMDPs (T-maze, Tiger, Paint, Cheese Maze and others) with closed-form gradients: pick the random policy with the largest λ-discrepancy, minimize it to learn a memory function, then run policy gradient. Normalized performance (random policy = 0, optimal belief-state policy = 1) climbs with memory size and clearly beats the no-memory-optimization baseline; 30 seeds, 95% confidence intervals.
| Environment | Scale | Result |
| Battleship | 2^100 states | Significantly beats RNN-PPO and memoryless PPO |
| Partially observable PacMan | 5400 states | Same; memory visualization reconstructs dot locations |
| RockSample (11,11)/(15,15) | — | Baselines plateau at the +10 exit reward; the λ-discrepancy agent keeps improving |
All four environments show significantly better final performance and learning speed, and the augmented agent never does worse than the single-critic recurrent baseline. Two controls add weight: the hyperparameter sweep favored widely separated λ values near 0 and 1, matching the theory; and a variant with two value functions but no auxiliary loss was never selected, so the gain comes from the loss rather than doubled critic capacity.
This is a diagnostic built from parts every RL system already has. The λ-discrepancy is estimated from sampled observation-action pairs, needs no knowledge of the hidden state space and no environment model, just value functions, which RL already trains. Deciding whether an environment is partially observable, or whether a representation is sufficient, used to be intuition plus ablations; now it is a differentiable metric that plugs into training as one auxiliary loss plus one extra critic. The scope is stated honestly too: it says memory is needed, not what to remember. That is left to the minimization itself.
The authors list the main ones in an appendix. First, some POMDPs evade detection: in the Parity Check environment the λ-discrepancy is zero for every policy, because symmetry projects the difference away on both the reward and observation sides. The authors argue such cases sit on a knife's edge, since perturbing the parameters or adding a randomly initialized memory breaks them, but they offer no quantitative measure of how rare the cases are. Second, the deep version departs from the theory: the theorems cover fixed points of converged value functions, while the auxiliary loss minimizes the difference between unconverged estimates. The authors admit they do not know why this works so well in practice and flag it as future work. Third, the "large" environments are still classic discrete POMDP benchmarks, with no pixel observations and no continuous control. Results are reported as learning curves rather than a final-score table, and the "never worse" claim rests on the best configuration from the sweep.