{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #f5f5f5; padding: 20px; border-radius: 10px; border: 1px solid #e0e0e0; text-align: center;\">\n",
    "  \n",
    "  <img src=\"https://i.ibb.co/mLyt3rk/SUT-Logo.png\" alt=\"SUT Logo\" width=\"100\">\n",
    "  \n",
    "  <h1 style=\"margin:10px 0 0 0; color:#333;\">Sharif University of Technology</h1>\n",
    "  <h2 style=\"margin:0; color:#555; font-weight:normal;\">Faculty of Mathematical Sciences</h2>\n",
    "  \n",
    "  <hr style=\"border: 1px solid #ccc; margin-top:15px; margin-bottom:15px;\">\n",
    "  \n",
    "  <h3 style=\"margin:10px 0 5px 0; color:#0056b3;\">HW4: Denoising Diffusion Models</h3>\n",
    "  <h4 style=\"margin:10px 0 5px 0; color:#0056b3;\">Prepared by: Sara Karimi</h4>\n",
    "  <br/>\n",
    "  \n",
    "  <div style=\"font-size: 1.1em; color: #333; line-height: 1.6;\">\n",
    "    <strong>Course:</strong> Generative Models<br>\n",
    "    <strong>Instructor:</strong> Dr. Fatemeh Seyyedsalehi<br>\n",
    "    <strong>Semester:</strong> Fall 1404-1405\n",
    "  </div>\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #f5f5f5; border: 1px solid #f5f5f5; padding: 20px; border-radius: 10px; margin-top: 25px; margin-bottom: 25px;\">\n",
    "  \n",
    "  <h3 style=\"text-align: center; margin-top: 0; color: #0056b3;\">Student Information</h3>\n",
    "  <p style=\"text-align: center; color: #555; margin-bottom: 20px; margin-top: 0;\">\n",
    "    Please fill in your details below.\n",
    "  </p>\n",
    "\n",
    "  <table style=\"width: 60%; margin-left: auto; margin-right: auto; border: none; font-size: 1.1em; color: #333;\">\n",
    "    <tbody>\n",
    "      <tr>\n",
    "        <td style=\"text-align: right; padding: 8px; width: 35%;\"><strong>Student Name:</strong></td>\n",
    "        <td style=\"text-align: left; padding: 8px;\"><em>Your Name Here</em></td> <!-- UPDATE THIS LINE -->\n",
    "      </tr>\n",
    "      <tr>\n",
    "        <td style=\"text-align: right; padding: 8px;\"><strong>Student ID:</strong></td>\n",
    "        <td style=\"text-align: left; padding: 8px;\"><em>Your Student ID Here</em></td>  <!-- UPDATE THIS LINE -->\n",
    "      </tr>\n",
    "    </tbody>\n",
    "  </table>\n",
    "  \n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# **HW4-P3: Trajectory Stitching with Human Reference Models**\n",
    "\n",
    "---\n",
    "\n",
    "**Assignment Overview:**\n",
    "\n",
    "1. Implement **classifier-guided diffusion** for trajectory generation\n",
    "2. Understand **negative behavior guidance** for distribution shift correction\n",
    "3. Implement **trajectory stitching** with conditional diffusion\n",
    "4. Integrate **Human Reference Models (HRM)** as planning priors\n",
    "\n",
    "<br/>\n",
    "\n",
    "> **Submission Note:** If the file size of your notebook becomes too large and you struggle to upload the final ZIP file to Quera, you may upload your solution to Colab or Kaggle instead. In this case, please include a `.txt` file in your Quera submission containing a link to your notebook (ensure the link settings allow anyone with the link to view it)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup & Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.patches import Rectangle\n",
    "from collections import deque\n",
    "import heapq\n",
    "from typing import List, Tuple, Optional, Dict\n",
    "from dataclasses import dataclass\n",
    "from scipy.ndimage import gaussian_filter\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "torch.manual_seed(42)\n",
    "np.random.seed(42)\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print(f\"Using device: {device}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Part A: Environment & HRM Setup (10 points)\n",
    "\n",
    "We use a 2D navigation environment with obstacles. The **Human Reference Model (HRM)** provides planning priors through:\n",
    "1. A* pathfinding for subgoal generation\n",
    "2. Heatmap representation of preferred trajectories\n",
    "3. Direction field for goal-directed navigation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class MazeConfig:\n",
    "    \"\"\"Configuration for the maze environment.\"\"\"\n",
    "    size: int = 10\n",
    "    start: Tuple[float, float] = (1.0, 1.0)\n",
    "    goal: Tuple[float, float] = (8.5, 8.5)\n",
    "    obstacles: List[Tuple[int, int]] = None\n",
    "    \n",
    "    def __post_init__(self):\n",
    "        if self.obstacles is None:\n",
    "            self.obstacles = [\n",
    "                (4, 4), (5, 4), (4, 5), (5, 5),  # Center block\n",
    "                (7, 2), (8, 2),                   # Lower-right\n",
    "                (2, 7), (2, 8),                   # Upper-left\n",
    "            ]\n",
    "\n",
    "class MazeEnvironment:\n",
    "    \"\"\"2D Navigation Environment with continuous state/action space.\"\"\"\n",
    "    \n",
    "    def __init__(self, config: MazeConfig):\n",
    "        self.config = config\n",
    "        self.size = config.size\n",
    "        self.start = np.array(config.start)\n",
    "        self.goal = np.array(config.goal)\n",
    "        self.obstacles = set(config.obstacles)\n",
    "        self.state = self.start.copy()\n",
    "        \n",
    "    def reset(self) -> np.ndarray:\n",
    "        \"\"\"Reset to start position.\"\"\"\n",
    "        self.state = self.start.copy()\n",
    "        return self.state.copy()\n",
    "    \n",
    "    def is_valid_position(self, pos: np.ndarray) -> bool:\n",
    "        \"\"\"Check if position is valid (in bounds and not in obstacle).\"\"\"\n",
    "        if pos[0] < 0.5 or pos[0] > self.size - 0.5:\n",
    "            return False\n",
    "        if pos[1] < 0.5 or pos[1] > self.size - 0.5:\n",
    "            return False\n",
    "        cell = (int(pos[0]), int(pos[1]))\n",
    "        return cell not in self.obstacles\n",
    "    \n",
    "    def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool]:\n",
    "        \"\"\"\n",
    "        Take action in environment.\n",
    "        Action: [dx, dy] velocity vector, clipped to max magnitude 0.5\n",
    "        Returns: (next_state, reward, done)\n",
    "        \"\"\"\n",
    "        action = np.clip(action, -0.5, 0.5)\n",
    "        \n",
    "        next_state = self.state + action\n",
    "        \n",
    "        if self.is_valid_position(next_state):\n",
    "            self.state = next_state\n",
    "            wall_penalty = 0.0\n",
    "        else:\n",
    "            wall_penalty = -0.5\n",
    "        \n",
    "        dist_to_goal = np.linalg.norm(self.state - self.goal)\n",
    "        reward = -0.01 * dist_to_goal + wall_penalty  \n",
    "        \n",
    "        done = dist_to_goal < 0.8\n",
    "        if done:\n",
    "            reward += 10.0  \n",
    "            \n",
    "        return self.state.copy(), reward, done\n",
    "    \n",
    "    def render(self, ax=None, trajectories=None, title=\"Maze Environment\"):\n",
    "        \"\"\"Visualize the environment.\"\"\"\n",
    "        if ax is None:\n",
    "            fig, ax = plt.subplots(figsize=(8, 8))\n",
    "        \n",
    "        ax.set_xlim(0, self.size)\n",
    "        ax.set_ylim(0, self.size)\n",
    "        ax.set_aspect('equal')\n",
    "        ax.grid(True, alpha=0.3)\n",
    "        \n",
    "        for (ox, oy) in self.obstacles:\n",
    "            rect = Rectangle((ox, oy), 1, 1, facecolor='gray', edgecolor='black')\n",
    "            ax.add_patch(rect)\n",
    "        \n",
    "        ax.plot(*self.start, 'go', markersize=15, label='Start', zorder=5)\n",
    "        ax.plot(*self.goal, 'y*', markersize=20, label='Goal', zorder=5)\n",
    "        \n",
    "        # Draw trajectories if provided\n",
    "        if trajectories is not None:\n",
    "            for traj in trajectories:\n",
    "                traj = np.array(traj)\n",
    "                ax.plot(traj[:, 0], traj[:, 1], 'b-', alpha=0.5, linewidth=1)\n",
    "        \n",
    "        ax.legend(loc='upper left')\n",
    "        ax.set_title(title)\n",
    "        return ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class HumanReferenceModel:\n",
    "    \"\"\"\n",
    "    Human Reference Model (HRM) for planning priors.\n",
    "\n",
    "    Provides:\n",
    "    1. A* pathfinding for subgoal generation\n",
    "    2. Heatmap of preferred trajectory regions\n",
    "    3. Direction field for navigation guidance\n",
    "    4. Scoring function for trajectory quality\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, env: MazeEnvironment, resolution: int = 50):\n",
    "        self.env = env\n",
    "        self.resolution = resolution\n",
    "        self.subgoals = self._compute_subgoals()\n",
    "        self.heatmap = self._generate_heatmap()\n",
    "        self.direction_field = self._compute_direction_field()\n",
    "\n",
    "    def _compute_subgoals(self) -> List[np.ndarray]:\n",
    "        \"\"\"Use A* to find path from start to goal.\"\"\"\n",
    "        start_cell = (int(self.env.start[0]), int(self.env.start[1]))\n",
    "        goal_cell = (int(self.env.goal[0]), int(self.env.goal[1]))\n",
    "\n",
    "        def heuristic(a, b):\n",
    "            return np.sqrt((a[0] - b[0])**2 + (a[1] - b[1])**2)\n",
    "\n",
    "        def get_neighbors(cell):\n",
    "            x, y = cell\n",
    "            neighbors = []\n",
    "            for dx, dy in [(-1,0), (1,0), (0,-1), (0,1), (-1,-1), (-1,1), (1,-1), (1,1)]:\n",
    "                nx, ny = x + dx, y + dy\n",
    "                if 0 <= nx < self.env.size and 0 <= ny < self.env.size:\n",
    "                    if (nx, ny) not in self.env.obstacles:\n",
    "                        neighbors.append((nx, ny))\n",
    "            return neighbors\n",
    "\n",
    "        frontier = [(0, start_cell)]\n",
    "        came_from = {start_cell: None}\n",
    "        cost_so_far = {start_cell: 0}\n",
    "\n",
    "        while frontier:\n",
    "            _, current = heapq.heappop(frontier)\n",
    "\n",
    "            if current == goal_cell:\n",
    "                break\n",
    "\n",
    "            for neighbor in get_neighbors(current):\n",
    "                new_cost = cost_so_far[current] + heuristic(current, neighbor)\n",
    "                if neighbor not in cost_so_far or new_cost < cost_so_far[neighbor]:\n",
    "                    cost_so_far[neighbor] = new_cost\n",
    "                    priority = new_cost + heuristic(neighbor, goal_cell)\n",
    "                    heapq.heappush(frontier, (priority, neighbor))\n",
    "                    came_from[neighbor] = current\n",
    "\n",
    "        path = []\n",
    "        current = goal_cell\n",
    "        while current is not None:\n",
    "            path.append(np.array([current[0] + 0.5, current[1] + 0.5]))\n",
    "            current = came_from.get(current)\n",
    "        path.reverse()\n",
    "\n",
    "        return path\n",
    "\n",
    "    def _generate_heatmap(self) -> np.ndarray:\n",
    "        \"\"\"Generate heatmap showing preferred trajectory regions.\"\"\"\n",
    "        heatmap = np.zeros((self.resolution, self.resolution))\n",
    "\n",
    "        for i in range(len(self.subgoals) - 1):\n",
    "            p1 = self.subgoals[i]\n",
    "            p2 = self.subgoals[i + 1]\n",
    "\n",
    "            for t in np.linspace(0, 1, 30):\n",
    "                p = p1 + t * (p2 - p1)\n",
    "                hx = int(p[0] / self.env.size * self.resolution)\n",
    "                hy = int(p[1] / self.env.size * self.resolution)\n",
    "                hx = np.clip(hx, 0, self.resolution - 1)\n",
    "                hy = np.clip(hy, 0, self.resolution - 1)\n",
    "                heatmap[hy, hx] = 1.0\n",
    "\n",
    "        heatmap = gaussian_filter(heatmap, sigma=4)\n",
    "\n",
    "        if heatmap.max() > 0:\n",
    "            heatmap = heatmap / heatmap.max()\n",
    "\n",
    "        return heatmap\n",
    "\n",
    "    def _compute_direction_field(self) -> np.ndarray:\n",
    "        \"\"\"Compute direction field pointing toward nearest subgoal.\"\"\"\n",
    "        field = np.zeros((self.resolution, self.resolution, 2))\n",
    "\n",
    "        for i in range(self.resolution):\n",
    "            for j in range(self.resolution):\n",
    "                x = (i + 0.5) / self.resolution * self.env.size\n",
    "                y = (j + 0.5) / self.resolution * self.env.size\n",
    "                pos = np.array([x, y])\n",
    "\n",
    "                min_dist = float('inf')\n",
    "                closest_sg_idx = -1\n",
    "                target = self.env.goal \n",
    "\n",
    "                for sg_idx, sg_item in enumerate(self.subgoals):\n",
    "                    dist = np.linalg.norm(pos - sg_item)\n",
    "                    if dist < min_dist:\n",
    "                        min_dist = dist\n",
    "                        closest_sg_idx = sg_idx\n",
    "\n",
    "                if closest_sg_idx != -1 and closest_sg_idx < len(self.subgoals) - 1:\n",
    "                    target = self.subgoals[closest_sg_idx + 1]\n",
    "\n",
    "                direction = target - pos\n",
    "                norm = np.linalg.norm(direction)\n",
    "                if norm > 1e-6:\n",
    "                    direction = direction / norm\n",
    "\n",
    "                field[j, i] = direction\n",
    "\n",
    "        return field\n",
    "\n",
    "    def get_heatmap_value(self, pos: np.ndarray) -> float:\n",
    "        \"\"\"Get heatmap value at continuous position.\"\"\"\n",
    "        hx = int(pos[0] / self.env.size * self.resolution)\n",
    "        hy = int(pos[1] / self.env.size * self.resolution)\n",
    "        hx = np.clip(hx, 0, self.resolution - 1)\n",
    "        hy = np.clip(hy, 0, self.resolution - 1)\n",
    "        return self.heatmap[hy, hx]\n",
    "\n",
    "    def get_direction(self, pos: np.ndarray) -> np.ndarray:\n",
    "        \"\"\"Get recommended direction at position.\"\"\"\n",
    "        hx = int(pos[0] / self.env.size * self.resolution)\n",
    "        hy = int(pos[1] / self.env.size * self.resolution)\n",
    "        hx = np.clip(hx, 0, self.resolution - 1)\n",
    "        hy = np.clip(hy, 0, self.resolution - 1)\n",
    "        return self.direction_field[hy, hx]\n",
    "\n",
    "    def get_direction_to_goal(self, pos: np.ndarray) -> np.ndarray:\n",
    "        \"\"\"Get unit direction vector toward goal.\"\"\"\n",
    "        direction = self.env.goal - pos\n",
    "        norm = np.linalg.norm(direction)\n",
    "        if norm > 1e-6:\n",
    "            return direction / norm\n",
    "        return np.zeros(2)\n",
    "\n",
    "    def score_trajectory(self, states: np.ndarray) -> float:\n",
    "        \"\"\"\n",
    "        Score a trajectory based on HRM criteria.\n",
    "        Higher score = better trajectory.\n",
    "        \"\"\"\n",
    "        score = 0.0\n",
    "\n",
    "        for i, s in enumerate(states):\n",
    "            score += self.get_heatmap_value(s) * 2.0\n",
    "\n",
    "            dist_to_goal = np.linalg.norm(s - self.env.goal)\n",
    "            score -= dist_to_goal * 0.1\n",
    "\n",
    "            if i > 0:\n",
    "                move_dir = states[i] - states[i-1]\n",
    "                move_norm = np.linalg.norm(move_dir)\n",
    "                if move_norm > 1e-6:\n",
    "                    move_dir = move_dir / move_norm\n",
    "                    ideal_dir = self.get_direction(states[i-1])\n",
    "                    alignment = np.dot(move_dir, ideal_dir)\n",
    "                    score += alignment * 1.0\n",
    "\n",
    "        final_dist = np.linalg.norm(states[-1] - self.env.goal)\n",
    "        if final_dist < 0.8:\n",
    "            score += 20.0\n",
    "\n",
    "        return score\n",
    "\n",
    "    def visualize(self):\n",
    "        \"\"\"Visualize HRM components.\"\"\"\n",
    "        fig, axes = plt.subplots(1, 2, figsize=(14, 6))\n",
    "\n",
    "        ax = axes[0]\n",
    "        im = ax.imshow(self.heatmap, origin='lower', extent=[0, self.env.size, 0, self.env.size],\n",
    "                      cmap='YlOrRd', alpha=0.7)\n",
    "        plt.colorbar(im, ax=ax, label='Preference')\n",
    "\n",
    "        for (ox, oy) in self.env.obstacles:\n",
    "            rect = Rectangle((ox, oy), 1, 1, facecolor='gray', edgecolor='black', alpha=0.8)\n",
    "            ax.add_patch(rect)\n",
    "\n",
    "        ax.plot(*self.env.start, 'go', markersize=12, label='Start')\n",
    "        ax.plot(*self.env.goal, 'y*', markersize=15, label='Goal')\n",
    "        ax.set_title('HRM Heatmap')\n",
    "        ax.legend()\n",
    "\n",
    "        ax = axes[1]\n",
    "        self.env.render(ax=ax, title='A* Path (Subgoals)')\n",
    "        subgoals = np.array(self.subgoals)\n",
    "        ax.plot(subgoals[:, 0], subgoals[:, 1], 'r-o', linewidth=2, markersize=8, label='A* Path')\n",
    "        ax.legend()\n",
    "\n",
    "        plt.tight_layout()\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize environment and HRM\n",
    "config = MazeConfig()\n",
    "env = MazeEnvironment(config)\n",
    "hrm = HumanReferenceModel(env)\n",
    "\n",
    "hrm.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Question A.1 (5 points)\n",
    "\n",
    "**Why is the HRM not a policy?** Explain the distinction between a planning prior (HRM) and a policy.\n",
    "\n",
    "---\n",
    "\n",
    "**YOUR ANSWER:**\n",
    "\n",
    "[Write your answer here]\n",
    "\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Part B: Behavior Policy & Data Collection (10 points)\n",
    "\n",
    "We define a **behavior policy β(a|s)** that explores the environment. This policy should:\n",
    "- Be suboptimal (to create distribution shift)\n",
    "- But still reach the goal sometimes (30-50% success) so diffusion learns useful patterns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GaussianPolicy(nn.Module):\n",
    "    \"\"\"\n",
    "    Gaussian Policy: π(a|s) = N(μ(s), σ²)\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, state_dim: int = 2, action_dim: int = 2, hidden_dim: int = 64,\n",
    "                 log_std_init: float = -1.0):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.network = nn.Sequential(\n",
    "            nn.Linear(state_dim, hidden_dim),\n",
    "            nn.Tanh(),\n",
    "            nn.Linear(hidden_dim, hidden_dim),\n",
    "            nn.Tanh(),\n",
    "            nn.Linear(hidden_dim, action_dim)\n",
    "        )\n",
    "        \n",
    "        self.log_std = nn.Parameter(torch.ones(action_dim) * log_std_init)\n",
    "        \n",
    "    def forward(self, state: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:\n",
    "        mean = self.network(state)\n",
    "        std = torch.exp(self.log_std).expand_as(mean)\n",
    "        return mean, std\n",
    "    \n",
    "    def sample(self, state: torch.Tensor) -> torch.Tensor:\n",
    "        mean, std = self.forward(state)\n",
    "        noise = torch.randn_like(mean)\n",
    "        action = mean + std * noise\n",
    "        return torch.clamp(action, -0.5, 0.5)\n",
    "    \n",
    "    def log_prob(self, state: torch.Tensor, action: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        TODO: Compute log π(a|s) for Gaussian policy.\n",
    "        \n",
    "        Formula: log N(a; μ, σ²) = -0.5 * [(a-μ)²/σ² + log(2πσ²)]\n",
    "        Sum over action dimensions.\n",
    "        \n",
    "        Returns: (batch,) tensor\n",
    "        \"\"\"\n",
    "        mean, std = self.forward(state)\n",
    "        var = std ** 2\n",
    "        \n",
    "        ### YOUR CODE HERE ###\n",
    "        \n",
    "        raise NotImplementedError(\"Implement log_prob\")\n",
    "        ### END YOUR CODE ###\n",
    "    \n",
    "    def score(self, state: torch.Tensor, action: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        TODO: Compute ∇_a log π(a|s) - the score function.\n",
    "        \n",
    "        For Gaussian: ∇_a log N(a; μ, σ²) = -(a - μ) / σ²\n",
    "        \n",
    "        Returns: (batch, action_dim) tensor\n",
    "        \"\"\"\n",
    "        mean, std = self.forward(state)\n",
    "        var = std ** 2\n",
    "        \n",
    "        ### YOUR CODE HERE ###\n",
    "        \n",
    "        raise NotImplementedError(\"Implement score\")\n",
    "        ### END YOUR CODE ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BehaviorPolicy:\n",
    "    \"\"\"\n",
    "    Behavior policy for data collection.\n",
    "    Uses noisy goal-directed strategy that reaches goal ~40% of time.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, env: MazeEnvironment, hrm: HumanReferenceModel):\n",
    "        self.env = env\n",
    "        self.hrm = hrm\n",
    "        self.policy = GaussianPolicy(state_dim=2, action_dim=2, log_std_init=-0.5)\n",
    "        self._pretrain_policy()\n",
    "        \n",
    "    def _pretrain_policy(self, n_steps: int = 3000):\n",
    "        \"\"\"Pre-train behavior policy to be suboptimal but functional.\"\"\"\n",
    "        optimizer = torch.optim.Adam(self.policy.parameters(), lr=1e-3)\n",
    "        \n",
    "        for step in range(n_steps):\n",
    "            states_np = np.random.rand(64, 2) * (self.env.size - 1) + 0.5\n",
    "            states = torch.tensor(states_np, dtype=torch.float32)\n",
    "            \n",
    "            target_actions = []\n",
    "            for s in states_np:\n",
    "                goal_dir = self.hrm.get_direction_to_goal(s)\n",
    "                hrm_dir = self.hrm.get_direction(s)\n",
    "                \n",
    "                direction = 0.6 * goal_dir + 0.4 * hrm_dir\n",
    "                direction = direction / (np.linalg.norm(direction) + 1e-6)\n",
    "                \n",
    "                speed = 0.25 + np.random.randn() * 0.05\n",
    "                action = direction * speed\n",
    "                \n",
    "                # Add noise to make suboptimal\n",
    "                action += np.random.randn(2) * 0.08\n",
    "                target_actions.append(action)\n",
    "            \n",
    "            target = torch.tensor(np.array(target_actions), dtype=torch.float32)\n",
    "            \n",
    "            mean, _ = self.policy(states)\n",
    "            loss = ((mean - target) ** 2).mean()\n",
    "            \n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "    \n",
    "    def get_action(self, state: np.ndarray) -> np.ndarray:\n",
    "        with torch.no_grad():\n",
    "            state_t = torch.tensor(state, dtype=torch.float32).unsqueeze(0)\n",
    "            action = self.policy.sample(state_t)\n",
    "            return action.squeeze(0).numpy()\n",
    "    \n",
    "    def log_prob(self, state: torch.Tensor, action: torch.Tensor) -> torch.Tensor:\n",
    "        return self.policy.log_prob(state, action)\n",
    "    \n",
    "    def score(self, state: torch.Tensor, action: torch.Tensor) -> torch.Tensor:\n",
    "        return self.policy.score(state, action)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collect_behavior_data(env: MazeEnvironment, behavior_policy: BehaviorPolicy,\n",
    "                          n_trajectories: int = 200, max_steps: int = 80) -> List[Dict]:\n",
    "    \"\"\"Collect trajectories from behavior policy.\"\"\"\n",
    "    trajectories = []\n",
    "    \n",
    "    for _ in range(n_trajectories):\n",
    "        state = env.reset()\n",
    "        states = [state.copy()]\n",
    "        actions = []\n",
    "        rewards = []\n",
    "        \n",
    "        for _ in range(max_steps):\n",
    "            action = behavior_policy.get_action(state)\n",
    "            next_state, reward, done = env.step(action)\n",
    "            \n",
    "            actions.append(action.copy())\n",
    "            rewards.append(reward)\n",
    "            states.append(next_state.copy())\n",
    "            state = next_state\n",
    "            \n",
    "            if done:\n",
    "                break\n",
    "        \n",
    "        trajectories.append({\n",
    "            'states': np.array(states),\n",
    "            'actions': np.array(actions),\n",
    "            'rewards': np.array(rewards)\n",
    "        })\n",
    "    \n",
    "    return trajectories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize behavior policy and collect data\n",
    "print(\"Training behavior policy...\")\n",
    "behavior_policy = BehaviorPolicy(env, hrm)\n",
    "\n",
    "print(\"Collecting behavior data...\")\n",
    "behavior_data = collect_behavior_data(env, behavior_policy, n_trajectories=300, max_steps=80)\n",
    "\n",
    "# Check success rate\n",
    "success_count = sum(1 for t in behavior_data if np.linalg.norm(t['states'][-1] - env.goal) < 0.8)\n",
    "print(f\"\\nCollected {len(behavior_data)} trajectories\")\n",
    "print(f\"Average trajectory length: {np.mean([len(t['rewards']) for t in behavior_data]):.1f}\")\n",
    "print(f\"Success rate: {success_count / len(behavior_data) * 100:.1f}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize behavior trajectories\n",
    "fig, ax = plt.subplots(figsize=(8, 8))\n",
    "\n",
    "# Color by success\n",
    "for t in behavior_data[:100]:\n",
    "    states = t['states']\n",
    "    success = np.linalg.norm(states[-1] - env.goal) < 0.8\n",
    "    color = 'green' if success else 'blue'\n",
    "    ax.plot(states[:, 0], states[:, 1], '-', color=color, alpha=0.3, linewidth=1)\n",
    "\n",
    "env.render(ax=ax, title=f'Behavior Policy (Success: {success_count/len(behavior_data)*100:.0f}%)')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Part C: Conditional Diffusion Model (20 points)\n",
    "\n",
    "We train a conditional diffusion model on sub-trajectories of length `w`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SubTrajectoryDataset(Dataset):\n",
    "    \"\"\"Dataset of sub-trajectories for training conditional diffusion.\"\"\"\n",
    "    \n",
    "    def __init__(self, trajectories: List[Dict], window_size: int = 8):\n",
    "        self.window_size = window_size\n",
    "        self.windows = []\n",
    "        self.conditions = []\n",
    "        \n",
    "        for traj in trajectories:\n",
    "            states = traj['states']\n",
    "            actions = traj['actions']\n",
    "            \n",
    "            for i in range(len(actions) - window_size + 1):\n",
    "                window_states = states[i:i+window_size+1]\n",
    "                window_actions = actions[i:i+window_size]\n",
    "                \n",
    "                # Flatten: [s0, a0, s1, a1, ..., s_w]\n",
    "                window = []\n",
    "                for j in range(window_size):\n",
    "                    window.extend(window_states[j].tolist())\n",
    "                    window.extend(window_actions[j].tolist())\n",
    "                window.extend(window_states[window_size].tolist())\n",
    "                \n",
    "                self.windows.append(np.array(window))\n",
    "                self.conditions.append(states[i].copy())\n",
    "        \n",
    "        self.windows = np.array(self.windows, dtype=np.float32)\n",
    "        self.conditions = np.array(self.conditions, dtype=np.float32)\n",
    "        \n",
    "        # Normalize\n",
    "        self.window_mean = self.windows.mean(axis=0)\n",
    "        self.window_std = self.windows.std(axis=0) + 1e-6\n",
    "        self.cond_mean = self.conditions.mean(axis=0)\n",
    "        self.cond_std = self.conditions.std(axis=0) + 1e-6\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.windows)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        \"\"\"\n",
    "        TODO 7: Return normalized (window, condition) pair.\n",
    "        \n",
    "        Normalization: x_norm = (x - mean) / std\n",
    "        \n",
    "        Steps:\n",
    "        1. Get window and condition at idx\n",
    "        2. Normalize window: (window - self.window_mean) / self.window_std\n",
    "        3. Normalize condition: (cond - self.cond_mean) / self.cond_std\n",
    "        4. Return as torch tensors\n",
    "        \n",
    "        Returns:\n",
    "            (window_tensor, condition_tensor) both normalized\n",
    "        \"\"\"\n",
    "        ### YOUR CODE HERE ###\n",
    "        \n",
    "        raise NotImplementedError(\"Implement __getitem__\")\n",
    "        ### END YOUR CODE ###\n",
    "    \n",
    "    def unnormalize_window(self, window: torch.Tensor) -> torch.Tensor:\n",
    "        mean = torch.tensor(self.window_mean, device=window.device, dtype=window.dtype)\n",
    "        std = torch.tensor(self.window_std, device=window.device, dtype=window.dtype)\n",
    "        return window * std + mean\n",
    "    \n",
    "    def normalize_condition(self, cond: torch.Tensor) -> torch.Tensor:\n",
    "        mean = torch.tensor(self.cond_mean, device=cond.device, dtype=cond.dtype)\n",
    "        std = torch.tensor(self.cond_std, device=cond.device, dtype=cond.dtype)\n",
    "        return (cond - mean) / std"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SinusoidalEmbedding(nn.Module):\n",
    "    def __init__(self, dim: int):\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "        \n",
    "    def forward(self, t: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        TODO 1: Implement sinusoidal positional embedding for diffusion timestep.\n",
    "        \n",
    "        This encodes the timestep t into a high-dimensional vector using sin/cos functions.\n",
    "        \n",
    "        Formula:\n",
    "            emb[2i] = sin(t / 10000^(2i/dim))\n",
    "            emb[2i+1] = cos(t / 10000^(2i/dim))\n",
    "        \n",
    "        Steps:\n",
    "        1. Compute half_dim = self.dim // 2\n",
    "        2. Create frequency: emb = log(10000) / (half_dim - 1)\n",
    "        3. Create position encoding: emb = torch.exp(torch.arange(half_dim, device=device) * -emb)\n",
    "        4. Multiply: emb = t[:, None] * emb[None, :]\n",
    "        5. Concatenate sin and cos: torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)\n",
    "        \n",
    "        Args:\n",
    "            t: (batch_size,) tensor of timesteps\n",
    "        Returns:\n",
    "            (batch_size, dim) tensor of embeddings\n",
    "        \"\"\"\n",
    "        device = t.device\n",
    "        \n",
    "        ### YOUR CODE HERE ###\n",
    "        \n",
    "        raise NotImplementedError(\"Implement sinusoidal embedding\")\n",
    "        ### END YOUR CODE ###\n",
    "\n",
    "\n",
    "class ConditionalDiffusionModel(nn.Module):\n",
    "    \"\"\"Conditional Diffusion Model for sub-trajectory generation.\"\"\"\n",
    "    \n",
    "    def __init__(self, trajectory_dim: int, cond_dim: int = 2, \n",
    "                 hidden_dim: int = 256, time_dim: int = 64):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.trajectory_dim = trajectory_dim\n",
    "        \n",
    "        self.time_embed = nn.Sequential(\n",
    "            SinusoidalEmbedding(time_dim),\n",
    "            nn.Linear(time_dim, hidden_dim),\n",
    "            nn.GELU(),\n",
    "            nn.Linear(hidden_dim, hidden_dim)\n",
    "        )\n",
    "        \n",
    "        self.cond_embed = nn.Sequential(\n",
    "            nn.Linear(cond_dim, hidden_dim),\n",
    "            nn.GELU(),\n",
    "            nn.Linear(hidden_dim, hidden_dim)\n",
    "        )\n",
    "        \n",
    "        self.input_proj = nn.Linear(trajectory_dim, hidden_dim)\n",
    "        \n",
    "        self.blocks = nn.ModuleList([\n",
    "            nn.Sequential(\n",
    "                nn.LayerNorm(hidden_dim),\n",
    "                nn.Linear(hidden_dim, hidden_dim * 2),\n",
    "                nn.GELU(),\n",
    "                nn.Linear(hidden_dim * 2, hidden_dim)\n",
    "            ) for _ in range(4)\n",
    "        ])\n",
    "        \n",
    "        self.output_proj = nn.Sequential(\n",
    "            nn.LayerNorm(hidden_dim),\n",
    "            nn.Linear(hidden_dim, trajectory_dim)\n",
    "        )\n",
    "        \n",
    "    def forward(self, x: torch.Tensor, t: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        TODO 2: Implement the forward pass of the conditional diffusion model.\n",
    "        \n",
    "        This predicts the noise ε given noisy trajectory x_t, timestep t, and condition.\n",
    "        \n",
    "        Steps:\n",
    "        1. Embed timestep: t_emb = self.time_embed(t.float())\n",
    "        2. Embed condition: c_emb = self.cond_embed(cond)\n",
    "        3. Project input: h = self.input_proj(x)\n",
    "        4. Add embeddings: h = h + t_emb + c_emb\n",
    "        5. Apply residual blocks: for block in self.blocks: h = h + block(h)\n",
    "        6. Output projection: return self.output_proj(h)\n",
    "        \n",
    "        Args:\n",
    "            x: (batch, traj_dim) noisy trajectory\n",
    "            t: (batch,) timesteps\n",
    "            cond: (batch, cond_dim) conditioning (initial state)\n",
    "        Returns:\n",
    "            (batch, traj_dim) predicted noise\n",
    "        \"\"\"\n",
    "        ### YOUR CODE HERE ###\n",
    "        \n",
    "        raise NotImplementedError(\"Implement conditional diffusion forward pass\")\n",
    "        ### END YOUR CODE ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DiffusionTrainer:\n",
    "    \"\"\"DDPM training for conditional diffusion model.\"\"\"\n",
    "    \n",
    "    def __init__(self, model: ConditionalDiffusionModel, \n",
    "                 n_steps: int = 100, beta_start: float = 1e-4, beta_end: float = 0.02):\n",
    "        self.model = model\n",
    "        self.n_steps = n_steps\n",
    "        \n",
    "        self.betas = torch.linspace(beta_start, beta_end, n_steps)\n",
    "        self.alphas = 1 - self.betas\n",
    "        self.alpha_bars = torch.cumprod(self.alphas, dim=0)\n",
    "        \n",
    "    def get_loss(self, x0: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        TODO 3: Implement DDPM training loss.\n",
    "        \n",
    "        The forward process adds noise: q(x_t | x_0) = N(x_t; √ᾱ_t x_0, (1-ᾱ_t)I)\n",
    "        \n",
    "        Steps:\n",
    "        1. Sample random timesteps: t ~ Uniform{0, 1, ..., n_steps-1}\n",
    "        2. Get noise schedule: alpha_bar = self.alpha_bars[t]\n",
    "        3. Sample noise: ε ~ N(0, I)\n",
    "        4. Create noisy sample: x_t = √(alpha_bar) * x_0 + √(1 - alpha_bar) * ε\n",
    "        5. Predict noise: ε_pred = model(x_t, t, cond)\n",
    "        6. Return MSE loss: ||ε - ε_pred||²\n",
    "        \n",
    "        Args:\n",
    "            x0: (batch, dim) clean trajectories\n",
    "            cond: (batch, cond_dim) conditioning\n",
    "        Returns:\n",
    "            scalar loss\n",
    "        \"\"\"\n",
    "        batch_size = x0.shape[0]\n",
    "        device = x0.device\n",
    "        \n",
    "        ### YOUR CODE HERE ###\n",
    "        \n",
    "        raise NotImplementedError(\"Implement DDPM training loss\")\n",
    "        ### END YOUR CODE ###\n",
    "    \n",
    "    def train(self, dataset: SubTrajectoryDataset, n_epochs: int = 150, \n",
    "              batch_size: int = 128, lr: float = 1e-3):\n",
    "        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)\n",
    "        optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr)\n",
    "        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, n_epochs)\n",
    "        \n",
    "        self.model.to(device)\n",
    "        losses = []\n",
    "        \n",
    "        for epoch in range(n_epochs):\n",
    "            epoch_loss = 0\n",
    "            for windows, conds in dataloader:\n",
    "                windows = windows.to(device)\n",
    "                conds = conds.to(device)\n",
    "                \n",
    "                loss = self.get_loss(windows, conds)\n",
    "                \n",
    "                optimizer.zero_grad()\n",
    "                loss.backward()\n",
    "                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)\n",
    "                optimizer.step()\n",
    "                \n",
    "                epoch_loss += loss.item()\n",
    "            \n",
    "            scheduler.step()\n",
    "            avg_loss = epoch_loss / len(dataloader)\n",
    "            losses.append(avg_loss)\n",
    "            \n",
    "            if (epoch + 1) % 30 == 0:\n",
    "                print(f\"Epoch {epoch+1}/{n_epochs}, Loss: {avg_loss:.4f}\")\n",
    "        \n",
    "        return losses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create dataset and train\n",
    "WINDOW_SIZE = 8\n",
    "\n",
    "dataset = SubTrajectoryDataset(behavior_data, window_size=WINDOW_SIZE)\n",
    "print(f\"Dataset size: {len(dataset)} windows\")\n",
    "print(f\"Window dimension: {dataset.windows.shape[1]}\")\n",
    "\n",
    "trajectory_dim = dataset.windows.shape[1]\n",
    "diffusion_model = ConditionalDiffusionModel(\n",
    "    trajectory_dim=trajectory_dim,\n",
    "    cond_dim=2,\n",
    "    hidden_dim=256\n",
    ")\n",
    "\n",
    "trainer = DiffusionTrainer(diffusion_model, n_steps=100)\n",
    "print(\"\\nTraining diffusion model...\")\n",
    "losses = trainer.train(dataset, n_epochs=150, batch_size=128, lr=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 4))\n",
    "plt.plot(losses)\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Loss')\n",
    "plt.title('Diffusion Model Training Loss')\n",
    "plt.grid(True, alpha=0.3)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Part D: Target Policy with HRM (15 points)\n",
    "\n",
    "The target policy is more aligned with HRM than the behavior policy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class HRMGuidedTargetPolicy(nn.Module):\n",
    "    \"\"\"Target policy π(a|s) guided by HRM - better than behavior policy.\"\"\"\n",
    "    \n",
    "    def __init__(self, env: MazeEnvironment, hrm: HumanReferenceModel,\n",
    "                 log_std: float = -1.5):\n",
    "        super().__init__()\n",
    "        self.env = env\n",
    "        self.hrm = hrm\n",
    "        \n",
    "        self.network = nn.Sequential(\n",
    "            nn.Linear(2, 64),\n",
    "            nn.Tanh(),\n",
    "            nn.Linear(64, 64),\n",
    "            nn.Tanh(),\n",
    "            nn.Linear(64, 2)\n",
    "        )\n",
    "        \n",
    "        self.log_std = nn.Parameter(torch.tensor([log_std, log_std]))\n",
    "        self._pretrain()\n",
    "        \n",
    "    def _pretrain(self, n_steps: int = 3000):\n",
    "        \"\"\"Pre-train target policy to follow HRM guidance closely.\"\"\"\n",
    "        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)\n",
    "        \n",
    "        for _ in range(n_steps):\n",
    "            states_np = np.random.rand(64, 2) * (self.env.size - 1) + 0.5\n",
    "            states = torch.tensor(states_np, dtype=torch.float32)\n",
    "            \n",
    "            target_actions = []\n",
    "            for s in states_np:\n",
    "                # Strong alignment with HRM direction\n",
    "                hrm_dir = self.hrm.get_direction(s)\n",
    "                goal_dir = self.hrm.get_direction_to_goal(s)\n",
    "                \n",
    "                # Mix: 70% HRM, 30% direct goal\n",
    "                direction = 0.7 * hrm_dir + 0.3 * goal_dir\n",
    "                direction = direction / (np.linalg.norm(direction) + 1e-6)\n",
    "                \n",
    "                # Faster speed than behavior (more optimal)\n",
    "                heatmap_val = self.hrm.get_heatmap_value(s)\n",
    "                speed = 0.35 + 0.1 * heatmap_val\n",
    "                action = direction * speed\n",
    "                target_actions.append(action)\n",
    "            \n",
    "            target = torch.tensor(np.array(target_actions), dtype=torch.float32)\n",
    "            mean = self.network(states)\n",
    "            loss = ((mean - target) ** 2).mean()\n",
    "            \n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "    \n",
    "    def forward(self, state: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:\n",
    "        mean = self.network(state)\n",
    "        std = torch.exp(self.log_std).expand_as(mean)\n",
    "        return mean, std\n",
    "    \n",
    "    def sample(self, state: torch.Tensor) -> torch.Tensor:\n",
    "        mean, std = self.forward(state)\n",
    "        noise = torch.randn_like(mean)\n",
    "        action = mean + std * noise\n",
    "        return torch.clamp(action, -0.5, 0.5)\n",
    "    \n",
    "    def log_prob(self, state: torch.Tensor, action: torch.Tensor) -> torch.Tensor:\n",
    "        mean, std = self.forward(state)\n",
    "        var = std ** 2\n",
    "        log_prob = -0.5 * (((action - mean) ** 2) / var + torch.log(2 * np.pi * var))\n",
    "        return log_prob.sum(dim=-1)\n",
    "    \n",
    "    def score(self, state: torch.Tensor, action: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"Compute ∇_a log π(a|s).\"\"\"\n",
    "        mean, std = self.forward(state)\n",
    "        var = std ** 2\n",
    "        return -(action - mean) / var"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize target policy\n",
    "print(\"Training target policy...\")\n",
    "target_policy = HRMGuidedTargetPolicy(env, hrm)\n",
    "target_policy.to(device)\n",
    "\n",
    "# Compare policies\n",
    "test_states = torch.tensor([[2.0, 2.0], [5.0, 5.0], [7.0, 7.0]], dtype=torch.float32)\n",
    "print(\"\\nPolicy comparison:\")\n",
    "for s in test_states:\n",
    "    s_np = s.numpy()\n",
    "    s_gpu = s.unsqueeze(0).to(device)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        beta_mean, _ = behavior_policy.policy(s.unsqueeze(0))\n",
    "        pi_mean, _ = target_policy(s_gpu)\n",
    "    \n",
    "    print(f\"  State {s_np}: β={beta_mean.squeeze().numpy()}, π={pi_mean.cpu().squeeze().numpy()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Part E: Guided Sampling - STITCH-OPE Core (20 points)\n",
    "\n",
    "The guidance function:\n",
    "$$g(\\tau) = \\alpha \\nabla_\\tau \\sum_t \\log \\pi(a_t|s_t) - \\lambda \\nabla_\\tau \\sum_t \\log \\beta(a_t|s_t)$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GuidedDiffusionSampler:\n",
    "    \"\"\"\n",
    "    Guided Diffusion Sampler implementing STITCH-OPE.\n",
    "    g(τ) = α∇_τ Σlog π(a|s) - λ∇_τ Σlog β(a|s)\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, \n",
    "                 diffusion_model: ConditionalDiffusionModel,\n",
    "                 target_policy: nn.Module,\n",
    "                 behavior_policy: BehaviorPolicy,\n",
    "                 dataset: SubTrajectoryDataset,\n",
    "                 trainer: DiffusionTrainer,\n",
    "                 hrm: HumanReferenceModel,\n",
    "                 window_size: int = 8,\n",
    "                 alpha: float = 1.0,\n",
    "                 lambda_: float = 0.3):\n",
    "        \n",
    "        self.model = diffusion_model\n",
    "        self.target_policy = target_policy\n",
    "        self.behavior_policy = behavior_policy\n",
    "        self.dataset = dataset\n",
    "        self.trainer = trainer\n",
    "        self.hrm = hrm\n",
    "        self.window_size = window_size\n",
    "        self.alpha = alpha\n",
    "        self.lambda_ = lambda_\n",
    "        \n",
    "    def extract_states_actions(self, window: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:\n",
    "        \"\"\"\n",
    "        TODO 4: Extract states and actions from flattened window tensor.\n",
    "        \n",
    "        Window format: [s0, a0, s1, a1, ..., s_{w-1}, a_{w-1}, s_w]\n",
    "        Where each s is 2D (x,y) and each a is 2D (dx,dy)\n",
    "        Total dim = w*4 + 2 (for w=8: 34)\n",
    "        \n",
    "        Steps:\n",
    "        1. Initialize lists for states and actions\n",
    "        2. For i in range(window_size):\n",
    "           - State at index i*4 : i*4+2\n",
    "           - Action at index i*4+2 : i*4+4\n",
    "        3. Final state at index -2:\n",
    "        4. Stack into tensors: (batch, w+1, 2) for states, (batch, w, 2) for actions\n",
    "        \n",
    "        Args:\n",
    "            window: (batch, traj_dim) flattened trajectory window\n",
    "        Returns:\n",
    "            states: (batch, w+1, 2) state sequence\n",
    "            actions: (batch, w, 2) action sequence\n",
    "        \"\"\"\n",
    "        batch_size = window.shape[0]\n",
    "        w = self.window_size\n",
    "        \n",
    "        ### YOUR CODE HERE ###\n",
    "        \n",
    "        raise NotImplementedError(\"Implement extract_states_actions\")\n",
    "        ### END YOUR CODE ###\n",
    "    \n",
    "    def compute_guidance(self, window: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        TODO 5: Compute STITCH-OPE classifier guidance.\n",
    "        \n",
    "        Paper Eq. 6: g(τ) = α∇_τ log π(τ) - λ∇_τ log β(τ)\n",
    "        \n",
    "        For trajectory τ = (s_0, a_0, ..., s_w):\n",
    "            ∇_τ log π(τ) = Σ_t ∇_{a_t} log π(a_t|s_t)\n",
    "        \n",
    "        The key insight: we guide TOWARD target π and AWAY from behavior β.\n",
    "        \n",
    "        Steps:\n",
    "        1. Unnormalize: window_unnorm = self.dataset.unnormalize_window(window)\n",
    "        2. Extract: states, actions = self.extract_states_actions(window_unnorm)\n",
    "        3. Initialize: guidance = torch.zeros_like(window)\n",
    "        4. For each timestep t:\n",
    "           - Get s_t, a_t\n",
    "           - Compute target score: pi_score = self.target_policy.score(s_t, a_t)\n",
    "           - Compute behavior score: beta_score = self.behavior_policy.policy.score(s_t, a_t)\n",
    "           - Combined: action_guidance = self.alpha * pi_score - self.lambda_ * beta_score\n",
    "           - Place at action indices: guidance[:, t*4+2 : t*4+4] = action_guidance\n",
    "        5. Normalize: guidance = guidance / (guidance.norm(dim=-1, keepdim=True) + 1e-6)\n",
    "        \n",
    "        Args:\n",
    "            window: (batch, traj_dim) current noisy trajectory\n",
    "        Returns:\n",
    "            (batch, traj_dim) guidance direction\n",
    "        \"\"\"\n",
    "        ### YOUR CODE HERE ###\n",
    "        \n",
    "        raise NotImplementedError(\"Implement STITCH-OPE guidance\")\n",
    "        ### END YOUR CODE ###\n",
    "    \n",
    "    @torch.no_grad()\n",
    "    def sample(self, initial_state: torch.Tensor, n_samples: int = 1,\n",
    "               guidance_scale: float = 0.5) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        TODO 6: Implement guided reverse diffusion sampling (DDPM).\n",
    "        \n",
    "        Reverse process: p(x_{t-1}|x_t) = N(x_{t-1}; μ_θ(x_t, t), σ_t²I)\n",
    "        \n",
    "        Where:\n",
    "            μ = (x_t - β_t/√(1-ᾱ_t) * ε_θ(x_t,t)) / √α_t\n",
    "            σ² = β_t * (1-ᾱ_{t-1}) / (1-ᾱ_t)\n",
    "        \n",
    "        With guidance:\n",
    "            μ_guided = μ + guidance_scale * σ * g(x_t)\n",
    "        \n",
    "        Steps:\n",
    "        1. Normalize condition: cond = self.dataset.normalize_condition(initial_state)\n",
    "        2. Start from noise: x = torch.randn(n_samples, traj_dim)\n",
    "        3. For t = T-1, ..., 0:\n",
    "           a. Get noise schedule values: alpha, alpha_bar, beta\n",
    "           b. Predict noise: noise_pred = model(x, t, cond)\n",
    "           c. Compute mean: mean = (x - β/√(1-ᾱ) * noise_pred) / √α\n",
    "           d. Compute variance: σ² = β * (1-ᾱ_{t-1}) / (1-ᾱ_t)\n",
    "           e. Apply guidance (if 10 < t < 80): mean += guidance_scale * σ * guidance\n",
    "           f. Sample: x = mean + σ * noise (if t > 0, else x = mean)\n",
    "        4. Return unnormalized trajectory\n",
    "        \n",
    "        Args:\n",
    "            initial_state: (2,) starting position\n",
    "            n_samples: number of trajectories to generate\n",
    "            guidance_scale: strength of classifier guidance\n",
    "        Returns:\n",
    "            (n_samples, traj_dim) generated trajectories\n",
    "        \"\"\"\n",
    "        self.model.eval()\n",
    "        \n",
    "        ### YOUR CODE HERE ###\n",
    "        \n",
    "        raise NotImplementedError(\"Implement guided reverse diffusion sampling\")\n",
    "        ### END YOUR CODE ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize sampler\n",
    "sampler = GuidedDiffusionSampler(\n",
    "    diffusion_model=diffusion_model,\n",
    "    target_policy=target_policy,\n",
    "    behavior_policy=behavior_policy,\n",
    "    dataset=dataset,\n",
    "    trainer=trainer,\n",
    "    hrm=hrm,\n",
    "    window_size=WINDOW_SIZE,\n",
    "    alpha=1.0,\n",
    "    lambda_=0.3\n",
    ")\n",
    "\n",
    "# Test sampling\n",
    "initial_state = torch.tensor(env.start, dtype=torch.float32, device=device)\n",
    "samples = sampler.sample(initial_state, n_samples=5, guidance_scale=0.5)\n",
    "print(f\"Sample shape: {samples.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Part F: Trajectory Stitching (10 points)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TrajectoryStitcher:\n",
    "    \"\"\"Stitch sub-trajectories using HRM scoring for candidate selection.\"\"\"\n",
    "    \n",
    "    def __init__(self, sampler: GuidedDiffusionSampler, env: MazeEnvironment,\n",
    "                 hrm: HumanReferenceModel, window_size: int = 8):\n",
    "        self.sampler = sampler\n",
    "        self.env = env\n",
    "        self.hrm = hrm\n",
    "        self.window_size = window_size\n",
    "        \n",
    "    def extract_trajectory_from_window(self, window: torch.Tensor) -> Tuple[np.ndarray, np.ndarray]:\n",
    "        window = window.cpu().numpy()\n",
    "        states = []\n",
    "        actions = []\n",
    "        \n",
    "        for i in range(self.window_size):\n",
    "            idx = i * 4\n",
    "            states.append(window[idx:idx+2])\n",
    "            actions.append(window[idx+2:idx+4])\n",
    "        states.append(window[-2:])\n",
    "        \n",
    "        return np.array(states), np.array(actions)\n",
    "    \n",
    "    def is_valid_trajectory(self, states: np.ndarray) -> bool:\n",
    "        \"\"\"Check if all states are valid.\"\"\"\n",
    "        for s in states:\n",
    "            if not self.env.is_valid_position(s):\n",
    "                return False\n",
    "        return True\n",
    "    \n",
    "    def stitch_trajectory(self, n_windows: int = 12, n_candidates: int = 8,\n",
    "                          guidance_scale: float = 0.5) -> Dict:\n",
    "        \"\"\"Generate trajectory by stitching windows.\"\"\"\n",
    "        all_states = [self.env.start.copy()]\n",
    "        all_actions = []\n",
    "        \n",
    "        current_state = torch.tensor(self.env.start, dtype=torch.float32, device=device)\n",
    "        \n",
    "        for w_idx in range(n_windows):\n",
    "            candidates = self.sampler.sample(current_state, n_samples=n_candidates,\n",
    "                                            guidance_scale=guidance_scale)\n",
    "            \n",
    "            best_score = float('-inf')\n",
    "            best_window = None\n",
    "            \n",
    "            for i in range(n_candidates):\n",
    "                window = candidates[i]\n",
    "                states, actions = self.extract_trajectory_from_window(window)\n",
    "                \n",
    "                if not self.is_valid_trajectory(states):\n",
    "                    continue\n",
    "                \n",
    "                score = self.hrm.score_trajectory(states)\n",
    "                \n",
    "                if score > best_score:\n",
    "                    best_score = score\n",
    "                    best_window = (states, actions)\n",
    "            \n",
    "            if best_window is None:\n",
    "                states, actions = self.extract_trajectory_from_window(candidates[0])\n",
    "                states = np.clip(states, 0.5, self.env.size - 0.5)\n",
    "                best_window = (states, actions)\n",
    "            \n",
    "            states, actions = best_window\n",
    "            all_states.extend(states[1:].tolist())\n",
    "            all_actions.extend(actions.tolist())\n",
    "            \n",
    "            current_state = torch.tensor(states[-1], dtype=torch.float32, device=device)\n",
    "            \n",
    "            if np.linalg.norm(states[-1] - self.env.goal) < 0.8:\n",
    "                break\n",
    "        \n",
    "        return {\n",
    "            'states': np.array(all_states),\n",
    "            'actions': np.array(all_actions),\n",
    "            'reached_goal': np.linalg.norm(all_states[-1] - self.env.goal) < 0.8\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize stitcher and generate trajectories\n",
    "stitcher = TrajectoryStitcher(sampler, env, hrm, window_size=WINDOW_SIZE)\n",
    "\n",
    "print(\"Generating trajectories...\")\n",
    "\n",
    "trajectories_no_guidance = []\n",
    "trajectories_with_guidance = []\n",
    "\n",
    "for i in range(15):\n",
    "    # No guidance\n",
    "    sampler.alpha = 0.0\n",
    "    sampler.lambda_ = 0.0\n",
    "    traj = stitcher.stitch_trajectory(n_windows=12, n_candidates=8, guidance_scale=0.0)\n",
    "    trajectories_no_guidance.append(traj)\n",
    "    \n",
    "    # With guidance\n",
    "    sampler.alpha = 1.0\n",
    "    sampler.lambda_ = 0.3\n",
    "    traj = stitcher.stitch_trajectory(n_windows=12, n_candidates=8, guidance_scale=0.5)\n",
    "    trajectories_with_guidance.append(traj)\n",
    "\n",
    "no_guid_success = np.mean([t['reached_goal'] for t in trajectories_no_guidance])\n",
    "with_guid_success = np.mean([t['reached_goal'] for t in trajectories_with_guidance])\n",
    "\n",
    "print(f\"\\nNo guidance success rate: {no_guid_success*100:.0f}%\")\n",
    "print(f\"With guidance success rate: {with_guid_success*100:.0f}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize\n",
    "fig, axes = plt.subplots(1, 2, figsize=(14, 6))\n",
    "\n",
    "ax = axes[0]\n",
    "env.render(ax=ax, title=f\"No Guidance (Success: {no_guid_success*100:.0f}%)\")\n",
    "for traj in trajectories_no_guidance:\n",
    "    states = traj['states']\n",
    "    color = 'green' if traj['reached_goal'] else 'blue'\n",
    "    ax.plot(states[:, 0], states[:, 1], '-', color=color, alpha=0.5, linewidth=1.5)\n",
    "\n",
    "ax = axes[1]\n",
    "env.render(ax=ax, title=f\"STITCH-OPE (Success: {with_guid_success*100:.0f}%)\")\n",
    "for traj in trajectories_with_guidance:\n",
    "    states = traj['states']\n",
    "    color = 'green' if traj['reached_goal'] else 'orange'\n",
    "    ax.plot(states[:, 0], states[:, 1], '-', color=color, alpha=0.5, linewidth=1.5)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Part G: Off-Policy Evaluation & Analysis (15 points)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_trajectory_return(states: np.ndarray, actions: np.ndarray, \n",
    "                               env: MazeEnvironment, gamma: float = 0.99) -> float:\n",
    "    total_return = 0\n",
    "    discount = 1.0\n",
    "    \n",
    "    for i in range(len(actions)):\n",
    "        s = states[i]\n",
    "        dist_to_goal = np.linalg.norm(s - env.goal)\n",
    "        reward = -0.01 * dist_to_goal\n",
    "        \n",
    "        if i + 1 < len(states):\n",
    "            if not env.is_valid_position(states[i + 1]):\n",
    "                reward -= 0.5\n",
    "        \n",
    "        if dist_to_goal < 0.8:\n",
    "            reward += 10.0\n",
    "        \n",
    "        total_return += discount * reward\n",
    "        discount *= gamma\n",
    "    \n",
    "    return total_return\n",
    "\n",
    "\n",
    "def get_ground_truth_returns(env: MazeEnvironment, target_policy: nn.Module,\n",
    "                              n_trajectories: int = 30, max_steps: int = 80):\n",
    "    returns = []\n",
    "    successes = []\n",
    "    \n",
    "    for _ in range(n_trajectories):\n",
    "        state = env.reset()\n",
    "        total_return = 0\n",
    "        discount = 1.0\n",
    "        \n",
    "        for step in range(max_steps):\n",
    "            state_t = torch.tensor(state, dtype=torch.float32, device=device).unsqueeze(0)\n",
    "            with torch.no_grad():\n",
    "                action = target_policy.sample(state_t).squeeze(0).cpu().numpy()\n",
    "            \n",
    "            next_state, reward, done = env.step(action)\n",
    "            total_return += discount * reward\n",
    "            discount *= 0.99\n",
    "            state = next_state\n",
    "            \n",
    "            if done:\n",
    "                successes.append(1)\n",
    "                break\n",
    "        else:\n",
    "            successes.append(0)\n",
    "        \n",
    "        returns.append(total_return)\n",
    "    \n",
    "    return returns, np.mean(successes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ground truth\n",
    "print(\"Computing ground truth...\")\n",
    "gt_returns, gt_success = get_ground_truth_returns(env, target_policy, n_trajectories=30)\n",
    "print(f\"Ground truth: Return = {np.mean(gt_returns):.2f} ± {np.std(gt_returns):.2f}, Success = {gt_success*100:.0f}%\")\n",
    "\n",
    "# Evaluate OPE methods\n",
    "def evaluate_ope(stitcher, alpha, lambda_, n_traj=15, guidance_scale=0.5):\n",
    "    stitcher.sampler.alpha = alpha\n",
    "    stitcher.sampler.lambda_ = lambda_\n",
    "    \n",
    "    returns = []\n",
    "    successes = []\n",
    "    \n",
    "    for _ in range(n_traj):\n",
    "        traj = stitcher.stitch_trajectory(n_windows=12, n_candidates=8, guidance_scale=guidance_scale)\n",
    "        ret = compute_trajectory_return(traj['states'], traj['actions'], env)\n",
    "        returns.append(ret)\n",
    "        successes.append(traj['reached_goal'])\n",
    "    \n",
    "    return np.mean(returns), np.std(returns), np.mean(successes)\n",
    "\n",
    "print(\"\\n\" + \"=\"*70)\n",
    "print(f\"{'Method':<35} {'Return':>12} {'Std':>10} {'Success':>10}\")\n",
    "print(\"-\"*70)\n",
    "print(f\"{'Ground Truth (Target Policy)':<35} {np.mean(gt_returns):>12.2f} {np.std(gt_returns):>10.2f} {gt_success*100:>9.0f}%\")\n",
    "print(\"-\"*70)\n",
    "\n",
    "ret, std, succ = evaluate_ope(stitcher, 0.0, 0.0, guidance_scale=0.0)\n",
    "print(f\"{'No Guidance':<35} {ret:>12.2f} {std:>10.2f} {succ*100:>9.0f}%\")\n",
    "\n",
    "ret, std, succ = evaluate_ope(stitcher, 1.0, 0.0, guidance_scale=0.5)\n",
    "print(f\"{'Target Only (α=1, λ=0)':<35} {ret:>12.2f} {std:>10.2f} {succ*100:>9.0f}%\")\n",
    "\n",
    "ret, std, succ = evaluate_ope(stitcher, 1.0, 0.3, guidance_scale=0.5)\n",
    "print(f\"{'STITCH-OPE (α=1, λ=0.3)':<35} {ret:>12.2f} {std:>10.2f} {succ*100:>9.0f}%\")\n",
    "\n",
    "ret, std, succ = evaluate_ope(stitcher, 1.0, 0.5, guidance_scale=0.5)\n",
    "print(f\"{'Higher λ (α=1, λ=0.5)':<35} {ret:>12.2f} {std:>10.2f} {succ*100:>9.0f}%\")\n",
    "\n",
    "print(\"=\"*70)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.9.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
