def test_wrapped_step_multi_step(env: LlvmEnv):
env = TimeLimit(env, max_episode_steps=5)
env.reset(benchmark="benchmark://cbench-v1/dijkstra")
env.step([0, 0, 0])
assert env.benchmark == "benchmark://cbench-v1/dijkstra"
assert env.actions == [0, 0, 0]
def test_wrapped_fork_type(env: LlvmEnv):
env = TimeLimit(env, max_episode_steps=5)
fkd = env.fork()
try:
assert isinstance(fkd, TimeLimit)
finally:
fkd.close()
def make_env():
FLAGS.env = "llvm-v0"
if not FLAGS.reward:
FLAGS.reward = "IrInstructionCountOz"
env = env_from_flags(benchmark=benchmark_from_flags())
env = ConstrainedCommandline(env, flags=FLAGS.flags)
env = TimeLimit(env, max_episode_steps=FLAGS.episode_len)
env = HistoryObservation(env)
return env
def make_env(self) -> CompilerEnv:
"""Construct a compiler environment from the given config."""
env = compiler_gym.make(self.id)
if self.observation_space:
env.observation_space = self.observation_space
if self.reward_space:
env.reward_space = self.reward_space
env = TimeLimit(env, max_episode_steps=self.max_episode_steps)
for wrapper in self.wrappers:
env = wrapper.wrap(env)
# Wrap the env to ignore errors during search.
env = JustKeepGoingEnv(env)
return env
def test_time_limit_reached(env: LlvmEnv):
env = TimeLimit(env, max_episode_steps=3)
env.reset()
_, _, done, info = env.step(0)
assert not done, info
_, _, done, info = env.step(0)
assert not done, info
_, _, done, info = env.step(0)
assert done, info
assert info["TimeLimit.truncated"], info
_, _, done, info = env.step(0)
assert done, info
assert info["TimeLimit.truncated"], info
def test_time_limit_fork(env: LlvmEnv):
"""Check that the time limit state is copied on fork()."""
env = TimeLimit(env, max_episode_steps=3)
env.reset()
_, _, done, info = env.step(0) # 1st step
assert not done, info
fkd = env.fork()
try:
_, _, done, info = env.step(0) # 2nd step
assert not done, info
_, _, done, info = fkd.step(0) # 2nd step
assert not done, info
_, _, done, info = env.step(0) # 3rd step
assert done, info
_, _, done, info = fkd.step(0) # 3rd step
assert done, info
finally:
fkd.close()
def test_wrapped_custom_step_args(env: LlvmEnv):
env = TimeLimit(env, max_episode_steps=5)
env.reset(benchmark="benchmark://cbench-v1/dijkstra")
(ic, ), _, _, _ = env.step(0, observations=["IrInstructionCount"])
assert isinstance(ic, int)
def test_wrapped_close(env: LlvmEnv):
env = TimeLimit(env, max_episode_steps=5)
env.close()
assert env.service is None