Derived observation spaces are unusable during construction time

[ChrisCummins]

🐛 Bug

Observation spaces that are registered using add_derived_observation_space() are unusable during construction time.

To Reproduce

Steps to reproduce the behavior:

In [1]: import compiler_gym

In [2]: env = compiler_gym.make("llvm-v0", observation_space="AutophaseDict")
---------------------------------------------------------------------------
LookupError                               Traceback (most recent call last)
<ipython-input-2-1bde98f62255> in <module>
----> 1 env = compiler_gym.make("llvm-v0", observation_space="AutophaseDict")

~/anaconda3/envs/compiler_gym/lib/python3.8/site-packages/compiler_gym-0.2.0-py3.8.egg/compiler_gym/util/registration.py in make(id, **kwargs)
     14 def make(id: str, **kwargs):
     15     """Equivalent to :code:`gym.make()`."""
---> 16     return gym.make(id, **kwargs)
     17 
     18 

~/anaconda3/envs/compiler_gym/lib/python3.8/site-packages/gym/envs/registration.py in make(id, **kwargs)
    143 
    144 def make(id, **kwargs):
--> 145     return registry.make(id, **kwargs)
    146 
    147 def spec(id):

~/anaconda3/envs/compiler_gym/lib/python3.8/site-packages/gym/envs/registration.py in make(self, path, **kwargs)
     88             logger.info('Making new env: %s', path)
     89         spec = self.spec(path)
---> 90         env = spec.make(**kwargs)
     91         # We used to have people override _reset/_step rather than
     92         # reset/step. Set _gym_disable_underscore_compat = True on

~/anaconda3/envs/compiler_gym/lib/python3.8/site-packages/gym/envs/registration.py in make(self, **kwargs)
     58         else:
     59             cls = load(self.entry_point)
---> 60             env = cls(**_kwargs)
     61 
     62         # Make the environment aware of which spec it came from.

~/anaconda3/envs/compiler_gym/lib/python3.8/site-packages/compiler_gym-0.2.0-py3.8.egg/compiler_gym/envs/llvm/llvm_env.py in __init__(self, benchmark, datasets_site_path, *args, **kwargs)
     69         # the LLVM service and are not included by the pip-installed package.
     70         download_llvm_files()
---> 71         super().__init__(
     72             *args,
     73             **kwargs,

~/anaconda3/envs/compiler_gym/lib/python3.8/site-packages/compiler_gym-0.2.0-py3.8.egg/compiler_gym/envs/compiler_env.py in __init__(self, service, rewards, datasets, benchmark, observation_space, reward_space, action_space, connection_settings, service_connection, logger)
    316         self.observation_space_spec: Optional[ObservationSpaceSpec] = None
    317         self.reward_space_spec: Optional[Reward] = None
--> 318         self.observation_space = observation_space
    319         self.reward_space = reward_space
    320 

~/anaconda3/envs/compiler_gym/lib/python3.8/site-packages/compiler_gym-0.2.0-py3.8.egg/compiler_gym/envs/compiler_env.py in observation_space(self, observation_space)
    541         if observation_space:
    542             if observation_space not in self.observation.spaces:
--> 543                 raise LookupError(f"Observation space not found: {observation_space}")
    544             self.observation_space_spec = self.observation.spaces[observation_space]
    545         else:

LookupError: Observation space not found: AutophaseDict

The workaround is to first construct the environment and then assign the observation space:

In [3]: env = compiler_gym.make("llvm-v0")

In [4]: env.observation_space = "AutophaseDict"

Expected behavior

In [1]: import compiler_gym

In [2]: env = compiler_gym.make("llvm-v0", observation_space="AutophaseDict")

Environment

Please fill in this checklist:

• CompilerGym: v0.2.0
• How you installed CompilerGym (conda, pip, source): source
• OS: macOS
• Python version: 3.8
• Build command you used (if compiling from source): make install