openpi 코드 레퍼런스

이 문서는 .git을 제외한 현재 워크트리의 모든 폴더와 파일을 기준으로 작성한 코드 인벤토리다. 현재 저장소에는 파일 143개, Python 파일 91개가 있다. third_party/aloha, third_party/libero는 submodule로 등록되어 있지만 현재 checkout되어 있지 않아 폴더만 존재한다.

1. 폴더 인벤토리

폴더	설명
`.`	프로젝트 루트. README, 라이선스, Python packaging, uv lock, lint/test config가 있다.
`.agents`	현재 파일이 없는 로컬 agent 메타데이터 폴더. 코드 실행 경로에는 쓰이지 않는다.
`.codex`	현재 파일이 없는 로컬 Codex 메타데이터 폴더. 코드 실행 경로에는 쓰이지 않는다.
`.github`	GitHub repository 설정.
`.github/workflows`	pre-commit, test CI workflow.
`.vscode`	VS Code Python/Ruff/Pytest 설정.
`docs`	사용자 문서. 기존 Docker/norm stats/remote inference 문서와 한국어 분석 문서가 있다.
`examples`	실제 실행 예제와 데이터 변환 script.
`examples/aloha_real`	실제 ALOHA 로봇 실행, ROS/Interbotix 환경, HDF5->LeRobot 변환.
`examples/aloha_sim`	`gym-aloha` simulator client/runtime.
`examples/droid`	DROID robot 실행, DROID 데이터 변환, full DROID training README.
`examples/libero`	LIBERO benchmark 변환/evaluation/Docker workflow.
`examples/simple_client`	로봇 없이 websocket policy server를 random observation으로 테스트.
`examples/ur5`	UR5 커스텀 로봇 적용 가이드.
`packages`	workspace package 모음.
`packages/openpi-client`	로봇/클라이언트 환경에서 설치하는 경량 client package.
`packages/openpi-client/src/openpi_client`	websocket client, action broker, runtime base class.
`packages/openpi-client/src/openpi_client/runtime`	agent/environment/subscriber/runtime loop abstraction.
`packages/openpi-client/src/openpi_client/runtime/agents`	policy를 runtime agent로 감싸는 구현.
`scripts`	CLI entrypoint. 학습, PyTorch 학습, norm stats, policy serving.
`scripts/docker`	openpi server용 Dockerfile/compose와 설치 helper.
`src`	main package source root.
`src/openpi`	openpi Python package.
`src/openpi/models`	JAX 모델과 tokenizer, SigLIP/Gemma/LoRA 구현.
`src/openpi/models/utils`	FSQ/LFQ action tokenizer baseline.
`src/openpi/models_pytorch`	PyTorch π0/π0.5와 transformers patch.
`src/openpi/models_pytorch/transformers_replace`	Hugging Face transformers 4.53.2에 복사하는 override 파일.
`src/openpi/models_pytorch/transformers_replace/models/gemma`	Gemma config/model patch.
`src/openpi/models_pytorch/transformers_replace/models/paligemma`	PaliGemma model patch.
`src/openpi/models_pytorch/transformers_replace/models/siglip`	SigLIP model patch와 설치 확인 helper.
`src/openpi/policies`	checkpoint policy 생성과 ALOHA/DROID/LIBERO 입출력 변환.
`src/openpi/serving`	websocket policy server.
`src/openpi/shared`	download/cache, normalization, image resize, typing, NNX helper.
`src/openpi/training`	config, data loader, checkpoint, optimizer, sharding, weight loader.
`src/openpi/training/misc`	RoboArena/PolaRiS 추가 연구 config.
`third_party`	외부 submodule root.
`third_party/aloha`	ALOHA submodule 위치. 현재 미초기화.
`third_party/libero`	LIBERO submodule 위치. 현재 미초기화.

2. 파일 인벤토리

