text: #   Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  #
  # Third party code
  # The following code are copied or modified from:
  # https://github.com/ray-project/ray/blob/master/python/ray/rllib/env/atari_wrappers.py
  
  import numpy as np
  from collections import deque
  import gym
  from gym import spaces
  import cv2
  cv2.ocl.setUseOpenCL(False)
  
  
  def get_wrapper_by_cls(env, cls):
      """Returns the gym env wrapper of the given class, or None."""
      currentenv = env
      while True:
          if isinstance(currentenv, cls):
              return currentenv
          elif isinstance(currentenv, gym.Wrapper):
              currentenv = currentenv.env
          else:
              return None
  
  
  class MonitorEnv(gym.Wrapper):
      def __init__(self, env=None):
          """Record episodes stats prior to EpisodicLifeEnv, etc."""
          gym.Wrapper.__init__(self, env)
          self._current_reward = None
          self._num_steps = None
          self._total_steps = None
          self._episode_rewards = []
          self._episode_lengths = []
          self._num_episodes = 0
          self._num_returned = 0
  
      def reset(self, **kwargs):
          obs = self.env.reset(**kwargs)
  
          if self._total_steps is None:
              self._total_steps = sum(self._episode_lengths)
  
          if self._current_reward is not None:
              self._episode_rewards.append(self._current_reward)
              self._episode_lengths.append(self._num_steps)
              self._num_episodes += 1
  
          self._current_reward = 0
          self._num_steps = 0
  
          return obs
  
      def step(self, action):
          obs, rew, done, info = self.env.step(action)
          self._current_reward += rew
          self._num_steps += 1
          self._total_steps += 1
          return (obs, rew, done, info)
  
      def get_episode_rewards(self):
          return self._episode_rewards
  
      def get_episode_lengths(self):
          return self._episode_lengths
  
      def get_total_steps(self):
          return self._total_steps
  
      def next_episode_results(self):
          for i in range(self._num_returned, len(self._episode_rewards)):
              yield (self._episode_rewards[i], self._episode_lengths[i])
          self._num_returned = len(self._episode_rewards)
  
  
  class NoopResetEnv(gym.Wrapper):
      def __init__(self, env, noop_max=30):
          """Sample initial states by taking random number of no-ops on reset.
          No-op is assumed to be action 0.
          """
          gym.Wrapper.__init__(self, env)
          self.noop_max = noop_max
          self.override_num_noops = None
          self.noop_action = 0
          assert env.unwrapped.get_action_meanings()[0] == 'NOOP'
  
      def reset(self, **kwargs):
          """ Do no-op action for a number of steps in [1, noop_max]."""
          self.env.reset(**kwargs)
          if self.override_num_noops is not None:
              noops = self.override_num_noops
          else:
              noops = self.unwrapped.np_random.randint(1, self.noop_max + 1)
          assert noops > 0
          obs = None
          for _ in range(noops):
              obs, _, done, _ = self.env.step(self.noop_action)
              if done:
                  obs = self.env.reset(**kwargs)
          return obs
  
      def step(self, ac):
          return self.env.step(ac)
  
  
  class ClipRewardEnv(gym.RewardWrapper):
      def __init__(self, env):
          gym.RewardWrapper.__init__(self, env)
  
      def reward(self, reward):
          """Bin reward to {+1, 0, -1} by its sign."""
          return np.sign(reward)
  
  
  class FireResetEnv(gym.Wrapper):
      def __init__(self, env):
          """Take action on reset.
  
          For environments that are fixed until firing."""
          gym.Wrapper.__init__(self, env)
          assert env.unwrapped.get_action_meanings()[1] == 'FIRE'
          assert len(env.unwrapped.get_action_meanings()) >= 3
  
      def reset(self, **kwargs):
          self.env.reset(**kwargs)
          obs, _, done, _ = self.env.step(1)
          if done:
              self.env.reset(**kwargs)
          obs, _, done, _ = self.env.step(2)
          if done:
              self.env.reset(**kwargs)
          return obs
  
      def step(self, ac):
          return self.env.step(ac)
  
  
  class EpisodicLifeEnv(gym.Wrapper):
      def __init__(self, env):
          """Make end-of-life == end-of-episode, but only reset on true game over.
          Done by DeepMind for the DQN and co. since it helps value estimation.
          """
          gym.Wrapper.__init__(self, env)
          self.lives = 0
          self.was_real_done = True
  
      def step(self, action):
          obs, reward, done, info = self.env.step(action)
          self.was_real_done = done
          # check current lives, make loss of life terminal,
          # then update lives to handle bonus lives
          lives = self.env.unwrapped.ale.lives()
          if lives < self.lives and lives > 0:
              # for Qbert sometimes we stay in lives == 0 condtion for a few fr
              # so its important to keep lives > 0, so that we only reset once
              # the environment advertises done.
              done = True
          self.lives = lives
          return obs, reward, done, info
  
      def reset(self, **kwargs):
          """Reset only when lives are exhausted.
          This way all states are still reachable even though lives are episodic,
          and the learner need not know about any of this behind-the-scenes.
          """
          if self.was_real_done:
              obs = self.env.reset(**kwargs)
          else:
              # no-op step to advance from terminal/lost life state
              obs, _, _, _ = self.env.step(0)
          self.lives = self.env.unwrapped.ale.lives()
          return obs
  
  
  class MaxAndSkipEnv(gym.Wrapper):
      def __init__(self, env, skip=4):
          """Return only every `skip`-th frame"""
          gym.Wrapper.__init__(self, env)
          # most recent raw observations (for max pooling across time steps)
          self._obs_buffer = np.zeros(
              (2, ) + env.observation_space.shape, dtype=np.uint8)
          self._skip = skip
  
      def step(self, action):
          """Repeat action, sum reward, and max over last observations."""
          total_reward = 0.0
          done = None
          for i in range(self._skip):
              obs, reward, done, info = self.env.step(action)
              if i == self._skip - 2:
                  self._obs_buffer[0] = obs
              if i == self._skip - 1:
                  self._obs_buffer[1] = obs
              total_reward += reward
              if done:
                  break
          # Note that the observation on the done=True frame
          # doesn't matter
          max_frame = self._obs_buffer.max(axis=0)
  
          return max_frame, total_reward, done, info
  
      def reset(self, **kwargs):
          return self.env.reset(**kwargs)
  
  
  class WarpFrame(gym.ObservationWrapper):
      def __init__(self, env, dim):
          """Warp frames to the specified size (dim x dim)."""
          gym.ObservationWrapper.__init__(self, env)
          self.width = dim
          self.height = dim
          self.observation_space = spaces.Box(
              low=0, high=255, shape=(self.height, self.width), dtype=np.uint8)
  
      def observation(self, frame):
          frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
          frame = cv2.resize(
              frame, (self.width, self.height), interpolation=cv2.INTER_AREA)
          return frame
  
  
  class FrameStack(gym.Wrapper):
      def __init__(self, env, k, obs_format='NHWC'):
          """Stack k last frames."""
          gym.Wrapper.__init__(self, env)
          self.k = k
          self.frames = deque([], maxlen=k)
          shp = env.observation_space.shape
          assert obs_format == 'NHWC' or obs_format == 'NCHW'
          self.obs_format = obs_format
          if obs_format == 'NHWC':
              obs_shape = (shp[0], shp[1], k)
          else:
              obs_shape = (k, shp[0], shp[1])
  
          self.observation_space = spaces.Box(
              low=0,
              high=255,
              shape=obs_shape,
              dtype=env.observation_space.dtype)
  
      def reset(self):
          ob = self.env.reset()
          for _ in range(self.k):
              self.frames.append(ob)
          return self._get_ob()
  
      def step(self, action):
          ob, reward, done, info = self.env.step(action)
          self.frames.append(ob)
          return self._get_ob(), reward, done, info
  
      def _get_ob(self):
          assert len(self.frames) == self.k
          if self.obs_format == 'NHWC':
              return np.stack(self.frames, axis=2)
          else:
              return np.array(self.frames)
  
  
  class TestEnv(gym.Wrapper):
      def __init__(self, env, eval_episodes=3):
          """ env wrapper for test and validation in atari environment.
  
          Args:
              env (ClipRewardEnv or FrameStack): the atari env
              eval_episodes (int): number of episodes for evaluation
          """
  
          gym.Wrapper.__init__(self, env)
          self._env = env
          self._monitor = get_wrapper_by_cls(env, MonitorEnv)
          self._eval_episodes = eval_episodes
          self._was_real_done = False
          self._eval_rewards = None
          self._end_episode = len(
              self._monitor.get_episode_rewards()) + eval_episodes
  
      def step(self, action):
          ob, reward, done, info = self._env.step(action)
          return ob, reward, done, info
  
      def reset(self, **kwargs):
          obs = self._env.reset(**kwargs)
          # During the noop reset in NoopResetEnv, env may be reset multiple times(may occur in mock env,
          # almost impossible in atari env), so the == condition may never be met. >= can avoid infinite loop.
          if self._get_curr_episode() >= self._end_episode:
              self._was_real_done = True
              self._eval_rewards = \
                  self._monitor.get_episode_rewards()[-self._eval_episodes:]
              self._end_episode = self._end_episode + self._eval_episodes
          else:
              self._was_real_done = False
              self._eval_rewards = None
  
          return obs
  
      def get_eval_rewards(self):
          return self._eval_rewards
  
      def get_real_done(self):
          return self._was_real_done
  
      def _get_curr_episode(self):
          return len(self._monitor.get_episode_rewards())
  
  
  def wrap_deepmind(env,
                    dim=84,
                    framestack=True,
                    obs_format='NHWC',
                    test=False,
                    eval_episodes=3):
      """Configure environment for DeepMind-style Atari.
  
      Args:
          dim (int): Dimension to resize observations to (dim x dim).
          framestack (bool): Whether to framestack observations.
          obs_format (str): observation output format
          test (bool): whether this is a test env
          eval_episodes (int): when test, number of episodes for each evaluation
      """
      env = MonitorEnv(env)
      env = NoopResetEnv(env, noop_max=30)
      if 'NoFrameskip' in env.spec.id:
          env = MaxAndSkipEnv(env, skip=4)
      env = EpisodicLifeEnv(env)
      if 'FIRE' in env.unwrapped.get_action_meanings():
          env = FireResetEnv(env)
      env = WarpFrame(env, dim)
      env = ClipRewardEnv(env)
      if framestack:
          env = FrameStack(env, 4, obs_format)
      if test:
          env = TestEnv(env, eval_episodes)
      return env
  