Segfault (exit 139) during import openpi.policies.policy_config— native library load order + workaround

Hi! First of all, thanks for this great project — it's been really helpful for my research. I hit a strange crash while following the LIBERO example (non-Docker path) and wanted to report it properly, with a minimal reproducer and the workaround that got me unblocked.

## Environment
- Cloud GPU container (AutoDL), Ubuntu 20.04, RTX 4090, NVIDIA driver 580.105.08
- Python 3.11.16 (created via `uv run` following the LIBERO example instructions, non-Docker path)
- openpi @ main (Aug 2026), `jax[cuda12]==0.5.3`, `torch==2.7.1`, `transformers`, `sentencepiece`, `orbax-checkpoint==0.11.13`

## Symptom
`uv run python -u scripts/serve_policy.py --env LIBERO` exits immediately with no output at all. Exit code is 139 (SIGSEGV), even with `-u` and `> log 2>&1` the log file stays empty.

## Minimal reproducer
This command reliably segfaults (5/5 in a fresh container):

```bash
uv run python -c "import jax; import orbax.checkpoint; import sentencepiece; from transformers import AutoProcessor; import openpi.training.checkpoints; print('ok')"; echo "exit=$?"
# -> exit=139 (no output)

Observations:

• import openpi.training.checkpoints alone works.
• Reversing the order works: import openpi.training.checkpoints; from openpi.models import tokenizer prints ok.
• Each single import from the chain + checkpoints also works. The crash only happens with the full order above (i.e., it is load-order dependent, not a broken package).
• A probe that imports the modules one by one pinpoints the crash at from openpi.policies import policy_config (which imports the chain above).

Workaround (confirmed)

Pre-loading openpi.training.checkpoints before anything else fixes it. Create the following sitecustomize.py in the venv site-packages (e.g. .venv/lib/python3.11/site-packages/sitecustomize.py):

try:
    import openpi.training.checkpoints
except Exception:
    pass
try:
    import torch
except Exception:
    pass

After this, the minimal reproducer passes 5/5 and serve_policy.py --env LIBERO starts and serves normally (verified end-to-end with the LIBERO client, ~90% success on libero_spatial 2-trial run).

Suspected cause (hypothesis)

A native-library load-order conflict: transformers pulls in torch; openpi.training.checkpoints pulls openpi.training.data_loader (dlimp/TensorFlow-related native libs). Loading them in the jax -> orbax -> sentencepiece -> transformers -> checkpoints order seems to segfault; loading checkpoints first avoids it. I did not fully root-cause which exact .so conflicts.

Possible proper fix

openpi/training/checkpoints.py imports openpi.training.data_loader at module top level, but _data_loader is only used by training code paths. Making that import lazy (inside the functions that need it) would avoid pulling TF-related native libs at import time for inference/serving and likely eliminate this crash class.

Possibly related: #831 and #827 (silent exits during training with no error output — maybe the same segfault).

Thanks for reading! Happy to run any additional tests if that helps narrow it down. (P.S. I used AI assistance to polish the English wording; every fact above was reproduced on my own machine.)