Particle-based simulators of physical environments

Abstract

A system that uses a graph neural network to determine a representation of a physical environment at a new time step The new time step can be before or after a current time step based on representations of the physical environment at the current time step and one or more other time steps, e.g. one or more time steps before and / or after the current time step. The representation of the physical environment at the new time step may, for example, be used to generate an image of the physical environment at the new time step. The system can be used for controlling a robot interacting with the physical environment. Some examples of the techniques are specifically adapted for implementation using hardware accelerator units.

Description

PARTICLE-BASED SIMULATORS OF PHYSICAL ENVIRONMENTSCROSS-REFERENCE TO RELATED APPLICATION

[0001] This application claims priority to U.S. Provisional Application No. 63 / 586,395, filed on September 28, 2023. The disclosure of the prior application is considered part of and is incorporated by reference in the disclosure of this application.BACKGROUND

[0002] This specification relates to processing data using machine learning models.

[0003] Machine learning models receive an input and generate an output, e.g., a predicted output, based on the received input. Some machine learning models are parametric models and generate the output based on the received input and on values of the parameters of the model.

[0004] Some machine learning models are deep models that employ multiple layers of models to generate an output for a received input. For example, a deep neural network is a deep machine learning model that includes an output layer and one or more hidden layers that each apply a non-linear transformation to a received input to generate an output.SUMMARY

[0005] This specification generally describes a system implemented as computer programs on one or more computers in one or more locations that uses a graph neural network to determine a representation of a physical environment at a new time step based on representations of the physical environment at a current time step. This may, for example, be used to generate an image of the physical environment at the new time step.

[0006] According to one aspect there is provided a method performed by one or more data processing apparatus for determining a (e.g., latent) representation of a physical environment at a new time step before or after a current time step (e.g., a next time step). The method comprises obtaining data defining respective representations of the physical environment at the current time step and one or more other time steps (e.g., previous time steps before the current time step). The representation at a given time step comprises a collection of particle nodes. Each particle node comprises data defining a corresponding spatial location in the physical environment and a set of featurescharacterizing the spatial location at the time step. The method further comprises constructing a graph comprising the particle nodes of the representations of the physical environment at the current time step and the one or more other time steps. The method also comprises processing the graph using a graph neural network (GNN) to update the set of features of each of the particle nodes of the representation of the physical environment at the current time step based on the sets of features of a respective subset of the particle nodes of the representations of the physical environment at the one or more other time steps. Each subset is determined based on a spatial proximity of the spatial locations of the particle nodes of the representations of the physical environment at the one or more other time steps to the spatial location of the corresponding particle node of the representation of the physical environment at the current time step. The method further comprises processing the updated sets of features of the particle nodes of the representation of the physical environment at the current time step to determine a representation of the physical environment at the new time step.

[0007] Representations of the physical environment determined using the systems and methods described in this specification can be used to predict previous or future properties, such as previous or future configurations, of the physical environment based on representations of the physical environment at later or earlier times. For example, the motion of one or more (interacting) objects in the physical environment can be predicted. The representations of the physical environment may be referred to as particle-based representations or “point clouds”. The representation of the physical environment at the new time step may, for example, be used to generate an image of the physical environment at the new time step.

[0008] In some implementations, a downstream process can process the representation determined for the physical environment at the new time step (or at multiple new time steps) to generate one or more predicted values (e.g., images) that can be compared with one or more observed values (e.g., observed images) of the physical environment to verify the simulation. The observed values may be obtained, for example, using sensor readings from one or more sensors in the physical environment (e.g., lidar or radar sensors).

[0009] In some implementations, the method may be used to generate representations of the physical environment at a time step that is between a pair of time steps for which representations of the physical environment have been obtained (e.g., from a video of the physical environment). Such representations may, for example, be used to generateimages that realistically (e.g., smoothly) interpolate between images of the physical environment at the pair of time steps.

[0010] The representations of the physical environment at the time step and the one or more other time steps can be generated from images of the physical environment at the time step and the one or more other time steps. For example, the sets of features of the particle nodes can, for example, be generated as the output of an embedding neural network that processes image data (e.g., video data) characterizing the physical environment at the time step and the one or more other time steps.

[0011] In general, the graph neural network that updates the sets of features of the particle nodes can use message-passing between the nodes of the graph to propagate information and iteratively update the node embeddings (i.e., sets of features) by the exchange of information with nodes to which they are connected by edges of the graph neural network.

[0012] In some implementations, the method is used for controlling a mechanical robot interacting with the physical environment. In particular, the representations of the physical environment at one or more new time steps can be used to control an agent (e.g., a mechanical agent, such as a robot) to interact with the physical environment to perform a task in the physical environment. For example, the representations of the physical environment may be used to generate images showing how the physical environment should change when the task is performed correctly. The agent may then process the images to determine what actions to take to perform the task, e.g., by selecting actions to minimize differences between predicted images of the physical environment obtained during performance of the task by the agent with corresponding generated images showing how the physical environment should change when the task is performed correctly.

[0013] Some examples of the described techniques are specifically adapted for implementation using hardware accelerator units.

[0014] Particular embodiments of the subject matter described in this specification can be implemented so as to realize one or more of the following advantages.

[0015] Realistic simulators of physical environments are invaluable to many scientific and engineering disciplines. However conventional simulation systems can be very expensive to create and use. Building a conventional simulator can entail years of engineering effort, and often must trade off generality for accuracy in a narrow range of settings. Furthermore, high-quality simulators require substantial computationalresources, which makes scaling up prohibitive. Even with significant effort and tuning, it is often impossible to capture the physics of any particular real-world physical environment perfectly.

[0016] Some simulators may require access to privileged “ground truth” information about the physical environment, such as the poses, 3D shapes and positions of all objects in the physical environment at all time steps. By contrast, the methods described in this specification can be used to simulate physical environments without access to privileged information about the physical environment. In particular, the methods can allow physical environments to be simulated at later time steps from sequences of images (e.g., videos) of the physical environment at earlier time steps. Thus, the present methods can be used to model dynamical physical behavior even for physical environments for which an analytic or physics-based simulator is not available or performs poorly, e.g., as may be the case for frictional interactions between objects in the physical environment.

[0017] As the representations of the physical system used by the method are particlebased, the representations may be edited to delete or modify some of the particle nodes, or to insert new particle nodes, to change the physical environment. For example, the particle nodes corresponding to a particular object in the physical environment may be deleted from the representation. An image of the physical environment generated from the representation may then not include the object. In some cases, particle nodes corresponding to an object in one physical environment may be included in a representation of another physical environment, thereby introducing the object into the other physical environment (such that the object would appear in an image generated using the representation). Thus, a representation of a physical environment may be edited or composed via changes / additions / deletions to the particle nodes of the representation.

[0018] Unlike some existing approaches, the systems and methods described in this specification can be trained using next-time step image prediction, without the need for particle tracking. For example, the graph neural network can be trained jointly (e.g., end-to-end) with an encoder neural network and an image generation neural network, such that the graph neural network can learn to generate updated particle nodes that accurately model the motion of objects, or other such changes to a physical environment from image data. In particular, it is not necessary for the system to determinecorrespondences between particle nodes explicitly to determine velocity features, unlike some existing techniques.

[0019] As mentioned in the aspect above, the subsets of the particle nodes whose sets of features are used to update the respective sets of features of the particle nodes of the physical environment at the current time step are selected based on spatial proximity. This avoids difficulties associated with maintaining data defining correspondences between the particle nodes at different time steps (i.e., tracking particles), which are encountered in other graph neural network based simulators, particularly when such correspondences may not exist e.g., for objects rotating or moving in and out of occlusions. For example, less memory and computing power (e.g. processor cycles) may be used.

[0020] In general, the systems and methods described in this specification can be used to generate simulations of physical environments with greater accuracy and using fewer computational resources (e.g., memory and computing power) than some conventional simulation systems.

[0021] The particle-based architecture and inductive biases of the systems and methods described in this specification may enable quicker generalization than existing methods. That is, the methods described in this specification may require fewer training examples to achieve comparable performance to existing methods.

[0022] The details of one or more embodiments of the subject matter of this specification are set forth in the accompanying drawings and the description below. Other features, aspects, and advantages of the subject matter will become apparent from the description, the drawings, and the claims.BRIEF DESCRIPTION OF THE DRAWINGS

[0023] FIG. 1 shows an example system for determining a representation of a physical environment.

