Test script to check whether data from the Atari Grand Challenge Dataset works with the Atari Learning Environment

ale-agcd-test-trajectory.py

#!/usr/bin/env python
# This assumes the ALE python libraries are in the path
# and that the dataset is in the current working directory

from __future__ import print_function

import sys
sys.path.insert(0, 'Arcade-Learning-Environment-0.5.0')

import csv
import pygame

from ale_python_interface import ALEInterface

ale = ALEInterface()

def reload(ale, seed):
    ale.setFloat('repeat_action_probability', 0)
    ale.setInt('random_seed', seed)
    ale.loadROM("montezuma_revenge.bin")
    ale.setInt('random_seed', seed)
    ale.setBool('display_screen', True)
    ale.setBool('sound', True)
    ale.reset_game()

def evaluate():
    with open('atari_v2_release/trajectories/revenge/0.txt') as fp:
        total_reward = 0
        fp.readline()
        fp.readline()
        reader = csv.reader(fp)
        for i, row in enumerate(reader):
            action = row[-1]
            reward = ale.act(action)
            if reward != 0:
                total_reward += reward
                print('Non-zero reward', reward, 'at frame', i)

            if ale.game_over():
                print('Game over at frame', i)
                ale.reset_game()
                break
    return total_reward

seed = 0
total_reward = 0
reload(ale, seed)
total_reward = evaluate()
print('Total reward:', total_reward, 'seed:', seed)