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