[0024] FIG. 2 illustrates processing performed by the system of FIG. 1 on exemplary input data representing a physical system.

[0025] FIG. 3 is a flow diagram of an example process for determining a representation of a physical environment.

[0026] FIG. 4 is a flow diagram illustrating operations performed by the graph neural network of the system of FIG. 1 on graph data.

[0027] FIG. 5 is a flow diagram of an example process for updating a representation of a physical environment, such as can be performed by the system of FIG. 1.

[0028] FIG. 6 is a flow diagram of an example process for controlling a robot interacting with a physical environment.

[0029] FIG. 7 shows examples of images generated by the described techniques.

[0030] FIG. 8 summarizes results for image quality as a function of rollout duration for sequences of images of exemplary physical environments obtained using the described techniques and using existing techniques for comparison.

[0031] Like reference numbers and designations in the various drawings indicate like elements.DETAILED DESCRIPTION

[0032] FIG. 1 shows a system implemented as computer programs on one or more computers in one or more locations that uses a graph neural network to determine a representation of a physical environment at a new time step. The system of FIG. 1 can be implemented as computer programs on one or more computers in one or more locations.

[0033] The system in FIG. 1 is a simulation system 100 that is configured to process input images 102 of a physical environment at a current timestep t and one or more other time steps, e.g., one or more time steps before the current time step Z-l, Z-2, so on. The system 100 generates, as an output, one or more predicted images 120 of the physical environment at a new time step before or after the current time step, e.g., a time step t +1 after the current time step.

[0034] In general, the physical environment is dynamic, e.g. it may comprise one or more moving objects and / or one or more moving or changing light sources, or the orientation of the physical environment relative to a viewing direction may change over time, or the spatial location of the physical environment relative to a view point may change. The physical environment is typically a three-dimensional physical environment.

[0035] The physical environment can be any appropriate environment, e.g., a real -world physical environment, e.g., a manufacturing environment, a warehouse environment, a roadway environment, and so forth. The physical environment can include one or more agents that interact with the environment, e.g., robotic agents, vehicles, and so forth.The physical environment can include any of a variety of objects, e.g., tools, packages,mechanical parts, electrical parts, walls, floors, roadway surfaces, conveyor belt surfaces, and so forth. The system 100 can be used to simulate how the objects in the physical environment interact, e.g., by translating, rotating, deforming, or otherwise changing over successive time steps.

[0036] The input and predicted images 102, 120 of the physical environment can, for example, comprise pixel values specifying one or more color values (or an intensity value in the case of a grayscale image) and a depth value for each of pixels in the image, e.g., the images may be RGB-D images in which the pixel values comprise values for each of a red, green and blue (RGB) color channel and a depth (D) value. The depth values may be obtained using lidar or radar, for example, or otherwise estimated from color values (e.g., RGB values) of the pixels in the image, e.g., using a depth estimation machine learning model. In some cases, the input images 102 may be obtained from a plurality of different viewpoints (viewing directions) of the physical environment. The input images 102 of the physical environment at the current time step and one or more other time steps can be obtained from one or more input videos. For example, each of the input images 102 can correspond to a respective frame in the one or more input videos.

[0037] The system 100 comprises an encoder neural network 106 that is configured to process the input images 102 to generate a respective representation of the physical environment at each of a plurality of time steps. Each representation comprises comprise a respective collection of particle nodes ?t108. Each of the particle nodes ?t108 comprises data defining a corresponding spatial location Xtin the physical environment and a set of features Ztcharacterizing the spatial location at a given time step t, i.e., each of the particle nodes108 comprises a latent point representation (Xt,Zt) for the given time step. The collection of particle nodes108 at a given time step may be referred to as a particle-based representation or “point cloud” in some instances. Each set of features (which can also be referred to as an embedding) can comprise an ordered collection of numerical values, e.g., a vector, matrix, or other tensor of numerical values. Updating the set of features may comprise updating one or more of the numerical values.

[0038] Each particle nodefor a given time step can, for example, correspond to a respective pixel in one of the one or more input images 102 for the time step. The spatial location Xtof each particle node ?tat the time step t can be determined from the location of the corresponding pixel in the input image 102 for the time step, togetherwith a depth value for the pixel. For example, the spatial location of a particle node can be determined from the (2D) location and depth value of the corresponding pixel using a transformation matrix for the camera that was used to obtain the input image 102.

[0039] The particle nodes 108 can, for example, represent one or more physical objects in the physical environment and each particle node can represent a point on an object located at the spatial location in the physical environment defined by the particle node and the set of features of the particle node may characterize the point on the object, e.g. one or more of a color, opacity, radiance or texture of the point. The term “object” can refer to any physical or tangible entity, including, e.g., surfaces, machinery, equipment, tools, products, or components thereof.

[0040] The encoder neural network 106 is configured to extract features from image data, and can have any suitable architecture, e.g., encoder neural network 106 can be a convolutional neural network, or a vision transformer neural network. As one example, the encoder neural network 106 can comprise a convolutional neural network having a U-Net architecture (i.e., a U-Net). Examples of U-Nets are described in Ronneberger et al. U-Net: Convolutional Networks for Biomedical Image Segmentation, arXiv: 1505.04597. In some implementations, the encoder neural network 106 processes the input images 102 to determine the respective sets of features Zt, Zt-x, ... , of the particle nodes Pt,... 108 without taking into account the depth values of the pixels in the input image 102, e.g., each RGB-D image may be processed by the encoder neural network 106 as an RGB image.

[0041] In cases where multiple input images are available for a given time step, the encoder neural network 106 can process each image separately to generate a respective collection of particle nodes for each image 102 that are then merged (e.g., concatenated) to generate the collection of particle nodes 108 for the time step.

[0042] To reduce memory requirements, in some implementations, the collections of particle nodes 108 at a given time step can be (e.g., randomly) sub-sampled, e.g., whilst maintaining approximately the same spatial distribution of particle nodes 108.

[0043] The simulation system 100 further comprises a representation update subsystem 110 that processes the particle nodes 108 at the current and one or more other time steps to obtain a collection of updated particle nodes ?t+1112 at a new time step. For example, the collection of particle nodes 108 comprise particle nodes 108 for the current time step ?tand particle nodes 108 for another time step that precedes the current time step 7’t-1, and the representation update subsystem 110 can be configuredto process the particle nodes 108 to obtain updated particle nodes Pt+1120 at a time step after the current time step.

[0044] In the present example, the representation update subsystem 110 is a dynamics model that predicts how the physical environment will evolve during a next time step, e.g., to obtain an updated representation of the physical environment at the new time step for use in planning a next action to be performed by an agent, such as a robot or mechanical agent. Other applications, such as retrodiction are also envisaged, however. In other words, the simulation system 100 is trained to produce predicted images 120 that, together with the input images 102, form a video sequence of the physical environment evolving in a dynamically consistent (e.g., physically realistic) way.

[0045] The representation update subsystem 110 processes the particle nodes 108 at the current and one or more other time steps to determine a graph 114, i.e., a data structure comprising a plurality of nodes and a plurality of edges, with each edge connecting a respective pair of nodes from the plurality of nodes. The graph 114 may not be fully connected in order to support processing of large numbers of particle nodes, which may be necessary for complex physical environments (e.g., physical environments comprising textures and / or complex geometries). That is, in some implementations, only some pairs of the nodes of the graph 114 are connected by a corresponding edge.

[0046] The representation update subsystem 110 comprises a graph neural network (GNN) 116 that is configured to processes the graph 114 to determine the updated particle nodes 112 at the new time step. The GNN 116 can, for example, be a hierarchical GNN (e.g., a 2-level hierarchical GNN) comprising multiple types of nodes, e.g., nodes other than the particle nodes 108 generated by the encoder neural network 106. Details of the graph 114 and the GNN 116 are described below in connection with FIGS. 4 and 5.

[0047] In general, the graph 114 is constructed so that the GNN 116 updates the set of features Ztof each of the particle nodes ?t108 at the current time step based on the sets of features Zt-rof a respective subset of the particle nodes 108 at the one or more other (e.g., earlier) time steps J’t-i. Each subset is determined based on a spatial proximity of the spatial locations of the particle nodes J>t-1108 at the one or more other time steps to the spatial location of the corresponding particle node108 at the current time step. The updated features of the particle nodes 108 at the current time step can then be processed to determine the updated particle nodes 112 to provide a representation of the physical environment at the new time step.

