occlusion manager changes

nuplan/planning/script/builders/occlusion_manager_builder.py

@@ -2,7 +2,7 @@
 from omegaconf import DictConfig
 
 from nuplan.planning.scenario_builder.abstract_scenario import AbstractScenario
-from nuplan.planning.simulation.occlusion.occlusion_manager import AbstractOcclusionManager
+from nuplan.planning.simulation.occlusion.abstract_occlusion_manager import AbstractOcclusionManager
 from nuplan.planning.simulation.occlusion.range_occlusion_manager import RangeOcclusionManager
 
 def build_occlusion_manager(occlusion_cfg: DictConfig, scenario: AbstractScenario) -> AbstractOcclusionManager:

nuplan/planning/script/builders/simulation_builder.py

@@ -17,7 +17,7 @@ from nuplan.planning.simulation.callback.metric_callback import MetricCallback
 from nuplan.planning.simulation.callback.multi_callback import MultiCallback
 from nuplan.planning.simulation.controller.abstract_controller import AbstractEgoController
 from nuplan.planning.simulation.observation.abstract_observation import AbstractObservation
-from nuplan.planning.simulation.occlusion.occlusion_manager import AbstractOcclusionManager
+from nuplan.planning.simulation.occlusion.abstract_occlusion_manager import AbstractOcclusionManager
 from nuplan.planning.simulation.planner.abstract_planner import AbstractPlanner
 from nuplan.planning.simulation.runner.simulations_runner import SimulationRunner
 from nuplan.planning.simulation.simulation import Simulation

nuplan/planning/simulation/occlusion/occlusion_manager.py → nuplan/planning/simulation/occlusion/abstract_occlusion_manager.py

@@ -11,24 +11,38 @@ from nuplan.planning.simulation.observation.observation_type import DetectionsTr
 
 
 class AbstractOcclusionManager(metaclass=ABCMeta):
+    """
+    Interface for a generic occlusion manager.
+    """
+
     def __init__(
         self,
         scenario: AbstractScenario
     ):
-        self._masks = {}
+        self._visible_agent_cache = {}
         self.scenario = scenario
 
     def reset(self) -> None:
-        self._masks = {}
+        """
+        Resets occlusion manager cache.
+        """
+        self._visible_agent_cache = {}
 
     def occlude_input(self, input_buffer: SimulationHistoryBuffer) -> SimulationHistoryBuffer:
+        """
+        Occludes SimulationHistoryBuffer input. Loops through each timestep defined by time_us,
+        checks to see if timestep is already contained in _visible_agent_cache and computes
+        occlusions if not, and occludes timestep using cached results. Repacks output in 
+        SimulationHistoryBuffer.
+        """
+
         ego_state_buffer = input_buffer.ego_state_buffer
         observations_buffer = input_buffer.observation_buffer
         sample_interval = input_buffer.sample_interval
 
         for ego_state, observations in zip(ego_state_buffer, observations_buffer):
-            if ego_state.time_us not in self._masks:
-                self._masks[ego_state.time_us] = self._compute_mask(ego_state, observations)
+            if ego_state.time_us not in self._visible_agent_cache:
+                self._visible_agent_cache[ego_state.time_us] = self._compute_mask(ego_state, observations)
                 
         output_buffer = SimulationHistoryBuffer(ego_state_buffer, \
                             deque([self._mask_input(ego_state.time_us, observations) for ego_state, observations in zip(ego_state_buffer, observations_buffer)]), \
@@ -41,14 +55,16 @@ class AbstractOcclusionManager(metaclass=ABCMeta):
         pass
 
     def _mask_input(self, time_us: int, observations: DetectionsTracks) -> DetectionsTracks:
-        assert time_us in self._masks, "Attempted to mask non-cached timestep!"
+        """
+        Occludes observations at timestep time_us based on cached occlusions.
+        """
+
+        assert time_us in self._visible_agent_cache, "Attempted to mask non-cached timestep!"
         assert isinstance(observations, DetectionsTracks), "Occlusions only support DetectionsTracks."
 
-        mask = self._masks[time_us]
+        mask = self._visible_agent_cache[time_us]
         tracks = observations.tracked_objects.tracked_objects
 
         visible_tracks = [track for track in tracks if track.metadata.track_token in mask]
 
         return DetectionsTracks(tracked_objects=TrackedObjects(visible_tracks))
-
-

nuplan/planning/simulation/occlusion/range_occlusion_manager.py

 from nuplan.common.actor_state.ego_state import EgoState
-from nuplan.planning.simulation.observation.observation_type import DetectionsTracks, Observation
-from nuplan.planning.simulation.occlusion.occlusion_manager import AbstractOcclusionManager
+from nuplan.planning.scenario_builder.abstract_scenario import AbstractScenario
+from nuplan.planning.simulation.observation.observation_type import DetectionsTracks
+from nuplan.planning.simulation.occlusion.abstract_occlusion_manager import AbstractOcclusionManager
 
 
 class RangeOcclusionManager(AbstractOcclusionManager):
+    """
+    Range occlusion manager. Occludes all objects outside of a given
+    range of the ego.
+    """
+
+    def __init__(
+        self,
+        scenario: AbstractScenario,
+        range_threshold: float = 25
+    ):
+        super().__init__(scenario)
+        self.range_threshold = range_threshold
+
     def _compute_mask(self, ego_state: EgoState, observations: DetectionsTracks) -> set:
         not_occluded = set()
         for track in observations.tracked_objects.tracked_objects:
-            if (ego_state.center.x - track.center.x) ** 2 + (ego_state.center.y - track.center.y) ** 2 <= 500:
+            if ((ego_state.center.x - track.center.x) ** 2 + \
+                (ego_state.center.y - track.center.y) ** 2) ** 0.5 <= self.range_threshold:
                 not_occluded.add(track.metadata.track_token)
 
-        return not_occluded
\ No newline at end of file
+        return not_occluded
+    
\ No newline at end of file

nuplan/planning/simulation/simulation_setup.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 from nuplan.planning.scenario_builder.abstract_scenario import AbstractScenario
 from nuplan.planning.simulation.controller.abstract_controller import AbstractEgoController
 from nuplan.planning.simulation.observation.abstract_observation import AbstractObservation
-from nuplan.planning.simulation.occlusion.occlusion_manager import AbstractOcclusionManager
+from nuplan.planning.simulation.occlusion.abstract_occlusion_manager import AbstractOcclusionManager
 from nuplan.planning.simulation.planner.abstract_planner import AbstractPlanner
 from nuplan.planning.simulation.simulation_time_controller.abstract_simulation_time_controller import (
     AbstractSimulationTimeController,