demo_design_optimization.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SyNZcVgQPsKs"
      },
      "source": [
        "Copyright 2022 DeepMind Technologies Limited\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "     https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I8K_e58H8S9d"
      },
      "source": [
        "# Demo design optimization for 3D WaterCourse environment\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1jNaCf3sM1Xm"
      },
      "outputs": [],
      "source": [
        "#@title Installation (if not running locally)\n",
        "# Note, this should be skipped if running locally.\n",
        "!mkdir /content/inverse_design\n",
        "!mkdir /content/inverse_design/src\n",
        "!touch /content/inverse_design/__init__.py\n",
        "!touch /content/inverse_design/src/__init__.py\n",
        "\n",
        "!wget -O /content/inverse_design/src/connectivity_utils.py https://raw.githubusercontent.com/deepmind/master/inverse_design/src/connectivity_utils.py\n",
        "!wget -O /content/inverse_design/src/graph_network.py https://raw.githubusercontent.com/deepmind/master/inverse_design/src/graph_network.py\n",
        "!wget -O /content/inverse_design/src/learned_simulator.py https://raw.githubusercontent.com/deepmind/master/inverse_design/src/learned_simulator.py\n",
        "!wget -O /content/inverse_design/src/model_utils.py https://raw.githubusercontent.com/deepmind/master/inverse_design/src/model_utils.py\n",
        "!wget -O /content/inverse_design/src/normalizers.py https://raw.githubusercontent.com/deepmind/master/inverse_design/src/normalizers.py\n",
        "!wget -O /content/inverse_design/src/watercourse_env.py https://raw.githubusercontent.com/deepmind/master/inverse_design/src/watercourse_env.py\n",
        "\n",
        "!wget -O /content/requirements.txt https://raw.githubusercontent.com/deepmind/master/inverse_design/requirements.txt\n",
        "!pip install -r requirements.txt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q08NCU1pds4j"
      },
      "outputs": [],
      "source": [
        "#@title Download Pickled Dataset \u0026 Params (if running in colab)\n",
        "# Note this can be skipped if following instructions for jupyter notebook\n",
        "from google.colab import auth\n",
        "auth.authenticate_user()\n",
        "\n",
        "!gsutil cp gs://dm_inverse_design_watercourse/init_sequence.pickle .\n",
        "!gsutil cp gs://dm_inverse_design_watercourse/gns_params.pickle .\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3gfBGlIrCkBf"
      },
      "outputs": [],
      "source": [
        "#@title Imports\n",
        "from inverse_design.src import learned_simulator\n",
        "from inverse_design.src import model_utils\n",
        "from inverse_design.src import watercourse_env"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SJVrIfyHnkOK"
      },
      "outputs": [],
      "source": [
        "#@title Open pickled parameters + dataset\n",
        "import pickle\n",
        "\n",
        "with open('init_sequence.pickle', \"rb\") as f:\n",
        "  pickled_data = pickle.loads(f.read())\n",
        "  gt_sequence = pickled_data['gt_sequence']\n",
        "  meta = pickled_data['meta']\n",
        "\n",
        "with open('gns_params.pickle', \"rb\") as f:\n",
        "  pickled_params = pickle.loads(f.read())\n",
        "  network = pickled_params['network']\n",
        "  plan = pickled_params['plan']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kdxb0hSLPRwH"
      },
      "outputs": [],
      "source": [
        "#@title make ID control/loss functions\n",
        "import jax\n",
        "import functools\n",
        "\n",
        "# maximum number of edges for single step of rollout to pad to\n",
        "MAX_EDGES = 2**16\n",
        "\n",
        "# define haiku model\n",
        "connectivity_radius = meta[\"connectivity_radius\"]\n",
        "flatten_fn = functools.partial(model_utils.flatten_features, **plan['flatten_kwargs'])\n",
        "haiku_model = functools.partial(learned_simulator.LearnedSimulator, connectivity_radius=connectivity_radius, flatten_features_fn=flatten_fn, **plan['model_kwargs'])\n",
        "\n",
        "# create initial landscape (obstacle) in the scene\n",
        "obstacle_pos = watercourse_env.make_plain_obstacles()\n",
        "for frame in gt_sequence:\n",
        "  pos = frame.nodes['world_position'].copy()\n",
        "  pos[:obstacle_pos.shape[0]] = obstacle_pos[:, None]\n",
        "  frame.nodes['world_position'] = pos\n",
        "\n",
        "\n",
        "# get initial sequence of particles from dataset for initial graph\n",
        "obstacle_edges, inflow_stack, initial_graph = watercourse_env.build_initial_graph(gt_sequence[15:], max_edges=MAX_EDGES)\n",
        "\n",
        "# infer the landscape size from the dataset (25 x 25)\n",
        "# note that this is not required, it is also possible to create a smaller\n",
        "# or larger landscape (obstacle) as the design space\n",
        "num_side = int(jax.numpy.sqrt(initial_graph.nodes['obstacle_mask'].sum()))\n",
        "n_obs = num_side**2\n",
        "\n",
        "# rollout length definition (final state taken for reward computation)\n",
        "length = 50\n",
        "# radius within which to connect particles\n",
        "radius = 0.1\n",
        "# smoothing factor for loss\n",
        "smoothing_factor = 1e2\n",
        "\n",
        "@jax.jit\n",
        "def run(vars):\n",
        "  # create landscape as graph from vars parameters\n",
        "  graph, raw_obs = watercourse_env.design_fn(vars, initial_graph)\n",
        "\n",
        "  # rollout\n",
        "  final_graph, traj = watercourse_env.rollout(\n",
        "      graph, inflow_stack[:length], network, haiku_model,\n",
        "      obstacle_edges, radius=radius)\n",
        "  \n",
        "  # losses\n",
        "  losses = {\n",
        "      'objective': watercourse_env.max_x_loss_fn(final_graph),\n",
        "      'smooth': smoothing_factor * watercourse_env.smooth_loss_fn(raw_obs, num_side=num_side),\n",
        "  }\n",
        "\n",
        "  # auxiliaries to keep track of for plotting\n",
        "  aux = {\n",
        "      'design': vars,\n",
        "      'losses': losses,\n",
        "      'traj': traj\n",
        "  }\n",
        "  return sum(losses.values()), aux\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TuuHYP3dOW45"
      },
      "outputs": [],
      "source": [
        "from IPython.display import clear_output\n",
        "import jax.numpy as jnp\n",
        "import matplotlib.pyplot as plt\n",
        "import optax\n",
        "\n",
        "# set learning rate and number of optimization steps\n",
        "LEARNING_RATE = 0.05\n",
        "num_opt_steps = 100\n",
        "\n",
        "# define optimizer as adam with learning rate\n",
        "optimizer = optax.adam(learning_rate=LEARNING_RATE)\n",
        "\n",
        "# initialize design parameters to be zeros (flat landscape)\n",
        "params = jnp.zeros(n_obs, dtype=jnp.float32)\n",
        "opt_state = optimizer.init(params)\n",
        "\n",
        "# initialize empty optimization trajectory (for tracking improvements to losses and design)\n",
        "opt_traj = []\n",
        "\n",
        "# optimization step with current design parameters\n",
        "@jax.jit\n",
        "def opt_step(params, opt_state):\n",
        "  grads, aux = jax.grad(run, has_aux=True)(params)\n",
        "  updates, opt_state = optimizer.update(grads, opt_state, params)\n",
        "  params = optax.apply_updates(params, updates)\n",
        "  return params, opt_state, aux\n",
        "\n",
        "# run optimization loop and track progress\n",
        "for i in range(num_opt_steps):\n",
        "  params, opt_state, aux = opt_step(params, opt_state)\n",
        "  opt_traj.append(aux)\n",
        "  clear_output(wait=True)\n",
        "  fig, ax = plt.subplots(1,1,figsize=(10,5))\n",
        "  for key in aux['losses'].keys():\n",
        "    ax.plot([t['losses'][key] for t in opt_traj])\n",
        "  ax.plot([sum(t['losses'].values()) for t in opt_traj])\n",
        "  ax.legend(list(aux['losses'].keys())+['total'])\n",
        "  plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "88voGdYuAPus"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "# plot design iterations (every 10 steps)\n",
        "n_sam = range(0, len(opt_traj), 10)\n",
        "fig, ax = plt.subplots(1,len(n_sam),figsize=(len(n_sam)*10, 10), squeeze=False)\n",
        "\n",
        "for fi, idx in enumerate(n_sam):\n",
        "  design = opt_traj[idx]['design']\n",
        "\n",
        "  # control function uses tanh as transformation, so mimic here to see heightfield\n",
        "  fld = np.tanh(design.reshape((num_side, num_side)))\n",
        "  ax[0, fi].imshow(fld, vmin=-1, vmax=1)\n",
        "  ax[0, fi].set_axis_off()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yPaWTUM0FL36"
      },
      "outputs": [],
      "source": [
        "from IPython.display import clear_output\n",
        "# plot video of how particles move for optimized design and initial design\n",
        "\n",
        "def _plt(ax, frame, i):\n",
        "  pos = frame['pos'][i]  \n",
        "  p = pos[frame['mask'][i]]\n",
        "  ax.scatter(p[:, 0], p[:, 2], p[:, 1], c='b',s=10)\n",
        "  obs = pos[:num_side**2]\n",
        "  ax.scatter(obs[:, 0], obs[:, 2], obs[:, 1], c='k',s=3)\n",
        "  ax.scatter([1.5],[1.5],[0], c='g',s=20)\n",
        "  ax.set_xlim([-0.6, 1.6])\n",
        "  ax.set_ylim([-0.1, 1.6])\n",
        "  ax.set_zlim([-0.1, 1.2])\n",
        "\n",
        "roll_fin0 = run(opt_traj[0]['design'])[1]['traj']\n",
        "roll_fin1 = run(opt_traj[-1]['design'])[1]['traj']\n",
        "\n",
        "for i in range(roll_fin0['pos'].shape[0]):\n",
        "  clear_output(wait=True)\n",
        "  fig = plt.figure(figsize=(20,10))\n",
        "  ax1 = fig.add_subplot(1, 2, 1, projection='3d')\n",
        "  ax1.set_title('Initial design, frame %d' % i)\n",
        "  _plt(ax1, roll_fin0, i)\n",
        "  ax2 = fig.add_subplot(1, 2, 2, projection='3d')\n",
        "  ax2.set_title('Design at final step')\n",
        "  _plt(ax2, roll_fin1, i)\n",
        "  plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "csTLqGTUPuAn"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "1ZYL6nDmJCvzc70qi5rwIQ7BSj9a2sdoL",
          "timestamp": 1667473449566
        },
        {
          "file_id": "1rOmaBHyQAVa6NY3nfsM1yba8GC0a8duM",
          "timestamp": 1667402678487
        },
        {
          "file_id": "1bHb2szUEMLP2gKErlnxrigutZyhZ4reh",
          "timestamp": 1667386305836
        },
        {
          "file_id": "1jrPkT2OlRGUpImN2cVetBtQ9PgFelNgy",
          "timestamp": 1665137721073
        },
        {
          "file_id": "144LaPYCJpaCSUX6jMoampo0dTMICcXWV",
          "timestamp": 1640626024159
        },
        {
          "file_id": "1B31Dl0yzZlW22l-X7qSR5S4w1nOHPAg0",
          "timestamp": 1640104963946
        },
        {
          "file_id": "1MXdt8a59t0hLm7Q1Prem3kdHpnixMOeJ",
          "timestamp": 1639481266768
        },
        {
          "file_id": "14OG1REDMq3i9phUABzuSWg3nQhUD_ozK",
          "timestamp": 1639156624846
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}