Don't import pygame when not necessary

documentation/source/action.ipynb

@@ -68,7 +68,7 @@
    ],
    "source": [
     "from metadrive.component.vehicle.base_vehicle import BaseVehicle\n",
-    "from metadrive.utils import print_source\n",
+    "from metadrive.utils.doc_utils import print_source\n",
     "print_source(BaseVehicle._set_action)\n",
     "print_source(BaseVehicle._apply_throttle_brake)"
    ]
@@ -96,7 +96,7 @@
    "source": [
     "from metadrive.envs.metadrive_env import MetaDriveEnv\n",
     "from metadrive.component.vehicle.vehicle_type import DefaultVehicle\n",
-    "from metadrive.utils import generate_gif\n",
+    "from metadrive.utils.doc_utils import generate_gif\n",
     "\n",
     "env=MetaDriveEnv(dict(map=\"S\", traffic_density=0))\n",
     "frames = []\n",
@@ -199,7 +199,7 @@
    ],
    "source": [
     "from metadrive.policy.env_input_policy import EnvInputPolicy\n",
-    "from metadrive.utils import print_source\n",
+    "from metadrive.utils.doc_utils import print_source\n",
     "print_source(EnvInputPolicy.get_input_space)"
    ]
   },
@@ -363,7 +363,7 @@
    "outputs": [],
    "source": [
     "from metadrive.envs.metadrive_env import MetaDriveEnv\n",
-    "from metadrive.utils import generate_gif\n",
+    "from metadrive.utils.doc_utils import generate_gif\n",
     "\n",
     "env=MetaDriveEnv(dict(map=\"S\",\n",
     "                      log_level=50,\n",
@@ -435,7 +435,7 @@
    "source": [
     "from metadrive.envs.metadrive_env import MetaDriveEnv\n",
     "from metadrive.policy.lange_change_policy import LaneChangePolicy\n",
-    "from metadrive.utils import generate_gif\n",
+    "from metadrive.utils.doc_utils import generate_gif\n",
     "\n",
     "env=MetaDriveEnv(dict(map=\"C\",\n",
     "                      discrete_action=True,\n",
@@ -545,7 +545,7 @@
    ],
    "source": [
     "from metadrive.policy.env_input_policy import ExtraEnvInputPolicy\n",
-    "from metadrive.utils import print_source\n",
+    "from metadrive.utils.doc_utils import print_source\n",
     "print_source(ExtraEnvInputPolicy.get_input_space)"
    ]
   },
@@ -882,7 +882,7 @@
     "from metadrive.engine.asset_loader import AssetLoader\n",
     "from metadrive.policy.replay_policy import ReplayEgoCarPolicy\n",
     "from metadrive.envs.scenario_env import ScenarioEnv\n",
-    "from metadrive.utils import generate_gif\n",
+    "from metadrive.utils.doc_utils import generate_gif\n",
     "import cv2\n",
     "from IPython.display import Image\n",
     "\n",

documentation/source/config_system.ipynb

@@ -147,7 +147,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from metadrive.utils import print_source\n",
+    "from metadrive.utils.doc_utils import print_source\n",
     "print_source(ScenarioEnv.default_config)"
    ]
   },
@@ -241,7 +241,7 @@
    "outputs": [],
    "source": [
     "import metadrive.envs.base_env as base_env\n",
-    "from metadrive.utils import print_source, CONFIG\n",
+    "from metadrive.utils.doc_utils import print_source, CONFIG\n",
     "module_source = print_source(base_env, [\"BASE_DEFAULT_CONFIG\", \")\\n\\n\"], colorscheme=CONFIG)"
    ]
   },

documentation/source/obs.ipynb

@@ -91,7 +91,7 @@
     "import cv2\n",
     "from metadrive.policy.idm_policy import IDMPolicy\n",
     "from IPython.display import Image\n",
-    "from metadrive.utils import generate_gif\n",
+    "from metadrive.utils.doc_utils import generate_gif\n",
     "import numpy as np\n",
     "import os\n",
     "sensor_size = (84, 60) if os.getenv('TEST_DOC') else (200, 100)\n",
@@ -340,7 +340,7 @@
    },
    "outputs": [],
    "source": [
-    "from metadrive import TopDownMetaDrive\n",
+    "from metadrive.envs.top_down_env import TopDownMetaDrive\n",
     "\n",
     "env = TopDownMetaDrive()\n",
     "try:\n",
@@ -373,7 +373,7 @@
    },
    "outputs": [],
    "source": [
-    "from metadrive.utils import print_source\n",
+    "from metadrive.utils.doc_utils import print_source\n",
     "from metadrive.envs.metadrive_env import MetaDriveEnv\n",
     "print_source(MetaDriveEnv.get_single_observation)"
    ]
