🤖 AI

A Differentiable Atari VCS:A Complex, Fully Known Ground Truth for Explainable AI

This paper introduces a differentiable, bit-for-bit accurate emulation of the Atari 2600 Video Computer System in Julia and JAX, providing a fully known, complex ground truth that enables the application of gradient-based methods for Explainable AI (XAI) to deep reinforcement learning tasks.

Original authors: Andreas Maier, Siming Bayer, Patrick Krauss

Published 2026-06-23

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Andreas Maier, Siming Bayer, Patrick Krauss

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: Explaining the Unexplainable

Imagine you are trying to explain how a magic trick works.

Scenario A: The trick is simple, like a card switch. You can see exactly how it's done. Explaining it is easy, but it's not very impressive.
Scenario B: The trick is incredibly complex, involving a giant, black-box machine with millions of gears. You can see the machine produce a result, but you have no idea what's happening inside. You can guess why it worked, but you have no way to prove if your guess is right or wrong.

This is the current state of Explainable AI (XAI). We have simple AI models we can explain perfectly, and complex AI models (like those playing video games) that we can't explain at all because we don't know their "ground truth" (the actual inner workings).

The Solution: Building a "Magic" Video Game Console

The authors decided to fix this by building a new kind of test object. They wanted something that is complex enough to be interesting (like a real computer) but fully known so we can check the answers.

They chose the Atari 2600, the classic video game console from the 1980s.

Why? It's a real computer with a real processor, memory, and graphics chip. It's complex enough to be a challenge, but small enough that we can map every single wire and switch inside it.
The Twist: They didn't just copy the game; they rebuilt the entire computer system from scratch in modern code, but with a superpower: it is "differentiable."

What Does "Differentiable" Mean? (The "Soft" Switch)

In a normal computer, decisions are "hard."

Example: "If the joystick is pushed right, move the sprite to column 10." It's either 10 or it isn't. If you try to calculate "how much" the joystick moved the sprite, the math breaks because the switch is binary (on/off).

The authors rebuilt the Atari so it can also run in "Soft Mode."

The Analogy: Imagine a dimmer switch instead of a light switch. Instead of the light being strictly "On" or "Off," it can be "50% On" or "73% On."
The Magic: They made the computer's logic (the CPU, the memory, the graphics chip) work like these dimmer switches.
- Forward Pass (The Result): The computer still plays the game perfectly. The screen looks exactly like the original Atari. The math proves that the "soft" version produces the exact same picture as the "hard" version.
- Backward Pass (The Explanation): Because the switches are "soft," the computer can now calculate gradients. It can answer questions like, "If I change this specific byte of memory by a tiny bit, how much does the pixel on the screen change?"

The Results: Two New Emulators

The team built two independent versions of this "Super Atari":

jutari: Written in the Julia programming language.
jaxtari: Written in JAX (a library for high-performance math).

Did it work?
Yes. They tested it against 64 different Atari games.

Memory Check: The memory (RAM) in their new versions was 100% identical to the original reference emulator, byte-for-byte.
Screen Check: The images on the screen were 100% identical, pixel-for-pixel.

This means their new "soft" computer is not a simulation; it is a perfect twin of the original, just with a special "explanation mode" turned on.

The "AI Assistant" Factor

The paper notes something surprising about how they built this.

The Claim: A project like this (rewriting a whole computer system) usually takes a human team many months or years.
The Reality: The authors built two complete, working versions in about 137 hours of active work (roughly 6 days).
How? They used AI coding agents (like advanced chatbots that write code) to do the heavy lifting. The humans acted as managers, checking the work and fixing bugs, while the AI wrote the bulk of the code.

Why This Matters (The "Ground Truth")

Before this paper, if you used an AI tool to explain why an Atari-playing AI made a move, you could never be sure if the tool was right. You could only say, "It looks plausible."

Now, because the Atari system is fully known and differentiable:

You can run an explanation tool on the system.
You can compare the tool's answer against the actual, known wiring of the computer.
You can finally say, "This explanation is correct," or "This explanation is wrong."

A Specific Example from the Paper

The authors tested this on the game Space Invaders.

The Problem: In the real game, moving the joystick right moves the cannon. But the computer does this by jumping to a specific memory address (a "hard" jump). A standard math tool would say, "The joystick has zero effect on the screen" because the math breaks at the jump.
The Fix: Using their "soft" version, they smoothed out that jump.
The Result: The math now correctly showed that pushing the joystick right lights up the edges of the cannon on the screen. They could prove the explanation was right because they knew exactly how the joystick was wired to the screen.

Summary

The authors built a perfectly known, complex video game computer that can also run in a "smooth" mode to allow for mathematical explanation. They proved that this smooth mode produces the exact same game as the real thing, but allows us to trace exactly how inputs turn into outputs. They built this incredibly fast using AI assistants, creating the first true "testbed" where we can verify if our AI explanations are actually telling the truth.

