Machine-learned component hybrid training and assistance of vehicle trajectory generation

An imitation learning-based machine-learned (ML) model to augment or replace the prediction and/or planner components of an autonomous vehicle may be trained using a two stage and multi-discipline approach. A first stage of training may include training the ML component to output a predicted action associated with a target vehicle and modifying the ML component to reduce a difference between the predicted action and the observed action taken by the target vehicle. A second stage may use reinforcement learning to further tun the ML component. The resultant model may be used on its own, with enough training data, or to rank or weight candidate trajectories generated by a planning component of the vehicle. The ML component may provide embeddings of environment features to first transformer(s) that output to a long short-term memory that outputs to second transformer(s) to determine the predicted action.