Fix memory leak when calling reset and close repeatedly (#589)

* memory leak tools

* randomizable clear

* clear no ramdom

* leak

* clear map memory

* close observation

* clear

* fix

* debug memory usage

* debug memory usage

* debug memory usage

* fix close reset memort leak

* format

* fix test bug

* engine core

* maintain obs all the time

* fix

metadrive/base_class/base_object.py

@@ -287,7 +287,6 @@ def destroy(self):
 
             self.dynamic_nodes.clear()
             self.static_nodes.clear()
-            self._config.clear()
 
     def set_position(self, position, height=None):
         """

metadrive/base_class/base_runnable.py

@@ -93,7 +93,13 @@ def sample_parameters(self):
         self.update_config(ret)
 
     def destroy(self):
+        """
+        Destroy base classes
+        """
         Configurable.destroy(self)
+        Randomizable.destroy(self)
+        Nameable.destroy(self)
+        self.PARAMETER_SPACE.destroy()
 
     @property
     def engine(self):

metadrive/base_class/configurable.py

@@ -31,7 +31,9 @@ def destroy(self):
         """
         Fully delete this element and release the memory
         """
-        self._config.clear()
+        if self._config is not None:
+            self._config.clear()
+            self._config = None
         # if hasattr(self, "engine"):
         #     self.engine = None
 

metadrive/base_class/nameable.py

@@ -31,3 +31,9 @@ def __str__(self):
     def rename(self, new_name):
         self.name = new_name
         self.id = self.name
+
+    def destroy(self):
+        """
+        Clear memory
+        """
+        self.name = self.id = None

metadrive/base_class/randomizable.py

@@ -19,3 +19,9 @@ def seed(self, random_seed):
 
     def generate_seed(self):
         return self.np_random.randint(0, self.MAX_RAND_INT)
+
+    def destroy(self):
+        """
+        Destroy random generator
+        """
+        self.np_random = None

metadrive/component/algorithm/BIG.py

@@ -165,3 +165,23 @@ def _destruct_current(self):
 
     def __del__(self):
         logging.debug("Destroy Big")
+
+    def destroy(self):
+        """
+        Destroy BIG
+        """
+        self.block_dist_config = None
+        self._block_sequence = None
+        self.random_seed = None
+        self.np_random = None
+        # Don't change this right now, since we need to make maps identical to old one
+        self._lane_num = None
+        self._lane_width = None
+        self.block_num = None
+        self._render_node_path = None
+        self._physics_world = None
+        self._global_network = None
+        self.blocks = []
+        self._exit_length = None
+        self.blocks = None
+        self.next_step = None

metadrive/component/lane/point_lane.py

@@ -151,7 +151,7 @@ def destroy(self):
         self.end = None
         self._polygon = None
         InterpolatingLine.destroy(self)
-        super(PointLane, self).destroy()
+        AbstractLane.destroy(self)
 
     @property
     def polygon(self):

metadrive/component/map/pg_map.py

@@ -75,7 +75,7 @@ def _big_generate(self, parent_node_path: NodePath, physics_world: PhysicsWorld)
         )
         big_map.generate(self._config[self.GENERATE_TYPE], self._config[self.GENERATE_CONFIG])
         self.blocks = big_map.blocks
-        del big_map
+        big_map.destroy()
 
     def _config_generate(self, blocks_config: List, parent_node_path: NodePath, physics_world: PhysicsWorld):
         assert len(self.road_network.graph) == 0, "These Map is not empty, please create a new map to read config"

metadrive/component/pg_space.py

@@ -61,8 +61,11 @@ def from_jsonable(self, sample_n):
         # By default, assume identity is JSONable
         return sample_n
 
-    def __del__(self):
-        del self.np_random
+    def destroy(self):
+        """
+        Clear memory
+        """
+        self.np_random = None
 
 
 class Dict(Space):

metadrive/engine/base_engine.py

@@ -537,6 +537,8 @@ def close(self):
             del self.top_down_renderer
             self.top_down_renderer = None
 
+        Randomizable.destroy(self)
+
     def __del__(self):
         logger.debug("{} is destroyed".format(self.__class__.__name__))
 

metadrive/engine/core/engine_core.py

@@ -389,6 +389,7 @@ def close_world(self):
         self.physics_world.destroy()
         self.destroy()
         close_asset_loader()
+        # EngineCore.global_config.clear()
         EngineCore.global_config = None
 
         import sys

metadrive/envs/base_env.py

@@ -645,10 +645,6 @@ def _get_step_return(self, actions, engine_info):
     def close(self):
         if self.engine is not None:
             close_engine()
-        # if self.logger is not None:
-        #     self.logger.handlers.clear()
-        #     del self.logger
-        # self.logger = None
 
     def force_close(self):
         print("Closing environment ... Please wait")

metadrive/manager/agent_manager.py

@@ -398,6 +398,8 @@ def destroy(self):
         self._dying_objects = {}
 
         # Dict[object_id: value], init for **only** once after spawning vehicle
+        for obs in self.observations.values():
+            obs.destroy()
         self.observations = {}
         self.observation_spaces = {}
         self.action_spaces = {}

metadrive/manager/base_manager.py

@@ -4,7 +4,6 @@
 class BaseManager(Randomizable):
     """
     Managers should be created and registered after launching BaseEngine
-    TODO LQY, inherit event manager so that we can control the process in a standard way
     """
     PRIORITY = 10  # the engine will call managers according to the priority
 
@@ -64,6 +63,8 @@ def destroy(self):
         Destroy manager
         """
         # self.engine = None
+        super(BaseManager, self).destroy()
+        self.clear_objects(list(self.spawned_objects.keys()), force_destroy=True)
         self.spawned_objects = None
 
     def spawn_object(self, object_class, **kwargs):

metadrive/manager/pg_map_manager.py

@@ -39,8 +39,9 @@ def unload_map(self, map):
             map.destroy()
 
     def destroy(self):
-        self.maps = None
         super(PGMapManager, self).destroy()
+        self.clear_stored_maps()
+        self.maps = None
 
     def before_reset(self):
         # remove map from world before adding
@@ -131,12 +132,14 @@ def load_all_maps(self, file_name):
         self.reset()
         return loaded_map_data
 
-    def clear_objects(self, *args, **kwargs):
+    def clear_stored_maps(self):
         """
-        As Map instance should not be recycled, we will forcefully destroy useless map instances.
+        Clear all stored maps
         """
-        return super(PGMapManager, self).clear_objects(force_destroy=True, *args, **kwargs)
-
-
-# For compatibility
-MapManager = PGMapManager
+        for m in self.maps.values():
+            if m is not None:
+                m.detach_from_world()
+                m.destroy()
+        start_seed = self.start_seed = self.engine.global_config["start_seed"]
+        env_num = self.env_num = self.engine.global_config["num_scenarios"]
+        self.maps = {_seed: None for _seed in range(start_seed, start_seed + env_num)}

metadrive/manager/scenario_data_manager.py

@@ -1,6 +1,6 @@
 import copy
 import os
-
+from metadrive.utils.config import Config
 import numpy as np
 
 from metadrive.manager.base_manager import BaseManager
@@ -188,3 +188,14 @@ def current_scenario_file_name(self):
     @property
     def data_coverage(self):
         return sum(self.coverage) / len(self.coverage) * self.engine.global_config["num_workers"]
+
+    def destroy(self):
+        """
+        Clear memory
+        """
+        super(ScenarioDataManager, self).destroy()
+        self._scenarios = {}
+        Config.clear_nested_dict(self.summary_dict)
+        self.summary_lookup.clear()
+        self.mapping.clear()
+        self.summary_dict, self.summary_lookup, self.mapping = None, None, None

metadrive/manager/scenario_map_manager.py

@@ -101,7 +101,8 @@ def unload_map(self, map):
             assert len(self.spawned_objects) == 0
 
     def destroy(self):
-        self.maps = None
+        self.clear_stored_maps()
+        self._stored_maps = None
         self.current_map = None
 
         self.sdc_start_point = None
@@ -122,12 +123,6 @@ def before_reset(self):
         self.current_sdc_route = None
         self.current_map = None
 
-    def clear_objects(self, *args, **kwargs):
-        """
-        As Map instance should not be recycled, we will forcefully destroy useless map instances.
-        """
-        return super(ScenarioMapManager, self).clear_objects(force_destroy=True, *args, **kwargs)
-
     def clear_stored_maps(self):
         for m in self._stored_maps.values():
             if m is not None:

metadrive/obs/image_obs.py

@@ -38,6 +38,11 @@ def observation_space(self):
     def observe(self, vehicle: BaseVehicle):
         return {self.IMAGE: self.img_obs.observe(vehicle), self.STATE: self.state_obs.observe(vehicle)}
 
+    def destroy(self):
+        super(ImageStateObservation, self).destroy()
+        self.img_obs.destroy()
+        self.state_obs.destroy()
+
 
 class ImageObservation(BaseObservation):
     """
@@ -88,3 +93,10 @@ def reset(self, env, vehicle=None):
         self.state = np.zeros(self.observation_space.shape, dtype=np.float32)
         if self.enable_cuda:
             self.state = cp.asarray(self.state)
+
+    def destroy(self):
+        """
+        Clear memory
+        """
+        super(ImageObservation, self).destroy()
+        self.state = None

metadrive/obs/observation_base.py

@@ -3,6 +3,7 @@
 import gymnasium as gym
 from copy import deepcopy
 from metadrive.engine.logger import get_logger
+from metadrive.utils.config import Config
 
 logger = get_logger()
 
@@ -31,6 +32,14 @@ def observe(self, *args, **kwargs):
     def reset(self, env, vehicle=None):
         pass
 
+    def destroy(self):
+        """
+        Clear allocated memory
+        """
+        pass
+        # Config.clear_nested_dict(self.config)
+        # self.config = None
+
 
 class DummyObservation(BaseObservation):
     """

metadrive/obs/state_obs.py

@@ -160,6 +160,9 @@ def get_line_detector_dim(self):
 
 
 class LidarStateObservation(BaseObservation):
+    """
+    This observation uses lidar to detect moving objects
+    """
     def __init__(self, config):
         self.state_obs = StateObservation(config)
         super(LidarStateObservation, self).__init__(config)
@@ -239,3 +242,12 @@ def _add_noise_to_cloud_points(self, points, gaussian_noise, dropout_prob):
             points[np.random.uniform(0, 1, size=points.shape) < dropout_prob] = 0.0
 
         return list(points)
+
+    def destroy(self):
+        """
+        Clear allocated memory
+        """
+        self.state_obs.destroy()
+        super(LidarStateObservation, self).destroy()
+        self.cloud_points = None
+        self.detected_objects = None

metadrive/obs/top_down_obs.py

@@ -228,3 +228,8 @@ def observe(self, vehicle: BaseVehicle):
         else:
             img = img.astype(np.uint8)
         return np.transpose(img, (1, 0, 2))
+
+    def destroy(self):
+        # scene
+        self.road_network = None
+        super(TopDownObservation, self).destroy()

metadrive/policy/base_policy.py

@@ -54,12 +54,16 @@ def reset(self):
         self.action_info.clear()
 
     def destroy(self):
+        """
+        Destroy Base class
+        """
         if self._debug_mark is not None:
             self.engine.taskMgr.remove(self._mark_update_task_name)
             self._debug_mark.removeNode()
             self._debug_mark = None
             self._mark_update_task_name = None
-        super(BasePolicy, self).destroy()
+        Configurable.destroy(self)
+        Randomizable.destroy(self)
         self.control_object = None
         logging.debug("{} is released".format(self.__class__.__name__))