[0048] In one example, the GNN 116 can process (or “decode”) the updated features of the particle nodes ?tat the current time step to determine a corresponding spatial offset AX and feature offset AZ for each particle node ?t108 at the current time step. Each of the updated particle nodes Pt+1112 can then be determined by applying the corresponding offsets (AX, AZ) to the spatial location and features (Xt, Zt) of the particle node ?tat the current time step before the particle nodes 108 were processed using by the GNN 116, e.g., Pt+1'. (Xt+1,Zt+1) = (Xt+ AX,Zt+ AZ). The predicted spatial offset AX may directly correspond to objects moving in the predicted video, while the predicted feature offset AZ may allow for appearance changes during the time step, such as may be caused by motion, such as a shadow moving beneath a flying object.

[0049] The process of updating the particle nodes 108 at the current time step can be applied recursively to the updated particle nodes 120 to obtain representations of the physical system at each of a plurality of new time steps. That is, the representation update subsystem can recursively update the particle nodes 108, 112 to obtain respective updated particle nodes at a plurality of new time steps. Such a process can be referred to as a multistep “rollout”. For example, the updated particle nodes 112 can be used to update the graph 114 such that the particle nodes 120 for the new time step are now the particle nodes for the current time step, and the particle nodes 108 at the previous current time step are the particle nodes at the other time step, and so on.

[0050] The simulation system 100 further comprises an image generation neural network (“renderer”) 118 that is configured to process the updated particle nodes 112 and to generate, as an output, one or more predicted images 120 of the physical environment at the new time step.

[0051] In some implementations, the image generation neural network 118 may be a scene representation neural network, that, after training, can generate synthetic images of the physical environment from arbitrary viewpoints. As one example, the scene representation neural network can be a Neural Radiance Fields (NeRF) model that represents radiance with a neural field that reproduces the geometric structure and appearance of a scene, allowing the use of backpropagation to reconstruct a set of input images. In particular, the NeRF model can predict the radiance and occupancy in space, i.e., the underlying space geometry, as part of rendering an image of a scene from a given viewpoint.

[0052] In particular, a NeRF model may take as input a camera pose and generates as output a synthetic image of the scene that appears as if the image was taken by a camera having the input camera pose. In some cases, the NeRF model also receives as input the camera intrinsics and generates as output a synthetic image that appears as if the image was taken by a camera having the input camera pose and having the input camera intrinsics.

[0053] The method can train any of a variety of NeRF models that make use of any of a variety of NeRF variants. Examples of such models and loss functions for training these models include those described in J. T. Barron, B. Mildenhall, D. Verbin, P. P. Srinivasan, and P. Hedman, “Mip-nerf 360: Unbounded anti-aliased neural radiance fields,” CoRR, vol. abs / 2111.12077, 2021. T. M'uller, A. Evans, C. Schied, and A. Keller, “Instant neural graphics primitives with a multiresolution hash encoding,” ACM Trans. Graph., vol. 41, pp. 102: 1-102: 15, July 2022 D. Verbin, P. Hedman, B. Mildenhall, T. E. Zickler, J. T. Barron, and P. P. Srinivasan, “Ref-nerf: Structured viewdependent appearance for neural radiance fields,” CoRR, vol. abs / 2112.03907, 2021 J. T. Barron, B. Mildenhall, M. Tancik, P. Hedman, R. Martin-Brualla, and P. P. Srinivasan, “Mip-nerf: A multiscale representation for antialiasing neural radiance fields,” CoRR, vol. abs / 2103.13415, 2021.

[0054] The updated particle nodes 112 can be processed to generate query data for querying the scene generation neural network. For example, the query data can comprise input feature data at a plurality of spatial locations along a ray through the physical environment and the scene representation neural network can then, for example, generate corresponding pixel data for the ray, e.g., color and / or density (occupancy) values for the ray. The input feature data at each spatial location along the ray can be determined by combining the features of a respective subset of the updated particle nodes 112 that have spatial locations proximate the spatial location along the ray, e.g., the k (approximate) nearest neighbors of the spatial location along the ray, where & is a positive integer (e.g., k = 16). Combining the respective features of the updated particle nodes 112 in the subset can be performed by averaging, e.g., using one or more kernels. For example, the respective features of each of the updated particle nodes 112 in the subset can be weighted according to the radial distance (e.g., using a Gaussian radial function) between the spatial location of the updated particle node 112 and the spatial location along the ray, and the weighted features then summed to obtain combined feature data for the spatial location along the ray. As one example, respectivecombined feature data can be obtained using one or more concentric kernels that each comprise a different respective Gaussian radial function parameterized in terms of a corresponding radius and bandwidth (or variance). The combined feature data for each of the kernels can then be combined, e.g., by concatenation, to obtain the input feature data for the spatial location along the ray.

[0055] As the simulation system 100 can be used for applications other than generating predicted images, it is not essential that the simulation system 100 comprises an image generation neural network 118. As one example, the updated particle nodes 112 can be processed by a policy neural network to generate an action distribution output that includes a respective action probability for each of a set of possible actions that can be performed by an agent (such as a mechanical agent or robot) interacting with the physical environment to perform a predetermined task in the physical environment.

[0056] In cases where the physical environment is a virtual physical environment, the representations of the physical environment may be used to train the agent to perform the task in the virtual physical environment, e.g., by generating images of the virtual physical environment showing the task being performed correctly by the agent. The trained agent may then be used to perform the task in a real-world physical environment. Such a procedure may be referred to as “sim-to-real” training.

[0057] FIG. 2 provides an illustration of how the simulation system 100 processes input images 202A, 202B at a current time step obtained from respective cameras 402A, 402B observing a physical environment 404 comprising one or more objects, in this example, a deformable block 404A that is vertically displaced above a planar surface 404B. For clarity, the input images of the physical environment 404 at a time step before the current time step are not shown in FIG. 2.

[0058] The input images 202A, 202B (and input images at an earlier time step before the current time step) are processed by the encoder neural network 106 of the simulation system 100 to obtain particle nodes 208 at the current time step (together with particle nodes 108 for the earlier time step), which are then updated using the representation update subsystem 110 to determine a collection of updated particle nodes 212 at a new (later) time step. In the present example, the GNN 116 of the representation update subsystem 110 has been trained to update the particle nodes 208 at the current time step in a physically realistic manner based on the particle nodes at an earlier time step, which in this case corresponds to the cube 404B falling under gravity to collide with the planar surface 404B such that it deforms on impact, as can be seen by comparing the particlenodes 208 at the current time step with the updated particle nodes 212 at the new time step.

[0059] The updated particle nodes 212 are then processed by the image generation neural network 118 to generate image 220A, 220B from different respective viewpoints (which in this case are different viewpoints from the viewpoints of the cameras 402A, 402B).

[0060] The systems described in this specification, such as the simulation system 100, can be trained using training data comprising “ground truth” representations of physical environments (e.g., representations obtained from video data) at a number of time steps during which a task is being performed. The system may then learn to predict how the physical environment evolves (e.g., as at task is being performed by a robot or mechanical agent) by generating representations of the physical environment that match or approximate the corresponding ground truth representations at each of the time steps. Thus, the system(s) may be trained from video data by (self-)supervised learning.

[0061] As one example, the simulation system 100 can be trained using video data (e.g. RGB-D video data) of an evolving physical environment, e.g., a real-world physical environment, or a simulated physical environment. The training can comprise adjusting the parameters of the encoder neural network 106, the GNN 114 and the image generation neural network 118 (e.g., scene representation neural network) to optimise an objective (loss) function that compares predicted images 120 generated by the simulation system 100 for a plurality of time steps from two or more input images 102 (video frames) with respective target images for the plurality of time steps from the video data. Thus, the simulation system 100 can be trained end-to-end using video data.

[0062] As one example, the objective function may comprise an average over the expected squared pixel reconstruction error for each of the one or more predicted images 120, e.g.:in which c- and c- are, respectively, a vector of predicted pixel values and a vector of ground truth pixel values for the z-th pixel, where i is an index that runs over the pixels of the predicted and ground truth images, t labels the time step, with t = 1, 2 corresponding to the time steps of the two input images 102, and T is the number of time steps. If the predicted images 120 and ground truth images at a given time step tcomprise images from different viewpoints then the pixel index z can extend over the different viewpoints.

[0063] FIG. 3 is a flow diagram of an example process 300 for determining a representation of a physical environment. The process of FIG. 3 may be implemented by one or more computers in one or more locations; for convenience the process is described with reference to the system of FIG. 1.

