{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e940562a-f560-4e30-aebc-e45f0cb743aa",
   "metadata": {},
   "source": [
    "# 01 — Exploring the CarRacing-v3 Environment\n",
    "\n",
    "### This notebook completes the step 1: understand the state, action, and rewardstructure of CarRacing-v3 before we start implementing PPO."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fb31fdc-63d5-450b-af3c-0881b490a440",
   "metadata": {},
   "source": [
    "## Cell 1：Verify GPU + Gymnasium"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "8bf88494-f116-435d-b70b-e31b4f4b21f6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PyTorch: 2.7.0+cu128\n",
      "CUDA available: True\n",
      "GPU: NVIDIA GeForce RTX 4060 Laptop GPU\n",
      "Gymnasium: 1.1.1\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import gymnasium as gym\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "print(\"PyTorch:\", torch.__version__)\n",
    "print(\"CUDA available:\", torch.cuda.is_available())\n",
    "print(\"GPU:\", torch.cuda.get_device_name(0))\n",
    "print(\"Gymnasium:\", gym.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50b52118-0cf5-4bc7-a7ca-c7923f954b22",
   "metadata": {},
   "source": [
    "## Cell 2：Create the environment and view the state and action space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "a7b00df7-904b-43be-ad36-25eb0034058c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Observation space: Box(0, 255, (96, 96, 3), uint8)\n",
      "Observation shape: (96, 96, 3) uint8\n",
      "Pixel value range: 0 to 253\n",
      "\n",
      "Action space: Discrete(5)\n",
      "Number of discrete actions: 5\n",
      "\n",
      "Info dict: {}\n"
     ]
    }
   ],
   "source": [
    "# continuous=False switches the action space to a discrete set of 5 actions\n",
    "env = gym.make(\"CarRacing-v3\", continuous=False, render_mode=\"rgb_array\")\n",
    "\n",
    "obs, info = env.reset(seed=42)\n",
    "\n",
    "print(\"Observation space:\", env.observation_space)\n",
    "print(\"Observation shape:\", obs.shape, obs.dtype)\n",
    "print(\"Pixel value range:\", obs.min(), \"to\", obs.max())\n",
    "print()\n",
    "print(\"Action space:\", env.action_space)\n",
    "print(\"Number of discrete actions:\", env.action_space.n)\n",
    "print()\n",
    "print(\"Info dict:\", info)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77b997c8-f5f9-48ec-b769-4e4b81138e44",
   "metadata": {},
   "source": [
    "## Cell 3：Visualize a few frames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "44f4a504-db59-4f41-93c5-7859ebf5cb09",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 60984 characters]",
      "text/plain": [
       "<Figure size 1600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axes = plt.subplots(1, 4, figsize=(16, 4))\n",
    "\n",
    "obs, _ = env.reset(seed=0)\n",
    "axes[0].imshow(obs)\n",
    "axes[0].set_title(\"Initial frame (step 0)\")\n",
    "axes[0].axis(\"off\")\n",
    "\n",
    "# Press the gas action 3 times in a row\n",
    "for i, action in enumerate([3, 3, 3]):\n",
    "  obs, reward, terminated, truncated, _ = env.step(action)\n",
    "  axes[i + 1].imshow(obs)\n",
    "  axes[i + 1].set_title(f\"Gas step {i + 1}\\nreward={reward:.2f}\")\n",
    "  axes[i + 1].axis(\"off\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33bcbb11-d1aa-4a5e-9fab-c0cc90631671",
   "metadata": {},
   "source": [
    "## Cell 4: Random Strategy baseline (1 episode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "e42c181b-3b2a-4c3f-9a53-bbd64d46e343",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Random policy, 1 episode:\n",
      "  steps        = 1000\n",
      "  total reward = -59.25\n",
      "  terminated   = False, truncated = True\n"
     ]
    }
   ],
   "source": [
    "obs, _ = env.reset(seed=0)\n",
    "total_reward = 0.0\n",
    "steps = 0\n",
    "done = False\n",
    "\n",
    "while not done:\n",
    "  action = env.action_space.sample()\n",
    "  obs, reward, term, trunc, _ = env.step(action)\n",
    "  total_reward += reward\n",
    "  steps += 1\n",
    "  done = term or trunc\n",
    "\n",
    "print(\"Random policy, 1 episode:\")\n",
    "print(f\"  steps        = {steps}\")\n",
    "print(f\"  total reward = {total_reward:.2f}\")\n",
    "print(f\"  terminated   = {term}, truncated = {trunc}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ca4ebb2-56eb-41a4-81e5-dc307d22c2a3",
   "metadata": {},
   "source": [
    "## Cell 5 : Calculate the average of the five episodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "d3198496-dbea-4a10-8152-ad9fe6fbb75b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 0: return = -52.98\n",
      "Episode 1: return = -52.73\n",
      "Episode 2: return = -64.18\n",
      "Episode 3: return = -48.34\n",
      "Episode 4: return = -52.73\n",
      "\n",
      "Random policy mean over 5 episodes: -54.19 +/- 5.29\n"
     ]
    }
   ],
   "source": [
    "returns = []\n",
    "for ep in range(5):\n",
    "  obs, _ = env.reset(seed=ep)\n",
    "  total_r = 0.0\n",
    "  done = False\n",
    "  while not done:\n",
    "      action = env.action_space.sample()\n",
    "      obs, r, term, trunc, _ = env.step(action)\n",
    "      total_r += r\n",
    "      done = term or trunc\n",
    "  returns.append(total_r)\n",
    "  print(f\"Episode {ep}: return = {total_r:.2f}\")\n",
    "\n",
    "print(f\"\\nRandom policy mean over 5 episodes: \"\n",
    "    f\"{np.mean(returns):.2f} +/- {np.std(returns):.2f}\")\n",
    "\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9856d21e-685e-49da-acca-01b1f4e73e64",
   "metadata": {},
   "source": [
    "## Cell 6 :Summary\n",
    "- **State**: (96, 96, 3) uint8 RGB top-down view of the track.\n",
    "- **Action**: `Discrete(5)` — `0=noop, 1=left, 2=right, 3=gas, 4=brake`.\n",
    "- **Reward**: `+1000/N` for each newly visited tile (N = total tiles), `-0.1` per frame.\n",
    "- **Random policy baseline**: roughly `-60` to `+30` (mean ~ `-30`).\n",
    "- **Goal**: train PPO to consistently reach `600+`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffc7f3db-bb34-449c-b330-710f25300a03",
   "metadata": {},
   "source": [
    "## Cell 7: Test the wrapped env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d0c7dda9-d5ce-452b-b9ca-3990fbe789ed",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wrapped observation space: Box(0, 255, (4, 84, 84), uint8)\n",
      "Wrapped observation shape: (4, 84, 84) uint8\n",
      "Wrapped action space: Discrete(5)\n",
      "Pixel range: 0 to 174\n"
     ]
    }
   ],
   "source": [
    "# Reload the kernel? If sys.path doesn't include the project root, do this once:\n",
    "import sys\n",
    "from pathlib import Path\n",
    "project_root = Path.cwd().parent\n",
    "if str(project_root) not in sys.path:\n",
    "  sys.path.insert(0, str(project_root))\n",
    "\n",
    "from src.env_wrappers import make_env\n",
    "\n",
    "wrapped = make_env(seed=42)\n",
    "obs, info = wrapped.reset(seed=42)\n",
    "\n",
    "print(\"Wrapped observation space:\", wrapped.observation_space)\n",
    "print(\"Wrapped observation shape:\", obs.shape, obs.dtype)\n",
    "print(\"Wrapped action space:\", wrapped.action_space)\n",
    "print(\"Pixel range:\", obs.min(), \"to\", obs.max())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf51239d-ae75-4ed0-b82e-cad7f385f202",
   "metadata": {},
   "source": [
    "## Cell 8: Visualize stacked frames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "07648c9a-09e9-433b-8e7e-c1bd296c5973",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 31396 characters]",
      "text/plain": [
       "<Figure size 1600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Show the 4 stacked frames side by side\n",
    "fig, axes = plt.subplots(1, 4, figsize=(16, 4))\n",
    "for i in range(4):\n",
    "  axes[i].imshow(obs[i], cmap=\"gray\")\n",
    "  axes[i].set_title(f\"Stacked frame {i}\")\n",
    "  axes[i].axis(\"off\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5037827a-8d11-46d2-91e2-47c101806f2d",
   "metadata": {},
   "source": [
    "## Cell 9: After a few skipped steps, frames should differ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e44d97ed-54bd-42b5-ad85-b4b42f2c31b3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 28776 characters]",
      "text/plain": [
       "<Figure size 1600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: 10 calls of step() = 40 raw env frames (skip=4)\n"
     ]
    }
   ],
   "source": [
    "# Take a few \"gas\" actions and see frames diverge\n",
    "obs, _ = wrapped.reset(seed=42)\n",
    "for _ in range(10):\n",
    "  obs, r, term, trunc, _ = wrapped.step(3)  # gas\n",
    "\n",
    "fig, axes = plt.subplots(1, 4, figsize=(16, 4))\n",
    "for i in range(4):\n",
    "  axes[i].imshow(obs[i], cmap=\"gray\")\n",
    "  axes[i].set_title(f\"After moving, frame {i}\")\n",
    "  axes[i].axis(\"off\")\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(\"Note: 10 calls of step() = 40 raw env frames (skip=4)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "308a650b-6c0b-440d-96c3-5af575693f53",
   "metadata": {},
   "source": [
    "## Cell 10: Random baseline with the wrapped env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "98e71435-c366-4a83-bef4-8864a6769585",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wrapped random episode 0: return = -59.26\n",
      "Wrapped random episode 1: return = -24.09\n",
      "Wrapped random episode 2: return = -28.32\n",
      "\n",
      "Wrapped random baseline (3 ep): -37.22\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "returns = []\n",
    "for ep in range(3):\n",
    "  obs, _ = wrapped.reset(seed=ep + 100)\n",
    "  total_r = 0.0\n",
    "  done = False\n",
    "  while not done:\n",
    "      action = wrapped.action_space.sample()\n",
    "      obs, r, term, trunc, _ = wrapped.step(action)\n",
    "      total_r += r\n",
    "      done = term or trunc\n",
    "  returns.append(total_r)\n",
    "  print(f\"Wrapped random episode {ep}: return = {total_r:.2f}\")\n",
    "\n",
    "print(f\"\\nWrapped random baseline (3 ep): {np.mean(returns):.2f}\")\n",
    "wrapped.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:pytorch]",
   "language": "python",
   "name": "conda-env-pytorch-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.21"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
