introduce graph_manager.phase_context; make sure that calls to graph_manager.train automatically set training phase

rl_coach/graph_managers/graph_manager.py

@@ -20,6 +20,7 @@ import time
 from collections import OrderedDict
 from distutils.dir_util import copy_tree, remove_tree
 from typing import List, Tuple
+import contextlib
 
 from rl_coach.base_parameters import iterable_to_items, TaskParameters, DistributedTaskParameters, \
     VisualizationParameters, \
@@ -285,6 +286,13 @@ class GraphManager(object):
         for environment in self.environments:
             environment.phase = val
 
+    @contextlib.contextmanager
+    def phase_context(self, phase):
+        old_phase = self.phase
+        self.phase = phase
+        yield
+        self.phase = old_phase
+
     def set_session(self, sess) -> None:
         """
         Set the deep learning framework session for all the modules in the graph
@@ -301,7 +309,7 @@ class GraphManager(object):
         self.verify_graph_was_created()
 
         if steps.num_steps > 0:
-            self.phase = RunPhase.HEATUP
+            with self.phase_context(RunPhase.HEATUP):
                 screen.log_title("{}: Starting heatup".format(self.name))
                 self.heatup_start_time = time.time()
 
@@ -313,9 +321,6 @@ class GraphManager(object):
                 while self.total_steps_counters[self.phase][steps.__class__] < count_end:
                     self.act(steps, continue_until_game_over=True, return_on_game_over=True)
 
-            # training phase
-            self.phase = RunPhase.UNDEFINED
-
     def handle_episode_ended(self) -> None:
         """
         End an episode and reset all the episodic parameters
@@ -333,9 +338,9 @@ class GraphManager(object):
         """
         self.verify_graph_was_created()
 
+        with self.phase_context(RunPhase.TRAIN):
             [manager.train() for manager in self.level_managers]
 
-
     def reset_internal_state(self, force_environment_reset=False) -> None:
         """
         Reset an episode for all the levels
@@ -417,7 +422,7 @@ class GraphManager(object):
 
         # perform several steps of training interleaved with acting
         if steps.num_steps > 0:
-            self.phase = RunPhase.TRAIN
+            with self.phase_context(RunPhase.TRAIN):
                 self.reset_internal_state(force_environment_reset=True)
                 #TODO - the below while loop should end with full episodes, so to avoid situations where we have partial
                 #  episodes in memory
@@ -428,7 +433,6 @@ class GraphManager(object):
                     self.act(EnvironmentSteps(1))
                     self.train()
                     self.occasionally_save_checkpoint()
-            self.phase = RunPhase.UNDEFINED
 
     def sync_graph(self) -> None:
         """

rl_coach/rollout_worker.py

@@ -89,8 +89,7 @@ def rollout_worker(graph_manager, checkpoint_dir, data_store, num_workers, polic
     task_parameters.__dict__['checkpoint_restore_dir'] = checkpoint_dir
     time.sleep(30)
     graph_manager.create_graph(task_parameters)
-    graph_manager.phase = RunPhase.TRAIN
-
+    with graph_manager.phase_context(RunPhase.TRAIN):
         error_compensation = 100
 
         last_checkpoint = 0
@@ -118,8 +117,6 @@ def rollout_worker(graph_manager, checkpoint_dir, data_store, num_workers, polic
 
             last_checkpoint = new_checkpoint
 
-    graph_manager.phase = RunPhase.UNDEFINED
-
 
 def main():
     parser = argparse.ArgumentParser()

rl_coach/tests/graph_managers/test_graph_manager.py

@@ -0,0 +1,47 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
+
+import pytest
+from rl_coach.graph_managers.graph_manager import GraphManager, ScheduleParameters
+from rl_coach.base_parameters import VisualizationParameters
+from rl_coach.core_types import RunPhase
+
+
+@pytest.mark.unit_test
+def test_phase_context():
+    graph_manager = GraphManager(name='', schedule_params=ScheduleParameters(), vis_params=VisualizationParameters())
+
+    assert graph_manager.phase == RunPhase.UNDEFINED
+    with graph_manager.phase_context(RunPhase.TRAIN):
+        assert graph_manager.phase == RunPhase.TRAIN
+    assert graph_manager.phase == RunPhase.UNDEFINED
+
+
+@pytest.mark.unit_test
+def test_phase_context_nested():
+    graph_manager = GraphManager(name='', schedule_params=ScheduleParameters(), vis_params=VisualizationParameters())
+
+    assert graph_manager.phase == RunPhase.UNDEFINED
+    with graph_manager.phase_context(RunPhase.TRAIN):
+        assert graph_manager.phase == RunPhase.TRAIN
+        with graph_manager.phase_context(RunPhase.TEST):
+            assert graph_manager.phase == RunPhase.TEST
+        assert graph_manager.phase == RunPhase.TRAIN
+    assert graph_manager.phase == RunPhase.UNDEFINED
+
+
+@pytest.mark.unit_test
+def test_phase_context_double_nested():
+    graph_manager = GraphManager(name='', schedule_params=ScheduleParameters(), vis_params=VisualizationParameters())
+
+    assert graph_manager.phase == RunPhase.UNDEFINED
+    with graph_manager.phase_context(RunPhase.TRAIN):
+        assert graph_manager.phase == RunPhase.TRAIN
+        with graph_manager.phase_context(RunPhase.TEST):
+            assert graph_manager.phase == RunPhase.TEST
+            with graph_manager.phase_context(RunPhase.UNDEFINED):
+                assert graph_manager.phase == RunPhase.UNDEFINED
+            assert graph_manager.phase == RunPhase.TEST
+        assert graph_manager.phase == RunPhase.TRAIN
+    assert graph_manager.phase == RunPhase.UNDEFINED