Merge pull request #12 from HumanCompatibleAI/black_jupyter

Make black work with jupyter

Author: PavelCz

requirements-dev.txt

@@ -1,5 +1,5 @@
 pytest==6.2.5
-black==22.10
+black[jupyter]==22.10
 flake8==3.9.2
 pytype==2021.8.24
 flake8-isort==4.0.0

src/notebooks/rl_util.ipynb

@@ -107,10 +107,15 @@
     "import numpy as np\n",
     "import torch as th\n",
     "from collections import OrderedDict\n",
+    "\n",
     "# from lucid.modelzoo.vision_base import Model\n",
     "# from lucid.scratch.rl_util import *\n",
     "# get_abbreviator defined during setup\n",
-    "from reward_preprocessing.ext.notebook_util import CategoricalPolicyGM, ImpalaModel, get_abbreviator\n",
+    "from reward_preprocessing.ext.notebook_util import (\n",
+    "    CategoricalPolicyGM,\n",
+    "    ImpalaModel,\n",
+    "    get_abbreviator,\n",
+    ")\n",
     "from reward_preprocessing.vis.reward_vis import LayerNMF\n",
     "from reward_preprocessing.vis.util import zoom_to"
    ],
@@ -157,14 +162,18 @@
     "model_state_dict = th.load(Path(model_path).expanduser(), device)[\"model_state_dict\"]\n",
     "\n",
     "embedder = ImpalaModel(in_channels=3)\n",
-    "model = CategoricalPolicyGM(embedder=embedder, action_size=model_state_dict[\"fc_policy.weight\"].shape[0])\n",
+    "model = CategoricalPolicyGM(\n",
+    "    embedder=embedder, action_size=model_state_dict[\"fc_policy.weight\"].shape[0]\n",
+    ")\n",
     "# Load data\n",
     "model.load_state_dict(model_state_dict)\n",
     "\n",
-    "value_function_name = 'fc_value'\n",
+    "value_function_name = \"fc_value\"\n",
     "\n",
     "# Load trajectories for dataset visualization.\n",
-    "trajectories = demonstrations.load_expert_trajs(str(Path(trajectories_path).expanduser()), n_expert_demos=None)\n",
+    "trajectories = demonstrations.load_expert_trajs(\n",
+    "    str(Path(trajectories_path).expanduser()), n_expert_demos=None\n",
+    ")\n",
     "trajectories = flatten_trajectories(trajectories)\n",
     "\n",
     "# Get observations from trajectories.\n",
@@ -174,9 +183,7 @@
     "\n",
     "layer_names = get_model_layers(model)\n",
     "abbreviator = get_abbreviator(layer_names)\n",
-    "layer_names = OrderedDict(\n",
-    "    [(name[abbreviator], name) for name in layer_names]\n",
-    ")\n"
+    "layer_names = OrderedDict([(name[abbreviator], name) for name in layer_names])"
    ],
    "execution_count": 2,
    "outputs": []
@@ -399,11 +406,17 @@
    },
    "source": [
     "model.eval()\n",
-    "layer = 'embedder_relu_after_convs'\n",
+    "layer = \"embedder_relu_after_convs\"\n",
     "# value_function_name = None\n",
     "# can take a couple of minutes\n",
     "# for the paper, we use observations[:], but this requires more memory\n",
-    "nmf = LayerNMF(model, layer, observations[:1024], features=None, attr_layer_name=value_function_name)"
+    "nmf = LayerNMF(\n",
+    "    model,\n",
+    "    layer,\n",
+    "    observations[:1024],\n",
+    "    features=None,\n",
+    "    attr_layer_name=value_function_name,\n",
+    ")"
    ],
    "execution_count": 12,
    "outputs": []
@@ -579,7 +592,15 @@
    },
    "source": [
     "# Show expects channels last, unlike the rest of lucent. Therefore we need to transpose here.\n",
-    "show([zoom_to(nmf.vis_dataset_thumbnail(i, num_mult=4, expand_mult=4, max_rep=np.inf)[0], 200).transpose(1,2,0) for i in range(nmf.features)])"
+    "show(\n",
+    "    [\n",
+    "        zoom_to(\n",
+    "            nmf.vis_dataset_thumbnail(i, num_mult=4, expand_mult=4, max_rep=np.inf)[0],\n",
+    "            200,\n",
+    "        ).transpose(1, 2, 0)\n",
+    "        for i in range(nmf.features)\n",
+    "    ]\n",
+    ")"
    ],
    "execution_count": 91,
    "outputs": [
@@ -784,8 +805,8 @@
     }
    },
    "source": [
-    "traj = trajectories['observations'][0][76:84]\n",
-    "attr = get_attr(model, value_function_name, layer_names['2b'], traj, integrate_steps=10)\n",
+    "traj = trajectories[\"observations\"][0][76:84]\n",
+    "attr = get_attr(model, value_function_name, layer_names[\"2b\"], traj, integrate_steps=10)\n",
     "attr.shape"
    ],
    "execution_count": 16,
