PPO_SharedControl_MK2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pygame 2.6.1 (SDL 2.28.4, Python 3.12.3)\n",
      "Hello from the pygame community. https://www.pygame.org/contribute.html\n",
      "Using device: cuda\n",
      "Episode 1/500, Reward: -191.15, Avg Reward (last 100): -191.15\n",
      "Episode 2/500, Reward: -382.16, Avg Reward (last 100): -286.66\n",
      "Episode 3/500, Reward: -63.07, Avg Reward (last 100): -212.13\n",
      "Episode 4/500, Reward: -79.46, Avg Reward (last 100): -178.96\n",
      "Episode 5/500, Reward: -180.01, Avg Reward (last 100): -179.17\n",
      "Episode 6/500, Reward: -254.00, Avg Reward (last 100): -191.64\n",
      "Episode 7/500, Reward: -85.35, Avg Reward (last 100): -176.46\n",
      "Episode 8/500, Reward: -153.69, Avg Reward (last 100): -173.61\n",
      "Episode 9/500, Reward: -365.54, Avg Reward (last 100): -194.94\n",
      "Episode 10/500, Reward: -338.55, Avg Reward (last 100): -209.30\n",
      "Episode 11/500, Reward: -481.17, Avg Reward (last 100): -234.01\n",
      "Episode 12/500, Reward: -200.40, Avg Reward (last 100): -231.21\n",
      "Episode 13/500, Reward: -118.10, Avg Reward (last 100): -222.51\n",
      "Episode 14/500, Reward: -469.54, Avg Reward (last 100): -240.16\n",
      "Episode 15/500, Reward: -8.51, Avg Reward (last 100): -224.71\n",
      "Episode 16/500, Reward: -139.82, Avg Reward (last 100): -219.41\n",
      "Episode 17/500, Reward: -39.37, Avg Reward (last 100): -208.82\n",
      "Episode 18/500, Reward: -34.68, Avg Reward (last 100): -199.14\n",
      "Episode 19/500, Reward: -229.15, Avg Reward (last 100): -200.72\n",
      "Episode 20/500, Reward: -80.25, Avg Reward (last 100): -194.70\n",
      "Episode 21/500, Reward: -207.22, Avg Reward (last 100): -195.29\n",
      "Episode 22/500, Reward: -302.36, Avg Reward (last 100): -200.16\n",
      "Episode 23/500, Reward: -39.49, Avg Reward (last 100): -193.18\n",
      "Episode 24/500, Reward: -250.16, Avg Reward (last 100): -195.55\n",
      "Episode 25/500, Reward: -90.95, Avg Reward (last 100): -191.37\n",
      "Episode 26/500, Reward: -47.68, Avg Reward (last 100): -185.84\n",
      "Episode 27/500, Reward: -74.41, Avg Reward (last 100): -181.71\n",
      "Episode 28/500, Reward: -120.48, Avg Reward (last 100): -179.53\n",
      "Episode 29/500, Reward: -135.97, Avg Reward (last 100): -178.02\n",
      "Episode 30/500, Reward: -186.08, Avg Reward (last 100): -178.29\n",
      "Episode 31/500, Reward: -216.68, Avg Reward (last 100): -179.53\n",
      "Episode 32/500, Reward: -101.35, Avg Reward (last 100): -177.09\n",
      "Episode 33/500, Reward: -228.48, Avg Reward (last 100): -178.64\n",
      "Episode 34/500, Reward: -23.50, Avg Reward (last 100): -174.08\n",
      "Episode 35/500, Reward: -6.31, Avg Reward (last 100): -169.29\n",
      "Episode 36/500, Reward: -135.08, Avg Reward (last 100): -168.34\n",
      "Episode 37/500, Reward: -278.29, Avg Reward (last 100): -171.31\n",
      "Episode 38/500, Reward: -88.68, Avg Reward (last 100): -169.14\n",
      "Episode 39/500, Reward: -348.17, Avg Reward (last 100): -173.73\n",
      "Episode 40/500, Reward: -148.23, Avg Reward (last 100): -173.09\n",
      "Episode 41/500, Reward: -122.21, Avg Reward (last 100): -171.85\n",
      "Episode 42/500, Reward: -306.08, Avg Reward (last 100): -175.04\n",
      "Episode 43/500, Reward: -137.60, Avg Reward (last 100): -174.17\n"
     ]
    }
   ],
   "source": [
    "import pygame\n",
    "import numpy as np\n",
    "import torch\n",
    "import math\n",
    "from collections import deque\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "import random\n",
    "import os\n",
    "from datetime import datetime\n",
    "\n",
    "##############################\n",
    "# 1. Select device (GPU or CPU)\n",
    "##############################\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(\"Using device:\", device)\n",
    "\n",
    "class SharedControlEnv:\n",
    "    \"\"\"Simplified shared-control environment without LIDAR or obstacles, single goal.\"\"\"\n",
    "    def __init__(self, window_size=(1200, 800), render_mode=None):\n",
    "        self.window_size = window_size\n",
    "        self.render_mode = render_mode\n",
    "\n",
    "        # Environment parameters\n",
    "        self.max_speed = 3\n",
    "        self.dot_radius = 30\n",
    "        self.target_radius = 10\n",
    "        self.goal_detection_radius = self.dot_radius + self.target_radius\n",
    "\n",
    "        # Initialize pygame if rendering\n",
    "        if self.render_mode == 'human':\n",
    "            try:\n",
    "                pygame.init()\n",
    "                self.screen = pygame.display.set_mode(window_size)\n",
    "                self.clock = pygame.time.Clock()\n",
    "            except pygame.error as e:\n",
    "                print(f\"Failed to initialize pygame: {e}\")\n",
    "                self.render_mode = None\n",
    "\n",
    "        # We'll keep some history of states if needed\n",
    "        self.state_history_len = 5\n",
    "        self.state_history = deque(maxlen=self.state_history_len)\n",
    "\n",
    "        # The state is: (dot_x, dot_y, goal_x, goal_y, human_dx, human_dy, gamma)\n",
    "        self.observation_dim = 7\n",
    "\n",
    "        # One-dimensional action: gamma in [0, 1]\n",
    "        self.action_dim = 1\n",
    "\n",
    "        # Environment state\n",
    "        self.dot_pos = None\n",
    "        self.goal_pos = None\n",
    "        self.reached_goal = False\n",
    "        self.current_gamma = 0.2\n",
    "\n",
    "        self.reset()\n",
    "\n",
    "    def get_state(self, human_input):\n",
    "        \"\"\"\n",
    "        Construct the state vector:\n",
    "        [dot_x, dot_y, goal_x, goal_y, h_in_x, h_in_y, gamma].\n",
    "        Positions normalized to [0,1].\n",
    "        Human input normalized by max_speed.\n",
    "        \"\"\"\n",
    "        # Normalize positions\n",
    "        norm_dot_pos = [\n",
    "            self.dot_pos[0] / self.window_size[0],\n",
    "            self.dot_pos[1] / self.window_size[1]\n",
    "        ]\n",
    "        norm_goal_pos = [\n",
    "            self.goal_pos[0] / self.window_size[0],\n",
    "            self.goal_pos[1] / self.window_size[1]\n",
    "        ]\n",
    "\n",
    "        # Normalize human input\n",
    "        norm_human_input = [\n",
    "            np.clip(human_input[0] / self.max_speed, -1, 1),\n",
    "            np.clip(human_input[1] / self.max_speed, -1, 1)\n",
    "        ]\n",
    "\n",
    "        state = np.array([\n",
    "            norm_dot_pos[0],\n",
    "            norm_dot_pos[1],\n",
    "            norm_goal_pos[0],\n",
    "            norm_goal_pos[1],\n",
    "            norm_human_input[0],\n",
    "            norm_human_input[1],\n",
    "            self.current_gamma\n",
    "        ], dtype=np.float32)\n",
    "\n",
    "        return state\n",
    "\n",
    "    def step(self, action, human_input):\n",
    "        \"\"\"\n",
    "        Step the environment forward by one.\n",
    "        action = gamma in [0, 1].\n",
    "        human_input = [dx, dy] with noise.\n",
    "        \"\"\"\n",
    "        # Clip gamma to [0, 1]\n",
    "        action = np.clip(float(action), 0.0, 1.0)\n",
    "        self.current_gamma = action\n",
    "\n",
    "        # Decompose human input\n",
    "        h_dx, h_dy = human_input\n",
    "        h_mag = math.hypot(h_dx, h_dy)\n",
    "        h_dir = [h_dx / h_mag, h_dy / h_mag] if h_mag > 0 else [0, 0]\n",
    "\n",
    "        # Direction to goal\n",
    "        w_dx = self.goal_pos[0] - self.dot_pos[0]\n",
    "        w_dy = self.goal_pos[1] - self.dot_pos[1]\n",
    "        w_mag = math.hypot(w_dx, w_dy)\n",
    "        w_dir = [w_dx / w_mag, w_dy / w_mag] if w_mag > 0 else [0, 0]\n",
    "\n",
    "        # Scale movement by max_speed\n",
    "        step_size = self.max_speed * min(max(h_mag / self.max_speed, 0), 1)\n",
    "\n",
    "        # Weighted movement: gamma for autopilot, (1 - gamma) for human\n",
    "        w_move = [\n",
    "            self.current_gamma * w_dir[0] * step_size,\n",
    "            self.current_gamma * w_dir[1] * step_size\n",
    "        ]\n",
    "        h_move = [\n",
    "            (1 - self.current_gamma) * h_dir[0] * step_size,\n",
    "            (1 - self.current_gamma) * h_dir[1] * step_size\n",
    "        ]\n",
    "\n",
    "        new_pos = [\n",
    "            self.dot_pos[0] + w_move[0] + h_move[0],\n",
    "            self.dot_pos[1] + w_move[1] + h_move[1]\n",
    "        ]\n",
    "\n",
    "        # Clip to window boundaries\n",
    "        self.dot_pos = [\n",
    "            max(0, min(self.window_size[0], new_pos[0])),\n",
    "            max(0, min(self.window_size[1], new_pos[1]))\n",
    "        ]\n",
    "\n",
    "        # Check if goal is reached\n",
    "        dist_to_goal = math.hypot(\n",
    "            self.dot_pos[0] - self.goal_pos[0],\n",
    "            self.dot_pos[1] - self.goal_pos[1]\n",
    "        )\n",
    "        self.reached_goal = (dist_to_goal < self.goal_detection_radius)\n",
    "\n",
    "        # Calculate reward\n",
    "        reward = self._compute_reward(dist_to_goal)\n",
    "\n",
    "        # Build next state\n",
    "        state = self.get_state(human_input)\n",
    "        self.state_history.append(state)\n",
    "\n",
    "        done = self.reached_goal\n",
    "\n",
    "        info = {\n",
    "            'distance_to_goal': dist_to_goal,\n",
    "            'reached_goal': self.reached_goal,\n",
    "            'gamma': self.current_gamma\n",
    "        }\n",
    "\n",
    "        return state, reward, done, info\n",
    "\n",
    "    def _compute_reward(self, dist_to_goal):\n",
    "        \"\"\"\n",
    "        Simple reward function:\n",
    "        - +1 if goal is reached\n",
    "        - -0.01 * dist_to_goal if not reached\n",
    "        - small penalty for deviation from gamma=0.5\n",
    "        \"\"\"\n",
    "        reward = 0.0\n",
    "        if self.reached_goal:\n",
    "            reward += 1.0\n",
    "        else:\n",
    "            reward -= 0.01 * dist_to_goal\n",
    "        reward -= 0.05 * abs(self.current_gamma - 0.5)\n",
    "        return reward\n",
    "\n",
    "    def reset(self):\n",
    "        \"\"\"Reset the environment: dot to center, new random goal.\"\"\"\n",
    "        self.dot_pos = [\n",
    "            self.window_size[0] // 2,\n",
    "            self.window_size[1] // 2\n",
    "        ]\n",
    "        margin = 100\n",
    "        self.goal_pos = [\n",
    "            random.randint(margin, self.window_size[0] - margin),\n",
    "            random.randint(margin, self.window_size[1] - margin)\n",
    "        ]\n",
    "        self.reached_goal = False\n",
    "        self.current_gamma = 0.2\n",
    "        self.state_history.clear()\n",
    "\n",
    "        init_state = self.get_state([0, 0])\n",
    "        self.state_history.append(init_state)\n",
    "        return init_state\n",
    "\n",
    "    def render(self):\n",
    "        \"\"\"Render environment if render_mode == 'human'.\"\"\"\n",
    "        if self.render_mode != 'human':\n",
    "            return\n",
    "        try:\n",
    "            self.screen.fill((255, 255, 255))\n",
    "\n",
    "            # Draw goal\n",
    "            pygame.draw.circle(\n",
    "                self.screen,\n",
    "                (255, 255, 0),\n",
    "                (int(self.goal_pos[0]), int(self.goal_pos[1])),\n",
    "                self.target_radius\n",
    "            )\n",
    "            pygame.draw.circle(\n",
    "                self.screen,\n",
    "                (0, 0, 0),\n",
    "                (int(self.goal_pos[0]), int(self.goal_pos[1])),\n",
    "                self.target_radius + 2, 2\n",
    "            )\n",
    "\n",
    "            # Draw dot\n",
    "            pygame.draw.circle(\n",
    "                self.screen,\n",
    "                (0, 0, 0),\n",
    "                (int(self.dot_pos[0]), int(self.dot_pos[1])),\n",
    "                self.dot_radius, 2\n",
    "            )\n",
    "\n",
    "            # Draw gamma\n",
    "            font = pygame.font.Font(None, 36)\n",
    "            gamma_text = font.render(f'γ: {self.current_gamma:.2f}', True, (0, 0, 0))\n",
    "            self.screen.blit(gamma_text, (10, 10))\n",
    "\n",
    "            pygame.display.flip()\n",
    "            self.clock.tick(60)\n",
    "        except pygame.error as e:\n",
    "            print(f\"Render error: {e}\")\n",
    "            self.render_mode = None\n",
    "\n",
    "    def close(self):\n",
    "        if self.render_mode == 'human':\n",
    "            pygame.quit()\n",
    "\n",
    "\n",
    "class ActorCritic(nn.Module):\n",
    "    \"\"\"Actor-Critic network (PPO-compatible).\"\"\"\n",
    "\n",
    "    def __init__(self, state_dim, action_dim, hidden_dim=128):\n",
    "        super(ActorCritic, self).__init__()\n",
    "\n",
    "        # Shared feature extractor\n",
    "        self.feature_net = nn.Sequential(\n",
    "            nn.Linear(state_dim, hidden_dim),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_dim, hidden_dim),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "\n",
    "        # Actor network\n",
    "        self.actor_mean = nn.Linear(hidden_dim, action_dim)\n",
    "        self.actor_log_std = nn.Parameter(torch.zeros(1, action_dim))\n",
    "\n",
    "        # Critic network\n",
    "        self.critic = nn.Linear(hidden_dim, 1)\n",
    "\n",
    "        # Initialize weights\n",
    "        self.apply(self._init_weights)\n",
    "\n",
    "    def _init_weights(self, module):\n",
    "        if isinstance(module, nn.Linear):\n",
    "            nn.init.orthogonal_(module.weight, gain=np.sqrt(2))\n",
    "            if module.bias is not None:\n",
    "                module.bias.data.zero_()\n",
    "\n",
    "    def forward(self, state):\n",
    "        features = self.feature_net(state)\n",
    "        action_mean = torch.sigmoid(self.actor_mean(features))  # ensures [0,1]\n",
    "        action_std = torch.exp(self.actor_log_std)\n",
    "        value = self.critic(features)\n",
    "        return action_mean, action_std, value\n",
    "\n",
    "    def get_action_distribution(self, state):\n",
    "        action_mean, action_std, _ = self(state)\n",
    "        return torch.distributions.Normal(action_mean, action_std)\n",
    "\n",
    "\n",
    "class PPOSharedControl:\n",
    "    \"\"\"Proximal Policy Optimization for the shared-control environment.\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self, \n",
    "        state_dim,\n",
    "        action_dim,\n",
    "        hidden_dim=128,\n",
    "        lr=1e-4,        # lower LR\n",
    "        gamma=0.99,\n",
    "        epsilon=0.2,\n",
    "        c1=1.0,\n",
    "        c2=0.01\n",
    "    ):\n",
    "        self.actor_critic = ActorCritic(state_dim, action_dim, hidden_dim).to(device)\n",
    "        self.optimizer = optim.Adam(self.actor_critic.parameters(), lr=lr)\n",
    "\n",
    "        self.gamma = gamma\n",
    "        self.epsilon = epsilon\n",
    "        self.c1 = c1\n",
    "        self.c2 = c2\n",
    "\n",
    "    def get_action(self, state):\n",
    "        \"\"\"Sample an action (gamma) from the current policy.\"\"\"\n",
    "        # Move state to GPU if available\n",
    "        state = state.to(device)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            dist = self.actor_critic.get_action_distribution(state)\n",
    "            action = dist.sample()\n",
    "            log_prob = dist.log_prob(action)\n",
    "            # Clip action to [0,1]\n",
    "            action = torch.clamp(action, 0.0, 1.0)\n",
    "        return action, log_prob\n",
    "\n",
    "    def get_value(self, state):\n",
    "        \"\"\"Get the critic's value estimate for a state.\"\"\"\n",
    "        state = state.to(device)\n",
    "        with torch.no_grad():\n",
    "            _, _, value = self.actor_critic(state)\n",
    "        return value\n",
    "\n",
    "    def update(self, states, actions, old_log_probs, returns, advantages,\n",
    "               epochs=10, batch_size=64):\n",
    "        # Move data to GPU\n",
    "        states = torch.FloatTensor(states).to(device)\n",
    "        actions = torch.FloatTensor(actions).to(device)\n",
    "        old_log_probs = torch.FloatTensor(old_log_probs).to(device)\n",
    "        returns = torch.FloatTensor(returns).to(device)\n",
    "        advantages = torch.FloatTensor(advantages).to(device)\n",
    "\n",
    "        # Normalize advantages safely\n",
    "        adv_std = advantages.std()\n",
    "        if adv_std < 1e-8:\n",
    "            advantages = advantages - advantages.mean()\n",
    "        else:\n",
    "            advantages = (advantages - advantages.mean()) / (adv_std + 1e-8)\n",
    "\n",
    "        for _ in range(epochs):\n",
    "            indices = torch.randperm(len(states))\n",
    "            for start_idx in range(0, len(states), batch_size):\n",
    "                idx = indices[start_idx : start_idx + batch_size]\n",
    "\n",
    "                batch_states = states[idx]\n",
    "                batch_actions = actions[idx]\n",
    "                batch_old_log_probs = old_log_probs[idx]\n",
    "                batch_returns = returns[idx]\n",
    "                batch_advantages = advantages[idx]\n",
    "\n",
    "                dist = self.actor_critic.get_action_distribution(batch_states)\n",
    "                _, _, values = self.actor_critic(batch_states)\n",
    "\n",
    "                new_log_probs = dist.log_prob(batch_actions)\n",
    "                ratio = torch.exp(new_log_probs - batch_old_log_probs)\n",
    "\n",
    "                # PPO objective\n",
    "                surr1 = ratio * batch_advantages\n",
    "                surr2 = torch.clamp(\n",
    "                    ratio,\n",
    "                    1.0 - self.epsilon,\n",
    "                    1.0 + self.epsilon\n",
    "                ) * batch_advantages\n",
    "                policy_loss = -torch.min(surr1, surr2).mean()\n",
    "\n",
    "                # Value function loss, flatten shapes\n",
    "                value_loss = F.mse_loss(values.view(-1), batch_returns.view(-1))\n",
    "\n",
    "                # Entropy bonus\n",
    "                entropy = dist.entropy().mean()\n",
    "\n",
    "                # Total loss\n",
    "                total_loss = policy_loss + self.c1 * value_loss - self.c2 * entropy\n",
    "\n",
    "                self.optimizer.zero_grad()\n",
    "                total_loss.backward()\n",
    "                # Stronger grad clipping\n",
    "                torch.nn.utils.clip_grad_norm_(self.actor_critic.parameters(), max_norm=0.1)\n",
    "                self.optimizer.step()\n",
    "\n",
    "    def save(self, path):\n",
    "        os.makedirs(os.path.dirname(path), exist_ok=True)\n",
    "        torch.save({\n",
    "            'actor_critic_state_dict': self.actor_critic.state_dict(),\n",
    "            'optimizer_state_dict': self.optimizer.state_dict(),\n",
    "        }, path)\n",
    "\n",
    "    def load(self, path):\n",
    "        checkpoint = torch.load(path, map_location=device)\n",
    "        self.actor_critic.load_state_dict(checkpoint['actor_critic_state_dict'])\n",
    "        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "\n",
    "\n",
    "def simulate_human_input(env):\n",
    "    \"\"\"\n",
    "    Simulate much noisier human input directed (roughly) toward the goal, \n",
    "    but not so large as to cause extreme instability.\n",
    "    \"\"\"\n",
    "    dx = env.goal_pos[0] - env.dot_pos[0]\n",
    "    dy = env.goal_pos[1] - env.dot_pos[1]\n",
    "\n",
    "    # Use a smaller std dev than 1000 to reduce instability\n",
    "    dx += np.random.normal(0, 10)\n",
    "    dy += np.random.normal(0, 10)\n",
    "\n",
    "    mag = math.hypot(dx, dy)\n",
    "    if mag > 0:\n",
    "        dx = dx / mag * env.max_speed\n",
    "        dy = dy / mag * env.max_speed\n",
    "\n",
    "    return np.array([dx, dy], dtype=np.float32)\n",
    "\n",
    "\n",
    "def compute_returns(rewards, gamma):\n",
    "    \"\"\"\n",
    "    Compute discounted returns, then safely normalize them.\n",
    "    \"\"\"\n",
    "    returns = []\n",
    "    R = 0\n",
    "    for r in reversed(rewards):\n",
    "        R = r + gamma * R\n",
    "        returns.insert(0, R)\n",
    "    returns = torch.tensor(returns, dtype=torch.float32)\n",
    "\n",
    "    # Safe normalization\n",
    "    returns_std = returns.std()\n",
    "    if returns_std < 1e-8:\n",
    "        returns = returns - returns.mean()\n",
    "    else:\n",
    "        returns = (returns - returns.mean()) / (returns_std + 1e-5)\n",
    "\n",
    "    return returns\n",
    "\n",
    "\n",
    "def compute_advantages(returns, states, agent):\n",
    "    \"\"\"\n",
    "    Advantages = returns - value estimates, done carefully to avoid NaNs.\n",
    "    \"\"\"\n",
    "    states_tensor = torch.FloatTensor(states)\n",
    "    values = agent.get_value(states_tensor).detach().squeeze()\n",
    "\n",
    "    # Move values back to CPU if needed\n",
    "    values = values.cpu()\n",
    "    advantages = returns - values\n",
    "\n",
    "    # If you want to do safe normalization here as well, you can,\n",
    "    # but it's also done inside agent.update(...) after concatenation.\n",
    "    return advantages\n",
    "\n",
    "\n",
    "def train_ppo(env, episodes=500, steps_per_episode=300, checkpoint_freq=50):\n",
    "    \"\"\"\n",
    "    Train a PPO agent on the simplified SharedControlEnv.\n",
    "    \"\"\"\n",
    "    state_dim = env.observation_dim\n",
    "    action_dim = env.action_dim\n",
    "\n",
    "    agent = PPOSharedControl(state_dim, action_dim)\n",
    "\n",
    "    # Create directory for checkpoints\n",
    "    checkpoint_dir = f'checkpoints_{datetime.now().strftime(\"%Y%m%d_%H%M%S\")}'\n",
    "    os.makedirs(checkpoint_dir, exist_ok=True)\n",
    "\n",
    "    best_reward = float('-inf')\n",
    "    episode_rewards = []\n",
    "\n",
    "    for episode in range(episodes):\n",
    "        state = env.reset()\n",
    "        episode_reward = 0\n",
    "\n",
    "        # Collect episode experience\n",
    "        states = []\n",
    "        actions = []\n",
    "        rewards = []\n",
    "        log_probs = []\n",
    "\n",
    "        for step in range(steps_per_episode):\n",
    "            state_tensor = torch.FloatTensor(state).unsqueeze(0)\n",
    "            human_input = simulate_human_input(env)\n",
    "            action, log_prob = agent.get_action(state_tensor)\n",
    "\n",
    "            # Convert action to a float for environment step\n",
    "            next_state, reward, done, info = env.step(action.item(), human_input)\n",
    "\n",
    "            states.append(state)\n",
    "            actions.append(action.squeeze().cpu().numpy())\n",
    "            rewards.append(reward)\n",
    "            log_probs.append(log_prob.squeeze().cpu().numpy())\n",
    "\n",
    "            episode_reward += reward\n",
    "            state = next_state\n",
    "\n",
    "            if env.render_mode == 'human':\n",
    "                env.render()\n",
    "\n",
    "            if done:\n",
    "                break\n",
    "\n",
    "        episode_rewards.append(episode_reward)\n",
    "\n",
    "        # Convert to arrays\n",
    "        states = np.array(states, dtype=np.float32)\n",
    "        actions = np.array(actions, dtype=np.float32)\n",
    "        rewards = np.array(rewards, dtype=np.float32)\n",
    "        log_probs = np.array(log_probs, dtype=np.float32)\n",
    "\n",
    "        # Compute returns and advantages\n",
    "        returns = compute_returns(rewards, agent.gamma)\n",
    "        advantages = compute_advantages(returns, states, agent)\n",
    "\n",
    "        # Update PPO\n",
    "        agent.update(states, actions, log_probs, returns, advantages)\n",
    "\n",
    "        # Checkpoint if best\n",
    "        if episode_reward > best_reward:\n",
    "            best_reward = episode_reward\n",
    "            agent.save(os.path.join(checkpoint_dir, 'best_model.pth'))\n",
    "\n",
    "        # Regular checkpoint\n",
    "        if (episode + 1) % checkpoint_freq == 0:\n",
    "            agent.save(os.path.join(checkpoint_dir, f'checkpoint_{episode+1}.pth'))\n",
    "\n",
    "        # Print progress\n",
    "        avg_reward = np.mean(episode_rewards[-100:]) if len(episode_rewards) >= 100 \\\n",
    "            else np.mean(episode_rewards)\n",
    "        print(f\"Episode {episode+1}/{episodes}, Reward: {episode_reward:.2f}, \"\n",
    "              f\"Avg Reward (last 100): {avg_reward:.2f}\")\n",
    "\n",
    "        # Early stop if solved\n",
    "        if avg_reward > 200 and len(episode_rewards) >= 100:\n",
    "            print(\"Environment solved!\")\n",
    "            agent.save(os.path.join(checkpoint_dir, 'solved_model.pth'))\n",
    "            break\n",
    "\n",
    "    return agent, episode_rewards\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    env = SharedControlEnv(render_mode='human')\n",
    "    try:\n",
    "        agent, rewards_history = train_ppo(env)\n",
    "        agent.save('final_model.pth')\n",
    "        np.save('training_rewards.npy', np.array(rewards_history))\n",
    "    except KeyboardInterrupt:\n",
    "        print(\"\\nTraining interrupted by user.\")\n",
    "    except Exception as e:\n",
    "        print(f\"Error during training: {e}\")\n",
    "    finally:\n",
    "        env.close()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}