Upgrade PyTorch version for python 3.9 (#5028)

.github/workflows/pytest.yml

         TEST_ENFORCE_BUFFER_KEY_TYPES: 1
     strategy:
       matrix:
-        python-version: [3.6.x, 3.7.x, 3.8.x]
+        python-version: [3.6.x, 3.7.x, 3.8.x, 3.9.x]
+        include:
+          - python-version: 3.6.x
+            pip_constraints: test_constraints_min_version.txt
+          - python-version: 3.7.x
+            pip_constraints: test_constraints_mid_version.txt
+          - python-version: 3.8.x
+            pip_constraints: test_constraints_mid_version.txt
+          - python-version: 3.9.x
+            pip_constraints: test_constraints_max_version.txt
     steps:
     - uses: actions/checkout@v2
     - name: Set up Python
         # This path is specific to Ubuntu
         path: ~/.cache/pip
         # Look to see if there is a cache hit for the corresponding requirements file
-        key: ${{ runner.os }}-pip-${{ hashFiles('ml-agents/setup.py', 'ml-agents-envs/setup.py', 'gym-unity/setup.py', 'test_requirements.txt') }}
+        key: ${{ runner.os }}-pip-${{ hashFiles('ml-agents/setup.py', 'ml-agents-envs/setup.py', 'gym-unity/setup.py', 'test_requirements.txt', matrix.pip_constraints) }}
         restore-keys: |
           ${{ runner.os }}-pip-
           ${{ runner.os }}-
       run: |
         python -m pip install --upgrade pip
         python -m pip install --upgrade setuptools
-        python -m pip install --progress-bar=off -e ./ml-agents-envs
-        python -m pip install --progress-bar=off -e ./ml-agents
-        python -m pip install --progress-bar=off -r test_requirements.txt
-        python -m pip install --progress-bar=off -e ./gym-unity
-        python -m pip install --progress-bar=off -e ./ml-agents-plugin-examples
+        python -m pip install --progress-bar=off -e ./ml-agents-envs -c ${{ matrix.pip_constraints }}
+        python -m pip install --progress-bar=off -e ./ml-agents -c ${{ matrix.pip_constraints }}
+        python -m pip install --progress-bar=off -r test_requirements.txt -c ${{ matrix.pip_constraints }}
+        python -m pip install --progress-bar=off -e ./gym-unity -c ${{ matrix.pip_constraints }}
+        python -m pip install --progress-bar=off -e ./ml-agents-plugin-examples -c ${{ matrix.pip_constraints }}
     - name: Save python dependencies
       run: |
         pip freeze > pip_versions-${{ matrix.python-version }}.txt

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

         [BRAIN_NAME], num_visual=num_visual, num_vector=0, action_sizes=(1, 1)
     )
     new_hyperparams = attr.evolve(
-        PPO_TORCH_CONFIG.hyperparameters, learning_rate=3.0e-4
+        PPO_TORCH_CONFIG.hyperparameters,
+        batch_size=64,
+        buffer_size=1024,
+        learning_rate=1e-4,
+    )
+    config = attr.evolve(
+        PPO_TORCH_CONFIG, hyperparameters=new_hyperparams, max_steps=8000
-    config = attr.evolve(PPO_TORCH_CONFIG, hyperparameters=new_hyperparams)
     check_environment_trains(env, {BRAIN_NAME: config}, training_seed=training_seed)
 
 
         buffer_init_steps=0,
     )
     config = attr.evolve(
-        SAC_TORCH_CONFIG, hyperparameters=new_hyperparams, max_steps=2200
+        SAC_TORCH_CONFIG, hyperparameters=new_hyperparams, max_steps=6000
     )
     check_environment_trains(env, {BRAIN_NAME: config}, success_threshold=0.9)
 
     new_hyperparams = attr.evolve(
         SAC_TORCH_CONFIG.hyperparameters,
         batch_size=256,
-        learning_rate=1e-3,
+        learning_rate=3e-4,
         buffer_init_steps=1000,
         steps_per_update=2,
     )
         network_settings=new_networksettings,
-        max_steps=3500,
+        max_steps=4000,
     )
     check_environment_trains(env, {BRAIN_NAME: config}, training_seed=1212)

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

     )
     bc_settings = BehavioralCloningSettings(demo_path=demo_path, steps=1500)
     reward_signals = {
-        RewardSignalType.GAIL: GAILSettings(encoding_size=32, demo_path=demo_path)
+        RewardSignalType.GAIL: GAILSettings(
+            gamma=0.8, encoding_size=32, demo_path=demo_path
+        )
-    hyperparams = attr.evolve(PPO_TORCH_CONFIG.hyperparameters, learning_rate=5e-3)
+    hyperparams = attr.evolve(PPO_TORCH_CONFIG.hyperparameters, learning_rate=1e-3)
     config = attr.evolve(
         PPO_TORCH_CONFIG,
         reward_signals=reward_signals,

ml-agents/setup.py

         "pyyaml>=3.1.0",
         # Windows ver. of PyTorch doesn't work from PyPi. Installation:
         # https://github.com/Unity-Technologies/ml-agents/blob/main/docs/Installation.md#windows-installing-pytorch
-        'torch>=1.6.0,<1.8.0;platform_system!="Windows"',
+        # Torch only working on python 3.9 for 1.8.0 and above. Details see:
+        # https://github.com/pytorch/pytorch/issues/50014
+        "torch>=1.8.0,<1.9.0;(platform_system!='Windows' and python_version>='3.9')",
+        "torch>=1.6.0,<1.9.0;(platform_system!='Windows' and python_version<'3.9')",
         "tensorboard>=1.15",
         # cattrs 1.1.0 dropped support for python 3.6, but 1.0.0 doesn't work for python 3.9
         # Since there's no version that supports both, we have to draw the line somwehere.

test_constraints_max_version.txt

+# pip constraints to use the *highest* versions allowed in ml-agents/setup.py
+# For projects with upper bounds, we should periodically update this list to the latest
+torch==1.8.0

test_constraints_mid_version.txt

+# pip constraints to use a version in the middle of allowed ranges in ml-agents/setup.py
+torch==1.7.0

test_constraints_min_version.txt

+# pip constraints to use the *lowest* versions allowed in ml-agents/setup.py
+torch==1.6.0