diff --git a/flatland/core/env_prediction_builder.py b/flatland/core/env_prediction_builder.py
index 756068efffa995968322bfecfefee178717a226b..e591d6cfae141aba474eb5424446515c12df89bb 100644
--- a/flatland/core/env_prediction_builder.py
+++ b/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.
 
diff --git a/flatland/envs/predictions.py b/flatland/envs/predictions.py
index 0b390f01c7d8ef876d3b30e58e0fd48f5aceecdf..77707b9f110376ddf2638b830830ff1a1c1edbf6 100644
--- a/flatland/envs/predictions.py
+++ b/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.