Fix pyright errors in envs and spaces folders

gym/envs/box2d/bipedal_walker.py

@@ -376,7 +376,7 @@ def _generate_terrain(self, hardcore):
             self.terrain_y.append(y)
             counter -= 1
             if counter == 0:
-                counter = self.np_random.integers(TERRAIN_GRASS / 2, TERRAIN_GRASS)
+                counter = self.np_random.integers(TERRAIN_GRASS // 2, TERRAIN_GRASS)
                 if state == GRASS and hardcore:
                     state = self.np_random.integers(1, _STATES_)
                     oneshot = True
@@ -617,7 +617,7 @@ def render(self):
         if self.screen is None and self.render_mode == "human":
             pygame.init()
             pygame.display.init()
-            self.screen = pygame.display.set_mode((VIEWPORT_W, VIEWPORT_H))
+            self.screen = pygame.display.set_mode((VIEWPORT_W, VIEWPORT_H))  # type: ignore
         if self.clock is None:
             self.clock = pygame.time.Clock()
 
@@ -795,10 +795,10 @@ def __init__(self):
         moving_s_base = 4 + 5 * moving_leg
         supporting_s_base = 4 + 5 * supporting_leg
 
-        hip_targ = [None, None]  # -0.8 .. +1.1
-        knee_targ = [None, None]  # -0.6 .. +0.9
-        hip_todo = [0.0, 0.0]
-        knee_todo = [0.0, 0.0]
+        hip_targ: List[Optional[float]] = [None, None]  # -0.8 .. +1.1
+        knee_targ: List[Optional[float]] = [None, None]  # -0.6 .. +0.9
+        hip_todo: List[float] = [0.0, 0.0]
+        knee_todo: List[float] = [0.0, 0.0]
 
         if state == STAY_ON_ONE_LEG:
             hip_targ[moving_leg] = 1.1

gym/envs/box2d/car_dynamics.py

@@ -261,7 +261,7 @@ def step(self, dt):
             )
 
     def draw(self, surface, zoom, translation, angle, draw_particles=True):
-        import pygame.draw
+        import pygame
 
         if draw_particles:
             for p in self.particles:
@@ -334,10 +334,10 @@ class Particle:
             pass
 
         p = Particle()
-        p.color = WHEEL_COLOR if not grass else MUD_COLOR
-        p.ttl = 1
-        p.poly = [(point1[0], point1[1]), (point2[0], point2[1])]
-        p.grass = grass
+        p.color = WHEEL_COLOR if not grass else MUD_COLOR  # type: ignore
+        p.ttl = 1  # type: ignore
+        p.poly = [(point1[0], point1[1]), (point2[0], point2[1])]  # type: ignore
+        p.grass = grass  # type: ignore
         self.particles.append(p)
         while len(self.particles) > 30:
             self.particles.pop(0)

gym/envs/box2d/car_racing.py

