MNIST-PRO: MNIST is Back as a Partially Observable World for AI Agents

Look at the zero above. With the whole image in view, recognizing it is easy. Now cover everything except a small square near the bottom-right stroke. That fragment could belong to several digits. To answer, you have to move the window, remember the earlier fragments, and work out how they fit together.

In the illustrated episode, the agent moves up along the stroke, reaches the top of the digit, and eventually answers 0. The interesting part is not the label. It is the work between the first glimpse and that answer: deciding where to look, keeping track of position, and deciding when the observations support a conclusion.

We built MNIST-PRO to study that work. In this note, I want to focus on what a wrong answer can tell us—and why collecting more evidence is not always the same as making better use of it.

Why return to handwritten digits?

Imagine navigating a room with only a flashlight. You need to combine successive views into a useful picture of your surroundings. But if a robot fails in that room, the cause may be difficult to identify: it could have misread an object, forgotten a location, collided with something, or failed to execute an action.

MNIST gives us a simpler setting. There is no physics to simulate and no object to manipulate. The agent moves a window over a static image and reports a digit. We can record exactly which parts it saw before it answered.

The window is 64 × 64 pixels and moves 32 pixels at a time, up, down, left or right. Level 1 places one digit on a 224 × 224 canvas. Level 2 places two digits side by side on a 224 × 448 canvas and asks for their ordered sequence: seeing a 5 and an 8 is not enough if the answer should be 58, not 85.

This keeps the recognition problem familiar while changing how the evidence arrives. It also gives us two separate questions to ask of an unsuccessful episode: did the agent look in the right places, and could it use what it had already seen?

What survives after the window moves?

Keeping every image in context sounds like an obvious solution. That is our Image Only baseline: the agent retains its previous glimpses and actions. But a sequence of crops is not a map. The model still has to work out which views overlap and where each fragment belongs.

We compare this with two ways of writing information down. In Textual State, only the current image is visible; descriptions of earlier observations carry information forward. In Metric Grid Map, the agent also records features against relative coordinates in a structured map. A description tells it what it saw. Coordinates are intended to help it remember where.

These representations impose different demands. A text description can omit a small visual detail. A map can contain an incorrect position. A long visual history can contain all the fragments without making their spatial relationship clear. None of these formats guarantees that the agent has a reliable picture of the digit.

The full-image control makes the contrast concrete. Gemini 3.1 Pro reaches 99% accuracy on the single-digit control but 38% in the Image Only condition. The digit has not become a different object; the model now has to gather and combine views of it.

Did it actually see enough?

Accuracy alone does not answer this. We also measure stroke coverage: how much of the digit’s foreground was exposed through the glimpse window. This is different from the number of moves. An agent can spend many steps returning to the same region.

The plot shows why coverage is worth measuring separately. Claude Fable 5 covers more of the strokes than Gemini 3.7 Flash in this comparison, yet answers fewer episodes correctly. In the Textual State condition specifically, Fable and Opus exceed 80% stroke coverage, while their Level 1 accuracies are 38% and 41%; their Level 2 accuracies are 18% and 13%.

That does not make coverage useless. Nor does it prove that memory is the sole cause of failure. A large covered area may still leave out the part that distinguishes two digits. What it tells us is that “the agent should look more” is not a complete diagnosis.

To go further, we need an intervention that changes how the evidence is presented without giving the agent another chance to explore.

Keep the path fixed; put the glimpses together

We take completed Image Only and Textual State trajectories and arrange the already observed glimpses on a coordinate-aligned canvas. The model then makes a final prediction from this consolidated view. The route stays fixed. No new region of the digit is revealed.

This matters because a better answer can no longer be attributed to a better search path. We have changed the organization of the evidence available at decision time.

All eight model settings plotted improve on average. Looking at the Textual State trajectories separately makes the size of the change easier to see:

These results support a narrower, useful conclusion: for these trajectories, much of the error can be reduced by organizing the observations differently. They do not show that exploration no longer matters. A canvas cannot recover a stroke that the agent never saw. And the intervention supplies spatial alignment, so it does not separate every possible failure in remembering, aligning and interpreting the fragments. The per-condition results are in Table 2 of the paper.

Some agents stop before reaching the second digit

The other side of the diagnosis is visible in the trajectories. In an audit of 1,600 Level 2 episodes using Textual State or Metric Grid Map, 323 predictions contained only one digit. In 238 of those episodes, the agent never looked at the second digit at all. In 266, it exposed less than a quarter of that digit. These single-digit answers arrived after roughly 15 steps on average, out of an available 78.

Here, reorganizing the existing evidence cannot solve the missing-observation problem. The agent has committed before checking the whole task.

We also see redundant exploration and failures to revise early hypotheses. For Gemini 3.7 Flash, adding a Metric Grid Map reduces the revisit rate from 45.1% to 13.6%. That is a substantial change in where it spends its moves, but a more efficient route is not itself a correct answer. Likewise, an early guess must remain revisable: later observations are useful only if they can change the prediction.

These failures call for different tests. An episode that never reaches the second digit needs a different explanation from one that observes both digits and still answers incorrectly.

What if the agent can write its own tools?

The offline canvas is something we construct for the model. We also tested a harness in which agents can write and execute Python and manage local files. Could they build useful representations during the task themselves?

In the Gemini 3.7 Flash runs, the agent wrote scripts such as stitch.py and render_mnist.py to assemble glimpses into a coordinate-aligned matrix, inspect an ASCII rendering, and guide further exploration. Its harness accuracy was 88% on Level 1 and 63% on Level 2.

This is a different experiment from replaying a fixed trajectory: tools can change both the representation and the search that follows. The result is encouraging, but it should not be read as an isolated measurement of memory quality. It shows a model using computation to make its observations easier to work with.

What I would check before adding more context

For an agent that fails on a partially observed task, I would first inspect what it had actually seen at the time it answered. Did it reach every relevant object? Did it revisit the same area? Did it stop while much of the task remained unseen?

Then I would hold that evidence fixed and change its representation. If the answer improves when the same observations are organized into a map, collecting additional observations is not the only route to improvement. If crucial evidence is absent, the search and stopping policy need attention too.

MNIST-PRO lets us make those checks in a small, controlled world. Handwritten digits do not establish how an agent will behave on a website or in a physical room. They do give us a way to tell apart errors that a single accuracy score would otherwise hide.

Citation

@misc{toh2026mnistpromnistpartiallyobservable,
  title={MNIST-PRO: MNIST is Back as a Partially Observable World for AI Agents},
  author={Vernon Toh and Navonil Majumder and Zhengyuan Liu and Nancy F. Chen and Soujanya Poria},
  year={2026},
  eprint={2608.31022},
  archivePrefix={arXiv},
  primaryClass={cs.AI},
  url={https://arxiv.org/abs/2608.31022}
}