forked from jknthn/unity-tennis-rl
/
run.py
31 lines (24 loc) · 1.11 KB
/
run.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
#!/usr/bin/env python
import argparse
from time import sleep
from unityagents import UnityEnvironment
from maddpg import MADDPGAgent
from config import config
from train import training_loop
from play import play_loop
parser = argparse.ArgumentParser()
parser.add_argument('--train', action='store_true', dest='train', help='Set the train mode')
parser.add_argument('--file_prefix', default=None, help='Set the file for agent to load weights with using prefix')
parser.add_argument('--playthroughs', default=10, type=int, help='Number of playthroughs played in a play mode')
parser.add_argument('--sleep', default=0, type=int, help='Time before environment starts in a play mode [seconds]')
arguments = parser.parse_args()
env = UnityEnvironment(file_name='./Tennis.app', seed=config.general.seed)
brain_name = env.brain_names[0]
agent = MADDPGAgent(config=config, file_prefix=arguments.file_prefix)
if arguments.train:
print('Train mode \n')
training_loop(env, brain_name, agent, config)
else:
print('Play mode \n')
sleep(arguments.sleep)
play_loop(env, brain_name, agent, playthrougs=arguments.playthroughs)