Save policies for Evolution Strategies (#638)

Save policies for evolution strategies.

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

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()