Clarify types

Author: PavelCz

src/reward_preprocessing/interpret.py

@@ -3,6 +3,7 @@
 
 from PIL import Image
 from imitation.scripts.common import common as common_config
+from imitation.util.logger import HierarchicalLogger
 from lucent.modelzoo.util import get_model_layers
 from lucent.optvis import transform
 import matplotlib
@@ -221,15 +222,15 @@ def interpret(
 
 
 def plot_img(
-    columns,
-    custom_logger,
-    feature_i,
-    fig: Optional,
-    img,
-    pyplot,
-    rows,
-    vis_scale,
-    wandb_logging,
+    columns: int,
+    custom_logger: HierarchicalLogger,
+    feature_i: int,
+    fig: Optional[matplotlib.figure.Figure],
+    img: np.ndarray,
+    pyplot: bool,
+    rows: int,
+    vis_scale: int,
+    wandb_logging: bool,
 ):
     """Plot the passed image to pyplot and wandb as appropriate."""
     _wandb_log(custom_logger, feature_i, img, vis_scale, wandb_logging)
@@ -239,7 +240,11 @@ def plot_img(
 
 
 def _wandb_log(
-    custom_logger, feature_i: int, img: np.ndarray, vis_scale: int, wandb_logging: bool
+    custom_logger: HierarchicalLogger,
+    feature_i: int,
+    img: np.ndarray,
+    vis_scale: int,
+    wandb_logging: bool,
 ):
     """Plot to wandb if wandb logging is enabled."""
     if wandb_logging: