{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inverse Planning from Scratch\n",
    "\n",
    "This tutorial explains how to use probabilistic programming to implement\n",
    "inverse planning for a simple gridworld agent. By modeling the agent as\n",
    "an *approximately rational planner* who takes efficient actions towards\n",
    "a goal, then inverting this generative process, inverse planning allows us\n",
    "to infer the agent's goal from their observed actions.\n",
    "\n",
    "We will see how to do this \"from scratch\", in the sense that we won't be using\n",
    "existing library code for planning or inverse planning. Instead we will\n",
    "implement everything using [Gen.jl](https://github.com/probcomp/Gen.jl), a\n",
    "general-purpose probabilistic programming framework that allows us to write\n",
    "generative models and inference algorithms with minimal math.\n",
    "\n",
    "Specifically, we will cover the following material:\n",
    " - Modeling agents with probabilistic programs\n",
    " - Simulating and visualizing our agent model\n",
    " - Inferring goals via inverse planning\n",
    " - Online inverse planning via incremental computation\n",
    " - TODO: Extending the model with compositional goals\n",
    " - TODO: Scaling inference to compositional goal spaces\n",
    "\n",
    "## Modeling agents with probabilistic programs\n",
    "\n",
    "Inverse planning treats goal inference — and the inference of mental states\n",
    "more generally — as an inverse problem: Given a generative model of an agent's\n",
    "goal-directed actions (that is, a model where agent's actions are generated\n",
    "conditional on a latent goal), we want to invert the model, inferring the goal\n",
    "from the actions we observe.\n",
    "\n",
    "What does this model look like? At a high level, the model should specify:\n",
    "- $P(g)$, a prior distribution over possible goals $g$ (*Goal Prior*)\n",
    "- $P(a_t | s_t, g)$, a model of how the agent takes an action $a_t$ given\n",
    "  their goal $g$ and the current environment state $s_t$ (*Action Policy*)\n",
    "- $P(s_{t+1} | s_t, a_t)$, a model of how the environment changes in\n",
    "  response to the agent's action $a_t$ (*Environment Dynamics*)\n",
    "\n",
    "Together, these components define a joint distribution over the agent's goal,\n",
    "actions, and environment states up to any timestep $T$:\n",
    "\n",
    "$$\n",
    "P(g, a_{1:T}, s_{1:T}) = P(g_1) \\prod_{t=1}^T P(a_t | s_t, g) P(s_{t+1} | s_t, a_t)\n",
    "$$\n",
    "\n",
    "In this tutorial, we'll start with a simple goal prior $P(g)$ that is\n",
    "uniform over a set of possible goals $\\mathcal{G}$. For the environment\n",
    "dynamics $P(s_{t+1} | s_t, a_t)$, we'll model a deterministic gridworld\n",
    "where the agent's location changes based on their action (i.e. if the **up**\n",
    "action is taken, the agent moves up unless there is an obstacle in the way).\n",
    "\n",
    "But what about the action policy $P(a_t | s_t, g)$? A key assumption in\n",
    "inverse planning is that agents are *approximately rational*. In particular,\n",
    "agents are more likely to take actions that lead them efficiently towards\n",
    "their goals. To capture this assumption, we'll model the agent as following\n",
    "a *Boltzmann-rational policy*:\n",
    "\n",
    "$$\n",
    "P(a_t | s_t, g) = \\frac{\\exp(\\beta \\cdot Q_g(s_t, a_t))}{\\sum_{a' \\in A(s_t)} \\exp(\\beta \\cdot Q_g(s_t, a'))}\n",
    "$$\n",
    "\n",
    "Here, $Q_g(s_t, a_t)$ is the expected utility of taking action $a_t$ at\n",
    "state $s_t$ when the agent's goal is $g$, also known as the *Q-value*.\n",
    "The parameter $\\beta$ is called the *inverse temperature* or\n",
    "*rationality coefficient*, and controls how optimal (vs. random) the agent is.\n",
    "When $\\beta$ is higher, the policy assigns higher probabilities to actions\n",
    "with higher expected utilities. $A(s_t)$ is the set of possible actions\n",
    "at state $s_t$.\n",
    "\n",
    "### Writing an agent model probabilistic program\n",
    "\n",
    "We can now specify our generative model of an agent as a probabilistic program,\n",
    "which includes each of the components described above (setting aside for now\n",
    "how Q-values are computed):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Gen\n",
    "using Gen.Distributions: softmax\n",
    "\n",
    "\"Uniform distribution over a list of labels.\"\n",
    "@dist labeled_uniform(labels) = labels[uniform_discrete(1, length(labels))]\n",
    "\n",
    "\"Categorical distribution over a list of labels.\"\n",
    "@dist labeled_categorical(labels, probs) = labels[categorical(probs)]\n",
    "\n",
    "\"Return the next gridworld environment state given the current state and action.\"\n",
    "function gridworld_env_step(grid, loc, act)\n",
    "    # Return next location if in-bounds and not a wall\n",
    "    next_loc = loc .+ act\n",
    "    if checkbounds(Bool, grid, next_loc...) && !grid[next_loc...]\n",
    "        return next_loc\n",
    "    else\n",
    "        return loc\n",
    "    end\n",
    "end\n",
    "\n",
    "\"Boltzmann-rational agent model in a gridworld environment.\"\n",
    "@gen function gridworld_agent_model(T, grid, init_loc, goal_locs, beta)\n",
    "    # Sample goal from uniform goal prior P(g)\n",
    "    goal = {:goal} ~ labeled_uniform(goal_locs)\n",
    "    loc = init_loc\n",
    "    loc_history = [loc]\n",
    "    act_history = Tuple{Int, Int}[]\n",
    "    # Iterate over timesteps\n",
    "    for t in 1:T\n",
    "        # Compute Q-values and Boltzmann action probabilities\n",
    "        actions, q_values = compute_q_values(grid, loc, goal)\n",
    "        act_probs = softmax(beta .* q_values)\n",
    "        # Sample action according to Boltzmann policy P(a_t | s_t, g)\n",
    "        act = {(:act, t)} ~ labeled_categorical(actions, act_probs)\n",
    "        # Update location following environment dynamics P(s_{t+1} | s_t, a_t)\n",
    "        loc = gridworld_env_step(grid, loc, act)\n",
    "        push!(loc_history, loc)\n",
    "        push!(act_history, act)\n",
    "    end\n",
    "    return goal, loc_history, act_history\n",
    "end;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Probabilistic programs are programs that make random choices during their\n",
    "execution, allowing us to specify generative models as code. When using\n",
    "Gen.jl, probabilistic programs look almost exactly like regular Julia functions,\n",
    "except they are annotated with the `@gen` macro at the start, and make random\n",
    "choices using the `~` operator. The `gridworld_agent_model` above makes random\n",
    "choices at two places: first, when sampling a goal $g$ frm a uniform prior\n",
    "$P(g)$ over goals:\n",
    "\n",
    "```julia\n",
    "goal = {:goal} ~ labeled_uniform(goal_locs)\n",
    "```\n",
    "\n",
    "and second, when sampling an action $a_t$ from a Boltzmann-rational policy\n",
    "$P(a_t | s_t, g)$:\n",
    "\n",
    "```julia\n",
    "act = {(:act, t)} ~ labeled_categorical(actions, act_probs)\n",
    "```\n",
    "\n",
    "Each of these random choices is associated with an *address*, which gives a\n",
    "unique name to the random variable being sampled. The result of the random\n",
    "choice can optionally be assigned to regular Julia variable. For example, when\n",
    "sampling the agent's action at step `t` of the for-loop, the sampled action\n",
    "is given the address `(:act, t)`, and is assigned to the local variable `act`.\n",
    "\n",
    "### Computing Q-values via planning\n",
    "\n",
    "To complete our probabilistic program, we need to specify how to compute the\n",
    "Q-values of each action via the `compute_q_values` function. To do this, let's\n",
    "first think about what $Q_g(s_t, a_t)$ represents: What is the expected utility\n",
    "of taking action $a_t$ at state $s_t$ when the agent's goal is $g$?\n",
    "\n",
    "In our simple gridworld, we assume that the agent wants to reach their goal $g$\n",
    "while minimizing the cost of the actions it takes to get there. So the expected\n",
    "utility of taking action $a_t$ is just the (negative) cost of action $a_t$\n",
    "itself, plus the expected utility of reaching $g$ from the next state $s_{t+1}$\n",
    "(denoted $V_g(s_{t+1})$). Recursively, we define $V_g(s_t)$ — also known as the\n",
    "*value function* or *cost-to-go function* — as the maximum expected utility\n",
    "(i.e. minimum expected cumulative cost) of taking any action from state $s_t$:\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "Q_g(s_t, a_t) &= -C(a_t) + V_g(s_{t+1}) \\\\\n",
    "V_g(s_t) &= \\max_a Q_g(s_t, a)\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "Let's assume that cost of any action $C(a_t)$ is 1. Together with the fact that\n",
    "our environment is deterministic, these definitions imply that $V_g(s_t)$ is\n",
    "just the *length of the shortest path* from $s_t$ to $g$ (but negated).\n",
    "Correspondingly, $Q_g(s_t, a_t)$ is the length of the shortest path from $s_t$\n",
    "to $g$ that involves first taking action $a_t$. This means we can compute\n",
    "$Q_g(s_t, a_t)$ for each action $a_t$ via a path planning algorithm. In our\n",
    "simple gridworld, it is enough to use [Dijkstra's algorithm](https://en.wikipedia.org/wiki/Dijkstra%27s_algorithm)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"Compute Q-values for each action given the current location and goal location.\"\n",
    "function compute_q_values(grid, loc, goal)\n",
    "    # Define list of possible actions (stay, up, down, left, right)\n",
    "    actions = ((0, 0), (-1, 0), (1, 0), (0, -1), (0, 1))\n",
    "    # Compute shortest path from each successor location to goal\n",
    "    q_values = zeros(length(actions))\n",
    "    for i in 1:length(actions)\n",
    "        next_loc = loc .+ actions[i]\n",
    "        if checkbounds(Bool, grid, next_loc...) && !grid[next_loc...]\n",
    "            # Run Dijkstra's algorithm from next_loc to goal\n",
    "            q_values[i] = -shortest_path_length(grid, next_loc, goal) - 1\n",
    "        else\n",
    "            # Negative infinite Q-value if next_loc is unreachable\n",
    "            q_values[i] = -Inf\n",
    "        end\n",
    "    end\n",
    "    return actions, q_values\n",
    "end\n",
    "\n",
    "\"Compute the shortest path length from start to goal via Dijkstra's algorithm.\"\n",
    "function shortest_path_length(grid, start, goal)\n",
    "    # Initialize distances to infinity\n",
    "    distances = fill(Inf, size(grid))\n",
    "    distances[start...] = 0\n",
    "    # Initialize priority queue\n",
    "    queue = [(0, start)]\n",
    "    while !isempty(queue)\n",
    "        dist, loc = popfirst!(queue)\n",
    "        if loc == goal\n",
    "            return dist\n",
    "        end\n",
    "        for d in ((-1, 0), (1, 0), (0, -1), (0, 1))\n",
    "            next_loc = loc .+ d\n",
    "            if checkbounds(Bool, grid, next_loc...) && !grid[next_loc...]\n",
    "                alt = dist + 1\n",
    "                if alt < distances[next_loc...]\n",
    "                    distances[next_loc...] = alt\n",
    "                    push!(queue, (alt, next_loc))\n",
    "                end\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "    return Inf\n",
    "end;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Question**: The approach above is not the most efficient way of computing\n",
    "Q-values, since it requires invoking Dijkstra's algorithm for each action.\n",
    "Can you think of an approach that requires invoking Dijkstra's algorithm only\n",
    "once for all actions?\n",
    "\n",
    "## Simulating and visualizing our agent model\n",
    "\n",
    "With our model fully defined, we can now run `gridworld_agent_model` with\n",
    "specific arguments to generate a possible trajectory for our agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "goal = (1, 5)\n",
      "loc_history = [(5, 3), (4, 3), (3, 3), (3, 4), (2, 4), (1, 4), (1, 5), (1, 5), (1, 5), (1, 5), (1, 5)]\n"
     ]
    }
   ],
   "source": [
    "grid = BitMatrix([ # 2D grid where 1s are walls\n",
    "    0 0 0 0 0\n",
    "    0 1 1 0 1\n",
    "    0 0 0 0 0\n",
    "    0 1 0 1 0\n",
    "    0 1 0 1 0\n",
    "])\n",
    "init_loc = (5, 3) # Initial location at row 5, column 3\n",
    "goal_locs = ((1, 2), (1, 5), (5, 1)) # 3 possible goal locations\n",
    "beta = 10.0 # Assume the agent is highly optimal\n",
    "T = 10 # Run the model for 10 timesteps\n",
    "goal, loc_history, act_history =\n",
    "    gridworld_agent_model(T, grid, init_loc, goal_locs, beta)\n",
    "@show goal;\n",
    "@show loc_history;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the agent starts from its initial location, then navigates the\n",
    "gridworld towards one of the three possible goals. The path the agent takes is\n",
    "highly optimal, due to `beta` being set to a high value. Once the agent reaches\n",
    "the goal, it stays in place.\n",
    "\n",
    "### Generating traces of probabilistic programs\n",
    "\n",
    "A key benefit of writing our model in a probabilistic programming language is\n",
    "that we can do more than just simulate executions of our model and collect\n",
    "their return values (as shown above). We can also *evaluate* the probability of\n",
    "a specific simulation, and run *constrained* simulations, where certain random\n",
    "variables are fixed to specific values. This can be done with the following\n",
    "functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "│\n",
      "├── (:act, 1) : (-1, 0)\n",
      "│\n",
      "├── (:act, 4) : (0, -1)\n",
      "│\n",
      "├── (:act, 3) : (0, -1)\n",
      "│\n",
      "├── (:act, 2) : (-1, 0)\n",
      "│\n",
      "├── :goal : (1, 2)\n",
      "│\n",
      "└── (:act, 5) : (-1, 0)\n",
      "logprob = -2.437251221207088\n",
      "\n",
      "Traced choices given constrained goal:\n",
      "│\n",
      "├── (:act, 1) : (-1, 0)\n",
      "│\n",
      "├── (:act, 4) : (0, -1)\n",
      "│\n",
      "├── (:act, 3) : (0, -1)\n",
      "│\n",
      "├── (:act, 2) : (-1, 0)\n",
      "│\n",
      "├── :goal : (1, 2)\n",
      "│\n",
      "└── (:act, 5) : (-1, 0)\n",
      "weight = -1.0986122886681098\n",
      "\n",
      "Traced choices given constrained actions:\n",
      "│\n",
      "├── (:act, 1) : (-1, 0)\n",
      "│\n",
      "├── (:act, 4) : (-1, 0)\n",
      "│\n",
      "├── (:act, 3) : (0, 1)\n",
      "│\n",
      "├── (:act, 2) : (-1, 0)\n",
      "│\n",
      "├── :goal : (1, 5)\n",
      "│\n",
      "└── (:act, 5) : (-1, 0)\n",
      "weight = -0.7301575864224461\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "-0.7301575864224461"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Simulate an execution trace from our agent model\n",
    "model_args = (5, grid, init_loc, goal_locs, 2.0)\n",
    "trace = Gen.simulate(gridworld_agent_model, model_args)\n",
    "show(stdout, \"text/plain\", get_choices(trace)) # Show the choices in this trace\n",
    "\n",
    "# Evaluate the log-probability (a.k.a. score) of the sampled trace\n",
    "logprob = Gen.get_score(trace)\n",
    "@show logprob\n",
    "\n",
    "# Generate a trace with the goal constrained to a specific value\n",
    "goal_constraint = choicemap((:goal, (1, 2)))\n",
    "trace, weight = Gen.generate(gridworld_agent_model, model_args, goal_constraint)\n",
    "println(\"\\n\", \"Traced choices given constrained goal:\")\n",
    "show(stdout, \"text/plain\", get_choices(trace)) # Goal is fixed, other choices are sampled\n",
    "@show weight # Log-prob. of constrained choices given ancestors (= log P(g))\n",
    "\n",
    "# Generate trace with actions constrained to observed values\n",
    "act_constraints = choicemap(\n",
    "    ((:act, 1), (-1, 0)), ((:act, 2), (-1, 0)), ((:act, 3), (0, 1)),\n",
    "    ((:act, 4), (-1, 0)), ((:act, 5), (-1, 0))\n",
    ")\n",
    "trace, weight = Gen.generate(gridworld_agent_model, model_args, act_constraints)\n",
    "println(\"\\n\", \"Traced choices given constrained actions:\")\n",
    "show(stdout, \"text/plain\", get_choices(trace))  # Actions are fixed, other choices are sampled\n",
    "@show weight # Log-prob. of constrained choices given ancestors (= log P(a_{1:T} | g))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `simulate`, `get_score`, and `generate` functions are basic primitives that\n",
    "can be used to implement a wide variety of inference algorithms (MCMC,\n",
    "particle filtering, etc.). All of them operate over **traces**: a representation\n",
    "of an execution trace of a probabilistic program, which records all the random\n",
    "choices sampled along the way (exposed by `get_choices`), and also the joint\n",
    "probability of those choices (exposed by `get_score`).\n",
    "\n",
    "Most users of Gen can make use of higher-level interfaces that call these\n",
    "functions under-the-hood, but being familiar with concept of a trace is crucial\n",
    "to using Gen. We'll also see later how the `Gen.generate` function can be used\n",
    "to implement an inference algorithm via constrained generation and weighting\n",
    "of traces.\n",
    "\n",
    "### Visually inspecting our model\n",
    "\n",
    "While it is possible work with and inspect traces entirely in code, it is good\n",
    "practice to visualize our model and the traces it generates. Visualization\n",
    "can surface modeling errors that are otherwise hard to notice, whether at the\n",
    "level of individual traces, or in terms of the overall distribution of traces\n",
    "sampled from the model. Below we provide code to visualize our agent model\n",
    "using the [Makie](https://docs.makie.org/stable/) plotting libary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "anim_traces"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "using GLMakie\n",
    "using Printf\n",
    "using FileIO\n",
    "\n",
    "GLMakie.activate!(scalefactor=1.0)\n",
    "Makie.inline!(true)\n",
    "\n",
    "include(\"videodisplay.jl\")\n",
    "\n",
    "AGENT_IMG = rotr90(FileIO.load(joinpath(@__DIR__, \"images\", \"human.png\")))\n",
    "GOAL_COLORS = [\n",
    "    colorant\"#D41159\", colorant\"#FFC20A\", colorant\"#1A85FF\", colorant\"#007D68\",\n",
    "    colorant\"#785EF0\", colorant\"#D55E00\", colorant\"#56B4E9\", colorant\"#CC79A7\"\n",
    "]\n",
    "\n",
    "\"Render a trace of the gridworld agent model to an existing axis.\"\n",
    "function render_trace!(ax, trace::Union{Trace, Observable}, t = nothing)\n",
    "    # Wrap trace in Observable to enable interactive plotting\n",
    "    trace = trace isa Observable ? trace : Observable(trace)\n",
    "    t_obs = isnothing(t) ?\n",
    "        @lift(Gen.get_args($trace)[1]) :\n",
    "        t isa Observable ? t : Observable(t)\n",
    "    # Plot heatmap showing grid\n",
    "    grid_obs = @lift reverse(transpose(Gen.get_args($trace)[2]), dims=2)\n",
    "    cmap = cgrad([:transparent, :black])\n",
    "    heatmap!(ax, grid_obs, colormap=cmap, colorrange=(0, 1))\n",
    "    hidedecorations!(ax, grid=false)\n",
    "    map!(h -> (1:h-1) .+ 0.5, ax.yticks, @lift(size($grid_obs)[1]))\n",
    "    map!(w -> (1:w-1) .+ 0.5, ax.xticks, @lift(size($grid_obs)[2]))\n",
    "    # Render goal locations\n",
    "    goal_locs_obs = Observable([(1, 1)])\n",
    "    goal_colors = Observable(GOAL_COLORS[1:1])\n",
    "    scatter!(ax, goal_locs_obs, color=goal_colors, marker='♦',\n",
    "             markerspace=:data, markersize=1.0)\n",
    "    # Render agent trajectory\n",
    "    loc_history_obs = Observable([(1, 1)])\n",
    "    scatter!(ax, loc_history_obs, color=:black, marker='⦿',\n",
    "             markerspace=:data, markersize=0.4)\n",
    "    # Render current agent location\n",
    "    agent_x_bounds_obs = Observable((0.5, 1.5))\n",
    "    agent_y_bounds_obs = Observable((0.5, 1.5))\n",
    "    image!(ax, agent_x_bounds_obs, agent_y_bounds_obs, AGENT_IMG)\n",
    "    # Update observables with values from trace\n",
    "    onany(trace, t_obs) do trace, t\n",
    "        # Extract argument and return values\n",
    "        T, grid, init_loc, goal_locs, beta = Gen.get_args(trace)\n",
    "        retvals = Gen.get_retval(trace)\n",
    "        if length(retvals) == 2 # Handle Unfold-model variants\n",
    "            goal, model_states = retvals\n",
    "            loc_history = [s.loc for s in model_states]\n",
    "            pushfirst!(loc_history, init_loc)\n",
    "        else\n",
    "            goal, loc_history, act_history = retvals\n",
    "        end\n",
    "        h, w = size(grid)\n",
    "        # Clamp rendered timestep to [0, T]\n",
    "        t = max(0, min(T, t))\n",
    "        # Update goal locations\n",
    "        empty!(goal_locs_obs[])\n",
    "        append!(goal_locs_obs[], [(x, h-y+1) for (y, x) in goal_locs])\n",
    "        empty!(goal_colors[])\n",
    "        append!(goal_colors[], GOAL_COLORS[1:length(goal_locs)])\n",
    "        notify(goal_locs_obs)\n",
    "        notify(goal_colors)\n",
    "        # Update location history (flipping y and x coordinates)\n",
    "        empty!(loc_history_obs[])\n",
    "        append!(loc_history_obs[], [(x, h-y+1) for (y, x) in loc_history[1:t+1]])\n",
    "        notify(loc_history_obs)\n",
    "        # Update agent location bounds\n",
    "        cur_loc = loc_history[t+1]\n",
    "        agent_x_bounds_obs[] = (cur_loc[2] - 0.5, cur_loc[2] + 0.5)\n",
    "        agent_y_bounds_obs[] = (h - cur_loc[1] + 0.5, h - cur_loc[1] + 1.5)\n",
    "    end\n",
    "    notify(trace)\n",
    "    return ax\n",
    "end\n",
    "\n",
    "\"Render a trace of the gridworld agent model to a new figure.\"\n",
    "function render_trace(trace::Union{Trace, Observable}, t = nothing;\n",
    "                      size = (512, 512))\n",
    "    figure = Figure(size=size)\n",
    "    ax = Axis(figure[1, 1], aspect=DataAspect())\n",
    "    render_trace!(ax, trace, t)\n",
    "    return figure\n",
    "end\n",
    "\n",
    "\"Animate a trace of the gridworld agent model over time.\"\n",
    "function anim_trace(trace; size = (512, 512), kwargs...)\n",
    "    T = trace isa Observable ? Gen.get_args(trace[])[1] : Gen.get_args(trace)[1]\n",
    "    t_obs = Observable(0)\n",
    "    figure = render_trace(trace, t_obs; size=size)\n",
    "    anim = Record(figure, 0:T; kwargs...) do t\n",
    "        t_obs[] = t\n",
    "    end\n",
    "    return VideoDisplay(anim)\n",
    "end\n",
    "\n",
    "\"Render a collection of traces as a grid of subplots in a new figure.\"\n",
    "function render_traces(\n",
    "    traces::AbstractVector, ts = nothing;\n",
    "    n_rows = ceil(Int, sqrt(length(traces))),\n",
    "    n_cols = ceil(Int, length(traces) / n_rows),\n",
    "    size = (192 * n_cols, 192 * n_rows)\n",
    ")\n",
    "    figure = Figure(size=size)\n",
    "    for n in 1:length(traces)\n",
    "        i = div(n - 1, n_cols) + 1\n",
    "        j = mod(n - 1, n_cols) + 1\n",
    "        ax = Axis(figure[i, j], aspect=DataAspect())\n",
    "        t = ts isa AbstractArray ? ts[n] : ts\n",
    "        render_trace!(ax, traces[n], t)\n",
    "    end\n",
    "    return figure\n",
    "end\n",
    "\n",
    "\"Animate a collection of traces over time.\"\n",
    "function anim_traces(\n",
    "    traces::AbstractVector;\n",
    "    n_rows = ceil(Int, sqrt(length(traces))),\n",
    "    n_cols = ceil(Int, length(traces) / n_rows),\n",
    "    size = (192 * n_cols, 192 * n_rows),\n",
    "    kwargs...\n",
    ")\n",
    "    T = maximum(traces) do trace\n",
    "        trace isa Observable ? Gen.get_args(trace[])[1] : Gen.get_args(trace)[1]\n",
    "    end\n",
    "    t_obs = Observable(0)\n",
    "    figure = render_traces(traces, t_obs; n_rows, n_cols, size)\n",
    "    anim = Record(figure, 0:T; kwargs...) do t\n",
    "        t_obs[] = t\n",
    "    end\n",
    "    return VideoDisplay(anim)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using this code, we can now visualize a single trace from our model:\n",
    "Specify arguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 25696 characters]",
      "text/html": [
       "<img width=512 height=512 style='object-fit: contain; height: auto;' src=\"data:image/png;base64, [encoded data, 25696 characters]\"/>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "grid = BitMatrix([ # 2D grid where 1s are walls\n",
    "    0 0 0 0 0 0 0 0\n",
    "    0 1 1 0 1 1 1 0\n",
    "    0 0 0 0 0 1 0 0\n",
    "    1 0 1 1 0 0 0 0\n",
    "    0 0 0 1 0 1 1 0\n",
    "    0 1 0 1 0 0 0 0\n",
    "    0 1 0 1 0 1 0 0\n",
    "    0 0 0 1 0 1 0 1\n",
    "])\n",
    "init_loc = (8, 7) # Initial location at row 5, column 3\n",
    "goal_locs = ((1, 1), (1, 5), (8, 2)) # 3 possible goal locations\n",
    "beta = 1.0 # Assume the agent is noisily rational\n",
    "T = 20 # Run the model for 20 timesteps\n",
    "model_args = (T, grid, init_loc, goal_locs, beta)\n",
    "\n",
    "# Simulate trace\n",
    "trace = Gen.simulate(gridworld_agent_model, model_args)\n",
    "\n",
    "# Visualize trace\n",
    "render_trace(trace)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also animate the trace over time (this allows us to more easily see\n",
    "actions when the agent randomly backtracks):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<video controls autoplay muted>\n",
       "    <source src=\"data:video/mp4;base64,[encoded data, 28772 characters]\"type=\"video/mp4\">\n",
       "</video>"
      ],
      "text/plain": [
       "VideoDisplay(VideoStream(Base.PipeEndpoint(Base.Libc.WindowsRawSocket(0x0000000000001974) open, 0 bytes waiting), Process(setenv(`\u001b[4m'C:\\Users\\Xuan\\.julia\\artifacts\\e53c880b3a155c54b452cbe94b72de51158aea2e\\bin\\ffmpeg.exe'\u001b[24m \u001b[4m-y\u001b[24m \u001b[4m-loglevel\u001b[24m \u001b[4mquiet\u001b[24m \u001b[4m-threads\u001b[24m \u001b[4m20\u001b[24m \u001b[4m-framerate\u001b[24m \u001b[4m3\u001b[24m \u001b[4m-pixel_format\u001b[24m \u001b[4mrgb24\u001b[24m \u001b[4m-f\u001b[24m \u001b[4mrawvideo\u001b[24m \u001b[4m-s:v\u001b[24m \u001b[4m512x512\u001b[24m \u001b[4m-r\u001b[24m \u001b[4m3\u001b[24m \u001b[4m-i\u001b[24m \u001b[4mpipe:0\u001b[24m \u001b[4m-vf\u001b[24m \u001b[4mvflip\u001b[24m \u001b[4m-profile:v\u001b[24m \u001b[4mhigh422\u001b[24m \u001b[4m-crf\u001b[24m \u001b[4m20\u001b[24m \u001b[4m-preset\u001b[24m \u001b[4mslow\u001b[24m \u001b[4m-c:v\u001b[24m \u001b[4mlibx264\u001b[24m \u001b[4m-pix_fmt\u001b[24m \u001b[4myuv420p\u001b[24m \u001b[4m-an\u001b[24m \u001b[4m'C:\\Users\\Xuan\\AppData\\Local\\Temp\\jl_dRoaeL\\##video#247.mp4'\u001b[24m`,[\"WINDIR=C:\\\\WINDOWS\", \"PATH=C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\3af036b12a076c2d4c042dde6d956ef4eac2c2ce\\\\bin;C:\\\\Program Files\\\\Julia\\\\1.10.2\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\13e3cf8e6f38d3b6418b758bb68cb8b149c28251\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\ce4261864aeaa78b205a3f6fb750be4c4afe9b4b\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\1da3b3a2698813adf2b925c46aa2d68ec8474449\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\58a307887a7b823f6a7cac7e6337710ec14aced2\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\df4cc9189a878ad66ad99fed70ece5cf4e45b246\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\c3147644202f8802c6131410ffd73f7e41db31f8\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\fcc05863e77d0e700007d714ef7304c87b39d10e\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\166f431d14e588224598f7952888bb847a545ac6\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\51e4c0b5a201939ab2352553b85f5d0932d5cf03\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\cb74fa09f359d430ea01be69a61cd98b50ec8c22\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\6e310796647a1c58e5873a7ab7ad32b8e60dea3f\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\d119ba6f3cc7e3fb09a8bb637f705bf9d9982000\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\ec3406e9ba84281719f20b0fb3a874c9e3bac027\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\8d9f87bf9e08a0e0db47c0ab531ee29fff143eee\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\41cbcbec1d1b6ca0484999b1db571f8066f29233\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\5c7df64500371e5ab4723fe2fc999b996aa8f3a7\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\81ad3b6f8af8fa057c1a9640ca91454b669b342b\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\8258f1ad1bf2f9853eb157e9d4b755f0e427a0ea\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\25bfd4c7adfcd3cb3399a7c509ec239a32c247f3\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\45c659a2370bbaa5bee3c0ed364427d4c33552db\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\98941fbb5eed27cef1d77fa8710290bc197dd24f\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\ba2dc75245508e0cf0c860081b2f2fd10bdf4357\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\a2b7bddf24b04fe5cf3e6d40e0527451b3439d23\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\7e207f56275705f9cbc88b826acd28ea0e40f1ae\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\e53c880b3a155c54b452cbe94b72de51158aea2e\\\\bin;C:\\\\Program Files\\\\Julia\\\\1.10.2\\\\bin\\\\..\\\\lib\\\\julia;C:\\\\Program Files\\\\Julia\\\\1.10.2\\\\bin\\\\..\\\\lib;C:\\\\Program Files\\\\Julia\\\\1.10.2\\\\bin;C:\\\\WINDOWS\\\\system32;C:\\\\WINDOWS;C:\\\\WINDOWS\\\\System32\\\\Wbem;C:\\\\WINDOWS\\\\System32\\\\WindowsPowerShell\\\\v1.0\\\\;C:\\\\WINDOWS\\\\System32\\\\OpenSSH\\\\;C:\\\\Program Files\\\\Microsoft VS Code\\\\bin;C:\\\\Program Files\\\\Git\\\\cmd;C:\\\\Program Files\\\\GitHub CLI\\\\;C:\\\\ProgramData\\\\chocolatey\\\\bin;C:\\\\Program Files\\\\Rust\\\\Rust-MSVC-1.82\\\\bin;C:\\\\Program Files\\\\Git LFS;C:\\\\Program Files\\\\Docker\\\\Docker\\\\resources\\\\bin;C:\\\\Users\\\\Xuan\\\\AppData\\\\Local\\\\Microsoft\\\\WindowsApps;C:\\\\Users\\\\Xuan\\\\AppData\\\\Local\\\\Programs\\\\MiKTeX\\\\miktex\\\\bin\\\\x64\\\\\", \"ELECTRON_RUN_AS_NODE=1\", \"USERDOMAIN_ROAMINGPROFILE=deep-titanium\", \"VSCODE_CODE_CACHE_PATH=C:\\\\Users\\\\Xuan\\\\AppData\\\\Roaming\\\\Code\\\\CachedData\\\\6f17636121051a53c88d3e605c491d22af2ba755\", \"ZES_ENABLE_SYSMAN=1\", \"LOCALAPPDATA=C:\\\\Users\\\\Xuan\\\\AppData\\\\Local\", \"HOMEPATH=\\\\Users\\\\Xuan\", \"VSCODE_NLS_CONFIG={\\\"userLocale\\\":\\\"en-us\\\",\\\"osLocale\\\":\\\"en-us\\\",\\\"resolvedLanguage\\\":\\\"en\\\",\\\"defaultMessagesFile\\\":\\\"C:\\\\\\\\Program Files\\\\\\\\Microsoft VS Code\\\\\\\\resources\\\\\\\\app\\\\\\\\out\\\\\\\\nls.messages.json\\\",\\\"locale\\\":\\\"en-us\\\",\\\"availableLanguages\\\":{}}\", \"PROCESSOR_IDENTIFIER=Intel64 Family 6 Model 186 Stepping 2, GenuineIntel\"  …  \"PROCESSOR_LEVEL=6\", \"EFC_7788_1592913036=1\", \"SYSTEMDRIVE=C:\", \"FPS_BROWSER_APP_PROFILE_STRING=Internet Explorer\", \"PROGRAMW6432=C:\\\\Program Files\", \"IGCCSVC_DB=AQAAANCMnd8BFdERjHoAwE/Cl+sBAAAAon9/u+k2YU+MR1UcImxbYAQAAAACAAAAAAAQZgAAAAEAACAAAAAVjZ9/NJxM1QJQJiNVkBcoSz9zpj/o0TIuggnp8vj0WQAAAAAOgAAAAAIAACAAAAApn7bxsL9jtBfibrJodoROHkEqcLS+9QzQAxmjKrM8pGAAAACkrW/4Uue8yhXVzWjF3YzLIlBgbO+7qPSD/vvkQwwOzTB8l2Hgp+TtRmZNO/uR2d63ivHhQYHH1g+32Yk4ZTNIijbj1aqe+ZCc/sMGnjUhCMoVlAYncUAKnlLvYPxZcTlAAAAANd4a+ne/noB73U12+A8hgkEQz97pmZXATj2GXQZGTv/oTJiXVLjUc4z6Fc/hLPv6NYVAIcPo3XuS6ZmL4W2oyQ==\", \"HOMEDRIVE=C:\", \"OPENBLAS_MAIN_FREE=1\", \"PROCESSOR_ARCHITECTURE=AMD64\", \"TEMP=C:\\\\Users\\\\Xuan\\\\AppData\\\\Local\\\\Temp\"]), ProcessRunning), GLMakie.Screen(...), Makie.TickController(Observable(Makie.Tick(Makie.OneTimeRenderTick, 21, 7.0, 0.3333333333333333)), 22, 0.3333333333333333, true, ObserverFunction defined at C:\\Users\\Xuan\\.julia\\packages\\Makie\\4JW9B\\src\\ffmpeg-util.jl:199 operating on Observable(Tick(OneTimeRenderTick, 21, 7.0, 0.333333))), ColorTypes.RGB{FixedPointNumbers.N0f8}[RGB(0.6, 0.416, 0.38) RGB(0.6, 0.416, 0.38) … RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0); RGB(0.6, 0.004, 0.0) RGB(0.6, 0.004, 0.0) … RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0); … ; RGB(0.38, 0.6, 0.004) RGB(0.38, 0.6, 0.004) … RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0); RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0) … RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0)], \"C:\\\\Users\\\\Xuan\\\\AppData\\\\Local\\\\Temp\\\\jl_dRoaeL\\\\##video#247.mp4\", Makie.VideoStreamOptions(\"mp4\", 3, 20, \"high422\", \"yuv420p\", 0, \"quiet\", \"pipe:0\", true)))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "anim = anim_trace(trace, framerate=3, format=\"mp4\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get a sense of the distribution over possible traces, we can also visualize\n",
    "multiple traces at once:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<video controls autoplay muted>\n",
       "    <source src=\"data:video/mp4;base64,[encoded data, 84164 characters]\"type=\"video/mp4\">\n",
       "</video>"
      ],
      "text/plain": [
       "VideoDisplay(VideoStream(Base.PipeEndpoint(Base.Libc.WindowsRawSocket(0x00000000000015bc) open, 0 bytes waiting), Process(setenv(`\u001b[4m'C:\\Users\\Xuan\\.julia\\artifacts\\e53c880b3a155c54b452cbe94b72de51158aea2e\\bin\\ffmpeg.exe'\u001b[24m \u001b[4m-y\u001b[24m \u001b[4m-loglevel\u001b[24m \u001b[4mquiet\u001b[24m \u001b[4m-threads\u001b[24m \u001b[4m20\u001b[24m \u001b[4m-framerate\u001b[24m \u001b[4m3\u001b[24m \u001b[4m-pixel_format\u001b[24m \u001b[4mrgb24\u001b[24m \u001b[4m-f\u001b[24m \u001b[4mrawvideo\u001b[24m \u001b[4m-s:v\u001b[24m \u001b[4m576x576\u001b[24m \u001b[4m-r\u001b[24m \u001b[4m3\u001b[24m \u001b[4m-i\u001b[24m \u001b[4mpipe:0\u001b[24m \u001b[4m-vf\u001b[24m \u001b[4mvflip\u001b[24m \u001b[4m-profile:v\u001b[24m \u001b[4mhigh422\u001b[24m \u001b[4m-crf\u001b[24m \u001b[4m20\u001b[24m \u001b[4m-preset\u001b[24m \u001b[4mslow\u001b[24m \u001b[4m-c:v\u001b[24m \u001b[4mlibx264\u001b[24m \u001b[4m-pix_fmt\u001b[24m \u001b[4myuv420p\u001b[24m \u001b[4m-an\u001b[24m \u001b[4m'C:\\Users\\Xuan\\AppData\\Local\\Temp\\jl_HzipJD\\##video#249.mp4'\u001b[24m`,[\"WINDIR=C:\\\\WINDOWS\", \"PATH=C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\3af036b12a076c2d4c042dde6d956ef4eac2c2ce\\\\bin;C:\\\\Program Files\\\\Julia\\\\1.10.2\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\13e3cf8e6f38d3b6418b758bb68cb8b149c28251\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\ce4261864aeaa78b205a3f6fb750be4c4afe9b4b\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\1da3b3a2698813adf2b925c46aa2d68ec8474449\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\58a307887a7b823f6a7cac7e6337710ec14aced2\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\df4cc9189a878ad66ad99fed70ece5cf4e45b246\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\c3147644202f8802c6131410ffd73f7e41db31f8\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\fcc05863e77d0e700007d714ef7304c87b39d10e\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\166f431d14e588224598f7952888bb847a545ac6\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\51e4c0b5a201939ab2352553b85f5d0932d5cf03\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\cb74fa09f359d430ea01be69a61cd98b50ec8c22\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\6e310796647a1c58e5873a7ab7ad32b8e60dea3f\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\d119ba6f3cc7e3fb09a8bb637f705bf9d9982000\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\ec3406e9ba84281719f20b0fb3a874c9e3bac027\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\8d9f87bf9e08a0e0db47c0ab531ee29fff143eee\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\41cbcbec1d1b6ca0484999b1db571f8066f29233\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\5c7df64500371e5ab4723fe2fc999b996aa8f3a7\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\81ad3b6f8af8fa057c1a9640ca91454b669b342b\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\8258f1ad1bf2f9853eb157e9d4b755f0e427a0ea\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\25bfd4c7adfcd3cb3399a7c509ec239a32c247f3\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\45c659a2370bbaa5bee3c0ed364427d4c33552db\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\98941fbb5eed27cef1d77fa8710290bc197dd24f\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\ba2dc75245508e0cf0c860081b2f2fd10bdf4357\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\a2b7bddf24b04fe5cf3e6d40e0527451b3439d23\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\7e207f56275705f9cbc88b826acd28ea0e40f1ae\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\e53c880b3a155c54b452cbe94b72de51158aea2e\\\\bin;C:\\\\Program Files\\\\Julia\\\\1.10.2\\\\bin\\\\..\\\\lib\\\\julia;C:\\\\Program Files\\\\Julia\\\\1.10.2\\\\bin\\\\..\\\\lib;C:\\\\Program Files\\\\Julia\\\\1.10.2\\\\bin;C:\\\\WINDOWS\\\\system32;C:\\\\WINDOWS;C:\\\\WINDOWS\\\\System32\\\\Wbem;C:\\\\WINDOWS\\\\System32\\\\WindowsPowerShell\\\\v1.0\\\\;C:\\\\WINDOWS\\\\System32\\\\OpenSSH\\\\;C:\\\\Program Files\\\\Microsoft VS Code\\\\bin;C:\\\\Program Files\\\\Git\\\\cmd;C:\\\\Program Files\\\\GitHub CLI\\\\;C:\\\\ProgramData\\\\chocolatey\\\\bin;C:\\\\Program Files\\\\Rust\\\\Rust-MSVC-1.82\\\\bin;C:\\\\Program Files\\\\Git LFS;C:\\\\Program Files\\\\Docker\\\\Docker\\\\resources\\\\bin;C:\\\\Users\\\\Xuan\\\\AppData\\\\Local\\\\Microsoft\\\\WindowsApps;C:\\\\Users\\\\Xuan\\\\AppData\\\\Local\\\\Programs\\\\MiKTeX\\\\miktex\\\\bin\\\\x64\\\\\", \"ELECTRON_RUN_AS_NODE=1\", \"USERDOMAIN_ROAMINGPROFILE=deep-titanium\", \"VSCODE_CODE_CACHE_PATH=C:\\\\Users\\\\Xuan\\\\AppData\\\\Roaming\\\\Code\\\\CachedData\\\\6f17636121051a53c88d3e605c491d22af2ba755\", \"ZES_ENABLE_SYSMAN=1\", \"LOCALAPPDATA=C:\\\\Users\\\\Xuan\\\\AppData\\\\Local\", \"HOMEPATH=\\\\Users\\\\Xuan\", \"VSCODE_NLS_CONFIG={\\\"userLocale\\\":\\\"en-us\\\",\\\"osLocale\\\":\\\"en-us\\\",\\\"resolvedLanguage\\\":\\\"en\\\",\\\"defaultMessagesFile\\\":\\\"C:\\\\\\\\Program Files\\\\\\\\Microsoft VS Code\\\\\\\\resources\\\\\\\\app\\\\\\\\out\\\\\\\\nls.messages.json\\\",\\\"locale\\\":\\\"en-us\\\",\\\"availableLanguages\\\":{}}\", \"PROCESSOR_IDENTIFIER=Intel64 Family 6 Model 186 Stepping 2, GenuineIntel\"  …  \"PROCESSOR_LEVEL=6\", \"EFC_7788_1592913036=1\", \"SYSTEMDRIVE=C:\", \"FPS_BROWSER_APP_PROFILE_STRING=Internet Explorer\", \"PROGRAMW6432=C:\\\\Program Files\", \"IGCCSVC_DB=AQAAANCMnd8BFdERjHoAwE/Cl+sBAAAAon9/u+k2YU+MR1UcImxbYAQAAAACAAAAAAAQZgAAAAEAACAAAAAVjZ9/NJxM1QJQJiNVkBcoSz9zpj/o0TIuggnp8vj0WQAAAAAOgAAAAAIAACAAAAApn7bxsL9jtBfibrJodoROHkEqcLS+9QzQAxmjKrM8pGAAAACkrW/4Uue8yhXVzWjF3YzLIlBgbO+7qPSD/vvkQwwOzTB8l2Hgp+TtRmZNO/uR2d63ivHhQYHH1g+32Yk4ZTNIijbj1aqe+ZCc/sMGnjUhCMoVlAYncUAKnlLvYPxZcTlAAAAANd4a+ne/noB73U12+A8hgkEQz97pmZXATj2GXQZGTv/oTJiXVLjUc4z6Fc/hLPv6NYVAIcPo3XuS6ZmL4W2oyQ==\", \"HOMEDRIVE=C:\", \"OPENBLAS_MAIN_FREE=1\", \"PROCESSOR_ARCHITECTURE=AMD64\", \"TEMP=C:\\\\Users\\\\Xuan\\\\AppData\\\\Local\\\\Temp\"]), ProcessRunning), GLMakie.Screen(...), Makie.TickController(Observable(Makie.Tick(Makie.OneTimeRenderTick, 21, 7.0, 0.3333333333333333)), 22, 0.3333333333333333, true, ObserverFunction defined at C:\\Users\\Xuan\\.julia\\packages\\Makie\\4JW9B\\src\\ffmpeg-util.jl:199 operating on Observable(Tick(OneTimeRenderTick, 21, 7.0, 0.333333))), ColorTypes.RGB{FixedPointNumbers.N0f8}[RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0) … RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0); RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0) … RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0); … ; RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0) … RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0); RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0) … RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0)], \"C:\\\\Users\\\\Xuan\\\\AppData\\\\Local\\\\Temp\\\\jl_HzipJD\\\\##video#249.mp4\", Makie.VideoStreamOptions(\"mp4\", 3, 20, \"high422\", \"yuv420p\", 0, \"quiet\", \"pipe:0\", true)))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Simulate 9 traces at random\n",
    "traces = [Gen.simulate(gridworld_agent_model, model_args) for _ in 1:9]\n",
    "# Animate all traces\n",
    "anim = anim_traces(traces, framerate=3, format=\"mp4\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As can be seen, the agent tends to take actions that eventually lead it to\n",
    "towards one of the goals, though there is a fair amount of stopping and\n",
    "backtracking due to the randomness associated with Boltzmann rationality.\n",
    "We can see this randomness more clearly if we visualize the distribution of\n",
    "trajectories towards a specific goal (e.g. the red goal at (1, 1))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<video controls autoplay muted>\n",
       "    <source src=\"data:video/mp4;base64,[encoded data, 82752 characters]\"type=\"video/mp4\">\n",
       "</video>"
      ],
      "text/plain": [
       "VideoDisplay(VideoStream(Base.PipeEndpoint(Base.Libc.WindowsRawSocket(0x0000000000002150) open, 0 bytes waiting), Process(setenv(`\u001b[4m'C:\\Users\\Xuan\\.julia\\artifacts\\e53c880b3a155c54b452cbe94b72de51158aea2e\\bin\\ffmpeg.exe'\u001b[24m \u001b[4m-y\u001b[24m \u001b[4m-loglevel\u001b[24m \u001b[4mquiet\u001b[24m \u001b[4m-threads\u001b[24m \u001b[4m20\u001b[24m \u001b[4m-framerate\u001b[24m \u001b[4m3\u001b[24m \u001b[4m-pixel_format\u001b[24m \u001b[4mrgb24\u001b[24m \u001b[4m-f\u001b[24m \u001b[4mrawvideo\u001b[24m \u001b[4m-s:v\u001b[24m \u001b[4m576x576\u001b[24m \u001b[4m-r\u001b[24m \u001b[4m3\u001b[24m \u001b[4m-i\u001b[24m \u001b[4mpipe:0\u001b[24m \u001b[4m-vf\u001b[24m \u001b[4mvflip\u001b[24m \u001b[4m-profile:v\u001b[24m \u001b[4mhigh422\u001b[24m \u001b[4m-crf\u001b[24m \u001b[4m20\u001b[24m \u001b[4m-preset\u001b[24m \u001b[4mslow\u001b[24m \u001b[4m-c:v\u001b[24m \u001b[4mlibx264\u001b[24m \u001b[4m-pix_fmt\u001b[24m \u001b[4myuv420p\u001b[24m \u001b[4m-an\u001b[24m \u001b[4m'C:\\Users\\Xuan\\AppData\\Local\\Temp\\jl_026vAL\\##video#251.mp4'\u001b[24m`,[\"WINDIR=C:\\\\WINDOWS\", \"PATH=C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\3af036b12a076c2d4c042dde6d956ef4eac2c2ce\\\\bin;C:\\\\Program Files\\\\Julia\\\\1.10.2\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\13e3cf8e6f38d3b6418b758bb68cb8b149c28251\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\ce4261864aeaa78b205a3f6fb750be4c4afe9b4b\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\1da3b3a2698813adf2b925c46aa2d68ec8474449\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\58a307887a7b823f6a7cac7e6337710ec14aced2\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\df4cc9189a878ad66ad99fed70ece5cf4e45b246\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\c3147644202f8802c6131410ffd73f7e41db31f8\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\fcc05863e77d0e700007d714ef7304c87b39d10e\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\166f431d14e588224598f7952888bb847a545ac6\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\51e4c0b5a201939ab2352553b85f5d0932d5cf03\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\cb74fa09f359d430ea01be69a61cd98b50ec8c22\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\6e310796647a1c58e5873a7ab7ad32b8e60dea3f\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\d119ba6f3cc7e3fb09a8bb637f705bf9d9982000\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\ec3406e9ba84281719f20b0fb3a874c9e3bac027\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\8d9f87bf9e08a0e0db47c0ab531ee29fff143eee\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\41cbcbec1d1b6ca0484999b1db571f8066f29233\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\5c7df64500371e5ab4723fe2fc999b996aa8f3a7\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\81ad3b6f8af8fa057c1a9640ca91454b669b342b\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\8258f1ad1bf2f9853eb157e9d4b755f0e427a0ea\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\25bfd4c7adfcd3cb3399a7c509ec239a32c247f3\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\45c659a2370bbaa5bee3c0ed364427d4c33552db\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\98941fbb5eed27cef1d77fa8710290bc197dd24f\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\ba2dc75245508e0cf0c860081b2f2fd10bdf4357\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\a2b7bddf24b04fe5cf3e6d40e0527451b3439d23\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\7e207f56275705f9cbc88b826acd28ea0e40f1ae\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\e53c880b3a155c54b452cbe94b72de51158aea2e\\\\bin;C:\\\\Program Files\\\\Julia\\\\1.10.2\\\\bin\\\\..\\\\lib\\\\julia;C:\\\\Program Files\\\\Julia\\\\1.10.2\\\\bin\\\\..\\\\lib;C:\\\\Program Files\\\\Julia\\\\1.10.2\\\\bin;C:\\\\WINDOWS\\\\system32;C:\\\\WINDOWS;C:\\\\WINDOWS\\\\System32\\\\Wbem;C:\\\\WINDOWS\\\\System32\\\\WindowsPowerShell\\\\v1.0\\\\;C:\\\\WINDOWS\\\\System32\\\\OpenSSH\\\\;C:\\\\Program Files\\\\Microsoft VS Code\\\\bin;C:\\\\Program Files\\\\Git\\\\cmd;C:\\\\Program Files\\\\GitHub CLI\\\\;C:\\\\ProgramData\\\\chocolatey\\\\bin;C:\\\\Program Files\\\\Rust\\\\Rust-MSVC-1.82\\\\bin;C:\\\\Program Files\\\\Git LFS;C:\\\\Program Files\\\\Docker\\\\Docker\\\\resources\\\\bin;C:\\\\Users\\\\Xuan\\\\AppData\\\\Local\\\\Microsoft\\\\WindowsApps;C:\\\\Users\\\\Xuan\\\\AppData\\\\Local\\\\Programs\\\\MiKTeX\\\\miktex\\\\bin\\\\x64\\\\\", \"ELECTRON_RUN_AS_NODE=1\", \"USERDOMAIN_ROAMINGPROFILE=deep-titanium\", \"VSCODE_CODE_CACHE_PATH=C:\\\\Users\\\\Xuan\\\\AppData\\\\Roaming\\\\Code\\\\CachedData\\\\6f17636121051a53c88d3e605c491d22af2ba755\", \"ZES_ENABLE_SYSMAN=1\", \"LOCALAPPDATA=C:\\\\Users\\\\Xuan\\\\AppData\\\\Local\", \"HOMEPATH=\\\\Users\\\\Xuan\", \"VSCODE_NLS_CONFIG={\\\"userLocale\\\":\\\"en-us\\\",\\\"osLocale\\\":\\\"en-us\\\",\\\"resolvedLanguage\\\":\\\"en\\\",\\\"defaultMessagesFile\\\":\\\"C:\\\\\\\\Program Files\\\\\\\\Microsoft VS Code\\\\\\\\resources\\\\\\\\app\\\\\\\\out\\\\\\\\nls.messages.json\\\",\\\"locale\\\":\\\"en-us\\\",\\\"availableLanguages\\\":{}}\", \"PROCESSOR_IDENTIFIER=Intel64 Family 6 Model 186 Stepping 2, GenuineIntel\"  …  \"PROCESSOR_LEVEL=6\", \"EFC_7788_1592913036=1\", \"SYSTEMDRIVE=C:\", \"FPS_BROWSER_APP_PROFILE_STRING=Internet Explorer\", \"PROGRAMW6432=C:\\\\Program Files\", \"IGCCSVC_DB=AQAAANCMnd8BFdERjHoAwE/Cl+sBAAAAon9/u+k2YU+MR1UcImxbYAQAAAACAAAAAAAQZgAAAAEAACAAAAAVjZ9/NJxM1QJQJiNVkBcoSz9zpj/o0TIuggnp8vj0WQAAAAAOgAAAAAIAACAAAAApn7bxsL9jtBfibrJodoROHkEqcLS+9QzQAxmjKrM8pGAAAACkrW/4Uue8yhXVzWjF3YzLIlBgbO+7qPSD/vvkQwwOzTB8l2Hgp+TtRmZNO/uR2d63ivHhQYHH1g+32Yk4ZTNIijbj1aqe+ZCc/sMGnjUhCMoVlAYncUAKnlLvYPxZcTlAAAAANd4a+ne/noB73U12+A8hgkEQz97pmZXATj2GXQZGTv/oTJiXVLjUc4z6Fc/hLPv6NYVAIcPo3XuS6ZmL4W2oyQ==\", \"HOMEDRIVE=C:\", \"OPENBLAS_MAIN_FREE=1\", \"PROCESSOR_ARCHITECTURE=AMD64\", \"TEMP=C:\\\\Users\\\\Xuan\\\\AppData\\\\Local\\\\Temp\"]), ProcessRunning), GLMakie.Screen(...), Makie.TickController(Observable(Makie.Tick(Makie.OneTimeRenderTick, 21, 7.0, 0.3333333333333333)), 22, 0.3333333333333333, true, ObserverFunction defined at C:\\Users\\Xuan\\.julia\\packages\\Makie\\4JW9B\\src\\ffmpeg-util.jl:199 operating on Observable(Tick(OneTimeRenderTick, 21, 7.0, 0.333333))), ColorTypes.RGB{FixedPointNumbers.N0f8}[RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0) … RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0); RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0) … RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0); … ; RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0) … RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0); RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0) … RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0)], \"C:\\\\Users\\\\Xuan\\\\AppData\\\\Local\\\\Temp\\\\jl_026vAL\\\\##video#251.mp4\", Makie.VideoStreamOptions(\"mp4\", 3, 20, \"high422\", \"yuv420p\", 0, \"quiet\", \"pipe:0\", true)))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Constrain goal to specific value\n",
    "goal_constraint = choicemap((:goal, (1, 1)))\n",
    "# Generate 9 traces with constrained goal\n",
    "traces = [Gen.generate(gridworld_agent_model, model_args, goal_constraint)[1] for _ in 1:9]\n",
    "# Animate all traces\n",
    "anim = anim_traces(traces, framerate=3, format=\"mp4\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As the animation demonstrates, the agent may take many different, approximately\n",
    "optimal paths towards the same goal.\n",
    "\n",
    "## Inferring goals via inverse planning\n",
    "\n",
    "Now that we have visualized our model and checked that it produces reasonable\n",
    "forward predictions of agent actions, let's turn to the inverse problem:\n",
    "inferring the agent's goal from their actions. Suppose we observe the following\n",
    "sequence of actions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<video controls autoplay muted>\n",
       "    <source src=\"data:video/mp4;base64,[encoded data, 20548 characters]\"type=\"video/mp4\">\n",
       "</video>"
      ],
      "text/plain": [
       "VideoDisplay(VideoStream(Base.PipeEndpoint(Base.Libc.WindowsRawSocket(0x0000000000001ff8) open, 0 bytes waiting), Process(setenv(`\u001b[4m'C:\\Users\\Xuan\\.julia\\artifacts\\e53c880b3a155c54b452cbe94b72de51158aea2e\\bin\\ffmpeg.exe'\u001b[24m \u001b[4m-y\u001b[24m \u001b[4m-loglevel\u001b[24m \u001b[4mquiet\u001b[24m \u001b[4m-threads\u001b[24m \u001b[4m20\u001b[24m \u001b[4m-framerate\u001b[24m \u001b[4m3\u001b[24m \u001b[4m-pixel_format\u001b[24m \u001b[4mrgb24\u001b[24m \u001b[4m-f\u001b[24m \u001b[4mrawvideo\u001b[24m \u001b[4m-s:v\u001b[24m \u001b[4m512x512\u001b[24m \u001b[4m-r\u001b[24m \u001b[4m3\u001b[24m \u001b[4m-i\u001b[24m \u001b[4mpipe:0\u001b[24m \u001b[4m-vf\u001b[24m \u001b[4mvflip\u001b[24m \u001b[4m-profile:v\u001b[24m \u001b[4mhigh422\u001b[24m \u001b[4m-crf\u001b[24m \u001b[4m20\u001b[24m \u001b[4m-preset\u001b[24m \u001b[4mslow\u001b[24m \u001b[4m-c:v\u001b[24m \u001b[4mlibx264\u001b[24m \u001b[4m-pix_fmt\u001b[24m \u001b[4myuv420p\u001b[24m \u001b[4m-an\u001b[24m \u001b[4m'C:\\Users\\Xuan\\AppData\\Local\\Temp\\jl_8u1qQ9\\##video#253.mp4'\u001b[24m`,[\"WINDIR=C:\\\\WINDOWS\", \"PATH=C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\3af036b12a076c2d4c042dde6d956ef4eac2c2ce\\\\bin;C:\\\\Program Files\\\\Julia\\\\1.10.2\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\13e3cf8e6f38d3b6418b758bb68cb8b149c28251\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\ce4261864aeaa78b205a3f6fb750be4c4afe9b4b\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\1da3b3a2698813adf2b925c46aa2d68ec8474449\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\58a307887a7b823f6a7cac7e6337710ec14aced2\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\df4cc9189a878ad66ad99fed70ece5cf4e45b246\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\c3147644202f8802c6131410ffd73f7e41db31f8\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\fcc05863e77d0e700007d714ef7304c87b39d10e\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\166f431d14e588224598f7952888bb847a545ac6\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\51e4c0b5a201939ab2352553b85f5d0932d5cf03\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\cb74fa09f359d430ea01be69a61cd98b50ec8c22\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\6e310796647a1c58e5873a7ab7ad32b8e60dea3f\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\d119ba6f3cc7e3fb09a8bb637f705bf9d9982000\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\ec3406e9ba84281719f20b0fb3a874c9e3bac027\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\8d9f87bf9e08a0e0db47c0ab531ee29fff143eee\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\41cbcbec1d1b6ca0484999b1db571f8066f29233\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\5c7df64500371e5ab4723fe2fc999b996aa8f3a7\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\81ad3b6f8af8fa057c1a9640ca91454b669b342b\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\8258f1ad1bf2f9853eb157e9d4b755f0e427a0ea\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\25bfd4c7adfcd3cb3399a7c509ec239a32c247f3\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\45c659a2370bbaa5bee3c0ed364427d4c33552db\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\98941fbb5eed27cef1d77fa8710290bc197dd24f\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\ba2dc75245508e0cf0c860081b2f2fd10bdf4357\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\a2b7bddf24b04fe5cf3e6d40e0527451b3439d23\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\7e207f56275705f9cbc88b826acd28ea0e40f1ae\\\\bin;C:\\\\Users\\\\Xuan\\\\.julia\\\\artifacts\\\\e53c880b3a155c54b452cbe94b72de51158aea2e\\\\bin;C:\\\\Program Files\\\\Julia\\\\1.10.2\\\\bin\\\\..\\\\lib\\\\julia;C:\\\\Program Files\\\\Julia\\\\1.10.2\\\\bin\\\\..\\\\lib;C:\\\\Program Files\\\\Julia\\\\1.10.2\\\\bin;C:\\\\WINDOWS\\\\system32;C:\\\\WINDOWS;C:\\\\WINDOWS\\\\System32\\\\Wbem;C:\\\\WINDOWS\\\\System32\\\\WindowsPowerShell\\\\v1.0\\\\;C:\\\\WINDOWS\\\\System32\\\\OpenSSH\\\\;C:\\\\Program Files\\\\Microsoft VS Code\\\\bin;C:\\\\Program Files\\\\Git\\\\cmd;C:\\\\Program Files\\\\GitHub CLI\\\\;C:\\\\ProgramData\\\\chocolatey\\\\bin;C:\\\\Program Files\\\\Rust\\\\Rust-MSVC-1.82\\\\bin;C:\\\\Program Files\\\\Git LFS;C:\\\\Program Files\\\\Docker\\\\Docker\\\\resources\\\\bin;C:\\\\Users\\\\Xuan\\\\AppData\\\\Local\\\\Microsoft\\\\WindowsApps;C:\\\\Users\\\\Xuan\\\\AppData\\\\Local\\\\Programs\\\\MiKTeX\\\\miktex\\\\bin\\\\x64\\\\\", \"ELECTRON_RUN_AS_NODE=1\", \"USERDOMAIN_ROAMINGPROFILE=deep-titanium\", \"VSCODE_CODE_CACHE_PATH=C:\\\\Users\\\\Xuan\\\\AppData\\\\Roaming\\\\Code\\\\CachedData\\\\6f17636121051a53c88d3e605c491d22af2ba755\", \"ZES_ENABLE_SYSMAN=1\", \"LOCALAPPDATA=C:\\\\Users\\\\Xuan\\\\AppData\\\\Local\", \"HOMEPATH=\\\\Users\\\\Xuan\", \"VSCODE_NLS_CONFIG={\\\"userLocale\\\":\\\"en-us\\\",\\\"osLocale\\\":\\\"en-us\\\",\\\"resolvedLanguage\\\":\\\"en\\\",\\\"defaultMessagesFile\\\":\\\"C:\\\\\\\\Program Files\\\\\\\\Microsoft VS Code\\\\\\\\resources\\\\\\\\app\\\\\\\\out\\\\\\\\nls.messages.json\\\",\\\"locale\\\":\\\"en-us\\\",\\\"availableLanguages\\\":{}}\", \"PROCESSOR_IDENTIFIER=Intel64 Family 6 Model 186 Stepping 2, GenuineIntel\"  …  \"PROCESSOR_LEVEL=6\", \"EFC_7788_1592913036=1\", \"SYSTEMDRIVE=C:\", \"FPS_BROWSER_APP_PROFILE_STRING=Internet Explorer\", \"PROGRAMW6432=C:\\\\Program Files\", \"IGCCSVC_DB=AQAAANCMnd8BFdERjHoAwE/Cl+sBAAAAon9/u+k2YU+MR1UcImxbYAQAAAACAAAAAAAQZgAAAAEAACAAAAAVjZ9/NJxM1QJQJiNVkBcoSz9zpj/o0TIuggnp8vj0WQAAAAAOgAAAAAIAACAAAAApn7bxsL9jtBfibrJodoROHkEqcLS+9QzQAxmjKrM8pGAAAACkrW/4Uue8yhXVzWjF3YzLIlBgbO+7qPSD/vvkQwwOzTB8l2Hgp+TtRmZNO/uR2d63ivHhQYHH1g+32Yk4ZTNIijbj1aqe+ZCc/sMGnjUhCMoVlAYncUAKnlLvYPxZcTlAAAAANd4a+ne/noB73U12+A8hgkEQz97pmZXATj2GXQZGTv/oTJiXVLjUc4z6Fc/hLPv6NYVAIcPo3XuS6ZmL4W2oyQ==\", \"HOMEDRIVE=C:\", \"OPENBLAS_MAIN_FREE=1\", \"PROCESSOR_ARCHITECTURE=AMD64\", \"TEMP=C:\\\\Users\\\\Xuan\\\\AppData\\\\Local\\\\Temp\"]), ProcessRunning), GLMakie.Screen(...), Makie.TickController(Observable(Makie.Tick(Makie.OneTimeRenderTick, 10, 3.333333333333333, 0.3333333333333333)), 11, 0.3333333333333333, true, ObserverFunction defined at C:\\Users\\Xuan\\.julia\\packages\\Makie\\4JW9B\\src\\ffmpeg-util.jl:199 operating on Observable(Tick(OneTimeRenderTick, 10, 3.33333, 0.333333))), ColorTypes.RGB{FixedPointNumbers.N0f8}[RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0) … RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0); RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0) … RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0); … ; RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0) … RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0); RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0) … RGB(0.0, 0.0, 0.0) RGB(0.0, 0.0, 0.0)], \"C:\\\\Users\\\\Xuan\\\\AppData\\\\Local\\\\Temp\\\\jl_8u1qQ9\\\\##video#253.mp4\", Makie.VideoStreamOptions(\"mp4\", 3, 20, \"high422\", \"yuv420p\", 0, \"quiet\", \"pipe:0\", true)))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Specify observed actions as a choicemap\n",
    "actions = [\n",
    "    (-1, 0), (-1, 0), # Up twice\n",
    "    (0, -1), (0, -1), # Left twice\n",
    "    (-1, 0), (-1, 0), (-1, 0), # Up thrice\n",
    "    (0, -1), (-1, 0)  # Left, then up\n",
    "]\n",
    "observations = choicemap([((:act, t), a) for (t, a) in enumerate(actions)]...)\n",
    "\n",
    "# Visualize trace with observed actions\n",
    "T = length(actions)\n",
    "model_args = (T, grid, init_loc, goal_locs, beta)\n",
    "obs_trace, _ = Gen.generate(gridworld_agent_model, model_args, observations)\n",
    "anim = anim_trace(obs_trace, framerate=3, format=\"mp4\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Which of the three goals do you think this agent is heading for? You've probably\n",
    "made some intuitive inferences at this point. With inverse planning, we can\n",
    "replicate such inferences. Given our agent model $P(g, a_{1:T}, s_{1:T})$, we\n",
    "wish to compute the posterior distribution $P(g | a_{1:T}, s_{1:T})$ over goals\n",
    "given actions $a_{1:T}$ and environment states $s_{1:T}$. This can be done via\n",
    "Bayes' rule:\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "P(g | a_{1:T}, s_{1:T})\n",
    "&= \\frac{P(g) P(a_{1:T}, s_{1:T} | g)}{P(a_{1:T}, s_{1:T})} \\\\\n",
    "&= \\frac{P(g, a_{1:T}, s_{1:T})}{\\sum_{g' \\in \\mathcal{G}} P(g', a_{1:T}, s_{1:T})}\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "### Inference via enumeration with Gen.jl\n",
    "\n",
    "If we had not written our model as a probabilistic program, we would have to\n",
    "manually implement the math above. With Gen.jl, we can instead use the\n",
    "`enumerative_inference` algorithm, which computes both the numerator and\n",
    "denominator of Bayes' rule by enumerating over each possible goal:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "goal = (1, 1)\tP(goal = (1, 1)) = 0.537\n",
      "goal = (1, 5)\tP(goal = (1, 5)) = 0.352\n",
      "goal = (8, 2)\tP(goal = (8, 2)) = 0.111\n",
      "log P(obs) = -4.906\n"
     ]
    }
   ],
   "source": [
    "# Specify latent choices (in this case, goals) to enumerate over\n",
    "choice_iter = choice_vol_grid((:goal, goal_locs))\n",
    "\n",
    "# Run enumerative inference given observations and the space of enumerated goals\n",
    "T = length(actions)\n",
    "beta = 1.0\n",
    "model_args = (T, grid, init_loc, goal_locs, beta)\n",
    "traces, log_posterior_probs, lml =\n",
    "    enumerative_inference(gridworld_agent_model, model_args,\n",
    "                          observations, choice_iter)\n",
    "\n",
    "# Display probability of each latent choice\n",
    "for (trace, log_p) in zip(traces, log_posterior_probs)\n",
    "    println(\"goal = \", trace[:goal], \"\\t\",\n",
    "            \"P(goal = $(trace[:goal])) = \", @sprintf(\"%.3f\", exp(log_p)))\n",
    "end\n",
    "\n",
    "# Display log marginal likelihood of observations (= log P(a_{1:T}, s_{1:T}))\n",
    "println(\"log P(obs) = \", @sprintf(\"%.3f\", lml))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also plot a bar chart next to our trace rendering to visualize the\n",
    "posterior distribution over goals:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 39696 characters]",
      "text/html": [
       "<img width=1024 height=512 style='object-fit: contain; height: auto;' src=\"data:image/png;base64, [encoded data, 39696 characters]\"/>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "GLMakie.Screen(...)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\"Plots bar chart of goal probabilities to an existing axis.\"\n",
    "function plot_goal_probs!(ax, goal_labels, goal_probs,\n",
    "                          goal_colors = GOAL_COLORS[1:length(goal_locs)])\n",
    "    n_goals = length(goal_labels)\n",
    "    barplot!(ax, 1:n_goals, goal_probs, color=goal_colors,\n",
    "             bar_labels = :y, label_formatter = s -> @sprintf(\"%.3f\", s))\n",
    "    ylims!(ax, 0.0, 1.0)\n",
    "    ax.xticks = (1:n_goals, [string(g) for g in goal_labels])\n",
    "    ax.xticklabelsize = 20\n",
    "    return ax\n",
    "end\n",
    "\n",
    "# Render trace and add new axis with goal probabilities\n",
    "figure = render_trace(obs_trace, size=(1024, 512))\n",
    "ax = Axis(figure[1, 2], title = \"P(goal | actions)\", titlesize = 24)\n",
    "goal_probs = exp.(log_posterior_probs)\n",
    "plot_goal_probs!(ax, goal_locs, goal_probs)\n",
    "display(figure)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The inferred goal probabilities here make intuitive sense. The agent appears to\n",
    "be heading towards either the red goal at (1, 1) or the yellow goal at (5, 1).\n",
    "However, the yellow goal is slightly less likely, since there is another path\n",
    "the agent could have taken to the right to reach the yellow goal. In contrast,\n",
    "most of the observed actions are part of *all* optimal paths from the start to\n",
    "the red goal. Finally, the fact that the agent moved up at the final step\n",
    "indicates that the blue goal at (8, 2) is unlikely.\n",
    "\n",
    "Recall that an important parameter in our Boltzmann-rational agent model is the\n",
    "rationality coefficient $\\beta$ (`beta`), which controls how optimal or noisy\n",
    "we assume the agent to be. Let's see what happens to our goal inferences when\n",
    "we compute them with different values of $\\beta$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 85724 characters]",
      "text/html": [
       "<img width=2048 height=512 style='object-fit: contain; height: auto;' src=\"data:image/png;base64, [encoded data, 85724 characters]\"/>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "GLMakie.Screen(...)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Render observed actions\n",
    "figure = render_trace(obs_trace, size=(512*4, 512))\n",
    "\n",
    "# Iterate over different values of beta and add plots\n",
    "betas = [0.5, 1.0, 2.0]\n",
    "for (i, beta) in enumerate(betas)\n",
    "    # Run enumerative inference with current beta\n",
    "    model_args = (T, grid, init_loc, goal_locs, beta)\n",
    "    traces, log_posterior_probs, lml =\n",
    "        enumerative_inference(gridworld_agent_model, model_args,\n",
    "                              observations, choice_iter)\n",
    "    # Add axis and bar chart of goal probabilities\n",
    "    ax = Axis(figure[1, i+1], titlesize = 24,\n",
    "              title = @sprintf(\"P(goal | actions; β = %.1f)\", beta))\n",
    "    goal_probs = exp.(log_posterior_probs)\n",
    "    plot_goal_probs!(ax, goal_locs, goal_probs)\n",
    "end\n",
    "display(figure)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As the figure above shows, lower values of $\\beta$ result in more graded\n",
    "inferences, whereas higher values of $\\beta$ result in sharper inferences.\n",
    "This is because suboptimal trajectories to a particular goal have a much lower\n",
    "likelihood when $\\beta$ is high.\n",
    "\n",
    "**Question**: Which value of $\\beta$ leads to goal inferences that best\n",
    "correspond with your own intuitive inferences? Now go back and use that value\n",
    "of $\\beta$ to simulate forward predictions of thhe agent's actions. Does the\n",
    "$\\beta$ that best capture your goal inferences also capture your intuitions\n",
    "about how an agent acts over time?\n",
    "\n",
    "### Understanding enumerative inference\n",
    "\n",
    "Let's take a moment to understand how enumerative inference actually works in\n",
    "Gen.jl, since this is useful for appreciating how it differs from the more\n",
    "efficient algorithms we'll use later. Under-the-hood, enumerative inference\n",
    "works as follows:\n",
    "- Use `Gen.generate` method to generate traces consistent with both the\n",
    "  observations and each possible latent choice.\n",
    "- Call `Gen.get_score` to compute the log joint probability of each trace.\n",
    "  This gives us the (log.) numerator of Bayes' rule for each latent choice.\n",
    "- Summing the probabilities of each trace to get the denominator of Bayes' rule,\n",
    "  and then dividing to get the posterior probability of each latent choice.\n",
    "\n",
    "Here's how that looks in code for our agent model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "goal = (1, 1)\tP(goal = (1, 1)) = 0.537\n",
      "goal = (1, 5)\tP(goal = (1, 5)) = 0.352\n",
      "goal = (8, 2)\tP(goal = (8, 2)) = 0.111\n",
      "log P(obs) = -4.906\n"
     ]
    }
   ],
   "source": [
    "# Specify model arguments and latent choices to enumerate over\n",
    "model_args = (T, grid, init_loc, goal_locs, beta)\n",
    "latent_choices = [choicemap((:goal, g)) for g in goal_locs]\n",
    "\n",
    "# Iterate over traces and record joint probabilities\n",
    "traces = []\n",
    "log_joint_probs = Float64[]\n",
    "for choices in latent_choices\n",
    "    # Combine observation constraints (a_{1:T}) with latent choice (g)\n",
    "    constraints = merge(observations, choices)\n",
    "    # Generate trace constrained to current latent choice and observations\n",
    "    trace, log_weight =\n",
    "        Gen.generate(gridworld_agent_model, model_args, constraints)\n",
    "    # Compute log joint probability (= log P(g, a_{1:T}, s_{1:T})) of the trace\n",
    "    log_joint_prob = Gen.get_score(trace)\n",
    "    # (Note: We could also use `log_weight` directly, w/o calling `get_score`)\n",
    "    push!(traces, trace)\n",
    "    push!(log_joint_probs, log_joint_prob)\n",
    "end\n",
    "\n",
    "# Compute log. denominator of Bayes' rule (= log P(a_{1:T}, s_{1:T}))\n",
    "log_total_prob = logsumexp(log_joint_probs)\n",
    "# Compute log. posterior probabilities (= log P(g | a_{1:T}, s_{1:T}))\n",
    "log_posterior_probs = log_joint_probs .- log_total_prob\n",
    "\n",
    "# Display probability of each latent choice\n",
    "for (trace, log_p) in zip(traces, log_posterior_probs)\n",
    "    println(\"goal = \", trace[:goal], \"\\t\",\n",
    "            \"P(goal = $(trace[:goal])) = \", @sprintf(\"%.3f\", exp(log_p)))\n",
    "end\n",
    "\n",
    "# Display log marginal likelihood of observations (= log P(a_{1:T}, s_{1:T}))\n",
    "println(\"log P(obs) = \", @sprintf(\"%.3f\", log_total_prob))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that our manual implementation of enumerative inference gives the same\n",
    "results as the `enumerative_inference` function provided by Gen.jl.\n",
    "\n",
    "As the implementation above makes clear, enumerative inference requires a\n",
    "for-loop over *every* possible value of a latent variable. In other words, for\n",
    "inference over discrete latent variables, enumerative inference is an *exact*\n",
    "algorithm, since it computes Bayes' rule exactly. This is fine for small\n",
    "hypothesis spaces, but scales poorly as the number of possible values or\n",
    "latent variables increases. We'll see later on how to scale more efficiently\n",
    "by turning to approximate, *sampling-based* inference algorithms.\n",
    "\n",
    "Another limitation of enumerative inference — at least as we implemented it\n",
    "above — is that we need to rerun the whole algorithm every time we receive\n",
    "a new observation. This makes it inefficient for online goal inference.\n",
    "In the next section, we'll see how to avoid this problem via *incremental\n",
    "computation*.\n",
    "\n",
    "## Online inverse planning via incremental computation\n",
    "\n",
    "In many scenarios, we want to infer an agent's goal (and other mental states)\n",
    "from a stream of incoming observations, updating a posterior distribution over\n",
    "the agent's goal as we receive them. For example, this is important for\n",
    "collaborative robots or AI assistants that have to interact with a user in\n",
    "real-time. Such systems should continually update their inferences in order to\n",
    "guide their actions.\n",
    "\n",
    "A naive way to achieve online goal inference would be to run an inference\n",
    "algorithm on the entire set of observations each time a new observation arrives.\n",
    "Let's try this with enumerative inference and measure the runtime:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct a series of 20 actions to the blue goal\n",
    "actions = [\n",
    "    (-1, 0), (-1, 0), (0, 1), # Up twice then right\n",
    "    (-1, 0), (-1, 0), # Up twice\n",
    "    (0, -1), (0, -1), (0, -1), (-1, 0), # Left thrice then up\n",
    "    (0, -1), (0, -1), (0, -1), # Left thrice\n",
    "    (1, 0), (1, 0), (0, 1), # Down twice then right\n",
    "    (1, 0), (1, 0), (1, 0), (0, -1), # Down thrice then left\n",
    "    (0, 0) # Stay in place\n",
    "]\n",
    "\n",
    "# Allocate buffers to store runtime and goal probabilities\n",
    "T = length(actions)\n",
    "runtimes_enum = Vector{Float64}(undef, T)\n",
    "goal_probs = Matrix{Float64}(undef, length(goal_locs), T)\n",
    "\n",
    "# Iterate over timesteps and run inference\n",
    "choice_iter = choice_vol_grid((:goal, goal_locs))\n",
    "observations = choicemap()\n",
    "for (t, act) in enumerate(actions)\n",
    "    # Add current action to observation choicemap\n",
    "    observations[(:act, t)] = act\n",
    "    # Run enumerative inference given all actions observed so far\n",
    "    model_args = (t, grid, init_loc, goal_locs, beta)\n",
    "    stats = @timed begin\n",
    "        traces, log_posterior_probs, lml =\n",
    "            enumerative_inference(gridworld_agent_model, model_args,\n",
    "                                  observations, choice_iter)\n",
    "    end\n",
    "    # Record runtime and goal probabilities\n",
    "    runtimes_enum[t] = stats.time\n",
    "    goal_probs[:, t] = exp.(log_posterior_probs)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's plot both the goal probabilities over time, and the runtime of the\n",
    "algorithm for each step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 110944 characters]",
      "text/html": [
       "<img width=1024 height=512 style='object-fit: contain; height: auto;' src=\"data:image/png;base64, [encoded data, 110944 characters]\"/>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "GLMakie.Screen(...)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Render observed actions\n",
    "obs_trace = Gen.generate(gridworld_agent_model, model_args, observations)[1]\n",
    "figure = render_trace(obs_trace; size = (1024, 512))\n",
    "\n",
    "# Plot goal probabilities over time\n",
    "layout = GridLayout(figure[1, 2])\n",
    "ax = Axis(layout[1, 1], titlesize = 24, title = \"p(goal | actions)\",\n",
    "          limits = (1, T, 0, 1), xticks = 1:T)\n",
    "series!(ax, 1:T, goal_probs, color = GOAL_COLORS[1:length(goal_locs)],\n",
    "        strokewidth = 2.0, markercolor = GOAL_COLORS[1:length(goal_locs)],\n",
    "        labels = string.(goal_locs))\n",
    "axislegend(ax, \"Goal\", position = :rt, framevisible = false)\n",
    "\n",
    "# Plot runtimes for each timestep\n",
    "ax = Axis(layout[2, 1], titlesize = 24, title = \"runtime (s)\",\n",
    "          limits = (1, T, 0, nothing), xticks = 1:T,\n",
    "          ytickformat = ys -> [@sprintf(\"%.1E\", y) for y in ys])\n",
    "series!(ax, 1:T, runtimes_enum', markersize = 10.0, solid_color=:black)\n",
    "display(figure)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the top-right plot, we see that the posterior distribution over goals evolves\n",
    "in a reasonable manner as more actions are observed. When the agent turns right\n",
    "at `t = 3`, the probability of the yellow goal increases significantly. But once\n",
    "the agent turns left again at `t = 6`, the yellow goal drops in probability.\n",
    "\n",
    "As for the bottom-right plot, we can see that while the runtime of inference\n",
    "is low (due to the simplicity of our gridworld environment), it increases more\n",
    "or less linearly over time. The cumulative runtime of all steps is thus\n",
    "*quadratic* in the number of observations. This is because each time we observe\n",
    "a new action, we are re-evaluating the likelihood of *all* actions observed so\n",
    "far, which requires recomputing the Q-values via planning for all steps.\n",
    "\n",
    "While a quadratic runtime for online goal inference is acceptable for small\n",
    "environments like the above, in more complex environments where planning is\n",
    "harder, this can lead to significant runtime costs. Ideally then, we would like\n",
    "to reuse our previous computations, and only compute the likelihood of the most\n",
    "recent action when it is observed. We'll address this by using *incremental\n",
    "computation* in the following ways:\n",
    "\n",
    "1. Rewriting our agent model probabilistic program using Gen's `Unfold`\n",
    "   combinator, enabling *incremental simulation* of our model using the\n",
    "   `Gen.update` function.\n",
    "2. Performing inference via a *Bayesian filtering* algorithm instead of\n",
    "   `enumerative_inference`, which allows us to update our inferences\n",
    "   step-by-step instead of re-running inference from scratch.\n",
    "3. Further accelerating inference with an *incremental planning* algorithm,\n",
    "   instead of rerunning Dijkstra's algorithm to compute Q-values at every step.\n",
    "\n",
    "### Incremental model simulation with the `Unfold` combinator\n",
    "\n",
    "Gen.jl provides the `Unfold` combinator, which allows us to replace a for-loop\n",
    "in a probabilistic program with a subroutine that can be incrementally simulated\n",
    "when the number of timesteps is increased (or decreased). To use the `Unfold`\n",
    "combinator we'll write a *kernel function* called `agent_step` that simulates\n",
    "just one timestep of the agent and environment (replacing the body of the\n",
    "original for-loop). We'll then apply the `Unfold` combinator to our kernel to\n",
    "construct a new probabilistic program (`agent_unfold = Unfold(agent_step)`), and\n",
    "call `agent_unfold` as a subroutine in our full agent model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"Kernel function for one step of the gridworld agent model.\"\n",
    "@gen function agent_step(t, model_state, grid, goal, beta)\n",
    "    loc = model_state.loc\n",
    "    # Compute Q-values and Boltzmann action probabilities\n",
    "    actions, q_values = compute_q_values(grid, loc, goal)\n",
    "    act_probs = softmax(beta .* q_values)\n",
    "    # Sample action according to Boltzmann policy\n",
    "    act = {:act} ~ labeled_categorical(actions, act_probs)\n",
    "    # Update location following environment dynamics\n",
    "    next_loc = gridworld_env_step(grid, loc, act)\n",
    "    # Return new state containing both location and action\n",
    "    next_model_state = (loc=next_loc, act=act)\n",
    "    return next_model_state\n",
    "end\n",
    "\n",
    "# Apply Unfold combinator, returning a new generative function\n",
    "agent_unfold = Gen.Unfold(agent_step)\n",
    "\n",
    "\"Agent model using the Unfold combinator for incremental simulation.\"\n",
    "@gen (static) function unfold_agent_model(T, grid, init_loc, goal_locs, beta)\n",
    "    # Sample goal from uniform goal prior\n",
    "    goal = {:goal} ~ labeled_uniform(goal_locs)\n",
    "    # Initial state contains location and a dummy action\n",
    "    init_state = NamedTuple{(:loc, :act)}((init_loc, (0, 0)))\n",
    "    # Simulate agent and environment for T steps\n",
    "    model_states = {:steps} ~ agent_unfold(T, init_state, grid, goal, beta)\n",
    "    # Return goal and history of locations and actions\n",
    "    return goal, model_states\n",
    "end;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the `(static)` annotation in front of `unfold_agent_model`, which declares\n",
    "it as a [static generative function](https://www.gen.dev/docs/dev/ref/modeling/sml/).\n",
    "This is currently required by Gen.jl to support incremental computation.\n",
    "\n",
    "By rewriting our model with the `Unfold` combinator, we can incrementally\n",
    "simulate our model using another of Gen's interface functions, `Gen.update`.\n",
    "Starting with a trace initialized with `T=0`, we'll repeatedly apply\n",
    "`Gen.update` to update the arguments passed to model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial trace at T=0:\n",
      "│\n",
      "├── :goal : (1, 1)\n",
      "│\n",
      "└── :steps\n",
      "\n",
      "After extending to T=20:\n",
      "│\n",
      "├── :goal : (1, 1)\n",
      "│\n",
      "└── :steps\n",
      "    │\n",
      "    ├── 1\n",
      "    │   │\n",
      "    │   └── :act : (0, 0)\n",
      "    │\n",
      "    ├── 2\n",
      "    │   │\n",
      "    │   └── :act : (0, 0)\n",
      "    │\n",
      "    ├── 3\n",
      "    │   │\n",
      "    │   └── :act : (0, 0)\n",
      "    │\n",
      "    ├── 4\n",
      "    │   │\n",
      "    │   └── :act : (-1, 0)\n",
      "    │\n",
      "    ├── 5\n",
      "    │   │\n",
      "    │   └── :act : (0, 0)\n",
      "    │\n",
      "    ├── 6\n",
      "    │   │\n",
      "    │   └── :act : (0, 0)\n",
      "    │\n",
      "    ├── 7\n",
      "    │   │\n",
      "    │   └── :act : (0, 1)\n",
      "    │\n",
      "    ├── 8\n",
      "    │   │\n",
      "    │   └── :act : (-1, 0)\n",
      "    │\n",
      "    ├── 9\n",
      "    │   │\n",
      "    │   └── :act : (-1, 0)\n",
      "    │\n",
      "    ├── 10\n",
      "    │   │\n",
      "    │   └── :act : (0, 0)\n",
      "    │\n",
      "    ├── 11\n",
      "    │   │\n",
      "    │   └── :act : (-1, 0)\n",
      "    │\n",
      "    ├── 12\n",
      "    │   │\n",
      "    │   └── :act : (-1, 0)\n",
      "    │\n",
      "    ├── 13\n",
      "    │   │\n",
      "    │   └── :act : (0, 0)\n",
      "    │\n",
      "    ├── 14\n",
      "    │   │\n",
      "    │   └── :act : (-1, 0)\n",
      "    │\n",
      "    ├── 15\n",
      "    │   │\n",
      "    │   └── :act : (-1, 0)\n",
      "    │\n",
      "    ├── 16\n",
      "    │   │\n",
      "    │   └── :act : (0, -1)\n",
      "    │\n",
      "    ├── 17\n",
      "    │   │\n",
      "    │   └── :act : (0, -1)\n",
      "    │\n",
      "    ├── 18\n",
      "    │   │\n",
      "    │   └── :act : (0, -1)\n",
      "    │\n",
      "    ├── 19\n",
      "    │   │\n",
      "    │   └── :act : (0, -1)\n",
      "    │\n",
      "    └── 20\n",
      "        │\n",
      "        └── :act : (0, 0)\n"
     ]
    }
   ],
   "source": [
    "# Start with T=0 (no actions taken yet)\n",
    "fixed_args = (grid, init_loc, goal_locs, beta)\n",
    "init_args = (0, fixed_args...)\n",
    "goal_constraint = choicemap((:goal, (1, 1)))\n",
    "trace, _ = Gen.generate(unfold_agent_model, init_args, goal_constraint)\n",
    "println(\"Initial trace at T=0:\")\n",
    "show(stdout, \"text/plain\", get_choices(trace))\n",
    "\n",
    "# Incrementally extend the trace for T=1, 2, 3, ...\n",
    "runtimes_unfold = Float64[]\n",
    "for T in 1:20\n",
    "    new_args = # Update the timestep argument\n",
    "        (T, fixed_args...)\n",
    "    argdiffs = # Inform Gen.update that no other arguments have changed\n",
    "        (UnknownChange(), NoChange(), NoChange(), NoChange(), NoChange())\n",
    "    stats = @timed begin\n",
    "        trace, _, _, _ = # Update the previous trace to get a new trace\n",
    "            Gen.update(trace, new_args, argdiffs, EmptyChoiceMap())\n",
    "    end\n",
    "    push!(runtimes_unfold, stats.time) # Record runtime\n",
    "end\n",
    "println(\"\\n\", \"After extending to T=20:\")\n",
    "show(stdout, \"text/plain\", get_choices(trace))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compared to our original model, the Unfold-based model has a slightly different\n",
    "address structure for its random choices: All addresses for the action variables\n",
    "are now nested under the address `:steps`, such that action $a_t$ has the\n",
    "address `:steps => t => :act`.\n",
    "\n",
    "If we plot the runtime of each update, we'll see that it is more or less\n",
    "constant across steps. This is because `Gen.update` only simulates the model for\n",
    "the most recent step when `T` is incremented."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 48352 characters]",
      "text/html": [
       "<img width=800 height=400 style='object-fit: contain; height: auto;' src=\"data:image/png;base64, [encoded data, 48352 characters]\"/>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "GLMakie.Screen(...)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "figure = Figure(size=(800, 400))\n",
    "ax = Axis(figure[1, 1], title=\"Runtime Improvement with Unfold Model\",\n",
    "          xlabel=\"Timestep\", ylabel=\"Runtime (s)\", xlabelsize=20, ylabelsize=20,\n",
    "          ytickformat = ys -> [@sprintf(\"%.1E\", y) for y in ys], titlesize=24)\n",
    "series!(ax, 1:20, runtimes_unfold', labels=[\"Unfold Model\"],\n",
    "        linewidth=2, markersize=10, solid_color=:black)\n",
    "display(figure)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In contrast, incremental computation is not supported by our original model,\n",
    "so `Gen.update` still has the rerun the entire model from `t=0` each time `T` is\n",
    "incremented. This results in a linearly increasing runtime:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 67488 characters]",
      "text/html": [
       "<img width=800 height=400 style='object-fit: contain; height: auto;' src=\"data:image/png;base64, [encoded data, 67488 characters]\"/>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "GLMakie.Screen(...)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Generate initial trace of our original model\n",
    "trace, _ = Gen.generate(gridworld_agent_model, init_args, goal_constraint)\n",
    "\n",
    "# Incrementally extend the trace for T=1, 2, 3, ...\n",
    "runtimes_original = Float64[]\n",
    "for T in 1:20\n",
    "    new_args = # Update the timestep argument\n",
    "        (T, fixed_args...)\n",
    "    argdiffs = # Inform Gen.update that no other arguments have changed\n",
    "        (UnknownChange(), NoChange(), NoChange(), NoChange(), NoChange())\n",
    "    stats = @timed begin\n",
    "        trace, _, _, _ = # Update the previous trace to get a new trace\n",
    "            Gen.update(trace, new_args, argdiffs, EmptyChoiceMap())\n",
    "    end\n",
    "    push!(runtimes_original, stats.time) # Record runtime\n",
    "end\n",
    "\n",
    "# Add runtimes for original model to plot\n",
    "series!(ax, 1:20, runtimes_original', labels=[\"Original Model\"],\n",
    "        linewidth=2, markersize=10, solid_color=:red)\n",
    "axislegend(ax, position=:rt, framevisible=false, labelsize=20)\n",
    "display(figure)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, while the runtime of the original model is similar for the first\n",
    "few timesteps, it grows to be many times larger.\n",
    "\n",
    "### Incremental goal inference via Bayesian filtering\n",
    "\n",
    "Now that we have a model that can be incrementally simulated, we can implement\n",
    "an inference algorithm that incrementally updates posterior beliefs about the\n",
    "agent's goal as new actions are observed. This is known as *Bayesian filtering*\n",
    "or *sequential Bayesian inference*. Instead of recomputing the full posterior\n",
    "from scratch each time we observe a new action $a_{t}$, we can update our\n",
    "posterior incrementally using Bayes' rule:\n",
    "\n",
    "$$P(g | a_{1:t}) \\propto P(g | a_{1:t-1}) P(a_{t} | g, a_{1:t-1}) $$\n",
    "\n",
    "where $P(g | a_{1:t-1})$ is our posterior belief after observing the first $t-1$\n",
    "actions, and $P(a_{t} | g, a_{1:t-1})$ is the likelihood of the new action\n",
    "given the goal and previous actions.\n",
    "\n",
    "To implement a Bayesian filter, we'll first define a `FilterState` struct that\n",
    "maintains traces for each possible latent value (in our case, the agent's goal)\n",
    "along with their associated weights (unnormalized posterior probabilities)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"State for Bayesian filtering, storing traces and weights for each hypothesis.\"\n",
    "struct FilterState\n",
    "    traces::Vector{Trace}\n",
    "    log_weights::Vector{Float64}\n",
    "end;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialization of our Bayesian filter works similarly to enumerative inference.\n",
    "Given an iterator over all possible latent values, we create a trace constrained\n",
    "to each latent choice and any initial observations with `Gen.generate`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"Initialize Bayesian filter with uniform prior over goals.\"\n",
    "function bayes_filter_init(model, model_args, choice_iter,\n",
    "                           obs_init = EmptyChoiceMap())\n",
    "    # Allocate buffers\n",
    "    traces = Vector{Trace}(undef, length(choice_iter))\n",
    "    log_weights = Vector{Float64}(undef, length(choice_iter))\n",
    "    # Create a trace for each possible latent choice\n",
    "    for (i, choices) in enumerate(choice_iter)\n",
    "        constraints = merge(choices, obs_init)\n",
    "        traces[i], log_weights[i] = Gen.generate(model, model_args, constraints)\n",
    "    end\n",
    "    return FilterState(traces, log_weights)\n",
    "end;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As with `enumerative_inference`, it is the user's responsibility to ensure that\n",
    "all latent values are enumerated over by `choice_iter`, and that  together with\n",
    "`init_obs`, each set of latent `choices` results in a fully constrained\n",
    "initial trace.\n",
    "\n",
    "Now we'll implement the filter update. Given a new observation provided as a\n",
    "choicemap `obs_new`, we extend each trace using `Gen.update`, which is called\n",
    "with the both the new timestep and the new observation as arguments:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"Update Bayesian filter with new timestep and observation.\"\n",
    "function bayes_filter_step!(filter_state::FilterState,\n",
    "                            t_new::Int, obs_new::ChoiceMap)\n",
    "    # Update model arguments (assume first argument is timestep)\n",
    "    old_args = Gen.get_args(filter_state.traces[1])\n",
    "    new_args = (t_new, old_args[2:end]...)\n",
    "    argdiffs = (UnknownChange(), map(x -> NoChange(), old_args[2:end])...)\n",
    "\n",
    "    # Update each trace with the new timestep and observation\n",
    "    for i in 1:length(filter_state.traces)\n",
    "        # Update trace to new timestep and incorporate new observation\n",
    "        trace, incremental_weight, _, _ = Gen.update(\n",
    "            filter_state.traces[i], new_args, argdiffs, obs_new\n",
    "        )\n",
    "        # Update stored trace and add log-likelihood of new observation\n",
    "        filter_state.traces[i] = trace\n",
    "        filter_state.log_weights[i] += incremental_weight\n",
    "    end\n",
    "    return filter_state\n",
    "end;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that in addition to returning an updated trace, `Gen.update` also returns\n",
    "the *incremental log-probability* (a.k.a. incremental weight) of the new trace\n",
    "relative to the old trace. For our model, this incremental weight is exactly\n",
    "equal to $\\log P(a_{t} | g, a_{1:t-1})$, the log-likelihood of the most recent\n",
    "action, since $a_{t}$ is the only random variable introduced when we move to\n",
    "timestep $t$.  As such, we can update our (unnormalized) posterior probabilities\n",
    "by incrementing the `log_weights` stored in our `FilterState`.\n",
    "\n",
    "Now let's test our Bayes filter on the same sequence of actions we used to\n",
    "study the runtime of enumerative inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Allocate buffers to store runtime and goal probabilities\n",
    "T = length(actions)\n",
    "runtimes_bfilter = Vector{Float64}(undef, T)\n",
    "goal_probs = Matrix{Float64}(undef, length(goal_locs), T)\n",
    "\n",
    "# Initialize Bayes filter\n",
    "init_args = (0, grid, init_loc, goal_locs, beta)\n",
    "choice_iter = (choicemap((:goal, g)) for g in goal_locs)\n",
    "filter_state = bayes_filter_init(unfold_agent_model, init_args, choice_iter)\n",
    "\n",
    "# Process each action incrementally\n",
    "for (t, act) in enumerate(actions)\n",
    "    # Create observation for this timestep\n",
    "    obs_t = choicemap((:steps => t => :act, act))\n",
    "    # Update Bayes filter and measure runtime\n",
    "    stats = @timed begin\n",
    "        bayes_filter_step!(filter_state, t, obs_t)\n",
    "    end\n",
    "    # Record runtime and goal probability\n",
    "    runtimes_bfilter[t] = stats.time\n",
    "    goal_probs[:, t] = softmax(filter_state.log_weights)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plotting the inferred goal probabilities and runtimes, we should see exactly\n",
    "the same inferences as produced by enumerative inference, but now without\n",
    "the linearly increasing runtime per step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 107668 characters]",
      "text/html": [
       "<img width=1024 height=512 style='object-fit: contain; height: auto;' src=\"data:image/png;base64, [encoded data, 107668 characters]\"/>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "GLMakie.Screen(...)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Render observed actions\n",
    "obs_trace = filter_state.traces[1]\n",
    "figure = render_trace(obs_trace; size = (1024, 512))\n",
    "\n",
    "# Plot goal probabilities over time\n",
    "layout = GridLayout(figure[1, 2])\n",
    "ax = Axis(layout[1, 1], titlesize = 24, title = \"p(goal | actions)\",\n",
    "          limits = (1, T, 0, 1), xticks = 1:T)\n",
    "series!(ax, 1:T, goal_probs, color = GOAL_COLORS[1:length(goal_locs)],\n",
    "        strokewidth = 2.0, markercolor = GOAL_COLORS[1:length(goal_locs)],\n",
    "        labels = string.(goal_locs))\n",
    "axislegend(ax, \"Goal\", position = :rt, framevisible = false)\n",
    "\n",
    "# Plot runtimes for each timestep\n",
    "ax = Axis(layout[2, 1], titlesize = 24, title = \"runtime (s)\",\n",
    "          limits = (1, T, 0, nothing), xticks = 1:T,\n",
    "          ytickformat = ys -> [@sprintf(\"%.1E\", y) for y in ys])\n",
    "series!(ax, 1:T, runtimes_bfilter', markersize = 10.0, solid_color=:black,\n",
    "        labels = [\"Bayes Filter w. Unfold Model\"])\n",
    "display(figure)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we add the runtimes for enumerative inference to the same plot, we see\n",
    "a marked difference compared to Bayesian filtering:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 125920 characters]",
      "text/html": [
       "<img width=1024 height=512 style='object-fit: contain; height: auto;' src=\"data:image/png;base64, [encoded data, 125920 characters]\"/>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "GLMakie.Screen(...)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "series!(ax, 1:T, runtimes_enum', markersize = 10.0, solid_color=:red,\n",
    "        labels = [\"Enum. Inference\"])\n",
    "axislegend(ax, position = :lt, framevisible = false)\n",
    "display(figure)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Through incremental computation, we have eliminated the unnecessary overhead\n",
    "of re-simulating the model from scratch at each timestep. While there are other\n",
    "ways we could have eliminated this overhead (e.g. implementing inference\n",
    "manually instead of using Gen.jl), doing this with Gen.jl allows us to maintain\n",
    "a clean separation between modeling and inference code, enabling future\n",
    "revisions to our model without having to reimplement inference. We'll consider\n",
    "one such revision in the next subsection.\n",
    "\n",
    "### TODO: Incremental planning for faster computation of Q-values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "*This notebook was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).*"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.10.2",
   "language": "julia",
   "name": "julia-1.10"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.10.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 3
}