@@ -840,10 +861,16 @@
     }
    },
    "source": [
-    "attr_reduced = nmf.transform(np.maximum(attr, 0)) - nmf.transform(np.maximum(-attr, 0))  # transform the positive and negative parts separately\n",
+    "attr_reduced = nmf.transform(np.maximum(attr, 0)) - nmf.transform(\n",
+    "    np.maximum(-attr, 0)\n",
+    ")  # transform the positive and negative parts separately\n",
     "nmf_norms = nmf.channel_dirs.sum(-1)\n",
-    "attr_reduced *= nmf_norms[None, None, None]  # multiply by the norms of the NMF directions, since the magnitudes of the NMF directions are not relevant\n",
-    "attr_reduced /= np.median(attr_reduced.max(axis=(-3, -2, -1)))  # globally normalize by the median max value to make the visualization balanced (a bit of a hack)\n",
+    "attr_reduced *= nmf_norms[\n",
+    "    None, None, None\n",
+    "]  # multiply by the norms of the NMF directions, since the magnitudes of the NMF directions are not relevant\n",
+    "attr_reduced /= np.median(\n",
+    "    attr_reduced.max(axis=(-3, -2, -1))\n",
+    ")  # globally normalize by the median max value to make the visualization balanced (a bit of a hack)\n",
     "attr_reduced.shape"
    ],
    "execution_count": 17,
@@ -1315,12 +1342,18 @@
     }
    },
    "source": [
-    "kernel_name = layer_names[\"3a\"].replace(\"Relu\", \"conv2d/kernel\")  # name of tensor of convolutional kernel of next layer\n",
+    "kernel_name = layer_names[\"3a\"].replace(\n",
+    "    \"Relu\", \"conv2d/kernel\"\n",
+    ")  # name of tensor of convolutional kernel of next layer\n",
     "kernel = editor[kernel_name]\n",
-    "saw_dir = nmf.channel_dirs[0][None, None, :, None]  # first NMF direction, corresponding to saw obstacle\n",
+    "saw_dir = nmf.channel_dirs[0][\n",
+    "    None, None, :, None\n",
+    "]  # first NMF direction, corresponding to saw obstacle\n",
     "saw_dir /= np.linalg.norm(saw_dir)\n",
     "# the kernel is left-multiplied by the activations from the previous layer, so we left-multiply the kernel by the projection matrix\n",
-    "kernel = kernel - saw_dir * (saw_dir * kernel).sum(axis=-2, keepdims=True)  # equivalently: kernel - saw_dir @ saw_dir.transpose((0, 1, 3, 2)) @ kernel\n",
+    "kernel = kernel - saw_dir * (saw_dir * kernel).sum(\n",
+    "    axis=-2, keepdims=True\n",
+    ")  # equivalently: kernel - saw_dir @ saw_dir.transpose((0, 1, 3, 2)) @ kernel\n",
     "editor[kernel_name] = kernel\n",
     "# note: this is not quite the same as the edit made for the paper, since we only used 1024 observations for the NMF calculation here"
    ],
@@ -1415,8 +1448,10 @@
     }
    },
    "source": [
-    "traj = trajectories['observations'][0][76:84]\n",
-    "attr = get_attr(edited_model, value_function_name, layer_names['2b'], traj, integrate_steps=10)\n",
+    "traj = trajectories[\"observations\"][0][76:84]\n",
+    "attr = get_attr(\n",
+    "    edited_model, value_function_name, layer_names[\"2b\"], traj, integrate_steps=10\n",
+    ")\n",
     "attr_reduced = nmf.transform(np.maximum(attr, 0)) - nmf.transform(np.maximum(-attr, 0))\n",
     "nmf_norms = nmf.channel_dirs.sum(-1)\n",
     "attr_reduced *= nmf_norms[None, None, None]\n",