def make_atari()

in jat/eval/rl/core.py [0:0]

• 21 lines of code
• 5 McCabe index (conditional complexity)

def make_atari(task_name: str, episodic_life: bool = True, clip_reward: bool = True, **kwargs) -> Env:
    kwargs = {"frameskip": 1, "repeat_action_probability": 0.0, **kwargs}
    if task_name == "atari-montezumarevenge":
        kwargs["max_episode_steps"] = 18_000
    env = gym.make(TASK_NAME_TO_ENV_ID[task_name], **kwargs)
    env.metadata["render_fps"] = 30
    env = gym.wrappers.RecordEpisodeStatistics(env)
    env = NoopResetEnv(env, noop_max=30)
    env = MaxAndSkipEnv(env, skip=4)
    if episodic_life:
        env = EpisodicLifeEnv(env)
    if "FIRE" in env.unwrapped.get_action_meanings():
        env = FireResetEnv(env)
    if clip_reward:
        env = ClipRewardEnv(env)
    env = gym.wrappers.ResizeObservation(env, (84, 84))
    env = gym.wrappers.GrayScaleObservation(env)
    env = gym.wrappers.FrameStack(env, num_stack=4)
    env = NumpyObsWrapper(env)
    env = AtariDictObservationWrapper(env)
    return env