Inferring the Future by Imagining the Past
Authors: Kartik Chandra
Tony Chen
Tzu-Mao Li
Jonathan Ragan-Kelley
Joshua B. Tenenbaum
Venue: NeurIPS 2023 spotlight (2023)
Keywords: cognitive science
cogsci
inverse planning
Bayesian inference
theory of mind
Monte Carlo
inverse reinforcement learning
TL;DR: We model how humans infer an agent's goal from a snapshot of its current state. We frame the problem as Monte Carlo path tracing, which allows us to apply ideas from computer graphics to design a cognitively-plausible sample-efficient algorithm.

Abstract: 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.

Inferring the Future by Imagining the Past
Kartik Chandra* MIT CSAIL
Tony Chen MIT Brain and Cognitive Sciences
Tzu-Mao Li UC San Diego
Jonathan Ragan-Kelley
Joshua Tenenbaum
MIT CSAIL
MIT Brain and Cognitive Sciences
Abstract
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 static snapshot of 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 wide variety of domains, while only using a small, cognitively-plausiin
number of samples. Our key technical insight 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 this seemingly-unreline theory mind task.
1 Introduction
Hemingway's shortest short story simply reads "For sale: baby shoes, never worn." There is no action in this sentence-it it is just description of a static scene-but 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 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? 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 al., 2015, Hamlin al., 2007], make rapid, flexible, and robust theory-of-minn judgments "out of the box": in novel domains they have never seen before, and without extensive pre- training on data. This remarkable ability has motivated decades of work in Al (see Section 4), which has successfully addressed the problem of inferring
agent's goal from trajectory of observed actions in a few- or even one-shot manner [Baker et al., 2017, 2009]. These methods cast the problem Bayesian inference: P(goal actions) x P(actions goal) P(goal), where P(actions goal) is modeled by comparing the observed actions to the optimal actions rational agent would take towards that goal.
of action understanding
However, if we only observe single state snapshot, these methods break down-thers 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- 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.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

