diff --git a/.ci/docker/requirements.txt b/.ci/docker/requirements.txt index f1f105122e..3d8507160a 100644 --- a/.ci/docker/requirements.txt +++ b/.ci/docker/requirements.txt @@ -38,8 +38,8 @@ tensorboard jinja2==3.1.3 pytorch-lightning torchx -torchrl==0.7.2 -tensordict==0.7.2 +torchrl==0.9.2 +tensordict==0.9.1 # For ax_multiobjective_nas_tutorial.py ax-platform>=0.4.0,<0.5.0 nbformat>=5.9.2 diff --git a/.jenkins/build.sh b/.jenkins/build.sh index 0187f6ba84..f25779f3f7 100755 --- a/.jenkins/build.sh +++ b/.jenkins/build.sh @@ -22,10 +22,13 @@ sudo apt-get install -y pandoc #Install PyTorch Nightly for test. # Nightly - pip install --pre torch torchvision torchaudio -f https://download.pytorch.org/whl/nightly/cu102/torch_nightly.html # Install 2.5 to merge all 2.4 PRs - uncomment to install nightly binaries (update the version as needed). -# sudo pip uninstall -y fbgemm-gpu torchrec -# sudo pip uninstall -y torch torchvision torchaudio torchtext torchdata torchrl tensordict -# sudo pip3 install fbgemm-gpu==1.1.0 torchrec==1.0.0 --no-cache-dir --index-url https://download.pytorch.org/whl/test/cu124 -# pip3 install torch==2.7.0 torchvision torchaudio --no-cache-dir --index-url https://download.pytorch.org/whl/test/cu126 +#sudo pip uninstall -y fbgemm-gpu torchrec +sudo pip uninstall -y torch torchvision torchaudio torchtext torchdata torchrl tensordict +pip3 install torch==2.8.0 torchvision torchaudio --no-cache-dir --index-url https://download.pytorch.org/whl/test/cu126 +pip3 install torchrl==0.9.1 tensordict==0.9.1 +# sudo pip3 install fbgemm-gpu==1.2.0 torchrec==1.2.0 --no-cache-dir --index-url https://download.pytorch.org/whl/test/cu126 + + # Install two language tokenizers for Translation with TorchText tutorial python -m spacy download en_core_web_sm python -m spacy download de_core_news_sm diff --git a/intermediate_source/dqn_with_rnn_tutorial.py b/intermediate_source/dqn_with_rnn_tutorial.py index bcc484f0a0..462415dcc7 100644 --- a/intermediate_source/dqn_with_rnn_tutorial.py +++ b/intermediate_source/dqn_with_rnn_tutorial.py @@ -342,7 +342,9 @@ # will return a new instance of the LSTM (with shared weights) that will # assume that the input data is sequential in nature. # -policy = Seq(feature, lstm.set_recurrent_mode(True), mlp, qval) +from torchrl.modules import set_recurrent_mode + +policy = Seq(feature, lstm, mlp, qval) ###################################################################### # Because we still have a couple of uninitialized parameters we should @@ -389,7 +391,10 @@ # For the sake of efficiency, we're only running a few thousands iterations # here. In a real setting, the total number of frames should be set to 1M. # -collector = SyncDataCollector(env, stoch_policy, frames_per_batch=50, total_frames=200, device=device) + +collector = SyncDataCollector( + env, stoch_policy, frames_per_batch=50, total_frames=200, device=device +) rb = TensorDictReplayBuffer( storage=LazyMemmapStorage(20_000), batch_size=4, prefetch=10 ) @@ -422,7 +427,8 @@ rb.extend(data.unsqueeze(0).to_tensordict().cpu()) for _ in range(utd): s = rb.sample().to(device, non_blocking=True) - loss_vals = loss_fn(s) + with set_recurrent_mode(True): + loss_vals = loss_fn(s) loss_vals["loss"].backward() optim.step() optim.zero_grad() @@ -464,5 +470,5 @@ # # Further Reading # --------------- -# +# # - The TorchRL documentation can be found `here `_.