@@ -395,7 +395,7 @@
    },
    "outputs": [],
    "source": [
-    "from metadrive.utils import print_source\n",
+    "from metadrive.utils.doc_utils import print_source\n",
     "from metadrive.envs.top_down_env import TopDownMetaDrive, TopDownSingleFrameMetaDriveEnv\n",
     "print_source(TopDownSingleFrameMetaDriveEnv)\n",
     "print_source(TopDownMetaDrive)"
@@ -511,7 +511,7 @@
     }
    ],
    "source": [
-    "from metadrive.utils import generate_gif\n",
+    "from metadrive.utils.doc_utils import generate_gif\n",
     "from IPython.display import Image\n",
     "\n",
     "frames = []\n",
@@ -605,7 +605,7 @@
     "from metadrive.obs.observation_base import BaseObservation\n",
     "from metadrive.obs.image_obs import ImageObservation\n",
     "import os\n",
-    "from metadrive.utils import generate_gif\n",
+    "from metadrive.utils.doc_utils import generate_gif\n",
     "from IPython.display import Image\n",
     "sensor_size = (1, 1) if os.getenv('TEST_DOC') else (200, 200)"
    ]

documentation/source/reward_cost_done.ipynb

@@ -85,7 +85,7 @@
    "source": [
     "from metadrive.envs.metadrive_env import MetaDriveEnv\n",
     "from metadrive.envs.scenario_env import ScenarioEnv\n",
-    "from metadrive.utils import print_source\n",
+    "from metadrive.utils.doc_utils import print_source\n",
     "print_source(MetaDriveEnv.reward_function)"
    ]
   },
@@ -290,7 +290,7 @@
     }
    ],
    "source": [
-    "from metadrive.utils import print_source\n",
+    "from metadrive.utils.doc_utils import print_source\n",
     "from metadrive.envs import MetaDriveEnv\n",
     "print_source(MetaDriveEnv.cost_function)"
    ]

documentation/source/rl_environments.ipynb

@@ -816,7 +816,7 @@
     "from metadrive.envs.varying_dynamics_env import VaryingDynamicsEnv\n",
     "import pygame\n",
     "import matplotlib.pyplot as plt\n",
-    "from metadrive.utils import generate_gif\n",
+    "from metadrive.utils.doc_utils import generate_gif\n",
     "import cv2\n",
     "from IPython.display import Image\n",
     "\n",
@@ -978,7 +978,7 @@
     }
    ],
    "source": [
-    "from metadrive.utils import print_source\n",
+    "from metadrive.utils.doc_utils import print_source\n",
     "from metadrive import MetaDriveEnv\n",
     "print_source(MetaDriveEnv.done_function)"
    ]
@@ -1299,7 +1299,7 @@
     "from metadrive.engine.asset_loader import AssetLoader\n",
     "from metadrive.policy.replay_policy import ReplayEgoCarPolicy\n",
     "from metadrive.envs.scenario_env import ScenarioEnv\n",
-    "from metadrive.utils import generate_gif\n",
+    "from metadrive.utils.doc_utils import generate_gif\n",
     "from IPython.display import Image, clear_output\n",
     "import cv2\n",
     "\n",
@@ -1384,7 +1384,7 @@
     "from metadrive.engine.asset_loader import AssetLoader\n",
     "from metadrive.policy.replay_policy import ReplayEgoCarPolicy\n",
     "from metadrive.envs.scenario_env import ScenarioEnv\n",
-    "from metadrive.utils import generate_gif\n",
+    "from metadrive.utils.doc_utils import generate_gif\n",
     "import cv2\n",
     "from IPython.display import Image\n",
     "\n",

documentation/source/system_design.ipynb

@@ -47,7 +47,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from metadrive.utils import print_source\n",
+    "from metadrive.utils.doc_utils import print_source\n",
     "from metadrive.envs.base_env import BaseEnv\n",
     "print_source(BaseEnv.step)"
    ]

documentation/source/top_down_render.ipynb

@@ -35,7 +35,7 @@
    "source": [
     "from metadrive.envs import MetaDriveEnv\n",
     "from IPython.display import Image\n",
-    "from metadrive.utils import print_source, get_source\n",
+    "from metadrive.utils.doc_utils import print_source, get_source\n",
     "import cv2"
    ]
   },
@@ -93,7 +93,7 @@
    "outputs": [],
    "source": [
     "from metadrive.engine.top_down_renderer import TopDownRenderer\n",
-    "from metadrive.utils import CONFIG, FUNC_2\n",
+    "from metadrive.utils.doc_utils import CONFIG, FUNC_2\n",
     "print_source(TopDownRenderer.__init__, [\"def\", \"# doc-end\"], colorscheme=FUNC_2)"
    ]
   },

