Skip to content

Commit 30a143c

Browse files
authored
[Doc] Document optional CUDA TorchRL wheels (#3817)
1 parent 28cf305 commit 30a143c

1 file changed

Lines changed: 17 additions & 1 deletion

File tree

README.md

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,8 @@ improvements that are worth surfacing up front:
5959
- stronger multi-agent coverage through MAPPO, IPPO, `MultiAgentGAE`,
6060
value-normalization utilities, and mixer configs;
6161
- better collector and replay-buffer ergonomics, including async prioritized
62-
writes, ordered storage access, compact observations, and HER;
62+
writes, ordered storage access, compact observations, HER, and optional CUDA
63+
wheels for CUDA-based prioritized replay-buffer kernels;
6364
- new transforms and value-estimator improvements such as `ActionScaling`,
6465
`FlattenAction`, `NextObservationDelta`, compact shifted estimators, and
6566
chunked forwards.
@@ -420,6 +421,21 @@ Install the stable release:
420421
pip install torchrl
421422
```
422423

424+
This standard PyPI wheel is the right default for most users, including CPU
425+
prioritized replay buffers and workloads that do not use prioritized replay.
426+
Starting with TorchRL 0.13, Linux CUDA wheels are also published for users who
427+
want the CUDA-based prioritized replay-buffer implementations. Install the
428+
CUDA wheel from the PyTorch wheel index that matches your PyTorch CUDA runtime
429+
(replace `cu128` with the CUDA build you use):
430+
431+
```bash
432+
pip install "torchrl==0.13.0+cu128" --extra-index-url https://download.pytorch.org/whl/cu128
433+
```
434+
435+
The CUDA wheel is optional: if you do not need CUDA prioritized replay buffers,
436+
or if your prioritized replay buffers run on CPU, keep using `pip install
437+
torchrl`.
438+
423439
Install common optional dependencies:
424440

425441
```bash

0 commit comments

Comments
 (0)