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 discusses the importance of forward prediction in physical reasoning, and particularly in the PHYRE benchmark: a dataset of physical tasks where the agent is asked to place a ball of a chosen radius in a 2d environment. The authors build a classifier that, given an initial state, predicts the probability of success in a given task. `Given the classifier, one can solve PHYRE tasks by sampling actions uniformly at random, and using actions that achieve good scores under the classifier. The classifier consists of an encoder, dynamics model, decoder, and a final classifier. The authors then evaluate two types of dynamics models (interaction nets and transformers) that act on two different types of representations (either images or object states).