[0064] The process 300 comprises, as an initial step 302, obtaining data defining respective representations of the physical environment at the current time step and one or more other time steps (e.g. previous time steps before the current time step). The representation at a given time step comprises a collection of particle nodes. Each particle node comprises data defining a corresponding spatial location in the physical environment and a set of features characterizing the spatial location at the time step.

[0065] The method further comprises constructing 304 a graph comprising the particle nodes of the representations of the physical environment at the current time step and the one or more other time steps.

[0066] The method also comprises processing 306 the graph using a graph neural network to update the set of features of each of the particle nodes of the representation of the physical environment at the current time step based on the sets of features of a respective subset of the particle nodes of the representations of the physical environment at the one or more other time steps. Each subset is determined based on a spatial proximity of the spatial locations of the particle nodes of the representations of the physical environment at the one or more other time steps to the spatial location of the corresponding particle node of the representation of the physical environment at the current time step.

[0067] The method further comprises processing 308 the updated sets of features of the particle nodes of the representation of the physical environment at the current time step to determine a representation of the physical environment at the new time step.

[0068] As noted above, the representations of the physical environment at the current time step and other time steps may be obtained from one or more sequences of images (e.g. one or more videos) of the physical environment at the current time step and other time steps. In some cases, the images may be obtained from a plurality of different viewpoints. The images may, for example, comprise pixel values specifying one or more color values (or an intensity value in the case of a grayscale image) and a depth value for each of pixels in the image, e.g. the images may be RGB-D images in whichthe pixel values comprise values for each of a red, green and blue (RGB) color channel and a depth (D) value. The depth values may be obtained using lidar or radar, for example, or otherwise estimated from color values (e.g. RGB values) of the pixels in the image, e.g. using a depth estimation machine learning model.

[0069] The subsets of the particle nodes of the representations of the physical environment at the one or more other time steps may be proper subsets of the particle nodes of the representations of the physical environment at the one or more other time steps, i.e., a subset that does not contain all of the particle nodes of the representations of the physical environment at the one or more other time steps. The subset can be selected based on spatial proximity according to whether a proximity criterion is satisfied for each of the particle nodes of the representations of the physical environment at the one or more other time steps, e.g., whether a distance between the spatial location of the particle node and the spatial location of the corresponding particle node of the representation of the physical environment at the current time step is less than a (pre-determined) threshold.

[0070] The method may also comprise generating an image of the physical environment at the new time step based on the representation of the physical environment at the new time step.

[0071] Representations of the physical environment may be determined for one or more additional time steps after (or before) the new time step by iteratively repeating the method using the representation that was determined for the new time step as the representation of the current time step (and so on). The representations can be used to generate a sequence of images of the physical environment at each of the time steps, e.g. a video of the physical environment. The process of obtaining the representations this way may be referred to as a “rollout”. Representations for any number of new time steps (e.g. from 1 to 64 time steps before or after an initial current time step) may be determined, although in some cases the quality of the representations may degrade for time steps far from an initial current time step.

[0072] The determined representation(s) and / or the generated image(s) for a given time step may be referred to as a simulation of the physical environment at the time step.

[0073] Representations of physical environments obtained using the techniques described in this specification may be used in a number of downstream processes.

[0074] In some implementations, as noted above, the techniques described in this specification can be used for real-world control such as controlling an electronic ormechanical agent (such as a robot) in a real-world environment to perform a task, e.g., using the representation of the physical environment at the new (e.g., next) time step for predictive control or as part of an optimal control system controlling the agent. As one example, the representation of the physical environment at the new time step may be used in this way to assist a robot in manipulating a deformable object.

[0075] In more detail, the physical environment may be or comprise a real-world environment including a physical object e.g. an object to be picked up, manipulated or avoided. A method of controlling the agent can include determining a representation of the physical environment at the new (e.g., next) time step and using the representation of the physical environment at the new time step to determine the state of the physical environment at the new time step by determining a predicted representation of the location, shape or configuration of the physical object e.g., when subject to a force or deformation e.g. from an actuator of a robot.

[0076] The method can involve controlling the robot using a predicted representation of the physical environment at the new (e.g. next) time step to manipulate the physical object, e.g. using the actuator. For example the method can involve controlling the robot using the predicted representation to manipulate the physical object towards a target location, shape or configuration of the physical object by controlling the robot to optimize an objective function dependent upon a difference between the predicted representation and the target location, shape or configuration of the physical object. Controlling the robot may involve providing control signals to the robot based on the predicted representation to cause the robot to perform actions, e.g., using an actuator of the robot, to manipulate the physical object to perform a task.

[0077] In some implementations, the above described systems and methods may be used for real-world control, in particular optimal control tasks, e.g., to assist a robot in manipulating a deformable or rigid object. Thus, as previously described, the physical environment may comprise a real-world environment including a physical object e.g., an object to be picked up or manipulated. The representation of the physical environment at the current or one or more other (e.g. previous) time steps may define a representation of a shape or configuration of the physical object at the time step. The representation of the physical environment at the new time step may define a predicted representation of the shape or configuration of the physical object e.g., when subject to a force or deformation e.g., from an actuator of a robot. The method may further comprise controlling the robot using the predicted representation to manipulate thephysical object, e.g., using the actuator, towards a target location, shape or configuration of the physical object by controlling the robot to optimize an objective function dependent upon a difference between the predicted representation and the target location, shape or configuration of the physical object. Controlling the robot may involve providing control signals to the robot based on the predicted representation to cause the robot to perform actions, e.g., using an actuator of the robot, to manipulate the physical object to perform a task. For example this may involve controlling the robot, e.g., the actuator, using a reinforcement learning process with a reward that is at least partly based on a value of the objective function, to learn to perform a task which involves manipulating the physical object.

[0078] In general, various orderings of the new time step, the current time step and the one or more other time steps are possible. For example, in some implementations the ordering may be: (i) the new time step is after the current time step and the one or more other time steps are before the current time step; or (ii) the new time step is before the current time step and the one or more other time steps are after the current time step; or (iii) the new time step is after the current time step and the one or more other time steps are after the new time step; or (iv) the new time step is before the current time step and the one or more other time steps are before the new time step.

[0079] For implementations where the new time step is after the current time step and the one or more other time steps are before the current time step, the representation of the physical environment at the new time step may be a forward prediction (e.g. an extrapolation) of the representations of the physical environment at the earlier time steps.

[0080] Conversely, for implementations where the new time step is before the current time step and the one or more other time steps are after the current time step, the representation of the physical environment at the new time step may be a backwards prediction (e.g. an extrapolation) of the representations of the physical environment at the later time steps. Such representations may, for example, be used for forensic purposes, e.g. to analyse how the physical environment at the current timestep may have been obtained / reached.

[0081] Alternatively, when the new time step is between the current time step and the one or more other time steps (orderings (iii) and (iv) in the preceding paragraph), the representation of the physical environment at the new time step may be an interpolation of the representations of the physical environment at the earlier and later time steps.Such representations may, for example, be used to generate respective images of the physical environment that can be inserted between two frames of a video to smoothly or realistically interpolate between the two frames, e.g. when a frame between the two frames is missing or corrupted. The representation of the physical environment at the new time step may be used for “in-filling” a portion of an image of the physical scene that has been temporarily occluded.

[0082] FIG. 4 illustrates a message passing process 400 performed on an exemplary graph 414 by the GNN 116 of the system of FIG. 1.

[0083] The graph 414 comprises a plurality of particle nodes at the current step 408A and a plurality of particle nodes at one or more other time steps 408B. The particle nodes at the current time step 408A are depicted in FIG. 4 using unfilled circles with a solid outline, whilst the particle nodes at the one or more other time steps 408B are depicted using unfilled circles with a broken outline.

[0084] Directly connecting the particle nodes to one another can generate an excessive number of edges in the graph 414. Therefore, a hierarchical structure can be used for the graph 414 to reduce the computational resources needed to maintain and process the graph 414. For example, in addition to the particle nodes 408A, 408B at the current and one or more other time steps, the graph 414 can also comprise auxiliary nodes 409A, 409B. Generally, the number of auxiliary nodes 409A, 409B is (significantly) less than the number of particle nodes at the current time step 408A. Each auxiliary node 409A, 409B comprises a corresponding spatial location in the physical environment and a corresponding set of features characterizing the spatial location. The spatial locations of each auxiliary node 409A, 409B can, for example, be determined by (randomly) sampling the spatial locations of the particle nodes 408A at the current time step, e.g., such that the spatial distribution of auxiliary nodes 409 approximates the spatial distribution of the particle nodes 408A at the current time step. The features for each auxiliary node 409A, 409B can, for example, be initialized to default values (e.g., zero).

