wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-28 11:25:04 +08:00
Code Issues Packages Projects Releases Wiki Activity

Small fix for Cuda Torch DQN. (#10177)

Browse Source
This commit is contained in:
Raphael Avalos
2020-08-19 13:28:05 -07:00
committed by GitHub
parent b70dce0d02
commit 8b704eb419
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rllib/agents/dqn/dqn_torch_policy.py
+4 -2
View File
@@ -215,7 +215,8 @@ def build_q_losses(policy, model, _, train_batch):
one_hot_selection = F.one_hot(train_batch[SampleBatch.ACTIONS],
policy.action_space.n)
q_t_selected = torch.sum(
torch.where(q_t > -float("inf"), q_t, torch.tensor(0.0)) *
torch.where(q_t > -float("inf"), q_t,
torch.tensor(0.0, device=policy.device)) *
one_hot_selection, 1)
q_logits_t_selected = torch.sum(
q_logits_t * torch.unsqueeze(one_hot_selection, -1), 1)
@@ -233,7 +234,8 @@ def build_q_losses(policy, model, _, train_batch):
q_tp1_best_one_hot_selection = F.one_hot(q_tp1_best_using_online_net,
policy.action_space.n)
q_tp1_best = torch.sum(
torch.where(q_tp1 > -float("inf"), q_tp1, torch.tensor(0.0)) *
torch.where(q_tp1 > -float("inf"), q_tp1,
torch.tensor(0.0, device=policy.device)) *
q_tp1_best_one_hot_selection, 1)
q_probs_tp1_best = torch.sum(
q_probs_tp1 * torch.unsqueeze(q_tp1_best_one_hot_selection, -1), 1)