파일	역할
`.dockerignore`	Docker build context에서 `.venv`, `checkpoints`, `data`를 제외한다.
`.github/CODEOWNERS`	코드 소유자 지정. 모델/학습 영역에 별도 owner가 있다.
`.github/workflows/pre-commit.yml`	CI에서 pre-commit을 실행한다.
`.github/workflows/test.yml`	CI에서 uv 환경을 만들고 pytest를 실행한다.
`.gitignore`	assets/checkpoints/data/wandb, Python cache/build/env 파일 제외.
`.gitmodules`	`third_party/aloha`, `third_party/libero` submodule 등록.
`.pre-commit-config.yaml`	uv lock, Ruff lint/format hook. `third_party` 제외.
`.python-version`	프로젝트 Python 버전 pin.
`.vscode/settings.json`	Ruff format on save와 pytest args 설정.
`CONTRIBUTING.md`	기여/이슈/PR 가이드.
`LICENSE`	openpi repository license.
`LICENSE_GEMMA.txt`	Gemma 관련 라이선스.
`README.md`	설치, checkpoint, inference, fine-tuning, PyTorch, troubleshooting의 중심 문서.
`docs/docker.md`	Docker 설치/실행 가이드.
`docs/norm_stats.md`	base norm stats asset id와 action space 설명.
`docs/remote_inference.md`	websocket remote inference 사용법.
`docs/openpi_deep_dive_ko.md`	한국어 전체 아키텍처/연구 배경 문서.
`docs/openpi_code_reference_ko.md`	현재 문서.
`docs/openpi_usage_line_by_line_ko.md`	설치/실행/핵심 코드 line-by-line 해설 문서.
`examples/aloha_real/Dockerfile`	ALOHA real client Docker image. ROS/Interbotix 의존성 포함.
`examples/aloha_real/README.md`	실제 ALOHA 실행, checkpoint, scene setup, training 안내.
`examples/aloha_real/compose.yml`	ALOHA real Docker compose.
`examples/aloha_real/constants.py`	ALOHA gripper/arm normalization 상수.
`examples/aloha_real/convert_aloha_data_to_lerobot.py`	ALOHA HDF5 demonstration을 LeRobot dataset으로 변환.
`examples/aloha_real/env.py`	`RealEnv`를 openpi-client `Environment`로 감싸 policy runtime에 맞춘다.
`examples/aloha_real/main.py`	ALOHA real runtime entrypoint. websocket policy + action chunk broker + real env.
`examples/aloha_real/real_env.py`	Interbotix/ROS 기반 실제 ALOHA robot environment.
`examples/aloha_real/requirements.in`	ALOHA real Python 의존성 source.
`examples/aloha_real/requirements.txt`	ALOHA real lock-style requirements.
`examples/aloha_real/robot_utils.py`	ROS image/joint recorder, arm/gripper command helper.
`examples/aloha_real/video_display.py`	runtime subscriber 형태의 OpenCV video display. 현재 main path에는 기본 연결되어 있지 않다.
`examples/aloha_sim/Dockerfile`	ALOHA simulator client Docker image.
`examples/aloha_sim/README.md`	ALOHA sim 실행법.
`examples/aloha_sim/compose.yml`	ALOHA sim Docker compose.
`examples/aloha_sim/env.py`	gym_aloha observation/action을 runtime Environment로 변환.
`examples/aloha_sim/main.py`	ALOHA sim runtime entrypoint.
`examples/aloha_sim/requirements.in`	ALOHA sim 의존성 source.
`examples/aloha_sim/requirements.txt`	ALOHA sim requirements.
`examples/aloha_sim/saver.py`	episode frame을 mp4로 저장하는 subscriber.
`examples/convert_jax_model_to_pytorch.py`	Orbax/JAX π0/π0.5 checkpoint를 PyTorch `model.safetensors`로 변환.
`examples/droid/README.md`	DROID robot inference/evaluation 안내.
`examples/droid/README_train.md`	full DROID RLDS training, idle filtering, norm stats, command 안내.
`examples/droid/compute_droid_nonidle_ranges.py`	DROID RLDS에서 idle 구간을 제외할 timestep range JSON 생성.
`examples/droid/convert_droid_data_to_lerobot.py`	DROID raw HDF5/MP4/annotation을 LeRobot dataset으로 변환.
`examples/droid/main.py`	실제 DROID robot inference loop. observation 추출, action chunk 실행, logging.
`examples/inference.ipynb`	checkpoint inference notebook.
`examples/libero/Dockerfile`	LIBERO evaluation Docker image.
`examples/libero/README.md`	LIBERO benchmark 실행, Docker workflow, 결과 안내.
`examples/libero/compose.yml`	LIBERO policy server/eval compose.
`examples/libero/convert_libero_data_to_lerobot.py`	LIBERO RLDS 데이터를 LeRobot dataset으로 변환.
`examples/libero/main.py`	LIBERO task suite evaluation loop.
`examples/libero/requirements.in`	LIBERO 의존성 source.
`examples/libero/requirements.txt`	LIBERO requirements.
`examples/policy_records.ipynb`	`PolicyRecorder`로 저장한 policy input/output 기록 분석 notebook.
`examples/simple_client/Dockerfile`	simple websocket client Docker image.
`examples/simple_client/README.md`	로봇 없이 random observation으로 server를 테스트하는 방법.
`examples/simple_client/compose.yml`	simple client compose.
`examples/simple_client/main.py`	random ALOHA/DROID/LIBERO observation 생성, websocket policy 호출, timing 측정.
`examples/simple_client/requirements.in`	simple client 의존성 source.
`examples/simple_client/requirements.txt`	simple client requirements.
`examples/ur5/README.md`	UR5 커스텀 policy/data config 작성 예시.
`packages/openpi-client/pyproject.toml`	openpi-client package metadata. Python `>=3.7`.
`packages/openpi-client/src/openpi_client/__init__.py`	openpi-client package init.
`packages/openpi-client/src/openpi_client/action_chunk_broker.py`	action chunk를 한 step씩 꺼내는 policy wrapper.
`packages/openpi-client/src/openpi_client/base_policy.py`	client-side `BasePolicy` abstract interface.
`packages/openpi-client/src/openpi_client/image_tools.py`	client-side PIL resize/pad, float-to-uint8 image helper.
`packages/openpi-client/src/openpi_client/image_tools_test.py`	client image resize test.
`packages/openpi-client/src/openpi_client/msgpack_numpy.py`	numpy ndarray/scalar를 msgpack으로 직렬화/역직렬화.
`packages/openpi-client/src/openpi_client/msgpack_numpy_test.py`	msgpack numpy roundtrip test.
`packages/openpi-client/src/openpi_client/runtime/agent.py`	runtime agent abstract class.
`packages/openpi-client/src/openpi_client/runtime/agents/policy_agent.py`	`BasePolicy`를 runtime `Agent`로 감싼다.
`packages/openpi-client/src/openpi_client/runtime/environment.py`	runtime environment abstract class.
`packages/openpi-client/src/openpi_client/runtime/runtime.py`	env-agent-subscriber loop. reset, step, max_hz 제어.
`packages/openpi-client/src/openpi_client/runtime/subscriber.py`	runtime subscriber abstract class.
`packages/openpi-client/src/openpi_client/websocket_client_policy.py`	websocket server에 observation을 보내고 action을 받는 sync policy.
`pyproject.toml`	openpi package metadata, dependency, uv source, Ruff, pytest 설정.
`scripts/__init__.py`	scripts package marker.
`scripts/compute_norm_stats.py`	config별 dataset에서 state/action normalization statistics 계산.
`scripts/docker/compose.yml`	openpi policy server compose. host network/GPU/mount 설정.
`scripts/docker/install_docker_ubuntu22.sh`	Ubuntu 22.04 Docker 설치 helper.
`scripts/docker/install_nvidia_container_toolkit.sh`	NVIDIA container toolkit 설치 helper.
`scripts/docker/serve_policy.Dockerfile`	CUDA runtime 기반 policy server Dockerfile.
`scripts/serve_policy.py`	checkpoint/default env policy를 websocket server로 띄우는 CLI.
`scripts/train.py`	JAX 학습 entrypoint. FSDP, EMA, Orbax checkpoint, wandb.
`scripts/train_pytorch.py`	PyTorch π0/π0.5 학습 entrypoint. DDP, safetensors checkpoint.
`scripts/train_test.py`	debug config 학습 smoke test.
`src/openpi/__init__.py`	openpi package init.
`src/openpi/conftest.py`	pytest/JAX platform setup. GPU가 없으면 CPU path로 테스트.
`src/openpi/models/__init__.py`	models package init.
`src/openpi/models/gemma.py`	π0/π0.5용 multi-expert Gemma 구현.
`src/openpi/models/gemma_fast.py`	π0-FAST용 single Gemma decoder 구현.
`src/openpi/models/lora.py`	LoRA config와 LoRA Einsum/FeedForward wrapper.
`src/openpi/models/lora_test.py`	LoRA shape/output equivalence test.
`src/openpi/models/model.py`	공통 `Observation`, `BaseModelConfig`, `BaseModel`, checkpoint restore.
`src/openpi/models/model_test.py`	모델 생성/loss/sample/restore smoke test.
`src/openpi/models/pi0.py`	π0/π0.5 JAX model. flow matching loss와 Euler sampling.
`src/openpi/models/pi0_config.py`	π0/π0.5 config, input spec, LoRA freeze filter.
`src/openpi/models/pi0_fast.py`	π0-FAST JAX model. autoregressive CE loss와 decoding.
`src/openpi/models/pi0_test.py`	π0 freeze filter/trainable parameter test.
`src/openpi/models/siglip.py`	SigLIP vision encoder.
`src/openpi/models/tokenizer.py`	PaliGemma/FAST/Binning/FSQ tokenizer wrapper.
`src/openpi/models/tokenizer_test.py`	tokenizer shape/roundtrip test.
`src/openpi/models/utils/fsq_tokenizer.py`	FSQ/LFQ action tokenizer baseline implementation.
`src/openpi/models/vit.py`	legacy ViT implementation. 현재 `openpi.models.resnet` 누락으로 메인 경로에서 사용하지 않는 파일로 봐야 한다.
`src/openpi/models_pytorch/gemma_pytorch.py`	PyTorch PaliGemma + action expert 결합 모델.
`src/openpi/models_pytorch/pi0_pytorch.py`	PyTorch π0/π0.5 flow matching model.
`src/openpi/models_pytorch/preprocessing_pytorch.py`	torch tensor image preprocess/augmentation.
`src/openpi/models_pytorch/transformers_replace/models/gemma/configuration_gemma.py`	transformers Gemma config patch.
`src/openpi/models_pytorch/transformers_replace/models/gemma/modeling_gemma.py`	transformers Gemma model patch.
`src/openpi/models_pytorch/transformers_replace/models/paligemma/modeling_paligemma.py`	transformers PaliGemma model patch.
`src/openpi/models_pytorch/transformers_replace/models/siglip/check.py`	transformers patch 설치 확인 helper.
`src/openpi/models_pytorch/transformers_replace/models/siglip/modeling_siglip.py`	transformers SigLIP model patch.
`src/openpi/policies/aloha_policy.py`	ALOHA observation/action space와 pi internal action space 변환.
`src/openpi/policies/droid_policy.py`	DROID observation/action transform.
`src/openpi/policies/libero_policy.py`	LIBERO observation/action transform.
`src/openpi/policies/policy.py`	runtime `Policy`, JAX/PyTorch inference dispatch, `PolicyRecorder`.
`src/openpi/policies/policy_config.py`	checkpoint + train config에서 inference `Policy` 생성.
`src/openpi/policies/policy_test.py`	policy inference와 action chunk broker test.
`src/openpi/py.typed`	package가 type hints를 제공함을 알리는 marker.
`src/openpi/serving/websocket_policy_server.py`	msgpack websocket server. metadata/healthz/infer 처리.
`src/openpi/shared/__init__.py`	shared package init.
`src/openpi/shared/array_typing.py`	jaxtyping/beartype helper와 pytree equality util.
`src/openpi/shared/download.py`	gs/fsspec/local download와 cache invalidation/permission 처리.
`src/openpi/shared/download_test.py`	download helper test. 일부 gs test는 네트워크/인증 영향 가능.
`src/openpi/shared/image_tools.py`	JAX/torch resize-with-pad 구현.
`src/openpi/shared/image_tools_test.py`	shared image resize test.
`src/openpi/shared/nnx_utils.py`	`module_jit`, regex path filter, state map helper.
`src/openpi/shared/normalize.py`	running stats, quantile stats, JSON serialization.
`src/openpi/shared/normalize_test.py`	normalization stats test.
`src/openpi/training/checkpoints.py`	Orbax checkpoint manager, EMA params split/merge, norm stats asset save/load.
`src/openpi/training/config.py`	모든 학습/추론 config registry와 data transform factory.
`src/openpi/training/data_loader.py`	Fake/LeRobot/RLDS dataset loader와 batch collation/sharding.
`src/openpi/training/data_loader_test.py`	data loader test. real dataset test는 manual 성격.
`src/openpi/training/droid_rlds_dataset.py`	dlimp/TFDS 기반 DROID RLDS pipeline.
`src/openpi/training/misc/polaris_config.py`	PolaRiS joint-position DROID co-training config.
`src/openpi/training/misc/roboarena_config.py`	RoboArena baseline config.
`src/openpi/training/optimizer.py`	cosine/rsqrt LR schedule, AdamW, SGD factory.
`src/openpi/training/sharding.py`	JAX mesh/FSDP sharding utilities.
`src/openpi/training/utils.py`	train state dataclass와 tree info printer.
`src/openpi/training/weight_loaders.py`	no-op/checkpoint/PaliGemma weight loader와 parameter merge.
`src/openpi/transforms.py`	transform DSL. repack, normalize, tokenize, delta/absolute actions, padding.
`src/openpi/transforms_test.py`	transform behavior test.
`uv.lock`	uv dependency lock. 연구 문서에는 필요 없지만 재현성에 중요하다.

