Repository: facebookresearch/CompilerGym

Issue:
--benchmark flag not working as expected
## 🐛 Bug

--benchmark flag not working correctly.  Results in error: `AssertionError: Must call reset() before step()`

## To Reproduce

Steps to reproduce the behavior:

Pass in a `--benchmark` and then run any action

```
(compiler_gym) bwasti@devfair0187:~/compiler_gym$ COMPILER_GYM_DEBUG=3 python -m compiler_gym.bin.manual_env --local_service_binary CompilerGym/examples/example_compiler_gym_service/service_py/example_service.py --benchmark bar
I0322 12:18:25.167602 140345274726208 example_service.py:215] Starting service on 42673 with working dir /dev/shm/compiler_gym/service-2021-03-22T12:18:24.112461-60ea
I0322 12:18:25.423372 140342182409984 example_service.py:146] GetSpaces()
Initialized environment in 1.314s
I0322 12:18:25.425686 140342182409984 example_service.py:157] GetBenchmarks()
I0322 12:18:25.426420 140342182409984 example_service.py:157] GetBenchmarks()
Welcome to the CompilerGym Shell!
---------------------------------
Type help or ? for more information.
The 'tutorial' command will give a step by step guide.

compilergym:bar> action -
Action b
Traceback (most recent call last):
  File "/private/home/bwasti/.conda/envs/compiler_gym/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/private/home/bwasti/.conda/envs/compiler_gym/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/private/home/bwasti/.conda/envs/compiler_gym/lib/python3.8/site-packages/compiler_gym-0.1.3-py3.8.egg/compiler_gym/bin/manual_env.py", line 909, in <module>
    main(sys.argv)
  File "/private/home/bwasti/.conda/envs/compiler_gym/lib/python3.8/site-packages/compiler_gym-0.1.3-py3.8.egg/compiler_gym/bin/manual_env.py", line 904, in main
    shell.cmdloop()
  File "/private/home/bwasti/.conda/envs/compiler_gym/lib/python3.8/cmd.py", line 138, in cmdloop
    stop = self.onecmd(line)
  File "/private/home/bwasti/.conda/envs/compiler_gym/lib/python3.8/cmd.py", line 217, in onecmd
    return func(arg)
  File "/private/home/bwasti/.conda/envs/compiler_gym/lib/python3.8/site-packages/compiler_gym-0.1.3-py3.8.egg/compiler_gym/bin/manual_env.py", line 495, in do_action
    observation, reward, done, info = self.env.step(index)
  File "/private/home/bwasti/.conda/envs/compiler_gym/lib/python3.8/site-packages/compiler_gym-0.1.3-py3.8.egg/compiler_gym/envs/compiler_env.py", line 761, in step
    assert self.in_episode, "Must call reset() before step()"
AssertionError: Must call reset() before step()
```



## Expected behavior

Same behavior as `set_benchmark`:

```
(compiler_gym) bwasti@devfair0187:~/compiler_gym$ COMPILER_GYM_DEBUG=3 python -m compiler_gym.bin.manual_env --local_service_binary CompilerGym/examples/example_compiler_gym_service/service_py/example_service.py
I0322 12:17:23.826539 140646749136704 example_service.py:215] Starting service on 41035 with working dir /dev/shm/compiler_gym/service-2021-03-22T12:17:22.682603-4f26
I0322 12:17:23.992188 140643586266880 example_service.py:146] GetSpaces()
Initialized environment in 1.311s
I0322 12:17:23.994066 140643586266880 example_service.py:157] GetBenchmarks()
I0322 12:17:23.994649 140643586266880 example_service.py:157] GetBenchmarks()
Welcome to the CompilerGym Shell!
---------------------------------
Type help or ? for more information.
The 'tutorial' command will give a step by step guide.

compilergym:NO-BENCHMARK> list_benchmarks
bar, foo

compilergym:NO-BENCHMARK> set_benchmark bar
I0322 12:17:38.067934 140643586266880 example_service.py:164] StartSession(benchmark=bar)
Reset bar environment in 1.4ms

compilergym:bar> action -
Action c
I0322 12:17:40.889100 140643586266880 example_service.py:194] Step()
I0322 12:17:40.889228 140643586266880 example_service.py:101] Apply action 2
Actions c in 4.1ms with reward 0.

compilergym:bar>
```

## Environment

Please fill in this checklist:

