> Markdown version of [/jobs/ext/3989875-staff-machine-learning-engineer-infrastructure](https://www.wearedevelopers.com/jobs/ext/3989875-staff-machine-learning-engineer-infrastructure). Every page supports `.md` or `Accept: text/markdown`. Links point to the HTML versions so they work for humans too. Agent guide: [/agents.md](https://www.wearedevelopers.com/agents.md).

---

# Staff Machine Learning Engineer, Infrastructure

- **Company:** Waymo LLC
- **Location:** San Francisco, CA, United States
- **Experience:** Experienced
- **Salary:** $251,000.0 - $310,000.0
- **Contract:** Permanent contract
- **Skills:** Business Logic, Profiling, Information Engineering, Programming Tools, Distributed Computing Environment, Machine Learning, Tensorflow, Software Engineering, Reinforcement Learning, Pytorch, Large Language Models, Multi-Agent Systems, Machine Learning Operations
- **Published:** October 28, 2026
- **Apply:** https://www.adzuna.com/details/5537302619

## About the Role

* 6+ years of professional software engineering experience, with at least 4 years focused on machine learning infrastructure (scaling, training, optimizing, and deploying large-scale ML systems).
 * Direct ML programming experience on TPU and GPU hardware using frameworks such as JAX, PyTorch, or TensorFlow.
 * Proven hands-on experience scaling large models using model parallelism, data parallelism, or distributed training techniques.
 * Strong understanding of state-of-the-art ML models (e.g., autoregressive transformers) and hands-on proficiency with ML accelerator profiling tools to diagnose bottlenecks.
 * Demonstrated ability to independently lead ambiguous technical initiatives end-to-end and build robust libraries, pipelines, and developer tooling.
 * Strong verbal and written communication skills to collaborate effectively across distributed, cross-functional teams.

Preferred Qualifications

 * Practical experience in Reinforcement Learning (RL), Sim2Real transfer, or Robotics.
 * Experience with distributed ML frameworks and accelerators like GPU/TPU.
 * Domain familiarity with Autonomous Driving systems, multi-agent simulations, or realistic world modeling.

## Description

Join our ML Infrastructure engineering team advancing state-of-the-art ultra-realistic multi-agent simulations using foundation models. In this role, you will work at the intersection of ML infrastructure, foundation models, and simulation engineering, with a specific focus on writing high-performance business and simulation logic in JAX/TensorFlow running directly on TPUs to power realistic environments for Reinforcement Learning (RL).

What You'll Do

 * Design, build, and optimize realistic simulation environments and business logic running on TPUs using JAX and TensorFlow. Implement and optimize large-scale model and data parallelism strategies for training and running foundation models on TPU hardware.
 * Collaborate closely with modeling teams to integrate foundation models into simulation pipelines.
 * Drive technical architectures and system designs from data engineering through simulation execution to meet business and performance objectives.
 * Profile systems, identify performance bottlenecks across ML accelerators, and optimize end-to-end execution speed.
 * Translate product and business goals into concrete technical requirements and system deliverables.