documentation/source/training.ipynb

@@ -53,7 +53,7 @@
     "import matplotlib.pyplot as plt\n",
     "from stable_baselines3.common.monitor import Monitor\n",
     "from metadrive.component.map.base_map import BaseMap\n",
-    "from metadrive.utils import generate_gif\n",
+    "from metadrive.utils.doc_utils import generate_gif\n",
     "from IPython.display import Image\n",
     "\n",
     "def create_env(need_monitor=False):\n",

metadrive/__init__.py

@@ -1,12 +1,6 @@
-from os import environ
-
-environ['PYGAME_HIDE_SUPPORT_PROMPT'] = '1'
-import pygame  # it is important to import pygame after that
+import os
 
-from metadrive.envs import MetaDriveEnv, TopDownMetaDrive, TopDownSingleFrameMetaDriveEnv, TopDownMetaDriveEnvV2, \
-    SafeMetaDriveEnv, MultiAgentRoundaboutEnv, MultiAgentIntersectionEnv, MultiAgentParkingLotEnv, \
-    MultiAgentTollgateEnv, MultiAgentBottleneckEnv, MultiAgentMetaDrive, ScenarioEnv
+from metadrive.envs import MetaDriveEnv, SafeMetaDriveEnv, ScenarioEnv
 from metadrive.utils.registry import get_metadrive_class
-import os
 
 MetaDrive_PACKAGE_DIR = os.path.dirname(os.path.abspath(__file__))

metadrive/engine/core/manual_controller.py

@@ -14,8 +14,11 @@
 
 from metadrive.utils import import_pygame
 
-pygame, gfxdraw = import_pygame()
-
+pygame = None
+def get_pygame():
+    global pygame
+    if not pygame:
+        pygame = import_pygame()
 
 class Controller:
     def process_input(self, vehicle):
@@ -39,6 +42,7 @@ class KeyboardController(Controller):
     def __init__(self, pygame_control):
         self.pygame_control = pygame_control
         if self.pygame_control:
+            get_pygame()
             pygame.init()
         else:
             self.inputs = InputState()
@@ -141,6 +145,7 @@ def __init__(self):
                 "Fail to load evdev, which is required for steering wheel control. "
                 "Install evdev via pip install evdev"
             )
+        get_pygame()
         pygame.display.init()
         pygame.joystick.init()
         assert not is_win(), "Steering Wheel is supported in linux and mac only"
@@ -223,6 +228,7 @@ def __init__(self):
                 "Fail to load evdev, which is required for steering wheel control. "
                 "Install evdev via pip install evdev"
             )
+        get_pygame()
         pygame.display.init()
         pygame.joystick.init()
         assert not is_win(), "Joystick is supported in linux and mac only"

metadrive/engine/top_down_renderer.py

@@ -1,7 +1,7 @@
 import copy
 from metadrive.engine.logger import get_logger
 
-from metadrive.utils import generate_gif
+from metadrive.utils.doc_utils import generate_gif
 import math
 from collections import deque
 from typing import Optional, Union, Iterable
@@ -16,7 +16,7 @@
 from metadrive.utils.utils import import_pygame
 from metadrive.utils.utils import is_map_related_instance
 
-pygame, gfxdraw = import_pygame()
+pygame = import_pygame()
 
 color_white = (255, 255, 255)
 

metadrive/envs/__init__.py

@@ -1,8 +1,5 @@
-from metadrive.envs.marl_envs import MultiAgentMetaDrive, MultiAgentBottleneckEnv, MultiAgentTollgateEnv, \
-    MultiAgentParkingLotEnv, MultiAgentIntersectionEnv, MultiAgentRoundaboutEnv
 from metadrive.envs.metadrive_env import MetaDriveEnv
 from metadrive.envs.safe_metadrive_env import SafeMetaDriveEnv
-from metadrive.envs.top_down_env import TopDownSingleFrameMetaDriveEnv, TopDownMetaDrive, TopDownMetaDriveEnvV2
 from metadrive.envs.varying_dynamics_env import VaryingDynamicsEnv
 from metadrive.envs.scenario_env import ScenarioEnv
 from metadrive.envs.base_env import BaseEnv

metadrive/examples/Basic_MetaDrive_Usages.ipynb

@@ -234,7 +234,7 @@
     "os.environ['SDL_VIDEODRIVER']='dummy'\n",
     "# Note: this step is only for cheating Colab\n",
     "\n",
