[P] Atari Learning Environment on the GPU (CuLE)
|
I recently spent some time modernizing NVIDIA’s CuLE (2019), the GPU-native Atari environment. It was tied to an older software stack (CUDA 10, The fork now builds with CUDA 12.x, PyTorch 2.x, Python 3.12, Gymnasium, and ale-py (so ROMs come bundled). I also compared it against EnvPool on my machine (RTX 4090 + i5-13600K). For PPO, CuLE reaches about 38k SPS vs EnvPool’s 21k (~1.8× faster). For raw environment stepping, the crossover is around 512 environments—below that EnvPool is faster, but above it CuLE continues to scale while EnvPool plateaus, reaching up to ~114k steps/s at 4,096 Breakout environments (~3× faster on pure stepping). To make benchmarking easier, I also added training scripts based on CleanRL and LeanRL (PPO, DQN, Rainbow, C51, PQN, SAC), alongside the original CuLE implementations. I’d be interested to hear if anyone here is still using CuLE or GPU-native Atari environments, or if there are additional benchmarks you’d like to see. Repository: submitted by /u/mehrdad96 |