Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Markov and Partially Observable Decision Processes › Implementation: Memory in a T-Maze

Michael BrenndoerferJune 7, 20268 min read

World Models HandbookAll chapters

Part 2 of 12Inference, Decisions, and Control

Chapter 7Markov and Partially Observable Decision Processes › Implementation: Memory in a T-Maze

Reading section 10 of 14 · Implementation: Memory in a T-Maze · 8 minSection 10 of 14 · 8 min

Section 10 of 14 · 8 min read

Markov and Partially Observable Decision Processes › Implementation: Memory in a T-MazeOne section · Sections in this chapter10 / 14 ↓

Sections in this chapter14

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

We now measure what partial observability costs and what belief tracking buys. The testbed is the classic T-maze, stripped to its essentials. This example uses the undiscounted sum over its finite two-action episode, equivalently $\gamma=1$ for that bounded horizon.

At the start of each episode, the agent stands at the base of a T-shaped maze and glimpses a noisy cue telling it which arm holds the reward: the cue is correct with probability 0.9. The agent then advances to the junction. Crucially, at the junction the observation carries no information about the cue; the glimpse is gone. Entering the correct arm pays 4, entering the wrong arm pays -4, and every action costs 0.1. The latent variable "which arm is good" is drawn fresh each episode, so nothing can be memorized across episodes; the only path to above-chance performance runs through within-episode memory.

This tiny world isolates the phenomenon from the previous section: the optimal action at the junction depends on information that exists only in the agent's history. Let us code four agents that differ in how much episode information they can use: one reads the latent state directly, one remembers the cue as a belief, and two choose without retaining the cue.

In[18]:

Code

CUE_ACCURACY = 0.9
STEP_COST = 0.1
ARM_REWARD = 4.0

OBS_CUE_LEFT, OBS_CUE_RIGHT, OBS_JUNCTION, OBS_LEFT_END, OBS_RIGHT_END = range(
    5
)


def tmaze_step(obs, action, good, rng):
    """Enforce the two-stage protocol: advance from a cue, then choose an arm."""
    if obs in (OBS_CUE_LEFT, OBS_CUE_RIGHT):
        if action != 0:
            raise ValueError(
                "An arm cannot be chosen before reaching the junction"
            )
        return OBS_JUNCTION, -STEP_COST, False
    if obs != OBS_JUNCTION:
        raise ValueError("The episode has already terminated")
    if action not in (1, 2):
        return OBS_JUNCTION, -STEP_COST, False
    reward = ARM_REWARD if (action - 1) == good else -ARM_REWARD
    return OBS_LEFT_END + (action - 1), reward - STEP_COST, True


class OracleAgent:
    """Knows the latent variable directly; upper bound on performance."""

    def __init__(self, good):
        self.good = good

    def act(self, obs):
        return 0 if obs in (OBS_CUE_LEFT, OBS_CUE_RIGHT) else 1 + self.good


class MemoryAgent:
    """Maintains the belief P(good = right) implied by the cue, then acts on it."""

    def __init__(self, good=None):
        self.belief_right = 0.5

    def act(self, obs):
        if obs in (OBS_CUE_LEFT, OBS_CUE_RIGHT):
            self.belief_right = (
                CUE_ACCURACY if obs == OBS_CUE_RIGHT else 1.0 - CUE_ACCURACY
            )
            return 0
        return 2 if self.belief_right >= 0.5 else 1


class ReactiveAgent:
    """Maps the current observation to an action; no memory of the cue."""

    def __init__(self, good=None, junction_action=1):
        self.junction_action = junction_action

    def act(self, obs):
        return (
            0 if obs in (OBS_CUE_LEFT, OBS_CUE_RIGHT) else self.junction_action
        )


class RandomAgent:
    """Chooses an arm uniformly at the junction."""

    def __init__(self, rng, good=None):
        self.rng = rng

    def act(self, obs):
        if obs in (OBS_CUE_LEFT, OBS_CUE_RIGHT):
            return 0
        return 1 + int(self.rng.random() < 0.5)

Look closely at the MemoryAgent. Its single internal variable is the posterior $P(\text{good} = \text{right} \mid \text{cue})$ , computed by Bayes rule with a uniform prior. With cue accuracy 0.9, seeing the right cue sets the belief to 0.9 and seeing the left cue sets it to 0.1. The agent always chooses the arm indicated by the cue, so its choice agrees with the cue with probability 1 and with the truly rewarded arm with probability 0.9. That one number, carried across the delay, is the entire difference between this agent and the reactive one. Notice also that the oracle and the reactive agent occupy opposite poles of an information-access spectrum: the oracle reads the latent variable straight from the simulator, an upper bound no real sensor-based agent can reach, while the reactive agent discards the cue before choosing an arm.

