-
Notifications
You must be signed in to change notification settings - Fork 1
/
run_simple_ctrl.py
executable file
·45 lines (37 loc) · 1.52 KB
/
run_simple_ctrl.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
42
43
44
45
#!/usr/bin/env python3
import os
import sys
sys.path.append(
os.path.abspath(
os.path.join(
os.path.abspath(os.path.join(os.getcwd(),os.pardir)), os.pardir)))
import tensorflow as tf
from baselines.common.cmd_util import gym_ctrl_arg_parser, make_gym_control_multi_env, make_gym_control_env
from baselines import logger
from baselines.common.cmd_util import make_mujoco_env, mujoco_arg_parser
from baselines.acktr.acktr_cont import learn
from baselines.acktr.policies import GaussianMlpPolicy
from baselines.acktr.value_functions import NeuralNetValueFunction
def train(env_id, num_timesteps, seed):
env = make_gym_control_env(env_id, seed)
with tf.Session(config=tf.ConfigProto()):
ob_dim = env.observation_space.shape[0]
ac_dim = env.action_space.shape[0]
with tf.variable_scope("vf"):
vf = NeuralNetValueFunction(ob_dim, ac_dim)
with tf.variable_scope("pi"):
policy = GaussianMlpPolicy(ob_dim, ac_dim)
learn(env, policy=policy, vf=vf,
gamma=0.99, lam=0.97, timesteps_per_batch=2500,
desired_kl=0.002,
num_timesteps=num_timesteps, animate=False)
env.close()
def main():
args = gym_ctrl_arg_parser().parse_args()
logger.configure(format_strs=["stdout",'log', 'csv'], log_suffix = "ACKTR-"+args.env)
import random
seed = 2
logger.log("Algorithm: ACKTR-"+args.env+"_seed_"+str(seed))
train(args.env, num_timesteps=args.num_timesteps, seed=seed)
if __name__ == '__main__':
main()