demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reveal Cancer Invasion Strategy with Reinforcement Learning\n",
    "We demonstrate how cancer cells learn to reproduce and invade using reinforcement learning. The idea originated from a proposal to Cell Fate Symposium 2018 (see proposal [here](https://www.dropbox.com/s/er54z58cbda2bkn/cell%20fate%20full%20proposal.pdf?dl=0)).\n",
    "\n",
    "At the beginning of each episode, one cancer cell at the center of a fixed square domain with even-distributed nutrients is initialized. We imagine that an agent is playing a game in which it takes cues from the microenvironment of a cancer cell and instructs the cell whether to reproduce or migrate. The cost of nutrient to reproduce is higher than moving around. The nutrient is limited and diffuse within the domain with a no-flux boundary condition. When the nutrient is below certain threshold the cells start dying. The goal is to maximize the number of offsprings that successfully exit the domain. For this goal, we design rewarding as follows:\n",
    "* reproduction results in a positive reward and\n",
    "* death leads to a negative reward of the same amount.\n",
    "* extra reward is given when a cancer cell is successfully exit the domain.\n",
    "\n",
    "Here we employed Policy Gradient algorithm to learn possible strategy of cancer invasion.  Policy function takes input of ambient nutrient level and stochasticaly maps to the action that can be taken by the cancer cell: moving to one of the nearby sites or reproduce. Here policy function is approximated as a neutral net with one hidden layer. Run the \"cell\" below to see how the cancer cells improve its reward with training with more and more episode. The bottom graph is the running average reward over the episodes. The two grid plots show details of several typical episodes. Note salient behavior changes after 100 episode the cancer cells compare to the 1st episode.\n",
    "\n",
    "This is a simple sketch of the idea. In the future, we hope to reveal interesting morphological patterns of cancer invasion by adding more hidden layers and more biologically meaningfully features to the current model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib nbagg\n",
    "\"\"\" Trains an agent with (stochastic) Policy Gradients on Cancer Invasion. \"\"\"\n",
    "import numpy as np\n",
    "import pickle\n",
    "from cancer_env import CancerEnv\n",
    "\n",
    "# hyperparameters\n",
    "H = 400 # number of hidden layer neurons\n",
    "batch_size = 10 # every how many episodes to do a param update?\n",
    "learning_rate = 1e-4\n",
    "gamma = 0.99 # discount factor for reward\n",
    "decay_rate = 0.99 # decay factor for RMSProp leaky sum of grad^2\n",
    "resume = False # resume from previous checkpoint?\n",
    "render = False\n",
    "\n",
    "# model initialization\n",
    "D = 3 * 3 # input dimensionality: 3x3 grid nut\n",
    "if resume:\n",
    "  model = pickle.load(open('save.p', 'rb'))\n",
    "else:\n",
    "  model = {}\n",
    "  model['W1'] = 0.1*np.random.randn(H,D) / np.sqrt(D) # \"Xavier\" initialization\n",
    "  model['W2'] = 0.1*np.random.randn(H,9) / np.sqrt(H) # 9 actions\n",
    "  \n",
    "grad_buffer = { k : np.zeros_like(v) for k,v in model.items() } # update buffers that add up gradients over a batch\n",
    "rmsprop_cache = { k : np.zeros_like(v) for k,v in model.items() } # rmsprop memory\n",
    "\n",
    "\n",
    "def discount_rewards(r):\n",
    "  \"\"\" take 1D float array of rewards and compute discounted reward \"\"\"\n",
    "  discounted_r = np.zeros_like(r)\n",
    "  running_add = 0\n",
    "  for t in reversed(range(0, r.size)):\n",
    "    if r[t] != 0: running_add = 0 # reset the sum, since this was a game boundary (pong specific!)\n",
    "    running_add = running_add * gamma + r[t]\n",
    "    discounted_r[t] = running_add\n",
    "  return discounted_r\n",
    "\n",
    "def policy_forward(x):\n",
    "  h = model['W1'] @ x\n",
    "  h[h<0] = 0 # ReLU nonlinearity\n",
    "  logp = h @ model['W2']\n",
    "  p = np.exp(logp) / np.exp(logp).sum()\n",
    "  return p, h # return probability of taking action 2, and hidden state\n",
    "\n",
    "def policy_backward(eph, epdlogp):\n",
    "  \"\"\" backward pass. (eph is array of intermediate hidden states) \"\"\"\n",
    "  dW2 = eph.T @ epdlogp\n",
    "  #dh = np.outer(epdlogp, model['W2'])\n",
    "  dh = epdlogp @ model['W2'].T\n",
    "  dh[eph <= 0] = 0 # backpro prelu\n",
    "  dW1 = np.dot(dh.T, epx)\n",
    "  return {'W1':dW1, 'W2':dW2}\n",
    "\n",
    "domain_width = 20\n",
    "env = CancerEnv(domain_width - 1)\n",
    "observation = env.reset()  # 3 by 3 neibnut\n",
    "xs,hs,dlogps,drs = [],[],[],[]\n",
    "running_reward = None\n",
    "reward_sum = 0\n",
    "episode_number = 0\n",
    "istep = 0\n",
    "max_episode_number = 1500\n",
    "while episode_number <= max_episode_number:\n",
    "  istep += 1\n",
    "  if render and istep % 2 == 0:\n",
    "    env.render()\n",
    "    istep = 0\n",
    "  # preprocess the observation, set input to network to be difference image\n",
    "  x = observation.ravel()\n",
    "\n",
    "  # forward the policy network and sample an action from the returned probability\n",
    "  aprob, h = policy_forward(x)\n",
    "  action = np.random.choice(range(9), p=aprob)  # roll the dice!\n",
    "\n",
    "  # record various intermediates (needed later for backprop)\n",
    "  xs.append(x) # observation\n",
    "  hs.append(h) # hidden state\n",
    "\n",
    "  y = np.zeros(9)\n",
    "  y[action] = 1\n",
    "  dlogps.append(y - aprob) # !grad that encourages the action that was taken to be taken (see http://cs231n.github.io/neural-networks-2/#losses if confused)\n",
    "\n",
    "  # step the environment and get new measurements\n",
    "  observation, reward, done = env.pg_step(action)\n",
    "  reward_sum += reward\n",
    "\n",
    "  drs.append(reward) # record reward (has to be done after we call step() to get reward for previous action)\n",
    "\n",
    "  if done: # an episode finished\n",
    "    #print(('ep %d: game finished, total reward: %f' % (episode_number, reward_sum)) + ('' if reward_sum < 11 else '!!!!'))\n",
    "    episode_number += 1\n",
    "    if episode_number in [1, 50, 100, 1000, 1500]:\n",
    "      render = True\n",
    "    else:\n",
    "      render = False\n",
    "    # stack together all inputs, hidden states, action gradients, and rewards for this episode\n",
    "    epx = np.vstack(xs)\n",
    "    eph = np.vstack(hs)\n",
    "    epdlogp = np.vstack(dlogps)\n",
    "    epr = np.vstack(drs)\n",
    "    xs,hs,dlogps,drs = [],[],[],[] # reset array memory\n",
    "\n",
    "    # compute the discounted reward backwards through time\n",
    "    discounted_epr = discount_rewards(epr)  # int????\n",
    "    # standardize the rewards to be unit normal (helps control the gradient estimator variance)\n",
    "    discounted_epr = discounted_epr - np.mean(discounted_epr)\n",
    "    discounted_epr /= np.std(discounted_epr)\n",
    "\n",
    "    epdlogp *= discounted_epr # !modulate the gradient with advantage (PG magic happens right here.)\n",
    "    grad = policy_backward(eph, epdlogp)\n",
    "    for k in model: grad_buffer[k] += grad[k] # accumulate grad over batch\n",
    "\n",
    "    # perform rmsprop parameter update every batch_size episodes\n",
    "    if episode_number % batch_size == 0:\n",
    "      for k,v in model.items():\n",
    "        g = grad_buffer[k] # gradient\n",
    "        rmsprop_cache[k] = decay_rate * rmsprop_cache[k] + (1 - decay_rate) * g**2\n",
    "        model[k] += learning_rate * g / (np.sqrt(rmsprop_cache[k]) + 1e-5)\n",
    "        grad_buffer[k] = np.zeros_like(v) # reset batch gradient buffer\n",
    "\n",
    "    # boring book-keeping\n",
    "    running_reward = reward_sum if running_reward is None else running_reward * 0.99 + reward_sum * 0.01\n",
    "    print('ep %d : resetting env. episode reward total was %f. running mean: %f' % (episode_number, reward_sum, running_reward))\n",
    "    env.plot_extra(episode_number, running_reward)\n",
    "    if episode_number % 100 == 0: pickle.dump(model, open('save.p', 'wb'))\n",
    "    reward_sum = 0\n",
    "    observation = env.reset() # reset env"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}