forked from fomorians/td-gammon
-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
41 lines (30 loc) · 1.22 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
import os
import tensorflow as tf
from model import Model
# modifying code to work with newer tensor flow
tf.compat.v1.disable_v2_behavior()
flags = tf.compat.v1.app.flags
FLAGS = flags.FLAGS
flags.DEFINE_boolean('test', False, 'If true, test against a random strategy.')
flags.DEFINE_boolean('play', False, 'If true, play against a trained TD-Gammon strategy.')
flags.DEFINE_boolean('restore', False, 'If true, restore a checkpoint before training.')
model_path = os.environ.get('MODEL_PATH', 'models/')
summary_path = os.environ.get('SUMMARY_PATH', 'summaries/')
checkpoint_path = os.environ.get('CHECKPOINT_PATH', 'checkpoints/')
if not os.path.exists(model_path):
os.makedirs(model_path)
if not os.path.exists(checkpoint_path):
os.makedirs(checkpoint_path)
if not os.path.exists(summary_path):
os.makedirs(summary_path)
if __name__ == '__main__':
graph = tf.compat.v1.Graph()
sess = tf.compat.v1.Session(graph=graph)
with sess.as_default(), graph.as_default():
model = Model(sess, model_path, summary_path, checkpoint_path, restore=FLAGS.restore)
if FLAGS.test:
model.test(episodes=1000)
elif FLAGS.play:
model.play()
else:
model.train()