Run pytest on GPU (#4865)

* make tests device-friendly

* mark all tests in test_simple_rl

ml-agents/mlagents/trainers/tests/torch/saver/test_saver.py

 import os
 
 import numpy as np
-from mlagents.torch_utils import torch
+from mlagents.torch_utils import torch, default_device
 from mlagents.trainers.policy.torch_policy import TorchPolicy
 from mlagents.trainers.ppo.optimizer_torch import TorchPPOOptimizer
 from mlagents.trainers.model_saver.torch_model_saver import TorchModelSaver
     """
     Make sure two policies have the same output for the same input.
     """
+    policy1.actor_critic = policy1.actor_critic.to(default_device())
+    policy2.actor_critic = policy2.actor_critic.to(default_device())
+
     decision_step, _ = mb.create_steps_from_behavior_spec(
         policy1.behavior_spec, num_agents=1
     )
             tensor_obs, masks=masks, memories=memories
         )
     np.testing.assert_array_equal(
-        log_probs1.all_discrete_tensor, log_probs2.all_discrete_tensor
+        ModelUtils.to_numpy(log_probs1.all_discrete_tensor),
+        ModelUtils.to_numpy(log_probs2.all_discrete_tensor),
     )
 
 

ml-agents/mlagents/trainers/tests/torch/test_action_model.py

     for _disc in log_probs.all_discrete_list:
         assert _disc.shape == (1, 2)
 
-    for clp in log_probs.continuous_tensor[0]:
+    for clp in log_probs.continuous_tensor[0].tolist():
-    for ent, val in zip(entropies[0], [1.4189, 0.6191, 0.6191]):
+    for ent, val in zip(entropies[0].tolist(), [1.4189, 0.6191, 0.6191]):
         assert ent == pytest.approx(val, abs=0.01)

ml-agents/mlagents/trainers/tests/torch/test_distributions.py

         optimizer.zero_grad()
         loss.backward()
         optimizer.step()
-    for prob in log_prob.flatten():
+    for prob in log_prob.flatten().tolist():
         assert prob == pytest.approx(-2, abs=0.1)
 
 
     dist_insts = gauss_dist(sample_embedding, masks=masks)
     for dist_inst in dist_insts:
         log_prob = dist_inst.all_log_prob()
-        assert log_prob.flatten()[-1] == pytest.approx(0, abs=0.001)
+        assert log_prob.flatten()[-1].tolist() == pytest.approx(0, abs=0.001)
 
 
 def test_gaussian_dist_instance():
     )
     action = dist_instance.sample()
     assert action.shape == (1, act_size)
-    for log_prob in dist_instance.log_prob(torch.zeros((1, act_size))).flatten():
+    for log_prob in (
+        dist_instance.log_prob(torch.zeros((1, act_size))).flatten().tolist()
+    ):
-    for ent in dist_instance.entropy().flatten():
+    for ent in dist_instance.entropy().flatten().tolist():
         # entropy of standard normal at 0, based on 1/2 + ln(sqrt(2pi)sigma)
         assert ent == pytest.approx(1.42, abs=0.01)
 

ml-agents/mlagents/trainers/tests/torch/test_encoders.py

     norm.update(vec_input3)
 
     # Test normalization
-    for val in norm(vec_input1)[0]:
+    for val in norm(vec_input1)[0].tolist():
         assert val == pytest.approx(0.707, abs=0.001)
 
     # Test copy normalization
     assert compare_models(norm, norm2)
-    for val in norm2(vec_input1)[0]:
+    for val in norm2(vec_input1)[0].tolist():
         assert val == pytest.approx(0.707, abs=0.001)
 
 

ml-agents/mlagents/trainers/tests/torch/test_hybrid.py

 SAC_TORCH_CONFIG = sac_dummy_config()
 
 
+@pytest.mark.check_environment_trains
 @pytest.mark.parametrize("action_size", [(1, 1), (2, 2), (1, 2), (2, 1)])
 def test_hybrid_ppo(action_size):
     env = SimpleEnvironment([BRAIN_NAME], action_sizes=action_size, step_size=0.8)
     check_environment_trains(env, {BRAIN_NAME: config}, success_threshold=0.9)
 
 
