Abstract: Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for generating trajectory predictions for one or more agents in an environment. In one aspect, a method comprises: obtaining scene context data characterizing a scene in an environment at a current time point and generating a respective predicted future trajectory for each of a plurality of agents in the scene at the current time point by sampling a sequence of discrete motion tokens that defines a joint future trajectory for the plurality of agents using a trajectory prediction neural network that is conditioned on the scene context data.