import numpy as np
from .space import Space


class Discrete(Space):
    r"""A discrete space in :math:`\{ 0, 1, \\dots, n-1 \}`.

    Example::


            

Reported by Pylint.

              import numpy as np
from .space import Space


class Discrete(Space):
    r"""A discrete space in :math:`\{ 0, 1, \\dots, n-1 \}`.

    Example::


            

Reported by Pylint.

              
    """

    def __init__(self, n):
        assert n >= 0
        self.n = n
        super(Discrete, self).__init__((), np.int64)

    def sample(self):

            

Reported by Pylint.

                  """

    def __init__(self, n):
        assert n >= 0
        self.n = n
        super(Discrete, self).__init__((), np.int64)

    def sample(self):
        return self.np_random.randint(self.n)

            

Reported by Bandit.

              
    def __init__(self, n):
        assert n >= 0
        self.n = n
        super(Discrete, self).__init__((), np.int64)

    def sample(self):
        return self.np_random.randint(self.n)


            

Reported by Pylint.

                  def __init__(self, n):
        assert n >= 0
        self.n = n
        super(Discrete, self).__init__((), np.int64)

    def sample(self):
        return self.np_random.randint(self.n)

    def contains(self, x):

            

Reported by Pylint.

                      self.n = n
        super(Discrete, self).__init__((), np.int64)

    def sample(self):
        return self.np_random.randint(self.n)

    def contains(self, x):
        if isinstance(x, int):
            as_int = x

            

Reported by Pylint.

                  def sample(self):
        return self.np_random.randint(self.n)

    def contains(self, x):
        if isinstance(x, int):
            as_int = x
        elif isinstance(x, (np.generic, np.ndarray)) and (
            x.dtype.char in np.typecodes["AllInteger"] and x.shape == ()
        ):

            

Reported by Pylint.

                  def sample(self):
        return self.np_random.randint(self.n)

    def contains(self, x):
        if isinstance(x, int):
            as_int = x
        elif isinstance(x, (np.generic, np.ndarray)) and (
            x.dtype.char in np.typecodes["AllInteger"] and x.shape == ()
        ):

            

Reported by Pylint.

                          as_int = int(x)
        else:
            return False
        return as_int >= 0 and as_int < self.n

    def __repr__(self):
        return "Discrete(%d)" % self.n

    def __eq__(self, other):

            

Reported by Pylint.

              import pytest

import numpy as np

import gym
from gym.wrappers import TransformReward


@pytest.mark.parametrize("env_id", ["CartPole-v1", "Pendulum-v0"])

            

Reported by Pylint.

                  scales = [0.1, 200]
    for scale in scales:
        env = gym.make(env_id)
        wrapped_env = TransformReward(gym.make(env_id), lambda r: scale * r)
        action = env.action_space.sample()

        env.seed(0)
        env.reset()
        wrapped_env.seed(0)

            

Reported by Pylint.

              
    # use case #3: sign
    env = gym.make(env_id)
    wrapped_env = TransformReward(gym.make(env_id), lambda r: np.sign(r))

    env.seed(0)
    env.reset()
    wrapped_env.seed(0)
    wrapped_env.reset()

            

Reported by Pylint.

              import pytest

import numpy as np

import gym
from gym.wrappers import TransformReward


@pytest.mark.parametrize("env_id", ["CartPole-v1", "Pendulum-v0"])

            

Reported by Pylint.

              

@pytest.mark.parametrize("env_id", ["CartPole-v1", "Pendulum-v0"])
def test_transform_reward(env_id):
    # use case #1: scale
    scales = [0.1, 200]
    for scale in scales:
        env = gym.make(env_id)
        wrapped_env = TransformReward(gym.make(env_id), lambda r: scale * r)

            

Reported by Pylint.

                      _, reward, _, _ = env.step(action)
        _, wrapped_reward, _, _ = wrapped_env.step(action)

        assert wrapped_reward == scale * reward
    del env, wrapped_env

    # use case #2: clip
    min_r = -0.0005
    max_r = 0.0002

            