[0085] The graph 414 may comprise: (i) edges 410 from the particle nodes at the current time step and the one or more other time steps 408A, 408B to the auxiliary nodes 409 A, 409B; (ii) edges 412A, 412B between pairs of the auxiliary nodes 409 A, 409B; and (iii) edges 415 from the auxiliary nodes 409A, 409B to the particle nodes of the representation of the physical environment at the current time step 408A. Each edge in the graph 414 can comprise a representation of a separation between the spatiallocations of the nodes joined by the edge. For example, each edge can encode the Euclidean distance between the nodes joined by the edge.

[0086] The edges 410 from the particle nodes 408 A, 408B to the auxiliary nodes 409 A, 409B can be generated based on the spatial proximity of the particle nodes to the auxiliary nodes, such that only a subset of (spatially proximate) particle nodes are connected to each auxiliary node. For example, each subset of spatially proximate particle nodes 408A, 408B can be selected such that each particle node in the subset is within a first pre-determined cut-off distance (e.g., radius) from the auxiliary node 409A, 409B. In some implementations, each particle node 408A, 408B is connected to at most a pre-determined number (e.g., two) auxiliary nodes 409 A, 409B.

[0087] The edges 412A, 412B between pairs of the auxiliary nodes 409 A, 409B can be generated based on the spatial proximity of the auxiliary nodes to one another, such that only spatially proximate auxiliary nodes are connected. For example, edges 412A, 412B can be generated only for pairs of auxiliary nodes 409A, 409B that are spatially separated by less than a second pre-determined cut-off distance. The second predetermined cut-off distance can be greater than the first pre-determined cut-off distance, for example, such that the auxiliary nodes 409A, 409B are able to transfer information over greater distances.

[0088] The GNN 116 can update the graph 414 by message passing, e.g., as illustrated in the process 400 of FIG. 4, which comprises a first round of message passing 402 in which the respective sets of features of the auxiliary nodes 409A, 409B are updated based on the respective sets of features of the particle nodes 408A, 408B in the graph 414 having an edge 410 from the particle node to the auxiliary node. For example, the message passing 402 can comprise: (i) updating the features of the edges 410 between the particle nodes 408A, 408B and the auxiliary nodes 409A, 409B based on the features of the particle nodes 408 A, 408B; and then (ii) updating each auxiliary node 409 A, 409B based on the features of the corresponding edges 410 between the particle nodes 408 A, 408B and that auxiliary node 409 A, 409B.

[0089] The process 400 can further comprise updating 404 the set of features of each of the auxiliary nodes 409A, 409B based on the sets of features of each of the auxiliary nodes in the graph having an edge 412A, 412B to the auxiliary node. The updating 404 can, for example, comprise a plurality of message passing steps between the auxiliary nodes 409A, 409B. Each message passing step can, for example, comprise, for each edge 412A, 412B between the auxiliary nodes 409 A, 409B: (i) updating the set offeatures of the edge based on the set of features of the auxiliary node 409A, 409B from which edge originates; and then (ii) updating the set of features of the auxiliary node 409A, 409B at which the edge terminates based on the set of features of the edge.

[0090] The process 400 can further comprise updating 406 the set of features of each of the particle nodes 408A at the current time step based on the sets of features of each of the auxiliary nodes 409 A, 409B in the graph 414 having an edge 415 to the particle node 408 A. For example, each particle node at the current time step 408 A can be updated by, for each of the auxiliary nodes having an edge to the particle node: (i) updating the set of features associated with the edge 415 from the auxiliary node 409 A, 409B to the particle node based on the set of features of the particle node and the set of features of the auxiliary node; and then (ii) updating the set of features of the particle node based on the sets of features associated with each of the edges from an auxiliary node to the particle node.

[0091] One suitable architecture for the GNN 116 is an Encode-Process-Decode architecture, e.g., as described in Pfaff el al. “Learning mesh-based simulation with graph networks" in International Conference on Learning Representations, 2021, although it will be appreciated that other architectures can be used.

[0092] Some example implementations that are adapted for hardware acceleration are now described.

[0093] One advantage of implementations of the described systems and methods is that, they can be configured for hardware acceleration. In such implementations, the method is performed by data processing apparatus comprising one or more computers and including one or more hardware accelerators units e.g. one or more GPUs (Graphics Processing Units) or TPUs (Tensor Processing Units). These implementations involve updating the graph 114, 414 (e.g., the auxiliary nodes of the graph) at each of one or more update iterations comprises updating the graph using a processor system comprising L message passing blocks, where each message passing block has the same neural network architecture and a separate set of neural network parameters. Then the method comprises applying the message passing blocks sequentially to process the data defining the graph over multiple iterations, and using the one or more hardware accelerators to apply the message passing blocks sequentially to process the data defining the graph. In implementations, the processing using the message passing blocks, i.e., the processor system, is distributed over the hardware accelerators.

[0094] Thus there is also provided herein a simulation method which is specifically adapted for implementation using hardware accelerator units, unlike some conventional approaches which are not capable of taking advantage of hardware acceleration. The one or more computers that perform the method may include multiple hardware accelerators, and wherein the method comprises distributing the processing using the message passing blocks over the hardware accelerators.

[0095] FIG. 5 shows an example process for updating particle nodes. The process of FIG. 5 may be implemented by one or more computers in one or more locations; for convenience the process is described with reference to system 100 of FIG. 1.

[0096] The process 500 of FIG. 5 comprises generating 502 a plurality of auxiliary nodes (which may alternatively be referred to as “abstract” nodes) to construct the graph 114, 414 which additionally comprises the particle nodes of the representation of the physical environment at the current time step and the one or more other time steps. Each auxiliary node can comprise data defining a corresponding spatial location in the physical environment and a set of features characterizing the spatial location. The graph 114, 414 may comprise edges: (i) from the particle nodes of the representations of the physical environment at the current time step and the one or more other time steps to the auxiliary nodes; (ii) between pairs of the auxiliary nodes; and (iii) from the auxiliary nodes to the particle nodes of the representation of the physical environment at the current time step.

[0097] For each of the edges from a particle node to an auxiliary node, a distance between the spatial location of the particle node and spatial location of the auxiliary node may be less than a first predetermined threshold. Additionally or alternatively, for each of the edges between pairs of auxiliary nodes, a distance between the respective spatial locations of the auxiliary nodes may be less than a second predetermined threshold. The spatial locations of the auxiliary nodes may be chosen by sampling the spatial locations of the particle nodes of the representation of the physical environment at the current time step, e.g. using uniform random (or pseudorandom) sampling.

[0098] The process 500 comprises processing the graph 114 using the graph neural network 116 to update the set of features of each of the particle nodes of the representation of the physical environment at the current time step. In particular, the process 500 may comprise, for each of the auxiliary nodes, updating 504 the set of features of the auxiliary node based on the set of features of each of the particle nodes in the graph having an edge from the particle node to the auxiliary node.

[0099] The process 500 can further comprises, for each of one or more update iterations (any integer number of iterations greater than zero, e.g. 10 iterations), updating the set of features of each of the auxiliary nodes based on the sets of features of each of the auxiliary nodes in the graph having an edge to the auxiliary node. The number of update iterations may be pre-defined or else determined based on a convergence criterion having been satisfied, e.g. a convergence criterion that is satisfied when values of the sets of features of the auxiliary nodes change by less than a pre-determined amount after the updating.

[0100] In some implementations, the data defining the auxiliary nodes may be stored on a different data processing apparatus from the data processing apparatus that stores the data defining the particle nodes. Thus, the one or more update iterations for updating the set of features of each of the auxiliary nodes based on the sets of features of each of the auxiliary nodes in the graph having an edge to the auxiliary node, may be carried out by the different data processing apparatus. Computational resources of the data processing apparatus that stores the data defining the particle nodes may therefore, in some implementations, not be used during the update iterations. The different data processing apparatus may, for example, be optimised or otherwise configured (e.g. by configuring one or more hardware accelerators, such as TPUs or GPUs) to perform (e.g. bidirectional) updates of the auxiliary nodes quickly and / or efficiently, e.g. more quickly and / or efficiently than the data processing apparatus that stores the data defining the particle nodes. In some cases, the particle nodes may be stored on a plurality of data processing apparatus that transmit the sets of features of the particle nodes to one or more data processing apparatus that store the auxiliary nodes. After updating the auxiliary nodes, the one or more data processing apparatus that store the auxiliary nodes may then transmit the sets of features of the auxiliary nodes to the one or more data processing apparatus that store the particle nodes of the representation of the physical environment at the current timestep.

