[BC] ExplorationWithPolicy

compiler_opt/rl/env.py

@@ -290,18 +290,23 @@ def _get_scores() -> dict[str, float]:
 def _get_clang_generator(
     clang_path: str,
     task_type: Type[MLGOTask],
+    interactive_only: bool = False,
 ) -> Generator[Optional[Tuple[ClangProcess, InteractiveClang]],
                Optional[corpus.LoadedModuleSpec], None]:
-  """Returns a generator for creating InteractiveClang objects.
-
-  TODO: fix this docstring
+  """Returns a tuple of generators for creating InteractiveClang objects.
 
   Args:
     clang_path: Path to the clang binary to use within InteractiveClang.
     task_type: Type of the MLGO task to use.
+    interactive_only: If set to true the returned tuple of generators is
+      iclang, iclang instead of iclang, clang
 
   Returns:
-    The generator for InteractiveClang objects.
+    A generator of tuples. Each element of the tuple is created with
+    clang_session. First argument of the tuple is always an interactive
+    clang session. The second argumnet is a default clang session if
+    interactive_only is False and otherwise the exact same interactive
+    clang session object as the first argument.
   """
   while True:
     # The following line should be type-hinted as follows:
@@ -311,9 +316,12 @@ def _get_clang_generator(
     module = yield
     with clang_session(
         clang_path, module, task_type, interactive=True) as iclang:
-      with clang_session(
-          clang_path, module, task_type, interactive=False) as clang:
-        yield iclang, clang
+      if interactive_only:
+        yield iclang, iclang
+      else:
+        with clang_session(
+            clang_path, module, task_type, interactive=False) as clang:
+          yield iclang, clang
 
 
 class MLGOEnvironmentBase:
@@ -332,8 +340,10 @@ def __init__(
       task_type: Type[MLGOTask],
       obs_spec,
       action_spec,
+      interactive_only: bool = False,
   ):
-    self._clang_generator = _get_clang_generator(clang_path, task_type)
+    self._clang_generator = _get_clang_generator(
+        clang_path, task_type, interactive_only=interactive_only)
     self._obs_spec = obs_spec
     self._action_spec = action_spec
 

compiler_opt/rl/env_test.py

@@ -214,6 +214,33 @@ def test_env(self, mock_popen):
 
       step = test_env.step(np.array([1], dtype=np.int64))
       self.assertEqual(step.step_type, env.StepType.LAST)
+      self.assertNotEqual(test_env._iclang, test_env._clang)  # pylint: disable=protected-access
+
+  @mock.patch('subprocess.Popen')
+  def test_env_interactive_only(self, mock_popen):
+    mock_popen.side_effect = mock_interactive_clang
+
+    test_env = env.MLGOEnvironmentBase(
+        clang_path=_CLANG_PATH,
+        task_type=MockTask,
+        obs_spec={},
+        action_spec={},
+        interactive_only=True,
+    )
+
+    for env_itr in range(3):
+      del env_itr
+      step = test_env.reset(_MOCK_MODULE)
+      self.assertEqual(step.step_type, env.StepType.FIRST)
+
+      for step_itr in range(_NUM_STEPS - 1):
+        del step_itr
+        step = test_env.step(np.array([1], dtype=np.int64))
+        self.assertEqual(step.step_type, env.StepType.MID)
+
+      step = test_env.step(np.array([1], dtype=np.int64))
+      self.assertEqual(step.step_type, env.StepType.LAST)
+      self.assertEqual(test_env._iclang, test_env._clang)  # pylint: disable=protected-access
 
 
 if __name__ == '__main__':

compiler_opt/rl/generate_bc_trajectories.py

@@ -0,0 +1,92 @@
+# coding=utf-8
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Module for running compilation and collect data for behavior cloning."""
+
+from typing import Callable, Dict, List, Optional
+
+import numpy as np
+import tensorflow as tf
+from tf_agents.trajectories import policy_step
+from tf_agents.trajectories import time_step
+
+
+class ExplorationWithPolicy:
+  """Policy which selects states for exploration.
+
+  Exploration is fascilitated in the following way. First the policy plays
+  all actions from the replay_prefix. At the following state the policy computes
+  a gap which is difference between the most likely action and the second most
+  likely action according to the randomized exploration policy (distr).
+  If the current gap is smaller than previously maintained gap, the gap is
+  updated and the exploration state is set to the current state.
+  The trajectory is completed by following following the policy from the
+  constructor.
+
+  Attributes:
+    replay_prefix: a replay buffer of actions
+    policy: policy to follow after exhausting the replay buffer
+    explore_policy: randomized policy which is used to compute the gap
+    curr_step: current step of the trajectory
+    explore_step: current candidate for exploration step
+    gap: current difference at explore step between probability of most likely
+      action according to explore_policy and second most likely action
+    explore_on_features: dict of feature names and functions which specify
+      when to explore on the respective feature
+  """
+
+  def __init__(
+      self,
+      replay_prefix: List[np.ndarray],
+      policy: Callable[[time_step.TimeStep], np.ndarray],
+      explore_policy: Callable[[time_step.TimeStep], policy_step.PolicyStep],
+      explore_on_features: Optional[Dict[str, Callable[[tf.Tensor],
+                                                       bool]]] = None,
+  ):
+    self.replay_prefix = replay_prefix
+    self.policy = policy
+    self.explore_policy = explore_policy
+    self.curr_step = 0
+    self.explore_step = 0
+    self.gap = np.inf
+    self.explore_on_features = explore_on_features
+    self._stop_exploration = False
+
+  def advice(self, state: time_step.TimeStep) -> np.ndarray:
+    """Action function for the policy.
+
+    Args:
+      state: current state in the trajectory
+
+    Returns:
+      policy_deca: action to take at the current state.
+
+    """
+    if self.curr_step < len(self.replay_prefix):
+      self.curr_step += 1
+      return np.array(self.replay_prefix[self.curr_step - 1])
+    policy_deca = self.policy(state)
+    distr = tf.nn.softmax(self.explore_policy(state).action.logits).numpy()[0]
+    if not self._stop_exploration and distr.shape[0] > 1 and self.gap > np.abs(
+        distr[0] - distr[1]):
+      self.gap = np.abs(distr[0] - distr[1])
+      self.explore_step = self.curr_step
+    if not self._stop_exploration and self.explore_on_features is not None:
+      for feature_name, explore_on_feature in self.explore_on_features.items():
+        if explore_on_feature(state.observation[feature_name]):
+          self.explore_step = self.curr_step
+          self._stop_exploration = True
+          break
+    self.curr_step += 1
+    return policy_deca

