urn.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We want to model an urn with multi-colored balls,\n",
    "# where we can change the composition by adding balls.\n",
    "# Furthermore, we want to be able to randomly draw\n",
    "# from this urn, and visualize the change in composition\n",
    "# (and therefore a change in the probability distribution.)\n",
    "import numpy.random\n",
    "from collections import Counter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's make a dictionary.\n",
    "urn = dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add initial balls.\n",
    "urn.update({\"blue\":3, \"yellow\":3, \"red\":3})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'blue': 3, 'yellow': 3, 'red': 3}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "urn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "urn['blue'] += 1\n",
    "urn['blue']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'blue': 4, 'yellow': 3, 'red': 3}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "urn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "total_num = sum([urn[color] for color in urn])\n",
    "total_num"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.4, 0.3, 0.3]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "probabilities = [urn[color]/total_num for color in urn]\n",
    "probabilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "zip?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'blue': 0.4, 'yellow': 0.3, 'red': 0.3}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# We can also make a dictionary to keep track of probabilities\n",
    "prob_dict = dict(zip(urn.keys(),probabilities))\n",
    "prob_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "#for color in urn:\n",
    "#    probabilities.append(urn[color]/total_num)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "numpy.random.choice?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Counter({'yellow': 3049, 'blue': 3980, 'red': 2971})"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Drawing from (static) urn (\"with replacement\")\n",
    "draw_ten_thousand = numpy.random.choice(list(urn.keys()), size=10000, p=probabilities)\n",
    "Counter(draw_ten_thousand)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now we have an urn, and we know how to draw from it.\n",
    "# A toy model of reinforcement learning will be a loop.\n",
    "# The loop will encode assumptions about reinforcement,\n",
    "# e.g. the goal, the amount of reinforcement (punishment)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reinforcement_loop(urn=None, goal=None, iterations=10):\n",
    "    '''\n",
    "    This function is a simple reinforcement loop.\n",
    "    \n",
    "    Takes three arguments.\n",
    "    \n",
    "    urn: dictionary\n",
    "    goal: key from urn\n",
    "    iterations: integer\n",
    "    \n",
    "    '''\n",
    "    \n",
    "    assert goal in urn.keys(), 'Please choose a color from the urn'\n",
    "    \n",
    "    # Number of loops\n",
    "    i = 0\n",
    "    \n",
    "    # Initialize probabilities\n",
    "    total_num = sum([urn[color] for color in urn])\n",
    "    probabilities = [urn[color]/total_num for color in urn]\n",
    "    \n",
    "    prob_dict = dict(zip(urn.keys(),probabilities))\n",
    "    \n",
    "    # Loop\n",
    "    while i < iterations:\n",
    "    \n",
    "        # Draw\n",
    "        draw = numpy.random.choice(list(urn.keys()), size=1, p=probabilities)\n",
    "        print(f\"I drew a {draw.item()} ball. There were {total_num} balls, and the probability of drawing {draw.item()} was {prob_dict[draw.item()]}.\")\n",
    "        \n",
    "        # Reinforce +1 if success!\n",
    "        if draw == goal:\n",
    "            urn[goal] += 1\n",
    "        \n",
    "        # Re-calculate total number of balls in urn\n",
    "        total_num = sum([urn[color] for color in urn])\n",
    "        print(f'Number of balls in urn after reinforcing: {total_num}')\n",
    "        \n",
    "        # Update probabilities\n",
    "        probabilities = [urn[color]/total_num for color in urn]\n",
    "        prob_dict = dict(zip(urn.keys(), probabilities))\n",
    "        \n",
    "        # Increment i\n",
    "        i += 1\n",
    "    \n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'blue': 4, 'yellow': 3, 'red': 3}"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "urn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I drew a blue ball. There were 9 balls, and the probability of drawing blue was 0.3333333333333333.\n",
      "Number of balls in urn after reinforcing: 10\n",
      "I drew a red ball. There were 10 balls, and the probability of drawing red was 0.3.\n",
      "Number of balls in urn after reinforcing: 10\n",
      "I drew a yellow ball. There were 10 balls, and the probability of drawing yellow was 0.3.\n",
      "Number of balls in urn after reinforcing: 10\n",
      "I drew a red ball. There were 10 balls, and the probability of drawing red was 0.3.\n",
      "Number of balls in urn after reinforcing: 10\n",
      "I drew a yellow ball. There were 10 balls, and the probability of drawing yellow was 0.3.\n",
      "Number of balls in urn after reinforcing: 10\n",
      "I drew a blue ball. There were 10 balls, and the probability of drawing blue was 0.4.\n",
      "Number of balls in urn after reinforcing: 11\n",
      "I drew a blue ball. There were 11 balls, and the probability of drawing blue was 0.45454545454545453.\n",
      "Number of balls in urn after reinforcing: 12\n",
      "I drew a blue ball. There were 12 balls, and the probability of drawing blue was 0.5.\n",
      "Number of balls in urn after reinforcing: 13\n",
      "I drew a yellow ball. There were 13 balls, and the probability of drawing yellow was 0.23076923076923078.\n",
      "Number of balls in urn after reinforcing: 13\n",
      "I drew a blue ball. There were 13 balls, and the probability of drawing blue was 0.5384615384615384.\n",
      "Number of balls in urn after reinforcing: 14\n",
      "I drew a blue ball. There were 14 balls, and the probability of drawing blue was 0.5714285714285714.\n",
      "Number of balls in urn after reinforcing: 15\n",
      "I drew a blue ball. There were 15 balls, and the probability of drawing blue was 0.6.\n",
      "Number of balls in urn after reinforcing: 16\n",
      "I drew a blue ball. There were 16 balls, and the probability of drawing blue was 0.625.\n",
      "Number of balls in urn after reinforcing: 17\n",
      "I drew a blue ball. There were 17 balls, and the probability of drawing blue was 0.6470588235294118.\n",
      "Number of balls in urn after reinforcing: 18\n",
      "I drew a blue ball. There were 18 balls, and the probability of drawing blue was 0.6666666666666666.\n",
      "Number of balls in urn after reinforcing: 19\n",
      "I drew a blue ball. There were 19 balls, and the probability of drawing blue was 0.6842105263157895.\n",
      "Number of balls in urn after reinforcing: 20\n",
      "I drew a red ball. There were 20 balls, and the probability of drawing red was 0.15.\n",
      "Number of balls in urn after reinforcing: 20\n",
      "I drew a blue ball. There were 20 balls, and the probability of drawing blue was 0.7.\n",
      "Number of balls in urn after reinforcing: 21\n",
      "I drew a yellow ball. There were 21 balls, and the probability of drawing yellow was 0.14285714285714285.\n",
      "Number of balls in urn after reinforcing: 21\n",
      "I drew a blue ball. There were 21 balls, and the probability of drawing blue was 0.7142857142857143.\n",
      "Number of balls in urn after reinforcing: 22\n",
      "I drew a blue ball. There were 22 balls, and the probability of drawing blue was 0.7272727272727273.\n",
      "Number of balls in urn after reinforcing: 23\n",
      "I drew a blue ball. There were 23 balls, and the probability of drawing blue was 0.7391304347826086.\n",
      "Number of balls in urn after reinforcing: 24\n",
      "I drew a blue ball. There were 24 balls, and the probability of drawing blue was 0.75.\n",
      "Number of balls in urn after reinforcing: 25\n",
      "I drew a blue ball. There were 25 balls, and the probability of drawing blue was 0.76.\n",
      "Number of balls in urn after reinforcing: 26\n",
      "I drew a blue ball. There were 26 balls, and the probability of drawing blue was 0.7692307692307693.\n",
      "Number of balls in urn after reinforcing: 27\n",
      "I drew a yellow ball. There were 27 balls, and the probability of drawing yellow was 0.1111111111111111.\n",
      "Number of balls in urn after reinforcing: 27\n",
      "I drew a blue ball. There were 27 balls, and the probability of drawing blue was 0.7777777777777778.\n",
      "Number of balls in urn after reinforcing: 28\n",
      "I drew a red ball. There were 28 balls, and the probability of drawing red was 0.10714285714285714.\n",
      "Number of balls in urn after reinforcing: 28\n",
      "I drew a blue ball. There were 28 balls, and the probability of drawing blue was 0.7857142857142857.\n",
      "Number of balls in urn after reinforcing: 29\n",
      "I drew a blue ball. There were 29 balls, and the probability of drawing blue was 0.7931034482758621.\n",
      "Number of balls in urn after reinforcing: 30\n",
      "I drew a blue ball. There were 30 balls, and the probability of drawing blue was 0.8.\n",
      "Number of balls in urn after reinforcing: 31\n",
      "I drew a blue ball. There were 31 balls, and the probability of drawing blue was 0.8064516129032258.\n",
      "Number of balls in urn after reinforcing: 32\n",
      "I drew a blue ball. There were 32 balls, and the probability of drawing blue was 0.8125.\n",
      "Number of balls in urn after reinforcing: 33\n",
      "I drew a blue ball. There were 33 balls, and the probability of drawing blue was 0.8181818181818182.\n",
      "Number of balls in urn after reinforcing: 34\n",
      "I drew a blue ball. There were 34 balls, and the probability of drawing blue was 0.8235294117647058.\n",
      "Number of balls in urn after reinforcing: 35\n",
      "I drew a blue ball. There were 35 balls, and the probability of drawing blue was 0.8285714285714286.\n",
      "Number of balls in urn after reinforcing: 36\n",
      "I drew a blue ball. There were 36 balls, and the probability of drawing blue was 0.8333333333333334.\n",
      "Number of balls in urn after reinforcing: 37\n",
      "I drew a blue ball. There were 37 balls, and the probability of drawing blue was 0.8378378378378378.\n",
      "Number of balls in urn after reinforcing: 38\n",
      "I drew a blue ball. There were 38 balls, and the probability of drawing blue was 0.8421052631578947.\n",
      "Number of balls in urn after reinforcing: 39\n",
      "I drew a blue ball. There were 39 balls, and the probability of drawing blue was 0.8461538461538461.\n",
      "Number of balls in urn after reinforcing: 40\n",
      "I drew a blue ball. There were 40 balls, and the probability of drawing blue was 0.85.\n",
      "Number of balls in urn after reinforcing: 41\n",
      "I drew a blue ball. There were 41 balls, and the probability of drawing blue was 0.8536585365853658.\n",
      "Number of balls in urn after reinforcing: 42\n",
      "I drew a blue ball. There were 42 balls, and the probability of drawing blue was 0.8571428571428571.\n",
      "Number of balls in urn after reinforcing: 43\n",
      "I drew a blue ball. There were 43 balls, and the probability of drawing blue was 0.8604651162790697.\n",
      "Number of balls in urn after reinforcing: 44\n",
      "I drew a blue ball. There were 44 balls, and the probability of drawing blue was 0.8636363636363636.\n",
      "Number of balls in urn after reinforcing: 45\n",
      "I drew a blue ball. There were 45 balls, and the probability of drawing blue was 0.8666666666666667.\n",
      "Number of balls in urn after reinforcing: 46\n",
      "I drew a blue ball. There were 46 balls, and the probability of drawing blue was 0.8695652173913043.\n",
      "Number of balls in urn after reinforcing: 47\n",
      "I drew a red ball. There were 47 balls, and the probability of drawing red was 0.06382978723404255.\n",
      "Number of balls in urn after reinforcing: 47\n",
      "I drew a blue ball. There were 47 balls, and the probability of drawing blue was 0.8723404255319149.\n",
      "Number of balls in urn after reinforcing: 48\n",
      "I drew a blue ball. There were 48 balls, and the probability of drawing blue was 0.875.\n",
      "Number of balls in urn after reinforcing: 49\n"
     ]
    }
   ],
   "source": [
    "# Reset urn since we are in a jupyter notebook and \n",
    "# changing the same dictionary object over and over.\n",
    "urn.update({\"blue\":3, \"yellow\":3, \"red\":3})\n",
    "reinforcement_loop(urn=urn, goal=\"blue\", iterations=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}