[0101] The process 500 also comprises, for each of the particle nodes of the representation of the physical environment at the current time step, updating 508 the set of features of the particle node based on the sets of features of each of the auxiliary nodes in the graph having an edge to the particle node, e.g., as described above in connection with FIG. 1

[0102] FIG. 6 shows an example process 600 for controlling a robot interacting with a physical environment. The process 600 comprises repeatedly performing the process300 of FIG. 3 (e.g., using the simulation system 100 of FIG. 1) to determine 602 representations of the physical environment for a sequence of time steps, the representations characterizing the physical environments that would be obtained when the robot performs a sequence of actions over the sequence of time steps. The process 600 further comprises processing 604 the representation of the physical environment at a final time step in the sequence of time steps to determine whether a success criterion is satisfied, and in response to determining that the success criterion is satisfied 606, causing the robot to perform the sequence of actions in the physical environment.

[0103] As one example, processing the representation of the physical environment at the final time step in the sequence of time steps to determine whether the success criterion is satisfied can comprise determining that respective locations, shapes, or configurations of one or more objects in the physical environment are within a tolerance of one or more target locations, shapes, or configurations of the one or more objects.

[0104] FIG. 7 shows examples of sequences (rollouts) of images 702, 704, 706 and images 708 from different viewpoints generated by the described techniques. The sequences of images 702, 704, 706 were obtained by providing two input images (frames) to the simulation system 100 to generate predicted images over a plurality of time steps (i.e., rollouts).

[0105] In the first sequence (“MuJoCo Block”) of images 702, the simulation system 100 was trained using videos of physics-based simulations (generated using the MuJoCo physics simulator) of a block being tossed onto a planar surface with different initial heights above the surface, velocities and angular velocities. The two input images of the first sequence of images 702 show the block falling towards a planar surface and the predicted images for times steps 3-31 show the block rolling on impact with the surface, thereby showing the ability of the described techniques to accurately represent rigid contact between objects in a physical environment.

[0106] In the second sequence (“Deformable Collision”) of images 704, the simulation system 100 was trained using rendered videos of deformable objects and the sequence shows on example of a rollout for a collision between a deformable toy duck and a block that are dropped onto the planar surface.

[0107] In the third sequence (“Kubric MOVi-C”) of images 706, the simulation system 100 was trained using videos that incorporate a range of visual complexities into the physics-based simulations, from Platonic solids to densely -textured scans of real objects and backgrounds. The two input images show three objects that have been thrown ontoa textured surface and the rollout shows the motion of the objects after collision with the surface, thereby testing the ability of the system to represent multi-object interactions in varied visual environments.

[0108] The images 708 from different viewpoints show a deformable torus and deformable block from different camera angles following a rollout over 15 time steps modelling the collision between the torus and block, thereby demonstrating the ability of the system to use the image generation neural network to accurately model different viewing directions.

[0109] FIG. 8 shows graphs of video quality, in terms of Peak Signal to Noise ratio, PSNR, (vertical axes) as a function of rollout duration (horizontal axes) for sequences of images generated following training of the system using each of the MuJoCo Block, Deformable Collision and Kubric MOVi-C datasets mentioned in connection with FIG. 7 (of which by the sequences of images 702, 704, 706 in FIG. 7 are representative). The results 802A, 804A and 806A shown for the techniques described in the present specification indicate generally superior video quality compared to: (i) the results 802B, 804B, 806B obtained using a compositional 2D object-based video prediction model (“SlotF ormer”, arXiv:2210.05861) that uses a slot-based image encoder and a transformer; (ii) the results 802C, 804C, 806C obtained using a system similar to NeRF- dy (arXiv:2107.04004), which uses a global latent vector per scene and models the particle dynamics using a multilayer perceptron, rather than a GNN; and (iii) the results 802D, 804D, 806D obtained using a version of the simulation system 100, in which the encoder neural network 106 and image generation neural network 118 are trained separately from the representation update subsystem 110 (GNN 116), i.e., a version of the system that is not trained end-to-end.

[0110] In this specification, the term "configured" is used in relation to computing systems and environments, as well as computer program components. A computing system or environment is considered "configured" to perform specific operations or actions when it possesses the necessary software, firmware, hardware, or a combination thereof, enabling it to carry out those operations or actions during operation. For instance, configuring a system might involve installing a software library with specific algorithms, updating firmware with new instructions for handling data, or adding a hardware component for enhanced processing capabilities. Similarly, one or more computer programs are "configured" to perform particular operations or actions whenthey contain instructions that, upon execution by a computing device or hardware, cause the device to perform those intended operations or actions.

[0111] The embodiments and functional operations described in this specification can be implemented in various forms, including digital electronic circuitry, software, firmware, computer hardware (encompassing the disclosed structures and their structural equivalents), or any combination thereof. The subject matter can be realized as one or more computer programs, essentially modules of computer program instructions encoded on a tangible non-transitory storage medium for execution by or to control the operation of a computing device or hardware. The storage medium can be a storage device such as a hard drive or solid-state drive (SSD), a storage medium, a random or serial access memory device, or a combination of these. Additionally or alternatively, the program instructions can be encoded on a transmitted signal, such as a machine-generated electrical, optical, or electromagnetic signal, designed to carry information for transmission to a receiving device or system for execution by a computing device or hardware. Furthermore, implementations may leverage emerging technologies like quantum computing or neuromorphic computing for specific applications, and may be deployed in distributed or cloud-based environments where components reside on different machines or within a cloud infrastructure.

[0112] The term "computing device or hardware" refers to the physical components involved in data processing and encompasses all types of devices and machines used for this purpose. Examples include processors or processing units, computers, multiple processors or computers working together, graphics processing units (GPUs), tensor processing units (TPUs), and specialized processing hardware such as field- programmable gate arrays (FPGAs) or application-specific integrated circuits (ASICs). In addition to hardware, a computing device or hardware may also include code that creates an execution environment for computer programs. This code can take the form of processor firmware, a protocol stack, a database management system, an operating system, or a combination of these elements. Embodiments may particularly benefit from utilizing the parallel processing capabilities of GPUs, in a General-Purpose computing on Graphics Processing Units (GPGPU) context, where code specifically designed for GPU execution, often called kernels or shaders, is employed. Similarly, TPUs excel at running optimized tensor operations crucial for many machine learning algorithms. By leveraging these accelerators and their specialized programming models, the system can achieve significant speedups and efficiency gains for tasks involvingartificial intelligence and machine learning, particularly in areas such as computer vision, natural language processing, and robotics.

[0113] A computer program, also referred to as software, an application, a module, a script, code, or simply a program, can be written in any programming language, including compiled or interpreted languages, and declarative or procedural languages. It can be deployed in various forms, such as a standalone program, a module, a component, a subroutine, or any other unit suitable for use within a computing environment. A program may or may not correspond to a single file in a file system and can be stored in various ways. This includes being embedded within a file containing other programs or data (e.g., scripts within a markup language document), residing in a dedicated file, or distributed across multiple coordinated files (e.g., files storing modules, subprograms, or code segments). A computer program can be executed on a single computer or across multiple computers, whether located at a single site or distributed across multiple sites and interconnected through a data communication network. The specific implementation of the computer programs may involve a combination of traditional programming languages and specialized languages or libraries designed for GPGPU programming or TPU utilization, depending on the chosen hardware platform and desired performance characteristics.

[0114] In this specification, the term "engine" broadly refers to a software-based system, subsystem, or process designed to perform one or more specific functions. An engine is typically implemented as one or more software modules or components installed on one or more computers, which can be located at a single site or distributed across multiple locations. In some instances, one or more dedicated computers may be used for a particular engine, while in other cases, multiple engines may operate concurrently on the same one or more computers. Examples of engine functions within the context of Al and machine learning could include data pre-processing and cleaning, feature engineering and extraction, model training and optimization, inference and prediction generation, and post-processing of results. The specific design and implementation of engines will depend on the overall architecture and the distribution of computational tasks across various hardware components, including CPUs, GPUs, TPUs, and other specialized processors.

