Trajectory prediction using diffusion models
Methods, systems, and apparatus, including computer programs encoded on computer storage media, for generating trajectory predictions for one or more target agents, e.g., a vehicle, a cyclist, or a pedestrian, in an environment. In one aspect, one of the methods include: obtaining scene context data characterizing a scene at a current time point in an environment that includes multiple target agents; generating, from the scene context data, an encoded representation of the scene in the environment; and generating, by a diffusion model based on the encoded representation, a respective trajectory prediction output that predicts a respective future trajectory for each of the multiple target agents after the current time point.
1 . A computer-implemented method comprising:
obtaining scene context data characterizing a scene at a current time point in an environment that includes multiple target agents;
generating, from the scene context data, an encoded representation of the scene in the environment;
generating, by using a diffusion model and based on the encoded representation, a respective trajectory prediction output that predicts a respective future trajectory for each of the multiple target agents after the current time point by repeatedly updating an intermediate representation of the respective trajectory prediction output across multiple steps in a generation process; and
controlling a vehicle based on the respective trajectory prediction output that predicts the respective future trajectory for each of the multiple target agents.
2 . The method of claim 1 , wherein generating the respective trajectory prediction output comprises, at each of the multiple steps in the generation process:
obtaining, for each of the multiple target agents, a current intermediate representation of the respective trajectory prediction output for the step; and
updating, for each of the multiple target agents, by the diffusion model based on the encoded representation, the current intermediate representation to generate an updated intermediate representation of the respective trajectory prediction output for the step.
3 . The method of claim 2 , wherein obtaining the current intermediate representation of the respective trajectory prediction output for the step comprises:
determining a noise level based on computing a noise level function for the step; and
determining a combination of the noise level and an updated intermediate representation of the respective trajectory prediction output for an immediately preceding step.
4 . The method of claim 3 , wherein determining the combination of the noise level and the updated intermediate representation comprises:
determining a concatenation of the noise level and the updated intermediate representation of the respective trajectory prediction output for an immediately preceding step; and
projecting the concatenation to have a projected dimensionality of the current intermediate representation.
5 . The method of claim 2 , wherein obtaining the intermediate representation of the respective trajectory prediction output for a beginning step comprises:
sampling from a Gaussian distributed noise with a predetermined standard deviation.
6 . The method of claim 2 , wherein the diffusion model is an attention-based diffusion model, and wherein updating the current intermediate representation to generate the updated intermediate representation of the respective trajectory prediction output for the step comprises:
generating, for each of the multiple target agents, a self-attended intermediate representation at least in part by applying a self-attention mechanism over the current intermediate representations;
generating, for each of the multiple target agents, a cross-attended intermediate representation at least in part by applying a cross-attention mechanism over the self-attended intermediate representations and the encoded representation; and
generating, for each of the multiple target agents and from the cross-attended intermediate representations, the updated intermediate representation of the respective trajectory prediction output for the step.
7 . The method of claim 2 , wherein for a particular target agent in the multiple target agents, future trajectories predicted by the updated intermediate representations of the respective trajectory prediction output for two or more different steps in the generation process are the same.
8 . The method of claim 2 , wherein the generating comprises, for each of the multiple target agents:
using the updated intermediate representation of the respective trajectory prediction output for a last step in the generation process as the respective trajectory prediction output for the target agent that predicts the future trajectory of the target agent after the current time point.
9 . The method of claim 1 , wherein the trajectory prediction outputs are generated in a compressed space that has a predetermined dimensionality.
10 . The method of claim 1 , wherein the trajectory prediction output defines a probability distribution over possible future trajectories of each of the multiple target agents after the current time point.
11 . The method of claim 1 , wherein:
the scene context data comprises data generated from data captured by one or more sensors of the vehicle, and
the target agents comprise agents in a vicinity of the vehicle in the environment.
12 . The method of claim 11 , further comprising:
providing (i) the trajectory prediction outputs for the target agents, (ii) data derived from the trajectory prediction outputs, or (iii) both to an on-board system of the vehicle for use in controlling the vehicle.
13 . The method of claim 12 , wherein the trajectory prediction outputs are generated on-board the vehicle.
14 . The method of claim 1 , wherein:
the vehicle is a simulated vehicle;
the scene context data comprises data generated from data that simulates data that would be captured by one or more sensors of the vehicle in the real-world environment, and
the target agents comprise simulated agents in a vicinity of the simulated autonomous vehicle in the computer simulation.
15 . The method of claim 14 , further comprising:
providing (i) the trajectory prediction outputs, (ii) data derived from the trajectory prediction outputs, or (iii) both for use in controlling the simulated vehicle in the computer simulation.
16 . The method of claim 1 , wherein the scene context data comprises target agent history context data characterizing current and previous states of the target agents.
17 . The method of claim 1 , wherein the scene context data comprises context agent history context data characterizing current and previous states of each of the one or more context agents.
18 . The method of claim 1 , wherein the scene context data comprises road graph context data characterizing road features in the scene.
19 . The method of claim 1 , wherein the scene context data comprises traffic signal context data characterizing at least respective current states of one or more traffic signals in the scene.
20 . A system comprising:
one or more computers; and
one or more storage devices storing instructions that, when executed by the one or more computers, cause the one or more computers to perform operations comprising:
obtaining scene context data characterizing a scene at a current time point in an environment that includes multiple target agents;
generating, from the scene context data, an encoded representation of the scene in the environment;
generating, by using a diffusion model and based on the encoded representation, a respective trajectory prediction output that predicts a respective future trajectory for each of the multiple target agents after the current time point by repeatedly updating an intermediate representation of the respective trajectory prediction output across multiple steps in a generation process; and
controlling a vehicle based on the respective trajectory prediction output that predicts the respective future trajectory for each of the multiple target agents.
21 . One or more non-transitory computer-readable storage media storing instructions that when executed by one or more computers cause the one or more computers to perform operations comprising:
obtaining scene context data characterizing a scene at a current time point in an environment that includes multiple target agents;
generating, from the scene context data, an encoded representation of the scene in the environment;
generating, by using a diffusion model and based on the encoded representation, a respective trajectory prediction output that predicts a respective future trajectory for each of the multiple target agents after the current time point by repeatedly updating an intermediate representation of the respective trajectory prediction output across multiple steps in a generation process; and
controlling a vehicle based on the respective trajectory prediction output that predicts the respective future trajectory for each of the multiple target agents.