cerberus-viz.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d635480c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pickle\n",
    "import matplotlib.pyplot as plt\n",
    "from copy import deepcopy\n",
    "from ipywidgets import interact\n",
    "from tqdm.notebook import tqdm\n",
    "from numba import njit\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import requests\n",
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b178b0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "891965ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "from kaggle_environments import make\n",
    "from kaggle_environments import structify"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a8059cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import lux_ai.rl_agent.rl_agent as rl_agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30016cc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# replay_id, player_id = 34014720, 0\n",
    "replay_id, player_id = 33947998, 1\n",
    "# replay_id, player_id = 34222068, 0\n",
    "\n",
    "replay = requests.get(f\"https://www.kaggleusercontent.com/episodes/{replay_id}.json\")\n",
    "replay = json.loads(replay.text)\n",
    "steps = replay[\"steps\"]\n",
    "config = replay[\"configuration\"]\n",
    "len(steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0eecbd5a",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "env = make(\"lux_ai_2021\", configuration=config, info=replay[\"info\"], steps=steps)\n",
    "env.render(mode=\"ipython\", width=1000, height=800)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48f27053",
   "metadata": {},
   "outputs": [],
   "source": [
    "from process_cerberus_replays import get_delta_with_cache, extract_obs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "618fcca8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_worker_build_prb(obs, out):\n",
    "    size = obs[\"width\"]\n",
    "    prb = out[\"policy_logits\"][\"city_tile\"]\n",
    "    prb = prb.squeeze(0).squeeze(0)[player_id]\n",
    "    prb = F.softmax(prb, dim=-1)\n",
    "    prb = np.array(prb)\n",
    "    return prb[:size, :size, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f62dd188",
   "metadata": {},
   "outputs": [],
   "source": [
    "states = list()\n",
    "values = list()\n",
    "deltas = list()\n",
    "outputs = list()\n",
    "worker_build_prbs = list()\n",
    "\n",
    "obs = extract_obs(steps[0], player_id)\n",
    "agent = rl_agent.RLAgent(obs, config)\n",
    "for step in tqdm(steps):\n",
    "    states.append(deepcopy(agent.game_state))\n",
    "    obs = extract_obs(step, player_id)\n",
    "    deltas.append(get_delta_with_cache(replay_id, agent, obs, config, skip_uncached=True))\n",
    "    out = agent(obs, config, True)\n",
    "    outputs.append(out)\n",
    "    values.append(float(out['baseline'][0][player_id]))\n",
    "    worker_build_prbs.append(get_worker_build_prb(obs, out))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e58cbce",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cerberus_viz import (\n",
    "    make_figure,\n",
    "    add_traces,\n",
    "    plot_array,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a8e41cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def heatmap_function(state):\n",
    "    if state.id:\n",
    "        return -deltas[state.turn]\n",
    "    return deltas[state.turn]\n",
    "\n",
    "def timeseries_function(state):\n",
    "    return {\"expected_value\": values[:state.turn + 1]}\n",
    "\n",
    "def cityhighlight_function(state):\n",
    "    return worker_build_prbs[state.turn] / 2 + 0.5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33151957",
   "metadata": {},
   "source": [
    "Widget that lets you step through the match:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0b5e42b",
   "metadata": {},
   "outputs": [],
   "source": [
    "@interact(view_step=(0, len(steps) - 1))\n",
    "def interactive_display(view_step=0):\n",
    "    state = states[view_step]\n",
    "    fig = make_figure(state, replay_id, player_id)\n",
    "    add_traces(fig, state, heatmap_function, timeseries_function, cityhighlight_function)\n",
    "    fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e88e516",
   "metadata": {},
   "source": [
    "Save all frames as PNGs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe42b3a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "for view_step in tqdm(range(len(steps) - 1)):\n",
    "    state = states[view_step]\n",
    "    fig = make_figure(state, replay_id, player_id)\n",
    "    add_traces(fig, state, heatmap_function, timeseries_function, cityhighlight_function)\n",
    "    fig.write_image(f\"cerberus_replays/pngs/{replay_id}-{view_step:03}.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38cb85af",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}