Windows/Pygame play mode may miss manual keyboard input in `src/play.py --pretrained`

[lidongyueis007-art]

Hi, thanks for releasing DIAMOND. I tested the pretrained Atari play mode on Windows and found a small input-handling issue that may affect manual control in the Pygame viewer.

Summary

When running:

python src/play.py --pretrained

and selecting Pong, the pretrained checkpoint and world model loaded correctly. The game window rendered normally, and switching between world model / real Atari environments worked.

However, after switching to human/manual control, manual keyboard input was not captured by Pygame. The UI stayed on noop, and the paddle did not respond.

Environment

• OS: Windows
• Python: 3.10.20
• PyTorch: 2.4.1+cu124
• GPU: NVIDIA GeForce RTX 4070 Laptop GPU
• Pygame: 2.5.2
• Gymnasium: 0.29.1
• ale-py: 0.9.0
• Launch mode: PowerShell, src/play.py --pretrained
• Tested game: Pong

Debugging notes

I added temporary logging around the input loop.

Programmatic action injection worked correctly:

ACTION total frames: 433
non-noop action frames: 433
action=2: 216
action=3: 217

So the model, environment, and action mapping were OK.

But manual keyboard input produced:

ACTION total frames: 2093
non-noop action frames: 0
KEYDOWN count: 0
KEYUP count: 0
action=0: 2093

This suggests the Pygame window rendered correctly, but did not receive manual keyboard events on this Windows setup.

Suggested improvement

It might help to make play mode more robust by:

1. Tracking held keys via KEYDOWN / KEYUP.
2. Adding optional keyboard/action debug logging.
3. Printing the active keymap for the selected Atari game.
4. Optionally adding a Windows fallback using GetAsyncKeyState when Pygame reports no key events.

In my local test, adding a Windows fallback via ctypes.windll.user32.GetAsyncKeyState(...) made manual left/right control work.

This appears to be an input-capture/window-focus issue, not a model-loading or training issue. Thanks again for the project!