area shaded repeesent pia)
$10
y
Suppose you observe an agent at x. If they started from some - p(s), what the likelibood they are going goal g?
Lopez- Brau al, sample and then paths from to However, most paths not pass through.... x...
We instend sample two kimds of paths: forward from and backwards from through the space of possible s's.
Figure 1: How can we infer what agent is trying to do, based on 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 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 up 30, 000 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 tuple (S, A. T. R), where S is the set of possible states, AA is the space of possible actions, T S A x S [0, 1] gives the transition probability, : and 7 E 10 is discount factor. We will specifically consider MDPs that have a single "goal state" g with positive reward, such as 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).
R the reward function
Consider an agent that begins in some initial state p(s). While the agent on its trajectory from s to g. we observe "snapshot" of the agent in some interim state x. Given only (and not s!), our aim is infer p(g x))
Applying Bayes' rule, we have p(g .)) x p( g)p(g), where p(g) is our prior belief about what 9 might be. The likelihood p(r g) can be computed by marginalizing over all possible start states s and all possible paths Ts.g from s to g. where path Ts:g is defined as a sequence of states and 8, = g. Performing this marginalization, and noting that
Ts: ($1,...60 Sn.) such that p(s = p(s) because s and are independent, we have:
likelihood of snapshot
p(æ 9) / p(x g)pss( g)ds - Kun
p( Ta:g 8, 9)
p( Ta:g s.9) p(s)
(1)
likelihood of path
To evaluate the likelihood of snapshot p( Tz:g: 8, g). we apply the size principle [Teneenbaum 1998, 1999, Griffiths and Tenenbaum, 2006], analogous to the generic viewpoint assumption in computer vision [Freeman, 1994, Albert and Hoffmaan, 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 snapshot conditioned indicates whether path и passes through æ, and |T] indicates the length n, then p(z: Ta:g 8, is given by 5(x E w ) | ssggl -1, a consequence of this definition, all paths that do not pass through contribute zero probability, and thus can be dropped from the integral.
path is inversely proportional to the length of the path. If 8(x x)

---

To evaluate the likelihood of a path p(T s, g,g), we apply the principle of rational action: agents are likelier to take actions that maximize their utility [Dennett, 1989, Jara-Ettinger 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 3. That is, p(æ + x' g) x E. exp($Qg(2 a))T(r r' a), where T(x + x' a)) is the transition probability from to x' if action is taken, and p([ g) x [1p((2 + x++1 gg)-
We now have all the ingredients we need to evaluate p(r g). However, to compute it exactly we would need to integrate over all possible initial states 8, and the set of paths =ag- 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 (Al- gorithm 1). They rejection-sampe sample paths Tæ: by sampling a candidate start state s( p(s),
simulating a rollout of the agent to sample path Tay p(msig (f) s()) ,)), and then averaging the
(i)
integrand over these samples. With N samples, their unbiased likelihood estimator is given by
p(x 9 Σi= 8(x
(6) Tayg) [ (1) -1.
Unfortunately, in practice this scheme is extremely slow: even in 7 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 series of algorithmic enhancements that allow for comparable inference quality with just 10 samples per goal (i.e. 30,000 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 8(x E π)) is extremely sparse- most paths likely do not pass through and so most naive path samples contribute zero to the estimator. We would like to only sample paths that pass through æ. Any such path can be partitioned at into two portions, Taw and Tz:g- Let us integrate separately over those portions.
p(1 ( 9) = 1
9)
pl 22g
(g) p(s) diw.gg disw ds
Ta
+
Faag
This suggests more efficient Monte Carlo sampling scheme. Rather than oroectionsamplingg paths (i)
Tag directly from s to g. we can independently sample two path segments: a "past" path (+) from (t) to x. and a "future" path Tang from x to g. Any such path is guaranteed to pass through æ, so no samples are wasted.
However, we now have two new problems. First, it is not clear how to sample paths Tee from to x. because rollouts of simulated agent are unlikely to pass through æ on their way to g. We could imagine using second planner just to chart paths from to z, 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. room only has one 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, show how to solve both of these problems by tracing paths backwards in time.
2.2 Second insight: sample T.. backwards in time
Our second insight is that we can collapse the first two integrals by jointly integrating over the domain of all paths v. that terminate at æ, no matter where they started from. Say a path u.. begins at n. [0]. Then, we can rewrite our likelihood as below.
f /
p(T:r 9) p(Taay g) p( [0]) drzzgg dr
p( 9)
(3)
Tix Mx:gl
This suggests that we should sample Tie backwards through time, starting from &. No matter how extend this path, we obtain a valid path from War [0] to We illustrate this in Figure 1.
(6)

---

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 light source in the scene and end at pixel on the camera's '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 into the past-8 corresponds to light source, each action taken by the agent corresponds to a stochastic surface interaction, and x corresponds to 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 however, that would lead to "unlikely" or irrational paths. Instead, we preferentially select predecessor state prey based on the likelihood of transitioning to the current state from Tprev- 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 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, weight subsequent samples appropriately to keep the estimator unbiased (see Algorithm 2, line 11).
Bidirectional path tracing The last concern 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 large but sparse-for for example, if we know the agent started from 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 random red house is not guaranteed to pass through the observed state. In computer graphics, this situation is analogous to rendering scene with many lights that are occluded from the camera. The classic solution bi-directional puth trucing [Lafortune and Willems, 1993, Veach and Guibas, 1995]: first, we "cache" some paths borward-simuland 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 correctness check, we compare the numerical results from our method and the rejection sampling approach of Lopez-Bras Bgpez-Brau& et al. [2020] in Appendix B, showing that likelihoods computed from the two methods match closely as expected.
Q
Figure 2: (left) In this example of the "grid" domain, we observe 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 Rejection sampling, as in prior work. Compare to our proposed method, Algorithm 2.
Input: 2, the agent's current state (e.g. position in gridworld)
9. the hypothesized goal
P(s s' g). the probability the agent will move to s' from
Pa the prior over the agent starting at
Output: ((rlg)
1: L 6- 0, 4- 0, sample Feurrent with probability & Paan(())
2: while Feurrent is not end state do
3:
if' Teurrent = then
4:
n+
5: sample nexx with probability Pchoice & P( + 11)
6:
Feurrent f- Feex and f- +
7: refurn 1 (f n > 0, otherwise 0
Algorithm 2 Our bidirectional likelihood sampler
Input: 2, 9. P(s s' g). (s)an(s) in Algorithm 1 a, the strength of importance sampling
d, average termination depth for Russian roulette
C, an optional bidirectional path tracing cache (see Algorithm 3)
Output: p(x1g)
1: 8
0-0
2:
Inext 0, Feurrent
Sample forward from g
3: while Zeurrent is not an end state do
4: sample nnext with probability Pchoice P(Teureet (1
5: Teurreat f- Enext and Enext f- nnxx +
6: fprev < Feurrent + +x, Pr +1
Sample backwards from
7:
while true do
8: if Feurrent E C then
b Check BDPT cache for available completions
9: sample (teache w) from Cl
return w (#C((aaumen))) HCC) Pr (Vanche + 1prev + (nexs)
11:
if flip() < 1/d then
to Russian roulette termination
return Patan(Faureen (Feurrem) Pn/ (Irrev + (nex)) 1/(1/d)
b Record sample starting al Faurrent
13:
Pa Px / (1 1/d)
>> Apply Russian roulette weight
14: sample &rrev with probability Pcholce x exp(a P(- + Teuneat ()) & Choose predecessor
15:
Pa Px (xEpre - Teurrent g) Pehoice
Apply importance sample weight
16:
Tecurent f- prev and prev +- [prev +
17:
return E
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 reinforeementleamming 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-irrelle 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 classical planning problems (i.e. domains with single "goal state" with positive reward), we use an online A planner instead of precomputing policies by value iteration. We define p(æ + z' g) to be softmax over the difference path costs between and x' to computed by the planner: p(a - x' g) En exp (B(C(x + g) C(x' + 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(z+-1 - - g). we run A-star backwards from the goal to the current state. This lets us re-use the bulk of intermediate computations (known distances, evaluations of the heuristic, etc.) between queries.

---

Word blocks (non-spatial) In this domain, the agent spells word out of the six letter blocks by picking and placing them in stacks on table. However, they are interrupted (e.g. by a fire alarm) and have leave the room before finishing. The inference tasks to look the blocks left behind and Algorithm 3 Grow the bidirectional path tracer's cache (to be called repeatedly)
Input: 9. P(s + s' g). (s) as Algorithm 1, d as in Algorithm 2, and C, a cache
I: 4 0, w t 1, sample Zeurrent with probability x Pstur( (-)
2: while Zeurrent is not an end state do
3:
add (t,d d w) to C ((Eeurrent)
4: sample Tnext with probability x P(Jcuurent + 11
5: if flip() < 1/dd then
6: break
7: W <- / (1-1/d)
8:
Teurrent f- nexx and t t- 1+1
3 Experiments
To evaluate our sampling algorithm, we chose suite of benchmark domains reflecting the variety of inferences humans make:
Simple gridworld We re-implement the 7 x gridworld domain from Lopez-Brau 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 uniform prior over the start state.
Doors, keys, and gems (multi-stage planning) This is by Zhi-Xun Chi-Xuan et al. [2020]. The agent blocked from its gem by doors, which can only be opened the agent is has the correct keys (the agent always starts empty hhadde-. The inference task is to look at a snapshot image and determine which gem the agent seeks. For example, we observe the snapshot in Figure 2, we might infer that the agent plans get the green key obtain the blue gem.
advanced 8 gridworld, inspired
determine (a) which word the agent was trying to spell, and (b) which blocks the agent has touched.
Additional domains Appendix 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 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. the inference algorithm produced no valid samples for cell, that cell is marked with an x 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 all-indeed, the blocks domain even 1,000 samples are not always enough to produce single non-rejected sample.
3.2 Quantitative analysis
xx) xz) 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 also show results if the true posterior is taken to be rejection sampling with 10,000 samples. Our results are shown in Table Across all domains, our algorithm substantially outperforms rejection sampling.
We report the total variation TV = 16(::

---

Table I: Qualitative comparison of inference algorithms. Cells are colored based on the posterior distribution over goals if the agent observed in that cell. Cells marked had all samples rejected. Gray shaded cells were excluded from analysis. 2 We show results for 10 samples and 1,000 (1k) samples, comparing rejection sampling as described by Lopez-Bre Brau et al. [2020]. our method, and human subjects. We produce neav-convergent human-like inferences with only 10 samples.
Rejection (10)
Rejection (1k)
Task
Ours (10)
Ours (Ik)
Humans
Previously shown to be matched by model; [Lopez-
Gridworld (two doors, in prior work)
al., 2022]
Gridworld (start anywhere)
Keys (holding no key)
Keys (holding pink key)
Keys (holding green key)
Pigoal green gem (x)
0.0
0.2
0.6
0.8
1.0
pfgoal blue
pigoal red gem
piacal gem
plgoal gam X)
3.3 Comparison to human judgements
We recruited N 200 participants and collected judgements for variety of "snapshots" in each of our domains (see experimental design in Appendix A). We found that our model predicts human intuitions well (see last columns of Tables and 2), with strong correlations across domains (see Figure 3). Thus, our work not only replicates the findings of previous work [Lopez- et al., 2020, 2022), but also shows that those findings continue hold in domains that previous inference algorithms could not scale to.
"Shaded cells in Row 3 are excluded because they are inaccessible without holding key. Shaded cells in Row 4 excluded because if the agent already picked up the pink key earlier, there is no reason then move

---

Table 2: Qualitative comparison of inference algorithms. Blocks are colored according to inferred probability of that block having been touched by the person stacking the blocks (red high probability, blue is low). We show results for 10 samples and 1,000 (1k) samples, comparing rejection sampling, our method, and human subjects. We produce near-convergent inferences with only 10 samples. In comparison, rejection sampling is unable to make any inference with 10 samples, and sometimes even fails with 1,000 samples. When it succeeds, its predictions are high-variance and overconfident.
Rejection (10)
Rejection (1k)
Ours (10)
Ours (1k)
Humans
(all samples rejected)
G
(all samples rejected)
(all samples
(all samples rejected)
rejected)
(all samples
rejected)
(all samples
rejected)
(all samples rejected)
(all samples rejected)
G
00
0.6
1.0
pitsuched
Related work
Cognitive science Human social cognition and "theory of mind" are well-modeled by Bayesian inverse planning [Baker et al., 2009, Jara-Fitinner, 2019, Baker et al., 2017], which infers an agent's goals from its observed actions. Inverse planning predicts human inferences about multiple agents et al., 201 6, Wu et al., 2021], social scenes [Ullman et al., 2009, Netanyahu et al., 2021], emotion [Ong et al., 2015], and agents who make mistakes [Zhi-Xuan al., 2020].
A growing body of work studies how people make inferences about the past from a "snapshot" of present physical evidence [Smith and Vul, 2014, Gerstenberg al., 2021, Lopez-Bree Brau and Jara- Ettinger, 2020, Paulun et al., 2015, Yildirim et al., 2017, 2019]. Recently, Lopez- Brau et al. [2020,
far away from the doors. Similarly, in Row S, there key. We found that querying repeatedly at those extreme locations confuses human participants about the task setup. The single cell above the green key in Row is an exception we wished to test: one might think from that position that the agent is trying to collect both keys, leading to uncertainty about the goal. But people (and the model) agree that the agent still wants the red gem.
that far left after picking up the green

---

Table 3: Quantitative comparison of inference algorithms (see Section 3.2). We show the average total variation distance (TV) of 10-sample posterior estimate against converged posterior, averaged over 100 trials. Lower is better. Our method does significantly better than rejection sampling [Lopez- Brae Brau al., 2022] for each task, whether the converged "ground truth" is computed with 1,000 samples using our method, or 10,000 samples using rejection sampling.
Benchmark (TV against Ours @ 1k)
Rejection @ 10
Ours @ 10
Grid (two doors)
0.063
0.0257
Grid (starting anywhere)
0.159
0.0538
Keys (observed holding key)
0.409
0.108
Keys (observed holding pink key)
0.388
0.157
Keys (observed holding green key)
0.381
0.119
Blocks
0.985
0.358
Benchmark (TV against Rejection @ 10k)
Rejection @ 10
Ours @ 10
Grid (two doors)
0.0705
0.0336
Grid (starting anywhere)
0.153
0.0836
Keys (observed holding key)
0.672
0.410
Keys (observed holding pink key)
0.456
0.382
Keys (observed holding green key)
0.463
0.312
Blocks
Rejection @ 10k did not converge
2022] asked how people make inferences about the past and future of agents from static physical evidence they leave behind, which is something even children can do [Pelz et al., 2020, Jacobs et al., 2021]. We build on this line of work by dramatically accelerating inference, and extending it to more sophisticated domains where previous methods could not scale.
Artificial intelligence Bayesian approaches have been successful in approaches to plan recognition, the problem of inferring an agent's s plan from observed actions [Ramirez and Geffner, 2009, 2010, Sohrabi et al., 2016, Charniak and Goldman, 1993]. Our work provides method for plan recognition from a single state snapshot, with no need observe actions.
The reinforcement learning community has long sought to learn an agent's 's reward function by observing the actions it takes via inverse reinforcement learning or IRL. [Ng et al., 2000, Arora and Doshi, 2021, Ziebart et al., 2008]. Recently, Shah et al. [2019] proposed "IRL from single state." Their method, "Reinforcement Learning by Simulating the Past" (RLSP), learns reward functions based on single observation of a human in the environment, assuming that the observation is taken end of a finite-horizon MDP fixed horizon T. We build on RLSP in three ways: (1) RLSP highly sensitive to the time horizon hyperparameter 7. We dispense with the fixed-horizon assumption altogether, integrating over past trajectories of all possible lengths. (2) Unlike Shah et al., we do not
Grid, r=0.94
Keys, r=0.81
Blocks, r=0.76
1.0
1.0
1.0
0.8
humans
0.6
Humans
0.6
humans
0.6
0.2
0.2
0.2
0.25
1.00
0.00
0.25
0.75 Model
0.50
1.00
0.25
0.50 model
0.75
1.00
model
Figure 3: Our inferences correlate well with human responses. In the "grid" plot, colors represent goal gem colors. In the "keys" plot, colors represent which keys (if any) the agent seen holding, and we show poaall is blue gem .r. In the "blocks" plot, show p(touched x)) for each block.

---

Sampling over goals In this paper. we showed how to scale inference to large set of possible initial states. However, we compute the posterior by enumeration over all possible goals, which takes linear time in the size of the space of goals. To scale to larger goal spaces as well, it may be possible use Markov Chain Monte Carlo to sample goals conditioned on the observed state, similar to what Veach and Guibas [1997] propose for path tracing in computer graphics. This may seem challenging because we only stochastically estimate the likelihood p(æ g). However, because our estimator transitions [Andrieu and assume the snapshot was taken at the end of the agent's journey-the agent can be observed partway. (3) Our sampling-basee method scales to significantly larger domains, because we do not have to integrate exhaustively over all possible trajectories.
The Deep RLSP algorithm of Lindner et al. [2020] also addresses point (3) above by sampling past trajectories. However, while Deep RLSP's -akiimu-likellioodd formulation designed to recover robot policies, our Bayesian approach is designed specifically to model human intuitions about agents' goals. Thus, unlike Deep RLSP, our method incorporates priors over start states in principled way, and produces human-like uncertainty quantification.
5 Limitations and future work
unbiased, it is still possible to use it to compute valid Metropolis- Roberts, 2009].
Cognitive plausibility Our method's sample efficiency suggests that it may resemble how humans actually do this task, similar how few-shot algorithms have shown promise in modeling human behavior in other domains [Vul et al., 2014]. Following previous work using eye-tracking studies to investigate cognitive processes underlying simulation and planning [Gerstenberg al., 2017], we hope to use eye-tracking to compare human strategies to our algorithm. In the language of Marr [1982], this would allow us to go beyond the computational account Lopez- Brau al. [2020] and take first step towards an algorithmic account of inverse planning with snapshot states.
Theoretical analysis The time complexity taking a single sample using our algorithm straight- forward: it scales linearly with the Russian roulette depth and the number of possible goals (although sub-linear runtime number of goals is possible through sampling as described above). However, the more pertinent question is that of sample complexity. While our algorithm is unbiased, analyzing rate of convergence as a function of the size of the sample space and complexity of the domain are topics for future work.
Investigating potential applications While our goal in this paper was to model how humans perform this inference task, extensions of our method could be applied to variety of practical problems: some examples include understanding dynamic action in photographs of sports, interpreting step-by-step instruction manuals, and forensic reconstruction based available static evidence.
From analysis synthesis We began this paper with Hemingway's short story-and indeed, more generally, artists have long depicted dynamic action through static scenes [McCloud, 1993]. In future work we hope to consider the inverse problem of designing interesting, evocative, or ambiguous scenes by optimizing over inference [Chandra et al., 2023a,b, 2022], and designing expressive and numan-interpretable robotic "gestures" similar to legible planning [Dragan et al., 2013].
Conclusion
In this paper, we offered cognitively-plausibe algorithm for making inferences about the past and the future based on the observed present. Building on prior work from both the cognitive science and AI communities, and drawing inspiration from the Monte Carlo rendering literature in computer graphics, we presented a highly sample-efficient method for performing such inferences. Finally, we showed that our method matched human intuitions on variety of challenging inference tasks to which previous methods could not scale.
10

---

Acknowledgments and Disclosure of Funding
We thank Joe Kwon, Tan Zhi-Xuan, Alicia Chen, Max Siegel, Lionel Wong, and John Cherian for thoughtful conversations as we developed these ideas. This research was funded by NSF grants #CCF-1231212 #CCF--173444 and #2238839, and ONR grant #00010803. Additionally, KC was supported by the Hertz Foundation, the Paul and Daisy Soros Fellowship. and an NSF Graduate Research Fellowship under grant #1745302, and TC was supported by an NDSEG Fellowship.
References
Marc K Albert and Donald D Hoffman. The generic-viewpoint assumption and illusory contours. Perception, 29(3):303-311 312, 2000. URL https:/jjurmallssan wwgepub.ccmm////ppdrrii 10.1008//311
Luke Anderson, Tzu-Mao Li, Jaakko Lehtinen, and Frédo Durand. Aether: An embedded domain specific sampling language for monte carlo rendering. ACM Transactions on Graphics (TOG), 36 (4):1-16, 2017. URL
Christophe Andrieu and Gareth O Roberts. The pseudo-marginge approach for efficient monte carlo computations. 2009. URL https//wwwwijsorrrr/ttta///3333
Saurabh Arora and Prashant Doshi. A survey of inverse reinforcement learning: Challenges, methods and progress. Artificial Intelligence, 297:103500, 2021. URL https//www.sccenedirrc.ccommmm
ssience/atieeemmmp///00077711111
James Arvo and David Kirk. Particle transport and image synthesis. Proceedings of the 17th annual conference on Computer graphics and interactive techniques, pages 63 66, 1990. URL
htts::/dl. htt.//ddllcm....o www.lllmmmmmmoo//...... www.w.//woo//....////// (45/97880099888)
Chris L Baker, Noah D Goodman, and Joshua B Tenenbaum. Theory-basee based social goal inference. Proceedings of the thirtieth annual conference of the cognitive science society, pages 1447- 1452. Cognitive Science Society Austin, TX, 2008. URL http:////oooo..ttttttooo..cd///e BakerEtA12008-Copee Cog
Chris I. Baker, Rehecca Saxe, and Joshua Tenenbaum. Action understanding inverse plan- ning. Cognition, 13(3):329-349, 2009. URL http:://wwwssdeeeeeeeeee.ccceeee S0010027709001007..
Chris L Baker, Julian Jara Ettinger, Rebecca Saxe, and Joshua B Tenenbaum. Rational quantitative attribution of beliefs, desires and percepts human mentalizing. Nature Human Behaviour, 1(4): 0064, 2017. URL
Leland Lavele Carter and Edmond Darrell Cashwell. Particle-transport simulation with the monte carlo method. Technical report, Los Alamos National (LANL,,,),, Los Alamos, NM (United States), 1975.
Kartik Chandra, Tzu-Mao Li, Joshua Tenenbaum, and Jonathan Ragan-Kelley. Designing percep- tual puzzles by differentiating probabilistic programs. In ACM SIGGRAPH 2022 Conference Proceedings, pages 1-9, 2022.
Kartik Chandra, Tzu-Mao Li, Joshua Tenenbaum. and Jonathan Ragan Kelllyy Acting as inverse inverse planning. In Special Interest Group on Computer Graphics and Interactive Techniques Conference Proceedings (SIGGRAPH '23 Conference Proceedings), aug 2023a. doi: 10.1145. 33884332.591.510.
Kartik Chandra, Tzu-Mao Li, Joshua Tenenbaum, and Jonathan Ragan-Kelley Storytelling as inverse inverse planning. Proceedings of the Annual Meeting of the Cognitive Science Society, volume 45, 2023b.
11

---

Eugene Charniak and Robert P Goldman. A bayesian model of plan recogmitionn, Artificial Intelligence. 64(1):53-79, 1993, URL Baapadee_o_ppaaeeeoogggeeeee 222330914
sacpe_oodelooooeeee of _lann wwopaiooonnnkss/77/0/09991266660 A- ooyeesianmmodel-off pla.oedeenooopla---cocoottiioooo
Baaaoliii__aayssianaaaaaaaa Sanyesian_mode___o
Marco F Cusumano-Towinn Feras A Saad, Alexander K Lew, and Vikash K Mansinghka. Gen: general-purpose probabilistic programming system with programmable inference. Proceedings of the 40th acm sigplan conference programming language design and implementation, pages 221-235 2019.
Daniel C Dennett. The intentional stance. MIT press, 1989.
Anca D Dragan, Kenton CT Lee, and Siddhartha S Srinivasa. Legibility and predictability of robot motion. In 2013 8th ACM/IEEE International Conference on Human-Robol Interaction (HRI), pages 301-308. IEFE, 2013. URL htt.:////..eeeeooccc...ee/////// pdf.
William T Freeman. The generic viewpoint assumption in framework for visual perception. Nature, 368(6471):542-544 1994. URL
György Gergely and Gergely Csibra. Teleological reasoning in infancy: The naive theory of rational action. Trends in cognitive sciences, 7(7):287-292, 2003.
Tobias Gerstenberg, Matthew F Peterson, Noah D Goodman, David Lagnado, and Joshua B Tenenbaum. Eye-tracking causality. Psychological science, 28(12):131-1174; 2017.
Tobias Gerstenberg, Max Siegel, and Joshua Tenenbaum. What happened? reconstructing the past through vision and sound. 2021. URL https://psyarxi..com/fffk
Thomas L Griffiths and Joshua B Tenenbaum. Optimal predictions in everyday cognition. Psy- chological science, 17(9)766-7773 2006. URL http://///mmmmmlssgagpuuuoooooop 1467-9280.2006.017880
J Kiley Hamlin, Karen Wynn, and Paul Bloom, Social evaluation by preverbal infants. Nature, 450 (7169):557-559, 2007.
Fritz. Heider and Marianne Simmel. An experimental study of apparent behavior. The American journal of psychology, 57(2):243-259, 1944.
Colin Jacobs, Michael Lopez-Brau and Julian Jara- Etttinet What happened here? children integrate physical reasoning to infer actions from indirect evidence. In Proceedings of the 43rd Annual Meeting of the Cognitive Science Society, volume 43, 2021.
Julian Jara-Etinnn Theory mind as inverse reinforcement learning. Current Opinion Behavioral Sciences, 29: 105-110, 2019.
Julian Jatanettinger, Hyowon Gweon, Joshua B Tenenbaum, and Laura E Schulz. Children's 'sss understanding of the costs and rewards underlying rational action. Cognition, 140:14-23, 2015.
Julian Jara- JttanEEtinger, Hyowon Gweon, Laura E Schulz, and Joshua B Tenenbaum. The naïve util- ity calculus: Computational principles underlying commonsense psychology. Trends in cog- nitive sciences, 20(8):589-604, 2016. URL
$136466131630053.
James T Kajiya. The rendering equation. In Proceedings of the 13th annual conference on Computer graphics and interactive techniques, pages 143-150, 1986. URL https://dl.acm/ http://dd.acc...//////// https///dl.cmmmorr//////////// 1145/1592..15022
Max Kleiman- Nleiman-Weiner, Mark K Ho, Joseph I. Austerweil, Michael L Littman, and Joshua B Tenen- baum. Coordinate to cooperate or compete: abstract goals and joint intentions social interaction. In CogSci, 2016.
12

---

Eric Lafortune and Yves Willems. Bi-directional path tracing. In Compugraphics' 93, pages 145-153, 1993. URL httpss//citeseerx.is.... c480a06abl 66a 28
David Lindner, Rohin Shah, Pieter Abbeel, and Anca Dragan. Learning what do by simulating the past. In International Conference on Learning Representations, 2020.
Michael Lopez-Brand Brau and Julian Jara- Eitingr, Physical pragmaties: Inferring the social meaning of objects. 2020. URL https:/pssaarii..oo/mm////
Michael Lopez-Brau, Joseph Kwon, and Julian Jara- lara-Ettinger Mental state inference from indirect evi- dence through bayesian event reconstruction. In CogSci, 2020. URL httpss//coonniivseeeeesscceee ogg/oogsii2//ppprr0008//0000 pdf.
Michael Lopez-Brau, Joseph Kwon, and Julian Jara-Eitinger. Social inferences from physical evidence via bayesian event reconstruction. Journal of Experimental Psychology: General, 2022.
URL https//s/pssyaii.com/////
David Marr. Vision: A computational investigation into the human representation and processing of visual information. MIT press, 1982.
Scott McCloud. Understanding comics: The invisible art. Northampton, Mass, 7:4, 1993.
Aviv Netanyahu, Tianmin Shu, Boris Katz, Andrei Barbu, and Joshua B Tenenbaum. Phase: mmysically-grounder abstract social events for machine social perception. In Proceedings of the AAAI Conference Artificial Intelligence, volume 35, pages 845-853, 2021. URL https://ojs.at aaai. wuggimmeeppp//AAA///attttieeeeie
Andrew Y Ng, Stuart Russell, et al. Algorithms for inverse reinforcement learning. In lcml, volume page 2, 2000. URL http::/wwwwweee....aarrddddeeede www.://wwww..ceeeooooeee////nnne jrl000. pdf.
Desmond C Ong, Jamil Zaki, and Noah D Goodman. Affective cognition: Exploring lay theories of emotion. Cognition, 143:141-162, 2015.
Stefan Palan and Christian Schitter. Prolific. subject pool for online experiments. Journal of Behavioral and Experimental Finance, 17:22-27, 2018.
Vivian C Paulun, Takahiro Kawabe, Shin's ya Nishida, and Roland W Fleming. Seeing liquids from static snapshots. Vision research, 115:1631774 2015.
Madeline Pelz, Laura Schulz, and Julian Jara Ettingr.
The signature all things: Children infer knowledge states from static images. In Proceedings of the 42nd Annual Meeting of the Cognitive Science Society, 2020. URL http://psyarxiv.com/62222
Matt Pharr, Wenzel Jakob, and Greg Humphreys. Physically based rendering: From theory to implementation. Morgan Kaufmann, 2016. URL
Integration.
Antonin Raffin, Ashley Hill, Maximilian Ernestus, Adam Gleave, Anssi Kanervisto, and Noah Dormann. Stable baselines3, 2019.
Miquel Ramırez and Hector Geffner. Plan recognition planning. In Proceedings of the s international joint conference on Artifical intelligence. Morgan Kaufmann Publishers Inc, pages 1778-1783. Citeseer, 2009. URL httpss//citesseex.if /suu.etu/doccmmee/reeidddzz) &ttp:pdddddpr@@ppepdpepet httpss//ccceeeerr...ssss.eeeeeee
Miquel Ramırez and Hector Geffner. Probabilistic plan recognition using off-the-shelf classical planners. In Proceedings of the Conference of the Association for the Advancement of Artificial Intelligence (AAAI 2010), pages 1121-1126. Citeseer, 2010. URL https:///ddnn.... onai.or//ojj/77745 11151-1-1174-11221111
13

---

htt.c//000044aaa11c7777999999999 John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint urXiv:1707.06347 2017.
Rohin Shah, Dmitrii Krasheninnikov, Jordan Alexander, Pieter Abbeel, and Anca Dragan. Preferences implicit in the state of the world. International Conference on Learning Representations (ICLR), 2019. URL https://arxivang///d//99222111
Kevin Smith and Edward Vul. Looking forwards and backwards: Similarities and differences in prediction and retrodiction. In Proceedings of the Annual Meeting the Cognitive Science Society, volume 36, 2014.
Shirin Sohrabi, Anton V Riabov, and Octavian Udrea. Plan recognition as planning revisited. In IJCAI, pages 3258-3264. New York, NY, 2016. URL http://www...oooooooooooooooooet Sobrabi-UJCAI-6.pp
Victoria Southgate and Gergely Csibra. Inferring the outcome of an ongoing novel action at 13 months. Developmental psychology, 45(6):1794, 2009.
Joshua Tenenbaum. Bayesian modeling of human concept learning. Advances in neural in- formation processing systems, 11, 1998. URL https//proceeiing.nnrii..cco 00109%%caSafffaaaccccc2222222663 Paper.....
Joshua Tenenbaum. Rules and similarity in concept learning. Advances in neural infor- mation processing systems, 12, 1999. URL Papead-d....
Christopher D Twigg and Doug L James. Backward steps in rigid body simulation. In ACM acm. wwtp://o/i//f///..11411333900900
SIGGRAPH 2008 papers, pages 1-10. 2008. URL 1360624.
Tomer Ullman, Chris Baker, Owen Macindoe, Owain Evans, Noah Goodman, and Joshua Tenenbaum. Help or hinder: Bayesian models of social goal inference. Advances in neural information processing systems, 22, 2009.
Eric Veach. Robust Monte Carlo methods for light transport simulation. PhD thesis. Stanford University, 1998. URL. http://wwwwppoposoooooooooooooop fromopenview-true.
Eric Veach and Leonidas Guibas. Bidirectional estimators for light transport. In Photorealistic Ren- dering Techniques, pages 145-167. Springer, 1995. URL lc815d05e9442115574
repid-repl http:///ppeppooocccoooooooooooed
Eric Veach and Leonidas J Guibas. Metropolis light transport. In Proceedings of the 24th annual conference Computer graphics and inferactive techniques, pages 65-76, 1997.
Edward Vul, Noah Goodman, Thomas L Griffiths, and Joshua B Tenenbaum. One and done? optimal decisions from very few samples. Cognitive science, 38(4):599-633,200) 2014, URL https:
1/11///o...12 12101.
Sarah A Wu, Rose E Wang, James A Evans, Joshua B Tenenbaum, David C Parkes, and Max Kleiman- Weiner. Too many cooks: Bayesian inference for coordinating multi-agent collaboration. Topics in Cognitive Science, 13(2):414 432, 2021.
lIker Yildirim, Tobias Gerstenberg, Basil Saeed, Mare Toussaint, and Josh Tenenbaum. Physical problem solving: Joint planning with symbolic, geometric, and dynamic constraints. arXiv preprint 1rXiv11707.82211 2017.
IIker Yildirim, Basil Saeed, Grace Sennett-Pierre Tobias Gerstenberg. Joshua Tenenbaum, and Hyowon Gweon. Explaining intuitive difficulty judgments by modeling physical effort and risk. arXiv preprint 1454441905.04445, 2019.
14

---

Tan Zhi-Xuan, Jordyn Mann, Tom Silver, Josh Tenenbaum, and Vikash Mansinghka. Online bayesian goal inference for boundedly rational planning agents. Advances in neural information processing systems, 33:19238-19250, 2020. URL
Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In AAAI, volume 8, pages 1433 1438. Chicago, IL, USA, 2008.
15

---

A Experimental design
For each our experiments reported in Section 3.3, we recruited N = 200 participants from Prolific paid $15 per hour ($1.25 total for blocks and grid
[Palan and Schitter, 2018]. Participants domains, and S2.00 for the keys domain), and our experiments were conducted with IRB approval.
Participants were first familiarized with the environment, through both text instructions and sample video of agent performing the task in the domain. Then, they were told that their objective was to infer the agent's s goal from single snapshot. They answered several questions to check their comprehension of both domain and the task they were asked to perform, and were not allowed to continue unless they answered the comprehension questions correctly. The full experimental design available in HTML format in the supplementary materials. No data was excluded from our analyses.
B Numerical test of correctness
Programming sophisticated importance sampling routines is challenging and bug-prone engineering effort [Cusumano- Towner et al., 2019, Anderson et al., 2017, Pharr et al., 2016]. To test that our algorithm unbiased, i.e. that it produces correct likelihoods in expectation, computed by rejection sampling and our sampler using converged estimates (25,000 samples each). For this experiment we used a uniform
compared likelihoods
4 grid-world, with the prior on start states being uniform along the first row (æ 0) and the goal being the far corner (3, 3). The results of this experiment are shown Figure Our estimator has dramatically different implementation than rejection sampling (compare Algorithms 1 and 2). However, the computed likelihoods are indistinguishable at every cell the grid, even in "comer-cass" cells such the goal cell itself. This provides strong check that our algorithm and its implementation are both indeed correct.
C Additional domains
We used our algorithm to perform inferences in three additional domains. The purpose of these domains is to show the remarkable flexibility our method: how it can make interesting inferences in wide variety of settings. Though did not collect human subject data for these domains, show results for cases where the inference task relatively straightforward.
Our sampler is unbiased and precisely matches rejection sampling
0.175
Ours
0.175
Ours
BDPT Rejection
BDPT
Likellhood pix [g=(3,3 3))
0.150
0.150
0.125
Likellhood (ours)
0.125
0.100
0.100
0.075
0.075
0.050
0.050
0.025
0.025
0.000
0.000
(0. 0)
(0. 1)
(0. 2)
(0. 3)
Observed location x (2. 3)
(1.1)
[1. 2)
(1. 3)
(2. 0)
[2. 1)
(2.2)
(3. 0)
(3. 1)
[3, 2)
[3. 3)
0.00
0.05
0.10
0.15
0.20
[1.
Likelihood (rejection)
Figure 4: Our sampler's likelihoods precisely match rejection sampling, with and without bidirectional
path tracing, giving a strong numerical check of our method's correctness (Appendix B).
16

---

K=L=M
KRL
LeKtn
Micke
L-MicK
K-MKi
K=LxM
KxLxMM
K=M-L
LcKKM
L<M-K
McKeL
McLcK
70 K7 UP MZ 7MM KLL. KIM UM MILL
beliof
Figure 5: The student is observed heading south around the wall. A rational inference that the student started at home, and went around the wall check what the far food truck was. Seeing that was Lebanese and not Mexican (their favorite), the student disappointedly turns around to make peace with the nearby Korean food. As shown the heatmap to the right, our model captures this joint belief-desire inference, predicting that the student now knows what is at both trucks, and reconstructing the student's likely preference ordering belief label "K X? ?" means that the student thinks the south-west parking spot has a Korean food truck parked, but is unsure about the north-east east parking spot. See Appendix C.I.
the three cuisines. Note: the
C.1 Food trucks (joint belief/desire inference)
The food trucks domain, taken from the cognitive science literature [Baker et al., 2017], is Partially Observable Markov Decision Process (POMDP). It consists of a 5 x 10 gridworld with an opaque wall in the middle. A hungry graduate student wakes up home (one side of the wall) and wishes to cat at food truck. There are two parking spots where food trucks usually park, and three kinds of food trucks that could be parked each of those spots: Korean, Lebanese, and Mexican (K, L, and M). The graduate student might have preferences among the cuisines, but might also be uncertain about which trucks are parked at each spot today. Thus, they might engage in behavior by looking behind the opaque wall, and then choosing a food truck to walk to based on their preferences. The inference task is to determine (a) the student's preferences over food trucks, and (b) the student's 's (current) belief state about which truck is at each parking spot.
commattio--seeiing
Using this domain, Baker al.'s inverse planning model was able jointly infer the student's beliefs and desires from an observed trajectory; those inferences closely matched responses from human subjects. Here, we perform the same type of inference, but from single observed snapshot.
For example, in the example in Figure 5, the student is observed moving south next to the wall. A Korean food truck is parked in the southwest parking spot, and a Lebanese food truck is parked in the northeast spot. Seeing this scene, reasonable inference is that the student went looking around the wall to see if the Mexican food truck (their favorite) was Lebanese food instead, the student turns around and makes peace with the nearby Korean food. Indeed, model captures this inference: in the joint posterior distribution over both beliefs and desires, our model confident that the student now knows that the northeast truck has Lebanese food, and furthermore that the student's 's favorite food is Mexican.
parked on the other side. Seeing that it
more sophisticated inference emerges if the student is observed moving north instead of south (Figure 6). Now, reasonable inference is that the student dislikes Korean food, and is going around the wall to check what is at the other truck. The model captures this: favors the hypothesis that the student is unsure what is the northeast truck, and also places high weight on Korean being the least favorite food option.
However, as is visible on the right half of the heatmap, the model also places some weight on the possibility that the student knows that there Lebanese food and prefers it, or that the student (mistakenly) believes there is Mexican food and prefers that.
17

---

McKe
MeL-X
40
1ot
ny.
MyT PpM kil
<ih IK 12M My. 19)
Current bele!
Figure 6: Here, the student is observed going north instead of south. A more sophisticated inference emerges, showing that the student is likely uncertain about which truck is parked behind the wall. See Appendix C.1
1.0
Blue agent is thief (Pink agent guard)
Pink agent thief (Blue agent is guard)
0.8
Model prediction
0.6
0.4
0.2
0.0
Bottom
Bottom right
left
left
right
Goal
Figure 7: Two agents are observed by security camera in an art museum. Who is the guard, who is the thief, and where the thief trying to escape to? Our model predicts that the guard the blue agent, the thief is the pink agent, and that the exit is in the bottom right. See Appendix C.2.
C.2 Heist (multi-agent domain)
In this multi-agent domain inspired by classic stimuli in cognitive science [Baker et al., 2008,
Southgate and Csibra, 2009, Heider and Simmel, 1944], two agents-blue and pink- a 7 x 7 gridworld representing an art museum. One of the agents is a "thief," whose objective is to escape the museum by reaching the exit, and the other is "guard." whose objective is to catch the thief. There are four doors in the room, only one of which is an exit, and the rest which are dead ends. Both agents know which door is the exit, but this information is not visible the observer (all doors rendered identically). The inference tasks are to look at snapshot of the two agents and jointly infer (a) which agent is the thief and which is the guard, and (b) which door is the exit.
the example in Figure 7, it is clear from snapshot that blue agent is the guard and is chasing the pink agent, the thief, to the bottom-right corner. The model reproduces this inference, though also acknowledges the possibility that the thief might actually heading onward past the bottom-right, to the bottom-left corner instead.
18

---

1.0
Blue agent is thief (Pink agent is guard)
Pink agent is thief (Blue agent is guard)
Model prediction
0.4
0.2
0.0
Top
Bottom left
Top right
Bottom right
left
Goal
1.0
Blue agent is thief
(Pink agent is guard)
Pink agent is thief (Blue agent is guard)
0.8
Model prediction
0.6
0.4
0.2
Top
Bottom
Top right
Bottom right
left
left
Goal
Figure 8: these examples, is unclear who the guard and thief are however, is clear where the exit is. The model reproduces this uncertainty desired. See Appendix C.2.
The next two examples (Figure 8) are ambiguous cases: the two agents are in symmetric positions, so it is unclear who is who. Here, the model can determine with high confidence where the exit is, but remains uncertain about who is the thief and who is the guard.
Finally, in the last example (Figure 9), it is unclear whether blue guard is blocking pink thief from heading to the top-right corner, or whether a pink guard blocking the blue thief from heading the bottom-righe right corner. Indeed, the model reproduces this ambiguity.
C.3 Cart-pole (continuous state space with physical dynamics)
The cart-pole domain a classic problem in reinforcement learning and optimal control. The goal to balance pole in upright position, by moving the cart left or right. The state space of this
19

---

1.0
Blue agent is thief (Pink agent is guard)
Pink agent is thief (Blue agent is guard)
Model prediction
0.4
0.2
0.0
Top left
Bottom
Top
Bottom right
left
right
Goal
Figure 9: In this example, it is unclear whether bluc guard is blocking pink thief from heading to the top-right corner, or whether a pink guard blocking the blue thief from heading the the bottom-right corner. The model reproduces this joint uncertainty as desired. See Appendix C.2.
domain consists of four continuous numbers: the horizontal position of the cart and its velocity, and the angle of the pole along with its angular velocity. The inference tasks are to look at a snapshot image-which only shows the cart position and the pole angle-and determine the velocity of the cart and the angular velocity of the pole. Note that rejection sampling cannot solve this task because the probability of ranpolly-sampled trace passing through the observed state is zero.
use an off-the-shelf pre-trained Proximal Policy Optimization (PPO) controller [Schulman et al., 2017] from stable-baselines3 [Raffin et al., 2019] to compute probability distribution over actions. Inference in this domain is complicated by the fact that computing backward dynamics physical simulation is challenging and often ill ill-posed. While previous work has proposed analytic approaches [Twigg and James, 2008], we instead train a neural network to approximate the reverse physical dynamics. We place unit Gaussian prior over the velocities, and use a von Mises distribution prior over the initial pole angle. We infer the velocities of the system by sampling candidate pairs of cart and pole velocities (stratified in 11 11 grid) and computing likelihoods using our algorithm.
The inferred posteriors are intuitive and track the relative stability the position in each snapshot (Figure C.3). For example, in part (a), the pole has almost completely fallen over, and our method infers that the pole has large negative angular velocity, and falling fast towards the ground. At same time, it infers that the cart is moving fast to the left, because the controller is likely attempting to re- re-reeelaee balance. In comparison, for part (f), the pole is nearly upright, so the model predicts that the pole is not rotating, and that the cart might be moving left or right to keep the pole balanced.
20

---

(a)
Angular velocity
5.0
(d)
Angular velocity
5.0
0.0
0.0
-5.0
-5.0
-2.0
0.0
2.0
-2.0
0.0
2.0
Cart velocity
Cart velocity
(b)
Angular velocity
5.0
(e)
Angular velocity
5.0
0.0
-5.0
-5.0
-2.0
0.0
2.0
-2.0
0.0
2.0
Cart velocity
Cart velocity
Angular velocity
5.0
Angular velocity
5.0
0.0
-5.0
-5.0
-2.0
0.0
2.0
-2.0
0.0
2.0
Cart velocity
Cart velocity
In each pair, the left image shows the cart-pole snapshot given to the algorithm, and the overlaid arrows summarize the model's predictions about how the system might evolve. The right heatmap shows our model's full joint distribution of inferred cart velocity (positive means moving to the right) and pole angular velocity (positive means clockwise), and the white expectations. When the pole near-horizontal, our algorithm infers that the pole is falling, and the cart is moving left to re-balance. When the pole is near-vertical, the algorithm infers that the pole is stationary, and the cart is making minor justments to keep the pole balanced.
Figure 10:
mark posterior
21


Reviews:
[{"review_id": "3YBmP6OMer", "reviewer": "Reviewer_iHfX", "summary": "This paper deals with the problem of inferring the goal state $g$ of an agent given only a single state $x$ in a trajectory, i.e., inferring $p(g|x)$.\nFor this, the authors claim that we need to integrate all possible initial states, and thus, sampling past trajectories is necessary for Monte Carlo estimation.\nThey propose several techniques to improve the sampling efficiency, and the main technique is to sample the past trajectory in reverse, starting from the given state $x$.", "questions": "- Please state the key assumptions\n- Can you prove that sampling the past trajectory is absolutely necessary to predict $p(g|x)$?", "limitations": "Yes", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The key intuition and motivation are clearly described.", "weaknesses": "### The key assumptions are not clearly stated\n\nI think the main text does not explain the key assumptions, such as\n- What is given? A fully-trained agent? A fully transparent environment with perfectly known dynamics?\n- How can we perform the past and future sampling?\n\n### Soundness of the main argument\n\nMost importantly, I am not convinced that explicitly sampling the past trajectory is necessary to infer $p(g|x)$.\nSince $x$ contains all the information that we have, sampling the past trajectory seems like an unnecessary complication.\nI think we can train a predictor that produces $p(g|x)$ directly from a decent amount of offline trajectories.\n\nMoreover, I believe there is a serious mistake in Eq.(2).\nThe left hand side of Eq.(2) is $p(x|g)$, but the right hand side assumes that every path passes through $x$ and does not consider the cases where $x$ is not included in the path.\nAs the authors noted in L81-82, most paths would not pass through $x$, so there should be a term that account for this fact.\n\n### Limited scope of the experiments\n\nThe experiments are limited to simple toy problems like grid world. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for raising several important points about our work. We believe we can address all of your concerns — we respond to them at the beginning of the **common response**."}, {"review_id": "Fn25u1A0A9", "reviewer": "Reviewer_yiSb", "summary": "This paper presents an algorithm for efficiently inferring the goal of an RL agent from just observing its current (single) state $x$. The method  improves substantially upon rejection sampling based prior work by 1) only sampling paths through $x$ by separating the path into past and future 2) sampling the past path backwards in time starting from $x$ and 3) using incremental A-star search for planning.\nThe efficacy of the approach is demonstrated in two gridworld domains and a word-blocks game (as well as three more in the supplementary). The posterior estimates over the goal of the agent qualitatively match human judgements, and converges several orders of magnitude faster than vanilla rejection sampling.", "questions": "* I am unclear about the role of A* in the proposed algorithm (sec 2.3). I assume that $p(x\\rightarrow x') \\propto \\sum_a \\exp (\\beta(C(x \\rightarrow g) - C(x' \\rightarrow g)) $ takes the role of the A* heuristic. But in A* this heuristic is used to decide which node to further expand  next, and I cannot find the equivalent of this search process in Algorithm 2. Could you please clarify the correspondence of A* with the proposed algorithm.\n* I am unclear about the shaded cells in in Table 1. The description states that \"Shaded cells were excluded from the analysis because it would be irrational for the agent to be there for any goal\".  I assume \"shaded cells\" refers to the gray cells in the rows corresponding to the keys world. But the shading of cells doesn't make sense to me. In row 3 it seems to correspond to inaccessible areas (which make sense to exclude). In row 4 the bottom third is shaded except for a single spot next to the green key. Why would that spot be less irrational than the surrounding ones? In row 5 the leftmost column is shaded which also seems odd. Is it impossible for the agent to start there? Also why are no cells shaded in row 1? Please clarify. \n", "limitations": "The authors discuss an important limitation of their work regarding the sampling over goals. \nAnother limitation that may be worth pointing out is that the separation of paths into independent paths for past and future depends on the environment to be Markovian (this *is* mentioned in the paper). More specifically that the current states $x$ contains all information connecting the past and the future. For more realistic environments this would for example have to include quantities like velocity or any relevant inner state of the agent in addition to the current position. Relaxing this assumption would require accounting for latent states and thus be equivalent to observing a distributions over current states $x$ instead of a concrete state. ", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "* The studied problem is interesting and potentially impactful for multi-agent RL or human-agent interactions.\n* The algorithmic modifications are well motivated and the efficiency improvement above simple rejection sampling is impressive. \n* The paper validated the results against human judgements and found a high correlation\n* The presentation is unusually clear, and a fun read.\n", "weaknesses": "The quantitative evaluation in Table 3 assumes the ground-truth to be a converged estimate of the presented method. This leaves open the possibility that the TV is better than the rejection sampling baseline because the presented method is biased in some way. The concern is partially addressed with the numerical validation check in the supplementary section B, which in my opinion should at least be mentioned in the main paper. Another way to strengthen the results in Tab 3 would be to also compare to a \"ground-truth\" estimate from (very?) many vanilla rejection samples. If the presented method is still better at approximating this ground-truth with few samples than rejection sampling with few samples, that would be strong evidence. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your careful feedback on our work. Please see the **common response** for our remarks on the role of A-star. We respond to the rest of your questions below.\n\n**Table 3 ground truth**  \nThank you for raising this concern, which we will address as follows:\n1. Mention the correctness checks in the main body of the paper, as suggested.\n2. Expand Table 3 with two new columns showing analogous TVs compared to converged posterior estimates from many rejection samples, as suggested. For example, for the start-anywhere grid-world, the numbers using 10,000 rejection samples are below (essentially the same, as expected):\n| Method | Ground truth | Total variation |\n|---|---|---|\n| Rejection | Ours, 1k samples | 0.159 (from paper) |\n| Rejection | Rejection, 10k samples | 0.153 **(new)** |\n| Ours | Ours, 1k samples       | 0.0538 (from paper) |\n| Ours | Rejection, 10k samples      | 0.0836 **(new)** |\n3. As an additional assurance, we can show scatterplots analogous to Fig 4 for our main benchmarks, to emphasize unbiasedness. This is essentially replotting the data in Tables 1/2. As an example, the scatterplot for grid-world is in the **attached PDF file** — the two candidate ground-truth posteriors match nearly perfectly.\n\n**Shaded cells in doors-keys-gems**  \nWe appreciate the detailed reading of our figures. Yes, the caption refers to the grayed cells in the keys domain, and the shaded cells in Row 3 are indeed shaded because they are inaccessible (we will make this clear when revising).\n\nAs for the other rows, we will update the paper to explain our reasoning:\n+ The shaded cells in Row 4 are excluded because if the agent already picked up the pink key earlier, there is no reason to _then_ move far _away_ from the doors. Similarly, in Row 5, there is no reason to move that far left after picking up the green key. We found that querying repeatedly at those extreme locations confuses human participants about the task setup. So, we excluded most such locations from our experiments.\n+ The single cell above the green key in Row 4 is an exception we wished to test. You might think from that position that the agent is trying to collect _both_ keys, leading to uncertainty about the goal. But people (and the model) agree that the agent still wants the red gem.\n\nWe will additionally clarify that the agent always starts empty-handed, which may be one source of confusion here.\n\n**Accounting for latent state information**  \nThank you for raising this limitation. We believe the cart-pole example from the supplement begins to address this: we make inferences when only positions, not velocities, are visible. It works exactly as you suggest, i.e. by observing a distribution over current states $x$. We will update the paper to explain this in more detail."}, {"review_id": "RcjhtcC0Gs", "reviewer": "Reviewer_cLNU", "summary": "This papers proposes a new approach to the inference problem of inferring an agent's goal state from information about a single state. This approach is based on the bidirectional Monte Carlo sampling of trajectory sequences, both from the agent's given state $x$ to the goal state $g$, and from an initial state $s$ to $x$. The authors show near convergent inference of the approach with 10 samples on two gridworld based environments, and a word game kind of environment when evaluated against human judgement.", "questions": "- Have you checked how the sampling scales with an increasing number of possible starting states, and an increasing number of possible goal states? It could potentially be interesting to see.\n- Why is incremental A-star search used for planning on-line?\n- I am wondering if its possible to emphasize more on the scope of the applications of the proposed sampling approach-- what other domains can it be used in, apart from tasks that look like spatial navigation? The word blocks is one good example, also the cartpole experiment. There are intuitive physics based experiments, or reasoning based tasks which involve inferring over a sequence of images say, can this approach be applied there?", "limitations": "The authors adequately address the limitations of their approach, especially when it comes to sampling over goals, or scaling up to a larger goal space. There is just one minor point here-- the approach is based on a combination of well-founded insights and ideas, which when put together, work well, but there seems to be a lot of hand-engineering/domain-specific engineering in assembling these together, which I think should be acknowledged in the paper.", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The evaluations show a strong \"near\" convergence result of only 10 samples, even for multi-stage planning. This indicates their sampling strategy is practically sample-efficient, and a huge improvement over prior approaches requiring more than 300,000 samples (reportedly).\n- The connection drawn from path tracing in computer graphics is a fresh insight, which the authors show applies well to the goal/previous states' inference problem, and can potentially have broader scope applications such a for spatial navigation tasks etc.\n- The use of human judgement (using 200 participants) for evaluation is also a very reasonable choice instead of hand-engineered reward functions designed by a single person/a few people.\n- The authors also show efficient sampling in a variety of other domains (in the appendix) which indicates it is more broadly applicable.", "weaknesses": "- _Experiments_: Something I might myself be unclear about, so let me know if that is the case-- there are no evaluations comparing the proposed model in this paper with the model shown by [Lopez-Brau et al 2022](https://psyarxiv.com/4zu9n), is this only because of sample-efficiency? It would be interesting to see how \"nearly\" convergent their proposed approach is in the cases where it can be applied.\n\n- _Theoretical analysis_: What is the complexity of the sampling with the size of the state space, the number of accessible states, the number of starting/goal states etc-- is it possible to perform some kind of theoretical analysis of this? \n\nI'll be very happy to increase my score on the addressal of these concerns/doubts.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful feedback on our paper. We address your questions below, and in the **common response**.\n\n**Did you compare to Lopez-Brau et al?**  \nYes, we indeed show these comparisons in Tables 1/2/3, under the heading of \"Rejection.\" We apologize for the confusion and will revise the heading to make clear that that column refers to Lopez-Brau et al's method.\n\n**Theoretical analysis of scaling**  \nWe agree it would be good to address this, and we will update the paper with a discussion.\n\nThe time complexity of taking a _single sample_ is straightforward: it scales linearly with the Russian Roulette termination depth and the number of possible goals, but is constant in the number of possible start states (because we sample backwards in time!). The plots in the **attached PDF file** establish this empirically using the grid-world domain; we will include such figures in the revision of our paper.\n\n(Runtime also scales linearly in the number of accessible \"neighbor\" states, because we softmax over possible actions. But in practice this number is fixed by the domain, and typically a small constant that contributes negligible overhead. For example, in our grid-worlds the number of accessible states is typically 4 for north/east/south/west.)\n\nFinally, another important factor is the _number of samples_ needed, which depends on how much variance is acceptable. Standard theoretical bounds for Monte Carlo integration hold here: variance reduces linearly with the number of samples, or even better via stratification [1, 2]. Such algorithms are well-studied in graphics [3, 4], and those tighter bounds may also be applicable to our setting.\n\n1. Hickernell. Koksma‐Hlawka Inequality. 2014.\n2. Bakhvalov. On the approximate calculation of multiple integrals. 2015.\n3. Singh et al. Analysis of sample correlations for Monte Carlo rendering. 2019.\n4. Subr et al. Fourier analysis of numerical integration in Monte Carlo rendering: theory and practice. 2016.\n\n**More non-spatial applications**  \nWe are glad the reviewers appreciated our non-spatial examples like blocks and cart-pole, and we are happy to expand our paper with a discussion of more such applications. As suggested, there is indeed scope for applications in intuitive physics and reasoning, such as:\n+ Determining the viscosity of a fluid from a static image of it being splashed\n+ Interpreting comic books or step-by-step instruction manuals, by inferring what happens between static panels\n\nAdditional domains where our method could be useful include:\n+ Robotics: walking into a kitchen and immediately recognizing what the chef is cooking, in order to help accordingly\n+ Vision: understanding dynamic action in static snapshots of sports games\n+ Forensics: inferring/reconstructing what occurred at a crime scene based on observable evidence\n\n**Acknowledging engineering work**  \nThank you for raising this - we will revise the paper accordingly. Please see the end of the **common response** for our full remarks on this point."}, {"review_id": "fFietiBo5Y", "reviewer": "Reviewer_PVLo", "summary": "This paper presents an efficient Monte-carlo algorithm to infer goals from single snapshots. The problem the authors consider is the same as in work by Lopez-Brau (2020, 2022) but previous solutions are slow as they apply rejection sampling. In this work, the authors use the insight that one can sample a valid path by sampling the part of the path \"before\" the current snapshot and \"after\" the snapshot independently, decresing the computational complexity by a large amount. Further, A-star search is used to find likley paths. The authors validate their algorithm on different grid worlds and provide results of a human study indicating that inferences of their algorithm coincide with human judgements.", "questions": "The regarded grid worlds all have a discrete state/action space. Is the algorithm applicable to the continuous domain?\n\nWhen introducing the algorithm, you write that the situation of the agent you consider is an MDP but you would relax this assumption later in the paper. Where do you do that?", "limitations": "Limitations are well discussed. To my understanding, only discrete state/action/goal spaces are considered. This point could be added.", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The problem of inferring goals from a single snapshot is interesting, as a major limitation of most IRL methods is that they only work in a certain context (state/actions of the full trajectory).\n\nThe paper is clearly written and easy to understand.\n\nThe presented method is much more efficient than previous methods in terms of scalability.\n\nThe plausability of the results found by the algorithm is supported by human experiments.", "weaknesses": "Quantities are not clearly defined. Are states/actions/goals continuous or discrete? In the formulas, one integrates over states and sums over actions. It might be very helpful for the reader to formally define the MDP and its sets. Further, what are \"paths\" formally, what does the indexing operator \"[0]\" mean for paths?\n\nWhile the main idea of the algorithm to sample the path \"before\" and \"after\" is very simple, the overall algorithm seems quite hacky to me: The basic idea is combined with other tweaks such as roulette termination and A-star search to avoid unlikely paths. With these, the considered tasks can be efficiently solved but there is no particularly nice linking theory.\n\nIn my opinion, the contribution is not strong enough for a NeurIPS submission as the main (conceptual) difference to previously published algorithms is not that large. The essential difference is that parts of the path before and after are sampled, together with other tweaks.\n\nI would find it helpful if the algorithm clearly indicated which quanitity it computes (p(x | g)), this is not direcly visible.\n\nThe plots in Table 1 do not have a colorbar, therefore the precise meaning and interval of the color's values is not clear to me.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful comments about our paper. We address your major concern about the lack of a linking theory in the **common response**. Here, we address your additional questions:\n\n**Is our method applicable to continuous domains?**  \nYes, it is: the cart-pole example in the supplement shows our method in a continuous state space. When revising, we will mention this in the main text. We will also update the paper to acknowledge that we do not address continuous _action_ spaces.\n\n**Where do we relax the MDP assumption?**  \nThank you for raising this — we will update the paper to clarify that we relax the MDP assumption in Sec 2.3. (Please see the common response regarding A-star for more information on why we do so, and how we will further revise that section based on the reviewers' comments.)\n\n**Improving presentation and definition of key quantities**  \nThank you for the suggestions for improving the paper's presentation. We will update the paper to:\n1. introduce notation that defines the state/action spaces of MDPs,\n2. formalize \"paths\" as finite ordered sequences of states, which are indexed and sliced the same as lists,\n3. clarify that the output of Algorithms 1 and 2 is an estimate of $p(x \\mid g)$, and\n4. include colorbars in Tables 1 and 2."}]

Author rebuttal:
We thank all the reviewers for their thoughtful feedback. We address some key concerns below, and the rest in individual responses.

**Is sampling really necessary? Why not fit a neural network? (iHfx)**  
Thank you for raising this important point. We realize that our paper's motivation was not fully clear, and we will revise the paper to explain the broader context of our work.

If our goal was only to solve this inference task in a particular environment, we could indeed just fit a model to predict $p(g\mid x)$ from datasets of offline trajectories generated in that environment. But here, we are interested specifically in how humans flexibly make such inferences _without_ extensive pre-training on data -- and how AI systems could do likewise.  A long line of empirical work in cognitive science, particularly in Theory of Mind, shows that people (even infants and young children) make rapid, flexible, and robust judgments "out of the box": in novel domains they have never seen before, and without extensive pre-training on data [1-5]. This remarkable one-shot ability is what motivates our work, and our sampling-based algorithm specifically seeks to capture that ability.

1. Gergely, G., and Csibra, G. "Teleological reasoning in infancy: The naive theory of rational action." _Trends in Cognitive Sciences_ 7.7 (2003): 287-292.
2. Baker, Chris L., Rebecca Saxe, and Joshua B. Tenenbaum. "Action understanding as inverse planning." _Cognition_ 113.3 (2009): 329-349.
3. Baker, Chris L., et al. "Rational quantitative attribution of beliefs, desires and percepts in human mentalizing." _Nature Human Behaviour_ 1.4 (2017): 0064.
4. Jara-Ettinger, Julian, et al. "Children’s understanding of the costs and rewards underlying rational action." _Cognition_ 140 (2015): 14-23
5. Hamlin, J. Kiley, Karen Wynn, and Paul Bloom. "Social evaluation by preverbal infants." _Nature_ 450.7169 (2007): 557-559.

**Limited scope of experiments? (iHfx)**  
As part of the revisions promised above, we will explain that we based our experiments on tasks studied by empirical cognitive science research. Our experiments are comparable in scale to contemporary related work, both in AI (Zhi-Xuan et al, NeurIPS 2020; Shah et al, ICLR 2019) and cognitive science (Lopez-Brau et al, CogSci 2020). Furthermore, while prior work typically evaluates on only 1-3 domains, we consider a total of 6 domains across a wide variety of conditions (partial observability, multi-agent, continuous physics, etc.).

**Possible mistake in Eqn 2? (iHfx)**  
Thank you for the attention to detail! This is a subtle point, but we still believe there is no mistake. We will update the paper to clarify: paths that do not pass through $x$ contribute zero likelihood of being observed at $x$, so we only need to integrate over paths that _do_ pass through $x$ (the "other term" would be zero). More generally, we hope that the numerical check in Supplement B assuages correctness concerns.

**Stating key assumptions (iHfx)**  
We will revise the paper to clarify our assumptions: we assume full knowledge of environmental dynamics, and access to a planner to model how an agent would act given goal $g$. Formally, this is captured by $P(s \rightarrow s^\prime \mid g)$, a key input to Algorithms 1 and 2. Finally, we assume the ability to enumerate the neighbors of the current state, which are used for past/future sampling.

**Role of incremental A-star? (yiSb, cLNU)**  
We realize that our motivation for using A-star was not clearly explained in Sec 2.3. We will revise the paper to clarify that we use A-star planning _only as an optimization_ — we can obtain the same results using value iteration, as in prior work, or for that matter any other planning algorithm (e.g. our cart-pole example uses deep RL).

We will additionally update the paper to better explain _why_ we use A-star:
1. Prior work formulates problems as MDPs, modeling agents' action choices by softmax over $Q(x, a)$. The drawback is that the full MDP must be solved offline before inference, e.g. by value iteration or deep RL. This is (a) cognitively implausible, and (b) wasteful if only some states are queried at inference time. Additionally, not all tasks are well-modeled by MDPs (e.g. multi-stage long-horizon planning).
2. Our solution is that where possible, we instead formulate problems as classical planning domains (c.f. PDDL) and weight actions by softmax over how much closer an action brings the agent to the goal. That is, for moving onto state $x^\prime$, we use the cost difference $C(x \rightarrow g) - C(x^\prime \rightarrow g)$.
3. The upshot is that we can compute the cost $C$ on-line by A-star search, avoiding the need for precomputation and only exploring relevant states. An additional optimization is to cache/memoize A-star's intermediate computations to avoid duplicate work if calls to $C$ are repeated.

We are happy to explain further if this remains unclear.

**Is this a "hacky" approach with no linking theory? (PVLo, cLNU)**  
We realize that our exposition confusingly mixes theory and engineering, and we thank the reviewers for raising this issue. We will revise the paper to explain that there is indeed a core linking theory at play: the theory of Monte Carlo path tracing.

Rather than a grab-bag of "tweaks," we see our paper as organized around one key idea: the analogy between light paths and agent paths. This idea is what allows us to (1) reframe our problem using the theory of Monte Carlo path tracing, and (2) import existing engineering techniques motivated by that theory (e.g. Russian Roulette, bidirectional tracing) and thus solve our problem efficiently.

In short, we agree with cLNU's assessment: our work combines well-founded insights with some engineering work. We will revise the paper to distinguish the theoretical insights and engineering contributions we offer.

Meta-review: This work leverages modeling from graphics to devise a cognitively plausible sampling method for inferring future states based on the previous trajectory.  The work compares to relevant baselines (including simple approaches like rejection sampling).  The reviewers are positive on the work with their primary concerns centering on framing.  Most issues are clarified in the author responses (e.g. issues around A*, the focus on cognitive plausibility, and clearly differentiating the novel modeling from engineering).  Personally, earlier discussion of the other domains (e.g. continuous cartpole) and what those imply about future broader application domains would strengthen the paper for those skeptical of the grid-world setup.
Decision: Accept (spotlight)