Modified based on code in https://github.com/openai/spinningup.