Scalable attention-based point cloud modeling

Abstract

A method performed by one or more data processing apparatus for updating a representation of an environment comprising a plurality of point embeddings. Each point embedding is derived from features characterizing a corresponding spatial location in the environment. The method comprises: updating each of the point embeddings using a respective proper subset of the point embeddings, each subset being selected based on the spatial location of the point embedding; and updating the point embeddings by applying a linear self-attention mechanism over the point embeddings. Applying the linear self-attention mechanism over the point embeddings comprises: using the point embeddings to determine a query matrix, a key matrix and a value matrix; and updating the point embeddings using a matrix product of the query matrix, the key matrix, and the value matrix.

Description

ATTORNEY DOCKET NO.45288-0439WO1 SCALABLE ATTENTION-BASED POINT CLOUD MODELING CROSS-REFERENCE TO RELATED APPLICATIONS

[0001] This application claims priority to U.S. Provisional Application No. 63 / 657,095, filed on June 6, 2024. 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, and a method, implemented as computer programs on one or more computers in one or more locations for updating a point cloud representation of an environment, i.e., a collection of point embeddings derived from features characterizing corresponding spatial locations in the environment. The method may be implemented using a deep neural network, for example, e.g. a point cloud transformer.

[0006] The updated representations of the environment determined using the system and method may be used for a wide variety of applications, including scene classification, object recognition or localization, action or gesture recognition, semantic segregation, image generation, control of a mechanical agent or robot, and so on.

[0007] As one example, the updated representations can be used predict or simulate future properties of the environment, such as future configurations of the environment. For example, the motions of interacting and / or deformable objects in the environment may beATTORNEY DOCKET NO.45288-0439WO1 predicted. The updated representations can be used to control a mechanical or electronic agent to perform a specified task, e.g., by processing the updated representations, e.g. using a policy neural network, to determine what action(s) the agent should perform to accomplish the task.

[0008] According to a first aspect there is provided a method performed by one or more data processing apparatus for updating a representation of an environment. The representation comprises a plurality of point embeddings, each point embedding being derived from features characterizing a corresponding spatial location in the environment. The method comprises updating each of the point embeddings using a respective proper subset of the point embeddings, each subset being selected based on the spatial location of the point embedding. The method can further comprise updating the point embeddings by applying a linear self-attention mechanism over the point embeddings. The linear self- attention mechanism can, for example, have a time and / or space complexity that is linear in the number of point embeddings (rather than e.g., a time and / or space complexity that is quadratic in the number of point embeddings, as in a conventional self-attention mechanism).

[0009] In some implementations, applying the linear self-attention mechanism over the point embeddings can comprise: using the point embeddings to determine a query matrix (Q’), a key matrix (K’) and a value matrix (V); and updating the point embeddings using a matrix product of the query matrix, the key matrix, and the value matrix, e.g. a matrix product of the form, Q’((K’)TV).

[0010] Although the query matrix and the key matrix can be considered as together defining an attention matrix, it is not necessary for elements of this attention matrix to be computed as the method is being performed. In particular, the method can be performed without explicitly constructing and storing the attention matrix, in contrast to conventional self-attention mechanisms. For example, in some implementations, updating the point embeddings using a matrix product of the query matrix (Q’), the key matrix (K’), and the value matrix (V) can comprise: multiplying the key matrix and the value matrix to obtain an key-value matrix ((K’)TV), multiplying the query matrix and the key-value matrix to obtain a query-key-value matrix Q((K’)TV); and updating the point embeddings using the query-key-value matrix. In some implementations, the queryATTORNEY DOCKET NO.45288-0439WO1 matrix and the key matrix have a dimension that is less than a number of point embeddings in the plurality of the point embeddings.

[0011] The use of a linear self-attention mechanism can allow larger point clouds to be processed and / or fewer computational resources to be required than for other approaches.

[0012] In general, the query matrix, key matrix and value matrix are generated by applying a respective learned transformation to the point embeddings, e.g., by multiplying a matrix comprising the point embeddings by a respective weight matrix comprising elements that have been optimized during training of the linear self-attention mechanism.

[0013] In some implementations, each of the query matrix and the key matrix has a dimension (^^^) that is smaller than the number of point embeddings (N). For example, the query matrix and the key matrix can each have a respective first dimension equal to the number of numerical values in each embedding (in general, each point embedding is an ordered collection of numerical values) and a second dimension (m) that is less than the number of point embeddings (N). The second dimension (m) is a hyper-parameter that is typically chosen to be much less than the number of point embeddings (N).

[0014] In some implementations, prior to determining the matrix product, a row-wise transformation (^) can be applied to the query matrix and the key matrix. The row-wise transformation can, for example, comprise an element-wise non-linear transformation, e.g. a Rectified Linear Unit (ReLU) or another activation function.

[0015] Thus, in some implementations, the query matrix and the key matrix each comprise a respective plurality of transformed rows (or equivalently, columns), each transformed row of the modified query matrix being obtained by applying a transformation (^) to a corresponding one of the rows of the query matrix and each transformed row of the (original) key matrix being obtained by applying a transformation (e.g., the same transformation, e.g., ^) to a corresponding one of the rows of the (original) key matrix. The transformation can allow the information associated with the query matrix and key matrix to be encoded more efficiently.

[0016] In some implementations, applying the transformation (^) to a row comprises applying a randomized mapping to the row, e.g., for each of a plurality of random functions, applying the random function to the row to generate a corresponding element of the transformed row. Applying each of the random functions to the row can, for example, comprise applying a corresponding random projection to the row, e.g., a randomATTORNEY DOCKET NO.45288-0439WO1 projection obtained by sampling values from a Gaussian distribution. Each of the random functions can be an exponential function of the corresponding random projection of the row, for example.

[0017] One example of a linear self-attention mechanism that can be used is the performer self-attention mechanism, e.g., as described in Choromanski et al. “Rethinking attention with performers” in 9thInternational Conference on Learning Representations, ICLR 2021.

[0018] In some implementations, updating each of the point embeddings (which may also be referred to as “feature vectors”) using a respective proper subset of the point embeddings can be implemented using a neural network, e.g. a Graph Neural Network (GNN). For example, the neural network may comprise one or more neighbor-attender layers configured to receive the point embeddings as input and to generate updated point embeddings as output. The neural network may provide a geometric structural inductive bias to the point embeddings, e.g. such that the point embeddings tend to attend to neighboring point embeddings more strongly than to point embeddings that are further away. That is, the neural network may modulate each point embedding using the other point embeddings based on the proximity of the spatial locations of the other point embeddings to the spatial location of the point embedding.

[0019] Thus, the method may efficiently model global point-to-point attention using the linear self-attention mechanism, whilst modulating the global attention with a local update mechanism.

[0020] In some implementations, each proper subset of the point embeddings is a group obtained by grouping the point embeddings by spatial location, e.g., using a clustering algorithm, such as a k nearest neighbors algorithm. The method may further comprise determining a corresponding group embedding for each group by combining the point embeddings in the group.

[0021] Updating each of the point embeddings can comprise updating the point embeddings using one or more of the group embeddings, e.g. by selecting a respective one of the groups based on the spatial location of the point embedding; and updating the point embedding using the group embedding of the selected group. Grouping the point embeddings by spatial location can, for example, comprise, for each group: sampling a spatial location from the point embeddings to use as an anchor point for the group; andATTORNEY DOCKET NO.45288-0439WO1 selecting the point embeddings of the group based on proximities of the spatial locations of the point embeddings to the anchor point. The number of anchor points sampled is generally less than the number of point embeddings, e.g. a quarter of the number of point embeddings. Selecting the point embeddings for the group can comprise selecting a predetermined number (k) of point embeddings having corresponding spatial locations that are nearest the anchor point, i.e. selecting the k nearest neighbors to the anchor point.

[0022] In some implementations, the anchor point of the selected group is the anchor point nearest the spatial location of point embedding. That is, the group embedding used to update the point embedding may be selected by finding the anchor point nearest the spatial location of the point embedding.

[0023] In some implementations, combining the point embeddings in the group can comprise determining a weighted sum of the point embeddings in the group. For example, a learnable multi-layer perceptron (MLP) can be used to compute the weighted sum of the point embeddings.

