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

Waymo•Mountain View, CA

1d•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 controllable generation and test-time compute algorithms, multimodal conditioning, and critic-guided flow/diffusion steering for foundation world models. The role also investigates multimodal sensor and semantic latent conditioning and representation alignment for holistic world simulation, and uses 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. The intern will write clean, modular, well-tested JAX code in core simulation and world-modeling codebases and collaborate with researchers across Sim Realism, Foundation Models, and Closed-Loop Training teams. Waymo interns work with industry leaders on impactful projects while developing their skills. This is a hybrid onsite internship. Applications are accepted on a rolling basis until the role is filled. The expected hourly rate for this full-time position is $85 USD. Interns are eligible to 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 date between December 2027 and August 2028, with 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/sequence modeling and inference-time scaling, or multimodal spatiotemporal representation learning and closed-loop simulation; working familiarity with the other area.
Research experience for the first area may include continuous flow matching/Rectified Flows, score-based diffusion guidance/steering, test-time compute/SMC, or parameter-efficient fine-tuning.
Research experience for the second area may include 4D/video tokenization, latent dynamics alignment such as REPA/DINO alignment, inverse RL/adversarial discriminators, or multi-agent trajectory modeling.

Nice To Haves

First-author publications at top-tier machine learning, computer vision, or robotics venues (e.g., 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, such as 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.