-    "from metadrive import MultiAgentRoundaboutEnv, MultiAgentBottleneckEnv, MultiAgentIntersectionEnv, MultiAgentParkingLotEnv, MultiAgentTollgateEnv\n",
+    "from metadrive.envs.marl_envs import MultiAgentRoundaboutEnv, MultiAgentBottleneckEnv, MultiAgentIntersectionEnv, MultiAgentParkingLotEnv, MultiAgentTollgateEnv\n",
     "from metadrive.examples import expert\n",
     "\n",
     "env_classes = [MultiAgentRoundaboutEnv, \n",

metadrive/examples/profile_metadrive_marl.py

@@ -5,7 +5,7 @@
 
 import numpy as np
 
-from metadrive import MultiAgentRoundaboutEnv
+from metadrive.envs.marl_envs import MultiAgentRoundaboutEnv
 from metadrive.utils import setup_logger
 
 if __name__ == '__main__':

metadrive/examples/top_down_metadrive.py

@@ -18,7 +18,7 @@
 
 import matplotlib.pyplot as plt
 
-from metadrive import TopDownMetaDrive
+from metadrive.envs.top_down_env import TopDownMetaDrive
 from metadrive.constants import HELP_MESSAGE
 from metadrive.examples.ppo_expert.numpy_expert import expert
 

metadrive/obs/top_down_obs.py

@@ -16,7 +16,7 @@
     ObjectGraphics, LaneGraphics
 from metadrive.utils import import_pygame
 
-pygame, gfxdraw = import_pygame()
+pygame = import_pygame()
 
 
 class TopDownObservation(BaseObservation):

metadrive/obs/top_down_obs_impl.py

@@ -11,7 +11,7 @@
 from collections import namedtuple
 
 PositionType = Union[Tuple[float, float], np.ndarray]
-pygame, gfxdraw = import_pygame()
+pygame = import_pygame()
 COLOR_BLACK = pygame.Color("black")
 history_object = namedtuple("history_object", "name position heading_theta WIDTH LENGTH color done type")
 

metadrive/obs/top_down_obs_multi_channel.py

@@ -19,7 +19,7 @@
 from metadrive.component.navigation_module.edge_network_navigation import EdgeNetworkNavigation
 from metadrive.component.navigation_module.trajectory_navigation import TrajectoryNavigation
 
-pygame, gfxdraw = import_pygame()
+pygame = import_pygame()
 COLOR_WHITE = pygame.Color("white")
 DEFAULT_TRAJECTORY_LANE_WIDTH = 3
 

metadrive/tests/scripts/profile_top_down_multi_channel_env.py

@@ -1,6 +1,6 @@
 import time
 
-from metadrive import TopDownMetaDriveEnvV2
+from metadrive.envs.top_down_env import TopDownMetaDriveEnvV2
 
 if __name__ == '__main__':
     env = TopDownMetaDriveEnvV2(dict(num_scenarios=10, frame_stack=10, frame_skip=3))

metadrive/tests/test_export_record_scenario/test_save_replay_episode.py

@@ -2,7 +2,7 @@
 from metadrive.utils.math import wrap_to_pi
 import numpy as np
 
-from metadrive import MultiAgentRoundaboutEnv
+from metadrive.envs.marl_envs import MultiAgentRoundaboutEnv
 from metadrive.component.map.base_map import BaseMap
 from metadrive.component.map.pg_map import MapGenerateMethod
 from metadrive.envs.safe_metadrive_env import SafeMetaDriveEnv

metadrive/utils/__init__.py

@@ -5,4 +5,3 @@
 from metadrive.utils.registry import get_metadrive_class
 from metadrive.utils.utils import is_mac, import_pygame, recursive_equal, setup_logger, merge_dicts, \
     concat_step_infos, is_win, time_me
-from metadrive.utils.doc_utils import print_source, list_files, get_source, generate_gif, CONFIG, FUNC, FUNC_2

metadrive/utils/draw_top_down_map.py

@@ -5,7 +5,7 @@
 from metadrive.engine.top_down_renderer import draw_top_down_map_native as native_draw
 from metadrive.utils.utils import import_pygame
 
-pygame, gfxdraw = import_pygame()
+pygame = import_pygame()
 
 
 def draw_top_down_map(map,

metadrive/utils/utils.py

@@ -27,8 +27,7 @@ def is_port_occupied(port, host='127.0.0.1'):
 def import_pygame():
     os.environ['PYGAME_HIDE_SUPPORT_PROMPT'] = "hide"
     import pygame
-    from pygame import gfxdraw
-    return pygame, gfxdraw
+    return pygame
 
 
 def get_time_str():