baselines/baselines/trpo_mpi/run_trpo_experiment.py at master · btaba/baselines · GitHub

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
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
#!/usr/bin/env python

# Mujoco must come before other imports. https://openai.slack.com/archives/C1H6P3R7B/p1492828680631850

import os
import sys
import time
import json
import subprocess
from functools import partial
from concurrent.futures import ProcessPoolExecutor
from baselines.bench.monitor import load_results
from baselines.bench.benchmarks import _BENCHMARKS


SEEDS = list(range(1, 100))


def train_mujoco(env_id, num_timesteps, seed, logdir):

    env = os.environ.copy()
    env["PATH"] = "/usr/sbin:/sbin:" + env["PATH"]
    env["OPENAI_LOGDIR"] = logdir
    python_path = sys.executable
    command = '{} -m baselines.trpo_mpi.run_mujoco --env {} --seed {} --num-timesteps {}'.format(
        python_path, env_id, seed, num_timesteps)
    p = subprocess.Popen(command, env=env, shell=True)
    out, err = p.communicate()


def train_atari(env_id, num_timesteps, seed):
    pass


def train(base_log_path, benchmark_name, task):
    results = []
    for trial in range(task['trials']):
        trial_logdir = os.path.join(
            base_log_path,
            '{}_{}_{}'.format(benchmark_name, task['env_id'], trial))
        os.makedirs(trial_logdir)

        if benchmark_name.lower().startswith('mujoco'):
            train_mujoco(
                task['env_id'],
                num_timesteps=task['num_timesteps'],
                seed=SEEDS[trial],
                logdir=trial_logdir)
        else:
            train_atari(
                task['env_id'],
                num_timesteps=task['num_timesteps'],
                seed=SEEDS[trial],
                logdir=trial_logdir)

        res = load_results(trial_logdir)
        res['trial'] = trial
        res['seed'] = SEEDS[trial]

        results.append(res)

    return results


def main():
    import argparse
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument(
        '--benchmark', help='benchmark name', default='Mujoco1M')
    parser.add_argument(
        '--logdir', help='logging directory')
    args = parser.parse_args()

    logdir = args.logdir
    assert logdir is not None

    # get benchmark tasks
    benchmark_name = args.benchmark
    benchmark_dict = dict(
        map(lambda x: (x[1]['name'], x[0]), enumerate(_BENCHMARKS)))
    assert benchmark_name in benchmark_dict
    benchmark_idx = benchmark_dict[benchmark_name]
    benchmark = _BENCHMARKS[benchmark_idx]

    # Make a master log directory
    path = time.strftime("{}_%d-%m-%y-%H-%M-%S_baseline".format(
        benchmark_name))
    base_log_path = os.path.join(os.path.expanduser(logdir), path)
    os.makedirs(base_log_path)

    # train all the benchmark tasks
    with ProcessPoolExecutor() as ex:
        train_func = partial(train, base_log_path, benchmark_name)
        for res in ex.map(
                train_func,
                benchmark['tasks']):

            for r in res:
                # f = open(os.path.join(base_log_path, 'logs.json'), 'a')
                # json.dump(r, f)
                # f.write('\n')
                print(r)


if __name__ == '__main__':
    main()