[0024] In some implementations, the group embedding comprises a weighted sum of embeddings of the spatial locations of the point embeddings in the group. The embedding of the spatial location of each point embedding may comprise the spatial location of the anchor point of the group and / or an embedding of the spatial location of the point embedding relative to the anchor point of the group. For example, each of the point embeddings in the group may be concatenated with a relative position feature vectordefined as (xi, xj, xi-xj, ^^^ − ^^^^), where xi and xj are respective embeddings (e.g.vectors) of the spatialthe anchor point and point embeddings, to give an embedding for an edge between the anchor point and the point embedding. The edge embeddings are then combined to determine the group embedding. Thus, in some implementations, the group embedding may comprise a weighted sum of embeddings of the spatial locations of the point embeddings in the group.

[0025] In some implementations, updating the point embedding using the group embedding of the selected group can comprise combining the point embedding with the group embedding, e.g., using a neural network (such as an MLP) trained to process an input comprising the point embedding and the group embedding to obtain an updated point embedding as output.ATTORNEY DOCKET NO.45288-0439WO1

[0026] In some implementations, the point embeddings comprise point embeddings that are each derived from features characterizing a corresponding spatial location in the environment at different respective times. Thus, the updated representation of the environment may be used to predict the state of the environment at a specified time after the different respective times. The method may be performed iteratively, for example, by combining or merging the point embeddings at a current time with the point embeddings obtained from a previous time. The method may additionally or alternatively be used to infer a state of the environment before the different respective times, i.e., for retrodiction, or between the different times, i.e., for interpolation.

[0027] In some implementations, the environment can be or comprise a real-world environment or a virtual (simulated) environment. Obtaining the representation of the environment can, for example, comprise processing one or more images of a real-world environment to determine the point embeddings. For example, each image can comprise a plurality of pixels, each pixel having an associated set comprising one or more pixel values. Processing the image to determine the point embeddings can comprise processing the pixel values to determine, for each of the point embeddings, the features characterizing the corresponding spatial location in the environment. For example, the pixel values may be processed using an encoder, such as a convolutional neural network (e.g. a U-Net), a vision transformer, or another image processing neural network. Each set of pixel values may, for example, comprise a respective depth value for the corresponding pixel, such that depth information can be included in the point embeddings. Additionally or alternatively, the set of pixel values can comprise one or more pixel values indicative of a color and / or intensity of the pixel. Thus, 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 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.

[0028] In some implementations, the updated representation can used to perform one or more machine learning tasks, such as one or more of:ATTORNEY DOCKET NO.45288-0439WO1 (i) a classification task in which the updated representation is used to determine an output classifying the environment into one or more categories; (ii) an object recognition or localization task in which the updated representation is used to determine an output indicative of whether one or more objects is present in the environment or an output defining coordinates of a respective region for one or more objects in the environment; (iii) an action or gesture recognition task in which the updated representation is used to determine an output indicative of whether one or more actions or gestures is being performed in the environment; (iv) a semantic segmentation task in which the updated representation is used to determine an output for assigning spatial locations in the environment to respective segmentation categories; (v) a keypoint detection task in which the updated representation is used to determine an output comprising coordinates of one or more keypoints in the environment; and (vi) an environment similarity task in which the updated representation is used to determine an output indicative of a similarity of the environment to one or more predetermined environments. Image Generation

[0029] In some implementations, the method can further comprise processing the updated representation using an image generation neural network to generate an image, e.g., an image of the environment at a current time, or a past or future time. As an example, the self-attention mechanism can be implemented by linear self-attention neural network layers that are trained with the image generation neural network using a loss function that compares sets of pixel values of the generated image of the environment with corresponding sets of pixel values of a target image of the environment.

[0030] In some examples, the image generation neural network can be a scene representation neural network and processing the updated representation can comprise: processing the updated representation to determine scene parameters of the scene representation neural network; and generating the image using the scene representation neural network configured with the parameters. In some cases, the scene parametersATTORNEY DOCKET NO.45288-0439WO1 encode a representation of a scene comprising the environment over a three-dimensional spatial volume. Generating the image using the scene representation neural network configured with the scene parameters can then comprise: 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 at the spatial location, to determine the one or more pixel values for the pixel in the image. Agent Control

[0031] In some implementations, the updated representation is used to determine one or more actions to be performed by an agent interacting with the environment to perform a specified task. For example, the environment can be a real-world environment and the agent can comprise (i) a mechanical agent or robot interacting with the real-world environment to perform the task, or (ii) an electronic agent controlling items of equipment in the real-world environment to perform the task. The updated representation can, for example, be processed to determine instructions or control signals for controlling the mechanical or electronic agent. The method can then comprise providing the instructions or control signals to the mechanical or electronic agent. In some examples, obtaining the representation can comprise processing sensor data obtained by one or more sensors in the real-world environment (e.g., sensors of the mechanical agent or robot) and processing the sensor data to determine, for each of the point embeddings, the features characterizing the corresponding spatial location in the environment. In some implementations, the updated representation can be provided as an input to a policy machine learning model (e.g., a policy neural network) that processes the updated representation in accordance with parameters of the policy machine learning model to determine an output comprising an action to be performed by the mechanical or electronic agent.ATTORNEY DOCKET NO.45288-0439WO1

[0032] In some examples, generating the updated representation of the environment is conditioned on an embedding of one or more actions that can be performed by the agent interacting with the environment. For example, the representation of the environment can be updated based on the embedding of the one or more actions such that the updated representation characterizes the environment that is predicted to be obtained if the agent performs the one or more actions. The updated representation can then be processed to determine an expected return for the one or more actions, e.g., using an action-value machine learning model (e.g. a Q-function neural network). Further details of how the updated representation can be conditioned according to an action are provided in the third aspect below.

[0033] In some examples, the updated representations of the environment may be used to generate one or more images showing how the environment should change when the task is performed correctly. The agent may then process the one or more images to determine what actions to take to perform the task, e.g., by selecting actions to minimize differences between images of the environment obtained during performance of the task by the agent with corresponding generated images showing how the environment should change when the task is performed correctly. Thus, in some implementations, the representation or the updated representation can be used to determine if the task has been completed or will be completed by the agent performing the one or more actions determined using the updated representation.

[0034] In some implementations, the representation characterizes the environment at a current time and the updated representation characterizes the environment at another time that is before or after the current time. In some examples, generating the updated representation of the environment is conditioned on each of one or more representations characterizing the environment at a respective time before or after the current time.

[0035] In cases where the environment is a virtual environment, the representations of the environment may be used to train the agent to perform the task in the virtual environment, e.g., by generating images of the virtual 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 environment. Such a procedure may be referred to as “sim-to- real” training.ATTORNEY DOCKET NO.45288-0439WO1

[0036] In a second aspect, there is provided a method of training a machine learning model. The method comprises: obtaining training examples that each comprise (i) one or more representations of an environment and (ii) a target output; and for each of the training examples: using the method of described above for the first aspect to update the one or more representations of the environment of the training example; and using the updated one or more representations and the target output of the training example to train the machine learning model. For example, the machine learning model can be an action selection policy for an agent interacting with the environment.

[0037] In some implementations, the systems described in this specification can be trained using training data comprising “ground truth” representations of environments (e.g., representations obtained from video data) at a number of time steps during which a task is being performed. The system can then learn to predict how the environment evolves as the task is being performed by generating representations of the environment that match or approximate the corresponding ground truth representations at each of the time steps. Thus, the system(s) may be trained by supervised learning. As used in this specification, the term “learned” means that an operation or a value has been adjusted during the training of the system.

[0038] In general, the linear self-attention mechanism can be implemented by one or more self-attention neural network layers and incorporated into a neural network. The neural network can also include other types of layers interspersed among the self-attention layers, e.g., feed-forward layers, e.g., fully-connected layers, other self-attention layers, convolutional layers, and so on. For example, during the updating of the representation of the environment, the self-attention layer can receive a layer input derived from the point embeddings and processes the layer input to generate a layer output. For example, the layer input can comprise a collection of the point embeddings, or embeddings derived from the point embeddings. The point embeddings can, for example, be generated by an encoder neural network that is trained together (i.e., jointly) with or separately from the linear self-attention mechanism.