Before simulating, we can compute the expected returns analytically. The oracle always earns $4 - 0.2 = 3.8$ . The memory agent picks the correct arm with probability 0.9, giving $0.9 \times 3.8 + 0.1 \times (-4.2) = 3.0$ . Among memoryless policies that commit to an arm at the junction, every arm-choice distribution wins half the time and therefore has expected return $0.5 \times 3.8 + 0.5 \times (-4.2) = -0.2$ . A policy that stalls instead of committing can do worse because it keeps paying the step cost. Simulation should reproduce the committing-policy values.

In[19]:

Code

def evaluate(make_agent, n_episodes=4000, seed=7):
    rng = np.random.default_rng(seed)
    action_rng = np.random.default_rng(seed + 1)
    totals = []
    for _ in range(n_episodes):
        good = int(rng.random() < 0.5)
        cue = good if rng.random() < CUE_ACCURACY else 1 - good
        agent = make_agent(good, action_rng)
        G, obs, done = 0.0, OBS_CUE_LEFT + cue, False
        for _ in range(10):
            G += 0.0
            obs, r, done = tmaze_step(obs, agent.act(obs), good, rng)
            G += r
            if done:
                break
        totals.append(G)
    return np.array(totals)


results = {
    "Oracle": evaluate(lambda g, rng: OracleAgent(g)),
    "Memory": evaluate(lambda g, rng: MemoryAgent()),
    "Reactive": evaluate(lambda g, rng: ReactiveAgent(junction_action=1)),
    "Random": evaluate(lambda g, rng: RandomAgent(rng)),
}

correct_arm_return = ARM_REWARD - 2 * STEP_COST
wrong_arm_return = -ARM_REWARD - 2 * STEP_COST
guess_probability = 0.5
theory = {
    "Oracle": correct_arm_return,
    "Memory": CUE_ACCURACY * correct_arm_return
    + (1 - CUE_ACCURACY) * wrong_arm_return,
    "Reactive": guess_probability * correct_arm_return
    + (1 - guess_probability) * wrong_arm_return,
    "Random": guess_probability * correct_arm_return
    + (1 - guess_probability) * wrong_arm_return,
}

Out[20]:

Console

Oracle    empirical +3.800   theory +3.800
Memory    empirical +2.982   theory +3.000
Reactive  empirical -0.214   theory -0.200
Random    empirical -0.246   theory -0.200

The empirical means match the analytical predictions within sampling error, and the gap between the memory agent and the reactive agent, more than three reward units per episode, is the measured price of forgetting. We predicted each agent's mean return in closed form before running the simulation. Agreement between analytic prediction and empirical sample checks both the environment code and our model of it, following the same discipline as the Monte Carlo check in the gridworld. Let us visualize the comparison.

Out[21]:

Visualization

Bar chart of average T-maze returns: memory agent near oracle, reactive agent negative. — Average return per episode in the T-maze across 4000 episodes (error bars show one standard deviation). Maintaining a belief from the cue recovers most of the oracle's performance, while reactive policies lose money on average.

The pattern in the bars tells the whole story of this chapter in one glance. The oracle bar is the ceiling imposed by the environment's latent variable. The memory bar nearly touches it, and the residual gap is $3.8-3.0=0.8$ reward units: exactly the cost of the cue being wrong on 10% of episodes, since each wrong arm changes the outcome by $3.8-(-4.2)=8.0$ . This is a perceptual limitation rather than a cognitive one. No amount of additional reasoning can extract information the sensor never delivered; this decomposition of performance loss into perceptual versus cognitive components is worth carrying into every applied project. The reactive and random bars sit below zero: an agent that cannot retain information is not merely suboptimal, it is actively losing, because half of its committed guesses walk into the punished arm. Suboptimality sounds benign; negative expected return does not, and the difference between the two words is the difference between a minor engineering flaw and a system you should switch off.

The error bars describe dispersion across episodes, not uncertainty in the plotted means. The oracle follows the same successful path every time, so its standard-deviation bar has zero width. The memory agent's return varies only with whether the cue is correct. The reactive agent always commits to the right arm, so its return varies with which arm the environment rewarded. For the random agent, a separate random-number stream makes its arm choice statistically independent of that rewarded-arm draw; the return still depends on both variables because it is determined by whether they match. Their intervals overlap, but that overlap should not be read as evidence that the policies have equal expected performance.

It is worth pausing on what the memory agent's belief is for. In this maze the belief never changes after the cue, because no later observation or action reveals anything about the latent variable; the Bayes update fires exactly once per episode and thereafter just holds steady. Real environments are kinder: actions usually generate informative observations, so beliefs evolve continuously as the agent explores. Tracking a belief over both position and latent variables, using the full predictor-corrector update, is precisely the business of the next chapter, Bayesian Filtering and Belief States. The T-maze is the degenerate but clarifying case where the filter reduces to remembering one number.