diff --git a/.circleci/config.yml b/.circleci/config.yml
index 9c6e7cc4..e4e0ac53 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -153,15 +153,13 @@ commands:
                 steps:
                   - run:
                       command: |
-                        pip install -e .[gym,test]
-                        pip install --pre torch torchvision -f https://download.pytorch.org/whl/nightly/cu101/torch_nightly.html
+                        pip install -e .[gym,test] --pre -f https://download.pytorch.org/whl/nightly/cu101/torch_nightly.html
             - unless:
                 condition: << parameters.is_ubuntu_gpu >>
                 steps:
                   - run:
                       command: |
-                        sudo pip install -e .[gym,test]
-                        sudo pip install --pre torch torchvision -f https://download.pytorch.org/whl/nightly/cu101/torch_nightly.html
+                        sudo pip install -e .[gym,test] --pre -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
 
   run_unittest:
     description: Run unittests, coverage and save results
@@ -172,7 +170,7 @@ commands:
       - run:
           no_output_timeout: 30m
           command: |
-            tox -e << parameters.tox_env >>
+            tox -vv -e << parameters.tox_env >>
             bash <(curl -s https://codecov.io/bash)
       - run: python setup.py bdist_wheel
       - store_artifacts:
diff --git a/pyproject.toml b/pyproject.toml
index bf54bd47..257c07f4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,8 @@
 [build-system]
 requires = [
-    "setuptools >= 35.0.2",
-    "setuptools_scm >= 2.0.0, <3"
+    "setuptools >= 42",
+    "setuptools_scm[toml] >= 3.4",
+    "wheel"
 ]
 build-backend = "setuptools.build_meta"
+[tool.setuptools_scm]
diff --git a/setup.cfg b/setup.cfg
index a1917c9a..9f1eb118 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -24,6 +24,7 @@ install_requires =
   petastorm>=0.9.0
   parameterized>=0.7.4
   pyspark==2.4.6
+  pytorch-lightning
   ruamel.yaml>=0.15.99
   scipy>=1.3.1
   tensorboard>=1.14
@@ -31,15 +32,16 @@ install_requires =
 
 [options.extras_require]
 gym =
-  gym[classic_control,box2d,atari]
+  # Some issue with https://github.com/openai/gym/pull/1974
+  # Remove the pinning when https://github.com/openai/gym/issues/2058 is fixed
+  gym[classic_control,box2d,atari]==0.17.2
   gym_minigrid
   recsim-no-tf
 
 test =
   coverage>=5.1
-  pytest-xdist==1.30.0
-  # Pinning due to https://github.com/pytest-dev/pytest/issues/6925
-  pytest==5.3
+  pytest-xdist>=1.30.0
+  pytest>=5.3
   spark-testing-base==0.10.0
   pytest-cov
 
diff --git a/tox.ini b/tox.ini
index cbce70c5..bbf75818 100644
--- a/tox.ini
+++ b/tox.ini
@@ -5,42 +5,45 @@
 
 [tox]
 envlist = py37
-isolated_build = True
 
 # install CUDA 10.1 Torch
 [ubuntu_gpu]
-commands_pre =
-    pip install --pre torch torchvision -f https://download.pytorch.org/whl/nightly/cu101/torch_nightly.html
+install_command =
+    pip install --pre -f https://download.pytorch.org/whl/nightly/cu101/torch_nightly.html {opts} {packages}
+
+[pytest]
+addopts = --verbose -d --tx popen --cov=reagent --cov-report=xml --cov-append --junitxml={envlogdir}/junit-{envname}.xml
 
 # Refer to https://docs.pytest.org/en/stable/example/markers.html
 # for how we include/exclude tests in pytest
 
 [testenv]
+# Install the latest pip, setuptools, wheel; this is needed for downloading opencv-python wheel,
+# instead of building from source (which is super slow).
+download = true
 extras =
     gym
     test
-setenv =
-    PYTEST_ADDOPTS=--verbose -d --tx popen --cov --cov-report=xml --cov-append --junitxml={envlogdir}/junit-{envname}.xml
-commands_pre =
-    pip install --pre torch torchvision -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
+install_command =
+    pip install --pre -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html {opts} {packages}
 commands =
     pytest -n 4 -m "(not serial) and (not seq2slate_long)"
     pytest -n0 -m "serial"
 
 [testenv:circleci_unittest]
-commands_pre = {[ubuntu_gpu]commands_pre}
+install_command = {[ubuntu_gpu]install_command}
 commands =
     pytest reagent/test -n auto -m "(not serial) and (not seq2slate_long)"
     pytest reagent/test -n0 -m "serial"
 
 [testenv:circleci_gym_unittest]
-commands_pre = {[ubuntu_gpu]commands_pre}
+install_command = {[ubuntu_gpu]install_command}
 commands =
     pytest reagent/gym/tests -n2 -m "(not serial) and (not seq2slate_long)"
     pytest reagent/gym/tests -n0 -m "serial"
 
 
 [testenv:circleci_seq2slate_unittest]
-commands_pre = {[ubuntu_gpu]commands_pre}
+install_command = {[ubuntu_gpu]install_command}
 commands =
     pytest reagent/test -n0 -m "seq2slate_long"