{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "V-tgUZphruDp"
      },
      "outputs": [],
      "source": [
        "# Import Required Libraries\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "\n",
        "# Create the Policy Network\n",
        "class PolicyNetwork(nn.Module):\n",
        "\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "\n",
        "        self.fc1 = nn.Linear(4, 16)\n",
        "        self.relu = nn.ReLU()\n",
        "        self.fc2 = nn.Linear(16, 2)\n",
        "        self.softmax = nn.Softmax(dim=-1)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.relu(self.fc1(x))\n",
        "        x = self.softmax(self.fc2(x))\n",
        "        return x"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Initialize the Model and Optimizer\n",
        "policy = PolicyNetwork()\n",
        "\n",
        "optimizer = optim.Adam(policy.parameters(), lr=0.01)\n",
        "\n",
        "# Define the Current State\n",
        "state = torch.tensor([0.5, 0.2, 0.8, 0.1], dtype=torch.float32)\n",
        "\n",
        "# Predict Action Probabilities\n",
        "action_probs = policy(state)\n",
        "\n",
        "print(action_probs)"
      ],
      "metadata": {
        "id": "4mSa3kO8r4WK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Sample an Action\n",
        "distribution = torch.distributions.Categorical(action_probs)\n",
        "\n",
        "action = distribution.sample()\n",
        "\n",
        "print(\"Selected Action:\", action.item())\n",
        "\n",
        "# Compute the Policy Loss\n",
        "reward = torch.tensor(2.0)\n",
        "\n",
        "loss = -distribution.log_prob(action) * reward\n",
        "\n",
        "print(loss)\n",
        "\n",
        "# Update the Policy Parameters\n",
        "optimizer.zero_grad()\n",
        "\n",
        "loss.backward()\n",
        "\n",
        "optimizer.step()"
      ],
      "metadata": {
        "id": "8BEMh2PisCyX"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}