jack_env.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "from tf_agents.environments.py_environment import PyEnvironment\n",
    "from tf_agents.environments.tf_py_environment import TFPyEnvironment\n",
    "from tf_agents.replay_buffers.tf_uniform_replay_buffer import TFUniformReplayBuffer\n",
    "from tf_agents.trajectories import time_step as ts\n",
    "from tf_agents.specs.array_spec import BoundedArraySpec,ArraySpec\n",
    "from tensorflow import TensorSpec\n",
    "from tensorflow.python.framework.tensor_spec import BoundedTensorSpec\n",
    "from tf_agents.trajectories.trajectory import Trajectory\n",
    "from tf_agents.trajectories import trajectory, policy_step\n",
    "from tf_agents.trajectories.time_step import TimeStep\n",
    "from tf_agents.environments import utils\n",
    "from tensorflow.keras.layers import *\n",
    "from tensorflow.keras.models import Sequential,load_model\n",
    "from tensorflow.keras.optimizers import Adam\n",
    "\n",
    "from ppo_agent import PPOAgent\n",
    "from yim.helpers.tf_helpers import append_tensor\n",
    "from yim.helpers.np_help import np_one_hot\n",
    "from common import whole_paths"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Blackjack Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "DECK_DIR='deck.npy'\n",
    "class CardGameEnv(PyEnvironment):\n",
    "    def __init__(self,shuffle_at=51):\n",
    "        #self.shuffle_at = np.array(shuffle_at,np.int32)\n",
    "        self.deck= np.load(DECK_DIR).astype(np.float32)\n",
    "        self.card_count = np.array(0,np.int32)\n",
    "        self.cards_seen = np.zeros(52,np.int32) -1\n",
    "        \n",
    "        self._action_spec = BoundedArraySpec(shape=(),minimum=0,maximum=2,dtype=np.int32)\n",
    "        self._observation_spec = (ArraySpec(shape=(52,),dtype=np.int32),\n",
    "                                 ArraySpec(shape=(),dtype=np.int32),\n",
    "                                 ArraySpec(shape=(),dtype=np.bool),\n",
    "                                 ArraySpec(shape=(),dtype=np.int32),\n",
    "                                 ArraySpec(shape=(),dtype=np.bool),\n",
    "                                 BoundedArraySpec(shape=(),minimum=0.0,maximum=1.0,dtype=np.float32))\n",
    "        self.start_cards = np.sum(self.deck)\n",
    "        self.out_of_cards=False\n",
    "    def action_spec(self):\n",
    "        return self._action_spec\n",
    "\n",
    "    def observation_spec(self):\n",
    "        return self._observation_spec\n",
    "    def _reset(self):\n",
    "        self.new_hand()\n",
    "        obs = self.get_obs()\n",
    "        return ts.restart(obs)\n",
    "    def _step(self,action):\n",
    "        #0 = stand\n",
    "        #1 = hit\n",
    "        #2 = doubledown\n",
    "        multiplier=1.0\n",
    "        if action==2:\n",
    "            multiplier=2.0\n",
    "        if action==0:\n",
    "            rew =self.stood(multiplier)\n",
    "            self.new_hand()\n",
    "            if self.out_of_cards:\n",
    "                obs = self.get_obs()\n",
    "                rew = np.array(0,np.float32)\n",
    "                return ts.termination(obs,rew)\n",
    "            obs = self.get_obs()\n",
    "            return ts.transition(obs,rew)\n",
    "        elif action==2:\n",
    "            self.deal_card(True)\n",
    "            if self.player_score> 21:\n",
    "                rew = np.array(-1,np.float32) * multiplier\n",
    "            else:\n",
    "                rew =self.stood(multiplier)\n",
    "            self.new_hand()\n",
    "            if self.out_of_cards:\n",
    "                obs = self.get_obs()\n",
    "                rew = np.array(0,np.float32)\n",
    "                return ts.termination(obs,rew)\n",
    "            obs = self.get_obs()\n",
    "            return ts.transition(obs,rew)\n",
    "        else:\n",
    "            rew = np.array(0,np.float32)\n",
    "            self.deal_card(True)\n",
    "            if self.player_score> 21:\n",
    "                rew = np.array(-1,np.float32) * multiplier\n",
    "                self.new_hand()\n",
    "            \n",
    "            if self.out_of_cards:\n",
    "                obs = self.get_obs()\n",
    "                rew = np.array(0,np.float32)\n",
    "                return ts.termination(obs,rew)\n",
    "            obs = self.get_obs()\n",
    "            return ts.transition(obs,rew)\n",
    "    def get_obs(self):\n",
    "        return (self.cards_seen,self.player_score,\n",
    "            self.player_has_ace,self.dealer_score,self.dealer_has_ace,(self.card_count).astype(np.float32)/self.start_cards)\n",
    "    def deal_card(self,player=True):\n",
    "        if np.sum(self.deck)<1:\n",
    "            self.out_of_cards=True\n",
    "            return\n",
    "        p = self.deck / np.sum(self.deck)\n",
    "        card = np.squeeze(np.random.choice(10,1,p=p),0)\n",
    "        self.deck[card]-=1\n",
    "        self.cards_seen[self.card_count]=card\n",
    "        self.card_count+=1\n",
    "        if card==0:\n",
    "            card=10\n",
    "        if player:\n",
    "            if card==1:\n",
    "                self.player_has_ace =True\n",
    "            self.player_score+=card\n",
    "        else:\n",
    "            if card==1:\n",
    "                self.dealer_has_ace=True\n",
    "            self.dealer_score+=card\n",
    "    def new_hand(self):\n",
    "        self.ended=False\n",
    "        self.player_score=np.array(0,np.int32)\n",
    "        self.dealer_score=np.array(0,np.int32)\n",
    "        self.player_has_ace=False\n",
    "        self.dealer_has_ace=False\n",
    "        self.deal_card(False)\n",
    "        self.deal_card(True)\n",
    "        self.deal_card(True)\n",
    "        \n",
    "    def stood(self,multiplier):\n",
    "        while not self.ended:\n",
    "            self.deal_card(False)\n",
    "            if self.out_of_cards:\n",
    "                self.ended=True\n",
    "            elif self.dealer_score > 21:\n",
    "                self.ended=True\n",
    "            elif self.dealer_score>=17:\n",
    "                self.ended=True\n",
    "            elif self.dealer_has_ace:\n",
    "                if self.dealer_score>8 and self.dealer_score<=11:\n",
    "                    self.ended=True\n",
    "        if self.player_has_ace:\n",
    "            if self.player_score+10<=21:\n",
    "                self.player_score+=10\n",
    "        if self.dealer_has_ace:\n",
    "            if self.dealer_score+10<=21:\n",
    "                self.dealer_score+=10\n",
    "        if (self.dealer_score>21) or(self.dealer_score<self.player_score) :\n",
    "            return np.array(1,np.float32) * multiplier\n",
    "        if self.dealer_score==self.player_score:\n",
    "            return np.array(0,np.float32)\n",
    "        return np.array(-1,np.float32) * multiplier\n",
    "\n",
    "environment = CardGameEnv()\n",
    "utils.validate_py_environment(environment, episodes=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualization of Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualization of Environment\n",
    "def print_info(time_step,startInd=0,endInd=0,action=0):\n",
    "    def to_string(cards):\n",
    "        if len(cards)<2:\n",
    "            the_str = '?, '+str(cards)\n",
    "        else:\n",
    "            the_str = str(cards)\n",
    "        return the_str.replace('1','A').replace('[','').replace(']','').replace('0','10')\n",
    "    newHand=False\n",
    "    obs = time_step.observation[0]\n",
    "    inds = obs==-1\n",
    "    obs = obs[~inds]\n",
    "    \n",
    "    obs = list(obs)\n",
    "    dealer_cards = []\n",
    "    dealer_cards.append(obs.pop(startInd))\n",
    "    \n",
    "    if endInd==0:\n",
    "        endInd=len(obs)\n",
    "    \n",
    "    if action==1 or action==2:\n",
    "        endInd+=1\n",
    "        \n",
    "    player_cards = obs[startInd:endInd]\n",
    "    if len(obs)>endInd:\n",
    "        \n",
    "\n",
    "        newHand=True\n",
    "        if time_step.is_last():\n",
    "            dealer_cards.append(obs[endInd:])\n",
    "        else:\n",
    "            startInd = endInd\n",
    "            while startInd<len(obs)-3:\n",
    "                dealer_cards.append(obs[startInd])\n",
    "                startInd+=1\n",
    "            \n",
    "    \n",
    "    \n",
    "        \n",
    "\n",
    "    print(\"Dealer:\",to_string(dealer_cards))\n",
    "    print(\"Player:\",to_string(player_cards))\n",
    "    if newHand:\n",
    "        print(\"Reward:\",time_step.reward)\n",
    "        if time_step.is_last():\n",
    "            print('-'*20)\n",
    "            print('END OF CARDS')\n",
    "            print('-'*20)\n",
    "            return None,None\n",
    "        print('-'*20)\n",
    "        print('New Hand')\n",
    "        print('-'*20)\n",
    "        return print_info(time_step,startInd=startInd+1,endInd=0)\n",
    "    \n",
    "    return startInd,endInd\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dealer: ?, 5\n",
      "Player: 10, 9\n"
     ]
    }
   ],
   "source": [
    "env = CardGameEnv()\n",
    "time_step=env._reset()   \n",
    "startInd,endInd=print_info(time_step,startInd=0,endInd=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dealer: 5, 6, 7\n",
      "Player: 10, 9\n",
      "Reward: 1.0\n",
      "--------------------\n",
      "New Hand\n",
      "--------------------\n",
      "Dealer: ?, 3\n",
      "Player: 4, 5\n"
     ]
    }
   ],
   "source": [
    "action=0 #[0: STAND] [1: HIT] [2:DoubleDown]\n",
    "\n",
    "time_step=env._step(action)\n",
    "#print(time_step.observation[0],time_step.observation[5])\n",
    "startInd,endInd=print_info(time_step,startInd=startInd,endInd=endInd,action=action)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: Logging before flag parsing goes to stderr.\n",
      "W0307 21:38:37.178660  8584 deprecation.py:323] From c:\\users\\ericy\\appdata\\local\\programs\\python\\python36\\lib\\site-packages\\tensorflow\\python\\keras\\backend.py:3868: add_dispatch_support.<locals>.wrapper (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.where in 2.0, which has the same broadcast rule as np.where\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2] [10] [[-2.1336157]] [[0.35737604 0.28994608 0.35267782]]\n",
      "[19] [2] [[-0.22117655]] [[0.8581143  0.09495622 0.04692944]]\n",
      "[19] [5] [[0.44305655]] [[0.9708848  0.02247752 0.00663763]]\n",
      "[20] [10] [[-0.3656309]] [[0.93128973 0.05583587 0.01287438]]\n",
      "[8] [10] [[-1.6748561]] [[0.3675769  0.37231967 0.26010343]]\n",
      "[15] [10] [[-2.1670158]] [[0.4792729  0.2612255  0.25950167]]\n",
      "[6] [10] [[-1.6730815]] [[0.4617107  0.23820426 0.3000851 ]]\n",
      "[20] [4] [[-0.525481]] [[0.9155629  0.0684029  0.01603422]]\n",
      "[20] [2] [[0.6782062]] [[0.7886682  0.15634415 0.05498765]]\n",
      "[10] [10] [[0.21357304]] [[0.19042571 0.2826331  0.5269412 ]]\n",
      "[17] [10] [[0.11500724]] [[0.36174363 0.43504918 0.20320714]]\n",
      "Run 0 Result: [3.]\n",
      "[16] [7] [[-1.9615806]] [[0.36137703 0.2822516  0.35637134]]\n",
      "[19] [10] [[-1.946772]] [[0.82751304 0.11430119 0.05818576]]\n",
      "[20] [8] [[0.0803768]] [[0.7303263  0.17001581 0.0996579 ]]\n",
      "[10] [6] [[-0.7833669]] [[0.350575   0.1622716  0.48715338]]\n",
      "[3] [4] [[-0.32007378]] [[0.26205596 0.3237459  0.41419813]]\n",
      "[18] [2] [[0.33877975]] [[0.6583075  0.18385024 0.15784225]]\n",
      "[13] [2] [[-0.05687359]] [[0.50417626 0.2048702  0.29095352]]\n",
      "[6] [6] [[0.16019386]] [[0.38941312 0.2292524  0.38133448]]\n",
      "[19] [2] [[0.6637912]] [[0.73499554 0.16276662 0.10223787]]\n",
      "[17] [10] [[-0.2892655]] [[0.49294078 0.24955067 0.25750852]]\n",
      "Run 1 Result: [3.]\n",
      "[10] [7] [[-0.4963767]] [[0.20615627 0.39663506 0.39720863]]\n",
      "[17] [4] [[-0.9072999]] [[0.66464216 0.23849796 0.09685988]]\n",
      "[19] [1] [[-1.744565]] [[0.7122972  0.11312152 0.17458126]]\n",
      "[5] [8] [[-0.291517]] [[0.22742751 0.48818466 0.28438786]]\n",
      "[8] [8] [[0.34370306]] [[0.17379119 0.32664767 0.4995611 ]]\n",
      "[14] [10] [[-2.045767]] [[0.4156099  0.26521417 0.3191759 ]]\n",
      "[14] [5] [[-1.7330033]] [[0.76657444 0.12613155 0.10729402]]\n",
      "[20] [2] [[-0.12784913]] [[0.79594284 0.1576679  0.04638925]]\n",
      "[20] [10] [[-0.9301635]] [[0.8306927  0.11883692 0.05047045]]\n",
      "[15] [2] [[0.03879569]] [[0.6038881  0.20124516 0.19486673]]\n",
      "[20] [10] [[0.32078147]] [[0.81182903 0.12619263 0.06197841]]\n",
      "Run 2 Result: [5.]\n",
      "Avg: 3.6666667\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "N_RUNS = 3 #each run includes collection and model update\n",
    "#Directory\n",
    "ACTOR_CHKP = 'saved_actor'\n",
    "VALUE_CHKP = 'saved_value'\n",
    "MOD_NAME='model'\n",
    "#Collection\n",
    "N_EPISODES = 1\n",
    "NUM_ACTIONS = 3\n",
    "\n",
    "#ACTOR_CHKP,VALUE_CHKP,DECK_DIR = whole_paths([ACTOR_CHKP,VALUE_CHKP,DECK_DIR])\n",
    "\n",
    "#==============================================================\n",
    "PRINTOUT=True\n",
    "def load_model_chkpoint(dir_name=None,ver_name='None',num_outs=1):#returns the model at given chkpoint\n",
    "    if ver_name =='None':\n",
    "        dir_name = tf.train.latest_checkpoint(dir_name)\n",
    "    else:\n",
    "        dir_name = os.path.join(dir_name,ver_name)\n",
    "\n",
    "    dummy_env = TFPyEnvironment(CardGameEnv())\n",
    "    time_step = dummy_env.reset()\n",
    "    temp = BlackJackModel(num_outs)\n",
    "    #initialize model shape by running an observation through\n",
    "    temp(time_step.observation)\n",
    "    checkpoint2 = tf.train.Checkpoint(module=temp)\n",
    "    status=checkpoint2.restore(dir_name)\n",
    "    return temp,checkpoint2\n",
    "def get_env_specs():\n",
    "    dummy_env = TFPyEnvironment(CardGameEnv())\n",
    "    return dummy_env.observation_spec(),dummy_env.action_spec()\n",
    "\n",
    "\n",
    "class BlackJackModel(tf.Module):\n",
    "    def __init__(self,num_outs):\n",
    "        super(BlackJackModel,self).__init__()\n",
    "        self.lstm_root= Sequential([\n",
    "            Masking(mask_value=0.),\n",
    "            LSTM(64)\n",
    "            ])\n",
    "            \n",
    "        self.dense1 = Dense(64,activation='relu')\n",
    "        self._droprate = 0.1\n",
    "        if num_outs>1:#num_outs 1 for value. num_outs=num_actions for actor\n",
    "            self.dense2 = Sequential([\n",
    "                Dense(64,activation='relu'),\n",
    "                Dense(num_outs,activation='softmax')\n",
    "                ])\n",
    "        else:\n",
    "            self.dense2 = Sequential([\n",
    "                Dense(64,activation='relu'),\n",
    "                Dense(num_outs)\n",
    "                ])\n",
    "    def preproc_obs(self,observation):#Preprocessing\n",
    "        def revise(score_ace):\n",
    "            score,ace = score_ace\n",
    "            if ace:\n",
    "                score2 = score+10\n",
    "            else:\n",
    "                score2=0\n",
    "            return tf.maximum(tf.one_hot(score-1,21),tf.one_hot(score2-1,21))\n",
    "        \n",
    "        a = tf.one_hot(observation[0],10,dtype=tf.float32)\n",
    "        b = tf.map_fn(revise,(observation[1],observation[2]),dtype=tf.float32)\n",
    "        c = tf.map_fn(revise,(observation[3],observation[4]),dtype=tf.float32)\n",
    "        d = tf.expand_dims(tf.cast(observation[5],dtype=tf.float32),-1)\n",
    "        return a,b,c,d\n",
    "    def __call__(self,all_obs,is_training=False):\n",
    "        x0,x1,x2,x3 = self.preproc_obs(all_obs)\n",
    "        \n",
    "        x0 = self.lstm_root(x0)\n",
    "        x_all = tf.concat((x0,x1,x2,x3),axis=-1)\n",
    "        x_all=self.dense1(x_all)\n",
    "        if is_training:\n",
    "            x_all = tf.nn.dropout(x_all,self._droprate)\n",
    "        x_all = self.dense2(x_all)\n",
    "        \n",
    "        return x_all\n",
    "\n",
    "def sample_policy_tf(probac,num_ac=2):\n",
    "    probac = tf.math.log(probac)\n",
    "    a = tf.random.categorical(logits=probac,num_samples=1,dtype=tf.int32)\n",
    "    a = tf.squeeze(a,0)\n",
    "    return a\n",
    "def get_action(the_trainer,the_obs,eval=True):\n",
    "    achead = the_trainer.actor_net(the_obs)\n",
    "    val = the_trainer._value_net(the_obs)\n",
    "    if PRINTOUT:\n",
    "        print(the_obs[1].numpy(),the_obs[3].numpy(),val.numpy(),achead.numpy())\n",
    "    if eval:\n",
    "        action = tf.argmax(achead,axis=-1)\n",
    "    else:\n",
    "        action = sample_policy_tf(achead)\n",
    "    return achead,tf.cast(action,dtype=tf.int32)\n",
    "def env_runner(the_trainer):\n",
    "    env=TFPyEnvironment(CardGameEnv())\n",
    "    time_step = env.reset()\n",
    "    rew=0.0\n",
    "    hands=0\n",
    "    while not tf.reduce_all(time_step.is_last()):\n",
    "        achead,action = get_action(the_trainer,time_step.observation)\n",
    "        next_time_step = env.step(action)\n",
    "        time_step = next_time_step\n",
    "        r = time_step.reward\n",
    "\n",
    "        if time_step.observation[5]>0.95:\n",
    "            r*=1\n",
    "        rew += r\n",
    "        \n",
    "    return rew\n",
    "def collection_run(trainer,n_episodes):\n",
    "    rewards = 0.0\n",
    "    for _ in range(n_episodes):\n",
    "        reward= env_runner(trainer)\n",
    "        rewards+=reward\n",
    "    return rewards /n_episodes\n",
    "\n",
    "\n",
    "\n",
    "if __name__==\"__main__\":\n",
    "    with tf.device('/CPU:0'):\n",
    "        #MAIN\n",
    "        value_net,checkpoint_val = load_model_chkpoint(dir_name=VALUE_CHKP,num_outs=1,ver_name='model-56')\n",
    "        actor_net,checkpoint_act = load_model_chkpoint(dir_name=ACTOR_CHKP,num_outs=NUM_ACTIONS,ver_name='model-56')\n",
    "        optimizer = Adam()\n",
    "        observation_spec,_ = get_env_specs()\n",
    "        ppo_trainer = PPOAgent(\n",
    "                   optimizer=optimizer,\n",
    "                   actor_net=actor_net,\n",
    "                   value_net=value_net,\n",
    "                   observation_spec=observation_spec,\n",
    "                   num_actions=NUM_ACTIONS)\n",
    "        all_rews = []\n",
    "        for n_step in range(N_RUNS):\n",
    "            #Collect\n",
    "            rew = collection_run(ppo_trainer,N_EPISODES)\n",
    "            rew = rew.numpy()\n",
    "            print(\"Run\",n_step,\"Result:\",rew)\n",
    "            all_rews.append(rew)\n",
    "        print(\"Avg:\",np.mean(all_rews))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_RUNS = 100 #each run includes collection and model update\n",
    "#Directory\n",
    "ACTOR_CHKP = 'model_actor'\n",
    "VALUE_CHKP = 'model_value'\n",
    "#Collection\n",
    "N_EPISODES = 1\n",
    "NUM_ACTIONS = 2\n",
    "#Train\n",
    "BATCH_SIZE=32\n",
    "N_BATCHES =8\n",
    "N_STEPS = 10\n",
    "N_EPOCHS = 15\n",
    "VALUE_LOSS_COEF= 0.5\n",
    "ENTROPY_REG_COEF=0.2\n",
    "LEARNING_RATE = 1e-3\n",
    "SAVE_INTERVAL=10\n",
    "\n",
    "ACTOR_CHKP,VALUE_CHKP = whole_paths([ACTOR_CHKP,VALUE_CHKP])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Buffer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Buffer\n",
    "def get_env_specs():\n",
    "    dummy_env = TFPyEnvironment(CardGameEnv())\n",
    "    return dummy_env.observation_spec(),dummy_env.action_spec()\n",
    "def get_tf_buffer(num_outs=2,max_length = 10000):\n",
    "    obs_spec,ac_spec = get_env_specs()\n",
    "    time_step_spec = ts.time_step_spec(obs_spec)\n",
    "    info = BoundedTensorSpec(shape=(num_outs,),\n",
    "                            dtype=np.float32, \n",
    "                            minimum=np.zeros(num_outs,dtype=np.float32),\n",
    "                            maximum=np.ones(num_outs,dtype=np.float32))\n",
    "    action_spec = policy_step.PolicyStep(ac_spec,info=info)\n",
    "    trajectory_spec = trajectory.from_transition(\n",
    "        time_step_spec, action_spec , time_step_spec)\n",
    "    the_replay_buffer = TFUniformReplayBuffer(\n",
    "        data_spec=trajectory_spec,\n",
    "        batch_size=1,\n",
    "        max_length=max_length)\n",
    "    \n",
    "    return the_replay_buffer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_model_chkpoint(dir_name=None,num_outs=1):#returns the model at given chkpoint\n",
    "    dir_name = tf.train.latest_checkpoint(dir_name)\n",
    "    dummy_env = TFPyEnvironment(CardGameEnv())\n",
    "    time_step = dummy_env.reset()\n",
    "    temp = BlackJackModel(num_outs)\n",
    "    #initialize model shape by running an observation through\n",
    "    temp(time_step.observation)\n",
    "    checkpoint2 = tf.train.Checkpoint(module=temp)\n",
    "    status=checkpoint2.restore(dir_name)\n",
    "    return temp,checkpoint2\n",
    "\n",
    "#Preprocessing\n",
    "def preproc_obs(observation):\n",
    "    def revise(score_ace):\n",
    "        score,ace = score_ace\n",
    "        if ace:\n",
    "            score2 = score+10\n",
    "        else:\n",
    "            score2=0\n",
    "        return tf.maximum(tf.one_hot(score-1,21),tf.one_hot(score2-1,21))\n",
    "\n",
    "    a = tf.one_hot(observation[0],10,dtype=tf.float32)\n",
    "    b = tf.map_fn(revise,(observation[1],observation[2]),dtype=tf.float32)\n",
    "    c = tf.map_fn(revise,(observation[3],observation[4]),dtype=tf.float32)\n",
    "    return a,b,c\n",
    "class BlackJackModel(tf.Module):\n",
    "    def __init__(self,num_outs):\n",
    "        super(BlackJackModel,self).__init__()\n",
    "        self.lstm_root= Sequential([\n",
    "            Masking(mask_value=0.),\n",
    "            LSTM(64)\n",
    "            ])\n",
    "            \n",
    "        self.dense1 = Dense(64,activation='relu')\n",
    "        self._droprate = 0.25\n",
    "        if num_outs>1:#num_outs 1 for value. num_outs=num_actions for actor\n",
    "            self.dense2 = Sequential([\n",
    "                Dense(64,activation='relu'),\n",
    "                Dense(num_outs,activation='softmax')\n",
    "                ])\n",
    "        else:\n",
    "            self.dense2 = Sequential([\n",
    "                Dense(64,activation='relu'),\n",
    "                Dense(num_outs)\n",
    "                ])\n",
    "    def __call__(self,all_obs,is_training=False):\n",
    "        x0,x1,x2 = preproc_obs(all_obs)\n",
    "        x0 = self.lstm_root(x0)\n",
    "        x_all = tf.concat((x0,x1,x2),axis=-1)\n",
    "        x_all=self.dense1(x_all)\n",
    "        if is_training:\n",
    "            x_all = tf.nn.dropout(x_all,self._droprate)\n",
    "        x_all = self.dense2(x_all)\n",
    "        \n",
    "        return x_all"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_policy_tf(probac,num_ac=2):\n",
    "    probac = tf.math.log(probac)\n",
    "    a = tf.random.categorical(logits=probac,num_samples=1,dtype=tf.int32)\n",
    "    a = tf.squeeze(a,0)\n",
    "    return a\n",
    "def get_action(the_model,the_obs,eval=False):\n",
    "    achead = the_model(the_obs)\n",
    "    if eval:\n",
    "        action = tf.argmax(achead,axis=-1)\n",
    "    else:\n",
    "        action = sample_policy_tf(achead)\n",
    "    return achead,tf.cast(action,dtype=tf.int32)\n",
    "def env_runner(buffer,model):\n",
    "    env=TFPyEnvironment(CardGameEnv())\n",
    "    time_step = env.reset()\n",
    "    rew=0.0\n",
    "    while not tf.reduce_all(time_step.is_last()):\n",
    "        achead,action = get_action(model,time_step.observation,eval=eval)\n",
    "        next_time_step = env.step(action)\n",
    "        pol_step = policy_step.PolicyStep(action,info=achead)\n",
    "        #tfagents ignores reward on terminal timestep... \n",
    "        #can include it by inserting a dummy step \n",
    "        if tf.reduce_all(next_time_step.is_last()):\n",
    "            dummy_time_step=ts.transition(next_time_step.observation,next_time_step.reward)\n",
    "            traj = trajectory.from_transition(time_step, pol_step, dummy_time_step)\n",
    "            buffer.add_batch(traj)\n",
    "            #dummy traj\n",
    "            traj = trajectory.from_transition(dummy_time_step, pol_step, next_time_step)\n",
    "            buffer.add_batch(traj)\n",
    "        else:\n",
    "            traj = trajectory.from_transition(time_step, pol_step, next_time_step)\n",
    "            buffer.add_batch(traj)\n",
    "        time_step = next_time_step\n",
    "        rew += time_step.reward\n",
    "    return buffer,rew\n",
    "def collection_run(model,buffer,n_episodes):\n",
    "    rewards = 0.0\n",
    "    for _ in range(n_episodes):\n",
    "        buffer,reward= env_runner(buffer,model)\n",
    "        rewards+=reward\n",
    "    return buffer,rewards /n_episodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#MAIN\n",
    "value_net,checkpoint_val = load_model_chkpoint(dir_name=VALUE_CHKP,num_outs=1)\n",
    "actor_net,checkpoint_act = load_model_chkpoint(dir_name=ACTOR_CHKP,num_outs=NUM_ACTIONS)\n",
    "optimizer = Adam(learning_rate = LEARNING_RATE)\n",
    "observation_spec,_ = get_env_specs()\n",
    "ppo_trainer = PPOAgent(\n",
    "               optimizer=optimizer,\n",
    "               actor_net=actor_net,\n",
    "               value_net=value_net,\n",
    "               observation_spec=observation_spec,\n",
    "               num_actions=NUM_ACTIONS,\n",
    "               importance_ratio_clipping=0.2,\n",
    "               lambda_value=0.95,\n",
    "               discount_factor=0.97,\n",
    "               entropy_regularization=ENTROPY_REG_COEF,\n",
    "               value_pred_loss_coef=VALUE_LOSS_COEF,\n",
    "               num_epochs=N_EPOCHS,\n",
    "               use_gae=True,\n",
    "               use_td_lambda_return=True,\n",
    "               normalize_rewards=True,\n",
    "               reward_norm_clipping=10.0,\n",
    "               normalize_observations=False,\n",
    "               gradient_clipping=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_ppo_agent(the_experience,the_weights,the_ppo_agent):\n",
    "    loss_info = the_ppo_agent._train(the_experience,the_weights)\n",
    "    return loss_info\n",
    "\n",
    "for n_step in range(N_RUNS):\n",
    "    #Collect\n",
    "    buffer = get_tf_buffer()\n",
    "    buffer,rew = collection_run(ppo_trainer.actor_net,buffer,N_EPISODES)\n",
    "    print(\"Run\",n_step,\"Result:\",rew.numpy())\n",
    "    ds = buffer.as_dataset(num_parallel_calls=2, sample_batch_size=BATCH_SIZE, num_steps=N_STEPS).prefetch(2).repeat(-1)\n",
    "    iterator = iter(ds)\n",
    "    for _ in range(N_BATCHES):\n",
    "        experience,_ = next(iterator)\n",
    "        weights = tf.ones((BATCH_SIZE,1),dtype=tf.float32)\n",
    "        loss_info = train_ppo_agent(experience,weights,ppo_trainer)\n",
    "    if (n_step+1)%SAVE_INTERVAL==0:\n",
    "        checkpoint_val.save(save_dir_val)\n",
    "        checkpoint_act.save(save_dir_act)\n",
    "        print(n_step+1,\"VLoss:\",loss_info.extra.value_estimation_loss.numpy())"
   ]
  }
 ],
 "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}