Senior / Staff Machine Learning Engineer - Behavior Models for Road Users

This role in the Offline Driving Intelligence team is responsible for developing/learning behavior models for road users such as cars, bicycles, and pedestrians. These agents populate Zoox's simulations and must be indistinguishable from real road users, yet fully controllable: promptable into rare, adversarial, safety-critical behaviors.

Responsibilities

• Develop new deep learning models that use imitation learning and reinforcement learning to generate driving plans for human-like driving agents.
• Work on novel techniques to estimate the quality of driving plans along dimensions of safety, progress, comfort, and realism.
• Build generative behavior models (e.g., autoregressive, diffusion) conditionable on scenario intent for targeted stress-testing.
• Leverage compute, infrastructure, and large datasets to push the boundaries of the field.
• Develop metrics and tools to analyze errors and track system improvements.
• Collaborate with engineers on Perception, Planning, Simulation, and Validation teams to solve the overall Autonomous Driving problem.

Qualifications

• PhD degree in computer science or related field and 4+ years of relevant professional experience, or master's degree and 7+ years of relevant professional experience.
• Experience in one of the following: Planning, Prediction, Reinforcement Learning, Imitation Learning, generative modeling (diffusion, autoregressive models).
• Experience with training and deploying transformer-based model architectures.
• Experience with production Machine Learning pipelines: dataset creation, training frameworks, metrics pipelines.
• Fluency in Python ML frameworks and a basic understanding of C++.

Bonus Qualifications

• Top tier publications (NeurIPS, ICML, CVPR).
• Experience with JAX.