[0039] In some implementations, the linear self-attention mechanism can comprise a multi-headed self-attention mechanism, e.g., such that there are multiple self-attention sublayers within the self-attention layer and the self-attention layer generates a final layerATTORNEY DOCKET NO.45288-0439WO1 output by combining the outputs generated by the sub-layers, e.g., by concatenating the sub-layer outputs and, optionally, applying a learned transformation, e.g., a linear transformation, to the concatenation. Use of a multi-headed self-attention mechanism may facilitate parallel processing of the point embeddings, e.g., by having the computations performed by the multi-headed self-attention layers performed by different respective processing units, e.g., accelerator units such as graphics processing units (GPUs) or tensor processing units (TPUs). A sequence of multiple self-attention layers can also be used to update the representation of the environment, i.e., a first self-attention layer may update the representation of the environment to obtain an updated representation which is then updated by one or more subsequent self-attention layers to generate a further updated representation, and so on.

[0040] The neural network is generally trained to perform one or more downstream tasks. Training the neural network to perform the one or more downstream tasks may be carried out using supervised learning, e.g., by gradient-based optimization methods, such as stochastic gradient descent. For example, the training may comprise backpropagating gradients of an objective function through the neural network and adjusting parameters of the neural network, including parameters of the self-attention layers, using the gradients. In some implementations, the system and methods described in this specification may be used to generate representations of a 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 generate images that realistically (e.g., smoothly) interpolate between images of the physical environment at the pair of time steps.

[0041] In a third aspect, there is provided a method of determining a representation of an environment at a specified time. The method comprises: obtaining a first representation comprising a first plurality of point embeddings, each point embedding being derived from features characterizing a corresponding spatial location in the environment at a first time; and obtaining a second representation comprising a second plurality of point embeddings, each point embedding being derived from features characterizing a corresponding spatial location in the environment at a second time different from the firstATTORNEY DOCKET NO.45288-0439WO1 time (e.g. the first time may precede the second time). The method further comprises merging the first plurality of point embeddings with the second plurality of point embeddings to obtain a merged representation comprising the merged plurality of point embeddings, and updating the merged representation to determine a representation of the environment at the specified time, wherein the merged plurality of point embeddings (i.e., both the first and second pluralities of point embeddings) are used to update at least the point embeddings of the first plurality of the point embeddings.

[0042] In some implementations, the merged plurality of point embeddings (i.e., both the first and second pluralities of point embeddings) are used to update the point embeddings of both the first plurality of the point embeddings and the second plurality of point embeddings.

[0043] In some implementations, the second plurality of point embeddings comprises point embeddings (“kinematic points” or “kinematic point embeddings”) having spatial locations defining a configuration of one or more objects in the environment at the specified time, i.e., the second time may be the specified time. For example, the second plurality of point embeddings may comprise point embeddings having spatial locations defining the respective configuration of the one or more objects in the environment at each of a plurality of times including the specified time. These point embeddings may each comprise an embedding of a respective one of the plurality of times and / or an embedding identifying a respective one of the one or more objects.

[0044] As one example, the one or more objects may comprise respective parts of a mechanical agent or robot, e.g. joints of the mechanical agent or robot, one or more grippers, one or more fingers of a gripper etc. For example, the (kinematic) point embeddings may define a “skeleton” of the mechanical agent or robot. For example, two of more of the point embeddings can be used to define an orientation of a joint or gripper. The point embeddings may be derived from sensor data from one or more sensors of the mechanical agent or robot and / or control signals provided to the mechanical agent or robot. The inclusion of these point embeddings may allow the point embeddings of the first plurality of point embeddings to be updated to take into account the known or predicted state of the mechanical agent or robot, i.e., the response of the environment to the known or predicted state of the mechanical agent or robot can be predicted. For example, a subset of the first plurality of point embeddings may represent an object in theATTORNEY DOCKET NO.45288-0439WO1 environment that is deformed and / or displaced by the mechanical agent or robot. The point embeddings of the second plurality of point embeddings may be updated to determine how the rest of the environment affects the mechanical agent or robot, e.g. fingers of a gripper of the mechanical agent or robot. Thus, the representation of the environment at the specified time can be used for planning, i.e. selecting actions to be performed by the mechanical agent or robot at each of a plurality of time steps to achieve a specified goal.

[0045] For example, the spatial locations defining the configuration of one or more objects in the environment at the specified time may be predicted based on an action to be performed by the mechanical agent or robot at a time preceding the specified time (e.g. at the first time). Thus, the (kinematic) point embeddings may collectively represent the expected state of the mechanical agent or robot that would be obtained after the mechanical agent or robot has performed the action, i.e., the point embeddings are an embedding of the action. For example, the (kinematic) point embeddings may be generated using a kinematic or physics model of the mechanical agent or robot.

[0046] The number of point embeddings of the second plurality of point embeddings may be less than the number of point embeddings of the first plurality of point embeddings.

[0047] In some implementations, the method can comprise, before merging the first plurality of point embeddings with the second plurality of point embeddings, updating the second representation by applying a self-attention mechanism over the second plurality of point embeddings. The self-attention mechanism may be a conventional “quadratic scaling” self-attention mechanism or a linear self-attention mechanism.

[0048] Updating the merged representation to determine a representation of the environment at the specified time may be performed using the method of the first or third aspect described above or the fourth aspect below, for example.

[0049] In a fourth aspect, there is provided a method performed by one or more data processing apparatus for updating a representation of an environment comprising a plurality of point embeddings. Each point embedding is derived from features characterizing a corresponding spatial location in the environment. The method comprises updating each of the point embeddings using a respective one of a plurality of groups of the point embeddings, the respective one of the plurality of groups of the pointATTORNEY DOCKET NO.45288-0439WO1 embeddings being selected based on the spatial location of the point embedding. Each group comprises a respective proper subset of the point embeddings obtained by grouping the point embeddings by spatial location. The method further comprises determining a corresponding group embedding for each group by combining the point embeddings in the group. Updating each of the point embeddings comprises updating the point embeddings using one or more of the group embeddings, e.g. by selecting a respective one of the groups based on the spatial location of the point embedding; and updating the point embedding using the group embedding of the selected group.

[0050] The method may be implemented by a neural network, e.g. a neural network comprising one or more “neighbor-attender” layers.

[0051] In some implementations, grouping the point embeddings by spatial location can, for example, comprise, for each group: sampling a spatial location from the point embeddings to use as an anchor point for the group; and selecting the point embeddings of the group based on proximities of the spatial locations of the point embeddings to the anchor point. The number of anchor points sampled is generally less than the number of point embeddings, e.g., a quarter of the number of point embeddings. Selecting the point embeddings for the group can comprise selecting a predetermined number (k) of point embeddings having corresponding spatial locations that are nearest the anchor point, i.e. selecting the k nearest neighbors to the anchor point.

[0052] In some implementations, the anchor point of the selected group is the anchor point nearest the spatial location of point embedding. That is, the group embedding used to update the point embedding may be selected by finding the anchor point nearest the spatial location of the point embedding.

[0053] In some implementations, combining the point embeddings in the group can comprise determining a weighted sum of the point embeddings in the group. For example, a learnable multi-layer perceptron (MLP) can be used to compute the weighted sum of the point embeddings.

[0054] In some implementations, the group embedding comprises a weighted sum of embeddings of the spatial locations of the point embeddings in the group. The embedding of the spatial location of each point embedding may comprise the spatial location of the anchor point of the group and / or an embedding of the spatial location of the point embedding relative to the anchor point of the group. For example, each of the pointATTORNEY DOCKET NO.45288-0439WO1 embeddings in the group may be concatenated with a relative position feature vectordefined as (xi, xj, xi-xj, ^^^ − ^^^^), where xi and xj are respective embeddings (e.g.vectors) the anchor point and point embeddings, to give ananchor point and the point embedding. The edge embeddings are then combined to determine the group embedding. Thus, in some implementations, the group embedding may comprise a weighted sum of embeddings of the spatial locations of the point embeddings in the group.

[0055] In some implementations, updating the point embedding using the group embedding of the selected group can comprise combining the point embedding with the group embedding, e.g., using a neural network (such as an MLP) trained to process an input comprising the point embedding and the group embedding to obtain an updated point embedding as output.

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

[0057] Point cloud data (point embeddings data) is often used to characterize physical environments for many tasks performed using machine learning methods, such as robotic manipulation, autonomous driving, scene recognition and understanding, and so on. Processing point cloud data is, however, a significant challenge for many machine learning methods because the computational resources required, such as processing and / or memory resources, often scale poorly with number of points in the point cloud. In particular, machine learning methods that use neural networks comprising self-attention layers (such as Transformer neural networks) may become prohibitively computationally expensive as the number of points in the point cloud is increased, e.g., to beyond around 100,000 points, particularly in applications where computational resources can be limited, such as robotics. For example, methods that use a self-attention mechanism may have a time- and space-complexity that is quadratic in the number of points in the point cloud, making them unsuitable for large point clouds.

