openpi-thor

Use openpi-thor -h to list commands, or openpi-thor <command> -h / --help to show all arguments for one command.

• openpi-thor doctor Checks whether the Jetson AGX Thor runtime has the required Python packages and host tools.
• openpi-thor convert-jax Converts a JAX checkpoint into a PyTorch bundle directory.
• openpi-thor export-onnx Exports an ONNX model into an existing bundle directory.
• openpi-thor build-engine Builds a TensorRT engine from the ONNX artifact stored in the bundle.
• openpi-thor prepare-engine High-level command that can export ONNX, build TensorRT, and optionally validate in one step.
• openpi-thor validate Compares PyTorch or TensorRT outputs against the JAX reference model on real dataset samples.
• openpi-thor validate-tensorrt Compares one TensorRT engine against another, such as fp8 or fp8+nvfp4 against the recommended fp16 engine.
• openpi-thor status Prints a compact summary of bundle contents, recommended engine, validation state, and report files.
• openpi-thor serve Starts the websocket inference server for the selected or recommended engine.

openpi_thor_bundle.json is intentionally a compact overview, not a full history log.

Use it to answer high-level questions such as:

• which artifacts exist
• which engine is recommended
• whether recent validations passed
• which detailed report file belongs to each phase

Detailed phase outputs stay under reports/, including export, build, validation, and debug payloads. Bundle-internal paths recorded in the manifest are bundle-relative, so the bundle remains easier to inspect and move as one directory.

For a quick overview, use:

openpi-thor status --bundle-dir /path/to/bundle

For the same overview plus loaded report payloads, use:

openpi-thor status --bundle-dir /path/to/bundle --verbose

Start with FP16 before spending time on FP8 or NVFP4:

openpi-thor prepare-engine \
  --config <PI05_TRAIN_CONFIG> \
  --bundle-dir /path/to/bundle \
  --precision fp16 \
  --validate \
  --reference-checkpoint-dir /path/to/jax-checkpoint

The Jetson AI Lab tutorial warns that pure fp16 export is unsupported for its tutorial exporter because Pi0.5 is BF16-native and naive fp16 overflows in Gemma. Our validated fp16 path is a custom openpi-thor extension, not a contradiction: it preserves fp32 stability islands during export and builds TensorRT engines with --strongly-typed by default.

openpi-thor prepare-engine \
  --config <PI05_TRAIN_CONFIG> \
  --bundle-dir /path/to/bundle \
  --precision fp8 \
  --validate \
  --reference-checkpoint-dir /path/to/jax-checkpoint

32 is the default calibration count. In our real-data sweeps, raising plain fp8 calibration to 128 did not improve accuracy, so openpi-thor keeps 32 as the fast default. Use 128 or 256 only when you explicitly want a slower comparison sweep.

--enable-llm-nvfp4 in openpi-thor currently enables broad Gemma attention-side NVFP4: all Gemma attention layers use NVFP4, and quantized attention matmuls are turned on automatically, while the Gemma MLP path stays on fp8. On the tested Jetson AGX Thor stack, this attention-only NVFP4 path outperformed plain fp8 without exceeding the fp8-vs-JAX error budget, while broader attention+MLP NVFP4 still drifted too much after TensorRT lowering.

Optional dataset overrides for calibration and validation:

• --dataset-repo-id <HF_DATASET_REPO>
• --dataset-root <LOCAL_LEROBOT_ROOT>

If omitted, openpi-thor uses the dataset specified by the training config.

Keep validate for JAX-based checks. Use validate-tensorrt when you want to compare one engine against another on the same real dataset samples.

openpi-thor validate-tensorrt \
  --config <PI05_TRAIN_CONFIG> \
  --bundle-dir /path/to/bundle \
  --candidate-engine-path /path/to/bundle/engine/model_fp8.engine

If --reference-engine-path is omitted, validate-tensorrt uses the bundle's recommended engine, which is normally the validated fp16 engine.

Normally serve uses the recommended engine recorded in the bundle. Use --engine-path only for A/B tests or debugging.

validate and validate-tensorrt always write a report, but serve treats fp8-style engines more strictly: an fp8 or fp8+nvfp4 engine must have a passing validation report for that exact artifact unless you explicitly override the safety gate with --no-require-validated.

Default pass/fail thresholds are currently:

