Separate train and eval environments for ddpg and sac.

PiperOrigin-RevId: 259407411 Change-Id: Ia9d3215d00e908613390b4cb0230df32e8250bc2
notadamking · Jul 22, 2019 · 11a875c · 11a875c
1 parent 80a2172
commit 11a875c
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 4 deletions.
diff --git a/tf_agents/agents/ddpg/examples/v1/train_eval.py b/tf_agents/agents/ddpg/examples/v1/train_eval.py
@@ -67,6 +67,7 @@
 def train_eval(
     root_dir,
     env_name='HalfCheetah-v2',
+    eval_env_name=None,
     env_load_fn=suite_mujoco.load,
     num_iterations=2000000,
     actor_fc_layers=(400, 300),
@@ -132,7 +133,8 @@ def train_eval(
               [lambda: env_load_fn(env_name)] * num_parallel_environments))
     else:
       tf_env = tf_py_environment.TFPyEnvironment(env_load_fn(env_name))
-    eval_py_env = env_load_fn(env_name)
+    eval_env_name = eval_env_name or env_name
+    eval_py_env = env_load_fn(eval_env_name)
 
     actor_net = actor_network.ActorNetwork(
         tf_env.time_step_spec().observation,

diff --git a/tf_agents/agents/ddpg/examples/v2/train_eval.py b/tf_agents/agents/ddpg/examples/v2/train_eval.py
@@ -69,6 +69,7 @@
 def train_eval(
     root_dir,
     env_name='HalfCheetah-v2',
+    eval_env_name=None,
     env_load_fn=suite_mujoco.load,
     num_iterations=2000000,
     actor_fc_layers=(400, 300),
@@ -132,7 +133,8 @@ def train_eval(
               [lambda: env_load_fn(env_name)] * num_parallel_environments))
     else:
       tf_env = tf_py_environment.TFPyEnvironment(env_load_fn(env_name))
-    eval_tf_env = tf_py_environment.TFPyEnvironment(env_load_fn(env_name))
+    eval_env_name = eval_env_name or env_name
+    eval_tf_env = tf_py_environment.TFPyEnvironment(env_load_fn(eval_env_name))
 
     actor_net = actor_network.ActorNetwork(
         tf_env.time_step_spec().observation,

diff --git a/tf_agents/agents/sac/examples/v1/train_eval.py b/tf_agents/agents/sac/examples/v1/train_eval.py
@@ -85,6 +85,7 @@ def normal_projection_net(action_spec,
 def train_eval(
     root_dir,
     env_name='HalfCheetah-v2',
+    eval_env_name=None,
     env_load_fn=suite_mujoco.load,
     num_iterations=1000000,
     actor_fc_layers=(256, 256),
@@ -144,7 +145,8 @@ def train_eval(
       lambda: tf.math.equal(global_step % summary_interval, 0)):
     # Create the environment.
     tf_env = tf_py_environment.TFPyEnvironment(env_load_fn(env_name))
-    eval_py_env = env_load_fn(env_name)
+    eval_env_name = eval_env_name or env_name
+    eval_py_env = env_load_fn(eval_env_name)
 
     # Get the data specs from the environment
     time_step_spec = tf_env.time_step_spec()

diff --git a/tf_agents/agents/sac/examples/v2/train_eval.py b/tf_agents/agents/sac/examples/v2/train_eval.py
@@ -85,6 +85,7 @@ def normal_projection_net(action_spec,
 def train_eval(
     root_dir,
     env_name='HalfCheetah-v2',
+    eval_env_name=None,
     env_load_fn=suite_mujoco.load,
     num_iterations=1000000,
     actor_fc_layers=(256, 256),
@@ -142,7 +143,8 @@ def train_eval(
   with tf.compat.v2.summary.record_if(
       lambda: tf.math.equal(global_step % summary_interval, 0)):
     tf_env = tf_py_environment.TFPyEnvironment(env_load_fn(env_name))
-    eval_tf_env = tf_py_environment.TFPyEnvironment(env_load_fn(env_name))
+    eval_env_name = eval_env_name or env_name
+    eval_tf_env = tf_py_environment.TFPyEnvironment(env_load_fn(eval_env_name))
 
     time_step_spec = tf_env.time_step_spec()
     observation_spec = time_step_spec.observation