Skip to content

[bug-fix] Move POCA critic to default device #5124

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

Merged
merged 6 commits into from
Mar 16, 2021
Merged

[bug-fix] Move POCA critic to default device #5124

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
b79d390
Move critic to default device
Mar 16, 2021
9c51601
Make sure to clone onto default device
Mar 16, 2021
deb0590
Add some debug stuff
Mar 16, 2021
5a1987c
Some more debug
Mar 16, 2021
dca2e41
Fix issue
Mar 16, 2021
b602f30
Fix bool tensor too
Mar 16, 2021
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
4 changes: 3 additions & 1 deletion ml-agents/mlagents/trainers/poca/optimizer_torch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
)
import numpy as np
import math
from mlagents.torch_utils import torch
from mlagents.torch_utils import torch, default_device

from mlagents.trainers.buffer import (
AgentBuffer,
Expand Down Expand Up @@ -155,6 +155,8 @@ def __init__(self, policy: TorchPolicy, trainer_settings: TrainerSettings):
network_settings=trainer_settings.network_settings,
action_spec=policy.behavior_spec.action_spec,
)
# Move to GPU if needed
self._critic.to(default_device())

params = list(self.policy.actor.parameters()) + list(self.critic.parameters())
self.hyperparameters: POCASettings = cast(
Expand Down
4 changes: 2 additions & 2 deletions ml-agents/mlagents/trainers/torch/networks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -268,7 +268,7 @@ def _get_masks_from_nans(self, obs_tensors: List[torch.Tensor]) -> torch.Tensor:
[_obs.flatten(start_dim=1)[:, 0] for _obs in only_first_obs], dim=1
)
# Get the mask from NaNs
attn_mask = only_first_obs_flat.isnan().type(torch.FloatTensor)
attn_mask = only_first_obs_flat.isnan().float()
return attn_mask

def _copy_and_remove_nans_from_obs(
Expand All @@ -283,7 +283,7 @@ def _copy_and_remove_nans_from_obs(
for obs in single_agent_obs:
new_obs = obs.clone()
new_obs[
attention_mask.type(torch.BoolTensor)[:, i_agent], ::
attention_mask.bool()[:, i_agent], ::
] = 0.0 # Remoove NaNs fast
no_nan_obs.append(new_obs)
obs_with_no_nans.append(no_nan_obs)
Expand Down