MR bugfix: TIA segfault fix

[josephdviviano]

Summary

• Clamp player position registers to [0, 159] before ourPlayerPositionResetWhenTable lookups in RESP0/RESP1 (the crash site from Segfault in Montezuma Revenge #11)
• Replace single-step HMOVE wrapping with proper modular arithmetic for all 5 position registers
• Validate positions after state restore (TIA::load) where int16_t values are deserialized with no bounds check

Context

The segfault in #11 occurs at this line in TIA::poke (case 0x11 — Reset Player 1):

int8_t when = ourPlayerPositionResetWhenTable[myNUSIZ1 & 7][myPOSP1][newx];

The table is int8_t[8][160][160], but myPOSP1 is int16_t. If it escapes [0, 159], the access is out-of-bounds → segfault.

Analysis

In an attempt to reproduce #11 I ran the following, but it did not fail on my build - macOS ARM64 even after 5M+ frames of random play with bounds-check assertions compiled in.

#!/usr/bin/env python
"""Reproduce the TIA segfault in Montezuma's Revenge.

From https://github.com/Farama-Foundation/Arcade-Learning-Environment/issues/11:
Sending NOOPs until frame 49126, then action 17 (DOWNLEFTFIRE), triggers a
segfault in TIA::poke at the ourPlayerPositionResetWhenTable lookup.

This script tries multiple approaches:
1. The exact sequence from the bug report
2. Action 17 on every frame (maximise RESP1 trigger chances)
3. NOOPs with action 17 injected at various frame windows
4. Random actions stress test

Usage:
    conda run -n ale python scripts/reproduce_segfault.py
"""

import os
import sys
import random
import ale_py

# Prefer local ROM in assets/roms/, fall back to installed package
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
LOCAL_ROM = os.path.join(REPO_ROOT, "assets", "roms", "montezuma_revenge.bin")


def get_rom():
    if os.path.exists(LOCAL_ROM):
        return LOCAL_ROM
    from ale_py.roms import get_rom_path
    return str(get_rom_path("montezuma_revenge"))


def make_ale():
    ale = ale_py.ALEInterface()
    ale.setInt("frame_skip", 1)
    ale.setFloat("repeat_action_probability", 0.0)
    ale.setInt("random_seed", 0)
    return ale


def run_exact_sequence(ale):
    """Reproduce the exact sequence from the bug report."""
    CRASH_FRAME = 49126
    print(f"\n=== Test 1: Exact sequence (NOOP x {CRASH_FRAME - 1}, then action 17) ===")
    ale.loadROM(get_rom())

    resets = 0
    for frame in range(1, CRASH_FRAME + 1):
        action = 17 if frame == CRASH_FRAME else 0
        ale.act(action)
        if ale.game_over():
            resets += 1
            ale.reset_game()
        if frame % 10000 == 0:
            print(f"  frame {frame}")

    print(f"  Completed. Resets: {resets}. No crash.")


def run_action17_spam(ale, num_frames=200_000):
    """Send action 17 on every frame to maximise RESP1 hits."""
    print(f"\n=== Test 2: Action 17 every frame for {num_frames} frames ===")
    ale.loadROM(get_rom())

    resets = 0
    for frame in range(1, num_frames + 1):
        ale.act(17)
        if ale.game_over():
            resets += 1
            ale.reset_game()
        if frame % 50000 == 0:
            print(f"  frame {frame}: resets={resets}")

    print(f"  Completed. Resets: {resets}. No crash.")


def run_noop_then_17_sweep(ale, window_size=100_000):
    """Try action 17 at every possible frame offset in a long NOOP run."""
    print(f"\n=== Test 3: NOOP run, action 17 at every frame in [1, {window_size}] ===")
    ale.loadROM(get_rom())

    resets = 0
    for frame in range(1, window_size + 1):
        # Alternate: NOOP most of the time, action 17 every 3rd frame
        action = 17 if frame % 3 == 0 else 0
        ale.act(action)
        if ale.game_over():
            resets += 1
            ale.reset_game()
        if frame % 25000 == 0:
            print(f"  frame {frame}: resets={resets}")

    print(f"  Completed. Resets: {resets}. No crash.")


def run_random_stress(ale, num_frames=2_000_000, seed=42):
    """Random actions stress test."""
    print(f"\n=== Test 4: Random actions for {num_frames} frames (seed={seed}) ===")
    ale.loadROM(get_rom())
    legal_actions = list(ale.getLegalActionSet())

    rng = random.Random(seed)
    resets = 0
    for frame in range(1, num_frames + 1):
        ale.act(rng.choice(legal_actions))
        if ale.game_over():
            resets += 1
            ale.reset_game()
        if frame % 500000 == 0:
            print(f"  frame {frame}: resets={resets}")

    print(f"  Completed. Resets: {resets}. No crash.")


def main():
    print("Reproduction attempts for TIA segfault (issue #11)")
    print("Build has bounds-check assertions — abort() on out-of-bounds.\n")

    ale = make_ale()
    run_exact_sequence(ale)
    run_action17_spam(ale)
    run_noop_then_17_sweep(ale)
    run_random_stress(ale)

    print("\nAll tests passed without crash. The bug did NOT reproduce.")
    return 0


if __name__ == "__main__":
    sys.exit(main())

Claude audited every code path that modifies the position registers. The HMOVE wrapping arithmetic is actually correct for the motion table's value range of [-15, +8] — the single-step correction (if ≥ 160 then -= 160) maintains the [0, 159] invariant.

I built from source, ran 5M+ frames of Montezuma's Revenge with bounds-check assertions, and nothing triggered. I can't be sure this will fix #11, but I thought you would it interesting regardless.

Since the int16_t indices are used in array accesses with no validation, it's possible this is the fix. One confirmed unsafe path is state restore (TIA::load), where positions are deserialized from getInt() with no bounds check — RL methods using cloneState/restoreState could trigger this. Platform-specific compiler behaviour (Linux x86_64, aggressive optimizations) may also surface the issue during normal gameplay.