{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Computational Models Lab: Q-Learning on Concurrent VI-VI Schedules\n",
    "\n",
    "In this lab, you will build a reinforcement learning agent from scratch and test whether a simple Q-learning algorithm produces choice allocation that approximates the **matching law** on concurrent variable-interval (VI) schedules.\n",
    "\n",
    "Recall from Week 2 that the generalized matching equation predicts:\n",
    "\n",
    "$$\\frac{B_1}{B_2} = b \\left(\\frac{r_1}{r_2}\\right)^a$$\n",
    "\n",
    "where $B$ is behavior allocation, $r$ is reinforcement rate, $a$ is sensitivity to reinforcement, and $b$ is bias.\n",
    "\n",
    "Here we ask: does an agent that learns purely from trial-and-error via Q-value updates converge on a steady-state allocation that looks like matching?\n",
    "\n",
    "**Objectives:**\n",
    "1. Implement a concurrent VI 30-s vs VI 60-s schedule environment\n",
    "2. Implement the Q-learning update rule\n",
    "3. Run simulations and analyze steady-state choice allocation\n",
    "4. Sweep learning rate and observe its effect on convergence\n",
    "5. Compare agent behavior to generalized matching predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Run the cell below to import the libraries you will need."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "np.random.seed(42)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 1: Define the Concurrent VI-VI Schedule Environment\n",
    "\n",
    "A variable-interval (VI) schedule arranges reinforcement for the first response after a variable amount of time has elapsed. On a **concurrent** VI-VI schedule, the organism can allocate responses to two alternatives, each operating on its own VI timer.\n",
    "\n",
    "Implement a `ConcurrentVIVI` class with the following behavior:\n",
    "\n",
    "- **Two VI timers**: one for the left alternative (VI 30 s) and one for the right (VI 60 s).\n",
    "- Each timer samples inter-reinforcement intervals from an **exponential distribution** with the appropriate mean.\n",
    "- At each time step (1 second), the timers tick down. If a timer has elapsed and the agent responds on that alternative, a reinforcer is delivered and the timer resets.\n",
    "- The environment should track: total responses to each alternative, total reinforcers from each alternative.\n",
    "\n",
    "The class needs:\n",
    "- `__init__(self, vi_left=30, vi_right=60)` \n",
    "- `reset(self)` -> returns initial state\n",
    "- `step(self, action)` -> returns `(next_state, reward, done)` where action is 0 (left) or 1 (right)\n",
    "\n",
    "Since concurrent VI schedules are stationary (no state transitions), you can use a single state. The episode ends after a fixed session duration (e.g., 1800 seconds = 30 min)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "class ConcurrentVIVI:\n",
    "    \"\"\"Concurrent VI-VI schedule environment.\"\"\"\n",
    "\n",
    "    def __init__(self, vi_left=30, vi_right=60, session_duration=1800):\n",
    "        self.vi_left = vi_left\n",
    "        self.vi_right = vi_right\n",
    "        self.session_duration = session_duration\n",
    "\n",
    "    def reset(self):\n",
    "        \"\"\"Reset the environment for a new session. Return initial state.\"\"\"\n",
    "        self.t = 0\n",
    "        # Time (s) until each schedule next sets up a reinforcer.\n",
    "        self.timer_left = np.random.exponential(self.vi_left)\n",
    "        self.timer_right = np.random.exponential(self.vi_right)\n",
    "        self.avail_left = False\n",
    "        self.avail_right = False\n",
    "        self.responses = [0, 0]\n",
    "        self.reinforcers = [0, 0]\n",
    "        return 0\n",
    "\n",
    "    def step(self, action):\n",
    "        self.t += 1\n",
    "        # VI timers \"hold\": they only run while no reinforcer is already set up.\n",
    "        if not self.avail_left:\n",
    "            self.timer_left -= 1\n",
    "            if self.timer_left <= 0:\n",
    "                self.avail_left = True\n",
    "                self.timer_left = np.random.exponential(self.vi_left)\n",
    "        if not self.avail_right:\n",
    "            self.timer_right -= 1\n",
    "            if self.timer_right <= 0:\n",
    "                self.avail_right = True\n",
    "                self.timer_right = np.random.exponential(self.vi_right)\n",
    "\n",
    "        # The agent emits one response on the chosen alternative this second.\n",
    "        self.responses[action] += 1\n",
    "        reward = 0.0\n",
    "        if action == 0 and self.avail_left:\n",
    "            reward = 1.0\n",
    "            self.avail_left = False\n",
    "            self.reinforcers[0] += 1\n",
    "        elif action == 1 and self.avail_right:\n",
    "            reward = 1.0\n",
    "            self.avail_right = False\n",
    "            self.reinforcers[1] += 1\n",
    "\n",
    "        done = self.t >= self.session_duration\n",
    "        return 0, reward, done"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 2: Implement the Q-Learning Agent\n",
    "\n",
    "Implement a Q-learning agent. Because our environment has only one state and two actions, the Q-table is simply a 1x2 array (or just a length-2 vector).\n",
    "\n",
    "The Q-value update rule is:\n",
    "\n",
    "$$Q(s, a) \\leftarrow Q(s, a) + \\alpha \\left[ r + \\gamma \\max_{a'} Q(s', a') - Q(s, a) \\right]$$\n",
    "\n",
    "where:\n",
    "- $\\alpha$ is the learning rate\n",
    "- $\\gamma$ is the discount factor\n",
    "- $r$ is the reward received\n",
    "- $s'$ is the next state\n",
    "\n",
    "For action selection, use a **softmax** policy:\n",
    "\n",
    "$$P(a) = \\frac{e^{Q(s,a) / \\tau}}{\\sum_{a'} e^{Q(s,a') / \\tau}}$$\n",
    "\n",
    "where $\\tau$ (tau) is a temperature parameter controlling exploration vs exploitation.\n",
    "\n",
    "Implement the `QLearningAgent` class with:\n",
    "- `__init__(self, n_actions=2, alpha=0.1, gamma=0.95, tau=0.5)`\n",
    "- `select_action(self, state)` -> action (using softmax)\n",
    "- `update(self, state, action, reward, next_state)` -> updates Q-values"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "class QLearningAgent:\n",
    "    \"\"\"Q-learning agent with softmax action selection.\"\"\"\n",
    "\n",
    "    def __init__(self, n_actions=2, alpha=0.1, gamma=0.95, tau=0.5):\n",
    "        self.n_actions = n_actions\n",
    "        self.alpha = alpha\n",
    "        self.gamma = gamma\n",
    "        self.tau = tau\n",
    "        self.Q = np.zeros(n_actions)\n",
    "\n",
    "    def select_action(self, state):\n",
    "        \"\"\"Select an action using a softmax policy.\"\"\"\n",
    "        z = self.Q / self.tau\n",
    "        z = z - np.max(z)            # subtract max for numerical stability\n",
    "        p = np.exp(z)\n",
    "        p = p / p.sum()\n",
    "        return np.random.choice(self.n_actions, p=p)\n",
    "\n",
    "    def update(self, state, action, reward, next_state):\n",
    "        \"\"\"Update Q-values using the Q-learning rule (single state).\"\"\"\n",
    "        best_next = np.max(self.Q)\n",
    "        td_target = reward + self.gamma * best_next\n",
    "        self.Q[action] += self.alpha * (td_target - self.Q[action])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 3: Run a Single Session\n",
    "\n",
    "Write a function `run_session(agent, env)` that:\n",
    "1. Resets the environment\n",
    "2. At each time step, the agent selects an action, the environment returns a reward, and the agent updates its Q-values\n",
    "3. Returns the total responses and reinforcers for each alternative\n",
    "\n",
    "Test it by running one session with default parameters and printing the results."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def run_session(agent, env):\n",
    "    \"\"\"Run one session and return choice and reinforcement counts.\"\"\"\n",
    "    state = env.reset()\n",
    "    done = False\n",
    "    while not done:\n",
    "        action = agent.select_action(state)\n",
    "        next_state, reward, done = env.step(action)\n",
    "        agent.update(state, action, reward, next_state)\n",
    "        state = next_state\n",
    "    return {\n",
    "        'responses_left': env.responses[0],\n",
    "        'responses_right': env.responses[1],\n",
    "        'reinforcers_left': env.reinforcers[0],\n",
    "        'reinforcers_right': env.reinforcers[1],\n",
    "        'q_values': agent.Q.copy(),\n",
    "    }\n",
    "\n",
    "\n",
    "# Test with one session\n",
    "env = ConcurrentVIVI()\n",
    "agent = QLearningAgent()\n",
    "result = run_session(agent, env)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 4: Run Multiple Sessions and Track Learning\n",
    "\n",
    "Now run the agent for **50 sessions** (the Q-values persist across sessions, simulating an organism learning over days).\n",
    "\n",
    "For each session, record:\n",
    "- The proportion of responses allocated to the left (richer) alternative: $B_L / (B_L + B_R)$\n",
    "- The proportion of reinforcers obtained from the left alternative: $r_L / (r_L + r_R)$\n",
    "- The Q-values at the end of each session\n",
    "\n",
    "Plot:\n",
    "1. **Learning curve**: Choice proportion (left) across sessions\n",
    "2. **Q-values**: Q(left) and Q(right) across sessions\n",
    "3. Add a horizontal dashed line at the **matching prediction** (for VI 30 vs VI 60, the matching law predicts $B_L / (B_L + B_R) = 0.667$)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "n_sessions = 50\n",
    "env = ConcurrentVIVI(vi_left=30, vi_right=60)\n",
    "agent = QLearningAgent(alpha=0.1, gamma=0.95, tau=0.5)\n",
    "\n",
    "choice_prop, reinf_prop, q_left, q_right = [], [], [], []\n",
    "for s in range(n_sessions):\n",
    "    res = run_session(agent, env)\n",
    "    bl, br = res['responses_left'], res['responses_right']\n",
    "    rl, rr = res['reinforcers_left'], res['reinforcers_right']\n",
    "    choice_prop.append(bl / (bl + br))\n",
    "    reinf_prop.append(rl / (rl + rr) if (rl + rr) > 0 else np.nan)\n",
    "    q_left.append(res['q_values'][0])\n",
    "    q_right.append(res['q_values'][1])\n",
    "\n",
    "# Matching prediction for VI 30 vs VI 60: rate_L / (rate_L + rate_R)\n",
    "matching_pred = (1 / 30) / (1 / 30 + 1 / 60)  # = 0.667\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n",
    "axes[0].plot(choice_prop, label='Choice proportion (left)')\n",
    "axes[0].plot(reinf_prop, label='Reinforcer proportion (left)', alpha=0.6)\n",
    "axes[0].axhline(matching_pred, ls='--', color='k',\n",
    "                label=f'Matching prediction ({matching_pred:.3f})')\n",
    "axes[0].set_xlabel('Session'); axes[0].set_ylabel('Proportion (left)')\n",
    "axes[0].set_title('Learning curve'); axes[0].set_ylim(0, 1); axes[0].legend()\n",
    "\n",
    "axes[1].plot(q_left, label='Q(left)')\n",
    "axes[1].plot(q_right, label='Q(right)')\n",
    "axes[1].set_xlabel('Session'); axes[1].set_ylabel('Q-value')\n",
    "axes[1].set_title('Q-values across sessions'); axes[1].legend()\n",
    "plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 5: Does the Agent Match?\n",
    "\n",
    "Analyze the agent's steady-state behavior (last 10 sessions):\n",
    "\n",
    "1. Compute the mean choice proportion for the left alternative over the last 10 sessions.\n",
    "2. Compute the mean reinforcer proportion for the left alternative over the last 10 sessions.\n",
    "3. Compare to the strict matching prediction: $B_L / B_R = r_L / r_R$ (i.e., $B_L/(B_L+B_R) \\approx 0.667$).\n",
    "4. Does the agent **undermatch**, **match**, or **overmatch**? \n",
    "\n",
    "Print the results and write 2-3 sentences interpreting what you find."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "last10_choice = np.mean(choice_prop[-10:])\n",
    "last10_reinf = np.nanmean(reinf_prop[-10:])\n",
    "print(f\"Mean choice proportion (left), last 10 sessions: {last10_choice:.3f}\")\n",
    "print(f\"Mean reinforcer proportion (left), last 10 sessions: {last10_reinf:.3f}\")\n",
    "print(f\"Strict matching prediction: {matching_pred:.3f}\")\n",
    "\n",
    "if last10_choice < matching_pred - 0.02:\n",
    "    verdict = \"undermatching (allocation less extreme than reinforcement)\"\n",
    "elif last10_choice > matching_pred + 0.02:\n",
    "    verdict = \"overmatching (allocation more extreme than reinforcement)\"\n",
    "else:\n",
    "    verdict = \"approximate matching\"\n",
    "print(f\"Verdict: the agent shows {verdict}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Write your interpretation here.*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 6: Learning Rate Sweep\n",
    "\n",
    "The learning rate $\\alpha$ controls how quickly the agent updates its value estimates. Sweep across **5 values** of alpha:\n",
    "\n",
    "$$\\alpha \\in \\{0.01, 0.05, 0.1, 0.3, 0.5\\}$$\n",
    "\n",
    "For each alpha value, run 50 sessions and record the choice proportion for the left alternative across sessions.\n",
    "\n",
    "Plot all 5 learning curves on the same figure (use different colors and a legend). Add the matching prediction line.\n",
    "\n",
    "Answer: How does alpha affect (a) the speed of convergence and (b) the stability of steady-state behavior?"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "alphas = [0.01, 0.05, 0.1, 0.3, 0.5]\n",
    "plt.figure(figsize=(10, 6))\n",
    "for a in alphas:\n",
    "    env = ConcurrentVIVI(vi_left=30, vi_right=60)\n",
    "    agent = QLearningAgent(alpha=a, gamma=0.95, tau=0.5)\n",
    "    cp = []\n",
    "    for s in range(50):\n",
    "        res = run_session(agent, env)\n",
    "        cp.append(res['responses_left'] / (res['responses_left'] + res['responses_right']))\n",
    "    plt.plot(cp, label=f'alpha = {a}')\n",
    "plt.axhline(matching_pred, ls='--', color='k', label=f'Matching ({matching_pred:.3f})')\n",
    "plt.xlabel('Session'); plt.ylabel('Choice proportion (left)')\n",
    "plt.title('Effect of learning rate on convergence'); plt.ylim(0, 1); plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Write your answer here.*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 7: Discount Factor Sweep\n",
    "\n",
    "Now hold alpha constant at 0.1 and sweep the discount factor:\n",
    "\n",
    "$$\\gamma \\in \\{0.0, 0.5, 0.9, 0.95, 0.99\\}$$\n",
    "\n",
    "Since this is a stationary single-state problem, think carefully about what role gamma plays. Does it matter here? Why or why not?\n",
    "\n",
    "Run the sweep and plot the results."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "gammas = [0.0, 0.5, 0.9, 0.95, 0.99]\n",
    "plt.figure(figsize=(10, 6))\n",
    "for gm in gammas:\n",
    "    env = ConcurrentVIVI(vi_left=30, vi_right=60)\n",
    "    agent = QLearningAgent(alpha=0.1, gamma=gm, tau=0.5)\n",
    "    cp = []\n",
    "    for s in range(50):\n",
    "        res = run_session(agent, env)\n",
    "        cp.append(res['responses_left'] / (res['responses_left'] + res['responses_right']))\n",
    "    plt.plot(cp, label=f'gamma = {gm}')\n",
    "plt.axhline(matching_pred, ls='--', color='k', label=f'Matching ({matching_pred:.3f})')\n",
    "plt.xlabel('Session'); plt.ylabel('Choice proportion (left)')\n",
    "plt.title('Effect of discount factor (single-state, stationary task)')\n",
    "plt.ylim(0, 1); plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Write your interpretation of gamma's role here.*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 8: Compare to the Generalized Matching Equation\n",
    "\n",
    "Recall from Week 2 that the generalized matching equation in logarithmic form is:\n",
    "\n",
    "$$\\log\\left(\\frac{B_1}{B_2}\\right) = a \\log\\left(\\frac{r_1}{r_2}\\right) + \\log(b)$$\n",
    "\n",
    "To test this more thoroughly, run the Q-learning agent on **multiple concurrent VI-VI schedule pairs**:\n",
    "\n",
    "| Condition | Left VI (s) | Right VI (s) | Reinforcer ratio (left:right) |\n",
    "|-----------|-------------|--------------|-------------------------------|\n",
    "| 1         | 20          | 60           | 3:1                           |\n",
    "| 2         | 30          | 60           | 2:1                           |\n",
    "| 3         | 30          | 30           | 1:1                           |\n",
    "| 4         | 60          | 30           | 1:2                           |\n",
    "| 5         | 60          | 20           | 1:3                           |\n",
    "\n",
    "For each condition:\n",
    "1. Run 50 sessions\n",
    "2. Compute the steady-state (last 10 sessions) log response ratio and log reinforcer ratio\n",
    "\n",
    "Then:\n",
    "1. Plot log(B_L/B_R) vs log(r_L/r_R)\n",
    "2. Fit a linear regression to get the sensitivity ($a$) and bias ($\\log b$) parameters\n",
    "3. Compare: does $a \\approx 1$ (strict matching)? Is $a < 1$ (undermatching)?"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "conditions = [\n",
    "    {'vi_left': 20, 'vi_right': 60},\n",
    "    {'vi_left': 30, 'vi_right': 60},\n",
    "    {'vi_left': 30, 'vi_right': 30},\n",
    "    {'vi_left': 60, 'vi_right': 30},\n",
    "    {'vi_left': 60, 'vi_right': 20},\n",
    "]\n",
    "\n",
    "log_B, log_r = [], []\n",
    "for cond in conditions:\n",
    "    env = ConcurrentVIVI(vi_left=cond['vi_left'], vi_right=cond['vi_right'])\n",
    "    agent = QLearningAgent(alpha=0.1, gamma=0.95, tau=0.5)\n",
    "    BL = BR = rL = rR = 0\n",
    "    for s in range(50):\n",
    "        res = run_session(agent, env)\n",
    "        if s >= 40:  # accumulate over the last 10 sessions\n",
    "            BL += res['responses_left']; BR += res['responses_right']\n",
    "            rL += res['reinforcers_left']; rR += res['reinforcers_right']\n",
    "    log_B.append(np.log(BL / BR))\n",
    "    log_r.append(np.log(rL / rR))\n",
    "\n",
    "log_B = np.array(log_B); log_r = np.array(log_r)\n",
    "slope, intercept = np.polyfit(log_r, log_B, 1)\n",
    "print(f\"Sensitivity (a):  {slope:.3f}\")\n",
    "print(f\"Bias (log b):     {intercept:.3f}\")\n",
    "\n",
    "plt.figure(figsize=(8, 8))\n",
    "plt.scatter(log_r, log_B, s=80, zorder=3)\n",
    "xline = np.linspace(log_r.min(), log_r.max(), 100)\n",
    "plt.plot(xline, slope * xline + intercept, 'r-',\n",
    "         label=f'Fit: a = {slope:.2f}, log b = {intercept:.2f}')\n",
    "plt.plot(xline, xline, 'k--', alpha=0.5, label='Strict matching (a = 1)')\n",
    "plt.axhline(0, color='gray', lw=0.5); plt.axvline(0, color='gray', lw=0.5)\n",
    "plt.xlabel('log(r_L / r_R)'); plt.ylabel('log(B_L / B_R)')\n",
    "plt.title('Generalized matching: Q-learning agent'); plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Interpret the sensitivity and bias parameters. How well does Q-learning approximate the generalized matching equation? What does the sensitivity parameter tell us about the agent's behavior?*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 9: Discussion\n",
    "\n",
    "Write a short discussion (3-5 paragraphs) addressing:\n",
    "\n",
    "1. **Emergence of matching**: Did the Q-learning agent produce behavior consistent with the matching law? Was it strict matching, undermatching, or overmatching? Why might a simple RL algorithm produce matching-like behavior?\n",
    "\n",
    "2. **Mechanism vs. description**: The matching law is a *descriptive* model -- it summarizes behavioral regularities. Q-learning is a *process* model -- it specifies a mechanism that generates behavior. What are the advantages and limitations of each approach?\n",
    "\n",
    "3. **Biological plausibility**: How might Q-learning relate to actual neural mechanisms of reinforcement learning (e.g., dopamine prediction error signaling)? What aspects of real behavior on concurrent schedules would be difficult for this simple Q-learning model to capture (e.g., changeover delay effects, momentary vs. molar matching)?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Write your discussion here.*"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}