+@pytest.mark.check_environment_trains
 @pytest.mark.parametrize("num_visual", [1, 2])
 def test_hybrid_visual_ppo(num_visual):
     env = SimpleEnvironment(
     check_environment_trains(env, {BRAIN_NAME: config}, training_seed=1336)
 
 
+@pytest.mark.check_environment_trains
 def test_hybrid_recurrent_ppo():
     env = MemoryEnvironment([BRAIN_NAME], action_sizes=(1, 1), step_size=0.5)
     new_network_settings = attr.evolve(
     check_environment_trains(env, {BRAIN_NAME: config}, success_threshold=0.9)
 
 
+@pytest.mark.check_environment_trains
 @pytest.mark.parametrize("action_size", [(1, 1), (2, 2), (1, 2), (2, 1)])
 def test_hybrid_sac(action_size):
     env = SimpleEnvironment([BRAIN_NAME], action_sizes=action_size, step_size=0.8)
     )
 
 
+@pytest.mark.check_environment_trains
 @pytest.mark.parametrize("num_visual", [1, 2])
 def test_hybrid_visual_sac(num_visual):
     env = SimpleEnvironment(
     check_environment_trains(env, {BRAIN_NAME: config})
 
 
+@pytest.mark.check_environment_trains
 def test_hybrid_recurrent_sac():
     env = MemoryEnvironment([BRAIN_NAME], action_sizes=(1, 1), step_size=0.5)
     new_networksettings = attr.evolve(

ml-agents/mlagents/trainers/tests/torch/test_networks.py

         loss.backward()
         optimizer.step()
     # In the last step, values should be close to 1
-    for _enc in encoded.flatten():
+    for _enc in encoded.flatten().tolist():
         assert _enc == pytest.approx(1.0, abs=0.1)
 
 
         loss.backward()
         optimizer.step()
     # In the last step, values should be close to 1
-    for _enc in encoded.flatten():
+    for _enc in encoded.flatten().tolist():
         assert _enc == pytest.approx(1.0, abs=0.1)
 
 
         loss.backward()
         optimizer.step()
     # In the last step, values should be close to 1
-    for _enc in encoded.flatten():
+    for _enc in encoded.flatten().tolist():
         assert _enc == pytest.approx(1.0, abs=0.1)
 
 
         optimizer.step()
     # In the last step, values should be close to 1
     for value in values.values():
-        for _out in value:
+        for _out in value.tolist():
             assert _out[0] == pytest.approx(1.0, abs=0.1)
 
 

ml-agents/mlagents/trainers/tests/torch/test_simple_rl.py

 PPO_TORCH_CONFIG = ppo_dummy_config()
 SAC_TORCH_CONFIG = sac_dummy_config()
 
+# tests in this file won't be tested on GPU machine
+pytestmark = pytest.mark.check_environment_trains
+
 
 @pytest.mark.parametrize("action_sizes", [(0, 1), (1, 0)])
 def test_simple_ppo(action_sizes):

ml-agents/mlagents/trainers/torch/encoders.py

         if not exporting_to_onnx.is_exporting():
             visual_obs = visual_obs.permute([0, 3, 1, 2])
         hidden = self.conv_layers(visual_obs)
-        hidden = torch.reshape(hidden, (-1, self.final_flat))
+        hidden = hidden.reshape(-1, self.final_flat)
         return self.dense(hidden)
 
 
         if not exporting_to_onnx.is_exporting():
             visual_obs = visual_obs.permute([0, 3, 1, 2])
         hidden = self.conv_layers(visual_obs)
-        hidden = torch.reshape(hidden, (-1, self.final_flat))
+        hidden = hidden.reshape(-1, self.final_flat)
         return self.dense(hidden)
 
 

.yamato/pytest-gpu.yml

+pytest_gpu:
+  name: Pytest GPU
+  agent:
+    type: Unity::VM::GPU
+    image: package-ci/ubuntu:stable
+    flavor: b1.large
+  commands:
+    - |
+      sudo apt-get update && sudo apt-get install -y python3-venv
+      python3 -m venv venv && source venv/bin/activate
+      python3 -m pip install pyyaml --index-url https://artifactory.prd.it.unity3d.com/artifactory/api/pypi/pypi/simple
+      python3 -u -m ml-agents.tests.yamato.setup_venv
+      python3 -m pip install --progress-bar=off -r test_requirements.txt --index-url https://artifactory.prd.it.unity3d.com/artifactory/api/pypi/pypi/simple
+      python3 -m pip install torch==1.7.1+cu101 torchvision==0.8.2+cu101 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html --index-url https://artifactory.prd.it.unity3d.com/artifactory/api/pypi/pypi/simple
+      python3 -m pytest -m "not check_environment_trains" --junitxml=junit/test-results.xml -p no:warnings
+  triggers:
+    cancel_old_ci: true
+    recurring:
+      - branch: master
+        frequency: daily
+  artifacts:
+    logs:
+      paths:
+        - "artifacts/standalone_build.txt"

pytest.ini

+[pytest]
+addopts = --strict-markers
+markers =
+    check_environment_trains: Slow training tests, do not run on yamato