2027 Summer Intern, PhD, Machine Learning - Simulation & Closed-Loop Training

Waymo

Posted 14 days ago

Internship

Mountain View, California

Hybrid

Smart Summary

Responsibilities

Design and implement reinforcement learning algorithms for foundation world models in simulation and autonomous driving. Develop learned reward models and conduct rigorous closed-loop ablation studies on large-scale TPU clusters.

Qualifications

You are a PhD student in Computer Science, Machine Learning, or a related field, with strong JAX/PyTorch skills and expertise in either Reinforcement Learning or Generative Modeling. You have hands-on experience training modern neural architectures and are passionate about contributing to autonomous driving technology.

Must Have Skills for ATS

JAX

PyTorch

Python

Reinforcement Learning

Generative Modeling

Transformers

Diffusion Models

Deep RL

TPU

XLA

Job Description

Waymo is an autonomous driving technology company with the mission to be the world's most trusted driver. Since its start as the Google Self-Driving Car Project in 2009, Waymo has focused on building the Waymo Driver—The World's Most Experienced Driver™—to improve access to mobility while saving thousands of lives now lost to traffic crashes. The Waymo Driver powers Waymo’s fully autonomous ride-hail service and can also be applied to a range of vehicle platforms and product use cases. The Waymo Driver has provided over ten million rider-only trips, enabled by its experience autonomously driving over 100 million miles on public roads and tens of billions in simulation across 15+ U.S. states.

Software Engineering builds the brains of Waymo's fully autonomous driving technology. Our software allows the Waymo Driver to perceive the world around it, make the right decision for every situation, and deliver people safely to their destinations. We think deeply and solve complex technical challenges in areas like robotics, perception, decision-making and deep learning, while collaborating with hardware and systems engineers. If you’re a software engineer or researcher who’s curious and passionate about Level 4 autonomous driving, we'd like to meet you. 

Waymo interns partner with leaders in the industry on projects that create impact to the company. We believe learning is a two-way street: applying your knowledge while providing you with opportunities to expand your skill-set. Interns are an important part of our culture and our recruiting pipeline. Join us at Waymo for a fun and rewarding internship!

You will:

  • Design, implement, and scale Reinforcement Learning post-training and distillation algorithms (such as flow or diffusion policy optimization, off-policy importance sampling, and constrained multi-objective RL) for foundation world models in simulation and autonomous driving
  • Develop learned reward models, realism critics, and discriminators within a closed-loop multi-agent simulator to align simulated agent and driving behaviors with human driving realism and safety
  • Run rigorous closed-loop ablation studies on large-scale TPU clusters—leveraging agentic auto-research workflows and fleet-scale driving behavior metrics (such as route progress, braking comfort, lateral stability, and lane selection)—to mitigate simulator exploitation and distribution shift
  • Write clean, modular, and well-tested JAX code in core simulation and autonomous driving codebases, collaborating closely with researchers and engineers across simulation, world modeling, and autonomous driving teams

You have:

  • Currently pursuing a PhD in Computer Science, Machine Learning, Robotics, or a related technical field, with an anticipated graduation date between December 2027 and August 2028 (returning to a degree program for at least one academic term after the internship)
  • Core Implementation Fluency (Mandatory): Strong tensor implementation and debugging skills in Python and JAX/Flax or PyTorch, with hands-on experience training and modifying modern neural architectures (such as Transformers, Diffusion/Flow models, or deep RL networks)
  • Deep Research Expertise in Either (a) Reinforcement Learning & Post-Training (hands-on research experience in policy gradient methods such as PPO/GRPO/REINFORCE, off-policy or offline RL such as importance sampling/AWR/DPO, constrained MDPs, inverse RL, or learned reward and critic modeling) OR (b) Generative World, Sequence, or Trajectory Modeling (hands-on research experience in continuous flow matching / Rectified Flows, score-based diffusion policies such as DDPO/DPPO/Diffuser, or spatiotemporal autoregressive world and trajectory models), with working familiarity in the other

We prefer:

  • First-author publications at top-tier machine learning or robotics venues (e.g., NeurIPS, ICML, ICLR, CVPR, CoRL, RSS, ICRA, or equivalent venues) in RL post-training, flow/diffusion policy optimization, generative world models, or autonomous driving
  • Experience with multi-agent closed-loop simulation, Sim2Real transfer, or reward-hacking mitigation (e.g., Waymax, WOMD, NuPlan, IsaacGym/MuJoCo, adversarial traffic generation, or discriminator stabilization)
  • Experience profiling and scaling distributed training workloads on TPUs or multi-node GPUs (such as XLA compilation, vectorized batched rollouts, and memory optimization) and familiarity with C++/Python interfaces

Note: This will be a hybrid onsite internship position. We will accept resumes on a rolling basis until the role is filled. To be in consideration for multiple roles, you will need to apply to each one individually - please apply to the top 3 roles you are interested in.

The expected hourly rate for this full-time position is listed below. Interns are also eligible to participate in the Company’s generous benefits programs, subject to eligibility requirements.Hourly PhD Pay$85—$85 USD

Waymo

On the journey to be the world's most trusted driver. With the Waymo Driver, we can improve mobility while saving thousands of lives. Download the Waymo One app and ride today. Waymo reaches out to candidates from official channels only (e.g. directly from @waymo.com email addresses, or through our recruiters or sourcers who are noted as such on LinkedIn). We do not contact candidates about career opportunities through instant messaging apps like Telegram, email addresses from domains other than waymo.com (such as Gmail addresses), direct messages on Twitter, Facebook, and Instagram, or text messages. Visit waymo.com to check out our official job listings. Need support? Send us a message – we're here to help.
Runway Icon
Boost Your Interview Chances

With Runway

See Your Fit for This Role

1-5 min

Your Score

?

Top Applicants

90%

Your Job Search Advantage

Key Gaps & Next Steps:

Address these in your resume & Interview

Top Strengths For This Role

Highlight these in your cover letter & interview

Your Interview Guide

A Personalized Interview Strategy

Freshest Opportunities

Never Miss a Good Fit

Get notified when jobs mach your criteria