Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve DQN Tutorial #2934

Merged
merged 4 commits into from
Jun 18, 2024
Merged

Improve DQN Tutorial #2934

Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
2ed58af
Add MPS support and original paper
alperenunlu Jun 16, 2024
5eb92fd
Merge branch 'main' into improve_dqn_tutorial
svekars Jun 17, 2024
39b20e7
Merge branch 'main' into improve_dqn_tutorial
svekars Jun 18, 2024
42fe906
Merge branch 'main' into improve_dqn_tutorial
svekars Jun 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 8 additions & 2 deletions intermediate_source/reinforcement_q_learning.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@
This tutorial shows how to use PyTorch to train a Deep Q Learning (DQN) agent
on the CartPole-v1 task from `Gymnasium <https://gymnasium.farama.org>`__.

You might find it helpful to read the original `Deep Q Learning (DQN) <https://arxiv.org/abs/1312.5602>`__ paper

**Task**

The agent has to decide between two actions - moving the cart left or
Expand Down Expand Up @@ -83,7 +85,11 @@
plt.ion()

# if GPU is to be used
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = torch.device(
"cuda" if torch.cuda.is_available() else
"mps" if torch.backends.mps.is_available() else
"cpu"
)


######################################################################
Expand Down Expand Up @@ -397,7 +403,7 @@ def optimize_model():
# can produce better results if convergence is not observed.
#

if torch.cuda.is_available():
if torch.cuda.is_available() or torch.backends.mps.is_available():
num_episodes = 600
else:
num_episodes = 50
Expand Down
Loading