Agent_pure.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "jp-MarkdownHeadingCollapsed": true,
        "tags": [],
        "id": "5IWe2Bd6vdQI"
      },
      "source": [
        "### Dataclasses"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CBWQeeGxvdQJ"
      },
      "outputs": [],
      "source": [
        "from dataclasses import dataclass, asdict\n",
        "from typing import Callable\n",
        "import json\n",
        "import yaml\n",
        "from yaml import Loader, Dumper\n",
        "\n",
        "@dataclass\n",
        "class HP_alg:\n",
        "  max_nactions: int = -1\n",
        "  critic_play: bool = False\n",
        "  actor_lr: float = 1e-4 # lr at time 0\n",
        "  critic_lr: float = 5e-5\n",
        "  train_condition_str: str = \"(lambda tr: tr[0] % 5 != 2)\"\n",
        "  td_gamma: float = 0.998\n",
        "  batch_size: int = 512\n",
        "  max_attention_length: int = 64\n",
        "  use_fork_discount: bool = False # True not supported by jar-file currently\n",
        "  batch_accumulation_steps: int = 1\n",
        "  epochs: int = 100\n",
        "  use_GAE: bool = False\n",
        "  use_double_atten: bool = False\n",
        "  use_FFM: bool = True\n",
        "  path_: str = '../Config/HP_alg.yaml'\n",
        "\n",
        "hp_alg = HP_alg(**yaml.load(open('../Config/HP_alg.yml').read(), Loader=Loader))\n",
        "hp_alg.max_nactions = int(1e9) if hp_alg.critic_play else hp_alg.max_attention_length\n",
        "\n",
        "\n",
        "@dataclass\n",
        "class HP_front:\n",
        "  critic_play: bool = False\n",
        "  state_scheme: tuple = tuple(json.load(open('../Config/schemes.json'))['stateScheme'])\n",
        "  feature_names_scheme: tuple = tuple(state_scheme[0])\n",
        "  trajectory_scheme: tuple = ('hash', 'trajectory', 'name', 'statementsCount', 'probabilities')\n",
        "  between_logs: int = 64\n",
        "  rnn_features_count: int = 33\n",
        "  gnn_features_count: int = 8\n",
        "  features_dim: int = len(feature_names_scheme)\n",
        "  use_cuda_if_available: bool = True\n",
        "  tr_hash_idx: int = trajectory_scheme.index('hash')\n",
        "  tr_path_idx: int = trajectory_scheme.index('trajectory')\n",
        "  tr_name_idx: int = trajectory_scheme.index('name')\n",
        "  tr_code_len_idx: int = trajectory_scheme.index('statementsCount')\n",
        "  tr_prev_probs_idx = trajectory_scheme.index('probabilities')\n",
        "  state_queue_idx: int = 0\n",
        "  state_chosenId_idx: int = 1\n",
        "  state_reward_idx: int = 2\n",
        "  cuda_sync: bool = False\n",
        "  json_path: str = '../Data/current_dataset.json'\n",
        "  jar_command: str = '/home/st-andrey-podivilov/java16/usr/lib/jvm/bellsoft-java16-amd64/bin/java -Dorg.jooq.no-logo=true -jar ../Game_env/usvm-jvm/build/libs/usvm-jvm-new.jar ../Game_env/jar_config.txt > ../Game_env/jar_log.txt'\n",
        "  actor_model_path: str = '../Game_env/actor_model.onnx'\n",
        "  path_: str = '../Config/HP_front.yaml'\n",
        "\n",
        "hp_front = HP_front(**yaml.load(open('../Config/HP_front.yml').read(), Loader=Loader))\n",
        "try:\n",
        "  hp_front.state_scheme = tuple(json.load(open('../Game_env/schemes.json'))['stateScheme'])\n",
        "except:\n",
        "  print('No scheme, using the default')\n",
        "hp_front.feature_names_scheme = tuple(hp_front.state_scheme[0])\n",
        "hp_front.features_dim = len(hp_front.feature_names_scheme)\n",
        "hp_front.tr_hash_idx = hp_front.trajectory_scheme.index('hash')\n",
        "hp_front.tr_path_idx = hp_front.trajectory_scheme.index('trajectory')\n",
        "hp_front.tr_name_idx = hp_front.trajectory_scheme.index('name')\n",
        "hp_front.tr_code_len_idx = hp_front.trajectory_scheme.index('statementsCount')\n",
        "\n",
        "\n",
        "@dataclass\n",
        "class HP_back:\n",
        "  shuffleTests: bool = False\n",
        "  maxAttentionLength: int = -1 # effectively it's a max number of actions to choose from during data gathering or inference (last in queue)\n",
        "  samplesPath: str = \"../Game_env/usvm-jvm/src/samples/java\"\n",
        "  gameEnvPath: str = \"../Game_env\"\n",
        "  dataPath: str = \"../Data\"\n",
        "  defaultAlgorithm: str = \"BFS\" # \"ForkDepthRandom\"\n",
        "  postprocessing: str = \"Argmax\" # \"Softmax\", \"None\"\n",
        "  mode: str = \"Both\" # \"Calculation\", \"Aggregation\"\n",
        "  inputShape: tuple = (-1, hp_front.features_dim) if hp_alg.critic_play else (1, -1, hp_front.features_dim)\n",
        "  dataConsumption: float = 100.0\n",
        "  hardTimeLimit: float = 30000\n",
        "  solverTimeLimit: float = 10000\n",
        "  maxConcurrency: int = 120\n",
        "  graphUpdate: str = \"Once\" # \"TestGeneration\"\n",
        "  logGraphFeatures: bool = False\n",
        "  gnnFeaturesCount: int = 8\n",
        "  rnnStateShape: tuple = (4, 1, 512)\n",
        "  useGnn: bool = True\n",
        "  useRnn: bool = True\n",
        "  path_: str = '../Config/HP_back.yaml'\n",
        "  inputJars: dict = None\n",
        "\n",
        "hp_back = HP_back(**yaml.load(open('../Config/HP_back.yml').read(), Loader=Loader))\n",
        "hp_back.inputShape = (-1, hp_front.features_dim) if hp_alg.critic_play else (1, -1, hp_front.features_dim)\n",
        "hp_back.maxAttentionLength = -1 if hp_alg.critic_play else hp_alg.max_attention_length\n",
        "\n",
        "\n",
        "def make_configs():\n",
        "  with open('../Config/HP_alg.yml', 'w') as outfile:\n",
        "      yaml.dump(asdict(hp_alg), outfile)\n",
        "  with open('../Config/HP_front.yml', 'w') as outfile:\n",
        "      yaml.dump(asdict(hp_front), outfile)\n",
        "  with open('../Config/HP_back.yml', 'w') as outfile:\n",
        "      yaml.dump(asdict(hp_back), outfile)\n",
        "  with open('../Game_env/jar_config.txt', 'w') as jar_config:\n",
        "    jar_config.write(json.dumps(asdict(hp_back)))\n",
        "\n",
        "make_configs()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t6GWfxrs9gFR",
        "jp-MarkdownHeadingCollapsed": true,
        "tags": []
      },
      "source": [
        "### Imports, meta"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BAP5ws3ofyyY",
        "tags": []
      },
      "outputs": [],
      "source": [
        "# %%capture\n",
        "\n",
        "# %env CUDA_DEVICE_ORDER=PCI_BUS_ID\n",
        "# %env CUDA_VISIBLE_DEVICES=4\n",
        "\n",
        "# from IPython.display import Javascript\n",
        "# def resize_colab_cell():\n",
        "#   display(Javascript('google.colab.output.setIframeHeight(0, true, {maxHeight: 600})'))\n",
        "# get_ipython().events.register('pre_run_cell', resize_colab_cell)\n",
        "\n",
        "import warnings\n",
        "import numpy as np\n",
        "from numpy import random\n",
        "import copy\n",
        "import inspect\n",
        "import torch\n",
        "from torch import nn\n",
        "import torch.onnx\n",
        "import json\n",
        "from tqdm import tqdm, trange\n",
        "from time import time\n",
        "import os\n",
        "import math\n",
        "from operator import itemgetter\n",
        "import torch.nn.functional as F\n",
        "from torch_geometric.loader import DataLoader as GraphDataLoader\n",
        "from torch_geometric.data import Data as GraphData\n",
        "from torch_geometric.nn import GCNConv\n",
        "\n",
        "\n",
        "# !pip install wandb\n",
        "import wandb\n",
        "\n",
        "# !pip install onnx==1.12\n",
        "# import onnx\n",
        "\n",
        "wandb.login()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VJbCP9HE9gFV",
        "jp-MarkdownHeadingCollapsed": true,
        "tags": []
      },
      "source": [
        "### Args/shortcuts\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z90K8IBjpUBu",
        "tags": []
      },
      "outputs": [],
      "source": [
        "batch_size = hp_alg.batch_size\n",
        "max_nactions = hp_alg.max_nactions\n",
        "device = 'cuda' if (torch.cuda.is_available() and hp_front.use_cuda_if_available) else 'cpu'\n",
        "td_gamma = hp_alg.td_gamma\n",
        "json_path = hp_front.json_path\n",
        "use_fork_discount = hp_alg.use_fork_discount\n",
        "batch_accumulation_steps = hp_alg.batch_accumulation_steps\n",
        "critic_play = hp_alg.critic_play\n",
        "features_dim = hp_front.features_dim\n",
        "train_condition = eval(hp_alg.train_condition_str)\n",
        "\n",
        "tr_hash_idx = hp_front.tr_hash_idx\n",
        "tr_path_idx = hp_front.tr_path_idx\n",
        "tr_name_idx = hp_front.tr_name_idx\n",
        "tr_prev_probs_idx = hp_front.tr_prev_probs_idx\n",
        "state_queue_idx = hp_front.state_queue_idx\n",
        "state_chosenId_idx = hp_front.state_chosenId_idx\n",
        "state_reward_idx = hp_front.state_reward_idx\n",
        "\n",
        "maybe_sync = torch.cuda.synchronize if hp_front.cuda_sync else (lambda *args: None)\n",
        "jar_command = hp_front.jar_command\n",
        "device"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "jp-MarkdownHeadingCollapsed": true,
        "tags": [],
        "id": "oRbyOLM1vdQP"
      },
      "source": [
        "### Utils"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eMNe1v3WvdQP"
      },
      "outputs": [],
      "source": [
        "def get_ass_dope_ids(l):\n",
        "  l = l.long()\n",
        "  if len(l.shape) == 1:\n",
        "    l = l[:, None]\n",
        "  ind = torch.LongTensor(np.indices(l.shape))\n",
        "  ind[-1] = l\n",
        "  return tuple(ind)\n",
        "\n",
        "def get_clip_eps(epoch):\n",
        "  if epoch < 5:\n",
        "    c = 0.3\n",
        "  elif epoch < 20:\n",
        "    c = 0.1\n",
        "  else:\n",
        "    c = 0.05\n",
        "  return c"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ndNrbEp79gFX",
        "jp-MarkdownHeadingCollapsed": true,
        "tags": []
      },
      "source": [
        "### Models, modules"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "abh-k_ytgHWT",
        "tags": []
      },
      "outputs": [],
      "source": [
        "class FFM_layer(torch.nn.Module):\n",
        "    def __init__(self, input_dim):\n",
        "      super().__init__()\n",
        "      # assert input_dim%2 == 0, 'even input_dim is more convenient'\n",
        "      # self.fourier_matrix = torch.nn.Linear(input_dim, int(input_dim), bias=False)\n",
        "      # nn.init.normal_(\n",
        "      #     self.fourier_matrix.weight,\n",
        "      #     std=1/np.sqrt(input_dim),\n",
        "      # )\n",
        "      # self.fourier_matrix.weight.requires_grad_(False)\n",
        "\n",
        "    def forward(self, x):\n",
        "      pre = x # self.fourier_matrix(x)\n",
        "      s = torch.sin(pre)\n",
        "      c = torch.cos(pre)\n",
        "      return torch.cat([x,s,c], dim=-1)\n",
        "\n",
        "\n",
        "class Attn_model(torch.nn.Module):\n",
        "  '''\n",
        "  PPO actor model.\n",
        "  '''\n",
        "  def __init__(\n",
        "    self,\n",
        "    d_model=512,\n",
        "    n_heads=8,\n",
        "    dim_feedforward=512,\n",
        "    dropout=0.0,\n",
        "  ):\n",
        "    super(Attn_model, self).__init__()\n",
        "    self.emb = nn.Sequential(\n",
        "      nn.LazyLinear(512),\n",
        "      nn.ReLU(),\n",
        "      nn.LayerNorm(512),\n",
        "      FFM_layer(512) if hp_alg.use_FFM else nn.Identity(),\n",
        "      nn.LazyLinear(d_model),\n",
        "      nn.ReLU(),\n",
        "    )\n",
        "    self.attn_EncoderLayer0 = nn.TransformerEncoderLayer(d_model, n_heads, dim_feedforward, dropout, batch_first=True)\n",
        "    if hp_alg.use_double_atten:\n",
        "      self.attn_EncoderLayer1 = nn.TransformerEncoderLayer(d_model, n_heads, dim_feedforward, dropout, batch_first=True)\n",
        "    self.head = nn.Sequential(\n",
        "      nn.LazyLinear(1),\n",
        "    )\n",
        "    self.sfmax = nn.Softmax(dim=-1)\n",
        "\n",
        "  def forward(self, x, mask=None):\n",
        "    x = self.emb(x)\n",
        "    x = self.attn_EncoderLayer0(x, src_key_padding_mask=mask)\n",
        "    if hp_alg.use_double_atten:\n",
        "      x = self.attn_EncoderLayer1(x, src_key_padding_mask=mask)\n",
        "    x = self.head(x).squeeze(-1)\n",
        "    if mask is None:\n",
        "      return self.sfmax(x)\n",
        "    inf_mask = mask.float().masked_fill(mask==True, -float('inf'))\n",
        "    x = self.sfmax(x + inf_mask)\n",
        "    return x\n",
        "\n",
        "\n",
        "class Q_Net(torch.nn.Module):\n",
        "  \"\"\"\n",
        "  Builds Q-network based on critic.\n",
        "  Currently is neither needed nor properly working (#reward_ind feature\n",
        "  is not a function of true reward. Might be solved, is not).\n",
        "  \"\"\"\n",
        "  def __init__(\n",
        "    self,\n",
        "    V_function,\n",
        "  ):\n",
        "    super(Q_Net, self).__init__()\n",
        "    self.V_function = V_function\n",
        "    self.reward_ind = hp_front.feature_names_scheme.index('logReward')\n",
        "\n",
        "  def forward(self, feature_batch):\n",
        "    v = self.V_function(feature_batch)\n",
        "    q = v + feature_batch[:, self.reward_ind][:, None].exp() - 1\n",
        "    return q\n",
        "\n",
        "\n",
        "class GNN_model(torch.nn.Module):\n",
        "  \"\"\"\n",
        "  GNN model used during jar-file calls.\n",
        "  Was trained once as a part of critic network.\n",
        "  \"\"\"\n",
        "  def __init__(\n",
        "    self,\n",
        "    num_input_features,\n",
        "    num_output_features\n",
        "  ):\n",
        "    super().__init__()\n",
        "    self.convs = nn.ModuleList([\n",
        "      GCNConv(num_input_features, 16),\n",
        "      GCNConv(16, 32),\n",
        "      GCNConv(32, 32),\n",
        "      GCNConv(32, 16),\n",
        "      GCNConv(16, num_output_features),\n",
        "    ])\n",
        "    self.dropout_probs = [\n",
        "      None,\n",
        "      None,\n",
        "      0.1,\n",
        "      0.3,\n",
        "    ]\n",
        "\n",
        "  def forward(self, data_x, data_edge_index):\n",
        "    x, edge_index = data_x, data_edge_index\n",
        "    x = self.convs[0](x, edge_index)\n",
        "    for conv, prob in zip(self.convs[1:], self.dropout_probs):\n",
        "      x = F.relu(x)\n",
        "      if prob is not None:\n",
        "        x = F.dropout(x, p=prob, training=self.training)\n",
        "      x = conv(x, edge_index)\n",
        "    return x\n",
        "\n",
        "\n",
        "class V_cell(torch.nn.Module):\n",
        "  \"\"\"\n",
        "  We use V_cell inner representations to enrich a world embedding.\n",
        "  Was trained once on a proxy seq2seq task (predicting Returns along a trajectory).\n",
        "  \"\"\"\n",
        "  def __init__(\n",
        "    self,\n",
        "    inner_hidden_size = 512,\n",
        "    top_hidden_size=32,\n",
        "    bias=True,\n",
        "    batch_first=True,\n",
        "  ):\n",
        "    super(V_cell, self).__init__()\n",
        "    self.inner_hidden_size = inner_hidden_size\n",
        "    self.top_hidden_size = top_hidden_size\n",
        "    self.num_layers = 2\n",
        "    self.emb = nn.Sequential(\n",
        "      nn.LazyLinear(128),\n",
        "      nn.LayerNorm(128),\n",
        "    )\n",
        "    self.lstm_cell = nn.LSTMCell(input_size=128, hidden_size=inner_hidden_size)\n",
        "    self.dropout = nn.Dropout(p=0.2)\n",
        "    self.lstm_cell2 = nn.LSTMCell(input_size=inner_hidden_size, hidden_size=top_hidden_size)\n",
        "    self.head = nn.Sequential(\n",
        "      nn.Dropout(p=0.1),\n",
        "      nn.LazyLinear(128),\n",
        "      nn.ReLU(),\n",
        "      nn.LazyLinear(1),\n",
        "    )\n",
        "\n",
        "  def forward(self, input_batch, prev_state_batch):\n",
        "    \"\"\"\n",
        "    input_batch: [batch_size, feature_size]\n",
        "    prev_state_batch: [2*num_layers, batch_size, inner_hidden_size]\n",
        "    Outputs: [batch_size, 1], [2*num_layers, batch_size, inner_hidden_size], [batch_size, top_hidden_size+1]\n",
        "\n",
        "    Weird stuff happens to top hidden h2 so that state_batch can be kept nice and simple.\n",
        "    \"\"\"\n",
        "    h_p = prev_state_batch[list(np.arange(self.num_layers) * 2)]\n",
        "    c_p = prev_state_batch[list(np.arange(self.num_layers) * 2 + 1)]\n",
        "    x = self.emb(input_batch)\n",
        "    h1, c1 = self.lstm_cell(x, (h_p[0], c_p[0]))\n",
        "    h1 = self.dropout(h1)\n",
        "    h2, c2 = self.lstm_cell2(h1, (h_p[1][:, :self.top_hidden_size], c_p[1][:, :self.top_hidden_size]))\n",
        "    v = self.head(h2)\n",
        "    h2_rep = h2.repeat(1, self.inner_hidden_size//self.top_hidden_size)\n",
        "    c2_rep = c2.repeat(1, self.inner_hidden_size//self.top_hidden_size)\n",
        "    return v, torch.stack((h1, c1, h2_rep, c2_rep), dim=0), torch.cat([h2, v], dim=-1).detach()\n",
        "\n",
        "\n",
        "def get_mlp_setup():\n",
        "  '''\n",
        "  Returns model and optimizer for Mr. critic of PPO.\n",
        "  '''\n",
        "    mlp = nn.Sequential(\n",
        "        nn.LazyLinear(512),\n",
        "        nn.ReLU(),\n",
        "        nn.Linear(512,256),\n",
        "        nn.LayerNorm(256),\n",
        "        FFM_layer(256) if hp_alg.use_FFM else nn.Identity(),\n",
        "        nn.LazyLinear(512),\n",
        "        nn.ReLU(),\n",
        "        nn.Linear(512,512),\n",
        "        nn.ReLU(),\n",
        "        nn.Linear(512,1),\n",
        "    ).to(device)\n",
        "    mlp_opt = torch.optim.AdamW(mlp.parameters(), lr=hp_alg.critic_lr, weight_decay=0.1, betas=(0.9, 0.999))\n",
        "    return mlp, mlp_opt\n",
        "\n",
        "\n",
        "def get_attn_setup(sched_total_iters=hp_alg.epochs):\n",
        "  '''\n",
        "  Returns model, optimizer and scheduler for Mr. actor of PPO.\n",
        "  '''\n",
        "  attn_model = Attn_model().to(device)\n",
        "  opt = torch.optim.AdamW(attn_model.parameters(), lr=hp_alg.actor_lr, weight_decay=1e-2,)\n",
        "  scheduler = torch.optim.lr_scheduler.LinearLR(opt, start_factor=1, end_factor=0.1, total_iters=sched_total_iters, verbose=False)\n",
        "  return attn_model, opt, scheduler"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "jp-MarkdownHeadingCollapsed": true,
        "tags": [],
        "id": "Lfn5Oj8qvdQR"
      },
      "source": [
        "### Logger"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9kKdAlkPvdQR"
      },
      "outputs": [],
      "source": [
        "class Logger:\n",
        "  \"\"\"\n",
        "  Supporting class, to be expanded.\n",
        "  Intended to store logging utils and relevant data.\n",
        "  \"\"\"\n",
        "  def __init__(\n",
        "      self,\n",
        "  ):\n",
        "    self.grad_a = None\n",
        "    self.grad_c = None\n",
        "    self.weight_a = None\n",
        "    self.weight_c = None\n",
        "    self.log_gamma = torch.tensor(0.95).to(device)\n",
        "    self.between_logs = hp_front.between_logs\n",
        "    self.timer = {}\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def link_models(self,):\n",
        "    self.grad_a = [p.grad.detach() for p in self.actor.parameters() if p.requires_grad]\n",
        "    self.weight_a = [p.detach() for p in self.actor.parameters()]\n",
        "    self.grad_c = [p.grad.detach() for p in self.critic.parameters() if p.requires_grad]\n",
        "    self.weight_c = [p.detach() for p in self.critic.parameters()]\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def list_norm(self, l, p=2):\n",
        "    n = 0\n",
        "    for t in l:\n",
        "      n += t.detach().norm(p) ** p\n",
        "    return n.item() ** (1/p)\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def list_cos_dist(self, a, b):\n",
        "    a_norm = self.list_norm(a, 2)\n",
        "    b_norm = self.list_norm(b, 2)\n",
        "    product = sum([torch.dot(torch.flatten(a[i]), torch.flatten(b[i])).item() for i in range(len(a))])\n",
        "    return product/(a_norm*b_norm)\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def on_list(self, a, b, operation):\n",
        "    assert len(a) == len(b), 'lists lengths differ'\n",
        "    return [operation(a[i], b[i]) for i in range(len(a))]\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def running_mean(self, a, b):\n",
        "    return [a[i].mul(self.log_gamma) + b[i].mul(1 - self.log_gamma) for i in range(len(a))]\n",
        "\n",
        "\n",
        "logger = Logger()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qBhLnmrnFDAu",
        "jp-MarkdownHeadingCollapsed": true,
        "tags": []
      },
      "source": [
        "### Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ov4TPknTgHSy",
        "tags": []
      },
      "outputs": [],
      "source": [
        "class Trajectories:\n",
        "  \"\"\"\n",
        "  Contains all kinds of data in a form of tensors.\n",
        "  realized are raw and derived features of visited states.\n",
        "  queues is a (3-dim) list of actions features by state.\n",
        "  Action and state embeddings only differ in RNN part, fyi.\n",
        "  \"\"\"\n",
        "  def __init__(self):\n",
        "    self.train_condition = eval(hp_alg.train_condition_str)\n",
        "    self.td_gamma = hp_alg.td_gamma\n",
        "    self.j_file = None\n",
        "    self.feature_names, self.feature_names2ids = None, None\n",
        "    # dict of train tensors f, f_n, r, R, is_last, queue_len, chosen_acts, tr_id +\n",
        "    # + list of queue tensors + list of previous probabilities tensors\n",
        "    self.realized, self.queues, self.prev_probs = None, None, None\n",
        "    # Psi is GAE critic estimates computed by trainer.GAE()\n",
        "    self.Psi = None\n",
        "    self.sampled_ids_list = []\n",
        "    self.sampled_queue_lengths_list = []\n",
        "\n",
        "\n",
        "  def gather_n_store(self, actor_model=None, json_path=hp_front.json_path):\n",
        "    \"\"\"\n",
        "    Plays, collects trajectories into json, then transforms json to tensors\n",
        "    \"\"\"\n",
        "    self.update_data_on_path(actor_model, json_path)\n",
        "    self.store_from_json(json_path)\n",
        "\n",
        "\n",
        "  def evaluate_val_train(self, verbose=True,):\n",
        "    \"\"\"\n",
        "    Evaluates val and train subset\n",
        "    \"\"\"\n",
        "    log_val = self.evaluate_data(eval_condition = (lambda tr: not self.train_condition(tr)),\n",
        "                                 wandb_prefix = 'val',\n",
        "                                 verbose=verbose,)\n",
        "    log_train = self.evaluate_data(eval_condition = self.train_condition,\n",
        "                                   wandb_prefix = 'train',\n",
        "                                   verbose=verbose,)\n",
        "    return log_val, log_train\n",
        "\n",
        "\n",
        "  def store_from_json(self, json_path=hp_front.json_path):\n",
        "    \"\"\"\n",
        "    Extracts trajectories from json to a torch-friendly form\n",
        "    \"\"\"\n",
        "    time_before = time()\n",
        "    self.j_file = json.load(open(json_path))\n",
        "    print('json loading time: ', time()-time_before)\n",
        "    self.feature_names = self.j_file['stateScheme'][0]\n",
        "    self.feature_names2ids = {self.feature_names[i]:i for i in range(len(self.feature_names))}\n",
        "    time_before = time()\n",
        "    self.realized, self.queues, self.prev_probs = self.j2torch(self.j_file)\n",
        "    print('j2torch time: ', time()-time_before)\n",
        "    self.Psi = None\n",
        "\n",
        "\n",
        "  def mean_code_length(self):\n",
        "    train_lines = 0\n",
        "    val_lines = 0\n",
        "    for tr in self.j_file['paths']:\n",
        "      if self.train_condition(tr):\n",
        "        train_lines += tr[hp_front.tr_code_len_idx]\n",
        "      else:\n",
        "        val_lines += tr[hp_front.tr_code_len_idx]\n",
        "    n_tr = self.get_properties()['number of traj-s']\n",
        "    n_val_tr = self.get_properties()['number of validation traj-s']\n",
        "    return train_lines/(n_tr-n_val_tr+1e-9), val_lines/(n_val_tr+1e-9)\n",
        "\n",
        "\n",
        "  def j2torch(self, j_file):\n",
        "    \"\"\"\n",
        "    Transforms json to data tensors.\n",
        "    Queues are flipped for reasons related to batch truncation and padding. Actions numeration is changed accordingly.\n",
        "    \"\"\"\n",
        "    features, features_next, rewards, Returns, is_last, queue_lengths, chosen_actions, queues = [], [], [], [], [], [], [], []\n",
        "    if not critic_play:\n",
        "      prev_probs = []\n",
        "    logger.timer = {\n",
        "      'tr_prev_probs': 0,\n",
        "      'tr_features':0,\n",
        "      'tr_queues':0,\n",
        "    }\n",
        "    for j, tr in enumerate(self.j_file['paths']):\n",
        "      if not self.train_condition(tr):\n",
        "        continue\n",
        "      if not critic_play:\n",
        "        tr_prev_probs = [torch.Tensor(dstrn).flip(dims=[0]).to(device) for dstrn in tr[hp_front.tr_prev_probs_idx]][1:] + [torch.tensor([1.0]).to(device)]\n",
        "        prev_probs += tr_prev_probs\n",
        "      tr = tr[hp_front.tr_path_idx]\n",
        "      if hp_alg.critic_play: # use features as an action emb, reward is what we get for that action\n",
        "        tr_rewards = [tr[i][hp_front.state_reward_idx] for i in range(len(tr))]\n",
        "      else: # use features as a state emb, reward is what we get the next step\n",
        "        tr_rewards = [tr[i][hp_front.state_reward_idx] for i in range(len(tr))][1:] + [0]\n",
        "      rewards += tr_rewards\n",
        "      do_discount = torch.ones(len(tr))\n",
        "      if use_fork_discount:\n",
        "        is_cfg_fork_idx = self.j_file['stateScheme'].index('is_cfg_fork')\n",
        "        do_discount = [tr[i][is_cfg_fork_idx] for i in range(len(tr))]\n",
        "      tr_Returns = self.tr_rewards_to_returns(tr_rewards, do_discount)\n",
        "      Returns += tr_Returns\n",
        "\n",
        "      tr_chosen_actions = [tr[i][hp_front.state_chosenId_idx] for i in range(len(tr))][1:] + [0]\n",
        "      chosen_actions += tr_chosen_actions\n",
        "\n",
        "      time0=time()\n",
        "      tr_features = [tr[i][hp_front.state_queue_idx][tr[i][hp_front.state_chosenId_idx]] for i in range(len(tr))]\n",
        "      is_last += [0]*(len(tr_features)-1) + [1]\n",
        "      features += tr_features\n",
        "      features_next += tr_features[1:] + [[-1]*hp_front.features_dim]\n",
        "      logger.timer['tr_features'] += time() - time0\n",
        "\n",
        "      time0=time()\n",
        "      tr_queues = [torch.Tensor(tr[i][hp_front.state_queue_idx]).flip(dims=[0]).to(device) for i in range(len(tr))][1:]\n",
        "      tr_queues += [torch.zeros_like(torch.Tensor([tr[0][hp_front.state_queue_idx][0]])).to(device)]\n",
        "      logger.timer['tr_queues'] += time()-time0\n",
        "\n",
        "      tr_queue_lengths = [len(q) for q in tr_queues]\n",
        "      queues += tr_queues\n",
        "      queue_lengths += tr_queue_lengths\n",
        "    rewards = torch.Tensor(rewards).to(device)\n",
        "    features = torch.Tensor(features).to(device)\n",
        "    features_next = torch.Tensor(features_next).to(device)\n",
        "    Returns = torch.Tensor(Returns).to(device)\n",
        "    is_last = torch.Tensor(is_last).to(device)\n",
        "    queue_lengths = torch.LongTensor(queue_lengths).to(device)\n",
        "    chosen_actions = queue_lengths - torch.LongTensor(chosen_actions).to(device) - 1 # we flip queue and numeration of actions\n",
        "    # print(logger.timer)\n",
        "    realized = {\n",
        "      'features': features,\n",
        "      'features_next': features_next,\n",
        "      'rewards': rewards,\n",
        "      'Returns': Returns,\n",
        "      'is_last': is_last,\n",
        "      'queue_lengths': queue_lengths,\n",
        "      'chosen_actions': chosen_actions,\n",
        "    }\n",
        "    if critic_play:\n",
        "      prev_probs = None\n",
        "    return realized, queues, prev_probs\n",
        "\n",
        "\n",
        "  def get_n_train_states(self):\n",
        "    return len(self.realized['features'])\n",
        "\n",
        "\n",
        "  def get_properties(self):\n",
        "    queue_lengths = np.array([q.shape[0] for q in self.queues])\n",
        "    longest_queue_ids = np.argmax(queue_lengths)\n",
        "    tr_lengths = np.array([len(tr[hp_front.tr_path_idx]) for tr in self.j_file['paths']])\n",
        "    prop = {\n",
        "            'traj_length mean, median, max': (f'{tr_lengths.mean():.2f}', np.median(tr_lengths), tr_lengths.max()),\n",
        "            'queue max length, idx': (self.queues[longest_queue_ids].shape[0], longest_queue_ids),\n",
        "            'number of train states': self.get_n_train_states(),\n",
        "            'number of traj-s': len(self.j_file['paths']),\n",
        "            'number of validation traj-s': sum([not self.train_condition(tr) for tr in self.j_file['paths']]),\n",
        "           }\n",
        "    return prop\n",
        "\n",
        "\n",
        "  def tr_rewards_to_returns(self, tr_rewards, do_discount=None):\n",
        "    if do_discount is None:\n",
        "      do_discount = [1]*(len(tr_rewards))\n",
        "    tr_R = [0]*(len(tr_rewards))\n",
        "    for i in range(len(tr_rewards)-2, -1, -1):\n",
        "        tr_R[i] = tr_rewards[i] + (self.td_gamma**do_discount[i]) * tr_R[i+1]\n",
        "    return tr_R\n",
        "\n",
        "\n",
        "  def sample_ids(self, batch_size=batch_size):\n",
        "    \"\"\"\n",
        "    Assembles ids for several consiquent batches based on queues lengths.\n",
        "    \"\"\"\n",
        "    if len(self.sampled_ids_list) == 0:\n",
        "      sampled_ids = torch.LongTensor(random.choice(self.get_n_train_states(), size=hp_alg.batch_size*hp_alg.batch_accumulation_steps)).to(device)\n",
        "      sampled_queue_lengths = self.realized['queue_lengths'][sampled_ids]\n",
        "      sampled_queue_lengths, sorting_ids = torch.sort(sampled_queue_lengths)\n",
        "      sampled_ids = sampled_ids[sorting_ids]\n",
        "      self.sampled_ids_list = list(torch.split(sampled_ids, batch_size))\n",
        "      self.sampled_queue_lengths_list = list(torch.split(sampled_queue_lengths, batch_size))\n",
        "    return self.sampled_ids_list.pop(0), self.sampled_queue_lengths_list.pop(0)\n",
        "\n",
        "\n",
        "  def sample_batch(self, batch_size=batch_size):\n",
        "    ids, queue_lengths = self.sample_ids(batch_size=batch_size)\n",
        "    sampled_realized = {k: self.realized[k][ids] for k in self.realized.keys()}\n",
        "    sampled_queues = itemgetter(*list(ids))(self.queues)\n",
        "    sampled_prev_probs = itemgetter(*list(ids))(self.prev_probs)\n",
        "    padded_length = torch.minimum(queue_lengths.max(), torch.tensor(hp_alg.max_nactions))\n",
        "    queues_tensor = torch.zeros(hp_alg.batch_size, padded_length, hp_front.features_dim).to(device)\n",
        "    prev_probs_tensor = torch.zeros(hp_alg.batch_size, padded_length).to(device)\n",
        "    pad_mask = torch.ones(hp_alg.batch_size, padded_length).to(device)\n",
        "\n",
        "    start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)\n",
        "    start.record()\n",
        "\n",
        "    for i, q in enumerate(sampled_queues):\n",
        "      l = min(q.shape[0], padded_length)\n",
        "      queues_tensor[i, 0 : l, :] = q[: l, :]\n",
        "      pad_mask[i, 0 : l] = 0\n",
        "    pad_mask = pad_mask.bool()\n",
        "\n",
        "    for i, p in enumerate(sampled_prev_probs):\n",
        "      l = min(p.shape[0], padded_length)\n",
        "      prev_probs_tensor[i, 0 : l] = p[: l]\n",
        "\n",
        "    end.record()\n",
        "    maybe_sync()\n",
        "    logger.timer['sample inside loop'] += start.elapsed_time(end)/1000\n",
        "    if not self.Psi is None:\n",
        "      Psi = self.Psi[ids]\n",
        "    else:\n",
        "      Psi = None\n",
        "    return sampled_realized, Psi, queues_tensor, pad_mask, prev_probs_tensor\n",
        "\n",
        "\n",
        "  def sample_realized_batch(self, n=hp_alg.batch_size):\n",
        "    '''\n",
        "    A lightened version of sample_batch, can be used in critic_play (Policy Iteration) mode.\n",
        "    '''\n",
        "    ids = torch.tensor(random.choice(self.get_n_train_states(), size=n)).long()\n",
        "    sampled = {k: self.realized[k][ids] for k in self.realized.keys()}\n",
        "    return sampled\n",
        "\n",
        "\n",
        "  def update_data_on_path(self, actor_model=None, path=hp_front.json_path):\n",
        "    \"\"\"\n",
        "    Communication with jar file on a server.\n",
        "    \"\"\"\n",
        "    if actor_model is None:\n",
        "      player_name = 'Heuristic' # using default heuristics\n",
        "      os.system('rm -f ' + hp_front.actor_model_path)\n",
        "    else:\n",
        "      player_name = 'NN'\n",
        "      total_features_dim = hp_front.features_dim\n",
        "      if hasattr(actor_model, 'sfmax'):\n",
        "        shape = [1, 1, total_features_dim]\n",
        "      else:\n",
        "        shape = [1, total_features_dim]\n",
        "      x = torch.randn(*shape, requires_grad=True).to(device)\n",
        "      torch_model = actor_model.eval()\n",
        "      torch_out = torch_model(x)\n",
        "      torch.onnx.export(torch_model,\n",
        "                        x,\n",
        "                        '../Game_env/actor_model.onnx',\n",
        "                        opset_version=15,\n",
        "                        export_params=True,\n",
        "                        input_names = ['input'],   # the model's input names\n",
        "                        output_names = ['output'],\n",
        "                        dynamic_axes={'input' : {0 : 'batch_size',\n",
        "                                                 1 : 'n_actions',\n",
        "                                                },    # variable length axes\n",
        "                                      'output' : {0 : 'batch_size',\n",
        "                                                  1 : 'n_actions',\n",
        "                                                },\n",
        "                                     },\n",
        "                        )\n",
        "    time_before = time()\n",
        "    os.system(hp_front.jar_command)\n",
        "    print(player_name, ' data gathering time: ', time() - time_before)\n",
        "\n",
        "\n",
        "  def evaluate_data(self,\n",
        "           eval_condition = (lambda tr: not self.train_condition(tr)),\n",
        "           wandb_prefix = 'val',\n",
        "           verbose=True,\n",
        "           factors=torch.Tensor([1, hp_alg.td_gamma]),\n",
        "           ):\n",
        "    for f in factors:\n",
        "      size = 0\n",
        "      tr_lengths = []\n",
        "      trs_R = []\n",
        "      for tr in self.j_file['paths']:\n",
        "        if not eval_condition(tr):\n",
        "          continue\n",
        "        tr=tr[1]\n",
        "        size += len(tr)\n",
        "        tr_lengths += [len(tr)]\n",
        "        tr_rewards = [tr[i][state_reward_idx] for i in range(len(tr))] # rewards list is not shifted, has a different purpose this time\n",
        "        trs_R += [0]\n",
        "        do_discount = torch.Tensor([1]*len(tr))\n",
        "        if use_fork_discount:\n",
        "          is_cfg_fork_idx = self.j_file['stateScheme'].index('is_cfg_fork')\n",
        "          do_discount = [tr[i][is_cfg_fork_idx] for i in range(len(tr))]\n",
        "        for i in range(len(tr_rewards)-1, -1, -1):\n",
        "          trs_R[-1] = tr_rewards[i] + (f ** do_discount[i]) * trs_R[-1]\n",
        "      log = {}\n",
        "      log[f'{wandb_prefix} size'] = size\n",
        "      log[f'{wandb_prefix}_eval/mean {f:.3f} discount '] = torch.Tensor(trs_R).mean()\n",
        "      log[f'{wandb_prefix}_eval/median {f:.3f} discount '] = torch.Tensor(trs_R).median()\n",
        "      log[f'{wandb_prefix} Return by trjs {f:.3f} (previous epoch)'] = wandb.Histogram(np_histogram=np.histogram(trs_R, bins=30, ))\n",
        "      log['code length train mean'], log['code length val mean'] = self.mean_code_length()\n",
        "      if verbose:\n",
        "          wandb.log(log.copy())\n",
        "      log['Returns'] = trs_R\n",
        "      log[f'{wandb_prefix} lengths'] = tr_lengths\n",
        "    return log"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tRSEdgzW9gFc",
        "jp-MarkdownHeadingCollapsed": true,
        "tags": []
      },
      "source": [
        "### Trainer\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NsLLouOFgHQZ",
        "tags": []
      },
      "outputs": [],
      "source": [
        "class NN_Trainer:\n",
        "  def __init__(\n",
        "      self,\n",
        "      NN_setup,\n",
        "      trajectories,\n",
        "      train_progress,\n",
        "      clip_eps = None,\n",
        "      batch_size=hp_alg.batch_size,\n",
        "      n_batches=1000,\n",
        "      target_update_steps=15,\n",
        "      td_gamma=hp_alg.td_gamma,\n",
        "      ):\n",
        "    self.n_batches = n_batches\n",
        "    self.batch_number = -1\n",
        "    self.td_gamma = td_gamma\n",
        "    self.gae_gamma = 0.7\n",
        "    self.clip_eps = clip_eps\n",
        "    self.actor = NN_setup['actor']\n",
        "    self.actor_opt = NN_setup['actor_opt']\n",
        "    self.actor_sched = NN_setup['actor_sched']\n",
        "    self.prev_actor = copy.deepcopy(self.actor).eval()\n",
        "    self.critic = NN_setup['critic'].train()\n",
        "    self.target_critic = copy.deepcopy(self.critic).eval()\n",
        "    self.critic_opt = NN_setup['critic_opt']\n",
        "    self.trajectories = trajectories\n",
        "    self.batch_size = batch_size\n",
        "    self.target_update_steps = target_update_steps\n",
        "    self.train_progress = train_progress\n",
        "    self.log = {}\n",
        "\n",
        "  def get_each_loss(self, sampled_batch,):\n",
        "    \"\"\"\n",
        "    Computes losses for actor, critic and exploration (loss_ent) within PPO algorithm.\n",
        "    Decisions were made to avoid python loops --\n",
        "    varying action space is not particularly batch-friendly.\n",
        "    \"\"\"\n",
        "    sampled_realized, Psi, queues_tensor, pad_mask, prev_probs_tensor = sampled_batch\n",
        "    self.log['Returns mean'] = sampled_realized['Returns'].mean().item()\n",
        "    self.log['rewards mean'] = sampled_realized['rewards'].mean().item()\n",
        "    self.log['queue length max'] = sampled_realized['queue_lengths'].max().item()\n",
        "    probs = self.actor(queues_tensor, mask=pad_mask)\n",
        "    self.log['max prob mean'] = probs.max(dim=-1).values.mean().item()\n",
        "    self.log['40 max prob quantile'] = torch.quantile(probs.max(dim=-1).values, 0.4).item()\n",
        "\n",
        "    start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)\n",
        "    start.record()\n",
        "\n",
        "    values = self.critic(sampled_realized['features']).squeeze(-1)\n",
        "    with torch.no_grad():\n",
        "      next_values = self.target_critic(sampled_realized['features_next']).squeeze(-1)\n",
        "\n",
        "    end.record()\n",
        "    maybe_sync()\n",
        "    logger.timer['values'] += start.elapsed_time(end)/1000\n",
        "\n",
        "    self.log['V-func mean'] = torch.mean(values.detach()).item()\n",
        "    self.log['V-func stdev'] = torch.std(values.detach()).item()\n",
        "    # hist = wandb.Histogram(np_histogram=np.histogram(values.detach().to('cpu'), bins=40, ))\n",
        "    # self.log['V-func hist'] = hist\n",
        "\n",
        "    # critic loss\n",
        "    TD = values - ((sampled_realized['rewards'] + next_values.detach() * self.td_gamma) * (1-sampled_realized['is_last']))\n",
        "    MC = (values - sampled_realized['Returns']).abs().mean()/100\n",
        "    loss_c = (TD**2).mean() + MC\n",
        "\n",
        "    self.log['TD loss'] = (TD**2).mean().item()\n",
        "    self.log['MC loss'] = MC.item()\n",
        "    # hist = wandb.Histogram(np_histogram=np.histogram(TD.detach().to('cpu'), bins=20, ))\n",
        "    # self.log['TD hist'] = hist\n",
        "\n",
        "    # entropy loss\n",
        "    t_entropy_loss = time()\n",
        "\n",
        "    entropies = - probs * torch.log(torch.max(torch.tensor(1e-40), probs))\n",
        "    entropy_regularizer = torch.log(torch.minimum(sampled_realized['queue_lengths']+1, torch.tensor(hp_alg.max_nactions+1))).to(device)\n",
        "    entropy_by_state_reg = torch.sum(entropies, dim=-1) / entropy_regularizer\n",
        "    loss_ent = -entropy_by_state_reg.mean()\n",
        "\n",
        "    logger.timer['entropy loss'] += time()-t_entropy_loss\n",
        "\n",
        "    # actor loss\n",
        "    start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)\n",
        "    start.record()\n",
        "\n",
        "    probs_chosen = probs[get_ass_dope_ids(sampled_realized['chosen_actions'])].squeeze(-1)\n",
        "    prev_probs_chosen = prev_probs_tensor[get_ass_dope_ids(sampled_realized['chosen_actions'])].squeeze(-1)\n",
        "    ratios = (probs_chosen / (prev_probs_chosen.detach()+1e-4)).to(device)\n",
        "\n",
        "    clipped = torch.clip(ratios, min=1-self.clip_eps, max=1+self.clip_eps)\n",
        "    Adv = - (TD * (1-sampled_realized['is_last'])).detach() # to not affect loss_a logs\n",
        "    self.log['Adv mean'] = Adv.sum() / (Adv.numel() - sampled_realized['is_last'].sum())\n",
        "    if self.trajectories.Psi is None:\n",
        "      loss_a = - torch.min(ratios*Adv, clipped*Adv).mean()\n",
        "    else:\n",
        "      loss_a = - torch.min(ratios*Psi, clipped*Psi).mean()\n",
        "    self.log['clipped to all'] = (clipped*Adv < ratios*Adv).long().sum().item() / ratios.numel()\n",
        "    self.log['ratios mean'] = ratios.mean().item()\n",
        "    end.record()\n",
        "    maybe_sync()\n",
        "    logger.timer['actor loss'] += start.elapsed_time(end)/1000\n",
        "    return loss_a, loss_c, loss_ent\n",
        "\n",
        "\n",
        "  def critic_loss(self,\n",
        "                  sampled_realized,\n",
        "    ):\n",
        "    \"\"\"\n",
        "    Temporal difference loss for PPO\n",
        "    \"\"\"\n",
        "    values = self.critic(sampled_realized['features']).squeeze(-1)\n",
        "    with torch.no_grad():\n",
        "      next_values = self.target_critic(sampled_realized['features_next']).squeeze(-1)\n",
        "    self.log['V-func mean'] = torch.mean(values.detach()).item()\n",
        "    TD = values - ((sampled_realized['rewards'] + next_values * self.td_gamma) * (1-sampled_realized['is_last']))\n",
        "    MC = (values - sampled_realized['Returns']).abs().mean()/100\n",
        "    critic_loss = (TD**2).mean() + MC\n",
        "    self.log['TD loss'] = (TD**2).mean().item()\n",
        "    self.log['MC loss'] = MC.item()\n",
        "    return critic_loss\n",
        "\n",
        "\n",
        "  def q_loss(self,\n",
        "                  sampled_realized,\n",
        "    ):\n",
        "    \"\"\"\n",
        "    Temporal difference loss for PI\n",
        "    \"\"\"\n",
        "    q = self.critic(sampled_realized['features']).squeeze(-1)\n",
        "    with torch.no_grad():\n",
        "      next_q = self.target_critic(sampled_realized['features_next']).squeeze(-1)\n",
        "    self.log['Q-func mean'] = torch.mean(q.detach()).item()\n",
        "    TD = q - (sampled_realized['rewards'] + next_q * self.td_gamma * (1-sampled_realized['is_last']))\n",
        "    MC = (q - sampled_realized['Returns']).abs().mean()/100\n",
        "    q_loss = (TD**2).mean() + MC\n",
        "    self.log['TD loss'] = (TD**2).mean().item()\n",
        "    self.log['MC loss'] = MC.item()\n",
        "    return q_loss\n",
        "\n",
        "\n",
        "  def learn_q(self, n_batches=None):\n",
        "    \"\"\"\n",
        "    Approximates Q-function of previous policy by iterating over collected data.\n",
        "    critic_play mode required (check realized['rewards'] lag in PPO setup).\n",
        "    \"\"\"\n",
        "    if n_batches is None:\n",
        "      n_batches = self.n_batches\n",
        "    for i in trange(n_batches):\n",
        "      self.batch_number = i\n",
        "      if self.batch_number % self.target_update_steps == 0:\n",
        "        self.target_critic = copy.deepcopy(self.critic).eval()\n",
        "      sampled_realized = self.trajectories.sample_realized_batch(n=2048)\n",
        "      loss = self.q_loss(sampled_realized)\n",
        "      self.critic_opt.zero_grad()\n",
        "      loss.backward()\n",
        "      torch.nn.utils.clip_grad_norm_(self.critic.parameters(), 30)\n",
        "      self.critic_opt.step()\n",
        "      if self.batch_number % (logger.between_logs+1) == 0:\n",
        "        self.log['weight critic'] = logger.list_norm([p for p in self.critic.parameters() if p.requires_grad], 2)\n",
        "        self.log['grad critic'] = logger.list_norm([p.grad for p in self.critic.parameters() if p.requires_grad], 2)\n",
        "        wandb.log(self.log)\n",
        "\n",
        "\n",
        "  def learn_new_policy(self, prev_actor=None):\n",
        "    \"\"\"\n",
        "    Implements one learning cycle over collected dataset.\n",
        "    \"\"\"\n",
        "    if prev_actor is None:\n",
        "      self.prev_actor = copy.deepcopy(self.actor).eval()\n",
        "    else:\n",
        "      self.prev_actor = prev_actor\n",
        "    logger.timer = {'values': 0,\n",
        "                    'entropy loss': 0,\n",
        "                    'actor loss': 0,\n",
        "                    'total loss': 0,\n",
        "                    'optimizers step': 0,\n",
        "                    'loss.backward': 0,\n",
        "                    'sample batch': 0,\n",
        "                    'sample ids': 0,\n",
        "                    'sample inside loop': 0,\n",
        "                    }\n",
        "    for i in trange(self.n_batches):\n",
        "      self.batch_number = i\n",
        "      if self.batch_number % self.target_update_steps == 0:\n",
        "        self.target_critic = copy.deepcopy(self.critic).eval()\n",
        "\n",
        "      start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)\n",
        "      start.record()\n",
        "\n",
        "      sampled_batch = self.trajectories.sample_batch(self.batch_size)\n",
        "\n",
        "      end.record()\n",
        "      maybe_sync()\n",
        "      logger.timer['sample batch'] += start.elapsed_time(end)/1000\n",
        "\n",
        "      start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)\n",
        "      start.record()\n",
        "\n",
        "      loss_a, loss_c, loss_ent = self.get_each_loss(sampled_batch)\n",
        "      loss = (loss_a + loss_c + loss_ent/(10 + self.train_progress*20)) / batch_accumulation_steps\n",
        "\n",
        "      end.record()\n",
        "      maybe_sync()\n",
        "      logger.timer['total loss'] += start.elapsed_time(end)/1000\n",
        "\n",
        "      start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)\n",
        "      start.record()\n",
        "\n",
        "      loss.backward()\n",
        "\n",
        "      end.record()\n",
        "      maybe_sync()\n",
        "      logger.timer['loss.backward'] += start.elapsed_time(end)/1000\n",
        "\n",
        "      if self.batch_number % batch_accumulation_steps == 0:\n",
        "        t_optimizers_step = time()\n",
        "        torch.nn.utils.clip_grad_norm_(self.critic.parameters(), 30)\n",
        "        torch.nn.utils.clip_grad_norm_(self.actor.parameters(), 5)\n",
        "        self.actor_opt.step()\n",
        "        self.critic_opt.step()\n",
        "        self.log.update({\n",
        "            'grad actor': logger.list_norm([p.grad for p in self.actor.parameters() if p.requires_grad], 2),\n",
        "            'grad critic': logger.list_norm([p.grad for p in self.critic.parameters() if p.requires_grad], 2),\n",
        "        })\n",
        "        self.critic_opt.zero_grad()\n",
        "        self.actor_opt.zero_grad()\n",
        "        logger.timer['optimizers step'] += time()-t_optimizers_step\n",
        "      if self.batch_number % (logger.between_logs+1) == 0:\n",
        "        self.log.update({\n",
        "            'loss actor': loss_a.item(),\n",
        "            'loss entropy': (loss_ent/(5 + self.train_progress*10)).item(),\n",
        "            'loss critic': loss_c.item(),\n",
        "            '-entropy by log(n)': loss_ent.item(),\n",
        "            'weight actor': logger.list_norm([p for p in self.actor.parameters() if p.requires_grad], 2),\n",
        "            'weight critic': logger.list_norm([p for p in self.critic.parameters() if p.requires_grad], 2),\n",
        "        })\n",
        "        wandb.log(self.log)\n",
        "        self.log = {}\n",
        "    self.actor_sched.step()\n",
        "\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def compute_GAE(self, critic_to_use=None):\n",
        "    \"\"\"\n",
        "    Computes GAE with passed/current critic.\n",
        "    GAEs could be used instead of Advantages in PPO actor loss\n",
        "    \"\"\"\n",
        "    if critic_to_use is None:\n",
        "      critic_to_use = copy.deepcopy(self.critic).eval()\n",
        "    features = self.trajectories.realized['features']\n",
        "    features_next = self.trajectories.realized['features_next']\n",
        "    rewards = self.trajectories.realized['rewards']\n",
        "    is_last = self.trajectories.realized['is_last']\n",
        "    Adv1 = []\n",
        "    ids_list = list(torch.split(torch.LongTensor(np.arange(len(is_last))), 4096))\n",
        "    for batch_ids in ids_list:\n",
        "      values = critic_to_use(features[batch_ids]).squeeze(-1)\n",
        "      values_next = critic_to_use(features_next[batch_ids]).squeeze(-1)\n",
        "      TD = values - ((rewards[batch_ids] + values_next * self.td_gamma) * (1-is_last[batch_ids]))\n",
        "      batch_Adv = - TD * (1-is_last[batch_ids])\n",
        "      Adv1 += [batch_Adv]\n",
        "    Adv1 = torch.cat(Adv1, dim=0)\n",
        "    Psi = torch.zeros_like(Adv1)\n",
        "\n",
        "    for i in range(len(Adv1)-1, -1, -1):\n",
        "      if is_last[i]:\n",
        "        continue\n",
        "      Psi[i] = Psi[i+1] * self.gae_gamma * self.td_gamma + Adv1[i]\n",
        "    self.trajectories.Psi = Psi\n",
        "    wandb.log({'Psi mean': Psi.mean().item(),\n",
        "               'Psi 95 perc': torch.quantile(Psi, 0.95).item(),\n",
        "               'Adv mean': Adv1.mean().item(),\n",
        "              })"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": true
        },
        "tags": [],
        "id": "y4PnKOg3vdQU"
      },
      "outputs": [],
      "source": [
        "class RNN_Trainer():\n",
        "  \"\"\"\n",
        "  Learns state history representation via a proxy task of V-function approximation.\n",
        "  \"\"\"\n",
        "  def __init__(self, v_cell, optimizer, trajectories, loss_horizon=10, rnn_batch_size=32):\n",
        "    self.rnn_cell = v_cell\n",
        "    self.optimizer = optimizer\n",
        "    self.loss_horizon = loss_horizon\n",
        "    self.retain_graph = True\n",
        "    self.trjs = trajectories\n",
        "    self.log = {}\n",
        "\n",
        "  def sample_ids_pairs(self, features_by_trjs, rnn_batch_size, n_steps):\n",
        "    \"\"\"\n",
        "    Returns [n_steps] list with [rnn_batch_size, 2] tensors\n",
        "    of pairs (traj idx, position in traj for rnn to start from)\n",
        "    \"\"\"\n",
        "    n_trjs = len(features_by_trjs)\n",
        "    residual_lengths = torch.tensor([len(f)-self.loss_horizon for f in features_by_trjs])\n",
        "    sampling_weights = np.log(residual_lengths) / np.log(residual_lengths).sum()\n",
        "    trjs_ids = torch.tensor(np.random.choice(np.arange(n_trjs), size=(n_steps, rnn_batch_size), p=sampling_weights))\n",
        "    position_samples = torch.rand(n_steps, rnn_batch_size)\n",
        "    residual_lengths_forall = torch.take(residual_lengths, trjs_ids)\n",
        "    start_ids = (residual_lengths_forall*position_samples).long()\n",
        "    stacked = torch.stack([trjs_ids, start_ids], dim=-1).to(device)\n",
        "    return list(stacked)\n",
        "\n",
        "  def train(self, n_steps=1000, rnn_batch_size=32, train_condition=eval(hp_alg.train_condition_str)):\n",
        "    is_last_ids = torch.nonzero(self.trjs.realized['is_last']).squeeze()\n",
        "    all_train_trjs_lengths = is_last_ids - torch.cat([torch.tensor([-1]).to(device), is_last_ids[:-1]])\n",
        "    not_long_enough_ids = (all_train_trjs_lengths < self.loss_horizon + 3).nonzero().squeeze()\n",
        "    donotlook_ids = torch.tensor([not train_condition(tr) for tr in self.trjs.j_file['paths']]).nonzero().squeeze()\n",
        "    features_by_trjs = list(torch.split(self.trjs.realized['features'], list(all_train_trjs_lengths)))\n",
        "    Returns_by_trjs = list(torch.split(self.trjs.realized['Returns'], list(all_train_trjs_lengths)))\n",
        "    for i in sorted(list(set(donotlook_ids.tolist()+not_long_enough_ids.tolist())), reverse=True):\n",
        "      # we train ony on long enough trajectories\n",
        "      popped = features_by_trjs.pop(i)\n",
        "      Returns_by_trjs.pop(i)\n",
        "    print(len(not_long_enough_ids), len(Returns_by_trjs))\n",
        "    sampled_ids = self.sample_ids_pairs(features_by_trjs, rnn_batch_size, n_steps)\n",
        "    for step in trange(n_steps):\n",
        "      batch_ids = sampled_ids.pop(0)\n",
        "      inputs = torch.zeros(rnn_batch_size, self.loss_horizon, features_dim).to(device)\n",
        "      targets = torch.zeros(rnn_batch_size, self.loss_horizon, 1).to(device)\n",
        "      for b in range(rnn_batch_size):\n",
        "        inputs[b] = features_by_trjs[batch_ids[b,0]][batch_ids[b,1]: batch_ids[b,1] + self.loss_horizon]\n",
        "        targets[b] = Returns_by_trjs[batch_ids[b,0]][batch_ids[b,1]: batch_ids[b,1] + self.loss_horizon][:, None]\n",
        "      states = [torch.zeros(2*self.rnn_cell.num_layers, rnn_batch_size, self.rnn_cell.inner_hidden_size).to(device)]\n",
        "      Vs = []\n",
        "      for rnn_step in range(self.loss_horizon):\n",
        "        v, new_state, h2 = self.rnn_cell(inputs[:, rnn_step, :], states[-1])\n",
        "        states += [new_state]\n",
        "        Vs += [v]\n",
        "      abs_difs = (torch.stack(Vs, dim=1) - targets).abs()\n",
        "      assert abs_difs.shape == (rnn_batch_size, self.loss_horizon, 1)\n",
        "      loss_weights = torch.arange(self.loss_horizon)**1.4 / (torch.arange(self.loss_horizon)**1.4).sum()\n",
        "      loss = (abs_difs * loss_weights[:, None].to(device)).mean()\n",
        "      self.optimizer.zero_grad()\n",
        "      loss.backward()\n",
        "      self.optimizer.step()\n",
        "      if step % logger.between_logs == 0:\n",
        "        self.log['rnn MC abs loss'] = loss.item()\n",
        "        self.log['rnn V mean'] = torch.cat(Vs).mean().item()\n",
        "        self.log['RNN weight'] = logger.list_norm([p for p in self.rnn_cell.parameters() if p.requires_grad], 2)\n",
        "        self.log['RNN grad'] = logger.list_norm([p.grad for p in self.rnn_cell.parameters() if p.requires_grad], 2)\n",
        "        hist = wandb.Histogram(np_histogram=np.histogram(torch.cat(Vs).flatten().detach().to('cpu'), bins=40, ))\n",
        "        self.log['V-func hist'] = hist\n",
        "        hist = wandb.Histogram(np_histogram=np.histogram(abs_difs[:, -1,:].flatten().detach().to('cpu'), bins=40, ))\n",
        "        self.log['abs difs -1'] = hist\n",
        "        hist = wandb.Histogram(np_histogram=np.histogram(abs_difs[:, 0,:].flatten().detach().to('cpu'), bins=40, ))\n",
        "        self.log['abs difs 0'] = hist\n",
        "        wandb.log(self.log)\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def collect_rnn_features(self):\n",
        "    \"\"\"\n",
        "    Happened to be not relevant in the end.\n",
        "    Collects rnn hidden and V for every state in train dataset\n",
        "    \"\"\"\n",
        "    features = self.trjs.realized['features']\n",
        "    is_last = self.trjs.realized['is_last']\n",
        "    hid_list = []\n",
        "    v_list = []\n",
        "    init_state = torch.zeros(2*self.rnn_cell.num_layers, 1, self.rnn_cell.inner_hidden_size).to(device)\n",
        "    state = init_state\n",
        "    for i, f in enumerate(features):\n",
        "      v, state = self.rnn_cell(f[None, :], state)\n",
        "      hid = state[2].squeeze(0)\n",
        "      hid_list += [hid.detach()]\n",
        "      v_list += [v.detach()]\n",
        "      if is_last[i]:\n",
        "        state = init_state\n",
        "    self.traj.realized['rnn_features'] = torch.stack(hid_list, dim=0)\n",
        "    self.traj.realized['rnn_v'] = torch.stack(v_list, dim=0)\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def concat_rnn_features(self):\n",
        "    \"\"\"\n",
        "    Happened to be not relevant in the end.\n",
        "    Concatenates state history embedding to embedding of that state and to embeddings of corresponding actions\n",
        "    \"\"\"\n",
        "    realized = self.trjs.realized\n",
        "    rnn_related = torch.cat([realized['rnn_features'], realized['rnn_v']], dim=-1)\n",
        "    realized['features'] = torch.cat([realized['features'], rnn_related], dim=-1)\n",
        "    queues = self.trjs.queues\n",
        "    time_before = time()\n",
        "    for i in range(len(queues)):\n",
        "      queues[i] = torch.cat([queues[i], rnn_related[i].repeat(len(queues[i]), 1)], dim=-1)\n",
        "    print('Concating rnn features to queues: ', time()-time_before)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UX4jPlzD9gFd",
        "tags": []
      },
      "source": [
        "### Procedures"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "tags": [],
        "id": "dxjWH1kcvdQV"
      },
      "outputs": [],
      "source": [
        "if hp_alg.critic_play:\n",
        "  run = wandb.init(\n",
        "          project=\"PS\",\n",
        "          name=f'PI',\n",
        "          config={}\n",
        "  )\n",
        "  trajectories = Trajectories()\n",
        "\n",
        "  actor, actor_opt, actor_sched = get_attn_setup() # not to be used\n",
        "  critic, critic_opt = get_mlp_setup()\n",
        "  q_net = Q_Net(V_function=critic)\n",
        "\n",
        "  trajectories.gather_n_store() # default heuristic play (ForkDeothRandom) to initialize\n",
        "  trajectories.evaluate_val_train()\n",
        "  print(trajectories.get_properties())\n",
        "\n",
        "  for epoch in range(hp_alg.epochs):\n",
        "      trainer = NN_Trainer(NN_setup={'actor': actor, 'actor_opt': actor_opt, 'actor_sched': actor_sched,\n",
        "                                     'critic': critic, 'critic_opt': critic_opt,},\n",
        "                           trajectories=trajectories,\n",
        "                           train_progress=epoch/hp_alg.epochs,\n",
        "                           )\n",
        "      trainer.learn_v(n_batches=2000 if epoch==0 else 400)\n",
        "      wandb.log({'epoch': epoch})\n",
        "      trajectories.gather_n_store(actor_model=q_net)\n",
        "      trajectories.evaluate_val_train()\n",
        "      print(trajectories.get_properties())\n",
        "\n",
        "  checkpoint = {\n",
        "    'critic':critic,\n",
        "  }\n",
        "  torch.save(checkpoint, os.path.join(wandb.run.dir, f'critic model'))\n",
        "  wandb.finish()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "Bi2uaButgHOM",
        "jupyter": {
          "outputs_hidden": true
        },
        "scrolled": true,
        "tags": []
      },
      "outputs": [],
      "source": [
        "if not hp_alg.critic_play:\n",
        "  actor, actor_opt, actor_sched = get_attn_setup()\n",
        "  critic, critic_opt = get_mlp_setup()\n",
        "  trajectories = Trajectories()\n",
        "  run = wandb.init(\n",
        "        project='PS',\n",
        "        name=f'PPO',\n",
        "        config={}\n",
        "  )\n",
        "  try:\n",
        "    checkpoint_path = '../Checkpoints/actor_critic27_08'\n",
        "    init_actor = torch.load(checkpoint_path)['actor']\n",
        "  except:\n",
        "    print(no checkpoint model, using random init)\n",
        "  else:\n",
        "    init_actor = actor\n",
        "  trajectories.gather_n_store(actor_model=init_actor) # init data gathering has to be performed by some policy\n",
        "  trajectories.evaluate_val_train()\n",
        "  print(trajectories.get_properties())\n",
        "\n",
        "  for epoch in range(hp_alg.epochs):\n",
        "      trainer = NN_Trainer(NN_setup={'actor': actor, 'actor_opt': actor_opt, 'actor_sched': actor_sched,\n",
        "                                     'critic': critic, 'critic_opt': critic_opt,},\n",
        "                           clip_eps = get_clip_eps(epoch),\n",
        "                           trajectories=trajectories,\n",
        "                           train_progress=epoch/hp_alg.epochs,\n",
        "                           n_batches=1000 if epoch==0 else 300,\n",
        "                           )\n",
        "      trainer.learn_v(2000 if epoch==0 else 50)\n",
        "      trainer.learn_new_policy()\n",
        "      wandb.log({'epoch': epoch})\n",
        "      trajectories.gather_n_store(actor_model=actor)\n",
        "      trajectories.evaluate_val_train()\n",
        "      print(trajectories.get_properties())\n",
        "  checkpoint = {\n",
        "    'actor': actor,\n",
        "  }\n",
        "  torch.save(checkpoint, os.path.join(wandb.run.dir, f'actor model'))\n",
        "  wandb.finish()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "tags": [],
        "id": "pAC8g7k-vdQW"
      },
      "outputs": [],
      "source": [
        "# exit()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "jp-MarkdownHeadingCollapsed": true,
        "tags": [],
        "id": "27fUUQetvdQW"
      },
      "source": [
        "### Data filter"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yJDIOyF3vdQW"
      },
      "outputs": [],
      "source": [
        "def data_filter(clear_blacklist=False, verbose=True):\n",
        "  '''\n",
        "  To raise training efficiency one might filter out useless trajectories.\n",
        "  '''\n",
        "  trajectories = Trajectories()\n",
        "  trajectories.gather_n_store()\n",
        "  print(trajectories.get_properties())\n",
        "\n",
        "  if clear_blacklist:\n",
        "    open('../Game_env/blacklist.txt', 'w').close()\n",
        "\n",
        "  paths = trajectories.j_file['paths']\n",
        "\n",
        "  lengths = [len(paths[i][1]) for i in range(len(paths))]\n",
        "  lengths_t = torch.Tensor(lengths)\n",
        "\n",
        "  mean_queue_len = [torch.Tensor([len(q[0]) for q in paths[i][1]]).mean() for i in range(len(paths))]\n",
        "  mean_queue_len_t = torch.Tensor(mean_queue_len)\n",
        "\n",
        "  ids = torch.nonzero((mean_queue_len_t<=1.1)).long().squeeze()\n",
        "  if verbose:\n",
        "    print('lengths: ', [lengths_t.quantile(10*i/100) for i in range(10)])\n",
        "    print('mean queue len: ', [mean_queue_len_t.quantile(10*i/100) for i in range(11)])\n",
        "    print('sum length where mean_queue_len_t<=1.1: ', lengths_t[torch.nonzero(mean_queue_len_t<=1.1).squeeze()].sum())\n",
        "    print('traj with short queues: ', len(ids))\n",
        "  bad_paths = [paths[i] for i in ids]\n",
        "\n",
        "  with open('../Game_env/blacklist.txt', 'a') as f:\n",
        "    for i in range(len(ids)):\n",
        "      s = bad_paths[i][2]\n",
        "      f.write(s + '\\n')\n",
        "\n",
        "  log = trajectories.evaluate_data(eval_condition = (lambda tr: True),\n",
        "                                       verbose=False,\n",
        "                                       factors=torch.Tensor([1]),\n",
        "                                      )\n",
        "  Returns = torch.Tensor(log['Returns'])\n",
        "  ids = torch.nonzero((Returns==0).long()).long().squeeze()\n",
        "  if verbose:\n",
        "    print('len Returns, len lengths: ', len(Returns), len(lengths))\n",
        "    print('sum length where Returns==0: ', lengths_t[torch.nonzero(Returns==0).squeeze()].sum())\n",
        "    print('traj with zero Return: ', len(ids))\n",
        "  bad_paths = [paths[i] for i in ids]\n",
        "\n",
        "  with open('../Game_env/blacklist.txt', 'a') as f:\n",
        "    for i in range(len(ids)):\n",
        "      s = bad_paths[i][2]\n",
        "      f.write(s + '\\n')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "jp-MarkdownHeadingCollapsed": true,
        "tags": [],
        "id": "nclEepwevdQX"
      },
      "source": [
        "### The end"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "5IWe2Bd6vdQI",
        "t6GWfxrs9gFR",
        "VJbCP9HE9gFV",
        "oRbyOLM1vdQP",
        "ndNrbEp79gFX",
        "Lfn5Oj8qvdQR",
        "qBhLnmrnFDAu",
        "tRSEdgzW9gFc",
        "UX4jPlzD9gFd",
        "27fUUQetvdQW"
      ],
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python [conda env:.conda-envatt] *",
      "language": "python",
      "name": "conda-env-.conda-envatt-py"
    },
    "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.0"
    },
    "varInspector": {
      "cols": {
        "lenName": 16,
        "lenType": 16,
        "lenVar": 40
      },
      "kernels_config": {
        "python": {
          "delete_cmd_postfix": "",
          "delete_cmd_prefix": "del ",
          "library": "var_list.py",
          "varRefreshCmd": "print(var_dic_list())"
        },
        "r": {
          "delete_cmd_postfix": ") ",
          "delete_cmd_prefix": "rm(",
          "library": "var_list.r",
          "varRefreshCmd": "cat(var_dic_list()) "
        }
      },
      "types_to_exclude": [
        "module",
        "function",
        "builtin_function_or_method",
        "instance",
        "_Feature"
      ],
      "window_display": false
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}