Technical Summary: A Differentiable Atari VCS

Problem Statement
The paper identifies a fundamental dichotomy in Explainable AI (XAI): systems are either simple and procedural (where the mechanism is known but trivial, rendering explanation tests meaningless) or genuinely complex (like deep neural networks) where XAI is needed but no ground truth of the inner functioning exists. Without ground truth, explanations can be plausible and confident yet entirely wrong, with no way to verify their correctness. The authors argue that the field lacks a study object that is simultaneously complex, fully specified (bit-for-bit known), and differentiable, which is necessary to validate gradient-based XAI methods against a known mechanism.

Methodology
To address this, the authors re-implemented the Atari 2600 Video Computer System (VCS)—the platform used for the first deep reinforcement learning demonstrations—as two independent, end-to-end differentiable emulators: jutari (in Julia) and jaxtari (in JAX).

Architecture Replication: The implementation covers the full VCS stack: the 6507 CPU (MOS 6502 core), the Television Interface Adapter (TIA), the RIOT chip (RAM, I/O, timer), cartridge bank-switching, and controllers. Both ports are validated against xitari, the C++ reference emulator used by DeepMind for the original DQN work.
Differentiable Formulation: The system treats the cartridge ROM as a weight tensor, the 128-byte RAM as an addressable "soft tape," and discrete control flow as gates. To enable gradient-based methods on this discrete hardware, three specific relaxations are applied:
- Memory Access: ROM and RAM reads are expressed as dot products against one-hot address vectors, making the gradient with respect to specific memory bytes defined.
- Opcode Dispatch: The selection of instruction handlers is modeled as a softmax over candidate outputs, controlled by a temperature parameter $T$ .
- Control Flow: Conditional branches are relaxed using sigmoid gates with a sharpness parameter $\alpha$ .
Forward-Backward Decoupling: The authors employ a Straight-Through Estimator (STE). In the forward pass, the system operates in "hard mode," producing bit-exact results identical to the original hardware. In the backward pass, the system operates in "soft mode," utilizing the gradients of the relaxed functions. This ensures the forward trajectory is exact while exposing surrogate gradients where the original bit logic has none.
Validation: The ports were subjected to rigorous conformance testing across all 64 supported Arcade Learning Environment (ALE) games, comparing internal states (RAM) and outputs (pixel buffers) byte-for-byte and pixel-for-pixel against xitari.

Key Contributions

Two Differentiable Ports: Independent Julia and JAX implementations that reproduce xitari exactly on all 64 games (64/64 RAM-exact and 64/64 pixel-exact).
Theoretical Guarantees:
- Theorem 1: Proves that the soft forward pass is bit-exact equal to the hard forward pass at any finite temperature when using the STE.
- Theorem 2: Provides an error bound for a fully relaxed variant (without STE) as temperature approaches zero and sharpness approaches infinity.
AI-Assisted Engineering: A quantified report demonstrating that a complex emulator port, traditionally requiring many person-months, was built in approximately 137 hours of active work over 29 days, with large portions written autonomously by coding agents.
Audit Tool: The project demonstrates that a bit-exact re-implementation can serve as an audit tool for the reference emulator itself, having uncovered a latent non-determinism bug in xitari related to uninitialized RAM seeding.
Qualitative Gradient Study: A demonstration that the substrate allows for checking attributions against known mechanisms, specifically showing how a differentiable sampler can recover the joystick-to-screen control path in Space Invaders, whereas a naive gradient vanishes.

Results

Conformance: Both ports achieved perfect conformance with the reference xitari across the entire 64-game ALE set, with zero differing bytes in RAM and zero differing pixels in the frame buffer.
Performance: The JAX port (jaxtari) enables batched, fully differentiable rollouts reaching approximately 3.1 million environment-steps per second on a single commodity GPU.
Gradient Validity: In a proof-of-concept study using the Space Invaders ROM, the system successfully generated gradients that correctly identified the moving cannon edges and the joystick control mapping. The naive gradient (without the differentiable sampler) was zero, correctly reflecting the discrete nature of the hardware, while the proposed method recovered the true causal link.

Significance and Claims
The paper positions this work as the construction of a foundation rather than a final XAI benchmark. Its primary significance is providing the first study object of genuine complexity that is fully specified and differentiable. This allows, for the first time, the verification of XAI attributions against a known mechanism.

The authors are modest in their claims:

They do not claim to rank XAI methods or evaluate learned agents (policies) in this paper.
They clarify that the gradients provided are surrogate estimators of the hard map, not the hard derivative itself.
They emphasize that the goal is to enable the community to develop and validate XAI tools against a known truth, moving beyond the "plausible but unverified" explanations that currently dominate the field.

The work also highlights the potential of AI-assisted engineering, showing that coding agents can significantly accelerate the development of complex, verified software systems when paired with strict, automated conformance oracles.