jobdescription_image_2x.png

Staff/Sr. Staff Machine Learning Engineer - World Modeling, Simulation

Mountain View, California, United States Full-Time Software Engineering 3394

Apply now

Staff/Sr. Staff Machine Learning Engineer - World Modeling, Simulation

  • 3394
  • On Site
  • Mountain View, California
  • Software Engineering
  • Full-Time

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.

The Simulator Team at Waymo builds state-of-the-art simulations of realistic environments for testing, training, and validation of the Waymo Driver. Our team is a multidisciplinary group of machine learning (ML) engineers, software engineers, and research scientists pushing the boundaries of autonomous systems. We develop industry-leading simulation platforms using large-scale generative and physical world modeling algorithms to model the physical world—encompassing dynamic agents, complex road topologies, multi-modal traffic systems, weather phenomena, and the vehicle’s full sensor suite (Camera, Lidar, Radar).

To pioneer the next frontier of fidelity, scalability, and physical controllability, we are expanding our focus on foundation world models. In this role, you will define and drive the technical roadmap for large-scale generative architectures, bridging cutting-edge multimodal pre-training (across 4D video, LiDAR, and Radar), spatiotemporal tokenization, and causal post-training/distillation to generate interactive, multi-platform, safety-critical driving worlds for closed-loop evaluation and onboard model training.

In this role, you will report to a Senior Staff Engineering Manager

You Will:

  • Own the Architectural Roadmap: Drive the technical vision and long-term roadmap for foundational 4D world modeling, unified spatiotemporal tokenizers (3D/4D VAEs), video generation architectures, and generative sensor simulation across Waymo -- enabling zero-shot generalization across next-generation vehicle platforms, sensor rigs, and international geographies.
  • Lead Large-Scale Pre-Training: Architect, train, and scale foundation video and physical world models (including Sparse MoE and long-context block-sparse attention architectures) across distributed TPU/GPU clusters, establishing efficient training recipes, multimodal data curation strategies, and scaling laws.
  • Drive Advanced Post-Training & Alignment: Design and implement robust post-training paradigms, including Supervised Fine-Tuning (SFT), Reinforcement Learning with AI/Critic Feedback (RLAIF), model merging/souping, and 3D/semantic conditioning to guarantee strict physical realism, counterfactual pose-divergence stability, and fine-grained control over safety-critical scene elements (signage, traffic signals, VRU gestures, and adverse weather).
  • Integrate Multimodal & Physical Signals: Develop action-conditioned, 3D-grounded physical world models and multi-view generative pipelines that couple explicit geometric and physical priors (4D scene graphs, occupancy, vehicle kinematics, and interactive agent dynamics) with neural synthesis across Camera, LiDAR, and Radar—faithfully modeling physical sensor transport, occlusions, and causal state transitions.
  • Optimize for Production & Serving: Scale inference and evaluation frameworks into production simulation infrastructure, applying few-step and causal autoregressive distillation (e.g., Distribution Matching Distillation [DMD], Self-Forcing, Consistency/Flow distillation), quantization, and KV-cached streaming inference for 10 Hz high-throughput closed-loop validation and RL gyms.
  • Foster Research-to-Production Impact: Partner closely with Onboard Perception, Planner, and AI Foundation labs across Waymo and Alphabet (e.g., Google DeepMind) to adapt frontier world models into deployment-ready simulation environments that achieve downstream perception/policy parity for rare-event reproduction, synthetic data training, and closed-loop RL.
  • Mentor & Elevate Engineering: Serve as a technical authority, mentoring senior engineers and driving high engineering standards across training stability, reproducibility, and system design.

You Have:

  • Education: PhD or MS in Computer Science, Machine Learning, Robotics, Computer Vision, or a related quantitative field (or equivalent practical research/industry tenure, including academic appointments looking to transition to industry).
  • Senior Technical Leadership: 8+ years of industry or relevant research lab experience (or tenure-track/associate professorship equivalent) developing, deploying, and leading complex ML architectures.
  • Large-Scale Pre-Training Track Record: Demonstrated experience training large-scale generative and physical world models from scratch (e.g., action-conditioned 3D/4D world models, large video diffusion models, autoregressive spatiotemporal transformers, flow matching, or vision-language-action models) on high-performance compute clusters (multi-node GPU/TPU).
  • Post-Training & Alignment Expertise: Hands-on experience designing post-training and student-conversion pipelines for generative models—including few-step diffusion/flow distillation (e.g., DMD, MMD, Consistency Models, Rectified Flow), autoregressive exposure-bias mitigation (e.g., Self-Forcing), SFT, and Reinforcement Learning (RLAIF / reward-model alignment)—to control and stabilize long-horizon rollouts.
  • Deep 3D/4D Vision Understanding: Deep technical foundation in 3D computer vision and physical world modeling—including action-conditioned scene dynamics, explicit 3D physical/geometric priors (4D occupancy, dynamic meshes, rigid/articulated kinematics), spatiotemporal tokenization (3D/4D causal VAEs), continuous ray/camera geometry (e.g., Plücker rays, cross-rig calibration), and generative 4D world representations (coupling explicit physical priors with neural generation).
  • Framework Proficiency: Expert-level proficiency in Python and modern distributed ML frameworks (such as JAX/Flax, XLA/Pallas, PyTorch, Megatron-LM, Pathways, or DeepSpeed) for multi-node model, sequence, and data parallelism.