[0115] The processes and logic flows described in this specification can be executed by one or more programmable computers running one or more computer programs to perform functions by operating on input data and generating output. Additionally,graphics processing units (GPUs) and tensor processing units (TPUs) can be utilized to enable concurrent execution of aspects of these processes and logic flows, significantly accelerating performance. This approach offers significant advantages for computationally intensive tasks often found in Al and machine learning applications, such as matrix multiplications, convolutions, and other operations that exhibit a high degree of parallelism. By leveraging the parallel processing capabilities of GPUs and TPUs, significant speedups and efficiency gains compared to relying solely on CPUs can be achieved. Alternatively or in combination with programmable computers and specialized processors, these processes and logic flows can also be implemented using specialized processing hardware, such as field-programmable gate arrays (FPGAs) or application-specific integrated circuits (ASICs), for even greater performance or energy efficiency in specific use cases.

[0116] Computers capable of executing a computer program can be based on general- purpose microprocessors, special-purpose microprocessors, or a combination of both. They can also utilize any other type of central processing unit (CPU). Additionally, graphics processing units (GPUs), tensor processing units (TPUs), and other machine learning accelerators can be employed to enhance performance, particularly for tasks involving artificial intelligence and machine learning. These accelerators often work in conjunction with CPUs, handling specialized computations while the CPU manages overall system operations and other tasks. Typically, a CPU receives instructions and data from read-only memory (ROM), random access memory (RAM), or both. The essential elements of a computer include a CPU for executing instructions and one or more memory devices for storing instructions and data. The specific configuration of processing units and memory will depend on factors like the complexity of the Al model, the volume of data being processed, and the desired performance and latency requirements. Embodiments can be implemented on a wide range of computing platforms, from small embedded devices with limited resources to large-scale data center systems with high-performance computing capabilities. The system may include storage devices like hard drives, SSDs, or flash memory for persistent data storage.

[0117] Computer-readable media suitable for storing computer program instructions and data encompass all forms of non-volatile memory, media, and memory devices. Examples include semiconductor memory devices such as read-only memory (ROM), solid-state drives (SSDs), and flash memory devices; hard disk drives (HDDs); optical media; and optical discs such as CDs, DVDs, and Blu-ray discs. The specific type ofcomputer-readable media used will depend on factors such as the size of the data, access speed requirements, cost considerations, and the desired level of portability or permanence.

[0118] To facilitate user interaction, embodiments of the subject matter described in this specification can be implemented on a computing device equipped with a display device, such as a liquid crystal display (LCD) or an organic light-emitting diode (OLED) display, for presenting information to the user. Input can be provided by the user through various means, including a keyboard), touchscreens, voice commands, gesture recognition, or other input modalities depending on the specific device and application. Additional input methods can include acoustic, speech, or tactile input, while feedback to the user can take the form of visual, auditory, or tactile feedback. Furthermore, computers can interact with users by exchanging documents with a user's device or application. This can involve sending web content or data in response to requests or sending and receiving text messages or other forms of messages through mobile devices or messaging platforms. The selection of input and output modalities will depend on the specific application and the desired form of user interaction.

[0119] Machine learning models can be implemented and deployed using machine learning frameworks, such as TensorFlow or JAX. These frameworks offer comprehensive tools and libraries that facilitate the development, training, and deployment of machine learning models.

[0120] Embodiments of the subject matter described in this specification can be implemented within a computing system comprising one or more components, depending on the specific application and requirements. These may include a back-end component, such as a back-end server or cloud-based infrastructure; an optional middleware component, such as a middleware server or application programming interface (API), to facilitate communication and data exchange; and a front-end component, such as a client device with a user interface, a web browser, or an app, through which a user can interact with the implemented subject matter. For instance, the described functionality could be implemented solely on a client device (e.g., for on- device machine learning) or deployed as a combination of front-end and back-end components for more complex applications. These components, when present, can be interconnected using any form or medium of digital data communication, such as a communication network like a local area network (LAN) or a wide area network (WAN) including the Internet. The specific system architecture and choice ofcomponents will depend on factors such as the scale of the application, the need for real-time processing, data security requirements, and the desired user experience.

[0121] The computing system can include clients and servers that may be geographically separated and interact through a communication network. The specific type of network, such as a local area network (LAN), a wide area network (WAN), or the Internet, will depend on the reach and scale of the application. The client-server relationship is established through computer programs running on the respective computers and designed to communicate with each other using appropriate protocols. These protocols may include HTTP, TCP / IP, or other specialized protocols depending on the nature of the data being exchanged and the security requirements of the system. In certain embodiments, a server transmits data or instructions to a user's device, such as a computer, smartphone, or tablet, acting as a client. The client device can then process the received information, display results to the user, and potentially send data or feedback back to the server for further processing or storage. This allows for dynamic interactions between the user and the system, enabling a wide range of applications and functionalities.

[0122] While this specification contains many specific implementation details, these should not be construed as limitations on the scope of any invention or on the scope of what may be claimed, but rather as descriptions of features that may be specific to particular embodiments of particular inventions. Certain features that are described in this specification in the context of separate embodiments can also be implemented in combination in a single embodiment. Conversely, various features that are described in the context of a single embodiment can also be implemented in multiple embodiments separately or in any suitable subcombination. Moreover, although features may be described above as acting in certain combinations and even initially be claimed as such, one or more features from a claimed combination can in some cases be excised from the combination, and the claimed combination may be directed to a subcombination or variation of a subcombination.

[0123] Similarly, while operations are depicted in the drawings and recited in the claims in a particular order, this should not be understood as requiring that such operations be performed in the particular order shown or in sequential order, or that all illustrated operations be performed, to achieve desirable results. In certain circumstances, multitasking and parallel processing may be advantageous. Moreover, the separation of various system modules and components in the embodiments described above shouldnot be understood as requiring such separation in all embodiments, and it should be understood that the described program components and systems can generally be integrated together in a single software product or packaged into multiple software products.

[0124] Particular embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions recited in the claims can be performed in a different order and still achieve desirable results. As one example, the processes depicted in the accompanying figures do not necessarily require the particular order shown, or sequential order, to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous.

Claims

CLAIMS:

1. A method performed by one or more data processing apparatus for determining a representation of a physical environment at a new time step before or after a current time step, the method comprising: obtaining data defining respective representations of the physical environment at the current time step and one or more other time steps, the representation at a given time step comprising a collection of particle nodes, each particle node comprising data defining a corresponding spatial location in the physical environment and a set of features characterizing the spatial location at the time step; constructing a graph comprising the particle nodes of the representations of the physical environment at the current time step and the one or more other time steps; processing the graph using a graph neural network to update the set of features of each of the particle nodes of the representation of the physical environment at the current time step based on the sets of features of a respective subset of the particle nodes of the representations of the physical environment at the one or more other time steps, wherein each subset is determined based on a spatial proximity of the spatial locations of the particle nodes of the representations of the physical environment at the one or more other time steps to the spatial location of the corresponding particle node of the representation of the physical environment at the current time step; processing the updated sets of features of the particle nodes of the representation of the physical environment at the current time step to determine a representation of the physical environment at the new time step.

2. The method of claim 1, wherein processing the updated sets of features of the particle nodes of the representation of the physical environment at the current time step to determine the representation of the physical environment at the new time step comprises, for each of the particle nodes in representation of the physical environment at the current time step: generating a corresponding particle node for the representation of the physical environment at the new time step, wherein the spatial location of the corresponding particle node is offset from the spatial location of the particle node by an amount determined from the updated sets of features of the particle nodes of the representationof the physical environment at the current time step, and the set of features of the corresponding particle node are determined from the updated sets of features of the particle nodes of the representation of the physical environment at the current time step.

3. The method of claim 1 or 2, wherein constructing the graph comprising the particle nodes of the representation of the physical environment at the current time step and the one or more other time steps comprises: generating a plurality of auxiliary nodes, wherein each auxiliary node comprises data defining a corresponding spatial location in the physical environment and a set of features characterizing the spatial location, and the graph comprises edges: (i) from the particle nodes of the representations of the physical environment at the current time step and the one or more other time steps to the auxiliary nodes; (ii) between pairs of the auxiliary nodes; and (iii) from the auxiliary nodes to the particle nodes of the representation of the physical environment at the current time step.

