Generalized Advantage Estimation (GAE) is widely used in RL, especially PPO. The computation of GAE involves a for-loop to iterate over the entire trajectory, which is expensive in Python and may become the training bottleneck.

This repository provides a simple implementation of GAE in CUDA, which can achive at most 2000x higher throughput than Python implementation.

Installation requires a CUDA-enabled GPU with nvcc and torch installed.

git clone https://github.com/garrett4wade/cugae
cd cugae && pip3 install -e .

After installation, run pytest -q -s test_cugae.py to run tests and validate your installation.

See cugae.py for detailed documentation of each implemented function.

This benchmark is performed using Python 3.10.12, CUDA 12.2 in WSL2 Unbuntu 22.04 on a laptop with Intel i7-12700H CPU and Nvidia 3070 GPU.