diff --git a/notebooks/direct_steering.ipynb b/notebooks/direct_steering.ipynb
index b987782c0..19b02eefc 100644
--- a/notebooks/direct_steering.ipynb
+++ b/notebooks/direct_steering.ipynb
@@ -9,9 +9,9 @@
     "\n",
     "This notebook can be directly downloaded {download}`here <./direct_steering.ipynb>` to run it locally.\n",
     "\n",
-    "It demonstrates the use of [direct steering](https://www.jupedsim.org/stable/concepts/routing.html#direct-steering). of agents.\n",
+    "It demonstrates the use of [direct steering](https://www.jupedsim.org/stable/concepts/routing.html#direct-steering) of agents.\n",
     "\n",
-    "Am agent (leader) embarks on a journey defined by specific waypoints and a final destination. Meanwhile, the remaining agents trail behind, constantly adjusting their course to align with the leader's current position."
+    "An agent (leader) embarks on a journey defined by specific waypoints and a final destination. Meanwhile, the remaining agents trail behind, constantly adjusting their course to align with the leader's current position."
    ]
   },
   {
@@ -249,10 +249,20 @@
     "direct_steering_journey_id = simulation.add_journey(direct_steering_journey)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "0d643d8f",
+   "metadata": {},
+   "source": [
+    "## Add agents\n",
+    "\n",
+    "First, add leader, then its followers."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "335348ab",
+   "id": "86acd974",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -287,37 +297,53 @@
     "        )\n",
     "        for pos in pos_in_spawning_area[1:]\n",
     "    ]\n",
-    ")\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "18018147",
+   "metadata": {},
+   "source": [
+    "## Simulation loop"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "335348ab",
+   "metadata": {},
+   "outputs": [],
+   "source": [
     "while simulation.agent_count() > 0:\n",
     "    # Find leader's position\n",
-    "    for agent in simulation.agents():\n",
-    "        if agent.id == leader_id:\n",
-    "            position_leader = agent.position\n",
+    "    if leader_id in simulation.removed_agents():\n",
+    "        leader_id = None\n",
+    "    if leader_id:\n",
+    "        position_leader = simulation.agent(leader_id).position\n",
     "\n",
     "    # Move followers towards leader\n",
-    "    for follower_id in followers_ids:\n",
-    "        for agent in simulation.agents():\n",
-    "            if agent.id == follower_id:\n",
-    "                # Define a target position near the leader with some randomness\n",
-    "                near_leader = (\n",
-    "                    position_leader[0] + random.normalvariate(1, 0.1),\n",
-    "                    position_leader[1] + random.normalvariate(1, 0.1),\n",
-    "                )\n",
-    "                near_leader_point = Point(near_leader[0], near_leader[1])\n",
+    "    for agent in simulation.agents():\n",
+    "        if agent.id == leader_id:\n",
+    "            continue\n",
+    "        # Define a target position near the leader with some randomness\n",
+    "        near_leader = (\n",
+    "            position_leader[0] + random.normalvariate(1, 0.1),\n",
+    "            position_leader[1] + random.normalvariate(1, 0.1),\n",
+    "        )\n",
+    "        near_leader_point = Point(near_leader[0], near_leader[1])\n",
     "\n",
-    "                # If the target position is inside the walkable area, set it as the agent's target\n",
-    "                target = (\n",
-    "                    near_leader\n",
-    "                    if any(\n",
-    "                        geom.contains(near_leader_point) for geom in area.geoms\n",
-    "                    )\n",
-    "                    else position_leader\n",
-    "                )\n",
-    "                agent.target = target\n",
+    "        # If the target position is inside the walkable area, set it as the agent's target\n",
+    "        target = (\n",
+    "            near_leader\n",
+    "            if any(geom.contains(near_leader_point) for geom in area.geoms)\n",
+    "            else position_leader\n",
+    "        )\n",
+    "        agent.target = target\n",
     "\n",
-    "                # Check if the agent reached the exit and mark it for removal if so\n",
-    "                if Point(agent.position).distance(exit_area.centroid) < 1:\n",
-    "                    simulation.mark_agent_for_removal(follower_id)\n",
+    "        # Check if the agent reached the exit and mark it for removal if so\n",
+    "        if Point(agent.position).distance(exit_area.centroid) < 1:\n",
+    "            simulation.mark_agent_for_removal(agent.id)\n",
     "\n",
     "    simulation.iterate()"
    ]