Questions about training speed

[zhuyaoye]

Thank you for the excellent work. I noticed that several efforts have been made to improve training speed. I’d like to ask: roughly how long does it take to train 1M steps now? Additionally, which part of the model is the main bottleneck limiting the training speed?

[nicklashansen]

Hi @zhuyaoye thank you for your interest in our work! Training speed depends quite a bit on your hardware, as well as task/simulator and training config (especially torch.compile and whether you use parallel envs or not). To give you some ballpark estimates, I'd say single-task state-based DMControl online RL should be approx 1.5hrs / 1M env steps with torch.compile and 1 env on a RTX 3090 GPU. If you instead use 10 envs (e.g. by switching to https://github.com/nicklashansen/newt) then it'll be about 1hr / 1M env steps, and if you in addition to that also upgrade to a RTX 5090 then it'll be closer to 30 min / 1M env steps.

[zhuyaoye]

Thank you very much for your prompt reply. I would like to further ask whether there are any plans to support both episode terminations and multi-environment training simultaneously in TD-MPC2 or NEWT. At the moment, these two features seem to be implemented on separate branches.

[nicklashansen]

It's not something I'm currently planning on adding support for, mostly because it complicates the implementation quite a bit. However I don't see a reason why it wouldn't be possible to do so! The newt codebase already supports multi-task online RL with tasks that have different (fixed) episode lengths, so it should be possible to add support for episode termination on top of that. There's some caveats to the current episode termination support in TD-MPC2 though, see the discussion in issue #62 for details!