diff --git a/notebooks/unit2/requirements-unit2.txt b/notebooks/unit2/requirements-unit2.txt
index 3209a69c..2900fc40 100644
--- a/notebooks/unit2/requirements-unit2.txt
+++ b/notebooks/unit2/requirements-unit2.txt
@@ -3,7 +3,6 @@ pygame
 numpy
 
 huggingface_hub
-pickle5
 pyyaml==6.0
 imageio
 imageio_ffmpeg
diff --git a/notebooks/unit2/unit2.ipynb b/notebooks/unit2/unit2.ipynb
index e9ae624c..8750137e 100644
--- a/notebooks/unit2/unit2.ipynb
+++ b/notebooks/unit2/unit2.ipynb
@@ -3,8 +3,8 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "view-in-github",
-        "colab_type": "text"
+        "colab_type": "text",
+        "id": "view-in-github"
       },
       "source": [
         "<a href=\"https://colab.research.google.com/github/huggingface/deep-rl-class/blob/main/notebooks/unit2/unit2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
@@ -36,6 +36,9 @@
     },
     {
       "cell_type": "markdown",
+      "metadata": {
+        "id": "DPTBOv9HYLZ2"
+      },
       "source": [
         "###🎮 Environments:\n",
         "\n",
@@ -48,10 +51,7 @@
         "- [Gymnasium](https://gymnasium.farama.org/)\n",
         "\n",
         "We're constantly trying to improve our tutorials, so **if you find some issues in this notebook**, please [open an issue on the GitHub Repo](https://github.com/huggingface/deep-rl-class/issues)."
-      ],
-      "metadata": {
-        "id": "DPTBOv9HYLZ2"
-      }
+      ]
     },
     {
       "cell_type": "markdown",
@@ -72,14 +72,14 @@
     },
     {
       "cell_type": "markdown",
+      "metadata": {
+        "id": "viNzVbVaYvY3"
+      },
       "source": [
         "## This notebook is from the Deep Reinforcement Learning Course\n",
         "\n",
         "<img src=\"https://huggingface.co/datasets/huggingface-deep-rl-course/course-images/resolve/main/en/notebooks/deep-rl-course-illustration.jpg\" alt=\"Deep RL Course illustration\"/>"
-      ],
-      "metadata": {
-        "id": "viNzVbVaYvY3"
-      }
+      ]
     },
     {
       "cell_type": "markdown",
@@ -156,28 +156,31 @@
     },
     {
       "cell_type": "markdown",
-      "source": [
-        "# Let's code our first Reinforcement Learning algorithm 🚀"
-      ],
       "metadata": {
         "id": "HEtx8Y8MqKfH"
-      }
+      },
+      "source": [
+        "# Let's code our first Reinforcement Learning algorithm 🚀"
+      ]
     },
     {
       "cell_type": "markdown",
+      "metadata": {
+        "id": "Kdxb1IhzTn0v"
+      },
       "source": [
         "To validate this hands-on for the [certification process](https://huggingface.co/deep-rl-course/en/unit0/introduction#certification-process), you need to push your trained Taxi model to the Hub and **get a result of >= 4.5**.\n",
         "\n",
         "To find your result, go to the [leaderboard](https://huggingface.co/spaces/huggingface-projects/Deep-Reinforcement-Learning-Leaderboard) and find your model, **the result = mean_reward - std of reward**\n",
         "\n",
         "For more information about the certification process, check this section 👉 https://huggingface.co/deep-rl-course/en/unit0/introduction#certification-process"
-      ],
-      "metadata": {
-        "id": "Kdxb1IhzTn0v"
-      }
+      ]
     },
     {
       "cell_type": "markdown",
+      "metadata": {
+        "id": "4gpxC1_kqUYe"
+      },
       "source": [
         "## Install dependencies and create a virtual display 🔽\n",
         "\n",
@@ -194,10 +197,7 @@
         "The Hugging Face Hub 🤗 works as a central place where anyone can share and explore models and datasets. It has versioning, metrics, visualizations and other features that will allow you to easily collaborate with others.\n",
         "\n",
         "You can see here all the Deep RL models available (if they use Q Learning) here 👉 https://huggingface.co/models?other=q-learning"
-      ],
-      "metadata": {
-        "id": "4gpxC1_kqUYe"
-      }
+      ]
     },
     {
       "cell_type": "code",
@@ -212,53 +212,53 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "n71uTX7qqzz2"
+      },
+      "outputs": [],
       "source": [
         "!sudo apt-get update\n",
         "!sudo apt-get install -y python3-opengl\n",
         "!apt install ffmpeg xvfb\n",
         "!pip3 install pyvirtualdisplay"
-      ],
-      "metadata": {
-        "id": "n71uTX7qqzz2"
-      },
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "markdown",
-      "source": [
-        "To make sure the new installed libraries are used, **sometimes it's required to restart the notebook runtime**. The next cell will force the **runtime to crash, so you'll need to connect again and run the code starting from here**. Thanks to this trick, **we will be able to run our virtual screen.**"
-      ],
       "metadata": {
         "id": "K6XC13pTfFiD"
-      }
+      },
+      "source": [
+        "To make sure the new installed libraries are used, **sometimes it's required to restart the notebook runtime**. The next cell will force the **runtime to crash, so you'll need to connect again and run the code starting from here**. Thanks to this trick, **we will be able to run our virtual screen.**"
+      ]
     },
     {
       "cell_type": "code",
-      "source": [
-        "import os\n",
-        "os.kill(os.getpid(), 9)"
-      ],
+      "execution_count": null,
       "metadata": {
         "id": "3kuZbWAkfHdg"
       },
-      "execution_count": null,
-      "outputs": []
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "os.kill(os.getpid(), 9)"
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "DaY1N4dBrabi"
+      },
+      "outputs": [],
       "source": [
         "# Virtual display\n",
         "from pyvirtualdisplay import Display\n",
         "\n",
         "virtual_display = Display(visible=0, size=(1400, 900))\n",
         "virtual_display.start()"
-      ],
-      "metadata": {
-        "id": "DaY1N4dBrabi"
-      },
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "markdown",
@@ -289,7 +289,7 @@
         "import os\n",
         "import tqdm\n",
         "\n",
-        "import pickle5 as pickle\n",
+        "import pickle\n",
         "from tqdm.notebook import tqdm"
       ]
     },
@@ -1065,6 +1065,11 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "U4mdUTKkGnUd"
+      },
+      "outputs": [],
       "source": [
         "def push_to_hub(\n",
         "    repo_id, model, env, video_fps=1, local_repo_path=\"hub\"\n",
@@ -1194,12 +1199,7 @@
         "    )\n",
         "\n",
         "    print(\"Your model is pushed to the Hub. You can view your model here: \", repo_url)"
-      ],
-      "metadata": {
-        "id": "U4mdUTKkGnUd"
-      },
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "markdown",
@@ -1748,15 +1748,15 @@
   ],
   "metadata": {
     "colab": {
-      "private_outputs": true,
-      "provenance": [],
       "collapsed_sections": [
         "67OdoKL63eDD",
         "B2_-8b8z5k54",
         "8R5ej1fS4P2V",
         "Pnpk2ePoem3r"
       ],
-      "include_colab_link": true
+      "include_colab_link": true,
+      "private_outputs": true,
+      "provenance": []
     },
     "gpuClass": "standard",
     "kernelspec": {
@@ -1769,4 +1769,4 @@
   },
   "nbformat": 4,
   "nbformat_minor": 0
-}
\ No newline at end of file
+}