From 637bf55da762a0a64ea3460e07df35dfd5b9bad5 Mon Sep 17 00:00:00 2001 From: Mathieu Poliquin Date: Wed, 31 Oct 2018 01:16:15 +0800 Subject: [PATCH] Use deepmind wrapper for retro (#685) * Use deepmind wrapper for retro * moved wrap_deepmind_retro after Monitor wrapper --- baselines/common/cmd_util.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/baselines/common/cmd_util.py b/baselines/common/cmd_util.py index 7c38a77..162e34d 100644 --- a/baselines/common/cmd_util.py +++ b/baselines/common/cmd_util.py @@ -60,12 +60,14 @@ def make_env(env_id, env_type, subrank=0, seed=None, reward_scale=1.0, gamestate allow_early_resets=True) if env_type == 'atari': - return wrap_deepmind(env, **wrapper_kwargs) - elif reward_scale != 1: - return retro_wrappers.RewardScaler(env, reward_scale) - else: - return env + env = wrap_deepmind(env, **wrapper_kwargs) + elif env_type == 'retro': + env = retro_wrappers.wrap_deepmind_retro(env, **wrapper_kwargs) + if reward_scale != 1: + env = retro_wrappers.RewardScaler(env, reward_scale) + + return env def make_mujoco_env(env_id, seed, reward_scale=1.0):