• PyTorch vs JAX: min_cosine >= 0.999, mean_abs_error <= 0.01, max_abs_error <= 0.05
• TensorRT fp16 vs JAX: min_cosine >= 0.995, mean_abs_error <= 0.03, max_abs_error <= 0.12
• Any comparison that involves fp8 or fp8+nvfp4: min_cosine >= 0.97, mean_abs_error <= 0.08, max_abs_error <= 0.3

So a validation report can exist and still block serving if it did not pass those thresholds.

Two design choices in openpi-thor are deliberate departures from a naive “export everything in fp16” or “turn on broad NVFP4 everywhere” approach.

The Jetson AI Lab tutorial is correct that a naive pure-fp16 export of Pi0.5 is unsafe: the model is BF16-native, and weakly typed TensorRT builds can overflow numerically sensitive Gemma operations such as RMSNorm and attention-related reductions. In our debugging on Jetson AGX Thor, that weakly typed fp16 path produced large regressions against JAX.

openpi-thor therefore uses a custom fp16 path instead:

• preserve selected fp32 stability islands during export
• build TensorRT with --strongly-typed by default

That combination was the difference between the broken and usable fp16 paths on the tested stack. In our real-sample validations, the corrected fp16 TensorRT engine matched JAX closely and became the recommended deployment path.

The original broad NVFP4 path that pushed both Gemma attention and Gemma MLP too far was much worse. Quantized PyTorch stayed reasonably close to the fp8 baseline, but the TensorRT engine drifted badly after ONNX/TensorRT lowering. In our use case, that broad attention+MLP fp8+nvfp4 path showed roughly a 30x larger mean absolute error than pure fp8.

That is why openpi-thor does not keep the broad full-layer-everywhere NVFP4 recipe. The current --enable-llm-nvfp4 path is:

• all Gemma attention layers use NVFP4
• attention matmuls are quantized explicitly
• the Gemma MLP path stays on fp8

That attention-only path brought the TensorRT result back into the same accuracy regime as pure fp8 on the tested Jetson AGX Thor stack, while also improving steady-state engine compute time.

If you debug this area further, the most reliable workflow is:

1. keep LayerNorm/RMSNorm and other stability-sensitive paths on the existing higher-precision export path
2. profile candidate engines with trtexec instead of relying on end-to-end Python timing
3. reject candidates whose TensorRT layer profiles are dominated by cast/dequant helper kernels
4. compare every candidate against both plain fp8 TensorRT and the JAX reference on the same real-data examples

The supported v1 companion-repo model assumes:

• checkout path is exactly packages/openpi-thor
• the host repo looks like upstream openpi
• the host repo still contains:
  • src/openpi
  • examples/convert_jax_model_to_pytorch.py
  • src/openpi/models_pytorch/transformers_replace/...

The patch script only edits host files that openpi-thor depends on. It ensures:

• tool.uv.override-dependencies contains:
  • ml-dtypes==0.5.1
  • tensorstore==0.1.74
• tool.uv.conflicts contains the Thor dependency-group conflict entries required by openpi-thor
• tool.uv.sources contains:
  • openpi = { workspace = true }
  • openpi-client = { workspace = true }
  • lerobot = { git = "...", tag = "v0.5.0" }
• if the host repo still uses the older upstream LeRobot import in src/openpi/training/data_loader.py, it patches that file to support newer LeRobot tags too
• if the host repo still uses the older upstream PromptFromLeRobotTask implementation in src/openpi/transforms.py, it patches that file so newer LeRobot task metadata also works

If your host repo already carries the newer LeRobot import and the DataFrame-aware PromptFromLeRobotTask, these source patch steps become a no-op.

These source patches matter most for real-dataset flows, especially:

• FP8 calibration on LeRobot data
• openpi-thor validate
• openpi-thor validate-tensorrt
• any config with prompt_from_task=True

If you want a single centralized host-side serving entrypoint, you can also patch the host repo's scripts/serve_policy.py to add a TensorRT-bundle mode that delegates to openpi_thor.runtime.load_tensorrt_policy(...). That is enough to let one host script serve either:

• a regular checkpoint-backed OpenPI policy
• or an openpi-thor TensorRT bundle

openpi-thor serve remains the canonical interface, but this optional patch can be convenient if you want one shared launcher in your host openpi fork.

Minimal shape:

@dataclasses.dataclass
class TensorRTBundle:
    config: str
    bundle_dir: str
    engine_path: str | None = None
    require_validated: bool = True