3. 핵심 패키지 상세

3.1 `src/openpi/models/model.py`

이 파일은 모델 구현이 반드시 지켜야 하는 contract다.

ModelType은 config dispatch의 핵심이다. training/config.py의 ModelTransformFactory가 이 enum으로 π0, π0.5, π0-FAST transform을 선택한다.
Observation은 모든 모델 입력의 normalized canonical form이다.
Observation.from_dict()는 transform chain 이후의 dict를 dataclass로 바꾼다. 이때 numpy uint8 image는 [-1, 1] float로 바뀐다. torch uint8 image는 NHWC에서 NCHW로 permute 후 정규화한다.
preprocess_observation()은 inference/train 공통 image resize를 담당한다. training이면 non-wrist camera에는 crop/resize/rotate, 전체 camera에는 color jitter를 적용한다.
BaseModelConfig.load()는 model instance를 만들고 Orbax restored params를 주입한다. checkpoint에 model보다 더 많은 parameter가 있어도 intersection path를 골라 넣을 수 있다.
BaseModelConfig.load_pytorch()는 model.safetensors를 읽고 PyTorch 모델에 state dict를 load한다.

수정 포인트:

새 모델 타입을 추가하면 ModelType, BaseModelConfig.model_type, ModelTransformFactory, policy/data config dispatch를 함께 수정해야 한다.
image key를 바꾸면 IMAGE_KEYS, robot-specific policy, model input spec이 모두 맞아야 한다.

