Back to Careers
Research Engineer - Reinforcement Learning
San Francisco, CA (In-Person)
General Diffusion is a foundational AI research lab establishing the scientific discipline of Compute Intelligence. We build frontier models that learn the physics of heterogeneous hardware, decoupling intelligence from infrastructure.
<br/>About the role
As a Research Engineer in Reinforcement Learning, you will apply RL techniques to the problem of compute scheduling. You will train agents that learn to predict optimal graph partitions and resource allocations in real-time, turning an NP-hard combinatorial optimization problem into an inference task.
<br/>What you might work on
- Training RL agents to optimize compute graph partitioning for heterogeneous clusters.
- Developing simulation environments that accurately model network latency and bandwidth constraints.
- Implementing "predictive auto-scaling" that anticipates workload spikes before they happen.
- Collaborating with the RS1 team to deploy RL-based schedulers into production.
What we’re looking for
- Strong background in Deep Reinforcement Learning (PPO, SAC, MCTS).
- Experience applying ML to systems problems (ML for Systems).
- Proficiency in PyTorch or JAX.
- Ability to design and implement custom gym environments.
Our culture
- Compute Intelligence. We are establishing a new scientific discipline.
- Silicon Neutrality. We build foundational models that run on any chip.
- Deep Work. We value long periods of uninterrupted focus.