[0058] Applying machine learning methods to point cloud data is also challenging because point clouds are inherently unstructured, the machine learning methods may lack a suitable inductive bias that takes into account locality, i.e., the proximities of the pointsATTORNEY DOCKET NO.45288-0439WO1 to one another. Suitable inductive biases can be created using graph neural networks, but these approaches typically scale poorly with the size of the point cloud.

[0059] To address these issues, the methods and systems described in this specification combine a linear self-attention mechanism with processing that provides an inductive bias that takes into account the distances between the points. In particular, the methods and systems implement a form of self-attention that includes both global self-attention (i.e., each point in the point cloud attends to all the other points, irrespective of locality) and local self-attention in which each particle attends only to a subset of neighboring particles.

[0060] Thus, the methods and systems are able to overcome the quadratic time- and space-complexity limitations imposed with conventional (“vanilla”) self-attention mechanisms, whilst also providing an inductive bias that takes into account distances between the points. Thus, they may be used to process very large point clouds, e.g., point clouds comprising more than 100,000 points, such as may be required for real-world robotics applications.

[0061] 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

[0062] FIG.1 shows an example prediction system.

[0063] FIG.2 shows an example dynamics model.

[0064] FIG.3 shows an example neighbor-attender layer.

[0065] FIG.4 illustrates schematically the operations performed by the neighbor-attender layer of FIG.3 to generate updated point embeddings.

[0066] FIG.5 shows another example prediction system.

[0067] FIG.6 is a flow diagram of an example process for updating point embeddings.

[0068] FIG. 7 is a flow diagram of another example process for updating point embeddings.

[0069] FIGS. 8A-C are graphs showing the performance of the system and methods described in this specification compared to other systems and methods.ATTORNEY DOCKET NO.45288-0439WO1

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

[0071] FIG.1 shows an example prediction system 100. The prediction system 100 is an example of a system implemented as computer programs on one or more computers in one or more locations in which the systems, components, and techniques described below are implemented.

[0072] The prediction system 100 is configured to process a plurality of images 102 of an environment, each image 104 corresponding to a respective time, to generate a predicted image 104 of the environment at a specified time, e.g., before, after, or between the respective time(s) corresponding to the images. For example, the environment can be a real-world environment, and the system can be configured to process the plurality of images 102 to predict how the real-world environment will evolve. As one example, the prediction system 100 can be used to predict how the real-world environment will change in response to one or more actions being performed in the environment, e.g., by a mechanical agent or robot.

[0073] Each image comprises a plurality of pixels, each pixel having an associated set of one or more pixel values. Each set of pixel values may, for example, comprise a respective depth value for the corresponding pixel, Additionally or alternatively, the set of pixel values can comprise one or more pixel values indicative of a color and / or intensity of the pixel. Thus, the images 102 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 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 or intensity values of the pixels in the image, e.g., using a depth estimation machine learning model.

[0074] The prediction system 100 comprises an encoder 108 configure to process the plurality of images 102 to generate a plurality of point embeddings 110. Each point embedding is derived from features characterizing a corresponding spatial location in the environment. For example, the pixel values of the images may be processed using the encoder 108 to generate the point embeddings 110. The encoder 108 can have anyATTORNEY DOCKET NO.45288-0439WO1 architecture appropriate for generating the point embeddings 110, such as a convolutional neural network, a vision transformer, or another image processing neural network. As one example, the encoder can comprise a U-Net neural network, see e.g., Ronneberger et al. arXiv: 1505.04597, configured to generate a per-pixel feature vector (i.e., a respective feature vector for each of the pixels). As a particular example, the encoder 108 can use a U-Net with [32, 64, 128, 256, 256, 128, 64, 32] channels applied to each input image to output 16-dimensional per-pixel feature vectors.

[0075] The encoder can generate a point cloud from the pixels of the images, e.g., using intrinsic and extrinsic parameters of the camera to transform (“unproject”) the 2D pixel coordinates into coordinates in the space of the 3D environment. Each point embedding can comprise a respective feature vector and a respective spatial location in the environment. For example, each point embedding can be a (location, feature)-tuple.

[0076] In some implementations, the plurality of images 102 can comprise images from multiple different viewpoints, e.g., from two or more cameras. Sets of point embeddings obtained from the respective images captured by each camera, and corresponding to the same respective time, can be combined to generate the plurality of point embeddings. In some examples, the plurality of point embeddings 110 can be obtained by sampling point embeddings from the sets of point embeddings. For example, point embeddings can be randomly sampled from the sets of point embeddings until a predetermined total number of point embeddings has been obtained.

[0077] The prediction system 100 further comprises a dynamics model 112 configured to process the plurality of point embeddings 110 to generate a plurality of predicted point embeddings 114 that characterize the environment at the specified time. In general, each point embedding is associated with the respective time at which the corresponding image 104 (from which the point embedding was generated) was captured. In some implementations, the dynamics model 112 can therefore use point embeddings for two or more times to generate the plurality of predicted point embeddings 114. For example,point embeddings at previous time steps ^ − 2 and ^ − 1 can be processed by thedynamics model 112 to generate predicted point embeddings for a next time step ^.

[0078] In some implementations, the dynamics model 112 can be configured to generate the predicted point embeddings by processing kinematic point embeddings 116 in addition to the plurality of point embeddings 114. For example, the kinematic pointATTORNEY DOCKET NO.45288-0439WO1 embeddings can have spatial locations defining a configuration of one or more objects in the environment at the specified time, or at each of a plurality of times including the specified time. As one particular example, the kinematic point embeddings can define positions and / or orientations of parts of a robot or mechanical agent, such as arms, grippers, fingertips, joints, and so on, in which case the kinematic point embeddings 116 can be regarded as defining a kinematic skeleton for the robot or mechanical agent at a given time. The kinematic point embeddings 116 can define estimated locations for the parts, e.g., in response to the robot or mechanical agent performing a particular one or more actions. The dynamics model 112 can then determine how the particle embeddings representing the rest of the environment, i.e., the parts of the environment other than the parts of the robot or mechanical agent defined by the kinematic point embeddings, change in response to the one or more actions being performed by the robot or mechanical agent.

[0079] In some implementations, the prediction system 100 can be configured to apply the dynamics model 112 to the predicted point embeddings 114, i.e., recursively, to generate predicted point embeddings at other times, e.g., at successive time steps.

[0080] One implementation of the dynamics model 112 is described below in more detail in connection with FIG.2.

[0022] The prediction system 100 further comprises a renderer 118 configured to process the predicted point embeddings 114 to generate one or more predicted images 120 of the environment at the specified time, e.g., at a next time step. In general, the renderer 118 can be any type of renderer that can generate an image from the point embeddings. As one example, the renderer 118 can comprise a machine learning model, such as an image generation neural network. In particular, the renderer 118 can comprise a scene representation neural network, that, after training, can generate synthetic images of the physical environment from arbitrary viewpoints.

[0023] 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. As one particular example, the NeRF model can comprise a Point-NeRF model, e.g., as described in Point-basedATTORNEY DOCKET NO.45288-0439WO1 neural radiance fields. in Proceedings of the IEEE / CVF Conference on Computer Vision and Pattern Recognition, pages 5438–5448, 2022.

[0024] At a high level, the model can calculate a plurality of rays through the scene and, at each of a plurality of sampled locations along the ray, determine a set of neighboring particles from which summary statistics describing the particles are determined. The model (e.g., a NeRF MLP) then predicts the color and density at the location in the scene from the summary statistics. Further details of a NeRF model that can be used for the renderer 118 are described in Whitney et al. arXiv:2312.05359.

[0025] The NeRF model can 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.

[0026] 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 view- dependent 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.