3.2 `src/openpi/models/pi0_config.py`

Pi0Config는 π0와 π0.5를 같은 dataclass로 표현한다.

중요 필드:

dtype: 기본 bfloat16.
paligemma_variant: 기본 gemma_2b, LoRA면 gemma_2b_lora.
action_expert_variant: 기본 gemma_300m, LoRA면 gemma_300m_lora.
action_dim: 기본 32. 로봇 action dim보다 클 수 있고 transform에서 padding한다.
action_horizon: 기본 50. DROID/LIBERO config에서는 10, 15, 16 등으로 줄인다.
max_token_len: π0는 48, π0.5는 200이 기본.
pi05: True면 ModelType.PI05.
discrete_state_input: 명시하지 않으면 π0.5에서 True.

get_freeze_filter()는 LoRA fine-tuning의 핵심이다. paligemma_variant나 action_expert_variant에 lora가 들어가면 LoRA path가 아닌 기존 LLM/action expert parameter를 freeze한다.

3.3 `src/openpi/models/pi0.py`

π0/π0.5의 실제 모델이다.

구성:

make_attn_mask(): prefix-LM 스타일 attention mask 생성.
posemb_sincos(): scalar timestep을 sinusoidal embedding으로 변환.
Pi0.__init__(): multi-expert Gemma와 SigLIP 생성, action projection layer 생성.
embed_prefix(): 모든 image view를 SigLIP token으로 바꾸고 prompt token embedding과 이어 붙인다.
embed_suffix(): state/action/time을 action expert token으로 만든다.
compute_loss(): flow matching 학습 loss.
sample_actions(): KV cache를 이용한 iterative denoising 추론.

