Does ColossalAI support reinforcement learning?

[ArcherShirou]

Could you give an example like DQN in CartPole to illustrate how to use ColossalAI in distributed training for reinforcement learning? thanks

[binmakeswell]

Hi, thanks for your idea. We will provide RL examples as soon as possible, maybe next week.

[number1roy]

Could you give an example like DQN in CartPole to illustrate how to use ColossalAI in distributed training for reinforcement learning? thanks

Hi, I suppose you want to use data-parallel for training DQN. In shallow models like DQN tensor-parallel, pipeline-parallel, and ZeRO would be less effective. Therefore, I write an example (https://github.com/number1roy/colossalai-DQN) that enables you to train DQN in CartPole or Atari Games environments on a single GPU or multiple GPUs. Hope you find it useful!