[rllib] Refactor save() / restore() code of agents and avoid O(n_workers) save size (#2982)

python/ray/tune/result.py

@@ -16,6 +16,12 @@ NODE_IP = "node_ip"
 # (Auto-filled) The pid of the training process.
 PID = "pid"
 
+# Number of timesteps in this iteration.
+EPISODES_THIS_ITER = "episodes_this_iter"
+
+# (Optional/Auto-filled) Accumulated time in seconds for this experiment.
+EPISODES_TOTAL = "episodes_total"
+
 # Number of timesteps in this iteration.
 TIMESTEPS_THIS_ITER = "timesteps_this_iter"
 

python/ray/tune/trainable.py

@@ -17,7 +17,8 @@ import uuid
 import ray
 from ray.tune.logger import UnifiedLogger
 from ray.tune.result import (DEFAULT_RESULTS_DIR, TIME_THIS_ITER_S,
-                             TIMESTEPS_THIS_ITER, DONE, TIMESTEPS_TOTAL)
+                             TIMESTEPS_THIS_ITER, DONE, TIMESTEPS_TOTAL,
+                             EPISODES_THIS_ITER, EPISODES_TOTAL)
 from ray.tune.trial import Resources
 
 logger = logging.getLogger(__name__)
@@ -77,6 +78,7 @@ class Trainable(object):
         self._iteration = 0
         self._time_total = 0.0
         self._timesteps_total = None
+        self._episodes_total = None
         self._time_since_restore = 0.0
         self._timesteps_since_restore = 0
         self._iterations_since_restore = 0
@@ -162,8 +164,15 @@ class Trainable(object):
             self._timesteps_total += result[TIMESTEPS_THIS_ITER]
             self._timesteps_since_restore += result[TIMESTEPS_THIS_ITER]
 
+        # self._timesteps_total should only be tracked if increments provided
+        if result.get(EPISODES_THIS_ITER):
+            if self._episodes_total is None:
+                self._episodes_total = 0
+            self._episodes_total += result[EPISODES_THIS_ITER]
+
         # self._timesteps_total should not override user-provided total
         result.setdefault(TIMESTEPS_TOTAL, self._timesteps_total)
+        result.setdefault(EPISODES_TOTAL, self._episodes_total)
 
         # Provides auto-filled neg_mean_loss for avoiding regressions
         if result.get("mean_loss"):
@@ -205,7 +214,7 @@ class Trainable(object):
         checkpoint_path = self._save(checkpoint_dir or self.logdir)
         pickle.dump([
             self._experiment_id, self._iteration, self._timesteps_total,
-            self._time_total
+            self._time_total, self._episodes_total
         ], open(checkpoint_path + ".tune_metadata", "wb"))
         return checkpoint_path
 
@@ -256,6 +265,7 @@ class Trainable(object):
         self._iteration = metadata[1]
         self._timesteps_total = metadata[2]
         self._time_total = metadata[3]
+        self._episodes_total = metadata[4]
         self._restored = True
 
     def restore_from_object(self, obj):