SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Kartik Chandra*
Tony Chen*
MIT Brain and Cognitive Sciences
Tzu-Mao Li
Jonathan Ragan-Kelley
A single panel of a comic book can say a lot: it can depict not only where the characters currently are, but also their motions, their motivations, their emotions, and what they might do next. More generally, humans routinely infer complex sequences of past and future events from a static snapshot of a dynamic scene, even in situations they have never seen before.
In this paper, we model how humans make such rapid and flexible inferences. Building on a long line of work in cognitive science, we offer a Monte Carlo algorithm whose inferences correlate well with human intuitions in a wide variety of domains, while only using a small, cognitively-plausible number of samples. Our key technical insight is a surprising connection between our inference problem and Monte Carlo path tracing, which allows us to apply decades of ideas from the computer graphics community to this seemingly-unrelated theory of mind task.
Hemingway's shortest short story simply reads "For sale: baby shoes, never worn." There is no action in this sentence—it is just a description of a static scene—but as readers, we nonetheless infer a complex and tragic backstory from the "snapshot" that Hemingway describes. This remarkable ability comes naturally to humans: we routinely reconstruct motives from evidence (e.g. at a crime scene), recognize intentions from unfinished tasks (e.g. grading incomplete homework), and enjoy artistic depictions of dynamic action in static media (e.g. a Renaissance "tableau" or sculpture).
How do we do it, so quickly and so intuitively? A long line of empirical work in cognitive science has shown that people, even infants and young children [Southgate and Csibra, 2009, Gergely and Csibra, 2003, Jara-Ettinger et al., 2015, Hamlin et al., 2007], make rapid, flexible, and robust theory-of-mind judgments "out of the box": in novel domains they have never seen before, and without extensive pretraining on data. This remarkable ability has motivated decades of work in AI (see Section 4), which has successfully addressed the problem of inferring an agent's goal from a trajectory of observed actions in a few- or even one-shot manner [Baker et al., 2017, 2009]. These methods cast the problem of action understanding as Bayesian inference:  $P(\text{goal} \mid \text{actions}) \propto P(\text{actions} \mid \text{goal}) P(\text{goal})$ , where  $P(\text{actions} \mid \text{goal})$  is modeled by comparing the observed actions to the optimal actions a rational agent would take towards that goal.
However, if we only observe a single state snapshot, these methods break down—there are simply no actions to condition on. Instead, we must jointly infer not only where the agent might be going, but also where it came from. Recently, Lopez-Brau et al. [2020, 2022] performed this inference by rejection-sampling possible paths taken by the agent. Their model's predictions are remarkably close to human judgements. However, rejection sampling is extremely inefficient—it is slow even on simple problems, and simply does not scale to more sophisticated problems, suggesting that there is more to how humans perform such inference.
Figure 1: How can we infer what an agent is trying to do, based on a snapshot of its current state?
In this paper, we propose a dramatically more efficient algorithm that recovers human-like judgements even on significantly more sophisticated problems. We take inspiration from the field of computer graphics, which has developed a wealth of Monte Carlo algorithms for "path tracing." Our key insight is an analogy between tracing paths of light through 3D scenes and paths of agents through planning domains. This insight allows us to import decades of algorithmic advances from computer graphics to the seemingly-unrelated social intelligence problem of snapshot inference.
In Section 2 we formalize this problem and present a new Monte Carlo algorithm for sampling approximate solutions. Then, in Section 3, we show that our algorithm is up to  $30,000 \times$  more efficient than prior work. Finally, in Section 3.3 we demonstrate via three behavioral studies that our model's predictions match human judgements on new, scaled-up tasks inaccessible to prior work.
2 Our proposed algorithm
We describe our method in the setting of Markov Decision Processes (MDPs). An MDP is a tuple  $\langle S, \mathcal{A}, \mathcal{T}, \mathcal{R}, \gamma \rangle$ , where  $\mathcal{S}$  is the set of possible states,  $\mathcal{A}$  is the space of possible actions,  $\mathcal{T}: S \times \mathcal{A} \times S \to [0,1]$  gives the transition probability,  $\mathcal{R}: S \times \mathcal{A} \to \mathbb{R}$  is the reward function and  $\gamma \in [0,1]$  is a discount factor. We will specifically consider MDPs that have a single "goal state"  $g$  with positive reward, such as a grid world with a target square the agent needs to reach. We will assume that  $g$  is known to the agent, but not to us (the observers).
Consider an agent that begins in some initial state  $s \sim p(s)$ . While the agent is on its trajectory from  $s$  to  $g$ , we observe a "snapshot" of the agent in some interim state  $x$ . Given only  $x$  (and not  $s!$ ), our aim is to infer  $p(g \mid x)$ .
Applying Bayes' rule, we have  $p(g \mid x) \propto p(x \mid g)p(g)$ , where  $p(g)$  is our prior belief about what  $g$  might be. The likelihood  $p(x \mid g)$  can be computed by marginalizing over all possible start states  $s$  and all possible paths  $\pi_{s:g}$  from  $s$  to  $g$ , where a path  $\pi_{s:g}$  is defined as a sequence of states  $\pi_{s:g} = (s_1, \ldots, s_n)$  such that  $s_1 = s$  and  $s_n = g$ . Performing this marginalization, and noting that  $p(s \mid g) = p(s)$  because  $s$  and  $g$  are independent, we have:
$$
p (x \mid g) = \int_ {s} p (x \mid s, g) p (s \mid g) d s = \int_ {s} \int_ {\pi_ {s: g}} \overbrace {p (x \mid \pi_ {s : g} , s , g)} ^ {\text {l i k e l i h o o d o f s n a p s h o t}} \underbrace {p \left(\pi_ {s : g} \mid s , g\right)} _ {\text {l i k e l i h o o d o f p a t h}} p (s) d \pi_ {s: g} d s \tag {1}
$$
To evaluate the likelihood of a snapshot  $p(x \mid \pi_{s:g}, s, g)$ , we apply the size principle [Tenenbaum, 1998, 1999, Griffiths and Tenenbaum, 2006], analogous to the generic viewpoint assumption in computer vision [Freeman, 1994, Albert and Hoffman, 2000]. In this case, the principle states that the snapshot was equally likely to have been taken anywhere along the path, and therefore the likelihood of a snapshot conditioned on a path is inversely proportional to the length of the path. If  $\delta(x \in \pi)$  indicates whether path  $\pi$  passes through  $x$ , and  $|\pi|$  indicates the length of  $\pi$ , then  $p(x \mid \pi_{s:g}, s, g)$  is given by  $\delta(x \in \pi_{s:g}) |\pi_{s:g}|^{-1}$ . As a consequence of this definition, all paths that do not pass through  $x$  contribute zero probability, and thus can be dropped from the integral.
To evaluate the likelihood of a path  $p(\pi_{s:g} \mid s, g)$ , we apply the principle of rational action: agents are likelier to take actions that maximize their utility [Dennett, 1989, Jara-Ettinger et al., 2016]. There are many ways to formalize this intuition. A common option, which we adopt, is to say that at each step, the agent chooses an action with probability proportional to the softmax over its  $Q$ -values at its current state, with some temperature  $\beta$ . That is,  $p(x \to x' \mid g) \propto \sum_{a} \exp(\beta Q_g(x, a)) T(x \to x' \mid a)$ , where  $T(x \to x' \mid a)$  is the transition probability from  $x$  to  $x'$  if action  $a$  is taken, and  $p(\pi \mid g) \propto \Pi_t p(x_t \to x_{t+1} \mid g)$ .
We now have all the ingredients we need to evaluate  $p(x \mid g)$ . However, to compute it exactly we would need to integrate over all possible initial states  $s$ , and the set of paths  $\pi_{s:g}$ , which could be infinite (agents might wander for arbitrarily long, albeit with vanishingly low probability). To approximate the likelihood in finite time, Lopez-Brau et al. turn to Monte Carlo sampling (Algorithm 1). They rejection-sample paths  $\pi_{s:g}$  by sampling a candidate start state  $s^{(i)} \sim p(s)$ , simulating a rollout of the agent to sample a path  $\pi_{s:g}^{(i)} \sim p(\pi_{s:g}^{(i)} \mid s^{(i)}, g)$ , and then averaging the integrand over these samples. With  $N$  samples, their unbiased likelihood estimator is given by  $\hat{p}(x \mid g) = \frac{1}{N} \sum_{i=1}^{N} \delta(x \in \pi_{s:g}^{(i)}) |\pi_{s:g}^{(i)}|^{-1}$ .
Unfortunately, in practice this scheme is extremely slow: even in a  $7 \times 7$  gridworld with fewer than 49 states (only 2 of which were possible initial states), Lopez-Brau et al. report taking over 300,000 trajectory samples per goal to perform inference. In the rest of this section, we will describe a series of algorithmic enhancements that allow for comparable inference quality with just 10 samples per goal (i.e.  $30,000 \times$  fewer). We will develop our algorithm (Algorithm 2) through three insights.
2.1 First insight: only sample paths through the observed state
Our first insight is that  $\delta(x \in \pi)$  is extremely sparse—most paths likely do not pass through  $x$ , and so most naive path samples contribute zero to the estimator. We would like to only sample paths that pass through  $x$ . Any such path can be partitioned at  $x$  into two portions,  $\pi_{s: x}$  and  $\pi_{x: g}$ . Let us integrate separately over those portions.
$$
p (x \mid g) = \int_ {s} \int_ {\pi_ {s: x}} \int_ {\pi_ {x: g}} \frac {p \left(\pi_ {s : x} \mid g\right) p \left(\pi_ {x : g} \mid g\right)}{\left| \pi_ {s : x} \right| + \left| \pi_ {x : g} \right|} p (s) d \pi_ {x: g} d \pi_ {s: x} d s \tag {2}
$$
This suggests a more efficient Monte Carlo sampling scheme. Rather than rejection-sampling paths  $\pi_{s:g}^{(i)}$  directly from  $s$  to  $g$ , we can independently sample two path segments: a "past" path  $\pi_{s:x}^{(i)}$  from  $s$  to  $x$ , and a "future" path  $\pi_{x:g}^{(i)}$  from  $x$  to  $g$ . Any such path is guaranteed to pass through  $x$ , so no samples are wasted.
However, we now have two new problems. First, it is not clear how to sample paths  $\pi_{s:x}^{(i)}$  from  $s$  to  $x$ , because rollouts of a simulated agent are unlikely to pass through  $x$  on their way to  $g$ . We could imagine using a second planner just to chart paths from  $s$  to  $x$ , but this would require a lot of additional planning work. Second, we still have to sample  $s$ . If the space of initial states is small (e.g., a room only has one or two doors), then this is no issue. However, in practice this space might be very large or even infinite. For example, if you observe someone driving to work in the morning, their home could be anywhere in the city. Furthermore, most of these states might be inaccessible or otherwise implausible, and it would be a waste of computational resources to consider them. In the next section, we show how to solve both of these problems by tracing paths backwards in time.
2.2 Second insight: sample  $\pi_{:x}$  backwards in time
Our second insight is that we can collapse the first two integrals by jointly integrating over the domain of all paths  $\pi_{:x}$  that terminate at  $x$ , no matter where they started from. Say a path  $\pi_{:x}$  begins at  $\pi_{:x}[0]$ . Then, we can rewrite our likelihood as below.
$$
p (x \mid g) = \int_ {\pi : x} \int_ {\pi_ {x: g}} \frac {p \left(\pi_ {: x} \mid g\right) p \left(\pi_ {x : g} \mid g\right)}{\left| \pi_ {: x} \right| + \left| \pi_ {x : g} \right|} p \left(\pi_ {: x} [ 0 ]\right) d \pi_ {x: g} d \pi_ {: x} \tag {3}
$$
This suggests that we should sample  $\pi_{:x}^{(i)}$  backwards through time, starting from  $x$ . No matter how we extend this path, we obtain a valid path from  $\pi_{:x}^{(i)}[0]$  to  $x$ . We illustrate this in Figure 1.
An analogy to path tracing in computer graphics is helpful here. When rendering a 3D scene, a renderer must integrate over all paths of light that begin at a light source in the scene and end at a pixel on the camera's film—a problem formalized by the rendering equation [Kajiya, 1986]. Of course, these paths may be reflected and refracted stochastically by several surface interactions along the way. Rather than starting at one of the millions of light sources in the scene and tracing a ray hoping to eventually reach the camera film, renderers instead start at the camera and trace rays backward into the scene until they inevitably reach a light source.
Similarly, here we trace paths backwards from  $x$  into the past— $s$  corresponds to a light source, each action taken by the agent corresponds to a stochastic surface interaction, and  $x$  corresponds to a pixel on the camera. Indeed, our integral is analogous to the rendering equation, bringing to our disposal the entire Monte Carlo light transport toolbox—a toolbox the rendering community has spent decades developing. (These techniques were pioneered by Veach [1998], though see Pharr et al. [2016] for an accessible review.) The particular ideas we borrow are the following:
Importance sampling The first idea is to be deliberate about how paths are extended backwards in time. We could sample predecessor states uniformly at random—however, that would lead to "unlikely" or irrational paths. Instead, we preferentially select a predecessor state  $x_{\mathrm{prev}}$  based on the likelihood of transitioning to the current state from  $x_{\mathrm{prev}}$ . Then, we re-weight our path sample appropriately so that our estimator remains unbiased (see Algorithm 2, line 14).
Russian roulette termination The next concern is when to stop extending a path into the past. In principle, paths could be infinitely long in some domains. However, at some point paths become so unlikely that extending them is not worth the computational effort of tracing them. An unbiased finite-time solution to this problem is given by the Russian roulette method [Carter and Cashwell, 1975, Arvo and Kirk, 1990]: at each step, we toss a weighted coin, and only continue extending the path if it comes up "heads." Then, we weight subsequent samples appropriately to keep the estimator unbiased (see Algorithm 2, line 11).
Bidirectional path tracing The last concern is that the path may not ever "find" the region of state space where the agent could have started. Consider a setting where the space of possible initial states is large but sparse—for example, if we know the agent started from a red house somewhere in the city. Importance sampling path predecessors with Russian roulette termination is not guaranteed to stop at a red house, whereas forward-sampling paths from a random red house is not guaranteed to pass through the observed state. In computer graphics, this situation is analogous to rendering a scene with many lights that are occluded from the camera. The classic solution is bi-directional path tracing [Lafortune and Willems, 1993, Veach and Guibas, 1995]: first, we "cache" some paths forward-simulated from randomly sampled lights (see Algorithm 3). Then, when backwards-tracing rays, we find opportunities to connect them to a continuation in the cache (see Algorithm 2, line 8).
As a correctness check, we compare the numerical results from our method and the rejection sampling approach of Lopez-Brau et al. [2020] in Appendix B, showing that likelihoods computed from the two methods match closely as expected.
Figure 2: (left) In this example of the "grid" domain, we observe an agent near the blue gem. Even though we do not know where the agent started from, our intuition says that the agent is heading towards the blue gem. (right) In this example of the "keys" domain, we observe an agent right next to the green key. Humans infer that the agent is heading towards the green key because it wants the blue gem. Our algorithm replicates both of these inferences with only 10 samples.
Algorithm 1 Rejection sampling, as in prior work. Compare to our proposed method, Algorithm 2.
Input:  $x$ , the agent's current state (e.g. position in gridworld)  
 $g$ , the hypothesized goal  
 $P(s \rightarrow s' \mid g)$ , the probability the agent will move to  $s'$  from  $s$ $P_{\text{start}}(s)$ , the prior over the agent starting at  $s$   
Output:  $\hat{p}(x|g)$   
1:  $t \gets 0, n \gets 0$ , sample  $x_{\text{current}}$  with probability  $\propto P_{\text{start}}(\cdot)$   
2: while  $x_{\text{current}}$  is not an end state do  
3: if  $x_{\text{current}} = x$  then  
4:  $\lfloor n \gets n + 1$ .  
5: sample  $x_{\text{next}}$  with probability  $p_{\text{choice}} \propto P(x_{\text{current}} \rightarrow \cdot \mid g)$   
6:  $x_{\text{current}} \gets x_{\text{next}}$  and  $t \gets t + 1$   
7: return  $1 / \ell$  if  $n > 0$ , otherwise 0
Algorithm 2 Our bidirectional likelihood sampler
Input:  $x,g,P(s\rightarrow s^{\prime}\mid g),P_{\mathrm{start}}(s)$  as in Algorithm 1  $\alpha$  , the strength of importance sampling  $d$  , an average termination depth for Russian roulette  $C$  , an optional bidirectional path tracing cache (see Algorithm 3)   
Output:  $\widehat{p} (x|g)$    
1:  $\ell \gets 0$    
2:  $t_\mathrm{next}\gets 0$ $x_{\mathrm{current}}\gets x$  ▷ Sample forward from  $x$  to  $g$    
3: while  $x_{\mathrm{current}}$  is not an end state do   
4: sample  $x_{\mathrm{next}}$  with probability  $p_{\mathrm{choice}}\propto P(x_{\mathrm{current}}\to \cdot \mid g)$    
5:  $x_{\mathrm{current}}\gets x_{\mathrm{next}}$  and  $t_\mathrm{next}\gets t_\mathrm{next} + 1$    
6:  $t_{\mathrm{prev}}\gets 1$ $x_{\mathrm{current}}\gets x$ $p_{\pi}\gets 1$  ▷ Sample backwards from  $x$    
7: while true do   
8: if  $x_{\mathrm{current}}\in C$  then ▷ Check BDPT cache for available completions   
9: sample  $(t_{\mathrm{cache}},w)$  from  $C[x_{\mathrm{current}}]$    
10: return  $w\cdot (\# C[x_{\mathrm{current}}] / \# C)\cdot p_{\pi} / (t_{\mathrm{cache}} + t_{\mathrm{prev}} + t_{\mathrm{next}})$    
11: if flip() < 1/d then ▷ Russian roulette termination   
12: return  $P_{\mathrm{start}}(x_{\mathrm{current}})\cdot p_{\pi} / (t_{\mathrm{prev}} + t_{\mathrm{next}})\cdot 1 / (1 / d)$  ▷ Record sample starting at  $x_{\mathrm{current}}$    
13:  $p_{\pi}\gets p_{\pi} / (1 - 1 / d)$  ▷ Apply Russian roulette weight   
14: sample  $x_{\mathrm{prev}}$  with probability  $p_{\mathrm{choice}}\propto \exp (\alpha \cdot P(\cdot \to x_{\mathrm{current}}|g))$  ▷ Choose predecessor   
15:  $p_{\pi}\gets p_{\pi}\cdot P(x_{\mathrm{prev}}\to x_{\mathrm{current}}|g) / p_{\mathrm{choice}}$  ▷ Apply importance sample weight   
16:  $x_{\mathrm{current}}\gets x_{\mathrm{prev}}$  and  $t_{\mathrm{prev}}\gets t_{\mathrm{prev}} + 1$    
17: return  $\ell$
2.3 Third insight: this algorithm allows for on-line planning
Our method is agnostic to the algorithm used to compute the agent's policy, which allows us to choose from any combination of reinforcement-learning or heuristic search approaches. Model-based methods like value iteration require an expensive pre-computation of the value function for all possible goals and states, even though some of those states may not ever be visited by our algorithm. It seems implausible that humans do this, because we make judgements quickly even in new domains. Moreover, adding obviously-irrelevant states should not make the problem harder to solve—we should simply not bother planning from those states until they somehow become relevant.
Motivated by this, in domains that can be expressed as classical planning problems (i.e. domains with a single "goal state" with positive reward), we use an online A* planner instead of precomputing policies by value iteration. We define  $p(x \to x' \mid g)$  to be a softmax over the difference in path costs between  $x$  and  $x'$  to  $g$  as computed by the planner:  $p(x \to x' \mid g) \propto \sum_{a} \exp(\beta(C(x \to g) - C(x' \to g))$ . That is, the agent is likelier to move to states that will bring it closer to the goal. To avoid re-planning from scratch for every evaluation of  $p(x_{t-1} \to x_t \mid g)$ , we run A-star backwards from the goal to the current state. This lets us re-use the bulk of its intermediate computations (known distances, evaluations of the heuristic, etc.) between queries.
Algorithm 3 Grow the bidirectional path tracer's cache (to be called repeatedly)
Input:  $g, P(s \to s' \mid g), P_{\text{start}}(s)$  as in Algorithm 1,  $d$  as in Algorithm 2, and  $C$ , a cache  
1:  $t \gets 0, w \gets 1$ , sample  $x_{\text{current}}$  with probability  $\propto P_{\text{start}}(\cdot)$   
2: while  $x_{\text{current}}$  is not an end state do  
3: add  $(t, d \cdot w)$  to  $C[x_{\text{current}}]$   
4: sample  $x_{\text{next}}$  with probability  $\propto P(x_{\text{current}} \to \cdot \mid g)$   
5: if flip() < 1/d then  
6:  $\square$  break  
7:  $w \gets w / (1 - 1/d)$   
8:  $x_{\text{current}} \gets x_{\text{next}}$  and  $t \gets t + 1$
To evaluate our sampling algorithm, we chose a suite of benchmark domains reflecting the variety of inferences humans make:
Simple gridworld We re-implement the  $7 \times 7$  gridworld domain from Lopez-Brau et al. (Figure 2). The agent seeks one of three gems and at every timestep can move north, south, east or west. The viewer's inference task is to determine which gem the agent seeks. While Lopez-Brau et al. fix two possible starting-points ("entryways") for the agent, our method can optionally relax this constraint and instead have a uniform prior over the start state.
Doors, keys, and gems (multi-stage planning) This is a more advanced  $8 \times 8$  gridworld, inspired by Zhi-Xuan et al. [2020]. The agent is blocked from its gem by doors, which can only be opened if the agent is has the correct keys (the agent always starts out empty-handed). The inference task is to look at a snapshot image and determine which gem the agent seeks. For example, if we observe the snapshot in Figure 2, we might infer that the agent plans to get the green key to obtain the blue gem.
Word blocks (non-spatial) In this domain, the agent spells a word out of the six letter blocks by picking and placing them in stacks on a table. However, they are interrupted (e.g. by a fire alarm) and have to leave the room before finishing. The inference tasks are to look at the blocks left behind and determine (a) which word the agent was trying to spell, and (b) which blocks the agent has touched.
Additional domains Appendix C shows results from additional domains from the cognitive science literature that further reflect the flexibility of our method. We demonstrate domains with partial observability (C.1), multiple agents (C.2), and continuous state spaces with physics (C.3). The breadth of our experiments reflects the typical scope of related work in cognitive science.
3.1 Qualitative analysis
Tables 1 and 2 show some example inferences made by our algorithm. Each cell is colored according to the posterior distribution over goals. If the inference algorithm produced no valid samples for a cell, that cell is marked with an  $\times$  symbol. With just 10 samples, our method's posterior inferences are near-convergent and align well with human responses. In comparison, with 10 samples rejection sampling typically produces extremely noisy predictions, and often simply fails to produce any non-rejected samples at all—indeed, in the blocks domain even 1,000 samples are not always enough to produce a single non-rejected sample.
3.2 Quantitative analysis
We report the total variation  $\mathrm{TV}(x) = \frac{1}{2}\sum_{g_i}|\hat{p} (g_i\mid x) - p(g_i\mid x)|$  between the true posterior and inferences made using 10 samples of both our method and rejection sampling, averaged for 100 trials and across all of the inference tasks in the benchmark. We take the true posterior to be our method's estimate with 1,000 samples, though we also show results if the true posterior is taken to be rejection sampling with 10,000 samples. Our results are shown in Table 3. Across all domains, our algorithm substantially outperforms rejection sampling.