수정 포인트:

새로운 camera view를 추가하려면 Observation.images dict를 바꾸는 transform과 model fake spec이 맞아야 한다.
action horizon/dim 변경은 config에서 하는 것이 원칙이다.
timestep sampling 분포나 solver step을 연구하려면 compute_loss()의 Beta sampling, sample_actions()의 Euler loop를 바꾸면 된다.

3.4 `src/openpi/models/pi0_fast.py`

π0-FAST autoregressive model이다.

구성:

make_attn_mask(): π0와 같은 누적 AR mask.
left_to_right_align(): batch마다 실제 prefix 길이가 다를 때 right-aligned decoding을 맞춘다.
put_along_last_axis(): JAX에 없는 last-axis put helper.
Pi0FASTConfig: action/token/image spec 정의.
Pi0FAST.embed_inputs(): image tokens와 tokenized prompt/action tokens를 하나로 합친다.
Pi0FAST.compute_loss(): next-token CE. action token 위치에만 token_loss_mask를 건다.
Pi0FAST.sample_actions(): autoregressive decode. 출력은 아직 raw token이며 transform이 action으로 decode한다.

수정 포인트:

max_token_len이 너무 작으면 FAST action tokens가 잘릴 수 있다.
action horizon/dim이 tokenizer decode shape와 맞아야 한다.
FAST는 PyTorch path가 없으므로 PyTorch 학습/추론 연구 대상이면 π0/π0.5로 시작해야 한다.