[0081] The prediction system 100 can be trained using sequences of images (e.g., videos) obtained using one or more cameras viewing a physical scene as it evolves. For example, the camera(s) can be used to record a sequence of images of a robot or mechanical agent performing a task. The sequence of images can then be used to generate a training example comprising a training input that comprises a first subset of the images in thesequence (e.g., corresponding to earlier times, e.g., for times ^ − 2 and ^ − 1) and aATTORNEY DOCKET NO.45288-0439WO1 target output that comprises a second subset of the images in the sequence (e.g.,corresponding to later times, e.g., for times ^, ^ + 1 … ^ + ^). The training input canadditionally comprise kinematic point embeddings of the robot or mechanical agent at each of the time steps of the task. The prediction system 100 can then be trained using a plurality of such training examples. For each training example, the prediction system 100 can process the one or more images in the training input (and optionally, in conjunction with the kinematic point embeddings) to generate one or more predicted images 120, which are then compared with the corresponding one or more images in the target output. That is, during training, a set of images (e.g., RGB-D images) are encoded into point embeddings 110 using the encoder 108, and the dynamics model 112 applied recursively to generate predicted images 120 at one or more time steps e.g., in the future. The target image and the predicted images can be compared using any suitable objective (loss) function for comparing images, such as a cross-entropy loss, a pixel-wise ^^loss (e.g., between observed and predicted RGB values), and so on. In some implementations, only selected (e.g., randomly sampled) pixels of the predicted images 120 may be generated by the renderer 118, which are then compared with the values of the corresponding images in the target images. Such an approach allows the prediction system 100 to be trained without the computational costs of generating complete images for each training example, and is facilitated by e.g., the use of a Point-NeRF model.

[0082] During training, parameters of the prediction system 100 are adjusted to optimize the objective function, e.g., using any suitable optimizer. For example, parameters of the dynamics model 112 and the encoder 108 can be adjusted to optimize the objective function.

[0083] As one particular example, the prediction system 100 can be trained with a batch size of 16 for a total of 1M steps, using a learning rate schedule of 3e-4 until step 1000, then 1e-4 until step 100K, then 3e-5 until the end of the training, and using the AdamW optimizer, see Loshchilov et al. arXiv:1711.05101v3, with weight decay 1e-3 and clipping the gradient norm to 0.01 (for example) to prevent outliers in the training data from destabilizing training.

[0084] FIG. 2 shows an implementation of the dynamics model 112 in which the dynamics model 112 comprises a Neighbor-Attender layer 202 that is configured to process the point embeddings 110 generated by the encoder 108 to generate updated pointATTORNEY DOCKET NO.45288-0439WO1 embeddings 204 in which each point embedding 110 is updated based on the features of neighboring point embeddings. The Neighbor-Attender layer 202 is described in more detail below in connection with FIGS.3 and 4.

[0085] The dynamics model 112 further comprises: a query layer 204 configured to process the updated point embeddings 204 to generate a query matrix 208; a key layer 210 configured to process the updated point embeddings 204 to generate a key matrix 212; and a value layer 214 configured to process the updated point embeddings 204 to generate a value matrix 216. The dynamics model 112 further comprises a linear self- attention mechanism 218 configured to process the query matrix 208, the key matrix 212, and the value matrix 216 to generate an (approximate) QKV matrix product that the dynamics model 112 uses to update the updated point embeddings 204 to generate the predicted point embeddings 114. In general, each of the query layer 206, the key layer 212 and the value layer comprises parameters that are optimized during training of the dynamics model 112.

[0086] As one particular example, the point embeddings 110 can be denoted ^ ^×^^^ ∈ ℝ ,where ^ is the number of points and ^ is the dimensionality of the embedding of eachpoint. The query matrix 208 can then be determined using ^ = ^^^^^, where ^^ ∈ℝ^×^ !is a matrix transformation performed by the query layer 208 and ^^^is the query / key dimensionality. Similarly, the key matrix 212 can thenusing" = ^ ^×^^^^^, where ^^ ∈ ℝ ! is a matrix transformation performed by the key layer210, and the value matrix 216 can be determined as # = ^ ^×^^^^$, where ^$ ∈ ℝ is amatrix transformation performed by the value layer 214. The linear self-attentionmechanism 218 can be configured to apply a row-wise transformation (^: ℝ^ ! → ℝ',in which the hyper-parameter ( ≪ ^) to the query matrix 208 and the key matrix 212.The row-wise transformation can comprise an element-wise non-linear transformation, e.g. a Rectified Linear Unit (ReLU) or another activation function, for example. Forexample, ^ can be a ReLU function applied element-wise, i.e., such that ^^^ = (. Thetransformed query matrix 208 and the transformed key matrix 212 can be denoted ^′ =^+^, and K′ = ^+",. The linear self-attention mechanism 218 can be configured todetermine the predicted point embeddings 114, denoted by ^ ^×^-. / ∈ ℝ by using:^ = 12 345"6789:;ATTORNEY DOCKET NO.45288-0439WO1in which, 0 = ^++"′,89<, and 9< is the ^ × ^ unit matrix.

[0087] FIG. 3 shows an implementation of the neighbor-attender layer 202, which comprises an anchor particle sampler 302 configured to select a proper subset of the point embeddings 110 to use as “anchor” particles 304, e.g., by uniformly sampling from the point embeddings 110. As one example, one quarter of the point embeddings, or thereabouts, can be selected as anchor particles 304.

[0088] The neighbor-attender layer 202 further comprises a neighbor particle selector 306 configured to process the anchor particles 304 and the point embeddings 110 to select a respective proper subset of the point embeddings 110 as neighbor particles 308 for each of the anchor particles 304. For example, the neighbor particles 308 can be the k nearest neighbor point embeddings for each anchor point (where k is a positive integer, e.g., k=16).

[0089] The neighbor-attender layer 202 further comprises an edge features generator 310 configured to generate a respective set of edge features 312 for each pair consisting of a neighbor particle 308 and the corresponding anchor particle 304 of the neighbor particle. For example, each set of edge features 312 can be obtained by concatenating the embedding of the neighbor particle 308 with a relative position feature vector determined from the spatial locations of the neighbor particle 308 and the anchor particle 304. Forexample, the relative position feature vector can comprise 4=>, =?, =@ − =?, ^=@ − =?^A;,where =@is a vector defining the spatial location of thea vector defining the spatial location of the anchor particle 304.

[0090] The neighbor-attender layer 202 further comprises an edge-particle update layer 314 configured to determine updates for the respective embeddings of each of the anchor particles 304 based on the respective sets of edge features 312 of the anchor particle 304 to generate updated anchor particle 316. For example, the edge-particle update layer 314 can determine an updated embedding B^3for the i-th anchor particle as a weighted sum of the edge features B^,Cin the set of edge features 312 D^for the anchor particle using a multilayer perceptron (MLP), e.g., parameterized by trainable parameters E2, using: exp 4MLP5B^,C; E27;B^,CATTORNEY DOCKET NO.45288-0439WO1 The updated embedding for each anchor particle may, for example, be referred to as a group embedding.

[0091] The neighbor-attender layer 202 further comprises an anchor particle selector 318 configured to determine, for each of the point embeddings 110, a corresponding anchor particle 320 S+T, of the anchor particles 304 that is closest to the spatial location of the point embedding 110.

[0092] The neighbor-attender layer 202 further comprises an output layer 322 configured to update each of the point embeddings 110 based on the updated embedding 316 of the corresponding anchor particle. For example, the output layer 322 can comprise a multi- layer perceptron configured to process the point embedding 110 B^of the updated embedding 316 BU3+^,of the corresponding anchor particle to generate updated point embeddings e.g.,B33 3^ = MLP5B^, BU+^, ; E^7where E^denotes the trainable parameters of the MLP.

[0093] FIG.4 illustrates the operations performed by the neighbor-attender layer 202 of FIG. 3 to generate the updated point embeddings 326. The operations comprise: a first operation 402 that includes selecting anchor particles 304 from the point embeddings 110; a second operation 404 in which the neighbor particles 308 for each anchor particle 304 are selected from the anchor particles 304; a third operation 406 in which respective edge features 312 are generated for each point embedding 110 and anchor particle 304 pair; a fourth operation 408 in which each anchor particle 304 is updated based on the edge features 312; and a fifth operation 410 in which the nearest anchor particles 320 for each point embedding 110 are determined and the point embeddings 110 updated using the updated embedding of the corresponding nearest anchor particle 304.

[0094] FIG. 5 shows another prediction system 500 configured to process a first pair ofimages 504A of an environment at a first time V − 1, a second pair of images 504B of theenvironment at a second time V − 2, kinematic point embeddings 516 at the first time, thesecond time and a current time V, to generate predicted point embeddings 514 for the current time, e.g., for processing by a renderer (not shown) to generate a predicted image of the environment at the current time.

