def wrap_deepmind(env,
downsample=True,
episode_life=True,
clip_rewards=True,
frame_stack=False,
scale=False,
color=False):
"""Configure environment for DeepMind-style Atari.
"""
if ("videopinball" in str(env.spec.id).lower()) or ('tennis' in str(
env.spec.id).lower()) or ('skiing' in str(env.spec.id).lower()):
env = WarpFrame(env, width=160, height=210, grayscale=False)
if episode_life:
env = EpisodicLifeEnv(env)
if 'FIRE' in env.unwrapped.get_action_meanings():
env = FireResetEnv(env)
if downsample:
env = WarpFrame(env, grayscale=False)
if not color:
env = GrayscaleWrapper(env)
if scale:
env = ScaledFloatFrame(env)
if clip_rewards:
env = ClipRewardEnv(env)
if frame_stack:
env = FrameStack(env, 4)
return env
def wrap_custom(env, clip_rewards=False, scale=True):
"""Configure environment for Openai procgen.
"""
if scale:
env = ScaledFloatFrame(env)
if clip_rewards:
env = ClipRewardEnv(env)
return env
def _wrap_deepmind_ram(env):
"""Applies various Atari-specific wrappers to make learning easier."""
env = EpisodicLifeEnv(env)
env = NoopResetEnv(env, noop_max=30)
env = MaxAndSkipEnv(env, skip=4)
if 'FIRE' in env.unwrapped.get_action_meanings():
env = FireResetEnv(env)
env = ClipRewardEnv(env)
return env
def wrap_env_ppo(env):
env = ThresholdResizeFrame(env)
# env = WarpFrame(env)
env = ClipRewardEnv(env)
# env = NoopResetEnv(env, noop_max=8)
env = MaxAndSkipEnv(env, skip=4)
env = Monitor(env, logger.get_dir())
env = DummyVecEnv([lambda: env])
env = VecFrameStack(env, 4)
return env
def wrap_deepmind_retro(env, scale=True, frame_stack=0):
"""
Configure environment for retro games, using config similar to DeepMind-style Atari in wrap_deepmind
"""
env = WarpFrame(env, grayscale=False)
env = ClipRewardEnv(env)
if frame_stack > 1:
env = FrameStack(env, frame_stack)
if scale:
env = ScaledFloatFrame(env)
return env
def _thunk():
env = gym.make(env_id)
env.seed(seed + rank)
env = bench.Monitor(env, logger.get_dir() and os.path.join(logger.get_dir(), str(rank)))
gym.logger.setLevel(logging.WARN)
#env = NoopResetEnv(env, noop_max=30)
#env = MaxAndSkipEnv(env, skip=4)
#if 'FIRE' in env.unwrapped.get_action_meanings():
# env = FireResetEnv(env)
#env = WarpFrame(env)
env = ClipRewardEnv(env)
return env
def make_podworld(env_id: str,
max_episode_steps=2000,
clip_rewards=True,
frame_stack=False):
env = PodWorldEnv(obs_mode='R', agent_ray_count=84, agent_obs_height=84)
if max_episode_steps is not None:
env = TimeLimit(env, max_episode_steps=max_episode_steps)
if clip_rewards:
env = ClipRewardEnv(env)
if frame_stack:
env = FrameStack(env, 4)
return env
def wrap_deepmind(env, episode_life=True, clip_rewards=True, frame_stack=False, scale=False):
"""Configure environment for DeepMind-style Atari.
"""
from baselines.common.atari_wrappers import ScaledFloatFrame, ClipRewardEnv, FrameStack
#env = WarpFrame(env)
if scale:
env = ScaledFloatFrame2(env)
if clip_rewards:
env = ClipRewardEnv(env)
if frame_stack:
env = FrameStack(env, 4)
return env
def _thunk():
env = gym.make(**env_base)
env = NoopResetEnv(env, noop_max=30)
env = MaxAndSkipEnv(env, skip=4)
env = RewardCollector(env)
env = EpisodicLifeEnv(env)
env = ClipRewardEnv(env)
env = WarpFrame(env)
env = ScaledFloatFrame(env)
env = TransposeImage(env)
env = UnrealEnvBaseWrapper(env)
return env
def create_env(self, env):
env = gym.make(env)
env = NoopResetEnv(env, noop_max=30)
env = MaxAndSkipEnv(env, skip=4)
env = RewardCollector(env)
env = EpisodicLifeEnv(env)
env = ClipRewardEnv(env)
env = WarpFrame(env)
env = FrameStack(env, 4)
env = ConvertToNumpy(env)
env = TransposeImage(env)
env = ScaledFloatFrame(env)
return env
def wrap_deepmind(env, episode_life=True, clip_rewards=True, frame_stack=False, scale=False):
"""Configure environment for DeepMind-style Atari.
"""
if episode_life:
env = EpisodicLifeEnv(env)
if 'FIRE' in env.unwrapped.get_action_meanings():
env = FireResetEnv(env)
env = WarpFrame(env)
if scale:
env = ScaledFloatFrame(env)
if clip_rewards:
env = ClipRewardEnv(env)
if frame_stack:
env = FrameStack(env, 4)
return env
def _thunk():
env = gym.make(env_id)
env.seed(seed + rank)
env = NoopResetEnv(env, noop_max=30)
env = MaxAndSkipEnv(env, skip=2)
env = WarpFrame(env)
# Janky Fix to Resize Environments to be 50x50
env.width = 50
env.height = 50
env = ScaledFloatFrame(env)
if not eval:
env = ClipRewardEnv(env)
env = EpisodicLifeEnv(env)
env = FrameStack(env, 3)
env = TransposeOb(env)
return env
def wrap_deepmind_custom(env,
episode_life=True,
clip_rewards=True,
frame_stack=frame_stack,
scale=False):
if episode_life:
env = EpisodicLifeEnv(env)
if 'FIRE' in env.unwrapped.get_action_meanings():
env = FireResetEnv(env)
env = WarpFrame(env, size=size)
if augment:
env = AugmentColor(env)
if scale:
env = ScaledFloatFrame(env)
if clip_rewards:
env = ClipRewardEnv(env)
if frame_stack:
env = FrameStack(env, frame_stack)
return env
def deepmind_wrap(atari_env,
episode_life=True,
clip_rewards=True,
frame_stack=False,
scale=False):
""" matching deepmind papers
"""
if episode_life:
env = EpisodicLifeEnv(atari_env)
if 'FIRE' in env.unwrapped.get_action_meanings():
env = FireResetEnv(env)
env = WarpFrame(env)
if scale:
env = ScaledFloatFrame(env)
if clip_rewards:
env = ClipRewardEnv(env)
if frame_stack:
env = FrameStack(env, 4)
return env
def make_env(args, env_id, seed, rank, log_dir, add_timestep, train=True,
natural=False, clip_rewards=True, loader=None):
# def _thunk():
if train:
vid_path = KINETICS_PATH
else:
vid_path = KINETICS_PATH_TEST
if env_id in IMG_ENVS:
if env_id == 'mnist':
channels = 2
else:
channels = 4
env = ImgEnv(env_id, max_steps=args.max_steps, channels=channels,
window=args.window, train=train)
elif env_id in ['cityscapes']:
env = DetectionEnv(env_id, max_steps=200, train=train)
elif env_id.startswith("dm"):
_, domain, task = env_id.split('.')
env = dm_control2gym.make(domain_name=domain, task_name=task)
else:
env = gym.make(env_id)
is_atari = hasattr(gym.envs, 'atari') and isinstance(
env.unwrapped, gym.envs.atari.atari_env.AtariEnv)
is_mujoco = hasattr(gym.envs, 'mujoco') and isinstance(
env.unwrapped, gym.envs.mujoco.MujocoEnv)
if is_atari:
env = make_atari(env_id)
if natural and is_atari:
env = ReplaceBackgroundEnv(
env,
BackgroundMattingWithColor((0, 0, 0)),
#RandomColorSource(shape2d)
#RandomImageSource(shape2d, glob.glob(COCO_PATH))
RandomVideoSource(env.observation_space.shape[:2],
glob.glob(vid_path)))
elif natural and is_mujoco:
env.observation_space = Box(
low=0, high=255, shape=(128, 128, 3), dtype=np.uint8)
env = ReplaceMuJoCoBackgroundEnv(
env,
BackgroundMattingWithColor((0, 0, 0)),
#RandomColorSource(shape2d)
#RandomImageSource(shape2d, glob.glob(COCO_PATH))
RandomVideoSource(env.observation_space.shape[:2],
glob.glob(vid_path)))
elif is_mujoco:
env.observation_space = Box(
low=0, high=255, shape=(128, 128, 3), dtype=np.uint8)
env = PixelMujoCoEnv(env)
env.seed(seed + rank)
obs_shape = env.observation_space.shape
if add_timestep and len(
obs_shape) == 1 and str(env).find('TimeLimit') > -1:
env = AddTimestep(env)
if log_dir is not None:
env = bench.Monitor(env, os.path.join(log_dir, str(rank)))
if is_atari:
env = wrap_deepmind(env, clip_rewards=clip_rewards)
if is_mujoco:
env = ClipRewardEnv(WarpFrame(MaxAndSkipEnv(env, skip=4)))
# If the input has shape (W,H,3), wrap for PyTorch convolutions
obs_shape = env.observation_space.shape
if len(obs_shape) == 3 and obs_shape[2] in [1, 3]:
env = WrapPyTorch(env)
return env
def _thunk():
env = gym.make(**env_kwargs)
env = RewardCollector(env)
env = ClipRewardEnv(env)
env = W(env)
return env