Refactoring: call constructor of BaseClass

flatland/core/env_prediction_builder.py

@@ -35,9 +35,6 @@ class PredictionBuilder:
 
         Parameters
         ----------
-        custom_args: dict
-            Implementation-dependent custom arguments, see the sub-classes.
-
         handle : int, optional
             Handle of the agent for which to compute the observation vector.
 

flatland/envs/predictions.py

@@ -24,8 +24,6 @@ class DummyPredictorForRailEnv(PredictionBuilder):
 
         Parameters
         ----------
-        custom_args: dict
-            Not used in this dummy implementation.
         handle : int, optional
             Handle of the agent for which to compute the observation vector.
 
@@ -87,9 +85,8 @@ class ShortestPathPredictorForRailEnv(PredictionBuilder):
     The prediction acts as if no other agent is in the environment and always takes the forward action.
     """
 
-    def __init__(self, max_depth=20):
-        # Initialize with depth 20
-        self.max_depth = max_depth
+    def __init__(self, max_depth: int = 20):
+        super().__init__(max_depth)
 
     def get(self, handle: int = None):
         """
@@ -98,8 +95,6 @@ class ShortestPathPredictorForRailEnv(PredictionBuilder):
 
         Parameters
         ----------
-        custom_args: dict
-            - distance_map : dict
         handle : int, optional
             Handle of the agent for which to compute the observation vector.