https://github.com/google-research/seed_rl/blob/eff7aaa7ab5843547fbf383fcc747b7a8ca67632/common/utils.py#L203