6_7920_Fall_2024_Homework_5_problem_2.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/hanklin3/code/blob/master/6_7920_Fall_2024_Homework_5_problem_2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Setup\n",
        "\n"
      ],
      "metadata": {
        "id": "wM6ks-aGKjJx"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this problem, we will consider the discounted stochastic LQ control problem, as described in Lecture 3. Specifically, we will consider using Q-learning to approximate the Q-function of this problem. Below, we provide an instant of this problem through the environment **LQEnv**. The environment is a python class, whose interface is very similar to that of OpenAI gym."
      ],
      "metadata": {
        "id": "fXiBmxYfofy3"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The code block below imports important and relevant packages, defnies the environment, and includes helper functions like a policy evaluation function."
      ],
      "metadata": {
        "id": "TSSjTKeTn29X"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Kerb: hanklin"
      ],
      "metadata": {
        "id": "a-hx6DyDvwkH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from tqdm import tqdm\n",
        "import gym\n",
        "from gym import error, spaces, utils\n",
        "from gym.utils import seeding\n",
        "\n",
        "dtype = np.float32\n",
        "class LQEnv(gym.Env):\n",
        "  def __init__(self, n, m, seed = 0, gamma = 0.9, sigma = 0.2):\n",
        "    \"\"\"\n",
        "    The LQ environment, initialized with the parameters\n",
        "    n: state dimension\n",
        "    m: action dimension\n",
        "    seed: random seed\n",
        "    gamma: discount factor\n",
        "    sigma: noise variance\n",
        "    \"\"\"\n",
        "    np.random.seed(seed)\n",
        "    self.n = n\n",
        "    self.action_space = spaces.Box(low = -np.inf, high = np.inf, shape =  [m,1], dtype = dtype )\n",
        "    self.observation_space = spaces.Box(low = -np.inf, high = np.inf, shape =  [n,1], dtype = dtype )\n",
        "    self.A = np.array([[0.0488135 , 0.21518937],\n",
        "                       [0.10276338, 0.04488318]\n",
        "                        ])\n",
        "    self.B =  np.array([[-0.0763452 ,  0.14589411],\n",
        "                        [-0.06241279,  0.391773  ]\n",
        "                       ])\n",
        "    self.R =  np.array([[1.57567331, 0.96575621],\n",
        "                        [0.96575621, 1.40655837]\n",
        "                        ])\n",
        "    self.Q =  np.array([[1.67940376, 0.12099823],\n",
        "                        [0.12099823, 0.51263764]\n",
        "                        ])\n",
        "    self.state = self.observation_space.sample()\n",
        "    self.sigma = sigma\n",
        "    self.gamma = gamma\n",
        "\n",
        "  def step(self, action):\n",
        "    \"\"\"\n",
        "    This method is the primary interface between environment and agent.\n",
        "    Paramters:\n",
        "        action: array of shape [m,1]\n",
        "\n",
        "    Returns:\n",
        "        output: (next state:array, cost:float, done:bool, None)\n",
        "\n",
        "    \"\"\"\n",
        "    err_msg = f\"{action!r} ({action.shape}) ({type(action)}) invalid\"\n",
        "    # assert self.action_space.contains(action), err_msg\n",
        "    assert self.state is not None, \"Call reset before using step method.\"\n",
        "    w = self.sigma*np.random.randn()\n",
        "\n",
        "    cost = (self.state.T.dot(self.Q)).dot(self.state) +  (action.T.dot(self.R)).dot(action)\n",
        "\n",
        "    self.state = self.A.dot(self.state) + self.B.dot(action) + w\n",
        "\n",
        "    done = False\n",
        "\n",
        "    return self.state, cost, done, {}\n",
        "\n",
        "  def reset(self, state = None):\n",
        "    \"\"\"\n",
        "    This method resets the environment to its initial values.\n",
        "    Paramters:\n",
        "        state array\n",
        "            set the env to specifc state (optional)\n",
        "    Returns:\n",
        "        observation:    array\n",
        "                        the initial state of the environment\n",
        "    \"\"\"\n",
        "    if state is None:\n",
        "        # sample from a guassian with zero mean and std of 10\n",
        "        self.state = 10*self.observation_space.sample()\n",
        "    else:\n",
        "        self.state = state\n",
        "    return  self.state\n",
        "\n",
        "  def close(self):\n",
        "    \"\"\"\n",
        "    This method provides the user with the option to perform any necessary cleanup.\n",
        "    \"\"\"\n",
        "    pass\n",
        "\n",
        "  def sample_random_action(self):\n",
        "      \"\"\"\n",
        "      sample actions from a normal dist with mean zero and std 10\n",
        "      \"\"\"\n",
        "      return  10*env.action_space.sample()\n",
        "\n",
        "\n",
        "def policy_evaluation(policy, env, T = 1000, N = 5):\n",
        "    \"\"\"\n",
        "    This method evaluate the performance of a specifc policy on the env through simulating\n",
        "    it on N trajectories each of length T.\n",
        "    Paramters:\n",
        "        policy: function that takes in state and return action\n",
        "        env: instant of LQEnv\n",
        "        T: number of iterations per trajectory\n",
        "        N: number of trajectories\n",
        "    Returns:\n",
        "        output: mean discounted total cost\n",
        "\n",
        "    \"\"\"\n",
        "    costs = []\n",
        "    for ite in tqdm(range(N)):\n",
        "        state = env.reset()\n",
        "        gamma = env.gamma\n",
        "        total_costs = 0\n",
        "        for t in range(T):\n",
        "            action = policy(state)\n",
        "            # print('\\npolicy_evaluation:action', action.shape, action)\n",
        "            state, cost, _, _ = env.step(action)\n",
        "            total_costs += cost * gamma**(t)\n",
        "        costs.append(total_costs)\n",
        "    return np.mean(costs)\n",
        "\n",
        "def lin_policy(K):\n",
        "    \"\"\"\n",
        "    helper function to define linear policies of the form u = L@x\n",
        "    \"\"\"\n",
        "    def policy(state):\n",
        "        return K.dot(state).astype(dtype)\n",
        "    return policy\n",
        "\n",
        "\n",
        "def Q_a(x,u, theta):\n",
        "    \"\"\"\n",
        "    Q-function parameterized by the tuple/list theta\n",
        "    \"\"\"\n",
        "    q =  x.T@theta[0].T@theta[0]@x\n",
        "    q+= u.T@theta[1].T@theta[1]@u\n",
        "    q+= 2*x.T@theta[2]@u + theta[3]\n",
        "    return q[0,0]\n",
        "\n"
      ],
      "metadata": {
        "id": "UAiHtkQUKpsv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Initialize environemnt\n",
        "\n",
        "Below, we initialize the environment that you need to interact with in this excercise."
      ],
      "metadata": {
        "id": "LmdZmBU4qwcf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "n = 2\n",
        "m = 2\n",
        "env = LQEnv(n,m)"
      ],
      "metadata": {
        "id": "CE8fn-nEqvVx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### (a) Assuming knowledge of the system matrices $A, B, R,$ and $Q$, compute $P$, the solution to the appropriate form of the Riccati equation, and the matrix $K$ that characterizes the optimal policy. Evaluate the policy performance over 100 trajectories each with length 1000.\n"
      ],
      "metadata": {
        "id": "VwnEuE3Jqmtg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "from scipy.linalg import solve_discrete_are\n",
        "\n",
        "# Solve the discrete-time algebraic Riccati equation\n",
        "P = solve_discrete_are(env.A, env.B, env.Q, env.R)\n",
        "\n",
        "# Compute the optimal policy L = -(B^T P B + R)^-1 B^T P A\n",
        "K = -np.linalg.inv(env.B.T @ P @ env.B + env.R) @ (env.B.T @ P @ env.A)\n",
        "\n",
        "\n",
        "print(\"Solution to the Riccati equation (P):\")\n",
        "print(P)\n",
        "print(\"\\nOptimal policy matrix (K):\")\n",
        "print(K)\n",
        "\n",
        "def optimal_policy(K):\n",
        "  def policy(state):\n",
        "     return -K_optimal @ state\n",
        "  return policy  # Using the optimal gain matrix K\n",
        "\n",
        "K = K.astype(dtype)\n",
        "# policy = lin_policy(K)\n",
        "policy = optimal_policy(K)\n",
        "print('policy', policy)\n",
        "\n",
        "# cost_mean = policy_evaluation(policy, env, T = 1000, N = 5)\n",
        "cost_mean = policy_evaluation(policy, env, T = 1000, N = 100)\n",
        "print('cost_mean', cost_mean)\n"
      ],
      "metadata": {
        "id": "31jSaZK4Kov0",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "92983816-e964-4ee5-848d-631d657dac0a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/ipykernel/ipkernel.py:283: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
            "  and should_run_async(code)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Solution to the Riccati equation (P):\n",
            "[[1.68872555 0.14018861]\n",
            " [0.14018861 0.58521217]]\n",
            "\n",
            "Optimal policy matrix (K):\n",
            "[[ 0.03457423  0.07474569]\n",
            " [-0.04677614 -0.09387158]]\n",
            "policy <function optimal_policy.<locals>.policy at 0x7b0e42b075b0>\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 100/100 [00:15<00:00,  6.59it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "cost_mean 244.43228610055655\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print('cost_mean', cost_mean)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AUIXroENSvYM",
        "outputId": "d63d1951-f915-418a-dbff-cde3fb339886"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "cost_mean 244.43228610055655\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a1 = env.action_space.sample()\n",
        "print(a1.shape, type(a1), a1)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uHXqZw-tOiz7",
        "outputId": "37c4e2dd-85c4-439c-b7e3-6dafe44b3f3b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(2, 1) <class 'numpy.ndarray'> [[-1.5008005 ]\n",
            " [ 0.45925373]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### (b) Given an arbitrary $\\Theta_0$, write the greedy policy with respect to $Q(x; u; \\Theta_0)$ in terms of the parameters $\\Theta_0$. What is the form of this greedy policy?"
      ],
      "metadata": {
        "id": "9LZ9Z16I3vPj"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "The greedy policy minimizes the Q-function with respect to the control input u.\n",
        "\n",
        "$dQ(x;u;\\theta_0)/du$ = 2$\\theta_2^T$$\\theta_2$$u$+2$\\theta_3^T$$x$ = 0\n",
        "\n",
        "$u = -(θ_{2}^Tθ_{2})^{-1}θ_{3}^Tx$"
      ],
      "metadata": {
        "id": "Bf912Htz39Df"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### (c) Implement Q-learning for this problem, where actions are chosen using the greedy policy (wrt to the current Q-function). Then train using samples chosen from a single arbitrarily long trajectory. Use the given function to initialize theta.\n",
        "\n",
        "Hint: If your states/parameters grow very large during training, consider making the learning rate very small"
      ],
      "metadata": {
        "id": "d38u384zrqVq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def init_theta(n,m, fun = np.random.randn, seed = 1):\n",
        "    np.random.seed(seed)\n",
        "    A = fun(n,n)\n",
        "    B = fun(m,m)\n",
        "    C = fun(n,m)\n",
        "    const = 0\n",
        "    return [A,B,C, const]\n"
      ],
      "metadata": {
        "id": "R5gRkrXOaEWb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "\n",
        "def Q_function(x, u, theta1, theta2, theta3, theta4):\n",
        "    \"\"\"Computes the quadratic approximation of the Q-function.\"\"\"\n",
        "    return (x.T @ theta1.T @ theta1 @ x + u.T @ theta2.T @ theta2 @ u\n",
        "            + 2 * x.T @ theta3 @ u + theta4)\n",
        "\n",
        "def greedy_policy(x, theta2, theta3, action_space):\n",
        "    \"\"\"Computes the greedy policy u based on the current parameters.\"\"\"\n",
        "    # u = -np.linalg.inv(theta2.T @ theta2) @ (theta3.T @ x)\n",
        "\n",
        "    # # Compute the matrix to invert with regularization to avoid singular\n",
        "    # regularized_matrix = theta2.T @ theta2 + reg * np.eye(theta2.shape[1])\n",
        "    # # Calculate the greedy action\n",
        "    # u = -np.linalg.inv(regularized_matrix) @ (theta3.T @ x)\n",
        "\n",
        "    # avoid singluar using pseudo inverse\n",
        "    u = -np.linalg.pinv(theta2.T @ theta2) @ (theta3.T @ x)\n",
        "\n",
        "    return u\n",
        "\n",
        "\n",
        "def Q_learning(env, theta1, theta2, theta3, theta4, alpha=0.001, gamma=0.9, num_steps=10000):\n",
        "    \"\"\"Q-learning with function approximation for the LQ control problem.\"\"\"\n",
        "    state = env.reset()  # Initialize the environment\n",
        "    for step in range(num_steps):\n",
        "        # Select action using greedy policy\n",
        "        action = greedy_policy(state, theta2, theta3, env.action_space)\n",
        "        # print(\"Action shape:\", action.shape)\n",
        "\n",
        "        # Perform the action and observe the next state and reward\n",
        "        next_state, reward, is_done, _ = env.step(action)\n",
        "\n",
        "        # Compute Q(s, a; Θ)\n",
        "        current_Q = Q_function(state, action, theta1, theta2, theta3, theta4)\n",
        "\n",
        "        # Compute the greedy action for the next state\n",
        "        next_action = greedy_policy(next_state, theta2, theta3, env.action_space)\n",
        "\n",
        "        # Compute Q(s', a'; Θ)\n",
        "        next_Q = Q_function(next_state, next_action, theta1, theta2, theta3, theta4)\n",
        "\n",
        "        # Update Q-learning parameters using the update rule\n",
        "        td_error = reward + gamma * next_Q - current_Q\n",
        "\n",
        "        # Update parameters\n",
        "        theta1 += alpha * td_error * (np.outer(state, state) @ theta1)\n",
        "        theta2 += alpha * td_error * (np.outer(action, action) @ theta2)\n",
        "        theta3 += alpha * td_error * (np.outer(state, action))\n",
        "        theta4 += alpha * td_error\n",
        "\n",
        "\n",
        "        # Move to the next state\n",
        "        state = next_state\n",
        "\n",
        "    return theta1, theta2, theta3, theta4\n",
        "\n",
        "# Initialize the environment and parameters\n",
        "env = LQEnv(n=2, m=2)\n",
        "theta1, theta2, theta3, theta4 = init_theta(2, 2)\n",
        "print(\"Initial Parameters:\")\n",
        "print(\"Theta1:\", theta1)\n",
        "print(\"Theta2:\", theta2)\n",
        "print(\"Theta3:\", theta3)\n",
        "print(\"Theta4:\", theta4)\n",
        "\n",
        "# Run Q-learning\n",
        "theta1, theta2, theta3, theta4 = Q_learning(env, theta1, theta2, theta3, theta4)\n",
        "\n",
        "# Evaluate the resulting policy\n",
        "print(\"Updated Parameters:\")\n",
        "print(\"Theta1:\", theta1)\n",
        "print(\"Theta2:\", theta2)\n",
        "print(\"Theta3:\", theta3)\n",
        "print(\"Theta4:\", theta4)"
      ],
      "metadata": {
        "id": "Arp3q5lNQNeu",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "744db0b8-2c30-4536-9d70-4ce84ee25ce8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Initial Parameters:\n",
            "Theta1: [[ 1.62434536 -0.61175641]\n",
            " [-0.52817175 -1.07296862]]\n",
            "Theta2: [[ 0.86540763 -2.3015387 ]\n",
            " [ 1.74481176 -0.7612069 ]]\n",
            "Theta3: [[ 0.3190391  -0.24937038]\n",
            " [ 1.46210794 -2.06014071]]\n",
            "Theta4: 0\n",
            "Updated Parameters:\n",
            "Theta1: [[ 1.68340042 -0.70572002]\n",
            " [-0.40154793 -1.05705392]]\n",
            "Theta2: [[ 0.86740957 -2.30426342]\n",
            " [ 1.73912616 -0.75998766]]\n",
            "Theta3: [[ 0.30515654 -0.23440247]\n",
            " [ 1.4641209  -2.07320247]]\n",
            "Theta4: [[0.63044645]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Do the parameters converge? Is the resulting policy the same as the policy obtained in part (a)? If not, why?\n"
      ],
      "metadata": {
        "id": "hGXFkSHExnoP"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Optimal feedback gain matrix K from Part a (using the Riccati equation)\n",
        "K_optimal = -np.linalg.inv(env.B.T @ P @ env.B + env.R) @ (env.B.T @ P @ env.A)\n",
        "\n",
        "# Learned parameters from Q-learning in Part c\n",
        "theta_2_learned = theta2  # Obtained from Q-learning\n",
        "theta_3_learned = theta3  # Obtained from Q-learning\n",
        "\n",
        "# Compute the learned policy gain matrix K_learned\n",
        "K_learned = np.linalg.inv(theta_2_learned.T @ theta_2_learned) @ theta_3_learned.T\n",
        "\n",
        "# Compare the two gain matrices\n",
        "diff = np.linalg.norm(K_learned - K_optimal, 'fro')\n",
        "\n",
        "print(\"Optimal K (from Part a):\")\n",
        "print(K_optimal)\n",
        "\n",
        "print(\"\\nLearned K (from Part c):\")\n",
        "print(K_learned)\n",
        "\n",
        "print(\"\\nDifference between K_optimal and K_learned (Frobenius norm):\", diff)\n",
        "\n",
        "# Threshold to determine if policies are similar\n",
        "if diff < 1e-3:  # You can adjust this threshold based on precision needs\n",
        "    print(\"\\nThe learned policy is approximately the same as the optimal policy from Part a.\")\n",
        "else:\n",
        "    print(\"\\nThe learned policy is different from the optimal policy from Part a.\")\n",
        "\n",
        "Answer = 'It did not converge likely due to it did not explore all states to find the optimal policy'\n",
        "print(Answer  )\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4_yecMupgKAo",
        "outputId": "702e1dad-b99d-47f2-e139-3ee012b20a16"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Optimal K (from Part a):\n",
            "[[ 0.03457423  0.07474569]\n",
            " [-0.04677614 -0.09387158]]\n",
            "\n",
            "Learned K (from Part c):\n",
            "[[ 0.09082672  0.15482267]\n",
            " [ 0.01141178 -0.26483182]]\n",
            "\n",
            "Difference between K_optimal and K_learned (Frobenius norm): 0.20540180977561184\n",
            "\n",
            "The learned policy is different from the optimal policy from Part a.\n",
            "It did not converge likely due to it did not explore all states to find the optimal policy\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Answer = 'It did not converge likely due to it did not explore all states to find the optimal policy, since we always do greedy and never explored.'"
      ],
      "metadata": {
        "id": "5ojzzW6DxscC"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### (d) Train using  samples chosen from multiple trajectories. Each with length 25.  Also choose actions using an $\\epsilon$-greedy policy for $\\epsilon = 0.3$. Do the parameters converge? Is the resulting policy the same as the policy obtained in part (a)? Evaluate the policy performance over 100 trajectories each with length 1000 and compare its average cost with that of the optimal policy.\n",
        "\n",
        "Hint: to sample actions randomly, use the method: env.sample_random_action()"
      ],
      "metadata": {
        "id": "Fhov62jOz28r"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "\n",
        "def epsilon_greedy_policy(x, theta1, theta2, theta3, epsilon, action_space):\n",
        "    \"\"\"Chooses an action using an epsilon-greedy policy.\"\"\"\n",
        "    if np.random.rand() < epsilon:\n",
        "        # Random action\n",
        "        return env.action_space.sample()\n",
        "    else:\n",
        "        # Greedy action\n",
        "        return greedy_policy(x, theta2, theta3, action_space)\n",
        "\n",
        "def Q_learning_multiple_trajectories(env, theta1, theta2, theta3, theta4, epsilon=0.3,\n",
        "                                     alpha=1e-8, gamma=0.9, num_trajectories=100, trajectory_length=25):\n",
        "    \"\"\"Q-learning with function approximation and multiple trajectories.\"\"\"\n",
        "    for trajectory in range(num_trajectories):\n",
        "        state = env.reset()  # Reset the environment at the start of each trajectory\n",
        "\n",
        "        for t in range(trajectory_length):\n",
        "            # Choose action using epsilon-greedy policy\n",
        "            action = epsilon_greedy_policy(state, theta1, theta2, theta3, epsilon, env.action_space)\n",
        "\n",
        "            # Perform action and observe the next state and reward\n",
        "            next_state, reward, is_done, _  = env.step(action)\n",
        "\n",
        "            # Compute Q(s, a; Θ)\n",
        "            current_Q = Q_function(state, action, theta1, theta2, theta3, theta4)\n",
        "\n",
        "            # Choose next action using greedy policy\n",
        "            next_action = greedy_policy(next_state, theta2, theta3, env.action_space)\n",
        "\n",
        "            # Compute Q(s', a'; Θ)\n",
        "            next_Q = Q_function(next_state, next_action, theta1, theta2, theta3, theta4)\n",
        "\n",
        "            # Compute TD error\n",
        "            td_error = reward + gamma * next_Q - current_Q\n",
        "\n",
        "            # # Update parameters using the TD error\n",
        "            # theta1 += alpha * td_error * (np.outer(state, state) @ theta1)  # Update rule for theta1\n",
        "            # theta2 += alpha * td_error * (np.outer(action, action) @ theta2)  # Update rule for theta2\n",
        "            # theta3 += alpha * td_error * (np.outer(state, action))  # Update rule for theta3\n",
        "            # theta4 += alpha * td_error * 1.0  # Update rule for theta4\n",
        "\n",
        "            # Clip the TD error to avoid thetas converges to NaN\n",
        "            clipped_td_error = np.clip(td_error, -1.0, 1.0)\n",
        "\n",
        "            # Update parameters\n",
        "            theta1 += alpha * clipped_td_error * (np.outer(state, state) @ theta1)\n",
        "            theta2 += alpha * clipped_td_error * (np.outer(action, action) @ theta2)\n",
        "            theta3 += alpha * clipped_td_error * (np.outer(state, action))\n",
        "            theta4 += alpha * clipped_td_error\n",
        "\n",
        "            # Move to the next state\n",
        "            state = next_state\n",
        "    return theta1, theta2, theta3, theta4"
      ],
      "metadata": {
        "id": "YVVYO38vXV7A"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import time\n",
        "\n",
        "def runs(epsilon, theta_seed=1):\n",
        "  n = 2\n",
        "  m = 2\n",
        "  env = LQEnv(n, m)\n",
        "  theta1, theta2, theta3, theta4 = init_theta(n, m, seed=theta_seed)\n",
        "\n",
        "  # Run Q-learning with multiple trajectories\n",
        "  time_start = time.time()\n",
        "  theta1, theta2, theta3, theta4 = Q_learning_multiple_trajectories(\n",
        "      env, theta1, theta2, theta3, theta4,\n",
        "      num_trajectories=100, trajectory_length=1000,\n",
        "      alpha=0.00001, epsilon=epsilon)\n",
        "  time_end = time.time()\n",
        "  print(\"Updated Parameters:\")\n",
        "  print(\"Theta1:\", theta1)\n",
        "  print(\"Theta2:\", theta2)\n",
        "  print(\"Theta3:\", theta3)\n",
        "  print(\"Theta4:\", theta4)\n",
        "\n",
        "  return theta1, theta2, theta3, theta4, time_end - time_start\n",
        "\n",
        "def optimal_cost(n = 2, m = 2):\n",
        "  env = LQEnv(n, m)\n",
        "  # Solve the discrete-time algebraic Riccati equation\n",
        "  P = solve_discrete_are(env.A, env.B, env.Q, env.R)\n",
        "  # Compare the resulting policy with the optimal policy\n",
        "  K_optimal = -np.linalg.inv(env.B.T @ P @ env.B + env.R) @ (env.B.T @ P @ env.A)  # Optimal policy from Part a\n",
        "  policy = lin_policy(K_optimal)\n",
        "  cost_mean = policy_evaluation(policy, env, T = 1000, N = 100)\n",
        "  return cost_mean\n",
        "\n",
        "def learned_cost(theta1, theta2, theta3, theta4):\n",
        "  K_learned = np.linalg.inv(theta2.T @ theta2) @ theta3.T  # Learned policy from Q-learning\n",
        "  policy = lin_policy(K_optimal)\n",
        "  cost_mean = policy_evaluation(policy, env, T = 1000, N = 100)\n",
        "  return cost_mean\n",
        "\n",
        "def print_compare_cost(epsilon, avg_cost_optimal, avg_cost_learned):\n",
        "  print('\\n')\n",
        "  print(\"Epsilon:\", epsilon)\n",
        "  # Print the results\n",
        "  print(\"Average Cost of Optimal Policy:\", avg_cost_optimal)\n",
        "  print(\"Average Cost of Learned Policy:\", avg_cost_learned)\n",
        "\n",
        "  # Compare the two average costs\n",
        "  if avg_cost_learned < avg_cost_optimal:\n",
        "      print(\"The learned policy is better than the optimal policy.\")\n",
        "  else:\n",
        "      print(\"The learned policy is worse than the optimal policy.\")\n",
        "\n",
        "\n",
        "def print_compare_policy(theta2, theta3, n = 2, m = 2):\n",
        "  # Compare the resulting policy with the optimal policy\n",
        "  env = LQEnv(n, m)\n",
        "  K_optimal = -np.linalg.inv(env.B.T @ P @ env.B + env.R) @ (env.B.T @ P @ env.A)  # Optimal policy from Part a\n",
        "  K_learned = np.linalg.inv(theta2.T @ theta2) @ theta3.T  # Learned policy from Q-learning\n",
        "\n",
        "  diff = np.linalg.norm(K_learned - K_optimal, 'fro')\n",
        "\n",
        "  print(\"Optimal feedback gain matrix K (from Part a):\")\n",
        "  print(K_optimal)\n",
        "\n",
        "  print(\"\\nLearned feedback gain matrix K (from Part d):\")\n",
        "  print(K_learned)\n",
        "\n",
        "  print(\"\\nDifference between K_optimal and K_learned (Frobenius norm):\", diff)\n",
        "\n",
        "  # Threshold to check if the policies are approximately the same\n",
        "  if diff < 1e-3:\n",
        "      print(\"\\nThe learned policy is approximately the same as the optimal policy from Part a.\")\n",
        "  else:\n",
        "      print(\"\\nThe learned policy is different from the optimal policy from Part a.\")"
      ],
      "metadata": {
        "id": "0okkkNdLjjV5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "avg_cost_optimal = optimal_cost()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vDy-Vt0U4yaU",
        "outputId": "ec7299e8-0d50-4581-d72c-cf86d12c0778"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 100/100 [00:14<00:00,  6.77it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print('avg_cost_optimal', avg_cost_optimal)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SMJkU8Gm7_3X",
        "outputId": "d2ffb1ca-fbc9-4d60-93c5-d8c0c22e5456"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "avg_cost_optimal 237.18853852775806\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "runtimes = {}"
      ],
      "metadata": {
        "id": "OiNwo1Gi6lrL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "epsilon = 0.3\n",
        "theta1, theta2, theta3, theta4, runtimes[epsilon] = runs(epsilon, theta_seed=1)\n",
        "avg_cost_learned = learned_cost(theta1, theta2, theta3, theta4)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VDDEQQMg6gwY",
        "outputId": "ccbdfae6-6590-4aba-abd6-1bff3c23deeb"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Updated Parameters:\n",
            "Theta1: [[ 1.51472038 -0.61475558]\n",
            " [-0.4552223  -1.0586355 ]]\n",
            "Theta2: [[ 0.9877118  -2.17478617]\n",
            " [ 1.57947942 -0.91222193]]\n",
            "Theta3: [[ 0.3134316  -0.2332111 ]\n",
            " [ 1.44995774 -2.04967514]]\n",
            "Theta4: [[-0.05280974]]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 100/100 [00:14<00:00,  6.90it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print_compare_cost(epsilon, avg_cost_optimal, avg_cost_learned)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6B3u1X0v6hpT",
        "outputId": "3f7eb045-efb6-4fed-d5b6-eed2f82feac0"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "\n",
            "Epsilon: 0.3\n",
            "Average Cost of Optimal Policy: 237.18853852775806\n",
            "Average Cost of Learned Policy: 190.43095456513208\n",
            "The learned policy is better than the optimal policy.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "V4cI40LbDNDe"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Do the parameters converge? Is the resulting policy the same as the policy obtained in part (a)?\n"
      ],
      "metadata": {
        "id": "IL8opJc258Q_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print_compare_policy(theta2, theta3)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "D7tmY2_87OZi",
        "outputId": "6b75f0c4-b8c7-4886-c389-435602160807"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Optimal feedback gain matrix K (from Part a):\n",
            "[[ 0.03457423  0.07474569]\n",
            " [-0.04677614 -0.09387158]]\n",
            "\n",
            "Learned feedback gain matrix K (from Part d):\n",
            "[[ 0.14113878  0.11031566]\n",
            " [ 0.04914226 -0.29734088]]\n",
            "\n",
            "Difference between K_optimal and K_learned (Frobenius norm): 0.25143851178909465\n",
            "\n",
            "The learned policy is different from the optimal policy from Part a.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "*type your answer here*\n",
        "It doesn't look like the policy matrix is exactly the same, but the cost is pretty close to the optimal policy cost.\n",
        "So YES, I would say it has converged."
      ],
      "metadata": {
        "id": "L7wufXko58R3"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Evaluate the policy performance over 100 trajectories each with length 1000 and compare its average cost with that of the optimal policy.\n"
      ],
      "metadata": {
        "id": "_IHOfPIP29R8"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### (e)  Try different values of epsilons $\\in \\{0.2,0.4,0.6, 0.8\\}$ (using the same initialization of $\\Theta$). With which value of epsilon does the algorithm converges faster?\n",
        "\n"
      ],
      "metadata": {
        "id": "EsCPN4i8-D1o"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "epsilon = 0.2\n",
        "theta1, theta2, theta3, theta4, runtimes[epsilon] = runs(epsilon, theta_seed=1)\n",
        "avg_cost_learned = learned_cost(theta1, theta2, theta3, theta4)\n",
        "print_compare_cost(epsilon, avg_cost_optimal, avg_cost_learned)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "etv1_RVewonu",
        "outputId": "e730ae13-63fd-405d-fb3d-5b5f719feed7"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Updated Parameters:\n",
            "Theta1: [[ 1.4499848  -0.67285978]\n",
            " [-0.34923284 -1.10862877]]\n",
            "Theta2: [[ 1.0294868  -2.14965543]\n",
            " [ 1.53084867 -0.95898962]]\n",
            "Theta3: [[ 0.30610166 -0.20934839]\n",
            " [ 1.43967336 -2.05060619]]\n",
            "Theta4: [[0.00176851]]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 100/100 [00:14<00:00,  6.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "\n",
            "Epsilon: 0.2\n",
            "Average Cost of Optimal Policy: 237.18853852775806\n",
            "Average Cost of Learned Policy: 218.5230098221142\n",
            "The learned policy is better than the optimal policy.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "epsilon = 0.4\n",
        "theta1, theta2, theta3, theta4, runtimes[epsilon] = runs(epsilon, theta_seed=1)\n",
        "print_compare_cost(epsilon, avg_cost_optimal, avg_cost_learned)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LUfSlpzdvw0N",
        "outputId": "e7efff82-2027-4527-951b-6893155b530b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Updated Parameters:\n",
            "Theta1: [[ 1.36153405 -0.62589886]\n",
            " [-0.33567115 -1.03539694]]\n",
            "Theta2: [[ 1.17960776 -2.09961113]\n",
            " [ 1.4418526  -1.14886842]]\n",
            "Theta3: [[ 0.31109595 -0.20980312]\n",
            " [ 1.43470282 -2.03158151]]\n",
            "Theta4: [[-0.08288427]]\n",
            "\n",
            "\n",
            "Epsilon: 0.4\n",
            "Average Cost of Optimal Policy: 237.18853852775806\n",
            "Average Cost of Learned Policy: 218.5230098221142\n",
            "The learned policy is better than the optimal policy.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "epsilon = 0.6\n",
        "theta1, theta2, theta3, theta4, runtimes[epsilon] = runs(epsilon, theta_seed=1)\n",
        "print_compare_cost(epsilon, avg_cost_optimal, avg_cost_learned)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bwpfiyNDvyux",
        "outputId": "84a143c1-f7bb-474a-d51d-2cbf586c140b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Updated Parameters:\n",
            "Theta1: [[ 1.43079401 -0.61688035]\n",
            " [-0.37646226 -0.99896519]]\n",
            "Theta2: [[ 1.35437888 -2.13973392]\n",
            " [ 1.45226702 -1.36868318]]\n",
            "Theta3: [[ 0.29516281 -0.22886941]\n",
            " [ 1.44078412 -2.00913236]]\n",
            "Theta4: [[-0.17331113]]\n",
            "\n",
            "\n",
            "Epsilon: 0.6\n",
            "Average Cost of Optimal Policy: 237.18853852775806\n",
            "Average Cost of Learned Policy: 218.5230098221142\n",
            "The learned policy is better than the optimal policy.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "epsilon = 0.8\n",
        "theta1, theta2, theta3, theta4, runtimes[epsilon] = runs(epsilon, theta_seed=1)\n",
        "print_compare_cost(epsilon, avg_cost_optimal, avg_cost_learned)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VwGHOfHLvz4z",
        "outputId": "b8837ca1-984f-4538-eec5-e7849a3ad457"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Updated Parameters:\n",
            "Theta1: [[ 1.32827875 -0.50416198]\n",
            " [-0.4277696  -0.88029103]]\n",
            "Theta2: [[ 1.47648568 -2.17259781]\n",
            " [ 1.43410515 -1.5007993 ]]\n",
            "Theta3: [[ 0.28610879 -0.2226254 ]\n",
            " [ 1.43863098 -1.99766938]]\n",
            "Theta4: [[-0.27432793]]\n",
            "\n",
            "\n",
            "Epsilon: 0.8\n",
            "Average Cost of Optimal Policy: 237.18853852775806\n",
            "Average Cost of Learned Policy: 218.5230098221142\n",
            "The learned policy is better than the optimal policy.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "DlCoXf4hxK48"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(runtimes)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Oy3JTCitxTEn",
        "outputId": "989ed8af-54a3-4234-de91-35abd45df846"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{0.2: 59.69315814971924, 0.4: 65.580899477005, 0.6: 75.0695641040802, 0.8: 76.54838824272156, 0.3: 62.00534701347351}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "Answer = '0.2 epsilon has the fastest runtimes'\n",
        "print(Answer)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EqIGuGgc4cGh",
        "outputId": "abfc6a08-b4a8-486a-b41e-ab6b2c13c696"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.2 epsilon has the fastest runtimes\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "###  (f) Is it fair to generalize the conclusions we have found about the optimal way to select epsilon and the trajectory length to other problems/environment? For example, if you have found that a single infinite trajectory does not work well, and larger values of epsilon works better, can you generalize this to other environments/problems like, for example, chess? why/why not?"
      ],
      "metadata": {
        "id": "s9kOnCafCO1W"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "*type your answer here*: The hyperparameters might not generalized to other environments/problems.\n",
        "A more complex environment, like large state space and non-convex, high epsilon results may result in too much exploration that lead to poor learning since random action are unlikely to provide valuable feedbacks.\n",
        "Smaller state space, and smoother and continuous dynamics (like LQ), a larger epsilon can help agent explore more and escape local optima.\n",
        "\n",
        "A short trajectories may limit the agent's ability to learn longer-term dependencies, espeically when the rewards/cost are sparse and near the end of episode (like winning a game).\n",
        "Longer trajectories on the other hands provides more data but is more computationally expensive."
      ],
      "metadata": {
        "id": "LFKqS71bXj28"
      }
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "RTrqhLELDtg4"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}