4. The method of claim 3, wherein processing the graph using the graph neural network to update the set of features of each of the particle nodes of the representation of the physical environment at the current time step comprises: for each of the auxiliary nodes: updating the set of features of the auxiliary node based on the set of features of each of the particle nodes in the graph having an edge from the particle node to the auxiliary node; and for each of one or more update iterations: updating the set of features of each of the auxiliary nodes based on the sets of features of each of the auxiliary nodes in the graph having an edge to the auxiliary node; and for each of the particle nodes of the representation of the physical environment at the current time step: updating the set of features of the particle node based on the sets of features of each of the auxiliary nodes in the graph having an edge to the particle node.

5. The method of claim 4, wherein updating the set of features of the auxiliary node based on the sets of features of each of the particle nodes in the graph having an edge from the particle node to the auxiliary node comprises: for each of the particle nodes in the graph having an edge from the particle node to the auxiliary node: generating a set of features associated with the edge from the particle node to the auxiliary node; and updating the set of features associated with the edge based on the set of features of the auxiliary node and the set of features of the particle node; and updating the set of features of the auxiliary node based on the sets of features associated with each of the edges from a particle node in the graph to the auxiliary node.

6. The method of claim 4 or 5, further comprising for each of the one or more update iterations: for each of the edges in the graph between pairs of auxiliary nodes: updating a set of features associated with the edge based on the set of features of each of the auxiliary nodes in the pair; and updating the set of features of each of the auxiliary nodes in the pair based on the sets of features of the edge.

7. The method of any one of claims 4 to 6, wherein updating the set of features of the particle node based on the sets of features of each of the auxiliary nodes in the graph having an edge to the particle node comprises: for each of the auxiliary nodes having an edge to the particle node: updating a respective set of features associated with the edge from the auxiliary node to the particle node based on the set of features of the particle node and the set of features of the auxiliary node; and updating the set of features of the particle node based on the sets of features associated with each of the edges from an auxiliary node to the particle node.

8. The method of any one of claims 4-7, wherein the number of auxiliary nodes is less than: (i) a number of particle nodes of the representation of the physical environment at the current time step; or (ii) a number of particle nodes of therepresentation of the physical environment at the one or more other time steps; or (iii) both.

9. The method of claim 8, wherein the spatial locations of the auxiliary nodes are chosen by sampling the spatial locations of the particle nodes of the representation of the physical environment at the current time step.

10. The method of claim 9, wherein the sampling is random or pseudorandom.

11. The method of any one of claims 3-10, wherein for each of the edges from a particle node to an auxiliary node, a distance between the spatial location of the particle node and spatial location of the auxiliary node is less than a first predetermined threshold.

12. The method of any one of claims 3-11, wherein for each of the edges between pairs of auxiliary nodes, a distance between the respective spatial locations of the auxiliary nodes is less than a second predetermined threshold.

13. The method of any one of claims 3-12, wherein each of the particle nodes has an edge to no more than a pre-defined number of auxiliary nodes.

14. The method of claim 13, wherein the pre-defined number of auxiliary nodes is two.

15. The method of any one of the preceding claims, wherein obtaining data defining respective representations of the physical environment at the current time step and the one or more other time steps comprises: obtaining a sequence of images, the sequence comprising images of the physical environment at the current time step and the one or more other time steps, the image for a given time step comprising a plurality of pixels, each pixel having an associated set of pixel values and corresponding to one of the particle nodes of the representation of the physical environment at the time step; for each of the images in the sequence:processing the pixels of the image to obtain data defining the spatial location in the physical environment for each of the particle nodes of the representation of the physical environment at the time step corresponding to the image; and processing the sets of pixel values associated with the pixels in the image using an encoder neural network to determine the set of features for each of the particle nodes of the representation of the physical environment at the time step corresponding to the image.

16. The method of claim 15, wherein processing the pixels of the image to obtain data defining the spatial location in the physical environment for each of the particle nodes of the representation of the physical environment at the time step corresponding to the image comprises: obtaining a depth value for each of the pixels, the depth value corresponding to a distance into the physical environment along a direction perpendicular to the plane of the image; and using the depth values and a transformation matrix associated with the image to map each of the pixels to a spatial location in the physical environment.

17. The method of claim 15 or 16, wherein each set of pixel values comprises a depth value for the pixel.

18. The method of any one of claims 15-17, wherein the set of pixel values comprises one or more values indicative of a color and / or intensity of the pixel.

19. The method of any one of claims 15-18, wherein the sequence of images comprises a respective plurality of images at each of the time steps, each plurality of images comprising images of the physical environment corresponding to: (i) different viewing directions of the physical environment; or (ii) different spatial locations in the physical environment; or (iii) both.

20. The method of any one of the preceding claims, further comprising generating an image of the physical environment at the new time step based on the representation of the physical environment at the new time step.

21. The method of claim 20, wherein generating an image of the physical environment at the new time step based on the representation of the physical environment at the new time step comprises: processing the representation of the physical environment at the new time step using an image generation neural network to obtain a set of pixel values for each of the pixels in the image.

22. The method of claim 21, wherein the image generation neural network is a scene representation neural network and generating an image of the physical environment at the new time step based on the representation of the physical environment at the new time step comprises: processing the representation of the physical environment at the new time step to determine scene parameters of the scene representation neural network for the physical environment at the new time step; and generating the image using the scene representation neural network configured with the parameters for the physical environment at the new time step.

23. The method of claim 22, wherein the scene parameters encode a representation of a scene comprising the physical environment at the new time step over a three dimensional spatial volume, and generating the image using the scene representation neural network configured with the scene parameters comprises: for each pixel in the image: determining a plurality of spatial locations along a ray into the scene from the pixel; for each of the spatial locations, processing a representation of a viewing direction from the pixel, and a representation of the spatial location, using the scene representation neural network, to generate a scene representation output, wherein the scene representation output defines the light level emitted from the spatial location along the viewing direction and the opacity or color at the spatial location; and combining, for the spatial locations along the ray, the light level emitted from the spatial location along the viewing direction and the opacity or color atthe spatial location, to determine the one or more pixel values for the pixel in the image.

24. The method of any one of claims 21-23 when dependent on claim 15, wherein the sequence of images comprises a target image of the physical environment at the new time step and the encoder neural network, the graph neural network and the image generation neural network are trained end-to-end using a loss function that compares sets of pixel values of the generated image with corresponding sets of pixel values of the target image.

25. The method of claim 24, wherein the loss function comprises a mean squared error between values of the sets of pixel values in the generated image and the corresponding values of the sets of pixel values in the target image.

26. The method of any one of the preceding claims, further comprising editing one or more of the representations of the physical system by performing one or more of: (i) deleting one or more of the particle nodes; (ii) modifying the spatial location or the set of features of one or more of the particle nodes; and (iii) introducing one or more additional particle nodes into the representation.

27. The method of any one of the preceding claims, wherein: (i) the new time step is after the current time step and the one or more other time steps are before the current time step; or (ii) the new time step is before the current time step and the one or more other time steps are after the current time step; or (iii) the new time step is after the current time step and the one or more other time steps are after the new time step; or (iv) the new time step is before the current time step and the one or more other time steps are before the new time step.

28. The method of any one of the preceding claims, wherein the physical environment comprises a real-world physical environment, and the method further comprises: using the representation of the physical environment at the new time step to determine one or more actions to be performed by an electronic or mechanical agent interacting with the physical environment to perform a task.

29. A method of controlling a robot interacting with a physical environment, the method comprising: repeatedly performing the method of any one of claims 1-27 to determine representations of a physical environment for a sequence of time steps, wherein the representations characterize the physical environments that would be obtained when the robot performs a sequence of actions over the sequence of time steps; processing the representation of the physical environment at a final time step in the sequence of time steps to determine whether a success criterion is satisfied; and in response to determining that the success criterion is satisfied, causing the robot to perform the sequence of actions in the physical environment.

30. The method of claim 29, wherein processing the representation of the physical environment at the final time step in the sequence of time steps to determine whether the success criterion is satisfied comprises: determining that respective locations, shapes, or configurations of one or more objects in the physical environment are within a tolerance of one or more target locations, shapes, or configurations of the one or more objects.

31. One or more non-transitory computer storage media storing instructions that when executed by one or more computers cause the one or more computers to perform operations of the respective method of any one of claims 1-30.

32. A system comprising: one or more computers; and one or more storage devices communicatively coupled to the one or more computers, wherein the one or more storage devices store instructions that, when executed by the one or more computers, cause the one or more computers to perform operations of the respective method of any one of claims 1-30.

33. One or more non-transitory computer storage media storing a representation of a physical environment at a new time step before or after a current time step obtained using the method of any one of claims 1-30.

Images