diff --git a/docs/tutorials/rlssm_quickstart.ipynb b/docs/tutorials/rlssm_quickstart.ipynb new file mode 100644 index 000000000..f832a4daf --- /dev/null +++ b/docs/tutorials/rlssm_quickstart.ipynb @@ -0,0 +1,712 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "3e9f45e8", + "metadata": {}, + "source": [ + "# Reinforcement Learning with HSSM\n", + "RLSSMs combine a reinforcement-learning update rule with a sequential sampling\n", + "model.\n", + "In practice, learned values are updated from feedback and then used to drive both\n", + "choices and response times across trials." + ] + }, + { + "cell_type": "markdown", + "id": "416c9a6f", + "metadata": {}, + "source": [ + "## RLSSM quickstart\n", + "\n", + "This notebook is for a first RLSSM fit in HSSM.\n", + "It is organized as a practical workflow rather than a full conceptual\n", + "introduction.\n", + "\n", + "By the end you will have:\n", + "- simulated a synthetic RLSSM dataset with [`ssm-simulators`](https://github.com/lnccbrown/ssm-simulators) (`ssms.rl`),\n", + "- fit a hierarchical (multi-participant) RLSSM in HSSM from the preset `2AB_RW_Angle`,\n", + "- checked a simple recovery summary, and\n", + "- run an RLSSM-aware posterior predictive check.\n", + "\n", + "> Other tutorials cover\n", + "> custom learning/decision models\n", + "> ([Custom models with ssms.rl](rlssm_advanced.ipynb) ·\n", + "> [Restless learner](rlssm_restless_learner.ipynb)) and HSSM-native\n", + "> registration ([Registering custom models in HSSM](rlssm_hssm_custom_models.ipynb)).\n" + ] + }, + { + "cell_type": "markdown", + "id": "9973c549", + "metadata": {}, + "source": [ + "## 1. What this model does\n", + "`2AB_RW_Angle` is a ready-made RLSSM for a two-armed bandit task.\n", + "At a high level, each trial has three moving parts:\n", + "1. updates option values from feedback with a Rescorla-Wagner learner;\n", + "2. converts the value difference into a drift rate with `scaler`;\n", + "3. turns that drift into a choice and response time with the angle SSM.\n", + "\n", + "For this tutorial, the main implementation detail to keep in mind is that the free\n", + "parameters are\n", + "`rl_alpha`, `scaler`, `a`, `z`, `t`, and `theta`.\n", + "The drift `v` is not sampled directly; it is computed from the learned values on each\n", + "trial.\n" + ] + }, + { + "cell_type": "markdown", + "id": "b982c7b2", + "metadata": {}, + "source": [ + "## 2. Setup\n", + "Import HSSM and the `ssms.rl` API, then set two notebook-wide defaults." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f090de8f", + "metadata": {}, + "outputs": [], + "source": [ + "import logging\n", + "import warnings\n", + "\n", + "import arviz as az\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import pandas as pd\n", + "from ssms import rl\n", + "\n", + "import hssm\n", + "\n", + "# We silence a few noisy library warnings so the notebook output stays readable\n", + "warnings.filterwarnings(\"ignore\")\n", + "logging.getLogger(\"jax._src.xla_bridge\").setLevel(logging.ERROR)\n", + "\n", + "hssm.set_floatX(\n", + " \"float32\", update_jax=True\n", + ") # set default float type to float32 for JAX operations\n", + "RANDOM_SEED = 20260704 # for reproducibility of pseudo-random draws used in simulation, posterior sampling, and PPC subsampling." + ] + }, + { + "cell_type": "markdown", + "id": "f1205a8b", + "metadata": {}, + "source": [ + "### Run size\n", + "This notebook uses a quick configuration so you can run the full workflow without a\n", + "long wait.\n", + "\n", + "> **Optional longer run (advanced):** after your first pass, increase the values\n", + "in the next cell (participants, trials, tune, and draws) to get tighter\n", + "recovery and PPC curves.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4a45b283", + "metadata": {}, + "outputs": [], + "source": [ + "N_PARTICIPANTS = 5\n", + "N_TRIALS = 70\n", + "N_CHAINS = 2\n", + "N_TUNE = 300\n", + "N_DRAWS = 300\n", + "\n", + "print(\n", + " f\"quick mode | participants={N_PARTICIPANTS} trials={N_TRIALS} \"\n", + " f\"tune={N_TUNE} draws={N_DRAWS}\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "ae6d444a", + "metadata": {}, + "source": [ + "## 3. Load the `2AB_RW_Angle` preset\n", + "`ssms.rl` presets bundle a task, a learning rule, and a decision model into one\n", + "ready-to-use configuration. Here we use `2AB_RW_Angle`: a two-armed bandit with\n", + "Rescorla-Wagner learning and an angle decision process.\n", + "The next two cells inspect the preset and confirm which parameters are sampled versus\n", + "computed.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "81a5570b", + "metadata": {}, + "outputs": [], + "source": [ + "print(rl.preset.info(\"2AB_RW_Angle\"))" + ] + }, + { + "cell_type": "markdown", + "id": "f4b8f53f", + "metadata": {}, + "source": [ + "For this tutorial, the key distinction is simple: `rl_alpha`, `scaler`, `a`, `z`,\n", + "`t`, and `theta` are sampled, while `v` is computed from the learning process.\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5a61fc70", + "metadata": {}, + "outputs": [], + "source": [ + "ssms_config = rl.preset.get(\"2AB_RW_Angle\")\n", + "assembled = ssms_config.assemble(backend=\"jax\")\n", + "print(\"computed params (driven by the learner):\", assembled.computed_params)\n", + "assert \"v\" in assembled.computed_params" + ] + }, + { + "cell_type": "markdown", + "id": "33fe1561", + "metadata": {}, + "source": [ + "## 4. Define ground-truth parameters\n", + "We simulate data from known parameters so we can later ask whether the hierarchical\n", + "fit recovers them.\n", + "Each parameter gets a group mean plus participant-level variation. `SDS` controls how\n", + "much participants differ, and `BOUNDS` clips sampled values into the supported\n", + "ranges.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6258c663", + "metadata": {}, + "outputs": [], + "source": [ + "GROUP_THETA = {\n", + " \"rl_alpha\": 0.08, # learning rate (small -> gradual, visible learning)\n", + " \"scaler\": 2.5, # value-difference -> drift gain\n", + " \"a\": 1.2, # boundary separation\n", + " \"z\": 0.5, # starting-point bias (0.5 = unbiased)\n", + " \"t\": 0.25, # non-decision time (s)\n", + " \"theta\": 0.35, # boundary collapse angle\n", + "}\n", + "# Between-participant SD for each parameter (individual differences to recover).\n", + "SDS = {\n", + " \"rl_alpha\": 0.03,\n", + " \"scaler\": 0.40,\n", + " \"a\": 0.20,\n", + " \"z\": 0.06,\n", + " \"t\": 0.05,\n", + " \"theta\": 0.10,\n", + "}\n", + "# Keep sampled values inside supported ranges.\n", + "BOUNDS = {\n", + " \"rl_alpha\": (0.01, 1.0),\n", + " \"scaler\": (0.1, 5.0),\n", + " \"a\": (0.3, 2.5),\n", + " \"z\": (0.1, 0.9),\n", + " \"t\": (0.05, 1.0),\n", + " \"theta\": (0.0, 1.2),\n", + "}\n", + "LIST_PARAMS = list(GROUP_THETA)\n", + "\n", + "rng = np.random.default_rng(RANDOM_SEED)\n", + "theta_arrays = {\n", + " name: np.clip(\n", + " rng.normal(GROUP_THETA[name], SDS[name], N_PARTICIPANTS), *BOUNDS[name]\n", + " )\n", + " for name in LIST_PARAMS\n", + "}\n", + "\n", + "# One row per participant: their true parameter values (for the recovery check later).\n", + "true_params = pd.DataFrame(theta_arrays)\n", + "true_params.index.name = \"participant_id\"\n", + "true_params.round(3)" + ] + }, + { + "cell_type": "markdown", + "id": "4f7c7b70", + "metadata": {}, + "source": [ + "## 5. Simulate the data\n", + "`rl.Simulator(...).simulate(...)` runs the full generative loop and returns one row\n", + "per trial.\n", + "Passing parameter arrays gives each participant their own parameter values while\n", + "keeping the dataset balanced across participants.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "887db61f", + "metadata": {}, + "outputs": [], + "source": [ + "data = rl.Simulator(ssms_config).simulate(\n", + " theta=theta_arrays,\n", + " n_trials=N_TRIALS,\n", + " n_participants=N_PARTICIPANTS,\n", + " random_state=RANDOM_SEED,\n", + ")\n", + "# Validate the panel matches what the model expects before doing anything else.\n", + "ssms_config.validate_data(data).raise_for_errors()\n", + "\n", + "print(\"rows:\", len(data), \"| columns:\", list(data.columns))\n", + "data.head()" + ] + }, + { + "cell_type": "markdown", + "id": "0b7a8618", + "metadata": {}, + "source": [ + "Each row is one trial. Column descriptions:\n", + "- `participant_id`, `trial_id`: who and when.\n", + "- `rt`: response time in seconds.\n", + "- `response`: the chosen arm (-1 or 1).\n", + "- `feedback`: reward (0 or 1), used by the learner.\n", + "\n", + "The next cell is a quick sanity check that the simulated data actually show\n", + "learning. It does three things:\n", + "1. bins trials into windows of 10;\n", + "2. computes choice accuracy as P(chose high-reward arm);\n", + "3. computes mean RT per bin.\n", + "\n", + "Expected pattern: accuracy should rise above chance, and RT should decrease as the\n", + "value difference becomes easier to act on.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "991e7115", + "metadata": {}, + "outputs": [], + "source": [ + "BIN = 10\n", + "learn = data[data[\"rt\"] > 0].copy()\n", + "learn[\"chose_high\"] = (learn[\"response\"] == -1).astype(float) # -1 == high-reward arm\n", + "learn[\"trial_bin\"] = (learn[\"trial_id\"] // BIN) * BIN\n", + "acc_curve = learn.groupby(\"trial_bin\")[\"chose_high\"].mean()\n", + "rt_curve = learn.groupby(\"trial_bin\")[\"rt\"].mean()\n", + "centers = acc_curve.index + BIN / 2\n", + "\n", + "fig, axes = plt.subplots(1, 2, figsize=(12, 4), constrained_layout=True)\n", + "axes[0].plot(centers, acc_curve.values, \"o-\", color=\"tab:green\")\n", + "axes[0].axhline(0.5, color=\"0.7\", ls=\"--\", lw=1, label=\"chance\")\n", + "axes[0].set(\n", + " xlabel=\"Trial\",\n", + " ylabel=\"P(chose high-reward arm)\",\n", + " title=\"Accuracy: choices shift to the good arm\",\n", + " ylim=(0, 1),\n", + ")\n", + "axes[0].legend(frameon=False)\n", + "\n", + "axes[1].plot(centers, rt_curve.values, \"o-\", color=\"tab:purple\")\n", + "axes[1].set(xlabel=\"Trial\", ylabel=\"Mean RT (s)\", title=\"Speed: responses get faster\")\n", + "fig.suptitle(\"Learning curves (simulated data)\")\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "db9b5799", + "metadata": {}, + "source": [ + "## 6. Build the HSSM model\n", + "We now fit the same preset directly in HSSM.\n" + ] + }, + { + "cell_type": "markdown", + "id": "f07adb27", + "metadata": {}, + "source": [ + "To keep the fit readable, we use one helper that applies the same hierarchical\n", + "template to every free parameter: one group intercept plus participant-level\n", + "deviations.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "52e648b3", + "metadata": {}, + "outputs": [], + "source": [ + "# Prior on the per-participant deviations: mean 0, with a learned spread (sigma).\n", + "PARTICIPANT_EFFECT_PRIOR = {\n", + " \"name\": \"Normal\",\n", + " \"mu\": 0,\n", + " \"sigma\": {\"name\": \"HalfNormal\", \"sigma\": 0.5},\n", + "}\n", + "\n", + "\n", + "def hierarchical_param(name, lower, upper, mu, sigma):\n", + " \"\"\"Build a group intercept (TruncatedNormal) + per-participant random effect.\"\"\"\n", + " return hssm.Param(\n", + " name,\n", + " formula=f\"{name} ~ 1 + (1|participant_id)\",\n", + " prior={\n", + " \"Intercept\": hssm.Prior(\n", + " \"TruncatedNormal\", lower=lower, upper=upper, mu=mu, sigma=sigma\n", + " ),\n", + " \"1|participant_id\": PARTICIPANT_EFFECT_PRIOR,\n", + " },\n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "f0cdcee8", + "metadata": {}, + "source": [ + "A few arguments deserve a quick note:\n", + "- **`model=\"2AB_RW_Angle\"`** tells HSSM to use the same preset we simulated from.\n", + "- **`include=[...]`** applies the same hierarchical prior template to each free parameter.\n", + "- **`p_outlier=0` / `lapse=None`** turn off the outlier/lapse mixture for this first fit.\n", + "- **`process_initvals=False`** is the main RLSSM-specific setting to remember in this\n", + " basic workflow.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f44b2a62", + "metadata": {}, + "outputs": [], + "source": [ + "model = hssm.RLSSM(\n", + " data=data,\n", + " model=\"2AB_RW_Angle\",\n", + " p_outlier=0,\n", + " lapse=None,\n", + " process_initvals=False,\n", + " include=[\n", + " hierarchical_param(\"rl_alpha\", 0.01, 1.0, 0.15, 0.15),\n", + " hierarchical_param(\"scaler\", 0.1, 5.0, 2.0, 0.8),\n", + " hierarchical_param(\"a\", 0.3, 2.5, 1.1, 0.3),\n", + " hierarchical_param(\"z\", 0.1, 0.9, 0.5, 0.15),\n", + " hierarchical_param(\"t\", 0.05, 1.0, 0.25, 0.1),\n", + " hierarchical_param(\"theta\", 0.0, 1.2, 0.35, 0.15),\n", + " ],\n", + ")\n", + "\n", + "print(\"participants:\", model.n_participants, \"| trials/participant:\", model.n_trials)\n", + "print(\"free parameters:\", list(model.params.keys()))\n", + "assert model.model_name == \"2AB_RW_Angle\"\n", + "assert \"rl_alpha\" in model.params\n", + "assert \"v\" not in model.params # computed by the learner, never sampled" + ] + }, + { + "cell_type": "markdown", + "id": "23ab0dd2", + "metadata": {}, + "source": [ + "## 7. Sample the posterior\n", + "This is the estimation step.\n", + "The run is intentionally short so the notebook stays usable as a first-pass\n", + "workflow example.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e4bb5cd3", + "metadata": {}, + "outputs": [], + "source": [ + "idata = model.sample(\n", + " sampler=\"numpyro\",\n", + " draws=N_DRAWS,\n", + " tune=N_TUNE,\n", + " chains=N_CHAINS,\n", + " cores=N_CHAINS,\n", + " target_accept=0.9,\n", + " random_seed=RANDOM_SEED,\n", + ")\n", + "idata" + ] + }, + { + "cell_type": "markdown", + "id": "ad91ad20", + "metadata": {}, + "source": [ + "## 8. A quick recovery check\n", + "First check: did the fit recover the group-level values we simulated from?\n", + "We keep this check at the group level because it is the fastest way to verify that\n", + "the workflow is behaving sensibly.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ae08ea78", + "metadata": {}, + "outputs": [], + "source": [ + "def group_recovery(idata, true_group):\n", + " \"\"\"Group intercept posterior vs. true group mean, per parameter.\"\"\"\n", + " names = [f\"{p}_Intercept\" for p in LIST_PARAMS]\n", + " summ = az.summary(\n", + " idata,\n", + " var_names=names,\n", + " kind=\"stats\",\n", + " ci_kind=\"hdi\",\n", + " ci_prob=0.94,\n", + " round_to=\"none\",\n", + " )\n", + " summ.index = LIST_PARAMS\n", + " summ[\"true\"] = [true_group[p] for p in LIST_PARAMS]\n", + "\n", + " fig, ax = plt.subplots(figsize=(8, 4.5))\n", + " y = np.arange(len(LIST_PARAMS))\n", + " ax.errorbar(\n", + " summ[\"mean\"],\n", + " y,\n", + " xerr=[summ[\"mean\"] - summ[\"hdi94_lb\"], summ[\"hdi94_ub\"] - summ[\"mean\"]],\n", + " fmt=\"o\",\n", + " capsize=4,\n", + " label=\"posterior (94% HDI)\",\n", + " )\n", + " ax.scatter(\n", + " summ[\"true\"], y, color=\"crimson\", marker=\"D\", zorder=5, label=\"true group mean\"\n", + " )\n", + " ax.set_yticks(y)\n", + " ax.set_yticklabels(LIST_PARAMS)\n", + " ax.invert_yaxis()\n", + " ax.set_title(\"Group-level recovery\")\n", + " ax.legend()\n", + " fig.tight_layout()\n", + " plt.show()\n", + " return summ" + ] + }, + { + "cell_type": "markdown", + "id": "f028a492", + "metadata": {}, + "source": [ + "### Group-level recovery\n", + "Each posterior interval (blue) should land reasonably close to the corresponding\n", + "true group mean (red diamond).\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "89a41288", + "metadata": {}, + "outputs": [], + "source": [ + "group_summary = group_recovery(idata, GROUP_THETA)\n", + "group_summary[[\"mean\", \"hdi94_lb\", \"hdi94_ub\", \"true\"]].round(3)" + ] + }, + { + "cell_type": "markdown", + "id": "055a0d0e", + "metadata": {}, + "source": [ + "## 9. Posterior predictive check (RLSSM-aware)\n", + "Second check: if we simulate new data from the fitted model, does it resemble the\n", + "observed dataset?\n", + "For RLSSMs, the PPC has to respect the learning history. The helper below replays each\n", + "participant's observed reward sequence before comparing predicted and observed\n", + "learning curves and RT distributions.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "433a3627", + "metadata": {}, + "outputs": [], + "source": [ + "def plot_rl_ppc(idata, observed_data, ssms_config, *, n_participants, random_seed):\n", + " \"\"\"Simulate an RL-aware PPC and plot the main checks.\"\"\"\n", + "\n", + " def draw_posterior_theta(draw_idx):\n", + " posterior = idata.posterior\n", + " if hasattr(posterior, \"to_dataset\"):\n", + " posterior = posterior.to_dataset() # PyMC 6 returns a DataTree node\n", + " post = posterior.stack(sample=(\"chain\", \"draw\"))\n", + " theta = {}\n", + " for name in LIST_PARAMS:\n", + " re = post[f\"{name}_1|participant_id\"]\n", + " pid_dim = [d for d in re.dims if d not in (\"sample\",)][0]\n", + " vals = (post[f\"{name}_Intercept\"] + re).isel(sample=draw_idx)\n", + " ids = [int(v) for v in re[pid_dim].values]\n", + " series = pd.Series(np.asarray(vals.values), index=ids).sort_index()\n", + " theta[name] = series.reindex(range(n_participants)).to_numpy()\n", + " return theta\n", + "\n", + " def learning_curve(df, bin_size=10):\n", + " d = df[df[\"rt\"] > -900].copy()\n", + " d[\"chose_high\"] = (d[\"response\"] == -1).astype(float)\n", + " d[\"trial_bin\"] = (d[\"trial_id\"] // bin_size) * bin_size\n", + " return d.groupby(\"trial_bin\")[\"chose_high\"].mean()\n", + "\n", + " def signed_rt(df):\n", + " d = df[df[\"rt\"] > -900].copy()\n", + " return np.where(\n", + " d[\"response\"].astype(int) == -1,\n", + " -d[\"rt\"].astype(float),\n", + " d[\"rt\"].astype(float),\n", + " )\n", + "\n", + " n_ppc_draws = 8\n", + " n_samples = idata.posterior.sizes[\"chain\"] * idata.posterior.sizes[\"draw\"]\n", + " ppc_rng = np.random.default_rng(random_seed + 1)\n", + " draw_ids = ppc_rng.choice(\n", + " n_samples, size=min(n_ppc_draws, n_samples), replace=False\n", + " )\n", + " ppc_frames = []\n", + " for offset, draw_id in enumerate(draw_ids):\n", + " theta_d = draw_posterior_theta(int(draw_id))\n", + " ppc_draw = rl.Simulator(ssms_config).simulate(\n", + " theta=theta_d,\n", + " mode=\"ppc\",\n", + " observed_data=observed_data,\n", + " random_state=random_seed + 100 + offset,\n", + " )\n", + " ppc_draw[\"ppc_draw\"] = offset\n", + " ppc_frames.append(ppc_draw)\n", + " ppc_data = pd.concat(ppc_frames, ignore_index=True)\n", + " fig, axes = plt.subplots(1, 2, figsize=(12, 4.5), constrained_layout=True)\n", + " obs_curve = learning_curve(observed_data)\n", + " ppc_curves = pd.concat(\n", + " [\n", + " learning_curve(group).rename(draw)\n", + " for draw, group in ppc_data.groupby(\"ppc_draw\")\n", + " ],\n", + " axis=1,\n", + " ).sort_index()\n", + " centers = obs_curve.index + 5\n", + " axes[0].fill_between(\n", + " ppc_curves.index + 5,\n", + " ppc_curves.quantile(0.03, axis=1),\n", + " ppc_curves.quantile(0.97, axis=1),\n", + " alpha=0.25,\n", + " color=\"tab:blue\",\n", + " label=\"PPC 94% band\",\n", + " )\n", + " axes[0].plot(\n", + " ppc_curves.index + 5,\n", + " ppc_curves.mean(axis=1),\n", + " color=\"tab:blue\",\n", + " lw=1.5,\n", + " label=\"PPC mean\",\n", + " )\n", + " axes[0].plot(centers, obs_curve.values, \"o-\", color=\"black\", label=\"observed\")\n", + " axes[0].axhline(0.5, color=\"0.7\", ls=\"--\", lw=1)\n", + " axes[0].set(\n", + " xlabel=\"Trial\",\n", + " ylabel=\"P(chose high-reward arm)\",\n", + " title=\"Learning-curve PPC\",\n", + " ylim=(0, 1),\n", + " )\n", + " axes[0].legend(frameon=False)\n", + " axes[1].hist(\n", + " signed_rt(observed_data),\n", + " bins=40,\n", + " density=True,\n", + " histtype=\"step\",\n", + " lw=1.8,\n", + " color=\"black\",\n", + " label=\"observed\",\n", + " )\n", + " axes[1].hist(\n", + " signed_rt(ppc_data),\n", + " bins=40,\n", + " density=True,\n", + " histtype=\"step\",\n", + " lw=1.8,\n", + " color=\"tab:blue\",\n", + " label=\"PPC\",\n", + " )\n", + " axes[1].axvline(0, color=\"0.6\", lw=1)\n", + " axes[1].set(\n", + " xlabel=\"Signed RT (negative = high-reward choice)\",\n", + " ylabel=\"density\",\n", + " title=\"Signed-RT PPC\",\n", + " )\n", + " axes[1].legend(frameon=False)\n", + " plt.show()\n", + " print(\"PPC datasets:\", len(draw_ids), \"| total rows:\", len(ppc_data))\n", + " return ppc_data" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4a32f6f7", + "metadata": {}, + "outputs": [], + "source": [ + "ppc_data = plot_rl_ppc(\n", + " idata,\n", + " observed_data=data,\n", + " ssms_config=ssms_config,\n", + " n_participants=N_PARTICIPANTS,\n", + " random_seed=RANDOM_SEED,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "a124489a", + "metadata": {}, + "source": [ + "## 10. Summary\n", + "You have run a compact RLSSM quickstart:\n", + "1. **Chose a model** — the `2AB_RW_Angle` preset (Rescorla–Wagner learning + angle SSM).\n", + "2. **Simulated** a hierarchical dataset from known parameters and confirmed learning.\n", + "3. **Fit** the same preset directly in HSSM with a hierarchical prior template.\n", + "4. **Sampled** the posterior with NumPyro using `process_initvals=False`.\n", + "5. **Checked** one recovery summary and an RL-aware PPC.\n", + "### Where to go next\n", + "- **[Custom models with ssms.rl](rlssm_advanced.ipynb)** — build your own task\n", + " environment and learning rule instead of using a preset.\n", + "- **[Restless learner](rlssm_restless_learner.ipynb)** — one learner driving *several*\n", + " decision parameters at once.\n", + "- **[Registering custom models in HSSM](rlssm_hssm_custom_models.ipynb)** — the\n", + " HSSM-native registry path.\n", + "> **Note:** HSSM also supports *choice-only* reinforcement-learning models (no RT).\n", + "> Those are documented separately once fully validated against the current release.\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "hssm (3.13.1.final.0)", + "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" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}