We Prefer:

  • Domain Pedigree: Prior experience from leading video generation labs, foundation model startups, advanced robotics institutes, or high-tier academic research groups focusing on physical AI, foundation video models, or world simulation.
  • Multi-Modal Sensor Simulation: Background in generative modeling and synthesis of non-RGB physical sensors—specifically 3D LiDAR point clouds/range images and 4D imaging radar—alongside multi-camera automotive rigs across diverse vehicle platforms, weather, and lighting conditions.
  • Embodied AI & World Dynamics: Hands-on experience in physical world modeling and embodied dynamics—including action-conditioned causal world models, differentiable or hybrid physics engines, rigid/articulated body kinematics, vehicle dynamics, interactive multi-agent traffic dynamics, or closed-loop robotics simulators.
  • High-Performance Infrastructure: Experience with low-level systems optimization, XLA/Pallas TPU kernels or custom CUDA/Triton kernels, block-sparse / Flash / Splash attention over 100k+ token sequences, and high-throughput KV-cached distributed serving.
  • Systems Engineering: Working proficiency in modern C++ for high-performance simulation integration.

The expected base salary range for this full-time position across US locations is listed below. Actual starting pay will be based on job-related factors, including exact work location, experience, relevant training and education, and skill level. Your recruiter can share more about the specific salary range for the role location or, if the role can be performed remote, the specific salary range for your preferred location, during the hiring process. 

Waymo employees are also eligible to participate in Waymo’s discretionary annual bonus program, equity incentive plan, and generous Company benefits program, subject to eligibility requirements. 

Salary Range
$251,000—$310,000 USD

We appreciate your interest in Waymo. Waymo is proud to be an equal opportunity employer, committed to creating a culture of belonging and maintaining a supportive workplace for all employees. We welcome applicants of all backgrounds, and employment decisions are based on a candidate’s qualifications, experience, and alignment with job requirements and business needs. Waymo does not discriminate against, and prohibits harassment of, any applicant or employee based on race, color, sex, sexual orientation, gender identity, religion, national origin, age, disability, military status, family status, pregnancy, genetic information or any other basis protected by applicable law. Waymo will also consider for employment qualified applicants with criminal records in accordance with applicable law. Waymo is committed to making sure our hiring process is accessible for all candidates. If you need assistance applying for a role or participating in the interview process due to a disability, please let the recruiting team know or email waymo-candidatesupport@google.com. (This email address is intended to be used only for requesting accommodations as part of the application process. Other inquiries will not receive a response.)

Some of our Benefits

Health and wellness

Our people are at the heart of everything we do. At Waymo, you can enjoy top-notch medical, dental and vision insurance, mental wellness support, gym membership, and special wellness programs.

Financial wellness

Your financial peace of mind is important to us. At Waymo, we offer competitive compensation, bonus opportunities, equity, employees provident fund, and lots of other perks and employee discounts.

Flexibility and time off

Take the time you need to relax and recharge. Enjoy the flexibility to work from another location for four weeks per year. We support an on-site or hybrid work model and offer remote working opportunities, paid time off, bereavement, sick, and parental leave. 

Supporting families

When it comes to growing your family or caring for your loved ones, you have our full support. Enhanced leave options include 26 weeks of paid leave for birthing parents and 18 weeks of paid leave for non-birthing parents.

Community and personal development

At Waymo, you'll find a range of opportunities to grow, connect, and give back. We offer education reimbursement, personal and professional development, mentorship, and other ways to connect through Employee Resource Groups (ERGs), other internal groups, and even time off to volunteer.

Cool perks

Access to Google offices, cafes, wellness centers, massages, and so much more. To support your wellbeing at home, you can enjoy at-home fitness and cooking classes, and more.

Ready to Apply?