[0095] The prediction system 500 comprises a U-Net (although other types of encoder can be used) configured to convert each pair of images 504A, B into a respective set ofATTORNEY DOCKET NO.45288-0439WO1 point embeddings 510A, 510B which are then processed by a neighbor-attender layer to generate updated point embeddings that are processed by a linear self-attention mechanism, in this case, a Performer point cloud transformer 530 (Performer-PCT). Although FIG. 5 appear to show each pair of images 504A, 504B being processed by a different U-Net, neighbor-attender layer 502 and Performer-PCT 530A, this is merely to illustrate the processing performed for each pair of images 504A, 504B, rather than duplication of these components.

[0096] The kinematic point embeddings 516 can be updated using a point cloud transformer (PCT) 532, which can be a “regular” PCT, e.g., with a quadratic scaling attention mechanism, rather than a Performer-PCT, as the number of kinematic point embeddings is generally small enough for this to be computational practicable. The PCT 532 can be incorporated into the prediction system 500 in some implementations.

[0097] The prediction system 500 further comprises a first merger layer 528 that is configured to combine (e.g., concatenate) the updated point embeddings output by the Performer-PCT 530A for the first time and the second time and the updated kinematic point embeddings 516 to generate a combined set of point embeddings.

[0098] The prediction system 500 further comprises a further neighbor-attender layer 504B configured to update the combined set of point embeddings, followed by a further performer PCT 530B configured to generate the predicted point embeddings 514 for the current time. The further performer PCT 530B is configured to determine respective changes in the spatial location and features for each point embedding 510A for the firsttime V − 1, which are then combined with the point embedding 510 for the first time togenerate a corresponding predicted point embedding 514 for the current time V.

[0099] FIG.6 is a flow diagram of an example process 600 for updating a representation of an environment. For convenience, the process 600 will be described as being performed by a system of one or more computers located in one or more locations. For example, a prediction system, e.g., the prediction system 100 of FIG. 1, appropriately programmed in accordance with this specification, can perform the process 600.

[0100] The representation comprises a plurality of point embeddings, each point embedding being derived from features characterizing a corresponding spatial location in the environment. The system updates 602 each of the point embeddings using a respective proper subset of the point embeddings, each subset being selected based on theATTORNEY DOCKET NO.45288-0439WO1 spatial location of the point embedding. The system then updates 604 the point embeddings by applying a linear self-attention mechanism over the point embeddings, which includes using 604 the point embeddings to determine a query matrix, a key matrix and a value matrix, and updating 606 the point embeddings using a matrix product of the query matrix, the key matrix, and the value matrix.

[0101] FIG. 7 is an example process 700 of determining a representation of an environment at a specified time. For convenience, the process 700 will be described as being performed by a system of one or more computers located in one or more locations. For example, a prediction system, e.g., the prediction system 100 of FIG.1, appropriately programmed in accordance with this specification, can perform the process 700.

[0102] The system obtains 702 a first representation comprising a first plurality of point embeddings, each point embedding being derived from features characterizing a corresponding spatial location in the environment at a first time. The system also obtains 704 a second representation comprising a second plurality of point embeddings, each point embedding being derived from features characterizing a corresponding spatial location in the environment at a second time different from the first time.

[0103] The system then merges 708 the first plurality of point embeddings with the second plurality of point embeddings to obtain a merged representation comprising the merged plurality of point embeddings.

[0104] The system updates 710 the merged representation to determine a representation of the environment at the specified time. The merged plurality of point embeddings are used to update at least the point embeddings of the first plurality of the point embeddings.

[0105] FIGS. 8A, 8B, and 8C compare results 800 obtained using the prediction system 500 of FIG. 5 with results from a Performer-PCT model 802 and results 804 from a hierarchical Graph Neural Network (GNN), see arXiv:2312.05359, as a function of the number of particles used. The GNN was unable to run with 65k or 131k particles due to memory limitations. Each number of particles on the horizontal axis is associated with a triplet of bars, with each bar showing the respective result for each model, in the following order from left to right (GNN model, Performer-PCT model, the prediction system 500).

[0106] FIG.8A shows prediction SSIM (structural similarity index measure). Prediction quality is observed to increase with the number of particles used, with the prediction system 500 performs best for the largest number of particles.ATTORNEY DOCKET NO.45288-0439WO1

[0107] FIG. 8B shows the number of steps performed per second by each model. The prediction system 500 is faster than the GNN system and, although slower than the Performer-PCT mode, achieves better results.

[0108] FIG.8C shows the amount of memory needed to train each model, with the prediction system 500 using less memory than the GNN model, allowing it to scale to larger point clouds.

[0109] The updated representation (“point cloud”) of the environment, such as the updated point embeddings generated by the prediction system 100 of FIG. 1, or the updated point embeddings of the process 600 of FIG.6, or the updated representation of the environment determined using the process 700 of FIG.7, can be used in a wide variety of downstream tasks, some examples of which now follow.

[0110] 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.

[0111] For example, the representation of the environment can characterize the environment at a current time and the updated representation can characterize the environment at another time that is after the current time, e.g., to simulate how the environment will evolve between the current time and the other time (i.e., prediction). Alternatively, the updated representation can characterize the environment at another time that is before the current time, e.g., to determine how the environment came to be characterized by the environment at the current time (i.e., retrodiction). For example, generating the updated representation of the environment can be conditioned on each of one or more representations characterizing the environment at a respective time before or after the current time. As one example, the method can be performed iteratively, starting from a representation of the environment at a starting time, to obtain an updated representation at each of one or more time steps after the starting time.

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

[0113] In some implementations, the techniques described in this specification can be used for real-world control such as controlling an electronic or mechanical agent (such as a robot) in a real-world environment to perform a specified 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.

[0114] 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.

[0115] 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.

[0116] 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,ATTORNEY DOCKET NO.45288-0439WO1 mechanical parts, electrical parts, walls, floors, roadway surfaces, conveyor belt surfaces, and so forth. 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 the physical 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.

[0117] Image(s) may be generated by processing the representation(s) of the physical environment at a given time step using a machine learning model, such as image generation neural network. In some cases, the image generation neural network may be a scene representation neural network, that, after training, can generate synthetic images of the physical environment from arbitrary viewpoints.

[0118] 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 canATTORNEY DOCKET NO.45288-0439WO1 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.

[0119] 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.

[0120] One advantage of implementations of the above-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). For example, the self-attention mechanism can be implemented as a multi-headed self-attention mechanism in which the representation is provided to each of a plurality of self-attention mechanisms (e.g. self- attention layers) and the respective outputs (e.g. updated representations) of the self- attention mechanisms combined to generate the updated representation. For example, different self-attention layers can be assigned to different hardware accelerator units such that the processing can be performed in parallel. The one or more computers that perform the method may therefore, in some examples, include multiple hardware accelerators configured to implement a self-attention mechanism as a multi-headed self-attention mechanism performed in parallel.

[0121] This specification uses the term “configured” in connection with systems and computer program components. For a system of one or more computers to be configured to perform particular operations or actions means that the system has installed on it software, firmware, hardware, or a combination of them that in operation cause the system to perform the operations or actions. For one or more computer programs to be configured to perform particular operations or actions means that the one or more programs include instructions that, when executed by data processing apparatus, cause the apparatus to perform the operations or actions.ATTORNEY DOCKET NO.45288-0439WO1

[0122] Embodiments of the subject matter and the functional operations described in this specification can be implemented in digital electronic circuitry, in tangibly- embodied computer software or firmware, in computer hardware, including the structures disclosed in this specification and their structural equivalents, or in combinations of one or more of them. Embodiments of the subject matter described in this specification can be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible non-transitory storage medium for execution by, or to control the operation of, data processing apparatus. The computer storage medium can be a machine-readable storage device, a machine-readable storage substrate, a random or serial access memory device, or a combination of one or more of them. Alternatively or in addition, the program instructions can be encoded on an artificially-generated propagated signal, e.g., a machine-generated electrical, optical, or electromagnetic signal, that is generated to encode information for transmission to suitable receiver apparatus for execution by a data processing apparatus.

