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

Robotics & Edge Computing Jetson & Embedded Systems Jetson Nano
1

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.
3

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
4

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!