Rename policy_saver.PolicySaver to (e.g.) policy_saver.MLGOPolicySaver (google#527)

Addresses issue google#309
Just a typo within the commit message 
It is indeed MLGOPolicySaver

Author: Harishankar14

compiler_opt/es/blackbox_learner_test.py

@@ -111,7 +111,7 @@ def setUp(self):
     init_params = policy_utils.get_vectorized_parameters_from_policy(policy)
 
     # save the policy
-    saver = policy_saver.PolicySaver({policy_name: policy})
+    saver = policy_saver.MLGOPolicySaver({policy_name: policy})
     policy_save_path = os.path.join(output_dir.full_path, 'temp_output',
                                     'policy')
     saver.save(policy_save_path)

compiler_opt/es/inlining/inlining_worker.py

@@ -42,7 +42,7 @@ class InliningWorker(worker.Worker):
   def _setup_base_policy(self):
     self._tf_base_temp_dir = tempfile.mkdtemp()
     policy = policy_utils.create_actor_policy()
-    saver = policy_saver.PolicySaver({"policy": policy})
+    saver = policy_saver.MLGOPolicySaver({"policy": policy})
     saver.save(self._tf_base_temp_dir)
     self._tf_base_policy_path = os.path.join(self._tf_base_temp_dir, "policy")
 

compiler_opt/es/policy_utils.py

@@ -121,7 +121,7 @@ def save_policy(policy: 'tf_policy.TFPolicy | HasModelVariables',
     policy_name: The value to name the policy.
   """
   set_vectorized_parameters_for_policy(policy, parameters)
-  saver = policy_saver.PolicySaver({policy_name: policy})
+  saver = policy_saver.MLGOPolicySaver({policy_name: policy})
   saver.save(save_folder)
 
 

compiler_opt/es/policy_utils_test.py

@@ -111,7 +111,8 @@ class VectorTest(absltest.TestCase):
   POLICY_NAME = 'test_policy_name'
 
   def _save_inlining_policy(
-      self) -> tuple[str, actor_policy.ActorPolicy, policy_saver.PolicySaver]:
+      self
+  ) -> tuple[str, actor_policy.ActorPolicy, policy_saver.MLGOPolicySaver]:
     problem_config = registry.get_configuration(
         implementation=inlining.InliningConfig)
     time_step_spec, action_spec = problem_config.get_signature_spec()
@@ -138,7 +139,7 @@ def _save_inlining_policy(
         actor_network=actor_network)
 
     # save the policy
-    saver = policy_saver.PolicySaver({VectorTest.POLICY_NAME: policy})
+    saver = policy_saver.MLGOPolicySaver({VectorTest.POLICY_NAME: policy})
     testing_path = self.create_tempdir()
     policy_save_path = os.path.join(testing_path, 'temp_output', 'policy')
     saver.save(policy_save_path)

compiler_opt/es/regalloc_trace/regalloc_trace_worker.py

@@ -56,7 +56,7 @@ class RegallocTraceWorker(worker.Worker):
   def _setup_base_policy(self):
     self._tf_base_temp_dir = tempfile.mkdtemp()
     policy = policy_utils.create_actor_policy()
-    saver = policy_saver.PolicySaver({"policy": policy})
+    saver = policy_saver.MLGOPolicySaver({"policy": policy})
     saver.save(self._tf_base_temp_dir)
     self._tf_base_policy_path = os.path.join(self._tf_base_temp_dir, "policy")
 

compiler_opt/rl/distributed/ppo_collect_lib.py

@@ -48,7 +48,7 @@ def _get_policy_bytes(agent):
   """Recover the collect_policy bytes from a TF agent"""
   policy_key = 'collect'
   with tempfile.TemporaryDirectory() as tmpdirname:
-    saver = policy_saver.PolicySaver(policy_dict={
+    saver = policy_saver.MLGOPolicySaver(policy_dict={
         policy_key: agent.collect_policy,
     })
     saver.save(tmpdirname)

compiler_opt/rl/distributed/ppo_eval_lib.py

@@ -121,7 +121,7 @@ def sequence_example_iterator_fn(seq_ex: list[str]):
     actions = []
     while True:
       with tempfile.TemporaryDirectory() as tmpdirname:
-        saver = policy_saver.PolicySaver(policy_dict=policy_dict)
+        saver = policy_saver.MLGOPolicySaver(policy_dict=policy_dict)
         saver.save(tmpdirname)
         policy_bytes = policy_saver.Policy.from_filesystem(
             os.path.join(tmpdirname, 'policy'))

compiler_opt/rl/imitation_learning/weighted_bc_trainer.py

@@ -69,7 +69,7 @@ def train():
         time_step_spec=expected_signature,
         action_spec=action_spec)
     policy_dict = {'tf_agents_policy': wrapped_keras_model}
-    saver = policy_saver.PolicySaver(policy_dict=policy_dict)
+    saver = policy_saver.MLGOPolicySaver(policy_dict=policy_dict)
     saver.save(_SAVE_MODEL_DIR.value)
 
 

compiler_opt/rl/policy_saver.py

@@ -142,11 +142,11 @@ def from_filesystem(location: str):
     return Policy(output_spec=output_spec, policy=policy)
 
 
-class PolicySaver:
+class MLGOPolicySaver:
   """Object that saves policy and model config file required by inference.
 
   ```python
-  policy_saver = PolicySaver(policy_dict, config)
+  policy_saver = MLGOPolicySaver(policy_dict, config)
   policy_saver.save(root_dir)
   ```
   """

compiler_opt/rl/policy_saver_test.py

@@ -28,7 +28,7 @@
 from compiler_opt.testing import model_test_utils
 
 
-class PolicySaverTest(tf.test.TestCase):
+class MLGOPolicySaverTest(tf.test.TestCase):
 
   def setUp(self):
     super().setUp()
@@ -54,7 +54,7 @@ def test_save_policy(self):
         'saved_policy': test_agent.policy,
         'saved_collect_policy': test_agent.collect_policy
     }
-    test_policy_saver = policy_saver.PolicySaver(policy_dict=policy_dict)
+    test_policy_saver = policy_saver.MLGOPolicySaver(policy_dict=policy_dict)
 
     root_dir = self.get_temp_dir()
     test_policy_saver.save(root_dir)