3.5 `src/openpi/models/gemma.py`

π0/π0.5의 multi-expert transformer다.

핵심 개념:

expert list를 받는다. 첫 expert는 PaliGemma backbone, 두 번째 expert는 action expert다.
_name(name, i)가 첫 expert는 원래 이름을 유지하고 두 번째부터 suffix를 붙인다. PaliGemma checkpoint loading과 action expert parameter 분리를 위해 중요하다.
Attention은 expert별 q/k/v projection을 만들고 attention 계산은 concatenated token sequence 위에서 수행한다.
RMSNorm은 π0.5에서 AdaRMS conditioning을 지원한다.
Block은 attention과 FFN을 반복하고 AdaRMS gate가 있으면 gated residual을 쓴다.
Module은 scan/remat, RoPE, KV cache를 포함한다.

수정 포인트:

action expert 크기 변경은 gemma.get_config() variant 추가와 checkpoint mapping까지 필요하다.
AdaRMS/knowledge-insulation 계열 실험은 여기와 pi0.py의 adarms_cond 경로를 함께 봐야 한다.

3.6 `src/openpi/models/gemma_fast.py`

π0-FAST 전용 single decoder다. gemma.py보다 단순하며 decode=True KV cache update path가 autoregressive generation에서 중요하다. LoRA는 attention/FFN에 들어갈 수 있다.

3.7 `src/openpi/models/siglip.py`

Big Vision 계열 SigLIP/ViT 구현이다. Module(variant="So400m/14", pool_type="none")로 호출되며, image를 patch embedding으로 바꾸고 transformer encoder를 통과시켜 token sequence를 반환한다. pool_type="none"이므로 class-level embedding이 아니라 dense visual token이 PaliGemma/Gemma로 넘어간다.

3.8 `src/openpi/models/tokenizer.py`

모델별 text/action tokenization 담당.

PaligemmaTokenizer: SentencePiece tokenizer를 사용한다. tokenizer model은 gs://big_vision/paligemma_tokenizer.model에서 cache된다.
π0 prompt는 보통 "<instruction>\n" 형태다.
π0.5의 discrete state prompt는 Task: ..., State: ...;\nAction: 형태로 state bins를 넣는다.
FASTTokenizer: physical-intelligence/fast AutoProcessor를 사용한다. training 때 action token을 PaliGemma vocab reserved slot으로 옮겨 넣고, inference output token을 다시 FAST token으로 되돌려 decode한다.
BinningTokenizer: OpenVLA/RT-2식 per-dimension binning baseline.
FSQTokenizer: FsqAttentionTokenizer checkpoint 기반 baseline.

3.9 `src/openpi/transforms.py`

저장소 전체에서 가장 중요한 비모델 파일 중 하나다.

주요 class:

Group: input/output transform 묶음.
CompositeTransform: transform sequence를 순서대로 실행.
RepackTransform: flatten key mapping으로 dict 구조 변환.
InjectDefaultPrompt: prompt가 없으면 기본 prompt 추가.
Normalize/Unnormalize: z-score 또는 quantile normalization.
ResizeImages: policy input image를 target resolution으로 바꾼다.
DeltaActions/AbsoluteActions: absolute action과 state-relative delta action 사이 변환.
TokenizePrompt: π0/π0.5 prompt tokenization.
TokenizeFASTInputs: π0-FAST prompt/state/action tokenization.
ExtractFASTActions: generated token을 continuous actions로 decode.
PromptFromLeRobotTask: LeRobot dataset의 task field를 prompt로 쓴다.
PadStatesAndActions: state/action 마지막 dim을 model action_dim까지 padding한다.