Reported by Bandit.

                  _, reward, _, _ = env.step(action)
    _, wrapped_reward, _, _ = wrapped_env.step(action)

    assert abs(wrapped_reward) < abs(reward)
    assert wrapped_reward == -0.0005 or wrapped_reward == 0.0002
    del env, wrapped_env

    # use case #3: sign
    env = gym.make(env_id)

            

Reported by Bandit.

                  _, wrapped_reward, _, _ = wrapped_env.step(action)

    assert abs(wrapped_reward) < abs(reward)
    assert wrapped_reward == -0.0005 or wrapped_reward == 0.0002
    del env, wrapped_env

    # use case #3: sign
    env = gym.make(env_id)
    wrapped_env = TransformReward(gym.make(env_id), lambda r: np.sign(r))

            

Reported by Pylint.

                  _, wrapped_reward, _, _ = wrapped_env.step(action)

    assert abs(wrapped_reward) < abs(reward)
    assert wrapped_reward == -0.0005 or wrapped_reward == 0.0002
    del env, wrapped_env

    # use case #3: sign
    env = gym.make(env_id)
    wrapped_env = TransformReward(gym.make(env_id), lambda r: np.sign(r))

            

Reported by Bandit.

                  for _ in range(1000):
        action = env.action_space.sample()
        _, wrapped_reward, done, _ = wrapped_env.step(action)
        assert wrapped_reward in [-1.0, 0.0, 1.0]
        if done:
            break
    del env, wrapped_env

            

Reported by Bandit.

                      self.closeables[next_id] = closeable
        return next_id

    def unregister(self, id):
        assert id is not None
        if id in self.closeables:
            del self.closeables[id]

    def close(self):

            

Reported by Pylint.

              import atexit
import threading
import weakref