[0123] The term “data processing apparatus” refers to data processing hardware and encompasses all kinds of apparatus, devices, and machines for processing data, including by way of example a programmable processor, a computer, or multiple processors or computers. The apparatus can also be, or further include, special purpose logic circuitry, e.g., an FPGA (field programmable gate array) or an ASIC (application-specific integrated circuit). The apparatus can optionally include, in addition to hardware, code that creates an execution environment for computer programs, e.g., code that constitutes processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of them. Thus a system, artificial neural network, or trained artificial neural network as described herein, can be implemented in hardware using electronic circuitry, e.g., in a physical box. Similarly computer code as described herein can be code to emulate such hardware or code for a hardware description language.

[0124] A computer program, which may also be referred to or described as a program, software, a software application, an app, a module, a software module, a script, or code, can be written in any form of programming language, including compiled or interpreted languages, or declarative or procedural languages; and it can be deployed in any form,ATTORNEY DOCKET NO.45288-0439WO1 including as a stand-alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may, but need not, correspond to a file in a file system. A program can be stored in a portion of a file that holds other programs or data, e.g., one or more scripts stored in a markup language document, in a single file dedicated to the program in question, or in multiple coordinated files, e.g., files that store one or more modules, sub-programs, or portions of code. A computer program can be deployed to be executed on one computer or on multiple computers that are located at one site or distributed across multiple sites and interconnected by a data communication network.

[0125] In this specification the term “engine” is used broadly to refer to a software- based system, subsystem, or process that is programmed to perform one or more specific functions. Generally, an engine will be implemented as one or more software modules or components, installed on one or more computers in one or more locations. In some cases, one or more computers will be dedicated to a particular engine; in other cases, multiple engines can be installed and running on the same computer or computers.

[0126] The processes and logic flows described in this specification can be performed by one or more programmable computers executing one or more computer programs to perform functions by operating on input data and generating output. The processes and logic flows can also be performed by special purpose logic circuitry, e.g., an FPGA or an ASIC, or by a combination of special purpose logic circuitry and one or more programmed computers.

[0127] Computers suitable for the execution of a computer program can be based on general or special purpose microprocessors or both, or any other kind of central processing unit. Generally, a central processing unit will receive instructions and data from a read-only memory or a random access memory or both. The elements of a computer are a central processing unit for performing or executing instructions and one or more memory devices for storing instructions and data. The central processing unit and the memory can be supplemented by, or incorporated in, special purpose logic circuitry. Generally, a computer will also include, or be operatively coupled to receive data from or transfer data to, or both, one or more mass storage devices for storing data, e.g., magnetic, magneto-optical disks, or optical disks. However, a computer need notATTORNEY DOCKET NO.45288-0439WO1 have such devices. Moreover, a computer can be embedded in another device, e.g., a mobile telephone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a Global Positioning System (GPS) receiver, or a portable storage device, e.g., a universal serial bus (USB) flash drive, to name just a few.

[0128] Computer-readable media suitable for storing computer program instructions and data include all forms of non-volatile memory, media and memory devices, including by way of example semiconductor memory devices, e.g., EPROM, EEPROM, and flash memory devices; magnetic disks, e.g., internal hard disks or removable disks; magneto-optical disks; and CD-ROM and DVD-ROM disks.

[0129] To provide for interaction with a user, embodiments of the subject matter described in this specification can be implemented on a computer having a display device, e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user and a keyboard and a pointing device, e.g., a mouse or a trackball, by which the user can provide input to the computer. Other kinds of devices can be used to provide for interaction with a user as well; for example, feedback provided to the user can be any form of sensory feedback, e.g., visual feedback, auditory feedback, or tactile feedback; and input from the user can be received in any form, including acoustic, speech, or tactile input. In addition, a computer can interact with a user by sending documents to and receiving documents from a device that is used by the user; for example, by sending web pages to a web browser on a user’s device in response to requests received from the web browser. Also, a computer can interact with a user by sending text messages or other forms of message to a personal device, e.g., a smartphone that is running a messaging application, and receiving responsive messages from the user in return.

[0130] Data processing apparatus for implementing machine learning models can also include, for example, special-purpose hardware accelerator units for processing common and compute-intensive parts of machine learning training or production, i.e., inference, workloads.

[0131] Machine learning models can be implemented and deployed using a machine learning framework, e.g., a TensorFlow framework, a Microsoft Cognitive Toolkit framework, an Apache Singa framework, or an Apache MXNet framework.ATTORNEY DOCKET NO.45288-0439WO1

[0132] Embodiments of the subject matter described in this specification can be implemented in a computing system that includes a back-end component, e.g., as a data server, or that includes a middleware component, e.g., an application server, or that includes a front-end component, e.g., a client computer having a graphical user interface, a web browser, or an app through which a user can interact with an implementation of the subject matter described in this specification, or any combination of one or more such back-end, middleware, or front-end components. The components of the system can be interconnected by any form or medium of digital data communication, e.g., a communication network. Examples of communication networks include a local area network (LAN) and a wide area network (WAN), e.g., the Internet.

[0133] The computing system can include clients and servers. A client and server are generally remote from each other and typically interact through a communication network. The relationship of client and server arises by virtue of computer programs running on the respective computers and having a client-server relationship to each other. In some embodiments, a server transmits data, e.g., an HTML page, to a user device, e.g., for purposes of displaying data to and receiving user input from a user interacting with the device, which acts as a client. Data generated at the user device, e.g., a result of the user interaction, can be received at the server from the device.

[0134] 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.

[0135] 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 beATTORNEY DOCKET NO.45288-0439WO1 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 should not 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.

[0136] 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.

[0137] What is claimed is:

Claims

ATTORNEY DOCKET NO.45288-0439WO1 CLAIMS:

1. A method performed by one or more data processing apparatus for updating arepresentation of an environment comprising a plurality of point embeddings, each point embedding being derived from features characterizing a corresponding spatial location in the environment, the method comprising: updating each of the point embeddings using a respective proper subset of the point embeddings, each subset being selected based on the spatial location of the point embedding; and updating the point embeddings by applying a linear self-attention mechanism over the point embeddings, wherein applying the linear self-attention mechanism over the point embeddings comprises: using the point embeddings to determine a query matrix, a key matrix and a value matrix; and updating the point embeddings using a matrix product of the query matrix, the key matrix, and the value matrix.

2. The method of claim 1, wherein the query matrix and the key matrix have adimension that is less than a number of point embeddings in the plurality of the point embeddings.

3. The method of claim 1 or 2, further comprising, prior to determining the matrixproduct, applying a row-wise transformation to the query matrix and the key matrix.

4. The method of claim 3, wherein the row-wise transformation comprises anelement-wise non-linear transformation.

5. The method of any preceding claim, wherein updating the point embeddings usinga matrix product of the query matrix, the key matrix, and the value matrix comprises: multiplying the key matrix and the value matrix to obtain a key-value matrix; multiplying the query matrix and the key-value matrix to obtain a query-key-value matrix; andATTORNEY DOCKET NO.45288-0439WO1 updating the point embeddings using the query-key-value matrix.

6. The method of any preceding claim, wherein each proper subset is a group obtained by grouping the point embeddings by spatial location and the method further comprises determining a corresponding group embedding for each group by combining the point embeddings in the group, and wherein updating each of the point embeddings comprises updating the point embeddings using one or more of the group embeddings.

7. The method of claim 6, wherein updating each of the point embeddings using one or more of the group embeddings comprises: selecting a respective one of the groups based on the spatial location of the point embedding; and updating the point embedding using the group embedding of the selected group.

8. The method of claim 6 or 7, wherein grouping the point embeddings by spatial location comprises, for each group: sampling a spatial location from the point embeddings to use as an anchor point for the group; and selecting the point embeddings of the group based on proximities of the spatial locations of the point embeddings to the anchor point.

9. The method of claim 8, wherein selecting the point embeddings for the group comprises selecting a predetermined number of point embeddings having corresponding spatial locations that are nearest the anchor point.

10. The method of claim 8 or claim 9, wherein the anchor point of the selected group is the anchor point nearest the spatial location of the point embedding.

11. The method of any of claims 6-10, wherein combining the point embeddings in the group comprises determining a weighted sum of the point embeddings in the group.ATTORNEY DOCKET NO.45288-0439WO1 12. The method of claim 10 or 11, wherein the group embedding comprises a weighted sum of embeddings of the spatial locations of the point embeddings in the group.

