google-deepmind / dqn_zoo

DQN Zoo is a collection of reference implementations of reinforcement learning agents developed at DeepMind based on the Deep Q-Network (DQN) agent.
Apache License 2.0
451 stars 78 forks source link

Running on CPU May Require GPU? #4

Closed RylanSchaeffer closed 3 years ago

RylanSchaeffer commented 4 years ago

Continuing my comments from Issue 2, it appears that the requirements

jax==0.1.72
jaxlib @ https://storage.googleapis.com/jax-releases/cuda101/jaxlib-0.1.49-cp36-none-linux_x86_64.whl

means that a GPU is required to run the code even with the flag --jax_platform_name=cpu.

jqdm commented 3 years ago

I've clarified in f6ea040cb5e55bf65da6cfa07d4daf8225245fa9 that dependencies may have to be adapted to suit the user's set up. In this case, if the user does not have a GPU then they can install jaxlib without GPU support.