Merge remote-tracking branch 'origin/main' into fix_copy_clone_transform

.github/pytorch-probot.yml

@@ -0,0 +1,5 @@
+# List of workflows that will be re-run in case of failures
+# https://github.com/pytorch/test-infra/blob/main/torchci/lib/bot/retryBot.ts
+retryable_workflows:
+- Build M1
+- Wheels

.github/scripts/m1_script.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+export BUILD_VERSION=0.2.1

.github/scripts/pre_build_script_m1.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+python3 -mpip install git+https://github.com/pytorch/tensordict.git

.github/unittest/linux/scripts/environment.yml

@@ -16,6 +16,7 @@ dependencies:
     - pytest-mock
     - pytest-instafail
     - pytest-rerunfailures
+    - pytest-timeout
     - expecttest
     - pyyaml
     - scipy
@@ -27,6 +28,6 @@ dependencies:
     - mlflow
     - av
     - coverage
-    - ray
+    - ray<2.8.0
     - transformers
     - ninja

.github/unittest/linux/scripts/run_all.sh

@@ -124,9 +124,9 @@ git submodule sync && git submodule update --init --recursive
 printf "Installing PyTorch with %s\n" "${CU_VERSION}"
 if [[ "$TORCH_VERSION" == "nightly" ]]; then
   if [ "${CU_VERSION:-}" == cpu ] ; then
-      pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+      pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cpu
   else
-      pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/$CU_VERSION
+      pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/$CU_VERSION
   fi
 elif [[ "$TORCH_VERSION" == "stable" ]]; then
     if [ "${CU_VERSION:-}" == cpu ] ; then
@@ -146,7 +146,7 @@ python -c "import functorch"
 pip3 install git+https://github.com/pytorch/torchsnapshot
 
 # install tensordict
-pip3 install git+https://github.com/pytorch-labs/tensordict.git
+pip3 install git+https://github.com/pytorch/tensordict.git
 
 printf "* Installing torchrl\n"
 python setup.py develop
@@ -184,10 +184,12 @@ 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'
 if [ "${CU_VERSION:-}" != cpu ] ; then
   python .github/unittest/helpers/coverage_run_parallel.py -m pytest test \
-    --instafail --durations 200 -vv --capture no --ignore test/test_rlhf.py
+    --instafail --durations 200 -vv --capture no --ignore test/test_rlhf.py \
+    --timeout=120
 else
   python .github/unittest/helpers/coverage_run_parallel.py -m pytest test \
-    --instafail --durations 200 -vv --capture no --ignore test/test_rlhf.py --ignore test/test_distributed.py
+    --instafail --durations 200 -vv --capture no --ignore test/test_rlhf.py --ignore test/test_distributed.py \
+    --timeout=120
 fi
 
 coverage combine

.github/unittest/linux_distributed/scripts/environment.yml

@@ -27,5 +27,5 @@ dependencies:
     - mlflow
     - av
     - coverage
-    - ray
+    - ray<2.8.0
     - virtualenv

.github/unittest/linux_distributed/scripts/install.sh

@@ -28,9 +28,9 @@ git submodule sync && git submodule update --init --recursive
 
 printf "Installing PyTorch with %s\n" "${CU_VERSION}"
 if [ "${CU_VERSION:-}" == cpu ] ; then
-    pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+    pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cpu
 else
-    pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/$CU_VERSION
+    pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/$CU_VERSION
 fi
 
 # smoke test
@@ -40,7 +40,7 @@ python -c "import functorch"
 pip install git+https://github.com/pytorch/torchsnapshot
 
 # install tensordict
-pip install git+https://github.com/pytorch-labs/tensordict.git
+pip install git+https://github.com/pytorch/tensordict.git
 
 printf "* Installing torchrl\n"
 python setup.py develop

.github/unittest/linux_examples/scripts/run_all.sh

