ALE State clone/restore: RAM and rendered frame are not updated correctly

[wookayin]

To dump(clone) or load(restore) underlying game state, one can use the following methods:

e_state = env.unwrapped.clone_state()    # returns a 1-D vector
env.unwrapped.restore_state(e_state)

The clone_state() and restore_state() functions are actually from gym (see atari_env.py). However it is simply calling ALE's cloneState() and decodeState().

The issue is: after restoring game state, screens and RAM entries are NOT updated. To my best of knowledge they are ONLY updated after reset() and step(). As a result, we only end up getting last frame/RAM information before reset.

Demonstration

e = make_atari('SeaquestNoFrameskip-v4')

# Initialize: 
e.reset()
for i in range(100):
    s, _, _, _ = e.step(0)  # stay, no-op

plt.imshow(s)
plt.show()

e_state = e.unwrapped.clone_state()
print("Saved state (the submarine is on the surface) :", e_state.shape)
print("RAM state: ", e.unwrapped.ale.getRAM()[:16])

print("------------------------------------\n\n")

for i in range(100):
    s, _, _, _ = e.step(5)   # go down

print("After 100 actions, now the submarine is at the floor.")

plt.imshow(s)
plt.show()

print("RAM state: ", e.unwrapped.ale.getRAM()[:16])
print("------------------------------------\n\n")


print("Now let's restore. The submarine should be on the surface")
e.unwrapped.restore_state(e_state)

s = e.render('rgb_array')         # not restored according to e_state
#s, _, _, _ = e.step(0)           # re-rendered to 1-step after e_state

plt.imshow(s)
plt.show()
print("RAM state: ", e.unwrapped.ale.getRAM()[:16])

One can see that the rendered frame (which is from getScreenRGB2()) is not updated. It should correspond to the observation of the dumped game state, but it is not. Same thing for RAM.

However, restoring game state itself works well. You can verify that by taking additional one environment step (e.g. e.step(0) as commented).

I believe this is an issue where we need to fix ALE C++ implementation, but for right now I am just reporting the issue. Any idea? Thanks!

[christopherhesse]

Fixing the ALE implementation may be fine, but it's possible that it makes more sense to just have a wrapper. Specifically the state is not e.unwrapped.clone_state(), but (e.unwrapped. clone_full_state(), e.unwrapped.getRAM(), e.unwrapped.getScreenRGB2()), and it should be restored with the full state as well.

There's an example wrapper I created for retro here: https://github.com/openai/retro/blob/master/retro/examples/determinism.py#L14

As long as you don't use any stateful wrappers, this may be sufficient. If you do, it may need to be changed to collect/restore state from all wrappers.

[JesseFarebro]

Upstream issue: Farama-Foundation/Arcade-Learning-Environment#417