diff --git a/examples/evolution_strategies/evolution_strategies.py b/examples/evolution_strategies/evolution_strategies.py index 22935123e..3d5296988 100644 --- a/examples/evolution_strategies/evolution_strategies.py +++ b/examples/evolution_strategies/evolution_strategies.py @@ -9,6 +9,7 @@ import argparse from collections import namedtuple import gym import numpy as np +import os import ray import time @@ -195,6 +196,8 @@ if __name__ == "__main__": timesteps_so_far = 0 tstart = time.time() + iteration = 0 + while True: step_tstart = time.time() theta = policy.get_trainable_flat() @@ -280,3 +283,12 @@ if __name__ == "__main__": tlogger.record_tabular("TimeElapsedThisIter", step_tend - step_tstart) tlogger.record_tabular("TimeElapsed", step_tend - tstart) tlogger.dump_tabular() + + if config.snapshot_freq != 0 and iteration % config.snapshot_freq == 0: + filename = os.path.join("/tmp", + "snapshot_iter{:05d}.h5".format(iteration)) + assert not os.path.exists(filename) + policy.save(filename) + tlogger.log("Saved snapshot {}".format(filename)) + + iteration += 1 diff --git a/examples/evolution_strategies/viz.py b/examples/evolution_strategies/viz.py new file mode 100644 index 000000000..7c86fc281 --- /dev/null +++ b/examples/evolution_strategies/viz.py @@ -0,0 +1,42 @@ +# Code in this file is copied and adapted from +# https://github.com/openai/evolution-strategies-starter. + +import click + + +@click.command() +@click.argument("env_id") +@click.argument("policy_file") +@click.option("--record", is_flag=True) +@click.option("--stochastic", is_flag=True) +@click.option("--extra_kwargs") +def main(env_id, policy_file, record, stochastic, extra_kwargs): + import gym + from gym import wrappers + import tensorflow as tf + from policies import MujocoPolicy + import numpy as np + + env = gym.make(env_id) + if record: + import uuid + env = wrappers.Monitor(env, "/tmp/" + str(uuid.uuid4()), force=True) + + if extra_kwargs: + import json + extra_kwargs = json.loads(extra_kwargs) + + with tf.Session(): + pi = MujocoPolicy.Load(policy_file, extra_kwargs=extra_kwargs) + while True: + rews, t = pi.rollout(env, render=True, + random_stream=np.random if stochastic else None) + print("return={:.4f} len={}".format(rews.sum(), t)) + + if record: + env.close() + return + + +if __name__ == "__main__": + main()