2027 Summer Intern, PhD, Machine Learning - Generative World Models & Controllable Simulation

Waymo•Mountain View, CA

2d•Hybrid

About The Position

Waymo is hiring a PhD Machine Learning intern to work on generative world models and controllable simulation for autonomous driving. The intern will develop methods for foundation world models, investigate multimodal sensor and semantic latent conditioning, and advance holistic world simulation. The work includes evaluating interactive realism, long-tail scenario synthesis, and multimodal consistency using learned critics and agentic research workflows on large-scale TPU clusters. The intern will contribute clean, modular, and well-tested JAX code to core simulation and world-modeling codebases and collaborate with researchers across Sim Realism, Foundation Models, and Closed-Loop Training teams. This is a hybrid onsite summer internship. Applicants must be pursuing a PhD in a relevant technical field, with an anticipated graduation date between December 2027 and August 2028, and must return to a degree program for at least one academic term after the internship. The expected hourly rate is $85. Interns may participate in company benefits programs, subject to eligibility requirements.

Requirements

Currently pursuing a PhD in Computer Science, Machine Learning, Robotics, or a related technical field.
Anticipated graduation between December 2027 and August 2028, with a return to a degree program for at least one academic term after the internship.
Strong tensor implementation and debugging skills in Python and JAX/Flax or PyTorch.
Hands-on experience training and modifying modern neural architectures, such as Transformers, Diffusion/Flow-Matching models, or deep RL/critic networks.
Hands-on research experience in either controllable generative world or sequence modeling and inference-time scaling—including continuous flow matching/Rectified Flows, score-based diffusion guidance/steering, test-time compute/SMC, or parameter-efficient finetuning—or multimodal spatiotemporal representation learning and closed-loop simulation—including 4D/video tokenization, latent dynamics alignment, inverse RL/adversarial discriminators, or multi-agent trajectory modeling.
Working familiarity with the other of the two specified research areas.

Nice To Haves

First-author publications at top-tier machine learning, computer vision, or robotics venues, such as NeurIPS, ICML, ICLR, CVPR, ECCV, CoRL, RSS, ICRA, or equivalent, in generative world models, flow/diffusion steering, controllable generation, or autonomous driving simulation.
Experience with multi-agent closed-loop simulation, counterfactual/adversarial traffic generation, or Sim2Real discriminator/critic stabilization, such as Waymax, WOMD, NuPlan, MGAIL, or reward-hacking mitigation.
Experience profiling and scaling distributed training and inference workloads on TPUs or multi-node GPUs, including XLA compilation, vectorized batched rollouts, and memory optimization.
Interest in external conference publication on Waymax/WOMD.

Responsibilities

Design, implement, and scale controllable generation and test-time compute algorithms, multimodal conditioning, and critic-guided flow/diffusion steering for foundation world models.
Investigate multimodal sensor and semantic latent conditioning, and representation alignment to advance holistic world simulation.
Use learned realism critics and agentic autoresearch workflows on large-scale TPU clusters to evaluate closed-loop interactive realism, long-tail scenario synthesis, and multimodal consistency.
Write clean, modular, and well-tested JAX code in core simulation and world-modeling codebases.
Collaborate with researchers across Sim Realism, Foundation Models, and Closed-Loop Training teams.