diff --git a/docs/getting_started/01_study.ipynb b/docs/getting_started/01_study.ipynb
index 5cf41b8..b76699d 100644
--- a/docs/getting_started/01_study.ipynb
+++ b/docs/getting_started/01_study.ipynb
@@ -155,7 +155,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 14,
+      "execution_count": 8,
       "id": "0df1511e-565d-45d0-93a8-adafbfaaaefa",
       "metadata": {
         "scrolled": true
@@ -163,41 +163,264 @@
       "outputs": [
         {
           "data": {
-            "application/vnd.jupyter.widget-view+json": {
-              "model_id": "bd157ce6c3034fd5a0533f8220e44c9e",
-              "version_major": 2,
-              "version_minor": 0
-            },
+            "text/html": [
+              "\n",
+              "          <iframe\n",
+              "                id=\"studysFvnWfhilS\"\n",
+              "                width=\"820\"\n",
+              "                height=\"620\"\n",
+              "                frameborder=\"0\"\n",
+              "                srcdoc=\"\n",
+              "&lt;!DOCTYPE html&gt;\n",
+              "&lt;html lang=&quot;en&quot;&gt;\n",
+              "  &lt;head&gt;\n",
+              "    &lt;title&gt;Network&lt;/title&gt;\n",
+              "    &lt;script\n",
+              "      type=&quot;text/javascript&quot;\n",
+              "      src=&quot;https://unpkg.com/vis-network/standalone/umd/vis-network.min.js&quot;\n",
+              "    &gt;&lt;/script&gt;\n",
+              "    &lt;style type=&quot;text/css&quot;&gt;\n",
+              "      #mynetwork {\n",
+              "        \n",
+              "        width: 800px;\n",
+              "        height: 600px;\n",
+              "        border: 1px solid lightgray;\n",
+              "        \n",
+              "      }\n",
+              "    &lt;/style&gt;\n",
+              "  &lt;/head&gt;\n",
+              "  &lt;body&gt;\n",
+              "    &lt;div id=&quot;mynetwork&quot;&gt;&lt;/div&gt;\n",
+              "    &lt;script type=&quot;text/javascript&quot;&gt;\n",
+              "        \n",
+              "        var nodes = new vis.DataSet([{ id: 0, label: `init`, group: &quot;states&quot;, x: -27, y: -203 },\n",
+              "{ id: 1, label: `studied\n",
+              "\u20ac\tR: 0`, group: &quot;states&quot;, x: -84, y: -6 },\n",
+              "{ id: 2, label: `didn&#x27;t study\n",
+              "\u20ac\tR: 15`, group: &quot;states&quot;, x: 57, y: -10 },\n",
+              "{ id: 3, label: `pass test\n",
+              "\u20ac\tR: 100`, group: &quot;states&quot;, x: -56, y: 102 },\n",
+              "{ id: 4, label: `fail test\n",
+              "\u20ac\tR: 0`, group: &quot;states&quot;, x: 58, y: 61 },\n",
+              "{ id: 5, label: `end\n",
+              "\u20ac\tR: 0`, group: &quot;states&quot;, x: 24, y: 187 },\n",
+              "{ id: 10000000000, label: `study\n",
+              "\u20ac\tR: 0`, group: &quot;actions&quot;, x: -136, y: -127 },\n",
+              "{ id: 10000000001, label: `don&#x27;t study\n",
+              "\u20ac\tR: 0`, group: &quot;actions&quot;, x: 77, y: -144 },\n",
+              "]);\n",
+              "        var edges = new vis.DataSet([{ from: 0, to: 10000000000 },\n",
+              "{ from: 10000000000, to: 1, label: &quot;1&quot; },\n",
+              "{ from: 0, to: 10000000001 },\n",
+              "{ from: 10000000001, to: 2, label: &quot;1&quot; },\n",
+              "{ from: 1, to: 3, label: &quot;9/10&quot; },\n",
+              "{ from: 1, to: 4, label: &quot;1/10&quot; },\n",
+              "{ from: 2, to: 3, label: &quot;2/5&quot; },\n",
+              "{ from: 2, to: 4, label: &quot;3/5&quot; },\n",
+              "{ from: 3, to: 5, label: &quot;1&quot; },\n",
+              "{ from: 4, to: 5, label: &quot;1&quot; },\n",
+              "]);\n",
+              "        var options = {\n",
+              "  &quot;__fake_macros&quot;: {\n",
+              "    &quot;__group_macro&quot;: {\n",
+              "      &quot;borderWidth&quot;: 1,\n",
+              "      &quot;color&quot;: {\n",
+              "        &quot;background&quot;: &quot;white&quot;,\n",
+              "        &quot;border&quot;: &quot;black&quot;,\n",
+              "        &quot;highlight&quot;: {\n",
+              "          &quot;background&quot;: &quot;white&quot;,\n",
+              "          &quot;border&quot;: &quot;red&quot;\n",
+              "        }\n",
+              "      },\n",
+              "      &quot;shape&quot;: &quot;ellipse&quot;,\n",
+              "      &quot;mass&quot;: 1,\n",
+              "      &quot;font&quot;: {\n",
+              "        &quot;color&quot;: &quot;black&quot;,\n",
+              "        &quot;size&quot;: 14\n",
+              "      }\n",
+              "    }\n",
+              "  },\n",
+              "  &quot;groups&quot;: {\n",
+              "    &quot;states&quot;: {\n",
+              "      &quot;borderWidth&quot;: 1,\n",
+              "      &quot;color&quot;: {\n",
+              "        &quot;background&quot;: &quot;#8ff0a4&quot;,\n",
+              "        &quot;border&quot;: &quot;black&quot;,\n",
+              "        &quot;highlight&quot;: {\n",
+              "          &quot;background&quot;: &quot;white&quot;,\n",
+              "          &quot;border&quot;: &quot;red&quot;\n",
+              "        }\n",
+              "      },\n",
+              "      &quot;shape&quot;: &quot;ellipse&quot;,\n",
+              "      &quot;mass&quot;: 1,\n",
+              "      &quot;font&quot;: {\n",
+              "        &quot;color&quot;: &quot;black&quot;,\n",
+              "        &quot;size&quot;: 14\n",
+              "      }\n",
+              "    },\n",
+              "    &quot;actions&quot;: {\n",
+              "      &quot;borderWidth&quot;: 1,\n",
+              "      &quot;color&quot;: {\n",
+              "        &quot;background&quot;: &quot;#ffc0cb&quot;,\n",
+              "        &quot;border&quot;: &quot;black&quot;,\n",
+              "        &quot;highlight&quot;: {\n",
+              "          &quot;background&quot;: &quot;white&quot;,\n",
+              "          &quot;border&quot;: &quot;red&quot;\n",
+              "        }\n",
+              "      },\n",
+              "      &quot;shape&quot;: &quot;box&quot;,\n",
+              "      &quot;mass&quot;: 1,\n",
+              "      &quot;font&quot;: {\n",
+              "        &quot;color&quot;: &quot;black&quot;,\n",
+              "        &quot;size&quot;: 14\n",
+              "      }\n",
+              "    },\n",
+              "    &quot;scheduled_actions&quot;: {\n",
+              "      &quot;borderWidth&quot;: 1,\n",
+              "      &quot;color&quot;: {\n",
+              "        &quot;background&quot;: &quot;pink&quot;,\n",
+              "        &quot;border&quot;: &quot;black&quot;,\n",
+              "        &quot;highlight&quot;: {\n",
+              "          &quot;background&quot;: &quot;white&quot;,\n",
+              "          &quot;border&quot;: &quot;red&quot;\n",
+              "        }\n",
+              "      },\n",
+              "      &quot;shape&quot;: &quot;box&quot;,\n",
+              "      &quot;mass&quot;: 1,\n",
+              "      &quot;font&quot;: {\n",
+              "        &quot;color&quot;: &quot;black&quot;,\n",
+              "        &quot;size&quot;: 14\n",
+              "      },\n",
+              "      &quot;schedColor&quot;: false\n",
+              "    }\n",
+              "  },\n",
+              "  &quot;reload_button&quot;: false,\n",
+              "  &quot;edges&quot;: {\n",
+              "    &quot;arrows&quot;: &quot;to&quot;,\n",
+              "    &quot;font&quot;: {\n",
+              "      &quot;color&quot;: &quot;black&quot;,\n",
+              "      &quot;size&quot;: 14\n",
+              "    },\n",
+              "    &quot;color&quot;: {\n",
+              "      &quot;color&quot;: &quot;black&quot;\n",
+              "    }\n",
+              "  },\n",
+              "  &quot;numbers&quot;: {\n",
+              "    &quot;fractions&quot;: true,\n",
+              "    &quot;digits&quot;: 5\n",
+              "  },\n",
+              "  &quot;state_properties&quot;: {\n",
+              "    &quot;show_results&quot;: true,\n",
+              "    &quot;result_symbol&quot;: &quot;\\u2606&quot;,\n",
+              "    &quot;show_rewards&quot;: true,\n",
+              "    &quot;reward_symbol&quot;: &quot;\\u20ac&quot;,\n",
+              "    &quot;show_zero_rewards&quot;: true,\n",
+              "    &quot;show_observations&quot;: true,\n",
+              "    &quot;observation_symbol&quot;: &quot;\\u0298&quot;\n",
+              "  },\n",
+              "  &quot;layout&quot;: {\n",
+              "    &quot;randomSeed&quot;: 5\n",
+              "  },\n",
+              "  &quot;misc&quot;: {\n",
+              "    &quot;enable_physics&quot;: true,\n",
+              "    &quot;width&quot;: 800,\n",
+              "    &quot;height&quot;: 600,\n",
+              "    &quot;explore&quot;: false\n",
+              "  },\n",
+              "  &quot;saving&quot;: {\n",
+              "    &quot;relative_path&quot;: true,\n",
+              "    &quot;filename&quot;: &quot;layouts/pinkgreen.json&quot;,\n",
+              "    &quot;save_button&quot;: false,\n",
+              "    &quot;load_button&quot;: false\n",
+              "  },\n",
+              "  &quot;positions&quot;: {\n",
+              "    &quot;0&quot;: {\n",
+              "      &quot;x&quot;: -27,\n",
+              "      &quot;y&quot;: -203\n",
+              "    },\n",
+              "    &quot;1&quot;: {\n",
+              "      &quot;x&quot;: -84,\n",
+              "      &quot;y&quot;: -6\n",
+              "    },\n",
+              "    &quot;2&quot;: {\n",
+              "      &quot;x&quot;: 57,\n",
+              "      &quot;y&quot;: -10\n",
+              "    },\n",
+              "    &quot;3&quot;: {\n",
+              "      &quot;x&quot;: -56,\n",
+              "      &quot;y&quot;: 102\n",
+              "    },\n",
+              "    &quot;4&quot;: {\n",
+              "      &quot;x&quot;: 58,\n",
+              "      &quot;y&quot;: 61\n",
+              "    },\n",
+              "    &quot;5&quot;: {\n",
+              "      &quot;x&quot;: 24,\n",
+              "      &quot;y&quot;: 187\n",
+              "    },\n",
+              "    &quot;10000000000&quot;: {\n",
+              "      &quot;x&quot;: -136,\n",
+              "      &quot;y&quot;: -127\n",
+              "    },\n",
+              "    &quot;10000000001&quot;: {\n",
+              "      &quot;x&quot;: 77,\n",
+              "      &quot;y&quot;: -144\n",
+              "    }\n",
+              "  },\n",
+              "  &quot;width&quot;: 800,\n",
+              "  &quot;height&quot;: 600,\n",
+              "  &quot;physics&quot;: true,\n",
+              "  &quot;results_and_rewards&quot;: {\n",
+              "    &quot;show_results&quot;: true,\n",
+              "    &quot;resultSymbol&quot;: &quot;\\u2606&quot;,\n",
+              "    &quot;show_rewards&quot;: true\n",
+              "  }\n",
+              "};\n",
+              "        //js\n",
+              "var container = document.getElementById(&quot;mynetwork&quot;);\n",
+              "var data = {\n",
+              "    nodes: nodes,\n",
+              "    edges: edges,\n",
+              "};\n",
+              "var network = new vis.Network(container, data, options);\n",
+              "function makeAllNodesInvisible() {\n",
+              "    ids = nodes.getIds();\n",
+              "    for (let i = 0; i &lt; ids.length; i++) {\n",
+              "        var nodeId = ids[i];\n",
+              "        var node = nodes.get(nodeId);\n",
+              "        node[&quot;hidden&quot;] = true;\n",
+              "        nodes.update(node);\n",
+              "    }\n",
+              "};\n",
+              "function makeNeighborsVisible(myNode) {\n",
+              "    var ids = network.getConnectedNodes(myNode, &#x27;to&#x27;);\n",
+              "    for (let i = 0; i &lt; ids.length; i++) {\n",
+              "        var nodeId = ids[i];\n",
+              "        var node = nodes.get(nodeId);\n",
+              "        node[&quot;hidden&quot;] = false;\n",
+              "        nodes.update(node);\n",
+              "    }\n",
+              "};\n",
+              "function makeNodeVisible(nodeId) {\n",
+              "    var node = nodes.get(nodeId);\n",
+              "    node[&quot;hidden&quot;] = false;\n",
+              "    nodes.update(node);\n",
+              "};\n",
+              "network.on( &#x27;click&#x27;, function(properties) {\n",
+              "    var nodeId = network.getNodeAt({x:properties.event.srcEvent.offsetX, y:properties.event.srcEvent.offsetY});\n",
+              "    makeNeighborsVisible(nodeId);\n",
+              "});\n",
+              "\n",
+              "    &lt;/script&gt;\n",
+              "  &lt;/body&gt;\n",
+              "&lt;/html&gt;\n",
+              "\"\n",
+              "                border:none !important;\n",
+              "                allowfullscreen webkitallowfullscreen mozallowfullscreen\n",
+              "          ></iframe>"
+            ],
             "text/plain": [
-              "Output()"
-            ]
-          },
-          "metadata": {},
-          "output_type": "display_data"
-        },
-        {
-          "data": {
-            "application/vnd.jupyter.widget-view+json": {
-              "model_id": "f7080afa9cce47deb47ef9c1b145ce71",
-              "version_major": 2,
-              "version_minor": 0
-            },
-            "text/plain": [
-              "Output()"
-            ]
-          },
-          "metadata": {},
-          "output_type": "display_data"
-        },
-        {
-          "data": {
-            "application/vnd.jupyter.widget-view+json": {
-              "model_id": "69b3da5b0e884431b9cb84d758067af4",
-              "version_major": 2,
-              "version_minor": 0
-            },
-            "text/plain": [
-              "HBox(children=(Output(), Output()))"
+              "<IPython.core.display.HTML object>"
             ]
           },
           "metadata": {},
@@ -208,10 +431,29 @@
         "vis = show(mdp, layout=Layout(\"layouts/pinkgreen.json\"), name=\"study\", show_editor=False)"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": 10,
+      "id": "6cb412a1-52fa-4663-b7ba-371e06001f7a",
+      "metadata": {},
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "{0: {Edge(from_=0, to_=10000000001, label=None)}, 10000000000: set(), 10000000001: set(), 1: {Edge(from_=1, to_=4, label='1/10')}, 2: {Edge(from_=2, to_=4, label='3/5')}, 3: set(), 4: set()}\n"
+          ]
+        }
+      ],
+      "source": [
+        "\n",
+        "print(vis.nt.edges)"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
-      "id": "a9e68b97-49d8-47dc-8ba1-5d179bcecc52",
+      "id": "f1a87851-75b5-4e49-91ab-9382f711e495",
       "metadata": {},
       "outputs": [],
       "source": []
diff --git a/docs/getting_started/study.html b/docs/getting_started/study.html
index e7edec5..2e637d8 100644
--- a/docs/getting_started/study.html
+++ b/docs/getting_started/study.html
@@ -1,6 +1,6 @@
 
           <iframe
-                id="studyNsXSOjuhzk"
+                id="studysFvnWfhilS"
                 width="820"
                 height="620"
                 frameborder="0"
diff --git a/stormvogel/visjs.py b/stormvogel/visjs.py
index a857886..56c1677 100644
--- a/stormvogel/visjs.py
+++ b/stormvogel/visjs.py
@@ -65,7 +65,6 @@ def __init__(
         self.nodes_js: str = ""
         self.edges_js: str = ""
         self.options_js: str = "{}"
-        self.new_nodes_hidden: bool = False
         if do_init_server:
             self.server: stormvogel.communication_server.CommunicationServer = (
                 stormvogel.communication_server.initialize_server()
@@ -75,23 +74,14 @@ def __init__(
             self.positions = {}
         else:
             self.positions = positions
-        self.nodes: dict[int, Node]
+        self.nodes: dict[int, Node] = {}
         self.edges: dict[int, set[Edge]] = {}
-        # Note that this refers to the same server as the global variable in stormvogel.communication_server.
-
-    def enable_exploration_mode(self, initial_node_id: int):
-        """Every node becomes invisible. You can then click any node to reveal all of its successors. Call before adding any nodes to the network."""
-        self.new_nodes_hidden = True
-        self.initial_node_id = initial_node_id
-
-    def update_exploration_mode(self, initial_node_id: int):
-        # Make all nodes invisible.
-        ipd.display(ipd.Javascript(self.content_window + ".makeAllNodesInvisible()"))
-        # Make the initial state visible.
-        ipd.display(
-            ipd.Javascript(self.content_window + f".makeNodeVisible({initial_node_id})")
-        )
-        # All future nodes to be added will be hidden as well.
+
+    def enable_exploration_mode(self):
+        """Make the initial node visible, and inject the JS that enables exploration by clicking."""
+
+    def make_node_visible(self, id):
+        """Make the node with this id visible."""
 
     def get_positions(self) -> dict:
         """Get the current positions of the nodes on the canvas. Returns empty dict if unsucessful.
@@ -123,8 +113,6 @@ def __add_node_pre(self, node: Node) -> None:
             current += f', group: "{node.group}"'
         if self.positions is not None and str(node.id) in self.positions:
             current += f', x: {self.positions[str(node.id)]["x"]}, y: {self.positions[str(node.id)]["y"]}'
-        if self.new_nodes_hidden and node.id != self.initial_node_id:
-            current += ", hidden: true"
         current += " },\n"
         self.nodes_js += current
 
@@ -157,7 +145,10 @@ def add_edge(
     ) -> None:
         """Add an edge. Only use before calling show."""
         edge = Edge(from_, to_, label)
-        self.edges[from_].add(edge)
+        if from_ in self.edges:
+            self.edges[from_].add(edge)
+        else:
+            self.edges[from_] = set()
         if self.nodes[from_].visible and self.nodes[from_].visible:
             self.add_edge_pre(edge)
 
diff --git a/stormvogel/visualization.py b/stormvogel/visualization.py
index 42f24e5..49e3938 100644
--- a/stormvogel/visualization.py
+++ b/stormvogel/visualization.py
@@ -105,7 +105,8 @@ def show(self) -> None:
             ipd.clear_output()
         self.__create_nt()
         if self.layout.layout["misc"]["explore"]:
-            self.nt.enable_exploration_mode(self.model.get_initial_state().id)
+            self.nt.enable_exploration_mode()
+            self.nt.make_node_visible(self.model.get_initial_state().id)
         self.layout.set_groups(self.separate_labels)
         self.__add_states()
         self.__add_transitions()
@@ -143,6 +144,7 @@ def __add_states(self) -> None:
                 state.id,
                 label=",".join(state.labels) + rewards + res + observations,
                 group=group,
+                visible=not self.layout.layout["misc"]["explore"],
             )
 
     def __add_transitions(self) -> None: