[Feature] VC1 integration (pytorch#1211)

ai4co · Jul 7, 2023 · 002a58a · 002a58a
1 parent 3e512ff
commit 002a58a
Show file tree

Hide file tree

Showing 14 changed files with 611 additions and 23 deletions.
diff --git a/.circleci/unittest/linux/scripts/install.sh b/.circleci/unittest/linux/scripts/install.sh
@@ -6,6 +6,7 @@ unset PYTORCH_VERSION
 # In fact, keeping PYTORCH_VERSION forces us to hardcode PyTorch version in config.
 
 set -e
+set -v
 
 eval "$(./conda/bin/conda shell.bash hook)"
 conda activate ./env
@@ -34,13 +35,13 @@ else
 fi
 
 # smoke test
-python -c "import functorch"
+python3 -c "import functorch"
 
 # install snapshot
-pip install git+https://github.com/pytorch/torchsnapshot
+pip3 install git+https://github.com/pytorch/torchsnapshot
 
 # install tensordict
-pip install git+https://github.com/pytorch-labs/tensordict.git
+pip3 install git+https://github.com/pytorch-labs/tensordict.git
 
 printf "* Installing torchrl\n"
-python setup.py develop
+python3 setup.py develop
diff --git a/.circleci/unittest/linux/scripts/run_all.sh b/.circleci/unittest/linux/scripts/run_all.sh
@@ -138,6 +138,21 @@ pip3 install git+https://github.com/pytorch-labs/tensordict.git
 printf "* Installing torchrl\n"
 python setup.py develop
 
+
+if [ "${CU_VERSION:-}" != cpu ] ; then
+  printf "* Installing VC1\n"
+  python3 -c """
+from torchrl.envs.transforms.vc1 import VC1Transform
+VC1Transform.install_vc_models(auto_exit=True)
+"""
+
+  python3 -c """
+import vc_models
+from vc_models.models.vit import model_utils
+print(model_utils)
+"""
+fi
+
 # ==================================================================================== #
 # ================================ Run tests ========================================= #
 
@@ -151,6 +166,7 @@ python -m torch.utils.collect_env
 export MKL_THREADING_LAYER=GNU
 export CKPT_BACKEND=torch
 
+
 pytest test/smoke_test.py -v --durations 200
 pytest test/smoke_test_deps.py -v --durations 200 -k 'test_gym or test_dm_control_pixels or test_dm_control or test_tb'
 python .circleci/unittest/helpers/coverage_run_parallel.py -m pytest test \

diff --git a/.circleci/unittest/linux/scripts/run_test.sh b/.circleci/unittest/linux/scripts/run_test.sh
@@ -1,12 +1,13 @@
 #!/usr/bin/env bash
 
 set -e
+set -v
 
 eval "$(./conda/bin/conda shell.bash hook)"
 conda activate ./env
 
 export PYTORCH_TEST_WITH_SLOW='1'
-python -m torch.utils.collect_env
+python3 -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
 
@@ -19,8 +20,21 @@ export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$lib_dir
 export MKL_THREADING_LAYER=GNU
 export CKPT_BACKEND=torch
 
-python .circleci/unittest/helpers/coverage_run_parallel.py -m pytest test/smoke_test.py -v --durations 200
-python .circleci/unittest/helpers/coverage_run_parallel.py -m pytest test/smoke_test_deps.py -v --durations 200 -k 'test_gym or test_dm_control_pixels or test_dm_control or test_tb'
-python .circleci/unittest/helpers/coverage_run_parallel.py -m pytest --instafail -v --durations 200 --ignore test/test_distributed.py --ignore test/test_rlhf.py
+# install vc1
+python3 -c """
+from torchrl.envs.transforms.vc1 import VC1Transform
+VC1Transform.install_vc_models(auto_exit=True)
+"""
+
+python3 -c """
+import vc_models
+from vc_models.models.vit import model_utils
+print(model_utils)
+"""
+
+python3 .circleci/unittest/helpers/coverage_run_parallel.py -m pytest test/smoke_test.py -v --durations 200
+python3 .circleci/unittest/helpers/coverage_run_parallel.py -m pytest test/smoke_test_deps.py -v --durations 200 -k 'test_gym or test_dm_control_pixels or test_dm_control or test_tb'
+python3 .circleci/unittest/helpers/coverage_run_parallel.py -m pytest --instafail -v --durations 200 --ignore test/test_distributed.py --ignore test/test_rlhf.py -k VC1
+python3 .circleci/unittest/helpers/coverage_run_parallel.py -m pytest --instafail -v --durations 200 --ignore test/test_distributed.py --ignore test/test_rlhf.py33
 coverage combine
 coverage xml -i
diff --git a/docs/source/reference/data.rst b/docs/source/reference/data.rst
@@ -42,6 +42,7 @@ We also give users the ability to compose a replay buffer using the following co
     TensorStorage
     Writer
     RoundRobinWriter
+    TensorDictRoundRobinWriter
 
 Storage choice is very influential on replay buffer sampling latency, especially in distributed reinforcement learning settings with larger data volumes.
 :class:`LazyMemmapStorage` is highly advised in distributed settings with shared storage due to the lower serialisation cost of MemmapTensors as well as the ability to specify file storage locations for improved node failure recovery.

diff --git a/docs/source/reference/envs.rst b/docs/source/reference/envs.rst
@@ -354,6 +354,7 @@ to be able to create this other composition:
     ToTensorImage
     UnsqueezeTransform
     VecNorm
+    VC1Transform
     VIPRewardTransform
     VIPTransform
 

diff --git a/setup.py b/setup.py
@@ -230,6 +230,7 @@ def _main(argv):
                 "tqdm",
                 "hydra-core>=1.1",
                 "hydra-submitit-launcher",
+                "git",
             ],
             "checkpointing": [
                 "torchsnapshot",