Added model checkpoint to the training result directory

Author: hedrox

ppo_agent.py

@@ -110,7 +110,7 @@ def __init__(self, *, scope,
                  update_ob_stats_every_step=True,
                  int_coeff=None,
                  ext_coeff=None,
-                 restore_model=False):
+                 restore_model_path=None):
 
         self.lr = lr
         self.ext_coeff = ext_coeff
@@ -201,11 +201,11 @@ def __init__(self, *, scope,
         if self.is_log_leader:
             tf_util.display_var_info(allvars)
 
-        model_path = os.path.join(os.getcwd(), 'saved_model')
+        model_path = os.path.join(logger.get_dir(), 'saved_model')
         self.model_path = os.path.join(model_path, 'ppo.ckpt')
 
-        if restore_model:
-            tf_util.load_state(model_path)
+        if restore_model_path:
+            tf_util.load_state(restore_model_path)
         else:
             #self.activate_graph_debugging()
             tf.get_default_session().run(tf.variables_initializer(allvars))

run_atari.py

@@ -74,7 +74,7 @@ def train(*, env_id, num_env, hps, num_timesteps, seed):
         update_ob_stats_every_step=hps.pop('update_ob_stats_every_step'),
         int_coeff=hps.pop('int_coeff'),
         ext_coeff=hps.pop('ext_coeff'),
-        restore_model=hps.pop('restore_model')
+        restore_model_path=hps.pop('restore_model_path')
     )
     agent.start_interaction([venv])
     if hps.pop('update_ob_stats_from_random_agent'):
@@ -139,7 +139,8 @@ def main():
     parser.add_argument('--ext_coeff', type=float, default=2.)
     parser.add_argument('--dynamics_bonus', type=int, default=0)
     parser.add_argument('--save_model', action='store_true')
-    parser.add_argument('--restore_model', action='store_true')
+    parser.add_argument('--restore_model_path', type=str, default='',
+                        help='Path to the saved_model dir containing the ppo checkpoint')
     parser.add_argument('--experiment', type=str, default='ego', choices=['baseline', 'attention', 'ego'])
 
     args = parser.parse_args()
@@ -175,7 +176,7 @@ def main():
         ext_coeff=args.ext_coeff,
         dynamics_bonus = args.dynamics_bonus,
         save_model=args.save_model,
-        restore_model=args.restore_model
+        restore_model_path=args.restore_model_path
     )
 
     tf_util.make_session(make_default=True)