--- abstract: | Tasks that involve complex interactions between objects with unknown dynamics make planning before execution difficult. These tasks require agents to iteratively improve their actions after actively exploring causes and effects in the environment. For these type of tasks, we propose `\ourName{}`{=latex}, a method that leverages Bayesian optimization to reason about causal interactions via a Physics-Informed Kernel to help guide efficient search for the best next action. Experimental results on Virtual Tools and PHYRE physical reasoning benchmarks show that `\ourName{}`{=latex} outperforms state-of-the-art results, requiring fewer actions to reach the goal. We also compare `\ourName{}`{=latex} to human studies, including results from a new user study we conducted on the PHYRE benchmark. We find that `\ourName{}`{=latex} remains competitive on tasks that are very challenging, even for human problem-solvers. bibliography: - ref.bib --- \newcommand{\theHalgorithm}{\arabic{algorithm}} \newcommand{\ourName}{Causal-PIK} \newcommand{\kernel}{Physics-Informed Kernel} \newcommand{\VirtualTools}{Virtual Tools} \newcommand{\taskType}{single-intervention physical reasoning tasks} \twocolumn[ \icmltitle{Causal-PIK: \\ Causality-based Physical Reasoning with a Physics-Informed Kernel} \begin{icmlauthorlist} \icmlauthor{Carlota Parés-Morlans}{st} \icmlauthor{Michelle Yi}{st} \icmlauthor{Claire Chen}{st} \icmlauthor{Sarah A. Wu}{stp} \icmlauthor{Rika Antonova}{st,cb} \icmlauthor{Tobias Gerstenberg}{stp} \icmlauthor{Jeannette Bohg}{st} \end{icmlauthorlist} \icmlaffiliation{st}{Department of Computer Science, Stanford University, CA, USA} \icmlaffiliation{stp}{Department of Psychology, Stanford University, CA, USA} \icmlaffiliation{cb}{Department of Computer Science and Technology, University of Cambridge, Cambridge, UK} \icmlcorrespondingauthor{Carlota Par\'{e}s-Morlans}{cpares@stanford.edu} \icmlkeywords{active exploration,physical reasoning,causality,bayesian optimization} ] `\printAffiliationsAndNotice{}`{=latex}
Example puzzles from the Virtual Tools and PHYRE benchmarks. For Virtual Tools, the objective is to place one of the blue action objects from the left to have a red ball fall into the green area. For PHYRE, the objective is to place a red ball of variable radius in the environment to have the green and blue objects touch for at least 3 seconds. The runs highlighted in red are failed attempts and the runs highlighted in green are successful attempts.
# Introduction Consider trying to solve a physical reasoning puzzle from the Virtual Tools benchmark, like those shown in Figure `\ref{fig:teaser}`{=latex}. Your goal is to have the red ball end up in the green region. To make that happen, you need to choose one of the three objects from the top left, place it somewhere in the scene, and let gravity and causality do the rest. Which object would you choose and where would you place it? Research in cognitive science suggests that humans solve tasks like these by building internal models of the physical world [@battaglia2013simulation; @smith2018different; @ota2021dataefficient; @zhou2023mental]. These models encode our causal understanding of the domain -- our beliefs about how one event leads to another. While predicting the exact sequence of events is difficult, our physical intuition allows us to estimate the immediate causal effects of our actions, such as which object will move and in which direction. In addition, we rapidly learn from the outcomes of previous actions [@allen2020rapid]. If the red ball almost landed in the goal region, we would slightly change our action and try again. If it was way off, we might try something completely different. Inspired by people's rapid learning in such tasks, we develop a method that leverages physical intuitions to efficiently solve physical reasoning tasks. Researchers have developed several benchmarks for evaluating the physical reasoning capabilities of AI agents [@melnik2023benchmarks]. These benchmarks range from predicting the stability of a stack of blocks from an image to recognizing violations of physical principles. In this work, we focus on `\taskType{}`{=latex}, where an agent chooses what action to take at the beginning of an episode and then observes the action's effects on the environment. We focus on the Virtual Tools [@allen2020rapid] and PHYRE [@bakhtin2019phyre] benchmarks, which despite looking deceivingly simple in their 2D form, are very hard to solve even for humans [@allen2020rapid]. These tasks involve complex physical interactions between objects, which without complete information about the environment dynamics make it impossible for agents to plan the exact solution without active exploration. Agents must interact with the environment to actively obtain information about the object dynamics. The main challenge is to reason about the causal consequences of one's actions, to leverage this knowledge in order to make effective decisions about what actions to try, and to rapidly learn from previous attempts. We propose `\ourName{}`{=latex}, a method that leverages physical intuition and causality to solve `\taskType{}`{=latex} in as few trials as possible. Our method uses Bayesian optimization to reason about causality via a `\kernel{}`{=latex} to obtain an expressive posterior distribution over the environment dynamics. This allows an agent to efficiently explore the search space and intelligently select what actions to evaluate next. `\ourName{}`{=latex} minimizes the number of physical interactions required. We demonstrate that `\ourName{}`{=latex} significantly outperforms state-of-the-art models on the `\taskType{}`{=latex} from the Virtual Tools [@allen2020rapid] and PHYRE [@bakhtin2019phyre] benchmarks, finding solutions in the fewest number of attempts. # Related Work {#sec:related} **Physical Reasoning:** Researchers have developed several benchmarks towards the goal of improving physical reasoning capabilities in machine-learning models [@melnik2023benchmarks]. We evaluate `\ourName{}`{=latex} on the PHYRE [@bakhtin2019phyre] and `\VirtualTools{}`{=latex} [@allen2020rapid] benchmarks. Both benchmarks use a physics simulator and capture a variety of complex interaction mechanisms. Furthermore, both contain a large set of puzzle variations, allowing us to study generalization across diverse environments. Both PHYRE and `\VirtualTools{}`{=latex} have inspired several methods for solving physical reasoning puzzles. A number of works [@harter2020solving; @girdhar2020forward; @ahmed2021physical; @qi2021learning; @li2022learningmechanismsphysicalreasoning] that evaluate on PHYRE leverage forward prediction models, also called dynamics models or world models, to predict the outcomes of actions. They use dynamics models to score actions and then execute the most promising actions until one succeeds. Like these works, `\ourName{}`{=latex} also uses a dynamics model to reason about the action outcomes. However, unlike these prior methods, `\ourName{}`{=latex} uses observations from previous trials to inform future action selection via Bayesian optimization. Instead of using a dynamics model to directly choose actions, we use dynamics predictions to instill physical intuition into kernel updates during Bayesian optimization. We show that our method solves PHYRE in fewer attempts than any of these prior methods. The \`Sample, Simulate, Update' model (SSUP) [@allen2020rapid] attempts to solve the `\VirtualTools{}`{=latex} benchmark. It samples actions from an object-based prior, simulates the sampled actions in a noisy physics engine to find the best action to try, executes the action, and updates the model's belief using information from both simulation and execution. SSUP uses a Gaussian mixture model for guidance, but it doesn't embed information about how actions are related via their effects, making SSUP search less efficient. In contrast, our work builds on active learning methods, such as Bayesian optimization, and successfully incorporates additional guidance for action selection by introducing a `\kernel{}`{=latex}. This significantly lowers how many trials are needed to succeed. **Bayesian optimization:** `\ourName{}`{=latex} uses Bayesian optimization (BO), a global search method for optimizing black-box functions [@shahriari2015taking; @wang2023recent]. Similar to how we use BO to efficiently find puzzle-solving actions, BO has been used in robotics to reduce the number of trials needed to be run on real robots [@feng2015optimization; @calandra17thesis; @jaquier2020bayesian; @berkenkamp2023bayesian]. Some prior works proposed to further increase data-efficiency by incorporating simulation-based information [@antonova2017deep; @marco2017virtual; @antonova2019bayesian]. However, they mostly considered continuous parametric controllers, while we need to accommodate a hybrid continuous/discrete action space.
An overview of . The Causality-Based Action Selection module proposes the best action to execute based on the predictions from the Gaussian Process. These predictions are generated considering the rewards from previous actions and the causal similarities between candidate actions computed by the . These similarities are derived from the individual effects of actions, as predicted by the learned dynamics model. They are expressed as the state change of object O caused by action a, denoted as $\dot{\bm{s}}^{O, a}$.
**Inferring Causality:** For the type of tasks we study in this work, the ability to reason about the outcome of an interaction is critical [@ullman2018learning; @bramley2018experiments]. Some works aim to reason about causality by focusing on the Violation of Expectations (VoE) paradigm [@smith2019modeling; @riochet2021intphys], where models learn physical plausibility by observing scenes that either follow or violate intuitive physics rules. These models can then judge if new scenes are plausible or not based on past observations. Other works utilize deep neural networks to predict the outcome of interactions from frames [@duan2022pip; @ye2018interpretable]. If an agent can reason about the outcome of interactions before actions are taken, then the reasoning can be used to guide smarter action selection, which is beneficial when taking actions in the real world is costly. [@battaglia2013simulation] and [@wu2017learning] introduce methods that conduct physical reasoning with the help of simulation after constructing a representation of the physical world from visual inputs. # Method {#sec:method} **Problem Formulation:** We consider `\taskType{}`{=latex} where an agent executes an action at the beginning of an episode, observes the action's effects on the environment over $T$ episode steps, and then receives a score at the final timestep. After each episode, the environment resets, and the agent may attempt the task again by trying another action. Formally, on each attempt, an agent executes an action $\bm{x}$ at the first state $\bm{s}_0$ and observes this action's effect on the environment's state at the remaining timesteps, $\bm{s}_t$ for $t=1, 2, ..., T$. The state evolves according to the environment's unknown dynamics, which we denote as the function $\mathbb{D}(\bm{s}_0,\bm{x})$, to give $\bm{s_{1,...T}} = \mathbb{D}(\bm{s}_0,\bm{x})$. At the end of the episode, the agent receives a score $y$, which is computed based on the observed states $\bm{s_{1,...T}}$ using a score function $\mathbb{S}(\bm{s}_{1,...T})$: $y = \mathbb{S}(\bm{s}_{1,...T})$. For a given initial state $\bm{s}_0$, the future states $\bm{s}_{1,...T}$ are a function of the action, so we can re-write the score function as a function of $\bm{x}$: $y = \mathbb{S}(\mathbb{D}(\bm{s}_0,\bm{x})) = f(\bm{x})$. Given an initial state $\bm{s}_0$, we aim to solve the task by finding the action that achieves the highest score, namely maximizing $f(\bm{x})$, in as few attempts as possible. However, the task's unknown dynamics make it difficult to find optima of $f(\bm{x})$. Therefore, we use Bayesian optimization to find actions that maximize the score. **Bayesian Optimization (BO)**: The goal of BO is to find the optimal $\bm{x}^*$ that maximizes a given objective $f(\bm{x})$. The objective captures characteristics of the desired outcome, e.g., in our setting, the task score. BO begins with a prior that expresses uncertainty over $f(\bm{x})$. After evaluating an $\bm{x}$, BO constructs a posterior based on the observed data $y$ obtained so far. It then uses an auxiliary acquisition function to pick the next $\bm{x}$ to evaluate, taking into account both the posterior mean and covariance. A Gaussian process (GP) is commonly used to model the objective function: $f(\bm{x}) \sim GP(\mu(\bm{x}), k(\cdot, \cdot))$. Its kernel function describes the covariance between the objective values of any pair of input points: $k(\bm{x}_1,\bm{x}_2)\!=\!cov(f(\bm{x}_1),f(\bm{x}_2))$. A common choice for the kernel is the Radial Basis Function (RBF): $k_{RBF}(\bm{x}_1, \bm{x}_2) \!=\! \sigma_k^2\text{exp}( -\tfrac{1}{2} (\bm{x}_1 - \bm{x}_2)^T \text{diag}(\bm{l})^{-2} (\bm{x}_1 - \bm{x}_2))$, where $\sigma_k^2$, $\bm{l}$ can be tuned automatically. The kernel plays a crucial role in BO, as it encodes an inductive bias about the properties of the underlying objective function to choose the next $\bm{x}$. The RBF kernel, in particular, assumes that values geometrically closer in the input space have a higher correlation (are more similar) than values geometrically farther in the input space. Selecting an appropriate kernel significantly impacts the accuracy of the objective function approximated by the GP, thus influencing the performance of BO. We develop a kernel that encodes an intuition of causality and physics, which we show solves physical reasoning tasks in fewer attempts than the RBF kernel. In the following section, we describe our full physics-informed BO algorithm `\ourName{}`{=latex}. ## `\ourName{}`{=latex} `\ourName{}`{=latex} uses BO and reasons about causality with a `\kernel{}`{=latex} to solve physical reasoning tasks. It iteratively proposes actions and uses their observed outcomes, in conjunction with an intuition of physics, to inform future action proposals. We encode physics intuition, such as effects of geometry, mass, and momentum, into the BO framework via the learned `\kernel{}`{=latex} used to build the GP that approximates the objective function. We will show that incorporating knowledge of physics through the kernel function allows BO to more quickly discover promising regions of the action space. \begin{algorithm} \caption{\ourName{}} \label{alg:overview} \begin{algorithmic}[1] {\small \STATE Attempt $i$ = $0$ \STATE $X$ = $\{\bm{x}^{\text{init}}_1,...,\bm{x^{\text{init}}_{n_{\text{initial}}}}\}$, $y$ = $\{\hat{y}^{\text{init}}_1,...,\hat{y}^{\text{init}}_{n_{\text{initial}}}\}$ \WHILE{success == False} \STATE $gp$ $\leftarrow$ \texttt{\textbf{GP}}$(X,y, \text{\texttt{PhysicsInformedKernel}})$ \STATE $\bm{x}^i \leftarrow \text{\texttt{CausalityBasedActionSelection}}(gp)$ \STATE $y^i, \text{success} \leftarrow \text{\texttt{Execute}}(\bm{x}^i)$ \STATE $X \leftarrow \text{concat}(X, \bm{x}^i)$, $y \leftarrow \text{concat}(y, y^i)$ \STATE $i$ += $1$ \ENDWHILE } \end{algorithmic} \end{algorithm} Algorithm `\ref{alg:overview}`{=latex} outlines `\ourName{}`{=latex}. First, we construct initial sets of actions $X$ and scores $y$ to initialize the GP prior using a probabilistic intuitive physics engine [@battaglia2013simulation]. After initializing $X$ and $y$, we begin the BO procedure, iteratively updating the physics-informed GP surrogate function, choosing the next action and executing it until the puzzle is solved. **GP Initialization (Alg. `\ref{alg:overview}`{=latex} L2)**: To initialize the GP for both Virtual Tools and PHYRE, we use $n_{initial}=9$ initial data points. This is the same number of initial data points that our Virtual Tools baseline SSUP [@allen2020rapid] uses to initialize their method. SSUP found that this value provided a good trade-off between the number of initial points, total attempts required, and convergence time. Since we use the same initialization framework as SSUP, we expect their parameter analysis to extend to our results. Importantly, like SSUP, we treat these initial noisy rollouts as warm-up samples that do not count towards the total attempt count. We choose the initial points following the same approach as SSUP: randomly select a dynamic object from the environment and sample a point from a Gaussian distribution centered at the object's center. As a result, each puzzle attempt has a unique set of $n_{initial}$ points. This heuristic helps the GP build a noisy prior that includes different areas. **Physics-Informed GP Update (Alg. `\ref{alg:overview}`{=latex} L4)**: On the $i$th attempt, we update the GP surrogate function with all attempted actions $X$ and their observed outcomes $y$ using our learned `\kernel{}`{=latex}. Conceptually, the `\kernel{}`{=latex} encodes similarities and differences between actions based on the predicted effects they will have on the environment. Using this kernel to construct the GP helps the GP use the observed outcomes of already-attempted actions to more accurately predict the outcomes of unexplored actions. We derive the `\kernel{}`{=latex} in Section `\ref{sec:kernel}`{=latex}. **Causality-Based Action Selection (Alg. `\ref{alg:overview}`{=latex} L5)**: On the $i$th attempt, we select the next most promising action $\bm{x}^i$ by finding the action that maximizes an Upper Confidence Bound (UCB) acquisition function. This acquisition function considers the mean and uncertainty of the physics-informed GP. First, we use a Sobol sequence generator to sample a set of $n_{\text{candidate}}=500$ candidate actions. Then, we evaluate the acquisition function at each of these $n_{\text{candidate}}$ actions. Adopting the intuitive physics procedure proposed by @allen2020rapid, we approximate the outcome of the $n_{\text{best}}=5$ candidate actions with the highest acquisition function values using a probabilistic simulation of the task. As per [@allen2020rapid], approximating the outcomes of actions mimics how humans use mental representations to imagine the potential effects of actions before committing to an action. Finally, from this set of $n_{\text{best}}$ actions, we select the action with the highest expected outcome as the next action. **Action Execution (Alg. `\ref{alg:overview}`{=latex} L6)**: We execute the selected action $\bm{x}^i$, observe its outcome, and compute the corresponding reward $y^i$. If $\bm{x}^i$ solves the task, the algorithm terminates. If $\bm{x}^i$ fails to solve the task, we append it along with the score $y^i$ to $X$ and $y$, respectively.
Illustration of how causal similarity is computed. We obviate the object superscript as there is only one object (red ball) in this example. (a): Cosine similarity (Eq. ). $\dot{\bm{s}}^{1}$ has a high cosine similarity with $\dot{\bm{s}}^{2}$ as they point in the same direction. $\dot{\bm{s}}^{3}$ points in the opposite direction of $\dot{\bm{s}}^{1}$, obtaining a negative similarity value. (b): Magnitude similarity (Eq. ). $\dot{\bm{s}}^{2}$ has a magnitude that is 40% smaller than $\dot{\bm{s}}^{1}$, receiving a low magnitude similarity score. $\dot{\bm{s}}^{3}$ has a magnitude that is very close to $\dot{\bm{s}}^{1}$, obtaining a high similarity value. (c) Final causal similarity combining cosine and magnitude similarities (Eq. ). We obtain a medium similarity score for $\dot{\bm{s}}^{2}$ relative to $\dot{\bm{s}}^{1}$ and low score for $\dot{\bm{s}}^{3}$ relative to $\dot{\bm{s}}^{1}$.
## `\kernel{}`{=latex} {#sec:kernel} In this section, we describe the `\kernel{}`{=latex} used to construct the GP in BO (Alg. `\ref{alg:overview}`{=latex} L4). We construct the `\kernel{}`{=latex} such that it captures two types of physics intuition that are important for solving physical reasoning tasks. The first type is the ability to reason about the causal effect of individual actions, or in other words, the changes that single actions induce in the environment. To capture this, we train a dynamics model to predict future states of the environment given an action and initial state. The second type of physics intuition that we encode into the kernel is the ability to reason about causal similarity, namely which actions cause similar changes in the environment. To capture this, we define a function that computes how similar two actions are based on their causal effects. As shown in Fig. `\ref{fig:overview}`{=latex}, the `\kernel{}`{=latex} function first uses the learned dynamics model to predict action outcomes. It then uses these dynamics model predictions to predict the similarity between actions. Using the predictions from the learned dynamics model to predict action similarity allows the `\kernel{}`{=latex} function to model the correlation between actions. In the following sections, we describe the learned dynamics model and the causal similarity function. Then, we show that the causal similarity function is a valid kernel function. **Predicting Causal Effects via Learned Dynamics** First, we train a dynamics model $\hat{\mathbb{D}}$ to predict the causal effects of individual actions. The physical reasoning tasks we consider involve placing an action object into an environment with $D$ dynamic objects. The action object and all dynamic objects move in the environment according to the unknown world's dynamics. We train the model to predict the effects that individual actions $\bm{x}$ have on these dynamic objects. Formally, at each timestep $t$ of an episode, the environment's state vector $\bm{s}_t$ consists of the state of the action object, $\bm{s}^{A}_t$, and the states of every dynamic object, $\bm{s}^{1}_t, \bm{s}^{2}_t, ..., \bm{s}^{D}_t$. For a given action $\bm{x}$ and initial environment state $\bm{s}_0$, the dynamics model $\hat{\mathbb{D}}$ predicts the state at the next $n_\text{pred}$ timesteps: $$\begin{align} [\bm{s}_1 \quad...\quad \bm{s}_{n_{\text{pred}}}] = \hat{\mathbb{D}} (\bm{s}_0, \bm{x}). \end{align}$$ **The Causal Similarity Function** We use predictions from the learned dynamics model $[\bm{s}_1 \quad...\quad \bm{s}_{n_{\text{pred}}}]$ to compute the causal similarity of actions, which captures how similar actions are based on their ability to cause similar outcomes in the environment. To achieve this, we begin by finding the timestep of the first causal event, where the action object interacts with any of the $D$ dynamic objects, which we call $t_{\text{event}}$. If the action object does not interact with any dynamic object during the $n_{\text{pred}}$ steps, we set $t_{\text{event}}$ to be the initial timestep, 0. Next, we quantify the causal effect of an action on a dynamic object $O$ based on its predicted motion as $$\begin{align} \label{eq:state_change} \dot{\bm{s}}^O = \frac{\bm{s}^O_{(t_{\text{event}} + \Delta t)} - \bm{s}^O_{t_{\text{event}}}}{\Delta t}, \end{align}$$ which computes the state change of object $O$ between $t_{\text{event}}$ and $\Delta t$ timesteps after $t_{\text{event}}$. To quantify similarity of two actions $a$ and $b$, we first look at the difference in the state changes, predicted by the dynamics model, they cause to a given object $O$. In computing this per-object similarity, we account for both the difference in the directions of the state changes caused by each action on the object (the cosine similarity), $$\begin{align} \label{eqn:cos-sim} \mathbb{sim}_{\text{cos}}(\dot{\bm{s}}^{O, a}, \dot{\bm{s}}^{O, b}) = \frac{\dot{\bm{s}}^{O, a} \cdot \dot{\bm{s}}^{O, b}}{||\dot{\bm{s}}^{O, a}|| \, ||\dot{\bm{s}}^{O, b}||} \in [-1,1], \end{align}$$ and the difference in the magnitudes of the state changes caused by each action on the object, $$\begin{align} \label{eqn:mag-sim} \mathbb{sim}_{\text{mag}}(\dot{\bm{s}}^{O, a}, \dot{\bm{s}}^{O, b}) = \frac{1}{1 + \big| \, ||\dot{\bm{s}}^{O, a}|| - ||\dot{\bm{s}}^{O, b}|| \,\big|} \in [0,1], \end{align}$$ where $\dot{\bm{s}}^{O, a}$ and $\dot{\bm{s}}^{O, b}$ are the state changes of object $O$ caused by actions $a$ and $b$, respectively (as computed by Eq. `\ref{eq:state_change}`{=latex}). Combining these, the full per-object similarity of actions $a$ and $b$ for an object $O$ is $$\begin{align} \mathbb{sim}_{\text{obj}}&(O, a, b) = \max\left[\right. \nonumber \\ 0, \quad &\mathbb{sim}_{\text{cos}}\left(\dot{\bm{s}}^{O, a}, \dot{\bm{s}}^{O, b}\right) \nonumber \mathbb{sim}_{\text{mag}}\left(\dot{\bm{s}}^{O, a}, \dot{\bm{s}}^{O, b}\right) \nonumber \\ &\left.\right] \in [0, 1]. \end{align}$$ Finally, we define the full causal similarity metric between actions $a$ and $b$ as: $$\begin{align} \label{eqn:act-sim} \mathbb{sim}_{\text{csl}}(a, b) &= \frac{1}{D} \sum_{O=1}^{D} \left[ \mathbb{sim}_{\text{obj}}(O, a, b) \right] \nonumber \\ &\cdot exp \left ( \left [ \frac{1}{D}\sum_{O=1}^{D} \mathbb{sim}_{\text{obj}}(O, a, b) \right ] - 1\right ). \end{align}$$ The above computes the mean per-object similarity over all $D$ dynamic objects scaled by the exponentiated mean to accentuate differences between state changes caused by actions. Figure `\ref{fig:act-sim}`{=latex} provides a visual example of Equations (`\ref{eqn:cos-sim}`{=latex})-(`\ref{eqn:act-sim}`{=latex}) and shows how the causal similarity metric captures both directional and magnitude similarity.
Comparison of posterior distributions f(x) generated by Causal-PIK and a BO agent with an RBF kernel. Given the same initial observations {x1, x2}, Causal-PIK produces a more expressive posterior, effectively clustering actions based on their causal effects - the red area contains those actions predicted to move the ball closer to the goal while the blue area contains those predicted to push it further away. In contrast, the BO agent with the RBF kernel requires more iterations to develop an informative posterior.
**From Causal Similarity to `\kernel{}`{=latex}** The causal similarity function in Equation `\ref{eqn:act-sim}`{=latex} uses predicted state changes to quantify the correlation between actions that have similar causal effects. Consequently, these scores possess significant potential to be used as the kernel function of a GP that models the objective function of physical reasoning tasks. For a function to be a valid kernel, it must satisfy two fundamental properties: symmetry and positive semi-definiteness. Analyzing Equations (`\ref{eqn:cos-sim}`{=latex})-(`\ref{eqn:act-sim}`{=latex}), the reader can verify that $\mathbb{sim}_{\text{csl}}(a, b) = \mathbb{sim}_{\text{csl}}(b, a)$, thereby fulfilling the symmetry criterion. Furthermore, for any arbitrary pair of actions $a$ and $b$, their similarity score $\mathbb{sim}_{\text{csl}}(a, b)$ is always non-negative. As such, causal similarity satisfies both criteria of valid kernel function, so we define the `\kernel{}`{=latex} used in Algorithm `\ref{alg:overview}`{=latex} Line 4 as $\mathbb{sim}_{\text{csl}}(a, b)$. ## `\ourName{}`{=latex} objective function We use `\ourName{}`{=latex} to find an optimal action $\bm{x}$ that maximizes the objective function $f(\bm{x})$. On each attempt, we use an objective function $f(\bm{x})$ that quantifies the progress that executing action $\bm{x}$ makes towards a goal state $\bm{s}^{g}$. We find that for tasks with complex dynamics, it is important for the objective function to capture how close an action gets to reaching the goal state at *any* timestep in the $T$-timestep long episode, not just at the final timestep. This insight aligns with the concept of \`\`almost" reaching a goal, introduced by Gerstenberg et al. [@gerstenberg2016understanding]. To capture this, we use the closest distance $d_{c} = \min_{t = 1,...,T} \, \text{dist}(\bm{s}_t, \bm{s}^{g})$ to the goal state achieved at *any* timestep $t$ in an episode to compute the objective function $f(\bm{x})$ (where $\text{dist}(\cdot)$ denotes a distance function): \begin{small} \begin{equation} f(\bm{x}) = \begin{dcases*} {\left( 1-\frac{d_{\text{c}}}{\text{dist}(\bm{s}_0, \bm{s}^{g})} \right) \text{exp}(\beta d_{c})} & if $d_{c} < \text{dist}(\bm{s}_0, \bm{s}^{g})$\\ \omit\hfil$0$\hfil & $\text{otherwise}$ \end{dcases*} \label{eqn:reward} \end{equation} \end{small} ## Leveraging Causality for Efficient Sample Learning Unlike methods that learn only from direct observations, our approach leverages physical reasoning to infer the outcomes of untested actions that are predicted to share the same causal effect as observed ones. As a result, a single rollout allows our model to update its belief not just about the executed action, but also about actions that are predicted to produce a similar physical outcome. This significantly enhances sample efficiency and enables reasoning about alternative scenarios without exhaustive exploration. Our `\kernel{}`{=latex} explicitly encodes causal dependencies between actions and their physical consequences. In contrast to standard approaches that tend to cluster actions based on proximity in the feature space, `\kernel{}`{=latex} compares them by assessing their impact on the environment. Using counterfactual reasoning, the kernel distinguishes between causal effects -- those directly attributable to actions -- and confounding factors arising from the dynamics of the environment. This distinction is achieved by evaluating action outcomes against a counterfactual baseline where no action object is placed in the environment. In Equation `\ref{eqn:reward}`{=latex}, the shortest observed distance $d_{\text{c}}$ is normalized by the counterfactual baseline distance $\text{dist}(\bm{s}_0, \bm{s}^{g})$. Similarly, when computing the causal similarity of actions in Equation `\ref{eqn:act-sim}`{=latex}, the counterfactual baseline helps cluster actions that have no effect on the environment and emphasizes the attributable effects of other actions. # Experimental Setup {#sec:results} **Benchmarks:** Virtual Tools: We test `\ourName{}`{=latex} on the 20 original puzzles from the Virtual Tools benchmark. These puzzles involve placing an action object in a 2D dynamic physical environment to guide a red ball into the green goal area. Once an action object is placed, gravity is activated, and its effect can be observed.. The three-dimensional action space consists of: $x_{pos} \in [0,600]$, $y_{pos} \in [0,600]$, and $action_{obj} \in \{tool_1, tool_2, tool_3\}$, resulting in a total of 1,080,000 possible actions. PHYRE: We test `\ourName{}`{=latex} on the 25 puzzles from the BALL tier in PHYRE, following the cross-generalization setting. These puzzles require selecting the radius of a ball and placing it in a 2D dynamic physical environment to have the blue and green objects stay in contact for at least 3 seconds. As in Virtual Tools, gravity is activated upon placement, after which the action's effect is observed. The three dimensional action space consists of: $x_{pos} \in [0,256]$, $y_{pos} \in [0,256]$, and $action_{obj} \in [2,32]$, yielding a total of 2,555,904 possible actions. **Baselines and Ablations:** For Virtual Tools, we compare `\ourName{}`{=latex} against SSUP [@allen2020rapid], the state-of-the-art method for this benchmark. We also compare against the two naive baselines from @allen2020rapid, RAND and DQN. For PHYRE, we compare Causal-PIK against three categories of prior work. The first category consists of a method that uses a reduced action space of size 1,000 [@girdhar2020forward]. The second category consists of methods that use a reduced action space of size 10,000 [@ahmed2021physical; @qi2021learning]. Critically, by reducing the action space from $\sim$``{=html}2.5 million actions, these methods simplify the problem dramatically. The third category consists of methods that operate over the full action space [@harter2020solving], as we do. Additionally, we report results for three naive baselines from the original PHYRE paper [@bakhtin2019phyre]: RAND, MEM, and DQN. Finally, to assess the impact of the `\kernel{}`{=latex}, we conduct an ablation study by replacing it with a standard RBF kernel. This baseline highlights the expressiveness of the `\kernel{}`{=latex}, particularly in tasks with complex dynamics, where capturing underlying physical principles is key to finding solutions efficiently. **Human Baseline:** In addition to model baselines, we compare Causal-PIK against a human baseline for both benchmarks. A human study with Virtual Tools was conducted by @allen2020rapid. Here, we conducted a human experiment with PHYRE by recruiting $n = 50$ participants from Prolific and asking each participant to solve one variation of each of the 25 puzzles. Participants had 10 attempts to solve each puzzle by using the mouse to draw and place a valid action object. As with the model setup, each puzzle was initially presented as a freeze-frame, and then gravity was activated once the participant submitted their action (see `\autoref{ap:human}`{=latex} for details). The experiment was approved by the Stanford Institutional Review Board (IRB). **Evaluation Metric:** We measure the performance of our method using the AUCCESS metric introduced by @bakhtin2019phyre. This metric aggregates the success percentages via a weighted average, placing more emphasis on solving tasks with fewer attempts. AUCCESS is computed as follows: $$\begin{align} & AUCCESS = \frac{\sum_k w_k \cdot s_k}{\sum_k w_k}, \, \text{where} \nonumber \\ & w_k = log(k + 1) - log(k) ,\, \text{for} \, k \in \{1,...,\text{MAX\_ATT}\} \nonumber \\ & s_k: \text{percentage of tasks solved within } k \text{ attempts} \end{align}$$ We compare overall AUCCESS scores for each model, as well as correlation between scores for humans and models across individual puzzles. **Dynamics Model:** We use Region Proposal Interaction Networks [@qi2021learning] as the architecture. We train the model on puzzles that share underlying physical concepts similar to the test puzzles. For Virtual Tools, we train the model on 10 variations for each original puzzle. For PHYRE, for each of the 10-fold splits from @bakhtin2019phyre, we train a model exclusively on the fold's training set, ensuring that `\ourName{}`{=latex} is tested on previously unseen puzzles. At inference time, an image of the test puzzle with the action embedded and the initial bounding boxes of all dynamic objects in the scene are fed into the model. The model outputs the bounding boxes of the next $n_{pred}$ time steps. We set $n_{pred}$ to capture the initial steps, but not the full roll-out. The choice of $n_{pred}$ helps the dynamic model focus on learning the immediate causal effect of an action, which is then used to compare the similarity between actions. For further details, refer to Appendix `\ref{ap:dyn}`{=latex}. # Results and Discussion **Virtual Tools:** We evaluate `\ourName{}`{=latex} on the 20 Virtual Tools puzzles with 100 tests per puzzle. For each test, we limit the maximum number of attempts to 10. Table `\ref{tab:vt-results}`{=latex} includes a summary of the AUCCESS rate achieved by the different agents. `\ourName{}`{=latex} is 7 points higher than the best-performing baseline, requiring a lower number of attempts to solve the tasks. **PHYRE:** We evaluate `\ourName{}`{=latex} on the PHYRE-1B Cross test set, which comprises 25 tasks distributed across 10 folds, with 10 variations per task and 10 tests per fold. For each test, we limit the maximum number of attempts to 100. As shown in Table `\ref{tab:phyre-results}`{=latex}, `\ourName{}`{=latex} achieves an AUCCESS rate over 10 points higher than the best-performing baseline. We also compare our results to baselines that use a reduced action space [@girdhar2020forward; @ahmed2021physical; @qi2021learning], which is guaranteed to lead to a 100 AUCCESS score for 10K actions if an optimal oracle is used. As shown in Table `\ref{tab:phyre-results}`{=latex}, `\ourName{}`{=latex} performs comparably to baselines that utilize a drastically reduced action space. Our approach tackles a significantly harder problem by considering any point in the action space. @bakhtin2019phyre analyzed how the number of actions ranked by agents at test time affects performance. As shown in Figure 4 of their work, the AUCCESS of the DQN agent decreases as the number of ranked actions increases by orders of magnitude, up to a maximum of 100,000, which is still far from the 2,555,904 possible actions per puzzle in the full action space. We argue that discretizing the environment for action selection is an unrealistic constraint when aiming to develop generalist algorithms capable of solving complex physical reasoning tasks. **`\kernel{}`{=latex} versus RBF Kernel:** We compare the posterior distributions of our Bayesian Optimization (BO) agent using the `\kernel{}`{=latex} and the RBF kernel. As shown in Figure `\ref{fig:kernel_comp}`{=latex}, given the same set of initial observations, the BO Agent with the `\kernel{}`{=latex} produces a more expressive posterior that focuses on high-likelihood actions. Given two actions, one that moved the red ball further from the goal and another one that brought it closer but overshoots, the GP with the `\kernel{}`{=latex} constructs a posterior with three distinct regions: a red area on the top left side of the red ball, clustering actions that move the ball closer to the goal; a blue area on the top right, clustering actions that push it further away; and a yellow area representing unexplored actions that cause no movement. In contrast, the BO agent with the RBF kernel requires more steps to build an informative posterior. ::: center +-----------------------------------+--------------------------------------+ | **Model** | **AUCCESS** $\uparrow$ | +:==================================+:====================================:+ | RAND | 16.0~$\pm$``{=html}20.0~ | +-----------------------------------+--------------------------------------+ | DQN | 25.0~$\pm$``{=html}24.0~ | +-----------------------------------+--------------------------------------+ | SSUP [@allen2020rapid] | 58.0~$\pm$``{=html}27.0~ | +-----------------------------------+--------------------------------------+ | \arrayrulecolor{black!50} | 42.0~$\pm$``{=html}33.0~ | | \specialrule{0.3pt}{2.5pt}{2.5pt} | | | | | | Ours RBF | | +-----------------------------------+--------------------------------------+ | Ours Causal-PIK | **65.0**~$\pm$``{=html}25.0~ | +-----------------------------------+--------------------------------------+ | Humans [@allen2020rapid] | 53.25~$\pm$``{=html}23~ | +-----------------------------------+--------------------------------------+ : AUCCESS scores of Causal-PIK, its RBF ablation variant, and state-of-the-art model on the Virtual Tools benchmark. Results are based on a maximum of 10 attempts per task, with higher scores indicating more efficient problem-solving. {#tab:vt-results} ::: ::: center +---------------------------------------------------------------------------------------------------------------------+--------------------------------------+ | **Model** | **AUCCESS** $\uparrow$ | +:====================================================================================================================+:====================================:+ | Dec \[Joint\] [@girdhar2020forward] $^*$ | 40.3~$\pm$``{=html}8~ | +---------------------------------------------------------------------------------------------------------------------+--------------------------------------+ | \arrayrulecolor{black!50} | 18.5~$\pm$``{=html}5.1~ | | \specialrule{0.3pt}{2.5pt}{2.5pt} | | | | | | MEM$^\dagger$ | | +---------------------------------------------------------------------------------------------------------------------+--------------------------------------+ | DQN$^\dagger$ | 36.8~$\pm$``{=html}9.7~ | +---------------------------------------------------------------------------------------------------------------------+--------------------------------------+ | @ahmed2021physical$^\dagger$ | 41.9~$\pm$``{=html}8.8~ | +---------------------------------------------------------------------------------------------------------------------+--------------------------------------+ | RPIN [@qi2021learning]$^\dagger$ | 42.2~$\pm$``{=html}7.1~ | +---------------------------------------------------------------------------------------------------------------------+--------------------------------------+ | \arrayrulecolor{black!50} | 13.0~$\pm$``{=html}5.0~ | | \specialrule{0.3pt}{2.5pt}{2.5pt} | | | | | | RAND | | +---------------------------------------------------------------------------------------------------------------------+--------------------------------------+ | @harter2020solving | 30.24~$\pm$``{=html}8.9~ | +---------------------------------------------------------------------------------------------------------------------+--------------------------------------+ | Ours RBF | 27.70~$\pm$``{=html}9.68~ | +---------------------------------------------------------------------------------------------------------------------+--------------------------------------+ | Ours Causal-PIK | **41.6**~$\pm$``{=html}9.33~ | +---------------------------------------------------------------------------------------------------------------------+--------------------------------------+ | Ours Causal-PIK \@10$^+$ | 24.8~$\pm$``{=html}9.22~ | +---------------------------------------------------------------------------------------------------------------------+--------------------------------------+ | Humans \@10$^+$ | 36.6~$\pm$``{=html}10.2~ | +---------------------------------------------------------------------------------------------------------------------+--------------------------------------+ | \specialrule{1pt}{2.5pt}{2.5pt} | | | | | | `\small `{=latex}$^*$``{=html}1K reduced action space. $^\dagger$``{=html}10K reduced action space. | | +---------------------------------------------------------------------------------------------------------------------+--------------------------------------+ | `\small `{=latex}$^+$ max of 10 attempts per task. | | +---------------------------------------------------------------------------------------------------------------------+--------------------------------------+ : AUCCESS scores of Causal-PIK, its RBF ablation variant, and state-of-the-art models on PHYRE-1B Cross. Results are based on a maximum of 100 attempts per task and averaged over 10 test folds. Higher scores indicate more efficient problem-solving. {#tab:phyre-results} ::: **Human performance:** Participants found the puzzles from both benchmarks to be challenging. Causal-PIK achieved a higher AUCCESS score than humans on both benchmarks, except when it was restricted to 10 attempts in PHYRE. The high variance in AUCCESS rates for both humans and models suggests that the puzzles had mixed levels of difficulty. We computed AUCCESS scores across individual puzzles. On Virtual Tools, human per-puzzle scores were most correlated with SSUP ($r=0.71$), followed by Causal-PIK ($r=0.63$), then DQN ($r=0.32$). Causal-PIK may be less correlated with humans but still have a higher overall AUCCESS rate because it is able to solve several puzzles that humans find very difficult. This lowers the per-puzzle correlation, but highlights the overall performance of our method. On PHYRE, human scores were most correlated with Causal-PIK ($r=0.73$), followed by Causal-PIK \@10, which has limited attempts ($r=0.71$), then the RBF kernel baseline ($r=0.64$), and finally @harter2020solving ($r=0.55$). Overall, the high correlation in scores between humans and our model, even when restricted to a maximum of 10 attempts per puzzle, suggests high alignment in the types of physical dynamics that were found to be easy or difficult to reason about. **Resistance to noisy predictions** We conduct an analysis to study the impact of the accuracy of the dynamics model on the performance of our method. To demonstrate that existing predictions are inherently noisy due to the characteristics of the train and test sets, we train the PHYRE dynamics model on tasks from the test templates, ensuring prior exposure to similar puzzles. As a result, the L2 error for object bounding boxes improved to 3.56, compared to 19.3 ± 4.55 when tested on entirely unseen puzzles. Despite this difference in prediction accuracy, our method achieved an AUCCESS of 45, which is only 4 points higher than the 41.6 ± 9.33 AUCCESS we reported for the case with unseen dynamics. This demonstrates that even with a substantial increase in prediction error, the performance drop is small, indicating that our method remains resilient to noisy dynamic predictions. While improved dynamic predictions can enhance performance, our approach does not rely on perfect predictions, retaining robustness even in the presence of inaccuracies. **On the shortcomings of RL agents:** Reinforcement learning (RL) agents like DQN have shown remarkable performance in strategic games such as Atari and Go after extensive training, but they struggle in tasks that require physical reasoning. In the context of the Virtual Tools and PHYRE physical reasoning tasks, DQN fails to generalize effectively to unseen puzzles due to its limited understanding of the physical properties of objects and the basic rules governing their interactions. Previous studies on the PHYRE and Virtual Tools benchmarks [@li2024iphyreinteractivephysicalreasoning; @allen2020rapid] have underscored this limitation, demonstrating that RL agents struggle to understand the underlying physics involved beyond their strong mapping between states and actions. This shortcoming leads to a broad and inefficient unguided exploration process, requiring numerous attempts to discover promising areas. Furthermore, DQN fails to learn from past attempts, often repeating similar unsuccessful actions. In contrast, our approach focuses on learning from each failed attempt and uses the causality and physical insights from the Physics-Informed Kernel to guide exploration towards more promising areas. # Limitations `\ourName{}`{=latex} currently does not share knowledge across tasks. Enabling agents to recognize similarities between tasks and leverage past observations from tasks requiring similar physical reasoning remains an avenue for future work. By identifying regions of the action space that share underlying dynamics, agents could integrate prior knowledge to solve new tasks more efficiently. Another limitation arises from the noise introduced by causal effect predictions, which directly impacts performance. Poor predictions introduce misleading similarities, potentially guiding the agent in the wrong direction. Improving these predictions would improve the expressivity of the `\kernel{}`{=latex}. However, our results demonstrate that `\ourName{}`{=latex} remains robust despite this noise, suggesting potential for future sim-to-real transfer. Additionally, the physical reasoning tasks considered in this study involve a three-dimensional action space. Scaling to higher-dimensional action spaces would require modifications to the causal effect predictor to accommodate the added complexity. However, the kernel equations (`\ref{eq:state_change}`{=latex}-`\ref{eqn:act-sim}`{=latex}) would remain unchanged, as the new dimensions would be encoded within the state returned by the model. Thereby, the fundamental structure of our approach remains unchanged, preserving the kernel's ability to compare the immediate effects of high-dimensional actions. Moreover, Bayesian Optimization (BO) is expected to remain effective in larger search spaces, as prior work has demonstrated its robustness in high-dimensional action spaces in real-world applications [@antonova2019bayesian]. # Conclusion {#sec:conclusion} We introduce `\ourName{}`{=latex}, a novel approach that integrates a `\kernel{}`{=latex} with BO to reason about causality in `\taskType{}`{=latex}. By leveraging information from past failed attempts, our method enables agents to efficiently search for optimal actions, reducing the number of trials needed to solve tasks with complex dynamics. Our experimental results demonstrate that `\ourName{}`{=latex} outperforms state-of-the-art baselines, requiring fewer attempts on average to solve the puzzles from the Virtual Tools and PHYRE benchmarks. \newpage # Acknowledgments {#acknowledgments .unnumbered} We thank Kelsey Allen and Kevin Smith for their insightful discussions and valuable feedback. We also thank the reviewers for their thoughtful feedback, which helped improve the quality of the manuscript. The Stanford Institute for Human-Centered Artificial Intelligence (HAI), the Sloan Research Foundation, and Intrinsic provided funds to support this work. TG was supported by a grant from Cooperative AI. # Impact Statement {#impact-statement .unnumbered} This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here. \bibliographystyle{icml2025} \newpage \appendix \onecolumn # Dynamics Model {#ap:dyn} We train the dynamics models to predict causal effects of actions in the Virtual Tools and PHYRE benchmarks [@allen2020rapid; @bakhtin2019phyre]. Our model is based on the Region Proposal Interaction Networks (RPIN) architecture [@qi2021learning] and is trained on puzzles that share similar underlying physical principles with the test puzzles. For the Virtual Tools benchmark, we generate 10 variations of each of the 20 original puzzles, modifying object sizes and relative positions to create diverse scenarios (see Figure `\ref{fig:variations}`{=latex}). None of the original puzzles are included in training. For each puzzle variation, 300 actions are generated, distributed evenly across the three action object types. At least 50% of actions result in collisions between the action object and a dynamic object, while 10% simulate the absence of an action object in the environment. The remaining actions are randomly sampled using a Sobol generator to determine object placement. The inclusion of \"no-action\" cases enables the model to implicitly learn that stationary objects remain static when no net external force is applied. For the PHYRE benchmark, we train 10 separate dynamics models, one per fold. Each model is trained on 20 out of the 25 puzzles assigned to the training set for that fold. For each puzzle, we generate 500 actions, of which 350 result in failed rollouts, 150 result in successful rollouts, and 50 simulate the absence of an action object. Actions are randomly drawn from the 100,000 pre-selected actions provided by [@bakhtin2019phyre]. At inference time, the model receives an image of the puzzle, along with the initial bounding boxes of all dynamic objects and the action object. The model outputs the bounding boxes for the next $n_{\text{pred}}$ time steps. We set $n_{\text{pred}}$ to 20, which usually captures one collision but not the full roll-out. The choice of $n_{\text{pred}}$ helps the dynamic model focus on learning the immediate causal effect of an action, which is then used to compare similarity between actions.
Puzzle variations for one of the original puzzles in the benchmark. The puzzle variations share a similar underlying physical concept with the original puzzle, but have different dynamics due to varying object sizes and relative positions.
# PHYRE Human Experiment {#ap:human} We adapted the PHYRE-1B benchmark into a suite of online games using Planck.js [@shakiba2017planckjs], a JavaScript rewrite of the Box2D physics engine used in PHYRE [@bakhtin2019phyre]. The experiment was posted on Prolific, an online crowd-sourcing research platform. We recruited $n = 50$ participants (*age*: M = 37, SD = 11; *gender*: 20 female, 27 male, 1 non-binary, and 2 undisclosed; *race*: 29 White, 9 Black, 5 Asian, 1 American Indian/Alaska Native, 4 Multiracial, and 2 undisclosed) and compensated them at a rate of \$12/hour. All participants were native English speakers residing in the US. Participants were first given instructions for the task, including a \`\`playground" environment where they learned about the different objects and the dynamics of the world. They then solved a simple practice puzzle before moving on to the main puzzles. During the main portion of the experiment, they were given one variation of each of the 25 puzzles in Phyre-1B in a random order. On each puzzle, participants were initially presented with a freeze-frame of the scene. They attempted to solve the puzzle by using the mouse to draw and place the action object (a red ball) in any valid location in the scene. On each attempt, they watched the simulation run until it either succeeded, after which they continued on to the next puzzle, or timed out or all objects stopped moving, at which point they could try again. If they ran out of attempts (maximum 10), then they were also directed to the next puzzle. We enforced all the same physics parameters and time limit described in the original PHYRE benchmark. We recorded participants' actions ($x_{pos}$, $y_{pos}$, and $action_{obj} = r_{ball}$) on each attempt. Overall, participants spent an average of 1.7 minutes (SD = 1.4) on each puzzle.