How can I apply ppo or other rl models based on torchrl and tensordict on a Orin NX?

My ENV is down below:

  • Jetson 6.2 CUDA12.6
  • Device: Orin NX
    Since it seems that no complied torchrl or tensordict whl for jetson6.2, but my model is written based on torchrl and tensordict, I have a little trouble implementing this model on NX. Do I have to rewrite the model using only torch or there is another simple solution to it. Thanks for answering.

Hi,

Although torchrl doesn’t have a prebuilt for Jetson, we can install torchdict via pip command directly:

$ pip install tensordict-nightly
Defaulting to user installation because normal site-packages is not writeable
Collecting tensordict-nightly
  Downloading tensordict_nightly-2023.6.8-py310-none-any.whl (133 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 133.8/133.8 KB 1.4 MB/s eta 0:00:00
Requirement already satisfied: torch in ./.local/lib/python3.10/site-packages (from tensordict-nightly) (2.6.0rc1)
Requirement already satisfied: numpy in /usr/lib/python3/dist-packages (from tensordict-nightly) (1.21.5)
Collecting cloudpickle
  Downloading cloudpickle-3.1.1-py3-none-any.whl (20 kB)
Requirement already satisfied: typing-extensions>=4.10.0 in ./.local/lib/python3.10/site-packages (from torch->tensordict-nightly) (4.12.2)
Requirement already satisfied: sympy==1.13.1 in ./.local/lib/python3.10/site-packages (from torch->tensordict-nightly) (1.13.1)
Requirement already satisfied: fsspec in ./.local/lib/python3.10/site-packages (from torch->tensordict-nightly) (2025.2.0)
Requirement already satisfied: filelock in ./.local/lib/python3.10/site-packages (from torch->tensordict-nightly) (3.17.0)
Requirement already satisfied: networkx in ./.local/lib/python3.10/site-packages (from torch->tensordict-nightly) (3.4.2)
Requirement already satisfied: jinja2 in ./.local/lib/python3.10/site-packages (from torch->tensordict-nightly) (3.1.5)
Requirement already satisfied: mpmath<1.4,>=1.1.0 in ./.local/lib/python3.10/site-packages (from sympy==1.13.1->torch->tensordict-nightly) (1.3.0)
Requirement already satisfied: MarkupSafe>=2.0 in /usr/lib/python3/dist-packages (from jinja2->torch->tensordict-nightly) (2.0.1)
Installing collected packages: cloudpickle, tensordict-nightly
Successfully installed cloudpickle-3.1.1 tensordict-nightly-2023.6.8

Thanks.

1 Like

Torchrl models can be exported to plain python / c++ with little effort, check the tutorial here
https://pytorch.org/rl/stable/tutorials/export.html
https://pytorch.org/tensordict/stable/tutorials/export.html

If you encounter any issue please let us know here or on the Torchrl GitHub repo!

This topic was automatically closed 14 days after the last reply. New replies are no longer allowed.