수정 포인트:

커스텀 데이터셋 key mismatch는 대부분 RepackTransform에서 해결한다.
action space mismatch는 DeltaActions, robot-specific outputs, PadStatesAndActions를 함께 확인한다.
norm stats key가 transform 이후 dict key와 일치해야 한다.

3.10 `src/openpi/training/config.py`

이 파일은 실행 가능한 모든 실험의 registry다.

중요 class:

AssetsConfig: norm stats 같은 asset 위치를 지정한다.
DataConfig: dataset repo id, norm stats, repack/data/model transforms, action sequence key, RLDS options를 담는다.
ModelTransformFactory: model type별 기본 transform을 생성한다.
FakeDataConfig: debug/test용 fake data.
SimpleDataConfig: lambda로 data/model transform을 넣는 범용 factory.
LeRobotAlohaDataConfig, LeRobotLiberoDataConfig, LeRobotDROIDDataConfig: dataset별 factory.
RLDSDroidDataConfig: full DROID RLDS training용 factory.
TrainConfig: model, data, optimizer, checkpoint, FSDP, wandb 등 전체 실험 설정.

새 config를 만들 때 가장 안전한 방법은 기존 config를 복사해 name, repo_id, data_transforms, weight_loader, num_train_steps, batch_size, assets만 바꾸는 것이다.

3.11 `src/openpi/training/data_loader.py`

학습 데이터 로딩 계층이다.

FakeDataset: model fake spec 기반 random data.
create_torch_dataset(): LeRobot dataset을 만든다.
create_rlds_dataset(): DROID RLDS iterable dataset을 만든다.
transform_dataset()/transform_iterable_dataset(): repack/data/norm/model transform을 dataset에 적용한다.
TorchDataLoader: PyTorch DataLoader 결과를 JAX sharded array 또는 torch tensor로 변환한다.
RLDSDataLoader: prebatched RLDS dataset을 감싼다.
DataLoaderImpl: (Observation, actions) tuple을 반환하는 최종 wrapper.

3.12 `src/openpi/training/droid_rlds_dataset.py`

full DROID RLDS training용 pipeline이다. TensorFlow가 GPU를 잡지 않도록 숨기고, dlimp/TFDS를 사용한다.

기능:

여러 RLDS dataset을 weight로 섞는다.
success trajectory 필터링.
idle filter JSON range 적용.
action/image/language instruction restructuring.
action chunk sampling.
image decode와 batch/prefetch.

3.13 `src/openpi/training/checkpoints.py`

Orbax checkpoint 관리.

checkpoint dir overwrite/resume 검증.
assets, train_state, params item 저장.
EMA가 있으면 inference params를 분리해 저장한다.
norm stats를 checkpoint asset으로 함께 저장해 inference가 학습 당시 normalization을 쓰게 한다.

3.14 `src/openpi/policies`

추론 path의 중심이다.

policy_config.create_trained_policy()는 checkpoint를 JAX/PyTorch로 자동 감지하고, model/data/norm/model transform을 조립해 Policy를 반환한다.
policy.Policy.infer()는 입력 dict를 transform하고 batch dimension을 추가한 뒤 model sample_actions()를 호출한다.
aloha_policy.py, droid_policy.py, libero_policy.py는 각각 로봇/벤치마크 observation 형식을 canonical model input으로 바꾼다.

3.15 `src/openpi/serving`와 `packages/openpi-client`

서버:

scripts/serve_policy.py가 config/checkpoint를 받아 policy를 만들고 server를 띄운다.
websocket_policy_server.py가 msgpack numpy protocol로 observation을 받고 action dict를 반환한다.

클라이언트:

WebsocketClientPolicy는 server metadata를 받고 infer() 호출마다 observation을 msgpack으로 전송한다.
ActionChunkBroker는 policy가 반환한 action chunk를 cache하고 step 단위로 나눠준다.
Runtime은 environment와 agent를 묶어 loop를 돈다.

4. 예제별 실행 경로

4.1 Simple client

examples/simple_client/main.py는 가장 안전한 smoke test다. 실제 robot 없이 random observation을 만들고 websocket server에 보낸다.

실행 경로:

