[bug-fix] Move POCA critic to default device

[ervteng]

Proposed change(s)

Move POCA critic to default device. Was failing Yamato GPU test without this. Needs to be cherry-picked into R15.

Types of change(s)

• Bug fix
• New feature
• Code refactor
• Breaking change
• Documentation update
• Other (please describe)

Checklist

• Added tests that prove my fix is effective or that my feature works
• Updated the changelog (if applicable)
• Updated the documentation (if applicable)
• Updated the migration guide (if applicable)

Other comments

[vincentpierre]

Do you manage to train on GPU after this fix ?

[ervteng]

Do you manage to train on GPU after this fix ?

It's still failing: (Yamato job here: https://yamato.cds.internal.unity3d.com/jobs/497-ml-agents/tree/develop-poca-gpu/.yamato%252Fpytest-gpu.yml%2523pytest_gpu/5821068/job) need to debug

[dongruoping]

One thing I noticed in networks.py:line 271
attn_mask = only_first_obs_flat.isnan().type(torch.FloatTensor)
this will move tensors on cpu.

Since we're using default tensor type we should avoid calling torch.FloatTensor and use .float() instead.
I've fixed several once before so I figured I should explicitly call this out.

[dongruoping]

Do you think we should enable gpu tests in all python PRs? I think we've seen couple times this fails in nightly CI and would be nice to catch them earlier.

[ervteng]

@dongruoping yep just caught that, the fix works 👍

[ervteng]

Do you think we should enable gpu tests in all python PRs? I think we've seen couple times this fails in nightly CI and would be nice to catch them earlier.

I agree. Spoke with @surfnerd about this; I've added it to the Release Checklist. I guess we have to weight the cost of spinning up a GPU machine vs. catching these earlier. Not sure what the tradeoff is.

[dongruoping]

@dongruoping yep just caught that, the fix works 👍

Can you also fix the torch.BoolTensor in the same file?