Forward Prediction for Physical Reasoning

1 INTRODUCTION . When presented with a picture of a Rube Goldberg machine , we can predict how the machine works . We do so by using our intuitive understanding of concepts such as force , mass , energy , collisions , etc. , to imagine how the machine state would evolve once released . This ability allows us to solve real world physical-reasoning tasks , such as how to hit a billiards cue such that the ball ends up in the pocket , or how to balance the weight of two children on a see-saw . In contrast , physical-reasoning abilities of machine-learning models have largely been limited to closed domains such as predicting dynamics of multi-body gravitational systems ( Battaglia et al. , 2016 ) , stability of block towers ( Lerer et al. , 2016 ) , or physical plausibility of observed dynamics ( Riochet et al. , 2018 ) . In this work , we explore the use of imaginative , forward-prediction approaches to solve complex physical-reasoning puzzles . We study modern object-based ( Battaglia et al. , 2016 ; Sanchez-Gonzalez et al. , 2020 ; Watters et al. , 2017 ) and pixel-based ( Finn et al. , 2016 ; Ye et al. , 2019 ; Hafner et al. , 2020 ) forward-prediction models in simple search-based agents on the PHYRE benchmark ( Bakhtin et al. , 2019 ) . PHYRE tasks involve placing one or two balls in a 2D world , such that the world reaches a state with a particular property ( e.g. , two balls are touching ) after being played forward . PHYRE tasks are very challenging because small changes in the action ( or the world ) can have a very large effect on the efficacy of an action ; see Figure 1 for an example . Moreover , PHYRE tests models ’ ability to generalize to completely new physical environments at test time , a significantly harder task than prior work that mostly varies number or properties of objects in the same environment . As a result , physical-reasoning agents may struggle even when their forward-prediction model works well . Nevertheless , our best agents substantially outperform the prior state-of-the-art on PHYRE . Specifically , we find that forward-prediction models can improve the performance of physical-reasoning agents when the models are trained on tasks that are very similar to the tasks that need to be solved at test time . However , we find forward-prediction based agents struggle to generalize to truly unseen tasks , presumably , because small deviations in forward predictions tend to compound over time . We also observe that better forward prediction does not always lead to better physical-reasoning performance on PHYRE ( c.f . Buesing et al . ( 2018 ) for similar observations in RL ) . In particular , we find that object-based forward-prediction models make more accurate forward predictions but pixel-based models are more helpful in physical reasoning . This observation may be the result of two key advantages of models using pixel-based state representations . First , it is easier to determine whether a task is solved in a pixel-based representation than in an object-based one , in fully observable 2D environments like PHYRE . Second , pixel-based models facilitate end-to-end training of the forward-prediction model and the task-solution model in a way that object-based models do not in the absence of a differentiable renderer ( Liu et al. , 2019 ; Loper & Black , 2014 ) . 2 RELATED WORK . Our study builds on a large body of prior research on forward prediction and physical reasoning . We discuss most closely related work in this section and report additional prior work in Appendix B . Forward prediction models attempt to predict the future state of objects in the world based on observations of past states . Such models operate either on object-based ( proprioceptive ) representations or on pixel-based state representations . A popular class of object-based models use graph neural networks to model interactions between objects ( Kipf et al. , 2018 ; Battaglia et al. , 2016 ) , for example , to simulate environments with thousands of particles ( Sanchez-Gonzalez et al. , 2020 ; Li et al. , 2019 ) . Another class of object-based models explicitly represents the Hamiltonian or Lagrangian of the physical system ( Greydanus et al. , 2019 ; Cranmer et al. , 2020 ; Chen et al. , 2019 ) . While promising , such models are currently limited to simple point objects and physical systems that conserve energy . Hence , they can not currently be used on PHYRE , which contains dissipative forces and extended objects . Modern pixel-based forward-prediction models extract state representations by applying a convolutional network on the observed frame ( s ) ( Watters et al. , 2017 ; Kipf et al. , 2020 ) or on object segments ( Ye et al. , 2019 ; Janner et al. , 2019 ) . The models perform forward prediction on the resulting state representation using graph neural networks ( Kipf et al. , 2020 ; Ye et al. , 2019 ; Li et al. , 2020 ) , recurrent neural networks ( Xingjian et al. , 2015 ; Hochreiter & Schmidhuber , 1997 ; Cho et al. , 2014 ; Finn et al. , 2016 ) , or a physics engine ( Wu et al. , 2017 ) . The models can be trained to predict object state ( Watters et al. , 2017 ) , perform pixel reconstruction ( Villegas et al. , 2017 ; Ye et al. , 2019 ) , transform the previous frames ( Ye et al. , 2018 ; 2019 ; Finn et al. , 2016 ) , or produce a contrastive state representation ( Kipf et al. , 2020 ; Hafner et al. , 2020 ) . Physical reasoning tasks gauge a system ’ s ability to intuitively reason about physical phenomena ( Battaglia et al. , 2013 ; Kubricht et al. , 2017 ) . Prior work has developed models that predict whether physical structures are stable ( Lerer et al. , 2016 ; Groth et al. , 2018 ; Li et al. , 2016 ) , predict whether physical phenomena are plausible ( Riochet et al. , 2018 ) , describe or answer questions about physical systems ( Yi et al. , 2020 ; Rajani et al. , 2020 ) , perform counterfactual prediction in physical worlds ( Baradel et al. , 2020 ) , predict effect of forces ( Mottaghi et al. , 2016 ; Wang et al. , 2018 ) , or solve physical puzzles/games ( Allen et al. , 2020 ; Bakhtin et al. , 2019 ; Du & Narasimhan , 2019 ) . Unlike other physical reasoning tasks , physical-puzzle benchmarks such as PHYRE ( Bakhtin et al. , 2019 ) and Tools ( Allen et al. , 2020 ) incorporate a full physics simulator , and contain a large set of physical environments to study generalization . This makes them particularly suitable for studying the effectiveness of forward prediction for physical reasoning , and we adopt the PHYRE benchmark in our study for that reason . Inferring object representations involve techniques like generative models and attention mechanisms to decompose scenes into objects ( Eslami et al. , 2016 ; Greff et al. , 2019 ; Burgess et al. , 2019 ; Engelcke et al. , 2019 ) . Many techniques also leverage the motion information for better decomposition or to implicitly learn object dynamics ( Kipf et al. , 2020 ; van Steenkiste et al. , 2018 ; Crawford & Pineau , 2020 ; Kosiorek et al. , 2018 ) . While relevant to our exploration of pixel-based methods as well , we leverage the simplicity of PHYRE visual world to extract object-like representations simply using connected component algorithm in our approaches ( c.f . STN in Section 4.1 ) . However , more sophisticated approaches could help further improve the performance , and would be especially useful for more visually complex and 3D environments . 3 PHYRE BENCHMARK . In PHYRE , each task consists of an initial state that is a 256× 256 image . Colors indicate object properties ; for instance , black objects are static while gray objects are dynamic and neither are involved in the goal state . PHYRE defines two task tiers ( B and 2B ) that differ in their action space . An action involves placing one ball ( in the B tier ) or two balls ( in the 2B tier ) in the image . Balls are parameterized by their position and radius , which determine the ball ’ s mass . An action solves the task if the blue or purple object touches the green object ( the goal state ) for a minimum of three seconds when the simulation is rolled out . Figure 1 illustrates the challenging nature of PHYRE tasks : small variations can change incorrect actions ( Figure 1 ( a ) and ( c ) ) into a correct solution ( Figure 1 ( b ) ) . Each tier in PHYRE contains 25 task templates . A task template contains 100 tasks that are structurally similar but that differ in the initial positions of the objects . Performance on PHYRE is measured in two settings . The within-template setting defines a train-test split over tasks , such that training and test tasks can contain different instantiations of the same template . The cross-template setting splits across templates , such that training and test tasks never correspond to the same template . A PHYRE agent can make multiple attempts at solving a task . The performance of the agent is measured by the area under the success curve ( AUCCESS ; ( Bakhtin et al. , 2019 ) ) , which ranges from 0 to 100 and is higher when the agent needs fewer attempts to solve a task . Performance is averaged over 10 random splits of tasks or templates . In addition to AUCCESS , we also measure a forward-prediction accuracy ( FPA ) that does not consider whether an action solves a task . We define FPA as the percentage of pixels that match the ground-truth in a 10-second rollout at 1 frame per second ( fps ) ; we only consider pixels that correspond to dynamic objects when computing forward-prediction accuracy . Please refer to Appendix E for exact implementation details . 4 METHODS . We develop physical-reasoning agents for PHYRE that use learned forward-prediction models in a search strategy to find actions that solve a task . The search strategy maximizes the score of a task-solution model that , given a world state , predicts whether that state will lead to a task solution . Figure 2 illustrates how our forward-prediction and task-solution models are combined . We describe both types of models , as well as the search strategy we use , separately below . 4.1 FORWARD-PREDICTION MODELS . At time t , a forward-prediction model F aims to predict the next state , x̂t+1 , of a physical system based on a series of τ past states of that system . F consists of a state encoder e , a forward-dynamics model f , and a state decoder d. The past states { xt−τ , . . . , xt } are first encoded into latent representations { zt−τ , . . . , zt } using a learned encoder e with parameters θe , i.e. , zt = e ( xt ; θe ) . The latent representations are then passed into the forward-dynamics model f with parameters θf : f ( { xt−τ , . . . , xt } , { zt−τ , . . . , zt } ; θf ) → ẑt+1 . Finally , the predicted future latent representation is decoded using the decoder d with parameters θd : d ( ẑt+1 ; θd ) → x̂t+1 . We learn the model parameters Θ = ( θe , θf , θd ) on a large training set of observations of the system ’ s dynamics . We experiment with forward-prediction models that use either object-based or pixel-based state representations . Object-based models . We experiment with two object-based forward-prediction models that capture interactions between objects : interaction networks ( Battaglia et al. , 2016 ) and transformers ( Vaswani et al. , 2017 ) . Both object-based forward-prediction models represent the system ’ s state as a set of tuples that contain object type ( ball , stick , etc . ) , location , size , color , and orientation . The models are trained by minimizing the mean squared error between the predicted and observed state . • Interaction networks ( IN ; Battaglia et al . ( 2016 ) ) maintain a vector representation for each object in the system at time t. Each vector captures information about the object ’ s type , position , and velocity . A relationship is computed for each ordered pair of objects , designating the first object as the sender and the second as the receiver of the relation . The relation is characterized by the concatenation of the two objects ’ feature vectors and a one-hot encoding representing the sender object ’ s attribute of static or dynamic . The dynamics model embeds the relations into “ effects ” per object using a multilayer perceptron ( MLP ) . The effects exerted on an object are summed into a single effect per object . This aggregated effect is concatenated with the object ’ s previous state vector , from a particular temporal offset , along with a placeholder for external effects , e.g. , gravity . The result is passed through another MLP to predict velocity of the object . We use two interaction networks with different temporal offsets ( Watters et al. , 2017 ) , and aggregate the results in an MLP to generate the final velocity prediction . The decoder then sums the object ’ s predicted velocity with the previous position to obtain the new position of the object . • Transformers ( Tx ; Vaswani et al . ( 2017 ) ) also maintain a representation per object : they encode each object ’ s state using a 2-layer MLP . In contrast to IN , the dynamics model f in Tx is a Transformer that uses self-attention layers over the latent representation to predict the future state . We add a sinusoidal temporal position encoding ( Vaswani et al. , 2017 ) of time t to the features of each object . The resulting representation is fed into a Transformer encoder with 6 layers and 8 heads . The output representation is decoded using a MLP and added to the previous state to obtain the future state prediction . Pixel-based models . In contrast to object-based models , pixel-based forward-prediction models do not assume direct access to the attribute values of the objects . Instead , they operate on images depicting the object configuration , and maintain a single , global world state that is extracted by an image encoder . Our image encoder e is a ResNet-18 network ( He et al. , 2016 ) that is clipped at the res4 block . Objects in PHYRE can have seven different colors ; hence , the input of the network consists of seven channels with binary values that indicate object presence , consistent with prior work ( Bakhtin et al. , 2019 ) . The representations extracted from the past τ frames are concatenated before being input into the two models we study . • Spatial transformer networks ( STN ; ( Jaderberg et al. , 2015 ) ) split the input frame into segments by detecting objects ( Ye et al. , 2019 ) , and then encode each object using the encoder e. Specifically , we use a simple connected components algorithm ( Weaver , 1985 ) to split each frame channel into object segments . The dynamics model concatenates the object channels for the τ input frames , and predicts a rotation and translation for each channel corresponding to the last frame using a small convolutional network . The decoder applies the predicted transformation to each channel . The resulting channels are combined into a single frame prediction by summing them . Inspired by modern keypoint localizers ( He et al. , 2017 ) , we train STNs by minimizing the spatial crossentropy , which sums the cross-entropies of H×W softmax predictions over all seven channels . • Deconvolutional networks ( Dec ) directly predict the pixels in the next frame using a deconvolutional network that does not rely on a segmentation of the input frame ( s ) . The representations for the last τ frames are concatenated along the channel dimension , and passed through a small convolutional network to generate a latent representation for the t+ 1th frame . Latent represen- tation ẑt+1 is then decoded to pixels using a deconvolutional network , implemented as series of five transposed-convolution and ( bilinear ) upsampling layers , with intermediate ReLU activation functions . We found Decs are best trained by minimizing the per-pixel cross-entropy , which sums the cross-entropy of seven-way softmax predictions at each pixel .

Summary: This paper investigates the performance of several state-of-the-art forward-prediction models in the complex physical-reasoning tasks of the PHYRE benchmark. The authors have provided thorough evaluations of the models by ablating on different ways of representing the state (object-based or pixel-based), forms of model class (Interaction network, Transformers, Spatial transformer networks, etc.), and specific evaluations settings (within-template or cross-template), from which they made several interesting observations. For example, forward predictors with better pixel accuracy do not necessarily lead to better physical-reasoning performance. Their best-performing model also sets a new state-of-the-art on the PHYRE benchmark.