scripts/serve_policy.py
  -> policy_config.create_trained_policy
  -> websocket_policy_server.serve_forever

examples/simple_client/main.py
  -> WebsocketClientPolicy
  -> random observation
  -> policy.infer over websocket

4.2 ALOHA sim

examples/aloha_sim/main.py는 gym_aloha environment를 AlohaSimEnvironment로 감싸고, ActionChunkBroker(WebsocketClientPolicy)를 runtime agent로 사용한다. VideoSaver가 episode 영상을 저장한다.

4.3 ALOHA real

실제 ROS/Interbotix 환경이다. real_env.py가 low-level robot command를 다루고 env.py가 runtime abstraction에 맞춘다. AlohaInputs/AlohaOutputs가 gripper와 joint convention을 pi internal convention으로 바꿀 수 있다.

4.4 DROID

examples/droid/main.py는 실제 DROID robot에서 camera/joint/gripper observation을 추출하고 policy server에 보내 action chunk를 받는다. open_loop_horizon마다 새 action chunk를 요청하고, action은 joint velocity + gripper command로 실행한다.

4.5 LIBERO

examples/libero/main.py는 LIBERO task suite를 돌며 policy server를 호출한다. image preprocessing에서 180도 회전 등 benchmark-specific 처리가 들어간다. Docker workflow가 권장된다.

4.6 JAX to PyTorch 변환

examples/convert_jax_model_to_pytorch.py는 JAX checkpoint tree key를 Hugging Face/PyTorch key로 mapping한다. SigLIP/PaliGemma/Gemma/action expert parameter 이름과 shape transpose가 많으므로, 새 model variant를 만들면 이 변환 script도 함께 수정해야 한다.

5. 테스트 범위

테스트는 src, scripts, packages 아래에 분산되어 있다.

모델 smoke: model_test.py, pi0_test.py, lora_test.py, tokenizer_test.py.
transform: transforms_test.py.
policy: policy_test.py.
data loader: data_loader_test.py.
training: scripts/train_test.py.
shared utilities: download_test.py, image_tools_test.py, normalize_test.py.
client utilities: packages/openpi-client/src/openpi_client/*_test.py.

문서만 수정했다면 필수 테스트는 아니지만, 코드 변경 후 기본 확인은 다음 순서가 좋다.

uv run pytest src/openpi/transforms_test.py
uv run pytest src/openpi/models/model_test.py
uv run pytest scripts/train_test.py

체크포인트 다운로드나 gs 접근이 필요한 테스트는 네트워크/인증/GPU 조건에 영향을 받는다.

openpi 코드 레퍼런스

1. 폴더 인벤토리

2. 파일 인벤토리

3. 핵심 패키지 상세

3.1 src/openpi/models/model.py

3.2 src/openpi/models/pi0_config.py

3.3 src/openpi/models/pi0.py

3.4 src/openpi/models/pi0_fast.py

3.5 src/openpi/models/gemma.py

3.6 src/openpi/models/gemma_fast.py

3.7 src/openpi/models/siglip.py

3.8 src/openpi/models/tokenizer.py

3.9 src/openpi/transforms.py

3.10 src/openpi/training/config.py

3.11 src/openpi/training/data_loader.py

3.12 src/openpi/training/droid_rlds_dataset.py

3.13 src/openpi/training/checkpoints.py

3.14 src/openpi/policies

3.15 src/openpi/serving와 packages/openpi-client

4. 예제별 실행 경로

4.1 Simple client

4.2 ALOHA sim

4.3 ALOHA real

4.4 DROID

4.5 LIBERO

4.6 JAX to PyTorch 변환

5. 테스트 범위

3.1 `src/openpi/models/model.py`

3.2 `src/openpi/models/pi0_config.py`

3.3 `src/openpi/models/pi0.py`

3.4 `src/openpi/models/pi0_fast.py`

3.5 `src/openpi/models/gemma.py`

3.6 `src/openpi/models/gemma_fast.py`

3.7 `src/openpi/models/siglip.py`

3.8 `src/openpi/models/tokenizer.py`

3.9 `src/openpi/transforms.py`

3.10 `src/openpi/training/config.py`

3.11 `src/openpi/training/data_loader.py`

3.12 `src/openpi/training/droid_rlds_dataset.py`

3.13 `src/openpi/training/checkpoints.py`

3.14 `src/openpi/policies`

3.15 `src/openpi/serving`와 `packages/openpi-client`