class Closer(object):
    """A registry that ensures your objects get closed, whether manually,
    upon garbage collection, or upon exit. To work properly, your
    objects need to cooperate and do something like the following:

            

Reported by Pylint.

              import weakref


class Closer(object):
    """A registry that ensures your objects get closed, whether manually,
    upon garbage collection, or upon exit. To work properly, your
    objects need to cooperate and do something like the following:

    ```

            

Reported by Pylint.

                      if atexit_register:
            atexit.register(self.close)

    def generate_next_id(self):
        with self.lock:
            self.next_id += 1
            return self.next_id

    def register(self, closeable):

            

Reported by Pylint.

                      """Registers an object with a 'close' method.

        Returns:
            int: The registration ID of this object. It is the caller's responsibility to save this ID if early closing is desired.
        """
        assert hasattr(closeable, "close"), "No close method for {}".format(closeable)

        next_id = self.generate_next_id()
        self.closeables[next_id] = closeable

            

Reported by Pylint.

                      Returns:
            int: The registration ID of this object. It is the caller's responsibility to save this ID if early closing is desired.
        """
        assert hasattr(closeable, "close"), "No close method for {}".format(closeable)

        next_id = self.generate_next_id()
        self.closeables[next_id] = closeable
        return next_id


            

Reported by Bandit.

                      self.closeables[next_id] = closeable
        return next_id

    def unregister(self, id):
        assert id is not None
        if id in self.closeables:
            del self.closeables[id]

    def close(self):

            

Reported by Pylint.

                      self.closeables[next_id] = closeable
        return next_id

    def unregister(self, id):
        assert id is not None
        if id in self.closeables:
            del self.closeables[id]

    def close(self):

            

Reported by Pylint.

                      return next_id

    def unregister(self, id):
        assert id is not None
        if id in self.closeables:
            del self.closeables[id]

    def close(self):
        # Explicitly fetch all monitors first so that they can't disappear while

            

Reported by Bandit.

                      if id in self.closeables:
            del self.closeables[id]

    def close(self):
        # Explicitly fetch all monitors first so that they can't disappear while
        # we iterate. cf. http://stackoverflow.com/a/12429620
        closeables = list(self.closeables.values())
        for closeable in closeables:
            closeable.close()

            

Reported by Pylint.

              import os
import gym
from typing import Callable

from gym.wrappers.monitoring import video_recorder


class RecordVideo(gym.Wrapper):
    def __init__(

            

Reported by Pylint.

              import os
import gym
from typing import Callable

from gym.wrappers.monitoring import video_recorder


class RecordVideo(gym.Wrapper):
    def __init__(

            

Reported by Pylint.

              from gym.wrappers.monitoring import video_recorder


class RecordVideo(gym.Wrapper):
    def __init__(
        self,
        env,
        video_folder: str,
        record_video_trigger: Callable[[int], bool],

            

Reported by Pylint.

              from gym.wrappers.monitoring import video_recorder


class RecordVideo(gym.Wrapper):
    def __init__(
        self,
        env,
        video_folder: str,
        record_video_trigger: Callable[[int], bool],

            

Reported by Pylint.

              

class RecordVideo(gym.Wrapper):
    def __init__(
        self,
        env,
        video_folder: str,
        record_video_trigger: Callable[[int], bool],
        video_length: int = 0,

            

Reported by Pylint.

                      video_length: int = 0,
        name_prefix: str = "rl-video",
    ):
        super(RecordVideo, self).__init__(env)
        self.record_video_trigger = record_video_trigger
        self.video_recorder = None

        self.video_folder = os.path.abspath(video_folder)
        # Create output folder if needed

            

Reported by Pylint.

                      self.is_vector_env = getattr(env, "is_vector_env", False)

    def reset(self, **kwargs):
        observations = super(RecordVideo, self).reset(**kwargs)
        self.start_video_recorder()
        return observations

    def start_video_recorder(self):
        self.close_video_recorder()

            

Reported by Pylint.

                      self.start_video_recorder()
        return observations

    def start_video_recorder(self):
        self.close_video_recorder()

        video_name = f"{self.name_prefix}-step-{self.step_id}"
        base_path = os.path.join(self.video_folder, video_name)
        self.video_recorder = video_recorder.VideoRecorder(

            

Reported by Pylint.

                      return self.record_video_trigger(self.step_id)

    def step(self, action):
        observations, rewards, dones, infos = super(RecordVideo, self).step(action)

        self.step_id += 1
        if self.recording:
            self.video_recorder.capture_frame()
            self.recorded_frames += 1

            

Reported by Pylint.

              
        return observations, rewards, dones if self.is_vector_env else dones[0], infos

    def close_video_recorder(self) -> None:
        if self.recording:
            self.video_recorder.close()
        self.recording = False
        self.recorded_frames = 1

            

Reported by Pylint.

              import numpy as np
from gym.envs.mujoco import mujoco_env
from gym import utils


DEFAULT_CAMERA_CONFIG = {
    "trackbodyid": 2,
    "distance": 3.0,
    "lookat": np.array((0.0, 0.0, 1.15)),

            

Reported by Pylint.

              }


class HopperEnv(mujoco_env.MujocoEnv, utils.EzPickle):
    def __init__(
        self,
        xml_file="hopper.xml",
        forward_reward_weight=1.0,
        ctrl_cost_weight=1e-3,

            

Reported by Pylint.

              }


class HopperEnv(mujoco_env.MujocoEnv, utils.EzPickle):
    def __init__(
        self,
        xml_file="hopper.xml",
        forward_reward_weight=1.0,
        ctrl_cost_weight=1e-3,

            

Reported by Pylint.

              

class HopperEnv(mujoco_env.MujocoEnv, utils.EzPickle):
    def __init__(
        self,
        xml_file="hopper.xml",
        forward_reward_weight=1.0,
        ctrl_cost_weight=1e-3,
        healthy_reward=1.0,

            

Reported by Pylint.

                      mujoco_env.MujocoEnv.__init__(self, xml_file, 4)

    @property
    def healthy_reward(self):
        return (
            float(self.is_healthy or self._terminate_when_unhealthy)
            * self._healthy_reward
        )


            

Reported by Pylint.

                          * self._healthy_reward
        )

    def control_cost(self, action):
        control_cost = self._ctrl_cost_weight * np.sum(np.square(action))
        return control_cost

    @property
    def is_healthy(self):

            

Reported by Pylint.

                      return control_cost

    @property
    def is_healthy(self):
        z, angle = self.sim.data.qpos[1:3]
        state = self.state_vector()[2:]

        min_state, max_state = self._healthy_state_range
        min_z, max_z = self._healthy_z_range

            

Reported by Pylint.

              
    @property
    def is_healthy(self):
        z, angle = self.sim.data.qpos[1:3]
        state = self.state_vector()[2:]

        min_state, max_state = self._healthy_state_range
        min_z, max_z = self._healthy_z_range
        min_angle, max_angle = self._healthy_angle_range

            

Reported by Pylint.

                      return is_healthy

    @property
    def done(self):
        done = not self.is_healthy if self._terminate_when_unhealthy else False
        return done

    def _get_obs(self):
        position = self.sim.data.qpos.flat.copy()

            

Reported by Pylint.

              import numpy as np
from gym import utils
from gym.envs.mujoco import mujoco_env


class InvertedDoublePendulumEnv(mujoco_env.MujocoEnv, utils.EzPickle):
    def __init__(self):
        mujoco_env.MujocoEnv.__init__(self, "inverted_double_pendulum.xml", 5)
        utils.EzPickle.__init__(self)

            

Reported by Pylint.

              from gym.envs.mujoco import mujoco_env


class InvertedDoublePendulumEnv(mujoco_env.MujocoEnv, utils.EzPickle):
    def __init__(self):
        mujoco_env.MujocoEnv.__init__(self, "inverted_double_pendulum.xml", 5)
        utils.EzPickle.__init__(self)

    def step(self, action):

            

Reported by Pylint.

              
    def step(self, action):
        self.do_simulation(action, self.frame_skip)
        ob = self._get_obs()
        x, _, y = self.sim.data.site_xpos[0]
        dist_penalty = 0.01 * x ** 2 + (y - 2) ** 2
        v1, v2 = self.sim.data.qvel[1:3]
        vel_penalty = 1e-3 * v1 ** 2 + 5e-3 * v2 ** 2
        alive_bonus = 10

            

Reported by Pylint.

                  def step(self, action):
        self.do_simulation(action, self.frame_skip)
        ob = self._get_obs()
        x, _, y = self.sim.data.site_xpos[0]
        dist_penalty = 0.01 * x ** 2 + (y - 2) ** 2
        v1, v2 = self.sim.data.qvel[1:3]
        vel_penalty = 1e-3 * v1 ** 2 + 5e-3 * v2 ** 2
        alive_bonus = 10
        r = alive_bonus - dist_penalty - vel_penalty

            

Reported by Pylint.

                  def step(self, action):
        self.do_simulation(action, self.frame_skip)
        ob = self._get_obs()
        x, _, y = self.sim.data.site_xpos[0]
        dist_penalty = 0.01 * x ** 2 + (y - 2) ** 2
        v1, v2 = self.sim.data.qvel[1:3]
        vel_penalty = 1e-3 * v1 ** 2 + 5e-3 * v2 ** 2
        alive_bonus = 10
        r = alive_bonus - dist_penalty - vel_penalty

            

Reported by Pylint.

                      ob = self._get_obs()
        x, _, y = self.sim.data.site_xpos[0]
        dist_penalty = 0.01 * x ** 2 + (y - 2) ** 2
        v1, v2 = self.sim.data.qvel[1:3]
        vel_penalty = 1e-3 * v1 ** 2 + 5e-3 * v2 ** 2
        alive_bonus = 10
        r = alive_bonus - dist_penalty - vel_penalty
        done = bool(y <= 1)
        return ob, r, done, {}

            

Reported by Pylint.

                      ob = self._get_obs()
        x, _, y = self.sim.data.site_xpos[0]
        dist_penalty = 0.01 * x ** 2 + (y - 2) ** 2
        v1, v2 = self.sim.data.qvel[1:3]
        vel_penalty = 1e-3 * v1 ** 2 + 5e-3 * v2 ** 2
        alive_bonus = 10
        r = alive_bonus - dist_penalty - vel_penalty
        done = bool(y <= 1)
        return ob, r, done, {}

            

Reported by Pylint.

                      v1, v2 = self.sim.data.qvel[1:3]
        vel_penalty = 1e-3 * v1 ** 2 + 5e-3 * v2 ** 2
        alive_bonus = 10
        r = alive_bonus - dist_penalty - vel_penalty
        done = bool(y <= 1)
        return ob, r, done, {}

    def _get_obs(self):
        return np.concatenate(

            

Reported by Pylint.

                      return self._get_obs()

    def viewer_setup(self):
        v = self.viewer
        v.cam.trackbodyid = 0
        v.cam.distance = self.model.stat.extent * 0.5
        v.cam.lookat[2] = 0.12250000000000005  # v.model.stat.center[2]

            

Reported by Pylint.

                      self.shape = (len(frames),) + self.frame_shape
        self.dtype = frames[0].dtype
        if lz4_compress:
            from lz4.block import compress

            frames = [compress(frame) for frame in frames]
        self._frames = frames
        self.lz4_compress = lz4_compress


            

Reported by Pylint.

              
    def _check_decompress(self, frame):
        if self.lz4_compress:
            from lz4.block import decompress

            return np.frombuffer(decompress(frame), dtype=self.dtype).reshape(
                self.frame_shape
            )
        return frame

            

Reported by Pylint.

              
    def reset(self, **kwargs):
        observation = self.env.reset(**kwargs)
        [self.frames.append(observation) for _ in range(self.num_stack)]
        return self._get_observation()

            

Reported by Pylint.

              from collections import deque
import numpy as np
from gym.spaces import Box
from gym import Wrapper


class LazyFrames(object):
    r"""Ensures common frames are only stored once to optimize memory use.


            

Reported by Pylint.

              from gym import Wrapper


class LazyFrames(object):
    r"""Ensures common frames are only stored once to optimize memory use.

    To further reduce the memory use, it is optionally to turn on lz4 to
    compress the observations.


            

Reported by Pylint.

                      self.shape = (len(frames),) + self.frame_shape
        self.dtype = frames[0].dtype
        if lz4_compress:
            from lz4.block import compress

            frames = [compress(frame) for frame in frames]
        self._frames = frames
        self.lz4_compress = lz4_compress


            

Reported by Pylint.

              
    def _check_decompress(self, frame):
        if self.lz4_compress:
            from lz4.block import decompress

            return np.frombuffer(decompress(frame), dtype=self.dtype).reshape(
                self.frame_shape
            )
        return frame

            

Reported by Pylint.

                  """

    def __init__(self, env, num_stack, lz4_compress=False):
        super(FrameStack, self).__init__(env)
        self.num_stack = num_stack
        self.lz4_compress = lz4_compress

        self.frames = deque(maxlen=num_stack)


            

Reported by Pylint.

                      )

    def _get_observation(self):
        assert len(self.frames) == self.num_stack, (len(self.frames), self.num_stack)
        return LazyFrames(list(self.frames), self.lz4_compress)

    def step(self, action):
        observation, reward, done, info = self.env.step(action)
        self.frames.append(observation)

            

Reported by Bandit.

              from gym import spaces


class RescaleAction(gym.ActionWrapper):
    r"""Rescales the continuous action space of the environment to a range [a,b].

    Example::

        >>> RescaleAction(env, a, b).action_space == Box(a,b)

            

Reported by Pylint.

              import numpy as np
import gym
from gym import spaces


class RescaleAction(gym.ActionWrapper):
    r"""Rescales the continuous action space of the environment to a range [a,b].

    Example::

            

Reported by Pylint.

                  """

    def __init__(self, env, a, b):
        assert isinstance(
            env.action_space, spaces.Box
        ), "expected Box action space, got {}".format(type(env.action_space))
        assert np.less_equal(a, b).all(), (a, b)

        super(RescaleAction, self).__init__(env)

            

Reported by Bandit.

                      assert isinstance(
            env.action_space, spaces.Box
        ), "expected Box action space, got {}".format(type(env.action_space))
        assert np.less_equal(a, b).all(), (a, b)

        super(RescaleAction, self).__init__(env)
        self.a = np.zeros(env.action_space.shape, dtype=env.action_space.dtype) + a
        self.b = np.zeros(env.action_space.shape, dtype=env.action_space.dtype) + b
        self.action_space = spaces.Box(

            

Reported by Bandit.

                      ), "expected Box action space, got {}".format(type(env.action_space))
        assert np.less_equal(a, b).all(), (a, b)

        super(RescaleAction, self).__init__(env)
        self.a = np.zeros(env.action_space.shape, dtype=env.action_space.dtype) + a
        self.b = np.zeros(env.action_space.shape, dtype=env.action_space.dtype) + b
        self.action_space = spaces.Box(
            low=a, high=b, shape=env.action_space.shape, dtype=env.action_space.dtype
        )

            

Reported by Pylint.

                      assert np.less_equal(a, b).all(), (a, b)

        super(RescaleAction, self).__init__(env)
        self.a = np.zeros(env.action_space.shape, dtype=env.action_space.dtype) + a
        self.b = np.zeros(env.action_space.shape, dtype=env.action_space.dtype) + b
        self.action_space = spaces.Box(
            low=a, high=b, shape=env.action_space.shape, dtype=env.action_space.dtype
        )


            

Reported by Pylint.

              
        super(RescaleAction, self).__init__(env)
        self.a = np.zeros(env.action_space.shape, dtype=env.action_space.dtype) + a
        self.b = np.zeros(env.action_space.shape, dtype=env.action_space.dtype) + b
        self.action_space = spaces.Box(
            low=a, high=b, shape=env.action_space.shape, dtype=env.action_space.dtype
        )

    def action(self, action):

            

Reported by Pylint.

                      )

    def action(self, action):
        assert np.all(np.greater_equal(action, self.a)), (action, self.a)
        assert np.all(np.less_equal(action, self.b)), (action, self.b)
        low = self.env.action_space.low
        high = self.env.action_space.high
        action = low + (high - low) * ((action - self.a) / (self.b - self.a))
        action = np.clip(action, low, high)

            

Reported by Bandit.

              
    def action(self, action):
        assert np.all(np.greater_equal(action, self.a)), (action, self.a)
        assert np.all(np.less_equal(action, self.b)), (action, self.b)
        low = self.env.action_space.low
        high = self.env.action_space.high
        action = low + (high - low) * ((action - self.a) / (self.b - self.a))
        action = np.clip(action, low, high)
        return action

            

Reported by Bandit.

                      self.num_envs = getattr(env, "num_envs", 1)
        self.t0 = (
            time.time()
        )  # TODO: use perf_counter when gym removes Python 2 support
        self.episode_count = 0
        self.episode_returns = None
        self.episode_lengths = None
        self.return_queue = deque(maxlen=deque_size)
        self.length_queue = deque(maxlen=deque_size)

            

Reported by Pylint.

              import time
from collections import deque
import numpy as np
import gym


class RecordEpisodeStatistics(gym.Wrapper):
    def __init__(self, env, deque_size=100):
        super(RecordEpisodeStatistics, self).__init__(env)

            

Reported by Pylint.

              import gym


class RecordEpisodeStatistics(gym.Wrapper):
    def __init__(self, env, deque_size=100):
        super(RecordEpisodeStatistics, self).__init__(env)
        self.num_envs = getattr(env, "num_envs", 1)
        self.t0 = (
            time.time()

            

Reported by Pylint.

              import gym


class RecordEpisodeStatistics(gym.Wrapper):
    def __init__(self, env, deque_size=100):
        super(RecordEpisodeStatistics, self).__init__(env)
        self.num_envs = getattr(env, "num_envs", 1)
        self.t0 = (
            time.time()

            

Reported by Pylint.

              
class RecordEpisodeStatistics(gym.Wrapper):
    def __init__(self, env, deque_size=100):
        super(RecordEpisodeStatistics, self).__init__(env)
        self.num_envs = getattr(env, "num_envs", 1)
        self.t0 = (
            time.time()
        )  # TODO: use perf_counter when gym removes Python 2 support
        self.episode_count = 0

            

Reported by Pylint.

                  def __init__(self, env, deque_size=100):
        super(RecordEpisodeStatistics, self).__init__(env)
        self.num_envs = getattr(env, "num_envs", 1)
        self.t0 = (
            time.time()
        )  # TODO: use perf_counter when gym removes Python 2 support
        self.episode_count = 0
        self.episode_returns = None
        self.episode_lengths = None

            

Reported by Pylint.

                      self.is_vector_env = getattr(env, "is_vector_env", False)

    def reset(self, **kwargs):
        observations = super(RecordEpisodeStatistics, self).reset(**kwargs)
        self.episode_returns = np.zeros(self.num_envs, dtype=np.float32)
        self.episode_lengths = np.zeros(self.num_envs, dtype=np.int32)
        return observations

    def step(self, action):

            

Reported by Pylint.

                      return observations

    def step(self, action):
        observations, rewards, dones, infos = super(RecordEpisodeStatistics, self).step(
            action
        )
        self.episode_returns += rewards
        self.episode_lengths += 1
        if not self.is_vector_env:

            

Reported by Pylint.

                      if not self.is_vector_env:
            infos = [infos]
            dones = [dones]
        for i in range(len(dones)):
            if dones[i]:
                infos[i] = infos[i].copy()
                episode_return = self.episode_returns[i]
                episode_length = self.episode_lengths[i]
                episode_info = {

            

Reported by Pylint.

              import pytest

import numpy as np

import gym
from gym.wrappers import GrayScaleObservation
from gym.wrappers import AtariPreprocessing

pytest.importorskip("atari_py")

            

Reported by Pylint.

              import pytest

import numpy as np

import gym
from gym.wrappers import GrayScaleObservation
from gym.wrappers import AtariPreprocessing

pytest.importorskip("atari_py")

            

Reported by Pylint.

              @pytest.mark.parametrize(
    "env_id", ["PongNoFrameskip-v0", "SpaceInvadersNoFrameskip-v0"]
)
@pytest.mark.parametrize("keep_dim", [True, False])
def test_gray_scale_observation(env_id, keep_dim):
    gray_env = AtariPreprocessing(gym.make(env_id), screen_size=84, grayscale_obs=True)
    rgb_env = AtariPreprocessing(gym.make(env_id), screen_size=84, grayscale_obs=False)
    wrapped_env = GrayScaleObservation(rgb_env, keep_dim=keep_dim)
    assert rgb_env.observation_space.shape[-1] == 3

            

Reported by Pylint.

                  gray_env = AtariPreprocessing(gym.make(env_id), screen_size=84, grayscale_obs=True)
    rgb_env = AtariPreprocessing(gym.make(env_id), screen_size=84, grayscale_obs=False)
    wrapped_env = GrayScaleObservation(rgb_env, keep_dim=keep_dim)
    assert rgb_env.observation_space.shape[-1] == 3

    seed = 0
    gray_env.seed(seed)
    wrapped_env.seed(seed)


            

Reported by Bandit.

                  wrapped_obs = wrapped_env.reset()

    if keep_dim:
        assert wrapped_env.observation_space.shape[-1] == 1
        assert len(wrapped_obs.shape) == 3
        wrapped_obs = wrapped_obs.squeeze(-1)
    else:
        assert len(wrapped_env.observation_space.shape) == 2
        assert len(wrapped_obs.shape) == 2

            

Reported by Bandit.

              
    if keep_dim:
        assert wrapped_env.observation_space.shape[-1] == 1
        assert len(wrapped_obs.shape) == 3
        wrapped_obs = wrapped_obs.squeeze(-1)
    else:
        assert len(wrapped_env.observation_space.shape) == 2
        assert len(wrapped_obs.shape) == 2


            

Reported by Bandit.

                      assert len(wrapped_obs.shape) == 3
        wrapped_obs = wrapped_obs.squeeze(-1)
    else:
        assert len(wrapped_env.observation_space.shape) == 2
        assert len(wrapped_obs.shape) == 2

    # ALE gray scale is slightly different, but no more than by one shade
    assert np.allclose(gray_obs.astype("int32"), wrapped_obs.astype("int32"), atol=1)

            

Reported by Bandit.

                      wrapped_obs = wrapped_obs.squeeze(-1)
    else:
        assert len(wrapped_env.observation_space.shape) == 2
        assert len(wrapped_obs.shape) == 2

    # ALE gray scale is slightly different, but no more than by one shade
    assert np.allclose(gray_obs.astype("int32"), wrapped_obs.astype("int32"), atol=1)

            

Reported by Bandit.

                      assert len(wrapped_obs.shape) == 2

    # ALE gray scale is slightly different, but no more than by one shade
    assert np.allclose(gray_obs.astype("int32"), wrapped_obs.astype("int32"), atol=1)

            

Reported by Bandit.