def create_policy(args: Args) -> _policy.Policy:
    match args.policy:
        case TensorRTBundle():
            from openpi_thor.runtime import load_tensorrt_policy

            return load_tensorrt_policy(
                _config.get_config(args.policy.config),
                args.policy.bundle_dir,
                engine_path=args.policy.engine_path,
                require_validated=args.policy.require_validated,
                default_prompt=args.default_prompt,
            )
        # existing checkpoint/default cases stay unchanged

Manual fallback:

• add those same tool.uv.override-dependencies, tool.uv.conflicts, and tool.uv.sources entries by hand if you do not want to run the patch script
• if needed, update src/openpi/training/data_loader.py from:

import lerobot.common.datasets.lerobot_dataset as lerobot_dataset

to:

try:
    import lerobot.datasets.lerobot_dataset as lerobot_dataset
except ImportError:
    import lerobot.common.datasets.lerobot_dataset as lerobot_dataset

• if needed, update PromptFromLeRobotTask in src/openpi/transforms.py so it supports both dict-style task mappings and newer DataFrame-style LeRobot task metadata
• if you apply the PromptFromLeRobotTask fallback manually, also make sure Any is imported from typing

The package is intended to be portable across Jetson AGX Thor setups, but the current flow was tested on:

• Ubuntu 24.04.4 LTS
• Python 3.12.3
• JetPack 7.0-b128
• L4T 38.2.2
• CUDA 13.0.48

Tested stack for the FP16 path:

• PyTorch 2.11.0
• torchvision 0.26.0
• JAX / jaxlib 0.5.3
• Orbax Checkpoint 0.11.13
• Transformers 4.53.2
• ONNX 1.18.0
• ONNX GraphSurgeon 0.5.8
• TensorRT Python 10.13.3.x

Additional packages for FP8 / NVFP4:

• nvidia-modelopt==0.33.1
• onnxruntime>=1.24.4

Important notes:

• openpi-thor uses the official lerobot v0.5.0 tag by default.
• The current dataset and calibration code paths are also compatible with official lerobot v0.4.4.
• lerobot is installed separately with --no-deps because both official tags publish dependency metadata that pins torch<2.11.0, which can downgrade the working Jetson AGX Thor PyTorch stack.
• openpi-thor expects host TensorRT and trtexec; it does not install those itself.
• openpi-thor builds TensorRT engines with --strongly-typed by default. This preserves explicit FP32 stability islands in Gemma-based models, such as RMSNorm, and avoids the large numerical drift seen with weakly typed builds on Jetson AGX Thor.
• The tutorial's warning about pure fp16 export applies to its own exporter. openpi-thor supports a custom fp16 path that preserves fp32 stability islands and uses strongly typed TensorRT builds.
• openpi-thor keeps the public NVFP4 path narrower than “full-layer everywhere”: it enables NVFP4 across Gemma attention with explicit attention-matmul quantization, while leaving the Gemma MLP path on fp8. That is an intentional deviation from the broadest tutorial recipe because it produced a better speed/accuracy tradeoff on the tested Jetson AGX Thor stack.
• FP8 calibration defaults to 32 real samples. Use 128 or 256 only for slower calibration sweeps and comparison experiments.

• If torch.cuda.is_available() becomes False after an install step, re-assert the tested Jetson AGX Thor PyTorch group.
• If prepare-engine fails very early on import, check that the runtime env includes both torch and torchvision.
• Weakly typed TensorRT builds can drift badly on Gemma-based models. openpi-thor uses --strongly-typed by default; only disable it with --no-strongly-typed for debugging.
• If FP8/NVFP4 export fails, confirm that nvidia-modelopt, onnxruntime, and host TensorRT tools are all present in the Jetson AGX Thor runtime.
• If you keep a backup of this repo, do not leave another openpi-thor checkout under packages/, because uv will see duplicate workspace members with the same package name.
• Norm stats not found in ... , skipping. usually means the host OpenPI config probed one assets directory first and then fell back to another path. That message is only a problem if no later Loaded norm stats from ... line appears.
• JAX version ... available. during TensorRT serving is an import-time info message from the Python stack, typically Hugging Face datasets detecting that JAX is installed. It does not mean the served policy is using JAX for inference.
• Use the repo dev environment for unit tests, not the Jetson AGX Thor runtime venv:

uv run --package openpi-thor pytest packages/openpi-thor/tests -m "not manual"