@@ -517,10 +517,11 @@ def reset(
             self.render()
         return self.step(None)[0], {}
 
-    def step(self, action: Union[np.ndarray, int]):
+    def step(self, action: Optional[Union[np.ndarray, int]]):
         assert self.car is not None
         if action is not None:
             if self.continuous:
+                assert isinstance(action, np.ndarray)
                 self.car.steer(-action[0])
                 self.car.gas(action[1])
                 self.car.brake(action[2])
@@ -565,6 +566,7 @@ def step(self, action: Union[np.ndarray, int]):
         return self.state, step_reward, terminated, truncated, {}
 
     def render(self):
+        assert self.render_mode is not None, "render_mode is not set"
         return self._render(self.render_mode)
 
     def _render(self, mode: str):
@@ -574,7 +576,7 @@ def _render(self, mode: str):
         if self.screen is None and mode == "human":
             pygame.init()
             pygame.display.init()
-            self.screen = pygame.display.set_mode((WINDOW_W, WINDOW_H))
+            self.screen = pygame.display.set_mode((WINDOW_W, WINDOW_H))  # type: ignore
         if self.clock is None:
             self.clock = pygame.time.Clock()
 
@@ -610,14 +612,14 @@ def _render(self, mode: str):
         font = pygame.font.Font(pygame.font.get_default_font(), 42)
         text = font.render("%04i" % self.reward, True, (255, 255, 255), (0, 0, 0))
         text_rect = text.get_rect()
-        text_rect.center = (60, WINDOW_H - WINDOW_H * 2.5 / 40.0)
+        text_rect.center = (60, WINDOW_H - WINDOW_H * 2.5 / 40.0)  # type: ignore
         self.surf.blit(text, text_rect)
 
         if mode == "human":
             pygame.event.pump()
             self.clock.tick(self.metadata["render_fps"])
             assert self.screen is not None
-            self.screen.fill(0)
+            self.screen.fill((0, 0, 0))  # black - (0,0,0)
             self.screen.blit(self.surf, (0, 0))
             pygame.display.flip()
 
@@ -671,6 +673,7 @@ def _render_indicators(self, W, H):
         h = H / 40.0
         color = (0, 0, 0)
         polygon = [(W, H), (W, H - 5 * h), (0, H - 5 * h), (0, H)]
+        assert self.surf is not None, "No valid pygame surface"
         pygame.draw.polygon(self.surf, color=color, points=polygon)
 
         def vertical_ind(place, val):
@@ -698,6 +701,7 @@ def horiz_ind(place, val):
         # simple wrapper to render if the indicator value is above a threshold
         def render_if_min(value, points, color):
             if abs(value) > 1e-4:
+                assert self.surf is not None, "No valid pygame surface"
                 pygame.draw.polygon(self.surf, points=points, color=color)
 
         render_if_min(true_speed, vertical_ind(5, 0.02 * true_speed), (255, 255, 255))
@@ -751,6 +755,7 @@ def _draw_colored_polygon(
             and (-MAX_SHAPE_DIM <= coord[1] <= WINDOW_H + MAX_SHAPE_DIM)
             for coord in poly
         ):
+            assert self.surf is not None, "No valid pygame surface"
             gfxdraw.aapolygon(self.surf, poly, color)
             gfxdraw.filled_polygon(self.surf, poly, color)
 

gym/envs/box2d/lunar_lander.py

@@ -2,7 +2,7 @@
 
 import math
 import warnings
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Optional, Union
 
 import numpy as np
 
@@ -234,12 +234,12 @@ def __init__(
         self.wind_idx = np.random.randint(-9999, 9999)
         self.torque_idx = np.random.randint(-9999, 9999)
 
-        self.screen: pygame.Surface = None
+        self.screen: Optional[pygame.Surface] = None
         self.clock = None
         self.isopen = True
         self.world = Box2D.b2World(gravity=(0, gravity))
         self.moon = None
-        self.lander: Optional[Box2D.b2Body] = None
+        self.lander: Optional[Box2D.b2Body] = None  # type: ignore
         self.particles = []
 
         self.prev_reward = None
@@ -441,7 +441,8 @@ def _clean_particles(self, all):
         while self.particles and (all or self.particles[0].ttl < 0):
             self.world.DestroyBody(self.particles.pop(0))
 
-    def step(self, action):
+    def step(self, action: Union[np.ndarray, int]):
+        # Todo: remove duplicate assert
         assert self.lander is not None
 
         # Update wind
@@ -477,7 +478,7 @@ def step(self, action):
             )
 
         if self.continuous:
-            action = np.clip(action, -1, +1).astype(np.float32)
+            action = np.clip(action, -1, +1).astype(np.float32)  # type: ignore
         else:
             assert self.action_space.contains(
                 action
@@ -489,11 +490,13 @@ def step(self, action):
         dispersion = [self.np_random.uniform(-1.0, +1.0) / SCALE for _ in range(2)]
 
         m_power = 0.0
-        if (self.continuous and action[0] > 0.0) or (
+
+        if (self.continuous and action[0] > 0.0) or (  # type: ignore
             not self.continuous and action == 2
         ):
             # Main engine
             if self.continuous:
+                assert isinstance(action, np.ndarray)
                 m_power = (np.clip(action[0], 0.0, 1.0) + 1.0) * 0.5  # 0.5..1.0
                 assert m_power >= 0.5 and m_power <= 1.0
             else:
@@ -520,11 +523,12 @@ def step(self, action):
             )
 
         s_power = 0.0
-        if (self.continuous and np.abs(action[1]) > 0.5) or (
+        if (self.continuous and np.abs(action[1]) > 0.5) or (  # type: ignore
             not self.continuous and action in [1, 3]
         ):
             # Orientation engines
             if self.continuous:
+                assert isinstance(action, np.ndarray)
                 direction = np.sign(action[1])
                 s_power = np.clip(np.abs(action[1]), 0.5, 1.0)
                 assert s_power >= 0.5 and s_power <= 1.0
@@ -611,7 +615,7 @@ def render(self):
         if self.screen is None and self.render_mode == "human":
             pygame.init()
             pygame.display.init()
-            self.screen = pygame.display.set_mode((VIEWPORT_W, VIEWPORT_H))
+            self.screen = pygame.display.set_mode((VIEWPORT_W, VIEWPORT_H))  # type: ignore
         if self.clock is None:
             self.clock = pygame.time.Clock()
 
@@ -776,7 +780,7 @@ def demo_heuristic_lander(env, seed=None, render=False):
     s, info = env.reset(seed=seed)
     while True:
         a = heuristic(env, s)
-        s, r, terminated, truncated, info = step_api_compatibility(env.step(a), True)
+        s, r, terminated, truncated, info = step_api_compatibility(env.step(a), True)  # type: ignore
         total_reward += r
 
         if render:

gym/envs/classic_control/pendulum.py

@@ -147,8 +147,8 @@ def reset(self, *, seed: Optional[int] = None, options: Optional[dict] = None):
             # state/observations.
             x = options.get("x_init") if "x_init" in options else DEFAULT_X
             y = options.get("y_init") if "y_init" in options else DEFAULT_Y
-            x = utils.verify_number_and_cast(x)
-            y = utils.verify_number_and_cast(y)
+            x = utils.verify_number_and_cast(x)  # type: ignore
+            y = utils.verify_number_and_cast(y)  # type: ignore
             high = np.array([x, y])
         low = -high  # We enforce symmetric limits.
         self.state = self.np_random.uniform(low=low, high=high)

gym/envs/classic_control/utils.py

@@ -36,8 +36,8 @@ def maybe_parse_reset_bounds(
     high = options.get("high") if "high" in options else default_high
 
     # We expect only numerical inputs.
-    low = verify_number_and_cast(low)
-    high = verify_number_and_cast(high)
+    low = verify_number_and_cast(low)  # type: ignore
+    high = verify_number_and_cast(high)  # type: ignore
     if low > high:
         raise ValueError(
             f"Lower bound ({low}) must be lower than higher bound ({high})."

gym/envs/mujoco/half_cheetah.py

@@ -28,7 +28,7 @@ def step(self, action):
         xposafter = self.sim.data.qpos[0]
 
         ob = self._get_obs()
-        reward_ctrl = -0.1 * np.square(action).sum()
+        reward_ctrl = -0.1 * np.square(action).sum()  # type: ignore
         reward_run = (xposafter - xposbefore) / self.dt
         reward = reward_ctrl + reward_run
         terminated = False

gym/envs/mujoco/mujoco_env.py

@@ -1,5 +1,5 @@
 from os import path
-from typing import Optional, Union
+from typing import Any, Dict, Optional, Union
 
 import numpy as np
 
@@ -48,11 +48,15 @@ def __init__(
 
         self.width = width
         self.height = height
+        self.data: Any  # to be set in _initilize_simulation
+        self.model: Any  # to be set in _initilize_simulation
         self._initialize_simulation()  # may use width and height
 
         self.init_qpos = self.data.qpos.ravel().copy()
         self.init_qvel = self.data.qvel.ravel().copy()
-        self._viewers = {}
+        self._viewers: Dict[
+            str, Union["mujoco_py.MjViewer", "mujoco_py.MjRenderContextOffscreen"]
+        ] = {}
 
         self.frame_skip = frame_skip
 
@@ -247,7 +251,7 @@ def render(self):
 
                 self._get_viewer(self.render_mode).render(
                     width, height, camera_id=camera_id
-                )
+                )  # type: ignore
 
         if self.render_mode == "rgb_array":
             data = self._get_viewer(self.render_mode).read_pixels(
@@ -372,12 +376,12 @@ def render(self):
                     camera_name,
                 )
 
-                self._get_viewer(self.render_mode).render(camera_id=camera_id)
+                self._get_viewer(self.render_mode).render(camera_id=camera_id)  # type: ignore
 
         if self.render_mode == "rgb_array":
             data = self._get_viewer(self.render_mode).read_pixels(depth=False)
             # original image is upside-down, so flip it
-            return data[::-1, :, :]
+            return data[::-1, :, :]  # type: ignore
         elif self.render_mode == "depth_array":
             self._get_viewer(self.render_mode).render()
             # Extract depth part of the read_pixels() tuple
@@ -397,7 +401,7 @@ def _get_viewer(
     ) -> Union[
         "gym.envs.mujoco.mujoco_rendering.Viewer",
         "gym.envs.mujoco.mujoco_rendering.RenderContextOffscreen",
-    ]:
+    ]:  # type: ignore
         self.viewer = self._viewers.get(mode)
         if self.viewer is None:
             if mode == "human":

gym/envs/mujoco/mujoco_rendering.py

@@ -2,6 +2,7 @@
 import os
 import time
 from threading import Lock
+from typing import Callable, Dict
 
 import glfw
 import imageio
@@ -27,7 +28,7 @@ def _import_osmesa(width, height):
     return GLContext(width, height)
 
 
-_ALL_RENDERERS = collections.OrderedDict(
+_ALL_RENDERERS: Dict[str, Callable] = collections.OrderedDict(
     [
         ("glfw", _import_glfw),
         ("egl", _import_egl),
@@ -197,8 +198,10 @@ def _add_marker_to_scene(self, marker):
                 attr[:] = np.asarray(value).reshape(attr.shape)
             elif isinstance(value, str):
                 assert key == "label", "Only label is a string in mjtGeom."
+                # Todo: we may consider removing this condition as 'value' cannot be
+                # None in this scope.
                 if value is None:
-                    g.label[0] = 0
+                    g.label[0] = 0  # type: ignore
                 else:
                     g.label = value
             elif hasattr(g, key):

gym/envs/registration.py

@@ -48,7 +48,7 @@
 )
 
 
-def load(name: str) -> callable:
+def load(name: str) -> Callable:
     """Loads an environment with name and returns an environment creation function
 
     Args:
@@ -492,7 +492,7 @@ def register(
         disable_env_checker=disable_env_checker,
         apply_api_compatibility=apply_api_compatibility,
         **kwargs,
-    )
+    )  # type: ignore
     _check_spec_register(new_spec)
     if new_spec.id in registry:
         logger.warn(f"Overriding environment {new_spec.id} already in registry.")

gym/envs/toy_text/cliffwalking.py

@@ -71,7 +71,7 @@ def __init__(self, render_mode: Optional[str] = None):
         self.shape = (4, 12)
         self.start_state_index = np.ravel_multi_index((3, 0), self.shape)
 
-        self.nS = np.prod(self.shape)
+        self.nS = int(np.prod(self.shape))
         self.nA = 4
 
         # Cliff Location
@@ -233,6 +233,8 @@ def _render_gui(self, mode):
 
         for s in range(self.nS):
             row, col = np.unravel_index(s, self.shape)
+            row = int(row)
+            col = int(col)
             pos = (col * self.cell_size[0], row * self.cell_size[1])
             check_board_mask = row % 2 ^ col % 2
             self.window_surface.blit(self.mountain_bg_img[check_board_mask], pos)