compiler_opt/rl/generate_bc_trajectories_test.py

@@ -0,0 +1,154 @@
+# coding=utf-8
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for compiler_opt.rl.generate_bc_trajectories."""
+
+from typing import List
+
+import numpy as np
+import tensorflow as tf
+import tensorflow_probability as tfp
+from tf_agents.trajectories import policy_step
+from tf_agents.trajectories import time_step
+
+from compiler_opt.rl import generate_bc_trajectories
+
+_eps = 1e-5
+
+
+def _get_state_list() -> List[time_step.TimeStep]:
+
+  state_0 = time_step.TimeStep(
+      discount=tf.constant(np.array([0.]), dtype=tf.float32),
+      observation={
+          'feature_1': tf.constant(np.array([0]), dtype=tf.int64),
+          'feature_2': tf.constant(np.array([50]), dtype=tf.int64),
+          'feature_3': tf.constant(np.array([0]), dtype=tf.int64),
+      },
+      reward=tf.constant(np.array([0]), dtype=tf.float32),
+      step_type=tf.constant(np.array([0]), dtype=tf.int32))
+  state_1 = time_step.TimeStep(
+      discount=tf.constant(np.array([0.]), dtype=tf.float32),
+      observation={
+          'feature_1': tf.constant(np.array([1]), dtype=tf.int64),
+          'feature_2': tf.constant(np.array([25]), dtype=tf.int64),
+          'feature_3': tf.constant(np.array([0]), dtype=tf.int64),
+      },
+      reward=tf.constant(np.array([0]), dtype=tf.float32),
+      step_type=tf.constant(np.array([0]), dtype=tf.int32))
+  state_2 = time_step.TimeStep(
+      discount=tf.constant(np.array([0.]), dtype=tf.float32),
+      observation={
+          'feature_1': tf.constant(np.array([0]), dtype=tf.int64),
+          'feature_2': tf.constant(np.array([25]), dtype=tf.int64),
+          'feature_3': tf.constant(np.array([1]), dtype=tf.int64),
+      },
+      reward=tf.constant(np.array([0]), dtype=tf.float32),
+      step_type=tf.constant(np.array([0]), dtype=tf.int32))
+  state_3 = time_step.TimeStep(
+      discount=tf.constant(np.array([0.]), dtype=tf.float32),
+      observation={
+          'feature_1': tf.constant(np.array([0]), dtype=tf.int64),
+          'feature_2': tf.constant(np.array([25]), dtype=tf.int64),
+          'feature_3': tf.constant(np.array([0]), dtype=tf.int64),
+      },
+      reward=tf.constant(np.array([0]), dtype=tf.float32),
+      step_type=tf.constant(np.array([0]), dtype=tf.int32))
+
+  return [state_0, state_1, state_2, state_3]
+
+
+def _policy(state: time_step.TimeStep) -> np.ndarray:
+  feature_sum = np.array([0])
+  for feature in state.observation.values():
+    feature_sum += feature.numpy()
+  return np.mod(feature_sum, 5)
+
+
+def _explore_policy(state: time_step.TimeStep) -> policy_step.PolicyStep:
+  probs = [
+      0.5 * float(state.observation['feature_3'].numpy()),
+      1 - 0.5 * float(state.observation['feature_3'].numpy())
+  ]
+  logits = [[0.0, tf.math.log(probs[1] / (1.0 - probs[1] + _eps))]]
+  return policy_step.PolicyStep(
+      action=tfp.distributions.Categorical(logits=logits))
+
+
+class ExplorationWithPolicyTest(tf.test.TestCase):
+
+  def test_explore_policy(self):
+    prob = 1.
+    state = _get_state_list()[3]
+    logits = [[0.0, tf.math.log(prob / (1.0 - prob + _eps))]]
+    action = tfp.distributions.Categorical(logits=logits)
+    self.assertAllClose(action.logits, _explore_policy(state).action.logits)
+
+  def test_explore_with_gap(self):
+    explore_with_policy = generate_bc_trajectories.ExplorationWithPolicy(
+        replay_prefix=[np.array([1])],
+        policy=_policy,
+        explore_policy=_explore_policy,
+    )
+    for state in _get_state_list():
+      _ = explore_with_policy.advice(state)[0]
+
+    self.assertAllClose(0, explore_with_policy.gap, atol=2 * _eps)
+    self.assertEqual(2, explore_with_policy.explore_step)
+
+    explore_with_policy = generate_bc_trajectories.ExplorationWithPolicy(
+        replay_prefix=[np.array([1]),
+                       np.array([1]),
+                       np.array([1])],
+        policy=_policy,
+        explore_policy=_explore_policy,
+    )
+    for state in _get_state_list():
+      _ = explore_with_policy.advice(state)[0]
+
+    self.assertAllClose(1, explore_with_policy.gap, atol=2 * _eps)
+    self.assertEqual(3, explore_with_policy.explore_step)
+
+  def test_explore_with_feature(self):
+
+    def explore_on_feature_1_val(feature_val):
+      return feature_val.numpy()[0] > 0
+
+    def explore_on_feature_2_val(feature_val):
+      return feature_val.numpy()[0] > 25
+
+    explore_on_features = {
+        'feature_1': explore_on_feature_1_val,
+        'feature_2': explore_on_feature_2_val
+    }
+
+    explore_with_policy = generate_bc_trajectories.ExplorationWithPolicy(
+        replay_prefix=[],
+        policy=_policy,
+        explore_policy=_explore_policy,
+        explore_on_features=explore_on_features)
+    for state in _get_state_list():
+      _ = explore_with_policy.advice(state)[0]
+    self.assertEqual(0, explore_with_policy.explore_step)
+
+    explore_with_policy = generate_bc_trajectories.ExplorationWithPolicy(
+        replay_prefix=[np.array([1])],
+        policy=_policy,
+        explore_policy=_explore_policy,
+        explore_on_features=explore_on_features,
+    )
+
+    for state in _get_state_list():
+      _ = explore_with_policy.advice(state)[0]
+    self.assertEqual(1, explore_with_policy.explore_step)