Merge branch 'try-feature-vis' into extend-interpret

Author: PavelCz

src/reward_preprocessing/common/utils.py

@@ -258,14 +258,13 @@ def log_img_wandb(
         pil_img = array_to_image(img, scale)
     elif isinstance(img, PIL.Image.Image):
         pil_img = img.resize(
-            # PIL expects tuple of (width, height), numpy's index 1 is width, 0 height.
+            # PIL expects tuple of (width, height), as opposed to numpy's
+            # (height, width).
             size=(img.width * scale, img.height * scale),
             resample=Image.NEAREST,
         )
     else:
-        raise ValueError(
-            f"img must be np.ndarray or PIL.Image.Image, {type(img)=}"
-        )
+        raise ValueError(f"img must be np.ndarray or PIL.Image.Image, {type(img)=}")
     wb_img = wandb.Image(pil_img, caption=caption)
     logger.record(wandb_key, wb_img)
     if step is not None:
@@ -275,7 +274,8 @@ def log_img_wandb(
 def array_to_image(arr: np.ndarray, scale: int) -> PIL.Image.Image:
     """Take numpy array on [0,1] scale, return PIL image."""
     return Image.fromarray(np.uint8(arr * 255), mode="RGB").resize(
-        # PIL expects tuple of (width, height), numpy's index 1 is width, 0 height.
+        # PIL expects tuple of (width, height), numpy's dimension 1 is width, and
+        # dimension 0 height.
         size=(arr.shape[1] * scale, arr.shape[0] * scale),
         resample=Image.NEAREST,
     )

src/reward_preprocessing/interpret.py

@@ -62,12 +62,13 @@ def _determine_features_are_actions(nmf: LayerNMF, layer_name: str) -> bool:
     human-understandable action name instead of the feature index."""
     # This is the heuristic for determining whether features are actions:
     # - If there is no dim reduction
-    # - If it is one of the layers from the list
+    # - If the name of the layer suggest that we are analysing the final layer of a
+    #   reward net
     # - If the number of features is 15 since that is the number of actions in all
     #   procgen games
     return (
         nmf.channel_dirs.shape[0] == nmf.channel_dirs.shape[1]
-        and layer_name in ["rew_net_cnn_dense_final"]
+        and layer_name.endswith("dense_final")
         and nmf.channel_dirs.shape[0] == 15
     )
 

src/reward_preprocessing/vis/attribution.py

@@ -15,11 +15,16 @@ def get_activations(
     # size. Generally the number of inputs would be in the thousands, so we don't want
     # to run the entire batch through the model at once.
     batch_size = 128
+    t_acts = []
     for i in range(0, len(model_inputs), batch_size):
         model(model_inputs[i : i + batch_size])
 
-    # Get activations at layer.
-    t_acts = hook(layer_name)
+        # Get activations at layer.
+        act_batch = hook(layer_name)
+        t_acts.append(act_batch)
+
+    t_acts = th.cat(t_acts, dim=0)
+    assert t_acts.shape[0] == len(model_inputs)
 
     # Reward activations might be 2 dimensional (scalar + batch dimension) e.g.
     # for linear layers. In this case we unsqueeze.