def get_player(viz=False, train=False):
pl = AtariPlayer(ROM_FILE, frame_skip=ACTION_REPEAT,
image_shape=IMAGE_SIZE[::-1], viz=viz, live_lost_as_eoe=train)
global NUM_ACTIONS
NUM_ACTIONS = pl.get_action_space().num_actions()
if not train:
pl = HistoryFramePlayer(pl, FRAME_HISTORY)
pl = PreventStuckPlayer(pl, 30, 1)
pl = LimitLengthPlayer(pl, 30000)
return pl
parser.add_argument('--algo',
help='algorithm',
choices=['DQN', 'Double', 'Dueling'],
default='Double')
args = parser.parse_args()
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
if args.task != 'train':
assert args.load is not None
ROM_FILE = args.rom
METHOD = args.algo
# set num_actions
pl = AtariPlayer(ROM_FILE, viz=False)
NUM_ACTIONS = pl.get_action_space().num_actions()
del pl
if args.task != 'train':
cfg = PredictConfig(model=Model(),
session_init=get_model_loader(args.load),
input_names=['state'],
output_names=['Qvalue'])
if args.task == 'play':
play_model(cfg, get_player(viz=0.01))
elif args.task == 'eval':
eval_model_multithread(cfg, EVAL_EPISODE, get_player)
else:
config = get_config()
if args.load:
config.session_init = SaverRestore(args.load)