13. The method of claim 12 when dependent on claim 8, wherein the embedding of the spatial location of each point embedding in the group comprises an embedding of the spatial location of the anchor point of the group.

14. The method of claim 12 or 13, when dependent on claim 8, wherein the embedding of the spatial location of each point embedding in the group comprises an embedding of the spatial location of the point embedding relative to the anchor point.

15. The method of any preceding claim, wherein the point embeddings comprise point embeddings that are each derived from features characterizing a corresponding spatial location in the environment at different respective times.

16. A method of determining a representation of an environment at a specified time, the method comprising: obtaining a first representation comprising a first plurality of point embeddings, each point embedding being derived from features characterizing a corresponding spatial location in the environment at a first time; obtaining a second representation comprising a second plurality of point embeddings, each point embedding being derived from features characterizing a corresponding spatial location in the environment at a second time different from the first time; merging the first plurality of point embeddings with the second plurality of point embeddings to obtain a merged representation comprising the merged plurality of point embeddings; updating the merged representation to determine a representation of the environment at the specified time, wherein the merged plurality of point embeddings are used to update at least the point embeddings of the first plurality of the point embeddings.ATTORNEY DOCKET NO.45288-0439WO1 17. The method of claim 16, wherein the second plurality of point embeddings comprises point embeddings having spatial locations defining a configuration of one or more objects in the environment at the specified time.

18. The method of claim 17, wherein the second plurality of point embeddings comprises point embeddings having spatial locations defining the respective configuration of the one or more objects in the environment at each of a plurality of times including the specified time.

19. The method of claim 18, wherein the point embeddings having spatial locations defining the respective configuration of the one or more objects in the environment each comprise an embedding of a respective one of the plurality of times.

20. The method of claim 19, wherein the point embeddings having spatial locations defining the respective configuration of the one or more objects in the environment each comprise an embedding identifying a respective one of the one or more objects.

21. The method of any of claims 17-20, further comprising, before merging the first plurality of point embeddings with the second plurality of point embeddings, updating the second representation by applying a self-attention mechanism over the second plurality of point embeddings.

22. The method of any of claims 17-21, wherein the one or more objects comprise respective parts of a mechanical agent or robot.

23. The method of claim 22, wherein the spatial locations defining the configuration of one or more objects in the environment at the specified time are predicted based on an action to be performed by the mechanical agent or robot at a time preceding the specified time.

24. The method of claim 23, wherein the time preceding the specified time is the first time.ATTORNEY DOCKET NO.45288-0439WO1 25. The method of any of claims 17-24, further comprising: obtaining a third representation comprising a third plurality of point embeddings, each point embedding being derived from features characterizing a corresponding spatial location in the environment at a third time different from the first time and the second time; and merging the third plurality of point embeddings with the merged plurality of point embeddings.

26. The method of any of claims 16-25, wherein updating the merged representation to determine a representation of the environment at the specified time comprises using the method of any of claims 1-14 to update the merged representation.

27. The method of any of claims 16-26, further comprising using the method of any one of claims 1-14 to update the first representation and / or the second representation before the first plurality of point embeddings is merged with the second plurality of point embeddings.

28. The method of any one of claims 16-27, wherein the first time is intermediate the second time and the specified time.

29. The method of any of claims 16-28, wherein updating the merged representation to determine the representation of the environment at the specified time comprises processing the merged plurality of point embeddings to determine, for each of the point embeddings of the first plurality of point embeddings, a respective change in spatial location between the first time and the specified time.

30. The method of any of claims 16-29, wherein updating the merged representation to determine the representation of the environment at the specified time comprises processing the merged plurality of point embeddings to determine, for each of the first plurality of point embeddings, respective element-wise differences between the point embedding at the specified time and the point embedding at the first time.ATTORNEY DOCKET NO.45288-0439WO1 31. The method of any one of the preceding claims, wherein the environment is a real- world environment and obtaining the representation of the environment comprises processing one or more images of the real-world environment to determine the point embeddings.

32. The method of claim 31, wherein each image comprises a plurality of pixels, each pixel having an associated set of one or more pixel values, and processing the image to determine the point embeddings comprises processing the pixel values to determine, for each of the point embeddings, the features characterizing the corresponding spatial location in the environment.

33. The method of claim 32, wherein each set of one or more pixel values comprises a respective depth value for the corresponding pixel.

34. The method of claim 32 or 33, wherein each set of one or more pixel values comprises one or more pixel values indicative of a color and / or intensity of the pixel.

35. The method of any of claims 31-34, further comprising using the updated representation to perform one or more of: (i) a classification task in which the updated representation is used to determine an output classifying the environment into one or more categories; (ii) an object recognition or localization task in which the updated representation is used to determine an output indicative of whether one or more objects is present in the environment or an output defining coordinates of a respective region for one or more objects in the environment; (iii) an action or gesture recognition task in which the updated representation is used to determine an output indicative of whether one or more actions or gestures is being performed in the environment; (iv) a semantic segmentation task in which the updated representation is used to determine an output for assigning spatial locations in the environment to respective segmentation categories;ATTORNEY DOCKET NO.45288-0439WO1 (v) a keypoint detection task in which the updated representation is used to determine an output comprising coordinates of one or more keypoints in the environment; and (vi) an environment similarity task in which the updated representation is used to determine an output indicative of a similarity of the environment to one or more predetermined environments.

36. The method of any preceding claim, further comprising processing the updated representation using an image generation neural network to generate an image.

37. The method of claim 36, wherein the linear self-attention mechanism is implemented by self-attention neural network layers that are trained with the image generation neural network using a loss function that compares sets of pixel values of the generated image of the environment with corresponding sets of pixel values of a target image of the environment.

38. The method of claim 36 or 37, wherein the image generation neural network is a scene representation neural network and processing the updated representation comprises: processing the updated representation to determine scene parameters of the scene representation neural network; and generating the image using the scene representation neural network configured with the parameters.

39. The method of claim 38, wherein the scene parameters encode a representation of a scene comprising the environment 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;ATTORNEY DOCKET NO.45288-0439WO1 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 at the spatial location, to determine the one or more pixel values for the pixel in the image.

40. The method of any preceding claim, further comprising using the updated representation to determine one or more actions to be performed by an agent interacting with the environment to perform a specified task.

41. The method of claim 40, wherein the environment is a real-world environment and the agent comprises (i) a mechanical agent or robot interacting with the real-world environment to perform the task, or (ii) an electronic agent controlling items of equipment in the real-world environment to perform the task.

42. The method of claim 41, wherein the or each representation is obtained by processing sensor data obtained by one or more sensors in the real-world environment to determine, for each of the point embeddings, the features characterizing the corresponding spatial location in the environment.

43. The method of any one of claims 40 to 42, wherein generating the updated representation of the environment is conditioned on an embedding of one or more actions that can be performed by the agent interacting with the environment.

44. The method of any one of claims 40 to 43, further comprising using the representation or the updated representation to determine that the task has been completedATTORNEY DOCKET NO.45288-0439WO1 or will be completed by the agent performing the one or more actions determined using the updated representation.

45. The method of any preceding claim, wherein the representation characterizes the environment at a current time and the updated representation characterizes the environment at another time that is before or after the current time.

46. The method of claim 45, wherein generating the updated representation of the environment is conditioned on each of one or more representations characterizing the environment at a respective time before or after the current time.

47. A method of training a machine learning model, the method comprising: obtaining training examples that each comprise (i) one or more representations of an environment and (ii) a target output; and for each of the training examples: using the method of any one of the preceding claims to update the one or more representations of the environment of the training example; and using the updated one or more representations and the target output of the training example to train the machine learning model.

48. The method of claim 47, wherein the machine learning model is an action selection policy for an agent interacting with the environment.

49. 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-48.ATTORNEY DOCKET NO.45288-0439WO1 50. 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-48.

51. 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 comprising: obtaining a representation of an environment comprising a plurality of point embeddings, each point embedding being derived from features characterizing a corresponding spatial location in the environment; updating each of the point embeddings using a respective proper subset of the point embeddings, each subset being selected based on the spatial location of the point embedding; and updating the point embeddings by applying a linear self-attention mechanism over the point embeddings, wherein applying the linear self-attention mechanism over the point embeddings comprises: using the point embeddings to determine a query matrix, a key matrix and a value matrix; and updating the point embeddings using a matrix product of the query matrix, the key matrix, and the value matrix.

Images