@@ -146,7 +146,7 @@ version="$(python -c "print('.'.join(\"${CUDA_VERSION}\".split('.')[:2]))")"
 git submodule sync && git submodule update --init --recursive
 
 printf "Installing PyTorch with %s\n" "${CU_VERSION}"
-pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/$CU_VERSION
+pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/$CU_VERSION
 
 # smoke test
 python -c "import functorch"
@@ -155,7 +155,7 @@ python -c "import functorch"
 pip install git+https://github.com/pytorch/torchsnapshot
 
 # install tensordict
-pip install git+https://github.com/pytorch-labs/tensordict.git
+pip install git+https://github.com/pytorch/tensordict.git
 
 printf "* Installing torchrl\n"
 python setup.py develop

.github/unittest/linux_examples/scripts/run_test.sh

@@ -37,17 +37,27 @@ python .github/unittest/helpers/coverage_run_parallel.py examples/decision_trans
   optim.updates_per_episode=3 \
   optim.warmup_steps=10 \
   optim.device=cuda:0 \
-  logger.backend=
+  logger.backend= \
+  env.backend=gymnasium \
+  env.name=HalfCheetah-v4
 python .github/unittest/helpers/coverage_run_parallel.py examples/decision_transformer/online_dt.py \
   optim.pretrain_gradient_steps=55 \
   optim.updates_per_episode=3 \
   optim.warmup_steps=10 \
   optim.device=cuda:0 \
-  logger.backend=
+  logger.backend= \
+  env.backend=gymnasium \
+  env.name=HalfCheetah-v4
 
 # ==================================================================================== #
 # ================================ Gymnasium ========================================= #
 
+python .github/unittest/helpers/coverage_run_parallel.py examples/impala/impala_single_node.py \
+  collector.total_frames=80 \
+  collector.frames_per_batch=20 \
+  collector.num_workers=1 \
+  logger.backend= \
+  logger.test_interval=10
 python .github/unittest/helpers/coverage_run_parallel.py examples/ppo/ppo_mujoco.py \
   env.env_name=HalfCheetah-v4 \
   collector.total_frames=40 \
@@ -66,13 +76,12 @@ python .github/unittest/helpers/coverage_run_parallel.py examples/ppo/ppo_atari.
 python .github/unittest/helpers/coverage_run_parallel.py examples/ddpg/ddpg.py \
   collector.total_frames=48 \
   collector.init_random_frames=10 \
-  optimization.batch_size=10 \
+  optim.batch_size=10 \
   collector.frames_per_batch=16 \
-  collector.num_workers=4 \
   collector.env_per_collector=2 \
-  collector.collector_device=cuda:0 \
+  collector.device=cuda:0 \
   network.device=cuda:0 \
-  optimization.utd_ratio=1 \
+  optim.utd_ratio=1 \
   replay_buffer.size=120 \
   env.name=Pendulum-v1 \
   logger.backend=
@@ -103,27 +112,37 @@ python .github/unittest/helpers/coverage_run_parallel.py examples/dqn/dqn.py \
   record_video=True \
   record_frames=4 \
   buffer_size=120
+python .github/unittest/helpers/coverage_run_parallel.py examples/cql/discrete_cql_online.py \
+  collector.total_frames=48 \
+  collector.init_random_frames=10 \
+  optim.batch_size=10 \
+  collector.frames_per_batch=16 \
+  collector.env_per_collector=2 \
+  collector.device=cuda:0 \
+  optim.optim_steps_per_batch=1 \
+  replay_buffer.size=120 \
+  logger.backend=
 python .github/unittest/helpers/coverage_run_parallel.py examples/redq/redq.py \
-  total_frames=48 \
-  init_random_frames=10 \
-  batch_size=10 \
-  frames_per_batch=16 \
   num_workers=4 \
-  env_per_collector=2 \
-  collector_device=cuda:0 \
-  optim_steps_per_batch=1 \
-  record_video=True \
-  record_frames=4 \
-  buffer_size=120
+  collector.total_frames=48 \
+  collector.init_random_frames=10 \
+  collector.frames_per_batch=16 \
+  collector.env_per_collector=2 \
+  collector.device=cuda:0 \
+  buffer.batch_size=10 \
+  optim.steps_per_batch=1 \
+  logger.record_video=True \
+  logger.record_frames=4 \
+  buffer.size=120 \
+  logger.backend=
 python .github/unittest/helpers/coverage_run_parallel.py examples/sac/sac.py \
   collector.total_frames=48 \
   collector.init_random_frames=10 \
   collector.frames_per_batch=16 \
-  collector.num_workers=4 \
   collector.env_per_collector=2 \
-  collector.collector_device=cuda:0 \
-  optimization.batch_size=10 \
-  optimization.utd_ratio=1 \
+  collector.device=cuda:0 \
+  optim.batch_size=10 \
+  optim.utd_ratio=1 \
   replay_buffer.size=120 \
   env.name=Pendulum-v1 \
   network.device=cuda:0 \
@@ -138,6 +157,7 @@ python .github/unittest/helpers/coverage_run_parallel.py examples/dreamer/dreame
   num_workers=4 \
   env_per_collector=2 \
   collector_device=cuda:0 \
+  model_device=cuda:0 \
   optim_steps_per_batch=1 \
   record_video=True \
   record_frames=4 \
@@ -146,25 +166,25 @@ python .github/unittest/helpers/coverage_run_parallel.py examples/dreamer/dreame
 python .github/unittest/helpers/coverage_run_parallel.py examples/td3/td3.py \
   collector.total_frames=48 \
   collector.init_random_frames=10 \
-  optimization.batch_size=10 \
+  optim.batch_size=10 \
   collector.frames_per_batch=16 \
   collector.num_workers=4 \
   collector.env_per_collector=2 \
-  collector.collector_device=cuda:0 \
+  collector.device=cuda:0 \
+  collector.device=cuda:0 \
   network.device=cuda:0 \
   logger.mode=offline \
   env.name=Pendulum-v1 \
   logger.backend=
 python .github/unittest/helpers/coverage_run_parallel.py examples/iql/iql_online.py \
-  total_frames=48 \
-  batch_size=10 \
-  frames_per_batch=16 \
-  num_workers=4 \
-  env_per_collector=2 \
-  collector_device=cuda:0 \
-  device=cuda:0 \
-  mode=offline \
-  logger=
+  collector.total_frames=48 \
+  buffer.batch_size=10 \
+  collector.frames_per_batch=16 \
+  collector.env_per_collector=2 \
+  collector.device=cuda:0 \
+  network.device=cuda:0 \
+  logger.mode=offline \
+  logger.backend=
 
 # With single envs
 python .github/unittest/helpers/coverage_run_parallel.py examples/dreamer/dreamer.py \
@@ -175,6 +195,7 @@ python .github/unittest/helpers/coverage_run_parallel.py examples/dreamer/dreame
   num_workers=2 \
   env_per_collector=1 \
   collector_device=cuda:0 \
+  model_device=cuda:0 \
   optim_steps_per_batch=1 \
   record_video=True \
   record_frames=4 \
@@ -183,13 +204,12 @@ python .github/unittest/helpers/coverage_run_parallel.py examples/dreamer/dreame
 python .github/unittest/helpers/coverage_run_parallel.py examples/ddpg/ddpg.py \
   collector.total_frames=48 \
   collector.init_random_frames=10 \
-  optimization.batch_size=10 \
+  optim.batch_size=10 \
   collector.frames_per_batch=16 \
-  collector.num_workers=2 \
   collector.env_per_collector=1 \
-  collector.collector_device=cuda:0 \
+  collector.device=cuda:0 \
   network.device=cuda:0 \
-  optimization.utd_ratio=1 \
+  optim.utd_ratio=1 \
   replay_buffer.size=120 \
   env.name=Pendulum-v1 \
   logger.backend=
@@ -208,52 +228,52 @@ python .github/unittest/helpers/coverage_run_parallel.py examples/dqn/dqn.py \
   record_frames=4 \
   buffer_size=120
 python .github/unittest/helpers/coverage_run_parallel.py examples/redq/redq.py \
-  total_frames=48 \
-  init_random_frames=10 \
-  batch_size=10 \
-  frames_per_batch=16 \
   num_workers=2 \
-  env_per_collector=1 \
-  collector_device=cuda:0 \
-  optim_steps_per_batch=1 \
-  record_video=True \
-  record_frames=4 \
-  buffer_size=120
+  collector.total_frames=48 \
+  collector.init_random_frames=10 \
+  collector.frames_per_batch=16 \
+  collector.env_per_collector=1 \
+  buffer.batch_size=10 \
+  collector.device=cuda:0 \
+  optim.steps_per_batch=1 \
+  logger.record_video=True \
+  logger.record_frames=4 \
+  buffer.size=120 \
+  logger.backend=
 python .github/unittest/helpers/coverage_run_parallel.py examples/sac/sac.py \
   collector.total_frames=48 \
   collector.init_random_frames=10 \
   collector.frames_per_batch=16 \
-  collector.num_workers=2 \
   collector.env_per_collector=1 \
-  collector.collector_device=cuda:0 \
+  collector.device=cuda:0 \
+  optim.batch_size=10 \
+  optim.utd_ratio=1 \
   network.device=cuda:0 \
-  optimization.batch_size=10 \
-  optimization.utd_ratio=1 \
+  optim.batch_size=10 \
+  optim.utd_ratio=1 \
   replay_buffer.size=120 \
   env.name=Pendulum-v1 \
   logger.backend=
-#  record_video=True \
-#  record_frames=4 \
 python .github/unittest/helpers/coverage_run_parallel.py examples/iql/iql_online.py \
-  total_frames=48 \
-  batch_size=10 \
-  frames_per_batch=16 \
-  num_workers=2 \
-  env_per_collector=1 \
-  mode=offline \
-  device=cuda:0 \
-  collector_device=cuda:0 \
-  logger=
+  collector.total_frames=48 \
+  collector.frames_per_batch=16 \
+  collector.env_per_collector=1 \
+  collector.device=cuda:0 \
+  network.device=cuda:0 \
+  buffer.batch_size=10 \
+  logger.mode=offline \
+  logger.backend=
 python .github/unittest/helpers/coverage_run_parallel.py examples/td3/td3.py \
   collector.total_frames=48 \
   collector.init_random_frames=10 \
-  optimization.batch_size=10 \
   collector.frames_per_batch=16 \
   collector.num_workers=2 \
   collector.env_per_collector=1 \
+  collector.device=cuda:0 \
   logger.mode=offline \
-  collector.collector_device=cuda:0 \
+  optim.batch_size=10 \
   env.name=Pendulum-v1 \
+  network.device=cuda:0 \
   logger.backend=
 python .github/unittest/helpers/coverage_run_parallel.py examples/multiagent/mappo_ippo.py \
   collector.n_iters=2 \
@@ -286,8 +306,10 @@ python .github/unittest/helpers/coverage_run_parallel.py examples/multiagent/sac
   train.minibatch_size=100 \
   logger.backend=
 
-
 python .github/unittest/helpers/coverage_run_parallel.py examples/bandits/dqn.py --n_steps=100
 
+## RLHF
+# RLHF tests are executed in the dedicated workflow
+
 coverage combine
 coverage xml -i

.github/unittest/linux_libs/scripts_brax/install.sh

@@ -30,19 +30,19 @@ if [ "${CU_VERSION:-}" == cpu ] ; then
     # conda install -y pytorch torchvision cpuonly -c pytorch-nightly
     # use pip to install pytorch as conda can frequently pick older release
 #    conda install -y pytorch cpuonly -c pytorch-nightly
-    pip3 install --pre torch --extra-index-url https://download.pytorch.org/whl/nightly/cpu --force-reinstall  --progress-bar off
+    pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu --force-reinstall  --progress-bar off
 else
-    pip3 install --pre torch --extra-index-url https://download.pytorch.org/whl/nightly/cu116 --force-reinstall  --progress-bar off
+    pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121 --force-reinstall  --progress-bar off
 fi
 
 # install tensordict
-pip install git+https://github.com/pytorch-labs/tensordict.git --progress-bar off
+pip install git+https://github.com/pytorch/tensordict.git --progress-bar off
 
 # smoke test
 python -c "import functorch;import tensordict"
 
 printf "* Installing torchrl\n"
-pip3 install -e .
+python setup.py develop
 
 # smoke test
 python -c "import torchrl"