- CompilerGym: 2623458424faa553df43e63f7b00bf5e45aade3f (Mon Mar 22 11:13:18 2021)
- How you installed CompilerGym (conda, pip, source): source
- OS: linux
- Python version: Python 3.8.8
- Build command you used (if compiling from source): make install; python setup.py install
- GCC/clang version (if compiling from source): gcc (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
- Bazel version (if compiling from source): not sure

[llvm] Add an InstCount observation space
## 🚀 Feature

Expose the raw instruction counts of [InstCount](https://llvm.org/doxygen/InstCount_8cpp_source.html) as an LLVM observation space.


Discussion:

Resolving patch:
diff --git a/examples/tabular_q.py b/examples/tabular_q.py
--- a/examples/tabular_q.py
+++ b/examples/tabular_q.py
@@ -47,7 +47,7 @@
     "Indices of Alphaphase features that are used to construct a state",
 )
 flags.DEFINE_float("learning_rate", 0.1, "learning rate of the q-learning.")
-flags.DEFINE_integer("episodes", 5000, "number of episodes used to learn.")
+flags.DEFINE_integer("episodes", 2000, "number of episodes used to learn.")
 flags.DEFINE_integer(
     "log_every", 50, "number of episode interval where progress is reported."
 )
@@ -86,7 +86,6 @@ def make_q_table_key(autophase_feature, action, step):
 
     Finally, we add the action index to the key.
     """
-
     return StateActionTuple(
         *autophase_feature[FLAGS.features_indices], step, FLAGS.actions.index(action)
     )
@@ -118,6 +117,8 @@ def rollout(qtable, env, printout=False):
         action_seq.append(a)
         observation, reward, done, info = env.step(env.action_space.flags.index(a))
         rewards.append(reward)
+        if done:
+            break
     if printout:
         print(
             "Resulting sequence: ", ",".join(action_seq), f"total reward {sum(rewards)}"
@@ -133,17 +134,19 @@ def train(q_table, env):
     # policy improvement happens directly after one another.
     for i in range(1, FLAGS.episodes + 1):
         current_length = 0
-        obs = env.reset()
+        observation = env.reset()
         while current_length < FLAGS.episode_length:
             # Run epsilon greedy policy to allow exploration.
-            a = select_action(q_table, obs, current_length, FLAGS.epsilon)
-            hashed = make_q_table_key(obs, a, current_length)
+            a = select_action(q_table, observation, current_length, FLAGS.epsilon)
+            hashed = make_q_table_key(observation, a, current_length)
             if hashed not in q_table:
                 q_table[hashed] = 0
             # Take a stap in the environment, record the reward and state transition.
             # Effectively we are evaluating the policy by taking a step in the
             # environment.
-            obs, reward, done, info = env.step(env.action_space.flags.index(a))
+            observation, reward, done, info = env.step(env.action_space.flags.index(a))
+            if done:
+                break
             current_length += 1
 
             # Compute the target value of the current state, by using the current
@@ -154,7 +157,7 @@ def train(q_table, env):
             # can be used to emphasize on immediate early rewards, and encourage
             # the agent to achieve higher rewards sooner than later.
             target = reward + FLAGS.discount * get_max_q_value(
-                q_table, obs, current_length
+                q_table, observation, current_length
             )
 
             # Update Q value. Instead of replacing the Q value at the current
@@ -166,7 +169,7 @@ def train(q_table, env):
                 + (1 - FLAGS.learning_rate) * q_table[hashed]
             )
 
-        if i % FLAGS.log_every == 0:
+        if FLAGS.log_every and i % FLAGS.log_every == 0:
 
             def compare_qs(q_old, q_new):
                 diff = [q_new[k] - v for k, v in q_old.items()]
@@ -186,7 +189,8 @@ def main(argv):
     q_table: Dict[StateActionTuple, float] = {}
     benchmark = benchmark_from_flags()
     assert benchmark, "You must specify a benchmark using the --benchmark flag"
-    env = gym.make("llvm-autophase-ic-v0", benchmark=benchmark)
+    env = gym.make("llvm-ic-v0", benchmark=benchmark)
+    env.observation_space = "Autophase"
 
     try:
         # Train a Q-table.
@@ -194,7 +198,7 @@ def main(argv):
             train(q_table, env)
 
         # Rollout resulting policy.
-        rollout(q_table, env, True)
+        rollout(q_table, env, printout=True)
 
     finally:
         env.close()
diff --git a/leaderboard/llvm_instcount/tabular_q/tabular_q_eval.py b/leaderboard/llvm_instcount/tabular_q/tabular_q_eval.py
new file mode 100644
--- /dev/null
+++ b/leaderboard/llvm_instcount/tabular_q/tabular_q_eval.py
@@ -0,0 +1,38 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+"""Evaluate tabular_q policy for leaderboard."""
+import os
+import sys
+from typing import Dict
+
+from absl import app, flags
+
+from compiler_gym.envs import LlvmEnv
+from compiler_gym.leaderboard.llvm_instcount import eval_llvm_instcount_policy
+
+sys.path.insert(0, os.path.dirname(os.path.realpath(__file__)) + "/../../../examples")
+from tabular_q import (  # noqa pylint: disable=wrong-import-position
+    StateActionTuple,
+    rollout,
+    train,
+)
+
+FLAGS = flags.FLAGS
+
+
+def train_and_run(env: LlvmEnv) -> None:
+    """ Run tabular Q learning on an environment """
+    FLAGS.log_every = 0  # Disable printing to stdout
+
+    q_table: Dict[StateActionTuple, float] = {}
+    env.observation_space = "Autophase"
+    training_env = env.fork()
+    train(q_table, training_env)
+    training_env.close()
+    rollout(q_table, env, printout=False)
+
+
+if __name__ == "__main__":
+    app.run(eval_llvm_instcount_policy(train_and_run))


Test patch:
diff --git a/leaderboard/llvm_instcount/tabular_q/tabular_q_test.py b/leaderboard/llvm_instcount/tabular_q/tabular_q_test.py
new file mode 100644
--- /dev/null
+++ b/leaderboard/llvm_instcount/tabular_q/tabular_q_test.py
@@ -0,0 +1,34 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+"""Tests for //leaderboard/llvm_instcount/tabular_q_eval."""
+import pytest
+from absl import flags
+
+from compiler_gym.leaderboard.llvm_instcount import eval_llvm_instcount_policy
+from leaderboard.llvm_instcount.tabular_q.tabular_q_eval import train_and_run
+from tests.test_main import main as _test_main
+
+FLAGS = flags.FLAGS
+
+pytest_plugins = ["tests.pytest_plugins.llvm"]
+
+
+def test_tabular_q():
+    FLAGS.unparse_flags()
+    FLAGS(
+        [
+            "argv0",
+            "--n=1",
+            "--max_benchmarks=1",
+            "--nproc=1",
+            "--novalidate",
+        ]
+    )
+    with pytest.raises(SystemExit):
+        eval_llvm_instcount_policy(train_and_run)
+
+
+if __name__ == "__main__":
+    _test_main()