Applicability to off-policy RL algorithms

[TATP-233]

Hi Lars, thanks for open-sourcing this work!

I’m Yufei Jia, the author of UniLab, which you cited when discussing massively parallel robot RL workloads. We are also working on scalable robot RL infrastructure.

I’m curious whether rl-triton could help accelerate off-policy algorithms such as SAC or TD3. If so, which part of their training pipeline do you think would benefit most?

I’d be happy to exchange notes, especially around massively parallel robot RL and AMD/MI300X.

Great work!

[simonsays1980]

@TATP-233 Hi Yufei - great to hear from you, and thanks! I cited UniLab exactly for the massively parallel robot RL regime, so it’s nice to connect.

For vanilla SAC/TD3, I don’t think rl-triton would move much directly. Their core critic updates use one-step Bellman targets sampled from replay, so there isn’t a long trajectory recurrence to parallelize.

The overlap gets more interesting with multi-step targets or trace-based corrections. Longer n-step returns use the same affine recurrence, and Retrace(λ) is already implemented in rl-triton if you use corrected multi-step off-policy targets. So I’d say the benefit depends more on how much multi-step/trace structure your training pipeline has than on SAC/TD3 specifically.

I’d definitely enjoy exchanging notes on massively parallel robot RL. The MI300X angle is especially interesting — Triton supports ROCm, but I’ve only tested rl-triton on NVIDIA so far, so I’d be very interested in seeing how the scan kernels behave there.

Happy to compare setups or hop on a call if useful.