training/auto_notebooks/rl/IntervalDiscrimination-v0.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/neurogym/ngym_usage/blob/master/supervised/auto_notebooks/rl/IntervalDiscrimination-v0.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install packages if on Colab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment following lines to install\n",
    "# %tensorflow_version 1.x\n",
    "# ! pip install --upgrade stable-baselines  # install latest version\n",
    "# ! pip install gym   # Install gym\n",
    "# ! git clone https://github.com/gyyang/neurogym.git  # Install neurogym\n",
    "# %cd neurogym/\n",
    "# ! pip install -e .\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!/usr/bin/env python3\n",
    "# -*- coding: utf-8 -*-\n",
    "\"\"\"\n",
    "Created on Wed Sep 23 09:33:08 2020\n",
    "\n",
    "@author: manuel\n",
    "\"\"\"\n",
    "\n",
    "import os\n",
    "from pathlib import Path\n",
    "import json\n",
    "import importlib\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.decomposition import PCA\n",
    "import matplotlib.pyplot as plt\n",
    "from neurogym.wrappers import ALL_WRAPPERS\n",
    "from stable_baselines.common.vec_env import SubprocVecEnv, DummyVecEnv\n",
    "from stable_baselines.common import set_global_seeds\n",
    "from stable_baselines.common.policies import LstmPolicy\n",
    "from stable_baselines.common.callbacks import CheckpointCallback\n",
    "import gym\n",
    "import glob\n",
    "import neurogym as ngym\n",
    "\n",
    "\n",
    "envid = 'IntervalDiscrimination-v0'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_modelpath(envid):\n",
    "    # Make a local file directories\n",
    "    path = Path('.') / 'files'\n",
    "    os.makedirs(path, exist_ok=True)\n",
    "    path = path / envid\n",
    "    os.makedirs(path, exist_ok=True)\n",
    "    return path\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_wrapper(env, wrap_string, params):\n",
    "    wrap_str = ALL_WRAPPERS[wrap_string]\n",
    "    wrap_module = importlib.import_module(wrap_str.split(\":\")[0])\n",
    "    wrap_method = getattr(wrap_module, wrap_str.split(\":\")[1])\n",
    "    return wrap_method(env, **params)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_env(env_id, rank, seed=0, wrapps={}, **kwargs):\n",
    "    \"\"\"\n",
    "    Utility function for multiprocessed env.\n",
    "    :param env_id: (str) the environment ID\n",
    "    :param rank: (int) index of the subprocess\n",
    "    :param seed: (int) the inital seed for RNG\n",
    "    \"\"\"\n",
    "    def _init():\n",
    "        env = gym.make(env_id, **kwargs)\n",
    "        env.seed(seed + rank)\n",
    "        for wrap in wrapps.keys():\n",
    "            if not (wrap == 'MonitorExtended-v0' and rank != 0):\n",
    "                env = apply_wrapper(env, wrap, wrapps[wrap])\n",
    "        return env\n",
    "    set_global_seeds(seed)\n",
    "    return _init\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_alg(alg):\n",
    "    if alg == \"A2C\":\n",
    "        from stable_baselines import A2C as algo\n",
    "    elif alg == \"ACER\":\n",
    "        from stable_baselines import ACER as algo\n",
    "    elif alg == \"ACKTR\":\n",
    "        from stable_baselines import ACKTR as algo\n",
    "    elif alg == \"PPO2\":\n",
    "        from stable_baselines import PPO2 as algo\n",
    "    return algo\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Supervised training networks.\n",
    "\n",
    "Save network in a path determined by environment ID.\n",
    "\n",
    "Args:\n",
    "    envid: str, environment ID.\n",
    "\"\"\"\n",
    "\n",
    "modelpath = get_modelpath(envid)\n",
    "config = {\n",
    "    'dt': 100,\n",
    "    'hidden_size': 64,\n",
    "    'lr': 1e-2,\n",
    "    'alg': 'ACER',\n",
    "    'rollout': 20,\n",
    "    'n_thrds': 1,\n",
    "    'wrappers_kwargs': {},\n",
    "    'alg_kwargs': {},\n",
    "    'seed': 0,\n",
    "    # 'num_steps': 100000,\n",
    "    'num_steps': 100,\n",
    "    'envid': envid,\n",
    "}\n",
    "\n",
    "env_kwargs = {'dt': config['dt']}\n",
    "config['env_kwargs'] = env_kwargs\n",
    "\n",
    "# Save config\n",
    "with open(modelpath / 'config.json', 'w') as f:\n",
    "    json.dump(config, f)\n",
    "algo = get_alg(config['alg'])\n",
    "# Make supervised dataset\n",
    "make_envs = [make_env(env_id=envid, rank=i, seed=config['seed'],\n",
    "                              wrapps=config['wrappers_kwargs'],\n",
    "                              **env_kwargs)\n",
    "                     for i in range(config['n_thrds'])]\n",
    "# env = SubprocVecEnv(make_envs)\n",
    "env = DummyVecEnv(make_envs)  # Less efficient but more robust\n",
    "model = algo(LstmPolicy, env, verbose=0, n_steps=config['rollout'],\n",
    "             n_cpu_tf_sess=config['n_thrds'], tensorboard_log=None,\n",
    "             policy_kwargs={\"feature_extraction\": \"mlp\",\n",
    "                            \"n_lstm\": config['hidden_size']},\n",
    "             **config['alg_kwargs'])\n",
    "chckpnt_cllbck = CheckpointCallback(save_freq=int(config['num_steps']/10),\n",
    "                                    save_path=modelpath,\n",
    "                                    name_prefix='model')\n",
    "model.learn(total_timesteps=config['num_steps'], callback=chckpnt_cllbck)\n",
    "print('Finished Training')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def infer_test_timing(env):\n",
    "    \"\"\"Infer timing of environment for testing.\"\"\"\n",
    "    timing = {}\n",
    "    for period in env.timing.keys():\n",
    "        period_times = [env.sample_time(period) for _ in range(100)]\n",
    "        timing[period] = np.median(period_times)\n",
    "    return timing\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extend_obs(ob, num_threads):\n",
    "    sh = ob.shape\n",
    "    return np.concatenate((ob, np.zeros((num_threads-sh[0], sh[1]))))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def order_by_sufix(file_list):\n",
    "    file_list = [os.path.basename(x) for x in file_list]\n",
    "    flag = 'model.zip' in file_list\n",
    "    file_list = [x for x in file_list if x != 'model.zip']\n",
    "    sfx = [int(x[x.find('_')+1:x.rfind('_')]) for x in file_list]\n",
    "    sorted_list = [x for _, x in sorted(zip(sfx, file_list))]\n",
    "    if flag:\n",
    "        sorted_list.append('model.zip')\n",
    "    return sorted_list, np.max(sfx)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run network after training for analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Run trained networks for analysis.\n",
    "\n",
    "Args:\n",
    "    envid: str, Environment ID\n",
    "\n",
    "Returns:\n",
    "    activity: a list of activity matrices\n",
    "    info: pandas dataframe, each row is information of a trial\n",
    "    config: dict of network, training configurations\n",
    "\"\"\"\n",
    "modelpath = get_modelpath(envid)\n",
    "files = glob.glob(str(modelpath)+'/model*')\n",
    "if len(files) > 0:\n",
    "    with open(modelpath / 'config.json') as f:\n",
    "        config = json.load(f)\n",
    "    env_kwargs = config['env_kwargs']\n",
    "    wrappers_kwargs = config['wrappers_kwargs']\n",
    "    seed = config['seed']\n",
    "    # Run network to get activity and info\n",
    "    sorted_models, last_model = order_by_sufix(files)\n",
    "    model_name = sorted_models[-1]\n",
    "    algo = get_alg(config['alg'])\n",
    "    model = algo.load(modelpath / model_name, tensorboard_log=None,\n",
    "                      custom_objects={'verbose': 0})\n",
    "\n",
    "    # Environment\n",
    "    env = make_env(env_id=envid, rank=0, seed=seed, wrapps=wrappers_kwargs,\n",
    "                   **env_kwargs)()\n",
    "    env.timing = infer_test_timing(env)\n",
    "    env.reset(no_step=True)\n",
    "    # Instantiate the network and print information\n",
    "    activity = list()\n",
    "    state_mat = []\n",
    "    ob = env.reset()\n",
    "    _states = None\n",
    "    done = False\n",
    "    info_df = pd.DataFrame()\n",
    "    # num_steps = 10 ** 5\n",
    "    num_steps = 10 ** 3\n",
    "    for stp in range(int(num_steps)):\n",
    "        ob = np.reshape(ob, (1, ob.shape[0]))\n",
    "        done = [done] + [False for _ in range(config['n_thrds']-1)]\n",
    "        action, _states = model.predict(extend_obs(ob, config['n_thrds']),\n",
    "                                        state=_states, mask=done)\n",
    "        action = action[0]\n",
    "        ob, rew, done, info = env.step(action)\n",
    "        if done:\n",
    "            env.reset()\n",
    "        if isinstance(info, (tuple, list)):\n",
    "            info = info[0]\n",
    "            action = action[0]\n",
    "        state_mat.append(_states[0, int(_states.shape[1]/2):])\n",
    "        if info['new_trial']:\n",
    "            gt = env.gt_now\n",
    "            correct = action == gt\n",
    "            # Log trial info\n",
    "            trial_info = env.trial\n",
    "            trial_info.update({'correct': correct, 'choice': action})\n",
    "            info_df = info_df.append(trial_info, ignore_index=True)\n",
    "            # Log stimulus period activity\n",
    "            state_mat = np.array(state_mat)\n",
    "            \n",
    "            # Excluding decision period if exists\n",
    "            if 'decision' in env.start_ind:\n",
    "                state_mat = state_mat[:env.start_ind['decision']]\n",
    "            \n",
    "            activity.append(state_mat)\n",
    "            state_mat = []\n",
    "    env.close()\n",
    "\n",
    "    activity = np.array(activity)\n",
    "    info = info_df\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### General analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def analysis_average_activity(activity, info, config):\n",
    "    # Load and preprocess results\n",
    "    plt.figure(figsize=(1.2, 0.8))\n",
    "    t_plot = np.arange(activity.shape[1]) * config['dt']\n",
    "    plt.plot(t_plot, activity.mean(axis=0).mean(axis=-1))\n",
    "\n",
    "analysis_average_activity(activity, info, config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_conditions(info):\n",
    "    \"\"\"Get a list of task conditions to plot.\"\"\"\n",
    "    conditions = info.columns\n",
    "    # This condition's unique value should be less than 5\n",
    "    new_conditions = list()\n",
    "    for c in conditions:\n",
    "        try:\n",
    "            n_cond = len(pd.unique(info[c]))\n",
    "            if 1 < n_cond < 5:\n",
    "                new_conditions.append(c)\n",
    "        except TypeError:\n",
    "            pass\n",
    "        \n",
    "    return new_conditions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def analysis_activity_by_condition(activity, info, config):\n",
    "    conditions = get_conditions(info)\n",
    "    for condition in conditions:\n",
    "        values = pd.unique(info[condition])\n",
    "        plt.figure(figsize=(1.2, 0.8))\n",
    "        t_plot = np.arange(activity.shape[1]) * config['dt']\n",
    "        for value in values:\n",
    "            a = activity[info[condition] == value]\n",
    "            plt.plot(t_plot, a.mean(axis=0).mean(axis=-1), label=str(value))\n",
    "        plt.legend(title=condition, loc='center left', bbox_to_anchor=(1.0, 0.5))\n",
    "\n",
    "analysis_activity_by_condition(activity, info, config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def analysis_example_units_by_condition(activity, info, config):\n",
    "    conditions = get_conditions(info)\n",
    "    if len(conditions) < 1:\n",
    "        return\n",
    "\n",
    "    example_ids = np.array([0, 1])    \n",
    "    for example_id in example_ids:        \n",
    "        example_activity = activity[:, :, example_id]\n",
    "        fig, axes = plt.subplots(\n",
    "                len(conditions), 1,  figsize=(1.2, 0.8 * len(conditions)),\n",
    "                sharex=True)\n",
    "        for i, condition in enumerate(conditions):\n",
    "            ax = axes[i]\n",
    "            values = pd.unique(info[condition])\n",
    "            t_plot = np.arange(activity.shape[1]) * config['dt']\n",
    "            for value in values:\n",
    "                a = example_activity[info[condition] == value]\n",
    "                ax.plot(t_plot, a.mean(axis=0), label=str(value))\n",
    "            ax.legend(title=condition, loc='center left', bbox_to_anchor=(1.0, 0.5))\n",
    "            ax.set_ylabel('Activity')\n",
    "            if i == len(conditions) - 1:\n",
    "                ax.set_xlabel('Time (ms)')\n",
    "            if i == 0:\n",
    "                ax.set_title('Unit {:d}'.format(example_id + 1))\n",
    "\n",
    "analysis_example_units_by_condition(activity, info, config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def analysis_pca_by_condition(activity, info, config):\n",
    "    # Reshape activity to (N_trial x N_time, N_neuron)\n",
    "    activity_reshape = np.reshape(activity, (-1, activity.shape[-1]))\n",
    "    pca = PCA(n_components=2)\n",
    "    pca.fit(activity_reshape)\n",
    "    \n",
    "    conditions = get_conditions(info)\n",
    "    for condition in conditions:\n",
    "        values = pd.unique(info[condition])\n",
    "        fig = plt.figure(figsize=(2.5, 2.5))\n",
    "        ax = fig.add_axes([0.2, 0.2, 0.7, 0.7])\n",
    "        for value in values:\n",
    "            # Get relevant trials, and average across them\n",
    "            a = activity[info[condition] == value].mean(axis=0)\n",
    "            a = pca.transform(a)  # (N_time, N_PC)\n",
    "            plt.plot(a[:, 0], a[:, 1], label=str(value))\n",
    "        plt.legend(title=condition, loc='center left', bbox_to_anchor=(1.0, 0.5))\n",
    "    \n",
    "        plt.xlabel('PC 1')\n",
    "        plt.ylabel('PC 2')\n",
    "\n",
    "analysis_pca_by_condition(activity, info, config)"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 4
}