main.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "AbYsqcZiCps9",
    "outputId": "39f4e4a4-ffa1-4210-a8d1-4f7353bdd185"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: swig in /usr/local/lib/python3.10/dist-packages (4.2.1)\n",
      "Requirement already satisfied: gym[box2d] in /usr/local/lib/python3.10/dist-packages (0.26.2)\n",
      "Requirement already satisfied: numpy>=1.18.0 in /usr/local/lib/python3.10/dist-packages (from gym[box2d]) (1.25.2)\n",
      "Requirement already satisfied: cloudpickle>=1.2.0 in /usr/local/lib/python3.10/dist-packages (from gym[box2d]) (2.2.1)\n",
      "Requirement already satisfied: gym-notices>=0.0.4 in /usr/local/lib/python3.10/dist-packages (from gym[box2d]) (0.0.8)\n",
      "Requirement already satisfied: box2d-py==2.3.5 in /usr/local/lib/python3.10/dist-packages (from gym[box2d]) (2.3.5)\n",
      "Requirement already satisfied: pygame==2.1.0 in /usr/local/lib/python3.10/dist-packages (from gym[box2d]) (2.1.0)\n",
      "Requirement already satisfied: swig==4.* in /usr/local/lib/python3.10/dist-packages (from gym[box2d]) (4.2.1)\n"
     ]
    }
   ],
   "source": [
    "#Update Gym\n",
    "!pip install swig\n",
    "!pip install gym[box2d] --upgrade"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "2sffYfT8yO1s",
    "outputId": "248afbf0-d10a-4b91-bb21-08f1e485f472"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.26.2\n"
     ]
    }
   ],
   "source": [
    "import gym\n",
    "print(gym.__version__)\n",
    "# NOTE: Version should be 0.26.02"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "j3-P-T38BFfm"
   },
   "source": [
    "# Define PPO Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "bgaZBJlxBDrW"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch.distributions import Categorical\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "\n",
    "class PolicyNetwork(torch.nn.Module):\n",
    "    def __init__(self, n=4, in_dim=128):\n",
    "        super(PolicyNetwork, self).__init__()\n",
    "\n",
    "        self.fc1 = torch.nn.Linear(in_dim, 128)\n",
    "        self.fc2 = torch.nn.Linear(128, 128)\n",
    "        self.fc3 = torch.nn.Linear(128, 128)\n",
    "        self.fc4 = torch.nn.Linear(128, n)\n",
    "        self.l_relu = torch.nn.LeakyReLU(0.1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.l_relu(self.fc1(x))\n",
    "        x = self.l_relu(self.fc2(x))\n",
    "        x = self.l_relu(self.fc3(x))\n",
    "        y = self.fc4(x)\n",
    "        y = F.softmax(y, dim=-1)\n",
    "        return y\n",
    "\n",
    "    def sample_action(self, state):\n",
    "\n",
    "        if not state is torch.Tensor:\n",
    "            state = torch.from_numpy(state).float().to(device)\n",
    "\n",
    "        if len(state.size()) == 1:\n",
    "            state = state.unsqueeze(0)\n",
    "\n",
    "        y = self(state)\n",
    "        dist = Categorical(y)\n",
    "        action = dist.sample()\n",
    "        log_probability = dist.log_prob(action)\n",
    "\n",
    "        return action.item(), log_probability.item()\n",
    "\n",
    "    def best_action(self, state):\n",
    "\n",
    "        if not state is torch.Tensor:\n",
    "            state = torch.from_numpy(state).float().to(device)\n",
    "\n",
    "        if len(state.size()) == 1:\n",
    "            state = state.unsqueeze(0)\n",
    "\n",
    "        y = self(state).squeeze()\n",
    "        action = torch.argmax(y)\n",
    "\n",
    "        return action.item()\n",
    "\n",
    "    def evaluate_actions(self, states, actions):\n",
    "        y = self(states)\n",
    "        dist = Categorical(y)\n",
    "        entropy = dist.entropy()\n",
    "        log_probabilities = dist.log_prob(actions)\n",
    "\n",
    "        return log_probabilities, entropy\n",
    "\n",
    "\n",
    "class ValueNetwork(torch.nn.Module):\n",
    "    def __init__(self, in_dim=128):\n",
    "        super(ValueNetwork, self).__init__()\n",
    "\n",
    "        self.fc1 = torch.nn.Linear(in_dim, 128)\n",
    "        self.fc2 = torch.nn.Linear(128, 128)\n",
    "        self.fc3 = torch.nn.Linear(128, 128)\n",
    "        self.fc4 = torch.nn.Linear(128, 1)\n",
    "        self.l_relu = torch.nn.LeakyReLU(0.1)\n",
    "\n",
    "    def forward(self, x):\n",
    "\n",
    "        x = self.l_relu(self.fc1(x))\n",
    "        x = self.l_relu(self.fc2(x))\n",
    "        x = self.l_relu(self.fc3(x))\n",
    "        y = self.fc4(x)\n",
    "\n",
    "        return y.squeeze(1)\n",
    "\n",
    "    def state_value(self, state):\n",
    "\n",
    "        if not state is torch.Tensor:\n",
    "            state = torch.from_numpy(state).float().to(device)\n",
    "\n",
    "        if len(state.size()) == 1:\n",
    "            state = state.unsqueeze(0)\n",
    "\n",
    "        y = self(state)\n",
    "\n",
    "        return y.item()\n",
    "\n",
    "\n",
    "def ac_loss_clipped(new_log_probabilities, old_log_probabilities, advantages, epsilon_clip=0.2):\n",
    "    probability_ratios = torch.exp(new_log_probabilities - old_log_probabilities)\n",
    "    clipped_probabiliy_ratios = torch.clamp(\n",
    "        probability_ratios, 1 - epsilon_clip, 1 + epsilon_clip\n",
    "    )\n",
    "\n",
    "    surrogate_1 = probability_ratios * advantages\n",
    "    surrogate_2 = clipped_probabiliy_ratios * advantages\n",
    "\n",
    "    return -torch.min(surrogate_1, surrogate_2)\n",
    "\n",
    "def train_combined_networks(policy_model, value_model, combined_optimizer, data_loader, epochs=40, clip=0.2):\n",
    "    c1 = 0.01  # Coefficient for entropy regularization\n",
    "    c2 = 0.5   # Coefficient for value loss weight\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        policy_losses = []\n",
    "        value_losses = []\n",
    "\n",
    "        for observations, actions, advantages, log_probabilities, rewards_to_go in data_loader:\n",
    "            observations = observations.float().to(device)\n",
    "            actions = actions.long().to(device)\n",
    "            advantages = advantages.float().to(device)\n",
    "            old_log_probabilities = log_probabilities.float().to(device)\n",
    "            rewards_to_go = rewards_to_go.float().to(device)\n",
    "\n",
    "            combined_optimizer.zero_grad()\n",
    "\n",
    "            new_log_probabilities, entropy = policy_model.evaluate_actions(observations, actions)\n",
    "            policy_loss = (\n",
    "                ac_loss_clipped(\n",
    "                    new_log_probabilities,\n",
    "                    old_log_probabilities,\n",
    "                    advantages,\n",
    "                    epsilon_clip=clip,\n",
    "                ).mean()\n",
    "                - c1 * entropy.mean()\n",
    "            )\n",
    "            policy_losses.append(policy_loss.item())\n",
    "\n",
    "            values = value_model(observations)\n",
    "            value_loss = c2 * F.mse_loss(values, rewards_to_go)\n",
    "            value_losses.append(value_loss.item())\n",
    "\n",
    "            total_loss = policy_loss + value_loss\n",
    "\n",
    "            total_loss.backward()\n",
    "            combined_optimizer.step()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "pgY5fBlQBL3X"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "\n",
    "def cumulative_sum(array, gamma=1.0):\n",
    "    curr = 0\n",
    "    cumulative_array = []\n",
    "\n",
    "    for a in array[::-1]:\n",
    "        curr = a + gamma * curr\n",
    "        cumulative_array.append(curr)\n",
    "\n",
    "    return cumulative_array[::-1]\n",
    "\n",
    "\n",
    "class Episode:\n",
    "    def __init__(self, gamma=0.99, lambd=0.95):\n",
    "        self.observations = []\n",
    "        self.actions = []\n",
    "        self.advantages = []\n",
    "        self.rewards = []\n",
    "        self.rewards_to_go = []\n",
    "        self.values = []\n",
    "        self.log_probabilities = []\n",
    "        self.gamma = gamma\n",
    "        self.lambd = lambd\n",
    "\n",
    "    def append(\n",
    "        self, observation, action, reward, value, log_probability, reward_scale=20\n",
    "    ):\n",
    "        self.observations.append(observation)\n",
    "        self.actions.append(action)\n",
    "        self.rewards.append(reward / reward_scale)\n",
    "        self.values.append(value)\n",
    "        self.log_probabilities.append(log_probability)\n",
    "\n",
    "    def end_episode(self, last_value):\n",
    "        rewards = np.array(self.rewards + [last_value])\n",
    "        values = np.array(self.values + [last_value])\n",
    "        deltas = rewards[:-1] + self.gamma * values[1:] - values[:-1]\n",
    "        self.advantages = cumulative_sum(deltas.tolist(), gamma=self.gamma * self.lambd)\n",
    "        self.rewards_to_go = cumulative_sum(rewards.tolist(), gamma=self.gamma)[:-1]\n",
    "\n",
    "\n",
    "def normalize_list(array):\n",
    "    array = np.array(array)\n",
    "    array = (array - np.mean(array)) / (np.std(array) + 1e-5)\n",
    "    return array.tolist()\n",
    "\n",
    "\n",
    "class History(Dataset):\n",
    "    def __init__(self):\n",
    "        self.episodes = []\n",
    "        self.observations = []\n",
    "        self.actions = []\n",
    "        self.advantages = []\n",
    "        self.rewards = []\n",
    "        self.rewards_to_go = []\n",
    "        self.log_probabilities = []\n",
    "\n",
    "    def free_memory(self):\n",
    "        del self.episodes[:]\n",
    "        del self.observations[:]\n",
    "        del self.actions[:]\n",
    "        del self.advantages[:]\n",
    "        del self.rewards[:]\n",
    "        del self.rewards_to_go[:]\n",
    "        del self.log_probabilities[:]\n",
    "\n",
    "    def add_episode(self, episode):\n",
    "        self.episodes.append(episode)\n",
    "\n",
    "    def build_dataset(self):\n",
    "        for episode in self.episodes:\n",
    "            self.observations += episode.observations\n",
    "            self.actions += episode.actions\n",
    "            self.advantages += episode.advantages\n",
    "            self.rewards += episode.rewards\n",
    "            self.rewards_to_go += episode.rewards_to_go\n",
    "            self.log_probabilities += episode.log_probabilities\n",
    "\n",
    "        assert (\n",
    "            len(\n",
    "                {\n",
    "                    len(self.observations),\n",
    "                    len(self.actions),\n",
    "                    len(self.advantages),\n",
    "                    len(self.rewards),\n",
    "                    len(self.rewards_to_go),\n",
    "                    len(self.log_probabilities),\n",
    "                }\n",
    "            )\n",
    "            == 1\n",
    "        )\n",
    "\n",
    "        self.advantages = normalize_list(self.advantages)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.observations)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return (\n",
    "            self.observations[idx],\n",
    "            self.actions[idx],\n",
    "            self.advantages[idx],\n",
    "            self.log_probabilities[idx],\n",
    "            self.rewards_to_go[idx],\n",
    "        )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aSRo2IPxBQQy"
   },
   "source": [
    "# TRAC Optimizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kyERnJ0idJI9"
   },
   "source": [
    "## Define Erfi function in Pytorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "1xN5ZOV0BSEo"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "# We depend on Erfi function, but python.special currently has no implementation.\n",
    "# We instead modify and rely on https://github.com/redsnic/torch_erf\n",
    "\n",
    "def polyval(x,coeffs):\n",
    "    \"\"\"Implementation of the Horner scheme to evaluate a polynomial\n",
    "\n",
    "    taken from https://discuss.pytorch.org/t/polynomial-evaluation-by-horner-rule/67124\n",
    "\n",
    "    Args:\n",
    "        x (torch.Tensor): variable\n",
    "        coeffs (torch.Tensor): coefficients of the polynomial\n",
    "    \"\"\"\n",
    "    curVal=0\n",
    "    for curValIndex in range(len(coeffs)-1):\n",
    "        curVal=(curVal+coeffs[curValIndex])*x[0]\n",
    "    return(curVal+coeffs[len(coeffs)-1])\n",
    "\n",
    "\n",
    "class ERF_1994(torch.nn.Module):\n",
    "    \"\"\"Class to compute the error function of a complex number (extends torch.special.erf behavior)\n",
    "\n",
    "    This class is based on the algorithm proposed in:\n",
    "    Weideman, J. Andre C. \"Computation of the complex error function.\" SIAM Journal on Numerical Analysis 31.5 (1994): 1497-1518\n",
    "    \"\"\"\n",
    "    def __init__(self, n_coefs):\n",
    "        \"\"\"Defaul constructor\n",
    "\n",
    "        Args:\n",
    "            n_coefs (integer): The number of polynomial coefficients to use in the approximation\n",
    "        \"\"\"\n",
    "        super(ERF_1994, self).__init__()\n",
    "        # compute polynomial coefficients and other constants\n",
    "        self.N = n_coefs\n",
    "        self.i = torch.complex(torch.tensor(0.),torch.tensor(1.))\n",
    "        self.M = 2*self.N\n",
    "        self.M2 = 2*self.M\n",
    "        self.k = torch.linspace(-self.M+1, self.M-1, self.M2-1)\n",
    "        self.L = torch.sqrt(self.N/torch.sqrt(torch.tensor(2.)))\n",
    "        self.theta = self.k*torch.pi/self.M\n",
    "        self.t = self.L*torch.tan(self.theta/2)\n",
    "        self.f = torch.exp(-self.t**2)*(self.L**2 + self.t**2)\n",
    "        self.a = torch.fft.fft(torch.fft.fftshift(self.f)).real/self.M2\n",
    "        self.a = torch.flipud(self.a[1:self.N+1])\n",
    "\n",
    "    def w_algorithm(self, z):\n",
    "        \"\"\"Compute the Faddeeva function of a complex number\n",
    "\n",
    "        The constant coefficients are computed in the constructor of the class.\n",
    "\n",
    "        Weideman, J. Andre C. \"Computation of the complex error function.\" SIAM Journal on Numerical Analysis 31.5 (1994): 1497-1518\n",
    "\n",
    "        Args:\n",
    "            z (torch.Tensor): A tensor of complex numbers (any shape is allowed)\n",
    "\n",
    "        Returns:\n",
    "            torch.Tensor: w(z) for each element of z\n",
    "        \"\"\"\n",
    "        Z = (self.L+self.i*z)/(self.L-self.i*z)\n",
    "        p = polyval(Z.unsqueeze(0), self.a)\n",
    "        w = 2*p/(self.L-self.i*z)**2+(1/torch.sqrt(torch.tensor(torch.pi)))/(self.L-self.i*z)\n",
    "        return w\n",
    "\n",
    "    def forward(self, z):\n",
    "        \"\"\"Compute the error function of a complex number\n",
    "\n",
    "        The result is computed by manipulating the Faddeeva function.\n",
    "\n",
    "        Args:\n",
    "            z (torch.Tensor): A tensor of complex numbers (any shape is allowed)\n",
    "\n",
    "        Returns:\n",
    "            torch.Tensor: erf(z) for each element of z\n",
    "        \"\"\"\n",
    "        # exploit the symmetry of the error function\n",
    "        # find the sign of the real part\n",
    "        sign_r = torch.sign(z.real)\n",
    "        sign_i = torch.sign(z.imag)\n",
    "        # flip sign of imaginary part if negative\n",
    "        z = torch.complex(torch.abs(z.real), torch.abs(z.imag))\n",
    "        out = -torch.exp(torch.log(self.w_algorithm(z*self.i)) - z**2) + 1\n",
    "        return torch.complex(out.real*sign_r, out.imag*sign_i)\n",
    "\n",
    "    def backward(self, z):\n",
    "        \"\"\"Compute the gradient of the error function of a complex number.\n",
    "\n",
    "        As we know the analytical derivative of the the error function, we can use it directly.\n",
    "\n",
    "        Args:\n",
    "            z (torch.Tensor): A tensor of complex numbers (any shape is allowed)\n",
    "        Returns:\n",
    "            torch.Tensor: grad(erf(z)) for each element of x\n",
    "        \"\"\"\n",
    "        return 2/torch.sqrt(torch.tensor(torch.pi))*torch.exp(-z**2)\n",
    "\n",
    "erf_torch = ERF_1994(128)\n",
    "\n",
    "def erfi(x):\n",
    "    if not torch.is_floating_point(x):\n",
    "        x = x.to(torch.float32)\n",
    "\n",
    "    # Convert x to a complex tensor where the real part is zero\n",
    "    ix = torch.complex(torch.zeros_like(x), x)\n",
    "\n",
    "    # Compute erf(ix) / i\n",
    "    erfi_x = erf_torch(ix).imag  # Extract the imaginary part of erf(ix)\n",
    "    return erfi_x\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "k3zE2Hd0dZJm"
   },
   "source": [
    "## TRAC Wrapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "2xccOJBKdTRU"
   },
   "outputs": [],
   "source": [
    "from typing import Tuple, Any, Callable, Dict\n",
    "import torch\n",
    "\n",
    "# We closely follow the meta-optimizer structure from the code in\n",
    "# Cutkosky et. al 2023\n",
    "def _init_state(\n",
    "        optimizer: torch.optim.Optimizer,\n",
    "        theta_ref: Dict[torch.Tensor, torch.Tensor],\n",
    "        betas: Tuple[float],\n",
    "        s_prev: float,\n",
    "        eps: float):\n",
    "    if '_trac' not in optimizer.state:\n",
    "        optimizer.state['_trac'] = {\n",
    "            'betas': torch.tensor(betas),\n",
    "            's_prev': torch.tensor(s_prev),\n",
    "            'eps': eps,\n",
    "            's': torch.zeros(len(betas)),\n",
    "            'theta_ref': {},\n",
    "            'variance': torch.zeros(len(betas)),\n",
    "            'sigma': torch.full((len(betas),), 1e-8),\n",
    "            'iter_count': 0,\n",
    "        }\n",
    "        _init_reference(optimizer, theta_ref)\n",
    "\n",
    "def _init_reference(\n",
    "        optimizer: torch.optim.Optimizer,\n",
    "        theta_ref: Dict[torch.Tensor, torch.Tensor],):\n",
    "    '''\n",
    "    Args:\n",
    "        optimizer: optimizer instance to store reference for.\n",
    "        theta_ref: mapping of parameters to their initial values at the start of optimization.\n",
    "    '''\n",
    "    for group in optimizer.param_groups:\n",
    "        for p in group['params']:\n",
    "            optimizer.state['_trac'][p] = {\n",
    "                'ref': theta_ref[p].clone(),\n",
    "            }\n",
    "            \n",
    "\n",
    "def _step(\n",
    "        optimizer: torch.optim.Optimizer,\n",
    "        base_step: Callable,\n",
    "        betas: Tuple[float],\n",
    "        s_prev: float,\n",
    "        eps: float,\n",
    "        ):\n",
    "    '''\n",
    "    Args:\n",
    "        optimizer: trac optimizer instance\n",
    "        base_step: The \"step\" function of the base optimizer\n",
    "        betas: list of beta values.\n",
    "        s_init: initial scale value.\n",
    "        eps: epsilon value.\n",
    "    '''\n",
    "\n",
    "    prev_grad = torch.is_grad_enabled()\n",
    "\n",
    "\n",
    "    torch.set_grad_enabled(False)\n",
    "    updates = {}\n",
    "    grads = {}\n",
    "    deltas = {}\n",
    "\n",
    "    for group in optimizer.param_groups:\n",
    "        for p in group['params']:\n",
    "\n",
    "            if p.grad is None:\n",
    "                grads[p] = None\n",
    "            else:\n",
    "                grads[p] = p.grad.clone()\n",
    "            updates[p] = p.data.clone()\n",
    "\n",
    "    torch.set_grad_enabled(prev_grad)\n",
    "    result = base_step(None)\n",
    "    torch.set_grad_enabled(False)\n",
    "    \n",
    "    _init_state(optimizer, updates, betas, s_prev, eps)\n",
    "    trac_state = optimizer.state['_trac']\n",
    "\n",
    "\n",
    "    for group in optimizer.param_groups:\n",
    "        for p in group['params']:\n",
    "            if grads[p] is None:\n",
    "                continue\n",
    "\n",
    "            theta_ref = trac_state[p]['ref']\n",
    "\n",
    "            deltas[p] = (updates[p] - theta_ref)/(torch.sum(trac_state['s']) + trac_state['eps'])\n",
    "\n",
    "            updates[p].copy_(p-updates[p])\n",
    "\n",
    "    h = 0.0\n",
    "    for group in optimizer.param_groups:\n",
    "        for p in group['params']:\n",
    "\n",
    "            if grads[p] is None:\n",
    "                continue\n",
    "\n",
    "            grad = grads[p]\n",
    "\n",
    "            delta = deltas[p]\n",
    "            product = torch.dot(delta.flatten(), grad.flatten())\n",
    "            if product.isnan():\n",
    "                raise ValueError(\"NaNs in product\")\n",
    "            h += product\n",
    "\n",
    "            delta.add_(updates[p])\n",
    "\n",
    "    device = h.device\n",
    "\n",
    "    for key in trac_state:\n",
    "        try:\n",
    "            if trac_state[key].device != device:\n",
    "                trac_state[key] = trac_state[key].to(device)\n",
    "        except:\n",
    "            pass\n",
    "\n",
    "    s = trac_state['s']\n",
    "    s_prev = trac_state['s_prev']\n",
    "    betas = trac_state['betas']\n",
    "    eps = trac_state['eps']\n",
    "    variance = trac_state['variance'] \n",
    "    sigma = trac_state['sigma']                                 \n",
    "    trac_state['iter_count'] += 1\n",
    "\n",
    "    variance.mul_(\n",
    "        betas**2).add_(torch.square(h))\n",
    "    sigma.mul_(betas).sub_(h)\n",
    "    f_term = s_prev / (erfi(torch.tensor(1.0) / torch.sqrt(torch.tensor(2.0))))\n",
    "    s_term = erfi(sigma / (torch.sqrt(torch.tensor(2.0)) * torch.sqrt(variance) + eps))\n",
    "    if (f_term * s_term).isnan().any():\n",
    "        raise ValueError(\"NaNs in s\")\n",
    "    s.copy_(f_term * s_term)\n",
    "\n",
    "    for group in optimizer.param_groups:\n",
    "        for p in group['params']:\n",
    "\n",
    "            if grads[p] is None:\n",
    "                continue\n",
    "\n",
    "            theta_ref = trac_state[p]['ref']\n",
    "            delta = deltas[p]\n",
    "            s_sum = torch.sum(s)\n",
    "\n",
    "            scale = max(s_sum, 0.0)\n",
    "            p.copy_(theta_ref + delta * scale)\n",
    "\n",
    "    log_data = {\n",
    "        'iter_count': trac_state['iter_count'],\n",
    "        's': torch.sum(s).item(),\n",
    "    }\n",
    "\n",
    "    torch.set_grad_enabled(prev_grad)\n",
    "    return result, log_data\n",
    "\n",
    "\n",
    "class trac:\n",
    "    pass\n",
    "\n",
    "def is_trac(opt):\n",
    "    return isinstance(opt, trac)\n",
    "\n",
    "def start_trac(\n",
    "        log_file,\n",
    "        Base: Any,\n",
    "        betas: Tuple[float] = (0.9, 0.99, 0.999, 0.9999,\n",
    "                               0.99999, 0.999999),\n",
    "        s_prev: float = 1e-8,\n",
    "        eps: float = 1e-8,\n",
    "        ):\n",
    "\n",
    "    class TRACOPT(Base, trac):\n",
    "        '''\n",
    "        Wraps the base opt with trac.\n",
    "        \n",
    "        '''\n",
    "\n",
    "        def step(self):\n",
    "            result, log_data = _step(self, super().step, betas, s_prev, eps)\n",
    "            with open (log_file, 'a') as f:\n",
    "                f.write(str(log_data) + '\\n')\n",
    "            return result\n",
    "\n",
    "    TRACOPT.__name__ += Base.__name__\n",
    "\n",
    "    return TRACOPT\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YZJjCxSEBqIn"
   },
   "source": [
    "# Lifelong Control Setting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "5qXh9tafBtrn"
   },
   "outputs": [],
   "source": [
    "import gym\n",
    "import numpy as np\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "from tqdm import tqdm\n",
    "import random\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "XcILqsPSDdQO"
   },
   "outputs": [],
   "source": [
    "# SET THE SEED\n",
    "seed = 2024\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "\n",
    "# num distribution shifts\n",
    "levels = 10\n",
    "\n",
    "# Hyperparameters\n",
    "lr = 0.01\n",
    "max_episodes = 2\n",
    "train_epochs = 5\n",
    "max_timesteps = 400\n",
    "state_scale = 1.0\n",
    "reward_scale = 20.0\n",
    "batch_size = 32\n",
    "\n",
    "# when to introduce distribution shift\n",
    "level_switch = 200\n",
    "max_iterations = levels * level_switch\n",
    "\n",
    "# peturbation range\n",
    "rp_range = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "ITlv3qL1B0Q7"
   },
   "outputs": [],
   "source": [
    "def get_peturbations(env_name, seed):\n",
    "  env = gym.make(env_name, render_mode=\"rgb_array\")\n",
    "  observation = env.reset()[0]\n",
    "  random_perturbations = [\n",
    "        np.random.normal(0, rp_range, observation.shape) for _ in range(levels)\n",
    "    ]\n",
    "  # make the first random perturbation zero\n",
    "  random_perturbations[0] = np.zeros(observation.shape)\n",
    "  return random_perturbations\n",
    "def train(env_name, opt_choice, random_perturbations):\n",
    "\n",
    "    # Create log txt files\n",
    "    trac_reward_log_file = f'logs/trac_reward_log_{env_name}_{seed}.txt'\n",
    "    base_reward_log_file = f'logs/base_reward_log_{env_name}_{seed}.txt'\n",
    "\n",
    "    # Setup env\n",
    "    env = gym.make(env_name, render_mode=\"rgb_array\")\n",
    "    observation = env.reset()[0]\n",
    "    n_actions = env.action_space.n\n",
    "    feature_dim = observation.size\n",
    "\n",
    "    tqdm_bar = tqdm(range(max_iterations), desc=\"Training\", unit=\"iteration\")\n",
    "\n",
    "    value_model = ValueNetwork(in_dim=feature_dim).to(device)\n",
    "    policy_model = PolicyNetwork(in_dim=feature_dim, n=n_actions).to(device)\n",
    "\n",
    "    trac_combined_optimizer = start_trac(log_file=f'logs/trac_{env_name}.text', Base=optim.Adam)(\n",
    "        [\n",
    "            {\"params\": policy_model.parameters(), \"lr\": lr},\n",
    "            {\"params\": value_model.parameters(), \"lr\": lr},\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    base_combined_optimizer = torch.optim.Adam(\n",
    "        [\n",
    "            {\"params\": policy_model.parameters(), \"lr\": lr},\n",
    "            {\"params\": value_model.parameters(), \"lr\": lr},\n",
    "        ]\n",
    "    )\n",
    "    if opt_choice == \"TRAC\":\n",
    "        combined_optimizer = trac_combined_optimizer\n",
    "        reward_log_file = trac_reward_log_file\n",
    "        print(\"USING TRAC.\")\n",
    "    if opt_choice == \"base\":\n",
    "        combined_optimizer = base_combined_optimizer\n",
    "        reward_log_file = base_reward_log_file\n",
    "    history = History()\n",
    "    level = 0\n",
    "    for ite in tqdm_bar:\n",
    "        # Switch perturbation level\n",
    "        if ite % level_switch == 0:\n",
    "            random_perturbation = random_perturbations[level]\n",
    "            level += 1\n",
    "\n",
    "        episodes_reward = []\n",
    "\n",
    "        for _ in range(max_episodes):\n",
    "            observation = env.reset()[0]\n",
    "            observation += random_perturbation\n",
    "            episode = Episode()\n",
    "\n",
    "            for timestep in range(max_timesteps):\n",
    "                action, log_probability = policy_model.sample_action(observation / state_scale)\n",
    "                value = value_model.state_value(observation / state_scale)\n",
    "\n",
    "                new_observation, reward, done, _, _ = env.step(action)\n",
    "                new_observation += random_perturbation\n",
    "\n",
    "                episode.append(\n",
    "                    observation=observation / state_scale,\n",
    "                    action=action,\n",
    "                    reward=reward,\n",
    "                    value=value,\n",
    "                    log_probability=log_probability,\n",
    "                    reward_scale=reward_scale,\n",
    "                )\n",
    "\n",
    "                observation = new_observation\n",
    "\n",
    "                if done:\n",
    "                    episode.end_episode(last_value=0)\n",
    "                    break\n",
    "\n",
    "                if timestep == max_timesteps - 1:\n",
    "                    value = value_model.state_value(observation / state_scale)\n",
    "                    episode.end_episode(last_value=value)\n",
    "\n",
    "            episodes_reward.append(reward_scale * np.sum(episode.rewards))\n",
    "            history.add_episode(episode)\n",
    "\n",
    "        mean_rewards = np.mean(episodes_reward)\n",
    "        tqdm_bar.set_postfix(mean_rewards=mean_rewards)\n",
    "\n",
    "        with open(reward_log_file, 'a') as f:\n",
    "            f.write(str(mean_rewards) + '\\n')\n",
    "        history.build_dataset()\n",
    "        data_loader = DataLoader(history, batch_size=batch_size, shuffle=True)\n",
    "        train_combined_networks(policy_model, value_model, combined_optimizer, data_loader, train_epochs)\n",
    "        history.free_memory()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4iBE7Qhs-Fuw"
   },
   "source": [
    "### Plotting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "c4H3Ws0v-Da8"
   },
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "def read_data(file_path):\n",
    "    with open(file_path, 'r') as file:\n",
    "        data = file.readlines()\n",
    "    return data\n",
    "\n",
    "def plot(env_name, seed):\n",
    "    # Read data from files\n",
    "    trac_data = read_data(f'logs/trac_reward_log_{env_name}_{seed}.txt')\n",
    "    base_data = read_data(f'logs/base_reward_log_{env_name}_{seed}.txt')\n",
    "\n",
    "    # Convert data to float\n",
    "    trac_data = [float(i) for i in trac_data]\n",
    "    base_data = [float(i) for i in base_data]\n",
    "\n",
    "\n",
    "    # Smooth trac and base data\n",
    "    window = 5\n",
    "    trac_data = np.convolve(trac_data, np.ones(window) / window, mode='valid')\n",
    "    base_data = np.convolve(base_data, np.ones(window) / window, mode='valid')\n",
    "\n",
    "    # Create a plot with seaborn\n",
    "    sns.set(style=\"darkgrid\")\n",
    "    plt.figure(figsize=(10, 6))\n",
    "\n",
    "    plt.plot(base_data, label='Adam PPO', color='#4a69bd')\n",
    "    plt.plot(trac_data, label='TRAC PPO', color='#b71540')\n",
    "\n",
    "    plt.xlabel('Timesteps')\n",
    "    plt.ylabel('Mean Episode Reward')\n",
    "    plt.title(f'{env_name}', fontsize=24)\n",
    "    plt.legend()\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pF8BK3-B6-_2"
   },
   "source": [
    "# Compare one seed results for Acrobot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5kjeR_wo7DxW"
   },
   "source": [
    "## Train with Adam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "YjGAMqwk7F6R",
    "outputId": "566868e2-157d-4a13-a3a4-8fc7e5c7c5d0"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Online Peturbations are\n",
      "[array([0., 0., 0., 0., 0., 0.]), array([-5.239924  , -2.65058914,  0.91997724,  0.20410329,  2.10710556,\n",
      "        3.24808522]), array([-3.00127004, -0.55566338,  2.38799004,  1.72363065, -0.83409209,\n",
      "       -0.49907283]), array([ 1.88735471, -1.53262128,  0.41645746,  2.81744585, -2.97820802,\n",
      "       -2.95161707]), array([ 1.98169263, -1.76646085, -0.73236775, -3.06941007, -0.70315102,\n",
      "        1.27983623]), array([ 1.37847833,  1.51450473, -2.86109954, -0.8577586 , -1.37002372,\n",
      "       -0.25128173]), array([ 2.28917448,  0.65441957, -0.27345488,  0.35838783,  1.93795414,\n",
      "        0.01174018]), array([ 1.18101088, -0.78495275,  0.07182293, -0.66150991,  1.61755215,\n",
      "        0.10662189]), array([-2.63780402, -2.15956139, -0.75193655,  0.29745353,  3.62983768,\n",
      "        1.10499787]), array([-1.07844989,  1.85410224,  2.39777336,  0.03593557,  2.03176482,\n",
      "       -2.36861302])]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training:   0%|          | 0/2000 [00:00<?, ?iteration/s]/n/home04/amuppidi/.conda/envs/torch/lib/python3.10/site-packages/gym/utils/passive_env_checker.py:233: DeprecationWarning: `np.bool8` is a deprecated alias for `np.bool_`.  (Deprecated NumPy 1.24)\n",
      "  if not isinstance(terminated, (bool, np.bool8)):\n",
      "Training: 100%|██████████| 2000/2000 [37:42<00:00,  1.13s/iteration, mean_rewards=-400]\n"
     ]
    }
   ],
   "source": [
    "peturbations = get_peturbations(\"Acrobot-v1\", seed)\n",
    "print(\"Online Peturbations are\")\n",
    "print(peturbations)\n",
    "env = \"Acrobot-v1\"\n",
    "opt = \"base\"\n",
    "train(env, opt, peturbations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VyCzCOmh7L6Z"
   },
   "source": [
    "## Train with TRAC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "2JieNwXI7NWy",
    "outputId": "e8c51e39-a390-4484-b04c-1f65f3993499"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Online Peturbations are\n",
      "[array([0., 0., 0., 0., 0., 0.]), array([-5.239924  , -2.65058914,  0.91997724,  0.20410329,  2.10710556,\n",
      "        3.24808522]), array([-3.00127004, -0.55566338,  2.38799004,  1.72363065, -0.83409209,\n",
      "       -0.49907283]), array([ 1.88735471, -1.53262128,  0.41645746,  2.81744585, -2.97820802,\n",
      "       -2.95161707]), array([ 1.98169263, -1.76646085, -0.73236775, -3.06941007, -0.70315102,\n",
      "        1.27983623]), array([ 1.37847833,  1.51450473, -2.86109954, -0.8577586 , -1.37002372,\n",
      "       -0.25128173]), array([ 2.28917448,  0.65441957, -0.27345488,  0.35838783,  1.93795414,\n",
      "        0.01174018]), array([ 1.18101088, -0.78495275,  0.07182293, -0.66150991,  1.61755215,\n",
      "        0.10662189]), array([-2.63780402, -2.15956139, -0.75193655,  0.29745353,  3.62983768,\n",
      "        1.10499787]), array([-1.07844989,  1.85410224,  2.39777336,  0.03593557,  2.03176482,\n",
      "       -2.36861302])]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training:   0%|          | 0/2000 [00:00<?, ?iteration/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "USING TRAC.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training:  89%|████████▉ | 1786/2000 [30:15<09:02,  2.54s/iteration, mean_rewards=-400] "
     ]
    }
   ],
   "source": [
    "print(\"Online Peturbations are\")\n",
    "print(peturbations)\n",
    "opt = \"TRAC\"\n",
    "train(env, opt, peturbations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WAkxrtxr7Pad"
   },
   "source": [
    "## Plot Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 584
    },
    "id": "xZGOsa1H7Q3Y",
    "outputId": "0007d167-694f-4524-dc51-a7da12b13d3e"
   },
   "outputs": [],
   "source": [
    "plot(env,seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ptPG30C4-R2J"
   },
   "source": [
    "# Compare one seed results for CartPole"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Cf69AqiBCUaN"
   },
   "source": [
    "## Train with Adam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "cJPMtCpNCT4K",
    "outputId": "2ad36023-85e4-4995-9e3c-c94ceca9ec2c"
   },
   "outputs": [],
   "source": [
    "peturbations = get_peturbations(\"CartPole-v1\", seed)\n",
    "print(\"Online Peturbations are\")\n",
    "print(peturbations)\n",
    "env = \"CartPole-v1\"\n",
    "opt = \"base\"\n",
    "train(env, opt, peturbations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "g_ZKyg2MDQ5j"
   },
   "source": [
    "## Train with TRAC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "rsc96nqxCZ9_",
    "outputId": "1cd1bfcb-c2ae-44fd-e3ad-8b4eb015ec7b"
   },
   "outputs": [],
   "source": [
    "print(\"Online Peturbations are\")\n",
    "print(peturbations)\n",
    "opt = \"TRAC\"\n",
    "train(env, opt, peturbations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PGj9Xlva-gZ_"
   },
   "source": [
    "## Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 584
    },
    "id": "BF41I9hz-hxl",
    "outputId": "2489d1f3-20d5-47ba-dd05-82ab01197469"
   },
   "outputs": [],
   "source": [
    "plot(env,seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DvrCDPnp8vhW"
   },
   "source": [
    "# Compare one seed results for LunarLander"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zl_ZQ9yy9JFc"
   },
   "source": [
    "## Train with Adam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true,
     "base_uri": "https://localhost:8080/"
    },
    "id": "FP3apiGk8ym1",
    "outputId": "56181e99-3b0c-4763-f0cd-fb677c13fda0"
   },
   "outputs": [],
   "source": [
    "env = \"LunarLander-v2\"\n",
    "peturbations = get_peturbations(env, seed)\n",
    "print(\"Online Peturbations are\")\n",
    "print(peturbations)\n",
    "opt = \"base\"\n",
    "train(env, opt, peturbations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mphqlL159LQ3"
   },
   "source": [
    "## Train with TRAC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true,
     "base_uri": "https://localhost:8080/"
    },
    "id": "y5nkoWbY83Sx",
    "outputId": "b2c32609-668c-4e3a-b9ed-37ba9d2746d5"
   },
   "outputs": [],
   "source": [
    "opt = \"TRAC\"\n",
    "train(env, opt, peturbations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cns3O9h59MxI"
   },
   "source": [
    "## Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "M8HLRgLk9OCT",
    "outputId": "02ff9ffb-db4e-4004-e8c2-7ff41331bf5b"
   },
   "outputs": [],
   "source": [
    "plot(env,seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Pytorch",
   "language": "python",
   "name": "torch"
  },
  "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}