Abstract: Systems and methods for generating attention masks are provided. In particular, a computing system can access sensor data and map data for an area around an autonomous vehicle. The computing system can generate a voxel grid representation of the sensor data and map data. The computing system can generate an attention mask based on the voxel grid representation. The computing system can generate, by using the voxel grid representation and the attention mask as input to a machine-learned model, an attention weighted feature map. The computing system can determine using the attention weighted feature map, a planning cost volume for an area around the autonomous vehicle. The computing system can select a trajectory for the autonomous vehicle based, at least in part, on the planning cost volume.