From d12c4ecdfaddf8dbd7f42f17395df569d360093f Mon Sep 17 00:00:00 2001 From: Chernobyllight Date: Thu, 3 Sep 2026 11:11:00 +0000 Subject: [PATCH 1/6] feat: integrate OpenPI LIBERO inference Integrate OpenPI pi0.5-LIBERO with local PyTorch inference, checkpoint conversion, runtime setup, and official single- and multi-GPU evaluation workflows. --- configs/openpi/pi05_libero.json | 15 + configs/openpi/pi05_libero_eval.json | 26 + lightx2v/infer.py | 2 + lightx2v/models/networks/openpi/NOTICE.md | 11 + lightx2v/models/networks/openpi/__init__.py | 7 + lightx2v/models/networks/openpi/config.py | 77 + lightx2v/models/networks/openpi/gemma.py | 244 ++++ .../models/networks/openpi/image_tools.py | 47 + .../models/networks/openpi/infer/__init__.py | 5 + .../networks/openpi/infer/post_infer.py | 26 + .../models/networks/openpi/infer/pre_infer.py | 125 ++ .../openpi/infer/transformer_infer.py | 24 + lightx2v/models/networks/openpi/model.py | 130 ++ .../models/networks/openpi/observation.py | 19 + lightx2v/models/networks/openpi/pi0.py | 423 ++++++ .../models/networks/openpi/preprocessing.py | 130 ++ .../models/gemma/configuration_gemma.py | 173 +++ .../models/gemma/modeling_gemma.py | 862 ++++++++++++ .../models/paligemma/modeling_paligemma.py | 622 +++++++++ .../models/siglip/check.py | 4 + .../models/siglip/modeling_siglip.py | 1237 +++++++++++++++++ .../networks/openpi/weights/__init__.py | 3 + .../models/networks/openpi/weights/loader.py | 59 + lightx2v/models/runners/openpi/__init__.py | 1 + lightx2v/models/runners/openpi/artifacts.py | 176 +++ .../models/runners/openpi/libero_evaluate.py | 342 +++++ .../models/runners/openpi/libero_protocol.py | 721 ++++++++++ .../models/runners/openpi/libero_rollout.py | 183 +++ .../models/runners/openpi/openpi_runner.py | 308 ++++ lightx2v/pipeline.py | 1 + pyproject.toml | 1 + .../1_convert_pi05_libero_to_pytorch.sh | 8 + scripts/openpi/2_setup_pytorch_runtime.sh | 14 + scripts/openpi/README.md | 213 +++ scripts/openpi/convert_jax_checkpoint.py | 207 +++ scripts/openpi/libero_summary.py | 145 ++ scripts/openpi/run_libero_evaluate_i2va.sh | 38 + .../run_libero_evaluate_parallel_i2va.sh | 137 ++ scripts/openpi/run_libero_i2va.sh | 41 + scripts/openpi/run_libero_task_i2va.sh | 42 + scripts/openpi/runtime.py | 498 +++++++ .../openpi/tests/test_task_inputs_manifest.py | 166 +++ .../openpi/tests/validate_pytorch_parity.py | 375 +++++ 43 files changed, 7888 insertions(+) create mode 100644 configs/openpi/pi05_libero.json create mode 100644 configs/openpi/pi05_libero_eval.json create mode 100644 lightx2v/models/networks/openpi/NOTICE.md create mode 100644 lightx2v/models/networks/openpi/__init__.py create mode 100644 lightx2v/models/networks/openpi/config.py create mode 100644 lightx2v/models/networks/openpi/gemma.py create mode 100644 lightx2v/models/networks/openpi/image_tools.py create mode 100644 lightx2v/models/networks/openpi/infer/__init__.py create mode 100644 lightx2v/models/networks/openpi/infer/post_infer.py create mode 100644 lightx2v/models/networks/openpi/infer/pre_infer.py create mode 100644 lightx2v/models/networks/openpi/infer/transformer_infer.py create mode 100644 lightx2v/models/networks/openpi/model.py create mode 100644 lightx2v/models/networks/openpi/observation.py create mode 100644 lightx2v/models/networks/openpi/pi0.py create mode 100644 lightx2v/models/networks/openpi/preprocessing.py create mode 100644 lightx2v/models/networks/openpi/transformers_replace/models/gemma/configuration_gemma.py create mode 100644 lightx2v/models/networks/openpi/transformers_replace/models/gemma/modeling_gemma.py create mode 100644 lightx2v/models/networks/openpi/transformers_replace/models/paligemma/modeling_paligemma.py create mode 100644 lightx2v/models/networks/openpi/transformers_replace/models/siglip/check.py create mode 100644 lightx2v/models/networks/openpi/transformers_replace/models/siglip/modeling_siglip.py create mode 100644 lightx2v/models/networks/openpi/weights/__init__.py create mode 100644 lightx2v/models/networks/openpi/weights/loader.py create mode 100644 lightx2v/models/runners/openpi/__init__.py create mode 100644 lightx2v/models/runners/openpi/artifacts.py create mode 100644 lightx2v/models/runners/openpi/libero_evaluate.py create mode 100644 lightx2v/models/runners/openpi/libero_protocol.py create mode 100644 lightx2v/models/runners/openpi/libero_rollout.py create mode 100644 lightx2v/models/runners/openpi/openpi_runner.py create mode 100755 scripts/openpi/1_convert_pi05_libero_to_pytorch.sh create mode 100755 scripts/openpi/2_setup_pytorch_runtime.sh create mode 100644 scripts/openpi/README.md create mode 100755 scripts/openpi/convert_jax_checkpoint.py create mode 100755 scripts/openpi/libero_summary.py create mode 100755 scripts/openpi/run_libero_evaluate_i2va.sh create mode 100755 scripts/openpi/run_libero_evaluate_parallel_i2va.sh create mode 100755 scripts/openpi/run_libero_i2va.sh create mode 100755 scripts/openpi/run_libero_task_i2va.sh create mode 100755 scripts/openpi/runtime.py create mode 100644 scripts/openpi/tests/test_task_inputs_manifest.py create mode 100644 scripts/openpi/tests/validate_pytorch_parity.py diff --git a/configs/openpi/pi05_libero.json b/configs/openpi/pi05_libero.json new file mode 100644 index 000000000..bac95ffe1 --- /dev/null +++ b/configs/openpi/pi05_libero.json @@ -0,0 +1,15 @@ +{ + "pi05": true, + "discrete_state_input": false, + "paligemma_variant": "gemma_2b", + "action_expert_variant": "gemma_300m", + "action_dim": 32, + "output_action_dim": 7, + "state_dim": 8, + "action_horizon": 10, + "max_token_len": 200, + "num_inference_steps": 10, + "device": "cuda", + "dtype": "bfloat16", + "pytorch_compile_mode": null +} diff --git a/configs/openpi/pi05_libero_eval.json b/configs/openpi/pi05_libero_eval.json new file mode 100644 index 000000000..3750c76f6 --- /dev/null +++ b/configs/openpi/pi05_libero_eval.json @@ -0,0 +1,26 @@ +{ + "benchmarks": [ + "libero_spatial", + "libero_object", + "libero_goal", + "libero_10" + ], + "task_ids": "all", + "num_trials_per_task": 50, + "env_seed": 7, + "policy_seed": 0, + "actions_per_plan": 5, + "num_steps_wait": 10, + "render_size": 256, + "video_fps": 10, + "video_policy": "none", + "save_actions": false, + "resume": true, + "fail_fast": false, + "max_steps": { + "libero_spatial": 220, + "libero_object": 280, + "libero_goal": 300, + "libero_10": 520 + } +} diff --git a/lightx2v/infer.py b/lightx2v/infer.py index f45729579..47251b499 100755 --- a/lightx2v/infer.py +++ b/lightx2v/infer.py @@ -24,6 +24,7 @@ from lightx2v.models.runners.minimax_h3.minimax_h3_runner import MiniMaxH3Runner # noqa: F401 from lightx2v.models.runners.motus.motus_runner import MotusRunner # noqa: F401 from lightx2v.models.runners.neopp.neopp_runner import NeoppRunner # noqa: F401 +from lightx2v.models.runners.openpi.openpi_runner import OpenPIRunner # noqa: F401 from lightx2v.models.runners.qwen_image.qwen_image_runner import QwenImageRunner # noqa: F401 from lightx2v.models.runners.seedvr.seedvr_runner import SeedVRRunner # noqa: F401 from lightx2v.models.runners.swiftvr.swiftvr_runner import SwiftVRRunner # noqa: F401 @@ -128,6 +129,7 @@ def main(): "seedvr2", "swiftvr", "neopp", + "openpi", "motus", "lingbot_world_fast", "worldmirror", diff --git a/lightx2v/models/networks/openpi/NOTICE.md b/lightx2v/models/networks/openpi/NOTICE.md new file mode 100644 index 000000000..ce898ed15 --- /dev/null +++ b/lightx2v/models/networks/openpi/NOTICE.md @@ -0,0 +1,11 @@ +# OpenPI attribution + +The `pi0.py`, `gemma.py`, and `preprocessing.py` implementation in this +directory is adapted from Physical Intelligence's OpenPI project at commit +`15a9616a00943ada6c20a0f158e3adb39df2ccac`. The files under +`transformers_replace/` are vendored from that revision without behavioral changes. + +OpenPI and the copied Hugging Face Transformers source files are distributed +under the Apache License 2.0. The localization changes replace OpenPI/JAX +imports with LightX2V-local, PyTorch-only modules while intentionally retaining +the official model parameter names for SafeTensors compatibility. diff --git a/lightx2v/models/networks/openpi/__init__.py b/lightx2v/models/networks/openpi/__init__.py new file mode 100644 index 000000000..2b1bd70e4 --- /dev/null +++ b/lightx2v/models/networks/openpi/__init__.py @@ -0,0 +1,7 @@ +"""Native PyTorch OpenPI network family for LightX2V.""" + +from .config import Pi0Config +from .model import OpenPIModel +from .observation import Observation + +__all__ = ["Observation", "OpenPIModel", "Pi0Config"] diff --git a/lightx2v/models/networks/openpi/config.py b/lightx2v/models/networks/openpi/config.py new file mode 100644 index 000000000..9b39a391d --- /dev/null +++ b/lightx2v/models/networks/openpi/config.py @@ -0,0 +1,77 @@ +"""Configuration for the PyTorch pi0.5-LIBERO backend.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any, Literal + + +@dataclass(frozen=True) +class GemmaConfig: + width: int + depth: int + mlp_dim: int + num_heads: int + num_kv_heads: int + head_dim: int + + +GemmaVariant = Literal["dummy", "gemma_300m", "gemma_2b"] + + +def get_config(variant: GemmaVariant) -> GemmaConfig: + if variant == "dummy": + return GemmaConfig(width=64, depth=4, mlp_dim=128, num_heads=8, num_kv_heads=1, head_dim=16) + if variant == "gemma_300m": + return GemmaConfig(width=1024, depth=18, mlp_dim=4096, num_heads=8, num_kv_heads=1, head_dim=256) + if variant == "gemma_2b": + return GemmaConfig(width=2048, depth=18, mlp_dim=16384, num_heads=8, num_kv_heads=1, head_dim=256) + raise ValueError(f"Unsupported OpenPI Gemma variant: {variant!r}") + + +@dataclass(frozen=True) +class Pi0Config: + """Dimensions and runtime options for pi0.5-LIBERO.""" + + action_dim: int = 32 + action_horizon: int = 10 + max_token_len: int = 200 + dtype: Literal["bfloat16", "float32"] = "bfloat16" + paligemma_variant: GemmaVariant = "gemma_2b" + action_expert_variant: GemmaVariant = "gemma_300m" + pi05: bool = True + discrete_state_input: bool = False + pytorch_compile_mode: str | None = None + + @classmethod + def from_mapping(cls, config: Mapping[str, Any]) -> "Pi0Config": + return cls( + action_dim=config["action_dim"], + action_horizon=config["action_horizon"], + max_token_len=config["max_token_len"], + dtype=config["dtype"], + paligemma_variant=config["paligemma_variant"], + action_expert_variant=config["action_expert_variant"], + pi05=config["pi05"], + discrete_state_input=config["discrete_state_input"], + pytorch_compile_mode=config["pytorch_compile_mode"], + ) + + def validate_pi05_libero(self) -> None: + expected = { + "pi05": True, + "paligemma_variant": "gemma_2b", + "action_expert_variant": "gemma_300m", + "action_dim": 32, + "action_horizon": 10, + "max_token_len": 200, + "discrete_state_input": False, + } + actual = {name: getattr(self, name) for name in expected} + wrong = {name: (actual[name], value) for name, value in expected.items() if actual[name] != value} + if wrong: + details = ", ".join(f"{name}={got!r} (expected {want!r})" for name, (got, want) in wrong.items()) + raise ValueError(f"Configuration does not match the released pi05_libero checkpoint: {details}") + if self.dtype not in {"bfloat16", "float32"}: + raise ValueError(f"Unsupported OpenPI dtype: {self.dtype!r}") diff --git a/lightx2v/models/networks/openpi/gemma.py b/lightx2v/models/networks/openpi/gemma.py new file mode 100644 index 000000000..21c1a05ea --- /dev/null +++ b/lightx2v/models/networks/openpi/gemma.py @@ -0,0 +1,244 @@ +# Adapted from Physical Intelligence OpenPI (Apache-2.0), commit 15a9616. +# Localized for the LightX2V OpenPI backend; no runtime OpenPI/JAX dependency. + +from typing import Literal + +import torch +from torch import nn +from transformers import GemmaForCausalLM, PaliGemmaForConditionalGeneration +from transformers.models.auto import CONFIG_MAPPING +from transformers.models.gemma import modeling_gemma + + +class PaliGemmaWithExpertModel(nn.Module): + def __init__( + self, + vlm_config, + action_expert_config, + use_adarms=None, + precision: Literal["bfloat16", "float32"] = "bfloat16", + ): + if use_adarms is None: + use_adarms = [False, False] + super().__init__() + + vlm_config_hf = CONFIG_MAPPING["paligemma"]() + vlm_config_hf._vocab_size = 257152 # noqa: SLF001 + vlm_config_hf.image_token_index = 257152 + vlm_config_hf.text_config.hidden_size = vlm_config.width + vlm_config_hf.text_config.intermediate_size = vlm_config.mlp_dim + vlm_config_hf.text_config.num_attention_heads = vlm_config.num_heads + vlm_config_hf.text_config.head_dim = vlm_config.head_dim + vlm_config_hf.text_config.num_hidden_layers = vlm_config.depth + vlm_config_hf.text_config.num_key_value_heads = vlm_config.num_kv_heads + vlm_config_hf.text_config.hidden_activation = "gelu_pytorch_tanh" + vlm_config_hf.text_config.torch_dtype = "float32" + vlm_config_hf.text_config.vocab_size = 257152 + vlm_config_hf.text_config.use_adarms = use_adarms[0] + vlm_config_hf.text_config.adarms_cond_dim = vlm_config.width if use_adarms[0] else None + vlm_config_hf.vision_config.intermediate_size = 4304 + vlm_config_hf.vision_config.projection_dim = 2048 + vlm_config_hf.vision_config.projector_hidden_act = "gelu_fast" + vlm_config_hf.vision_config.torch_dtype = "float32" + + action_expert_config_hf = CONFIG_MAPPING["gemma"]( + head_dim=action_expert_config.head_dim, + hidden_size=action_expert_config.width, + intermediate_size=action_expert_config.mlp_dim, + num_attention_heads=action_expert_config.num_heads, + num_hidden_layers=action_expert_config.depth, + num_key_value_heads=action_expert_config.num_kv_heads, + vocab_size=257152, + hidden_activation="gelu_pytorch_tanh", + torch_dtype="float32", + use_adarms=use_adarms[1], + adarms_cond_dim=action_expert_config.width if use_adarms[1] else None, + ) + + self.paligemma = PaliGemmaForConditionalGeneration(config=vlm_config_hf) + self.gemma_expert = GemmaForCausalLM(config=action_expert_config_hf) + self.gemma_expert.model.embed_tokens = None + + self.to_bfloat16_for_selected_params(precision) + + def to_bfloat16_for_selected_params(self, precision: Literal["bfloat16", "float32"] = "bfloat16"): + if precision == "bfloat16": + self.to(dtype=torch.bfloat16) + elif precision == "float32": + self.to(dtype=torch.float32) + return + else: + raise ValueError(f"Invalid precision: {precision}") + + params_to_keep_float32 = [ + "vision_tower.vision_model.embeddings.patch_embedding.weight", + "vision_tower.vision_model.embeddings.patch_embedding.bias", + "vision_tower.vision_model.embeddings.position_embedding.weight", + "input_layernorm", + "post_attention_layernorm", + "model.norm", + ] + + for name, param in self.named_parameters(): + if any(selector in name for selector in params_to_keep_float32): + param.data = param.data.to(dtype=torch.float32) + + def embed_image(self, image: torch.Tensor): + return self.paligemma.model.get_image_features(image) + + def embed_language_tokens(self, tokens: torch.Tensor): + return self.paligemma.language_model.embed_tokens(tokens) + + def forward( + self, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + inputs_embeds: list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + adarms_cond: list[torch.Tensor] | None = None, + ): + if adarms_cond is None: + adarms_cond = [None, None] + if inputs_embeds[1] is None: + prefix_output = self.paligemma.language_model.forward( + inputs_embeds=inputs_embeds[0], + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + adarms_cond=adarms_cond[0] if adarms_cond is not None else None, + ) + prefix_past_key_values = prefix_output.past_key_values + prefix_output = prefix_output.last_hidden_state + suffix_output = None + elif inputs_embeds[0] is None: + suffix_output = self.gemma_expert.model.forward( + inputs_embeds=inputs_embeds[1], + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + adarms_cond=adarms_cond[1] if adarms_cond is not None else None, + ) + suffix_output = suffix_output.last_hidden_state + prefix_output = None + prefix_past_key_values = None + else: + models = [self.paligemma.language_model, self.gemma_expert.model] + num_layers = self.paligemma.config.text_config.num_hidden_layers + + use_gradient_checkpointing = (hasattr(self.gemma_expert.model, "gradient_checkpointing") and self.gemma_expert.model.gradient_checkpointing and self.training) or ( + hasattr(self, "gradient_checkpointing") and self.gradient_checkpointing and self.training + ) + + if self.training and hasattr(self.gemma_expert.model, "gradient_checkpointing"): + if not self.gemma_expert.model.gradient_checkpointing: + self.gemma_expert.model.gradient_checkpointing = True + use_gradient_checkpointing = True + + def compute_layer_complete(layer_idx, inputs_embeds, attention_mask, position_ids, adarms_cond): + models = [self.paligemma.language_model, self.gemma_expert.model] + + query_states = [] + key_states = [] + value_states = [] + gates = [] + for i, hidden_states in enumerate(inputs_embeds): + layer = models[i].layers[layer_idx] + hidden_states, gate = layer.input_layernorm(hidden_states, cond=adarms_cond[i]) # noqa: PLW2901 + gates.append(gate) + + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, layer.self_attn.head_dim) + query_state = layer.self_attn.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_state = layer.self_attn.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_state = layer.self_attn.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + query_states.append(query_state) + key_states.append(key_state) + value_states.append(value_state) + + query_states = torch.cat(query_states, dim=2) + key_states = torch.cat(key_states, dim=2) + value_states = torch.cat(value_states, dim=2) + + dummy_tensor = torch.zeros( + query_states.shape[0], + query_states.shape[2], + query_states.shape[-1], + device=query_states.device, + dtype=query_states.dtype, + ) + cos, sin = self.paligemma.model.language_model.rotary_emb(dummy_tensor, position_ids) + query_states, key_states = modeling_gemma.apply_rotary_pos_emb(query_states, key_states, cos, sin, unsqueeze_dim=1) + + batch_size = query_states.shape[0] + scaling = self.paligemma.language_model.layers[layer_idx].self_attn.scaling + + att_output, _ = modeling_gemma.eager_attention_forward( + self.paligemma.language_model.layers[layer_idx].self_attn, + query_states, + key_states, + value_states, + attention_mask, + scaling, + ) + head_dim = self.paligemma.language_model.layers[layer_idx].self_attn.head_dim + att_output = att_output.reshape(batch_size, -1, 1 * 8 * head_dim) + + outputs_embeds = [] + start_pos = 0 + for i, hidden_states in enumerate(inputs_embeds): + layer = models[i].layers[layer_idx] + end_pos = start_pos + hidden_states.shape[1] + + if att_output.dtype != layer.self_attn.o_proj.weight.dtype: + att_output = att_output.to(layer.self_attn.o_proj.weight.dtype) + out_emb = layer.self_attn.o_proj(att_output[:, start_pos:end_pos]) + + out_emb = modeling_gemma._gated_residual(hidden_states, out_emb, gates[i]) # noqa: SLF001 + after_first_residual = out_emb.clone() + out_emb, gate = layer.post_attention_layernorm(out_emb, cond=adarms_cond[i]) + if layer.mlp.up_proj.weight.dtype == torch.bfloat16: + out_emb = out_emb.to(dtype=torch.bfloat16) + + out_emb = layer.mlp(out_emb) + out_emb = modeling_gemma._gated_residual(after_first_residual, out_emb, gate) # noqa: SLF001 + outputs_embeds.append(out_emb) + start_pos = end_pos + + return outputs_embeds + + for layer_idx in range(num_layers): + if use_gradient_checkpointing: + inputs_embeds = torch.utils.checkpoint.checkpoint( + compute_layer_complete, + layer_idx, + inputs_embeds, + attention_mask, + position_ids, + adarms_cond, + use_reentrant=False, + preserve_rng_state=False, + ) + else: + inputs_embeds = compute_layer_complete(layer_idx, inputs_embeds, attention_mask, position_ids, adarms_cond) + + def compute_final_norms(inputs_embeds, adarms_cond): + outputs_embeds = [] + for i, hidden_states in enumerate(inputs_embeds): + out_emb, _ = models[i].norm(hidden_states, cond=adarms_cond[i]) + outputs_embeds.append(out_emb) + return outputs_embeds + + if use_gradient_checkpointing: + outputs_embeds = torch.utils.checkpoint.checkpoint(compute_final_norms, inputs_embeds, adarms_cond, use_reentrant=False, preserve_rng_state=False) + else: + outputs_embeds = compute_final_norms(inputs_embeds, adarms_cond) + + prefix_output = outputs_embeds[0] + suffix_output = outputs_embeds[1] + prefix_past_key_values = None + + return [prefix_output, suffix_output], prefix_past_key_values diff --git a/lightx2v/models/networks/openpi/image_tools.py b/lightx2v/models/networks/openpi/image_tools.py new file mode 100644 index 000000000..7d03e2957 --- /dev/null +++ b/lightx2v/models/networks/openpi/image_tools.py @@ -0,0 +1,47 @@ +"""Image helpers copied from OpenPI's torch preprocessing path (Apache-2.0).""" + +import torch +import torch.nn.functional as F # noqa: N812 + + +def resize_with_pad_torch(images: torch.Tensor, height: int, width: int, mode: str = "bilinear") -> torch.Tensor: + """Resize without distortion and pad with black / -1, preserving layout.""" + input_was_unbatched = images.dim() == 3 + channels_last = images.shape[-1] <= 4 + if input_was_unbatched: + images = images.unsqueeze(0) + if channels_last: + images = images.permute(0, 3, 1, 2) + + _, _, current_height, current_width = images.shape + ratio = max(current_width / width, current_height / height) + resized_height = int(current_height / ratio) + resized_width = int(current_width / ratio) + resized = F.interpolate( + images, + size=(resized_height, resized_width), + mode=mode, + align_corners=False if mode == "bilinear" else None, + ) + if images.dtype == torch.uint8: + resized = torch.round(resized).clamp(0, 255).to(torch.uint8) + pad_value = 0 + elif images.dtype == torch.float32: + resized = resized.clamp(-1.0, 1.0) + pad_value = -1.0 + else: + raise ValueError(f"Unsupported image dtype: {images.dtype}") + + pad_h0, remainder_h = divmod(height - resized_height, 2) + pad_w0, remainder_w = divmod(width - resized_width, 2) + resized = F.pad( + resized, + (pad_w0, pad_w0 + remainder_w, pad_h0, pad_h0 + remainder_h), + mode="constant", + value=pad_value, + ) + if channels_last: + resized = resized.permute(0, 2, 3, 1) + if input_was_unbatched: + resized = resized.squeeze(0) + return resized diff --git a/lightx2v/models/networks/openpi/infer/__init__.py b/lightx2v/models/networks/openpi/infer/__init__.py new file mode 100644 index 000000000..370eaf0f9 --- /dev/null +++ b/lightx2v/models/networks/openpi/infer/__init__.py @@ -0,0 +1,5 @@ +from .post_infer import OpenPIPostInfer +from .pre_infer import OpenPIPreInfer +from .transformer_infer import OpenPITransformerInfer + +__all__ = ["OpenPIPostInfer", "OpenPIPreInfer", "OpenPITransformerInfer"] diff --git a/lightx2v/models/networks/openpi/infer/post_infer.py b/lightx2v/models/networks/openpi/infer/post_infer.py new file mode 100644 index 000000000..609f3d47b --- /dev/null +++ b/lightx2v/models/networks/openpi/infer/post_infer.py @@ -0,0 +1,26 @@ +"""Convert normalized model actions to LIBERO's 7-D action space.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import torch + +from .pre_infer import load_norm_stats + +LIBERO_ACTION_DIM = 7 + + +class OpenPIPostInfer: + def __init__(self, norm_stats_path: str | Path): + self.stats = load_norm_stats(norm_stats_path)["actions"] + + def infer(self, actions: torch.Tensor) -> np.ndarray: + if actions.ndim != 3 or actions.shape[0] != 1 or actions.shape[-1] < LIBERO_ACTION_DIM: + raise ValueError(f"Expected actions with shape [1, horizon, padded_action_dim], got {tuple(actions.shape)}") + normalized = actions[0, :, :LIBERO_ACTION_DIM].detach().to(torch.float32).cpu().numpy() + q01 = self.stats["q01"][:LIBERO_ACTION_DIM] + q99 = self.stats["q99"][:LIBERO_ACTION_DIM] + physical = (normalized + 1.0) / 2.0 * (q99 - q01 + 1e-6) + q01 + return np.asarray(physical) diff --git a/lightx2v/models/networks/openpi/infer/pre_infer.py b/lightx2v/models/networks/openpi/infer/pre_infer.py new file mode 100644 index 000000000..9624b5af6 --- /dev/null +++ b/lightx2v/models/networks/openpi/infer/pre_infer.py @@ -0,0 +1,125 @@ +"""LIBERO input construction for the native PyTorch OpenPI backend.""" + +from __future__ import annotations + +import json +import logging +from pathlib import Path + +import numpy as np +import sentencepiece +import torch +from PIL import Image + +from ..observation import Observation + +LOGGER = logging.getLogger(__name__) +IMAGE_SIZE = 224 +LIBERO_STATE_DIM = 8 + + +def load_norm_stats(path: str | Path) -> dict[str, dict[str, np.ndarray]]: + with Path(path).open("r", encoding="utf-8") as handle: + payload = json.load(handle)["norm_stats"] + if "state" not in payload or "actions" not in payload: + raise ValueError(f"Invalid LIBERO norm_stats file: {path}") + # Match upstream by preserving the JSON quantiles as float64. + return {key: {stat: np.asarray(value) for stat, value in stats.items()} for key, stats in payload.items()} + + +def normalize_quantile(value: np.ndarray, stats: dict[str, np.ndarray]) -> np.ndarray: + q01 = stats["q01"][: value.shape[-1]] + q99 = stats["q99"][: value.shape[-1]] + return (value - q01) / (q99 - q01 + 1e-6) * 2.0 - 1.0 + + +def _require_rgb(image: np.ndarray) -> np.ndarray: + array = np.asarray(image) + if array.dtype != np.uint8 or array.ndim != 3 or array.shape[-1] != 3: + raise ValueError(f"Expected an HWC uint8 RGB image, got shape={array.shape}, dtype={array.dtype}") + return np.ascontiguousarray(array) + + +def _resize_with_pad(image: np.ndarray, size: int = IMAGE_SIZE) -> np.ndarray: + height, width = image.shape[:2] + if (height, width) == (size, size): + return np.array(image, copy=True) + ratio = max(width / size, height / size) + resized_height = int(height / ratio) + resized_width = int(width / ratio) + resized = Image.fromarray(image, mode="RGB").resize((resized_width, resized_height), resample=Image.BILINEAR) + canvas = Image.new("RGB", (size, size), 0) + canvas.paste(resized, ((size - resized_width) // 2, (size - resized_height) // 2)) + return np.asarray(canvas, dtype=np.uint8).copy() + + +class PaligemmaTokenizer: + """PaliGemma SentencePiece tokenizer backed by a local model file.""" + + def __init__(self, model_path: str | Path, max_len: int = 200): + self.max_len = max_len + model_path = Path(model_path) + if not model_path.is_file(): + raise FileNotFoundError(f"PaliGemma tokenizer not found: {model_path}") + self.processor = sentencepiece.SentencePieceProcessor(model_proto=model_path.read_bytes()) + + def tokenize(self, prompt: str) -> tuple[np.ndarray, np.ndarray]: + cleaned = prompt.strip().replace("_", " ").replace("\n", " ") + tokens = self.processor.encode(cleaned, add_bos=True) + self.processor.encode("\n") + if len(tokens) > self.max_len: + LOGGER.warning("Prompt uses %d tokens; truncating to %d", len(tokens), self.max_len) + tokens = tokens[: self.max_len] + mask = [True] * len(tokens) + padding = self.max_len - len(tokens) + tokens += [0] * padding + mask += [False] * padding + return np.asarray(tokens, dtype=np.int64), np.asarray(mask, dtype=np.bool_) + + +class OpenPIPreInfer: + """Convert LIBERO observations to padded model tensors.""" + + def __init__( + self, + norm_stats_path: str | Path, + tokenizer_path: str | Path, + device: torch.device | str, + action_dim: int = 32, + max_token_len: int = 200, + ): + self.norm_stats = load_norm_stats(norm_stats_path) + self.tokenizer = PaligemmaTokenizer(tokenizer_path, max_len=max_token_len) + self.device = torch.device(device) + self.action_dim = action_dim + + def infer(self, images: dict[str, np.ndarray], state: np.ndarray, task_description: str) -> Observation: + base = _resize_with_pad(_require_rgb(images["agentview"])) + left_wrist = _resize_with_pad(_require_rgb(images["wrist"])) + right_wrist = np.zeros_like(base) + + # Quantile normalization runs at simulator input precision upstream. + raw_state = np.asarray(state).reshape(-1) + if raw_state.shape != (LIBERO_STATE_DIM,): + raise ValueError(f"pi05_libero expects an 8-D state, got {raw_state.shape}") + normalized_state = normalize_quantile(raw_state, self.norm_stats["state"]) + padded_state = np.pad(normalized_state, (0, self.action_dim - LIBERO_STATE_DIM)) + + tokens, token_mask = self.tokenizer.tokenize(task_description) + + image_arrays = { + "base_0_rgb": base, + "left_wrist_0_rgb": left_wrist, + "right_wrist_0_rgb": right_wrist, + } + images = {key: torch.from_numpy(value).unsqueeze(0).to(self.device, dtype=torch.float32).permute(0, 3, 1, 2) / 255.0 * 2.0 - 1.0 for key, value in image_arrays.items()} + return Observation( + images=images, + image_masks={ + "base_0_rgb": torch.ones(1, dtype=torch.bool, device=self.device), + "left_wrist_0_rgb": torch.ones(1, dtype=torch.bool, device=self.device), + "right_wrist_0_rgb": torch.zeros(1, dtype=torch.bool, device=self.device), + }, + state=torch.from_numpy(padded_state).unsqueeze(0).to(self.device), + tokenized_prompt=torch.from_numpy(tokens).unsqueeze(0).to(self.device), + tokenized_prompt_mask=torch.from_numpy(token_mask).unsqueeze(0).to(self.device), + ) diff --git a/lightx2v/models/networks/openpi/infer/transformer_infer.py b/lightx2v/models/networks/openpi/infer/transformer_infer.py new file mode 100644 index 000000000..f0b909723 --- /dev/null +++ b/lightx2v/models/networks/openpi/infer/transformer_infer.py @@ -0,0 +1,24 @@ +"""Flow-matching action sampler for OpenPI.""" + +from __future__ import annotations + +import torch + +from ..observation import Observation + + +class OpenPITransformerInfer: + def __init__(self, num_steps: int = 10): + if num_steps <= 0: + raise ValueError("num_steps must be positive") + self.num_steps = num_steps + + @torch.no_grad() + def infer( + self, + model, + observation: Observation, + device: torch.device | str, + noise: torch.Tensor | None = None, + ) -> torch.Tensor: + return model.sample_actions(device, observation, noise=noise, num_steps=self.num_steps) diff --git a/lightx2v/models/networks/openpi/model.py b/lightx2v/models/networks/openpi/model.py new file mode 100644 index 000000000..89bc16759 --- /dev/null +++ b/lightx2v/models/networks/openpi/model.py @@ -0,0 +1,130 @@ +"""LightX2V-native wrapper around the official PyTorch pi0.5 architecture.""" + +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from torch import nn + +from .config import Pi0Config +from .infer import OpenPIPostInfer, OpenPIPreInfer, OpenPITransformerInfer +from .weights import load_pi05_libero_weights + + +class OpenPIModel(nn.Module): + """PyTorch pi0.5 model with LightX2V inference stages.""" + + def __init__( + self, + core_model: nn.Module, + pre_infer: OpenPIPreInfer, + transformer_infer: OpenPITransformerInfer, + post_infer: OpenPIPostInfer, + model_config: Pi0Config, + device: torch.device | str, + seed: int = 0, + ): + super().__init__() + self.core_model = core_model + self.pre_infer = pre_infer + self.transformer_infer = transformer_infer + self.post_infer = post_infer + self.model_config = model_config + self.device = torch.device(device) + self.seed = seed + self._generator: torch.Generator + self.reset() + + @classmethod + def from_config(cls, config: Mapping[str, Any]) -> "OpenPIModel": + values = dict(config) + model_config = Pi0Config.from_mapping(values) + + checkpoint_dir = Path(values["model_path"]).expanduser().resolve() + weight_path = checkpoint_dir / "model.safetensors" + norm_stats_path = checkpoint_dir / "assets/physical-intelligence/libero/norm_stats.json" + tokenizer_path = checkpoint_dir / "assets/paligemma_tokenizer.model" + + device = torch.device(values["device"]) + if device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("OpenPI config requests CUDA, but torch.cuda.is_available() is false") + + core_model = load_pi05_libero_weights(weight_path, model_config, device) + return cls( + core_model=core_model, + pre_infer=OpenPIPreInfer( + norm_stats_path=norm_stats_path, + tokenizer_path=tokenizer_path, + device=device, + action_dim=model_config.action_dim, + max_token_len=model_config.max_token_len, + ), + transformer_infer=OpenPITransformerInfer(num_steps=values["num_inference_steps"]), + post_infer=OpenPIPostInfer(norm_stats_path), + model_config=model_config, + device=device, + seed=values.get("seed", 0), + ) + + def _make_generator(self, seed: int) -> torch.Generator: + generator = torch.Generator(device=self.device) + generator.manual_seed(seed) + return generator + + def reset(self) -> None: + self._generator = self._make_generator(self.seed) + + def get_rng_state(self) -> torch.Tensor: + return self._generator.get_state() + + def set_rng_state(self, state: torch.Tensor) -> None: + self._generator.set_state(state) + + def _sample_noise(self, seed: int | None = None) -> torch.Tensor: + if seed is not None: + generator = self._make_generator(seed) + else: + generator = self._generator + return torch.randn( + (1, self.model_config.action_horizon, self.model_config.action_dim), + dtype=torch.float32, + device=self.device, + generator=generator, + ) + + @torch.no_grad() + def predict_normalized_action_chunk( + self, + images: dict[str, np.ndarray], + state: np.ndarray, + task_description: str, + *, + seed: int | None = None, + noise: torch.Tensor | np.ndarray | None = None, + ) -> torch.Tensor: + observation = self.pre_infer.infer(images, state, task_description) + if noise is None: + noise_tensor = self._sample_noise(seed) + else: + noise_tensor = torch.as_tensor(noise, dtype=torch.float32, device=self.device) + if noise_tensor.ndim == 2: + noise_tensor = noise_tensor.unsqueeze(0) + expected = (1, self.model_config.action_horizon, self.model_config.action_dim) + if tuple(noise_tensor.shape) != expected: + raise ValueError(f"Noise must have shape {expected}, got {tuple(noise_tensor.shape)}") + return self.transformer_infer.infer(self.core_model, observation, self.device, noise=noise_tensor) + + @torch.no_grad() + def predict_action_chunk( + self, + images: dict[str, np.ndarray], + state: np.ndarray, + task_description: str, + seed: int | None = None, + ) -> np.ndarray: + normalized = self.predict_normalized_action_chunk(images, state, task_description, seed=seed) + return self.post_infer.infer(normalized) diff --git a/lightx2v/models/networks/openpi/observation.py b/lightx2v/models/networks/openpi/observation.py new file mode 100644 index 000000000..5e415513f --- /dev/null +++ b/lightx2v/models/networks/openpi/observation.py @@ -0,0 +1,19 @@ +"""Torch-only observation container matching OpenPI's public tensor layout.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass +class Observation: + # Images are float32 in [-1, 1] with BCHW layout. + images: dict[str, torch.Tensor] + image_masks: dict[str, torch.Tensor] + state: torch.Tensor + tokenized_prompt: torch.Tensor + tokenized_prompt_mask: torch.Tensor + token_ar_mask: torch.Tensor | None = None + token_loss_mask: torch.Tensor | None = None diff --git a/lightx2v/models/networks/openpi/pi0.py b/lightx2v/models/networks/openpi/pi0.py new file mode 100644 index 000000000..8513da856 --- /dev/null +++ b/lightx2v/models/networks/openpi/pi0.py @@ -0,0 +1,423 @@ +# Adapted from Physical Intelligence OpenPI (Apache-2.0), commit 15a9616. +# Localized for the LightX2V OpenPI backend; parameter names intentionally unchanged. + +import logging +import math + +import torch +import torch.nn.functional as F # noqa: N812 +from torch import Tensor, nn + +from . import config as _gemma +from . import preprocessing as _preprocessing +from .gemma import PaliGemmaWithExpertModel + + +def get_safe_dtype(target_dtype, device_type): + if device_type == "cpu": + if target_dtype == torch.bfloat16: + return torch.float32 + if target_dtype == torch.float64: + return torch.float64 + return target_dtype + + +def create_sinusoidal_pos_embedding(time: torch.tensor, dimension: int, min_period: float, max_period: float, device="cpu") -> Tensor: + """Computes sine-cosine positional embedding vectors for scalar positions.""" + if dimension % 2 != 0: + raise ValueError(f"dimension ({dimension}) must be divisible by 2") + + if time.ndim != 1: + raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.") + + dtype = get_safe_dtype(torch.float64, device.type) + fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device) + period = min_period * (max_period / min_period) ** fraction + + scaling_factor = 1.0 / period * 2 * math.pi + sin_input = scaling_factor[None, :] * time[:, None] + return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1) + + +def sample_beta(alpha, beta, bsize, device): + alpha_t = torch.as_tensor(alpha, dtype=torch.float32, device=device) + beta_t = torch.as_tensor(beta, dtype=torch.float32, device=device) + dist = torch.distributions.Beta(alpha_t, beta_t) + return dist.sample((bsize,)) + + +def make_att_2d_masks(pad_masks, att_masks): + """Copied from big_vision. + + Tokens can attend to valid inputs tokens which have a cumulative mask_ar + smaller or equal to theirs. This way `mask_ar` int[B, N] can be used to + setup several types of attention, for example: + + [[1 1 1 1 1 1]]: pure causal attention. + + [[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between + themselves and the last 3 tokens have a causal attention. The first + entry could also be a 1 without changing behaviour. + + [[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a + block can attend all previous blocks and all tokens on the same block. + + Args: + input_mask: bool[B, N] true if its part of the input, false if padding. + mask_ar: int32[B, N] mask that's 1 where previous tokens cannot depend on + it and 0 where it shares the same attention mask as the previous token. + """ + if att_masks.ndim != 2: + raise ValueError(att_masks.ndim) + if pad_masks.ndim != 2: + raise ValueError(pad_masks.ndim) + + cumsum = torch.cumsum(att_masks, dim=1) + att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None] + pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None] + return att_2d_masks & pad_2d_masks + + +class PI0Pytorch(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.pi05 = config.pi05 + + paligemma_config = _gemma.get_config(config.paligemma_variant) + action_expert_config = _gemma.get_config(config.action_expert_variant) + + self.paligemma_with_expert = PaliGemmaWithExpertModel( + paligemma_config, + action_expert_config, + use_adarms=[False, True] if self.pi05 else [False, False], + precision=config.dtype, + ) + + self.action_in_proj = nn.Linear(config.action_dim, action_expert_config.width) + self.action_out_proj = nn.Linear(action_expert_config.width, config.action_dim) + + if self.pi05: + self.time_mlp_in = nn.Linear(action_expert_config.width, action_expert_config.width) + self.time_mlp_out = nn.Linear(action_expert_config.width, action_expert_config.width) + else: + self.state_proj = nn.Linear(config.action_dim, action_expert_config.width) + self.action_time_mlp_in = nn.Linear(2 * action_expert_config.width, action_expert_config.width) + self.action_time_mlp_out = nn.Linear(action_expert_config.width, action_expert_config.width) + + torch.set_float32_matmul_precision("high") + if config.pytorch_compile_mode is not None: + self.sample_actions = torch.compile(self.sample_actions, mode=config.pytorch_compile_mode) + + self.gradient_checkpointing_enabled = False + + msg = "OpenPI's patched transformers==4.53.2 runtime is not active. Run scripts/openpi/2_setup_pytorch_runtime.sh and prepend OPENPI_TRANSFORMERS_RUNTIME_PATH to PYTHONPATH." + try: + from transformers.models.siglip import check + + if not check.check_whether_transformers_replace_is_installed_correctly(): + raise ValueError(msg) + except ImportError: + raise ValueError(msg) from None + + def gradient_checkpointing_enable(self): + self.gradient_checkpointing_enabled = True + self.paligemma_with_expert.paligemma.language_model.gradient_checkpointing = True + self.paligemma_with_expert.paligemma.vision_tower.gradient_checkpointing = True + self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = True + + logging.info("Enabled gradient checkpointing for PI0Pytorch model") + + def gradient_checkpointing_disable(self): + self.gradient_checkpointing_enabled = False + self.paligemma_with_expert.paligemma.language_model.gradient_checkpointing = False + self.paligemma_with_expert.paligemma.vision_tower.gradient_checkpointing = False + self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = False + + logging.info("Disabled gradient checkpointing for PI0Pytorch model") + + def is_gradient_checkpointing_enabled(self): + return self.gradient_checkpointing_enabled + + def _apply_checkpoint(self, func, *args, **kwargs): + if self.gradient_checkpointing_enabled and self.training: + return torch.utils.checkpoint.checkpoint(func, *args, use_reentrant=False, preserve_rng_state=False, **kwargs) + return func(*args, **kwargs) + + def _prepare_attention_masks_4d(self, att_2d_masks): + att_2d_masks_4d = att_2d_masks[:, None, :, :] + return torch.where(att_2d_masks_4d, 0.0, -2.3819763e38) + + def _preprocess_observation(self, observation, *, train=True): + observation = _preprocessing.preprocess_observation_pytorch(observation, train=train) + return ( + list(observation.images.values()), + list(observation.image_masks.values()), + observation.tokenized_prompt, + observation.tokenized_prompt_mask, + observation.state, + ) + + def sample_noise(self, shape, device): + return torch.normal( + mean=0.0, + std=1.0, + size=shape, + dtype=torch.float32, + device=device, + ) + + def sample_time(self, bsize, device): + time_beta = sample_beta(1.5, 1.0, bsize, device) + time = time_beta * 0.999 + 0.001 + return time.to(dtype=torch.float32, device=device) + + def embed_prefix(self, images, img_masks, lang_tokens, lang_masks) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Embed images with SigLIP and language tokens with embedding layer to prepare + for PaliGemma transformer processing. + """ + embs = [] + pad_masks = [] + att_masks = [] + + for img, img_mask in zip(images, img_masks, strict=True): + + def image_embed_func(img): + return self.paligemma_with_expert.embed_image(img) + + img_emb = self._apply_checkpoint(image_embed_func, img) + + bsize, num_img_embs = img_emb.shape[:2] + + embs.append(img_emb) + pad_masks.append(img_mask[:, None].expand(bsize, num_img_embs)) + + att_masks += [0] * num_img_embs + + def lang_embed_func(lang_tokens): + lang_emb = self.paligemma_with_expert.embed_language_tokens(lang_tokens) + lang_emb_dim = lang_emb.shape[-1] + return lang_emb * math.sqrt(lang_emb_dim) + + lang_emb = self._apply_checkpoint(lang_embed_func, lang_tokens) + + embs.append(lang_emb) + pad_masks.append(lang_masks) + + num_lang_embs = lang_emb.shape[1] + att_masks += [0] * num_lang_embs + + embs = torch.cat(embs, dim=1) + pad_masks = torch.cat(pad_masks, dim=1) + att_masks = torch.tensor(att_masks, dtype=torch.bool, device=pad_masks.device) + + bsize = pad_masks.shape[0] + att_masks = att_masks[None, :].expand(bsize, len(att_masks)) + + return embs, pad_masks, att_masks + + def embed_suffix(self, state, noisy_actions, timestep): + """Embed state, noisy_actions, timestep to prepare for Expert Gemma processing.""" + embs = [] + pad_masks = [] + att_masks = [] + + if not self.pi05: + if self.state_proj.weight.dtype == torch.float32: + state = state.to(torch.float32) + + def state_proj_func(state): + return self.state_proj(state) + + state_emb = self._apply_checkpoint(state_proj_func, state) + + embs.append(state_emb[:, None, :]) + bsize = state_emb.shape[0] + device = state_emb.device + + state_mask = torch.ones(bsize, 1, dtype=torch.bool, device=device) + pad_masks.append(state_mask) + + att_masks += [1] + + time_emb = create_sinusoidal_pos_embedding(timestep, self.action_in_proj.out_features, min_period=4e-3, max_period=4.0, device=timestep.device) + time_emb = time_emb.type(dtype=timestep.dtype) + + def action_proj_func(noisy_actions): + return self.action_in_proj(noisy_actions) + + action_emb = self._apply_checkpoint(action_proj_func, noisy_actions) + + if not self.pi05: + time_emb = time_emb[:, None, :].expand_as(action_emb) + action_time_emb = torch.cat([action_emb, time_emb], dim=2) + + def mlp_func(action_time_emb): + x = self.action_time_mlp_in(action_time_emb) + x = F.silu(x) + return self.action_time_mlp_out(x) + + action_time_emb = self._apply_checkpoint(mlp_func, action_time_emb) + adarms_cond = None + else: + + def time_mlp_func(time_emb): + x = self.time_mlp_in(time_emb) + x = F.silu(x) + x = self.time_mlp_out(x) + return F.silu(x) + + time_emb = self._apply_checkpoint(time_mlp_func, time_emb) + action_time_emb = action_emb + adarms_cond = time_emb + + embs.append(action_time_emb) + + bsize, action_time_dim = action_time_emb.shape[:2] + action_time_mask = torch.ones(bsize, action_time_dim, dtype=torch.bool, device=timestep.device) + pad_masks.append(action_time_mask) + + att_masks += [1] + ([0] * (self.config.action_horizon - 1)) + + embs = torch.cat(embs, dim=1) + pad_masks = torch.cat(pad_masks, dim=1) + att_masks = torch.tensor(att_masks, dtype=embs.dtype, device=embs.device) + att_masks = att_masks[None, :].expand(bsize, len(att_masks)) + + return embs, pad_masks, att_masks, adarms_cond + + def forward(self, observation, actions, noise=None, time=None) -> Tensor: + """Do a full training forward pass and compute the loss (batch_size x num_steps x num_motors)""" + images, img_masks, lang_tokens, lang_masks, state = self._preprocess_observation(observation, train=True) + + if noise is None: + noise = self.sample_noise(actions.shape, actions.device) + + if time is None: + time = self.sample_time(actions.shape[0], actions.device) + + time_expanded = time[:, None, None] + x_t = time_expanded * noise + (1 - time_expanded) * actions + u_t = noise - actions + + prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(images, img_masks, lang_tokens, lang_masks) + suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = self.embed_suffix(state, x_t, time) + if self.paligemma_with_expert.paligemma.language_model.layers[0].self_attn.q_proj.weight.dtype == torch.bfloat16: + suffix_embs = suffix_embs.to(dtype=torch.bfloat16) + prefix_embs = prefix_embs.to(dtype=torch.bfloat16) + + pad_masks = torch.cat([prefix_pad_masks, suffix_pad_masks], dim=1) + att_masks = torch.cat([prefix_att_masks, suffix_att_masks], dim=1) + + att_2d_masks = make_att_2d_masks(pad_masks, att_masks) + position_ids = torch.cumsum(pad_masks, dim=1) - 1 + + att_2d_masks_4d = self._prepare_attention_masks_4d(att_2d_masks) + + def forward_func(prefix_embs, suffix_embs, att_2d_masks_4d, position_ids, adarms_cond): + (_, suffix_out), _ = self.paligemma_with_expert.forward( + attention_mask=att_2d_masks_4d, + position_ids=position_ids, + past_key_values=None, + inputs_embeds=[prefix_embs, suffix_embs], + use_cache=False, + adarms_cond=[None, adarms_cond], + ) + return suffix_out + + suffix_out = self._apply_checkpoint(forward_func, prefix_embs, suffix_embs, att_2d_masks_4d, position_ids, adarms_cond) + + suffix_out = suffix_out[:, -self.config.action_horizon :] + suffix_out = suffix_out.to(dtype=torch.float32) + + def action_out_proj_func(suffix_out): + return self.action_out_proj(suffix_out) + + v_t = self._apply_checkpoint(action_out_proj_func, suffix_out) + + return F.mse_loss(u_t, v_t, reduction="none") + + @torch.no_grad() + def sample_actions(self, device, observation, noise=None, num_steps=10) -> Tensor: + """Do a full inference forward and compute the action (batch_size x num_steps x num_motors)""" + bsize = observation.state.shape[0] + if noise is None: + actions_shape = (bsize, self.config.action_horizon, self.config.action_dim) + noise = self.sample_noise(actions_shape, device) + + images, img_masks, lang_tokens, lang_masks, state = self._preprocess_observation(observation, train=False) + + prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(images, img_masks, lang_tokens, lang_masks) + prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks) + prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1 + + prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(prefix_att_2d_masks) + self.paligemma_with_expert.paligemma.language_model.config._attn_implementation = "eager" # noqa: SLF001 + + _, past_key_values = self.paligemma_with_expert.forward( + attention_mask=prefix_att_2d_masks_4d, + position_ids=prefix_position_ids, + past_key_values=None, + inputs_embeds=[prefix_embs, None], + use_cache=True, + ) + + dt = -1.0 / num_steps + dt = torch.tensor(dt, dtype=torch.float32, device=device) + + x_t = noise + time = torch.tensor(1.0, dtype=torch.float32, device=device) + while time >= -dt / 2: + expanded_time = time.expand(bsize) + v_t = self.denoise_step( + state, + prefix_pad_masks, + past_key_values, + x_t, + expanded_time, + ) + + x_t = x_t + dt * v_t + time += dt + return x_t + + def denoise_step( + self, + state, + prefix_pad_masks, + past_key_values, + x_t, + timestep, + ): + """Apply one denoising step of the noise `x_t` at a given timestep.""" + suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = self.embed_suffix(state, x_t, timestep) + + suffix_len = suffix_pad_masks.shape[1] + batch_size = prefix_pad_masks.shape[0] + prefix_len = prefix_pad_masks.shape[1] + + prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, suffix_len, prefix_len) + + suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks) + + full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2) + + prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None] + position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1 + + full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks) + self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager" # noqa: SLF001 + + outputs_embeds, _ = self.paligemma_with_expert.forward( + attention_mask=full_att_2d_masks_4d, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=[None, suffix_embs], + use_cache=False, + adarms_cond=[None, adarms_cond], + ) + + suffix_out = outputs_embeds[1] + suffix_out = suffix_out[:, -self.config.action_horizon :] + suffix_out = suffix_out.to(dtype=torch.float32) + return self.action_out_proj(suffix_out) diff --git a/lightx2v/models/networks/openpi/preprocessing.py b/lightx2v/models/networks/openpi/preprocessing.py new file mode 100644 index 000000000..07771a5fa --- /dev/null +++ b/lightx2v/models/networks/openpi/preprocessing.py @@ -0,0 +1,130 @@ +# Adapted from Physical Intelligence OpenPI (Apache-2.0), commit 15a9616. +# Localized for the LightX2V OpenPI backend. + +import logging +from collections.abc import Sequence + +import torch + +from . import image_tools + +logger = logging.getLogger("openpi") + +IMAGE_KEYS = ( + "base_0_rgb", + "left_wrist_0_rgb", + "right_wrist_0_rgb", +) + +IMAGE_RESOLUTION = (224, 224) + + +def preprocess_observation_pytorch( + observation, + *, + train: bool = False, + image_keys: Sequence[str] = IMAGE_KEYS, + image_resolution: tuple[int, int] = IMAGE_RESOLUTION, +): + if not set(image_keys).issubset(observation.images): + raise ValueError(f"images dict missing keys: expected {image_keys}, got {list(observation.images)}") + + batch_shape = observation.state.shape[:-1] + + out_images = {} + for key in image_keys: + image = observation.images[key] + + is_channels_first = image.shape[1] == 3 + + if is_channels_first: + image = image.permute(0, 2, 3, 1) + + if image.shape[1:3] != image_resolution: + logger.info(f"Resizing image {key} from {image.shape[1:3]} to {image_resolution}") + image = image_tools.resize_with_pad_torch(image, *image_resolution) + + if train: + image = image / 2.0 + 0.5 + + if "wrist" not in key: + height, width = image.shape[1:3] + + crop_height = int(height * 0.95) + crop_width = int(width * 0.95) + + max_h = height - crop_height + max_w = width - crop_width + if max_h > 0 and max_w > 0: + start_h = torch.randint(0, max_h + 1, (1,), device=image.device) + start_w = torch.randint(0, max_w + 1, (1,), device=image.device) + image = image[:, start_h : start_h + crop_height, start_w : start_w + crop_width, :] + + image = torch.nn.functional.interpolate( + image.permute(0, 3, 1, 2), + size=(height, width), + mode="bilinear", + align_corners=False, + ).permute(0, 2, 3, 1) + + angle = torch.rand(1, device=image.device) * 10 - 5 + if torch.abs(angle) > 0.1: + angle_rad = angle * torch.pi / 180.0 + cos_a = torch.cos(angle_rad) + sin_a = torch.sin(angle_rad) + grid_x = torch.linspace(-1, 1, width, device=image.device) + grid_y = torch.linspace(-1, 1, height, device=image.device) + grid_y, grid_x = torch.meshgrid(grid_y, grid_x, indexing="ij") + grid_x = grid_x.unsqueeze(0).expand(image.shape[0], -1, -1) + grid_y = grid_y.unsqueeze(0).expand(image.shape[0], -1, -1) + grid_x_rot = grid_x * cos_a - grid_y * sin_a + grid_y_rot = grid_x * sin_a + grid_y * cos_a + grid = torch.stack([grid_x_rot, grid_y_rot], dim=-1) + + image = torch.nn.functional.grid_sample( + image.permute(0, 3, 1, 2), + grid, + mode="bilinear", + padding_mode="zeros", + align_corners=False, + ).permute(0, 2, 3, 1) + + brightness_factor = 0.7 + torch.rand(1, device=image.device) * 0.6 + image = image * brightness_factor + + contrast_factor = 0.6 + torch.rand(1, device=image.device) * 0.8 + mean = image.mean(dim=[1, 2, 3], keepdim=True) + image = (image - mean) * contrast_factor + mean + + saturation_factor = 0.5 + torch.rand(1, device=image.device) * 1.0 + gray = image.mean(dim=-1, keepdim=True) + image = gray + (image - gray) * saturation_factor + image = torch.clamp(image, 0, 1) + image = image * 2.0 - 1.0 + + if is_channels_first: + image = image.permute(0, 3, 1, 2) + + out_images[key] = image + + out_masks = {} + for key in out_images: + if key not in observation.image_masks: + out_masks[key] = torch.ones(batch_shape, dtype=torch.bool, device=observation.state.device) + else: + out_masks[key] = observation.image_masks[key] + + class SimpleProcessedObservation: + def __init__(self, **kwargs): + for key, value in kwargs.items(): + setattr(self, key, value) + + return SimpleProcessedObservation( + images=out_images, + image_masks=out_masks, + state=observation.state, + tokenized_prompt=observation.tokenized_prompt, + tokenized_prompt_mask=observation.tokenized_prompt_mask, + token_ar_mask=observation.token_ar_mask, + token_loss_mask=observation.token_loss_mask, + ) diff --git a/lightx2v/models/networks/openpi/transformers_replace/models/gemma/configuration_gemma.py b/lightx2v/models/networks/openpi/transformers_replace/models/gemma/configuration_gemma.py new file mode 100644 index 000000000..3d9353bd2 --- /dev/null +++ b/lightx2v/models/networks/openpi/transformers_replace/models/gemma/configuration_gemma.py @@ -0,0 +1,173 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/gemma/modular_gemma.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_gemma.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# coding=utf-8 +# Copyright 2024 Google Inc. HuggingFace Inc. team. All rights reserved. +# +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Optional +from ...configuration_utils import PretrainedConfig + + +class GemmaConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`GemmaModel`]. It is used to instantiate an Gemma + model according to the specified arguments, defining the model architecture. Instantiating a configuration with the + defaults will yield a similar configuration to that of the Gemma-7B. + e.g. [google/gemma-7b](https://huggingface.co/google/gemma-7b) + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + Args: + vocab_size (`int`, *optional*, defaults to 256000): + Vocabulary size of the Gemma model. Defines the number of different tokens that can be represented by the + `inputs_ids` passed when calling [`GemmaModel`] + hidden_size (`int`, *optional*, defaults to 3072): + Dimension of the hidden representations. + intermediate_size (`int`, *optional*, defaults to 24576): + Dimension of the MLP representations. + num_hidden_layers (`int`, *optional*, defaults to 28): + Number of hidden layers in the Transformer decoder. + num_attention_heads (`int`, *optional*, defaults to 16): + Number of attention heads for each attention layer in the Transformer decoder. + num_key_value_heads (`int`, *optional*, defaults to 16): + This is the number of key_value heads that should be used to implement Grouped Query Attention. If + `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if + `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When + converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed + by meanpooling all the original heads within that group. For more details, check out [this + paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to + `num_attention_heads`. + head_dim (`int`, *optional*, defaults to 256): + The attention head dimension. + hidden_act (`str` or `function`, *optional*, defaults to `"gelu_pytorch_tanh"`): + The legacy activation function. It is overwritten by the `hidden_activation`. + hidden_activation (`str` or `function`, *optional*): + The non-linear activation function (function or string) in the decoder. Will default to `"gelu_pytorch_tanh"` + if not specified. `"gelu_pytorch_tanh"` uses an approximation of the `"gelu"` activation function. + max_position_embeddings (`int`, *optional*, defaults to 8192): + The maximum sequence length that this model might ever be used with. + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + rms_norm_eps (`float`, *optional*, defaults to 1e-06): + The epsilon used by the rms normalization layers. + use_cache (`bool`, *optional*, defaults to `True`): + Whether or not the model should return the last key/values attentions (not used by all models). Only + relevant if `config.is_decoder=True`. + pad_token_id (`int`, *optional*, defaults to 0): + Padding token id. + eos_token_id (`int`, *optional*, defaults to 1): + End of stream token id. + bos_token_id (`int`, *optional*, defaults to 2): + Beginning of stream token id. + tie_word_embeddings (`bool`, *optional*, defaults to `True`): + Whether to tie weight embeddings + rope_theta (`float`, *optional*, defaults to 10000.0): + The base period of the RoPE embeddings. + attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`): + Whether to use a bias in the query, key, value and output projection layers during self-attention. + attention_dropout (`float`, *optional*, defaults to 0.0): + The dropout ratio for the attention probabilities. + use_adarms (`bool`, *optional*, defaults to `False`): + Whether to use ADARMS. + adarms_cond_dim (`int`, *optional*, defaults to `None`): + The dimension of the ADARMS condition. + ```python + >>> from transformers import GemmaModel, GemmaConfig + >>> # Initializing a Gemma gemma-7b style configuration + >>> configuration = GemmaConfig() + >>> # Initializing a model from the gemma-7b style configuration + >>> model = GemmaModel(configuration) + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "gemma" + keys_to_ignore_at_inference = ["past_key_values"] + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + def __init__( + self, + vocab_size=256000, + hidden_size=3072, + intermediate_size=24576, + num_hidden_layers=28, + num_attention_heads=16, + num_key_value_heads=16, + head_dim=256, + hidden_act="gelu_pytorch_tanh", + hidden_activation=None, + max_position_embeddings=8192, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=0, + eos_token_id=1, + bos_token_id=2, + tie_word_embeddings=True, + rope_theta=10000.0, + attention_bias=False, + attention_dropout=0.0, + use_adarms: bool = False, + adarms_cond_dim: Optional[int] = None, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.head_dim = head_dim + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.hidden_activation = hidden_activation + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.use_adarms = use_adarms + self.adarms_cond_dim = adarms_cond_dim + + # Set default for adarms_cond_dim if use_adarms is True + if self.use_adarms and self.adarms_cond_dim is None: + self.adarms_cond_dim = self.hidden_size + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + +__all__ = ["GemmaConfig"] diff --git a/lightx2v/models/networks/openpi/transformers_replace/models/gemma/modeling_gemma.py b/lightx2v/models/networks/openpi/transformers_replace/models/gemma/modeling_gemma.py new file mode 100644 index 000000000..dec6439a6 --- /dev/null +++ b/lightx2v/models/networks/openpi/transformers_replace/models/gemma/modeling_gemma.py @@ -0,0 +1,862 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/gemma/modular_gemma.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_gemma.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# coding=utf-8 +# Copyright 2024 Google Inc. HuggingFace Inc. team. All rights reserved. +# +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Callable, Optional, Union + +import torch +from torch import nn + +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...masking_utils import create_causal_mask +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, + SequenceClassifierOutputWithPast, + TokenClassifierOutput, +) +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import LossKwargs, auto_docstring, can_return_tuple, logging +from .configuration_gemma import GemmaConfig + + +logger = logging.get_logger(__name__) + + +class GemmaRMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6, cond_dim: Optional[int] = None): + super().__init__() + self.eps = eps + self.dim = dim + self.cond_dim = cond_dim + + # Dense layer for adaptive normalization (if cond_dim is provided) + if cond_dim is not None: + #self.dense = nn.Linear(cond_dim, dim * 3, bias=True, dtype=torch.bfloat16) + self.dense = nn.Linear(cond_dim, dim * 3, bias=True) + # Initialize with zeros (matches source implementation) + nn.init.zeros_(self.dense.weight) + else: + self.weight = nn.Parameter(torch.zeros(dim, dtype=torch.bfloat16)) + self.dense = None + + def _norm(self, x): + # Compute variance in float32 (like the source implementation) + var = torch.mean(torch.square(x.float()), dim=-1, keepdim=True) + # Compute normalization in float32 + normed_inputs = x * torch.rsqrt(var + self.eps) + return normed_inputs + + def forward(self, x, cond=None): + dtype = x.dtype # original dtype, could be half-precision + normed_inputs = self._norm(x) + + if cond is None or self.dense is None: + # regular RMSNorm + # scale by learned parameter in float32 (matches source implementation) + normed_inputs = normed_inputs * (1.0 + self.weight.float()) + return normed_inputs.to(dtype), None # return in original dtype with None gate + + # adaptive RMSNorm (if cond is provided and dense layer exists) + if cond.shape[-1] != self.cond_dim: + raise ValueError(f"Expected cond dimension {self.cond_dim}, got {cond.shape[-1]}") + + #self.dense.to(dtype=torch.bfloat16).to(dtype=torch.float32) + modulation = self.dense(cond) + # Reshape modulation to broadcast properly: [batch, 1, features] for [batch, seq, features] + if len(x.shape) == 3: # [batch, seq, features] + modulation = modulation.unsqueeze(1) + + scale, shift, gate = torch.chunk(modulation, 3, dim=-1) + + # Apply adaptive normalization: use model weight dtype to ensure compatibility + # model_dtype = self.dense.weight.dtype # Use the model's dtype (bfloat16) + # scale = scale.to(model_dtype) + # shift = shift.to(model_dtype) + # gate = gate.to(model_dtype) + # normed_inputs = normed_inputs.to(model_dtype) # Convert normed_inputs to model dtype + + normed_inputs = normed_inputs * (1 + scale.to(torch.float32)) + shift.to(torch.float32) + + return normed_inputs.to(dtype), gate.to(dtype) + + def extra_repr(self): + repr_str = f"{tuple(self.weight.shape)}, eps={self.eps}" + if self.dense is not None: + repr_str += f", adaptive=True, cond_dim={self.cond_dim}" + return repr_str + + +class GemmaMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +class GemmaRotaryEmbedding(nn.Module): + def __init__(self, config: GemmaConfig, device=None): + super().__init__() + # BC: "rope_type" was originally "type" + if hasattr(config, "rope_scaling") and config.rope_scaling is not None: + self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) + else: + self.rope_type = "default" + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.original_inv_freq = self.inv_freq + + @torch.no_grad() + @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) + def forward(self, x, position_ids): + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + position_ids_expanded = position_ids[:, None, :].float() + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`, *optional*): + Deprecated and unused. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def _gated_residual(x, y, gate): + """ + Applies gated residual connection with optional gate parameter. + + Args: + x: Input tensor (residual) + y: Output tensor to be added + gate: Optional gate tensor to modulate the addition + + Returns: + x + y if gate is None, otherwise x + y * gate + """ + if x is None and y is None: + return None + if x is None or y is None: + return x if x is not None else y + if gate is None: + return x + y + return x + y * gate + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + scaling: float, + dropout: float = 0.0, + **kwargs, +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +class GemmaAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: GemmaConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = True + + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], + past_key_value: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + use_cache: bool = False, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + # Use cache if provided + if past_key_value is not None: + if use_cache: + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + else: + key_states = torch.cat([past_key_value[self.layer_idx][0], key_states], dim=2) + value_states = torch.cat([past_key_value[self.layer_idx][1], value_states], dim=2) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class GemmaDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: GemmaConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = GemmaAttention(config=config, layer_idx=layer_idx) + + self.mlp = GemmaMLP(config) + cond_dim = getattr(config, 'adarms_cond_dim', None) if getattr(config, 'use_adarms', False) else None + self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps, cond_dim=cond_dim) + self.post_attention_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps, cond_dim=cond_dim) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + adarms_cond: Optional[torch.Tensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: + residual = hidden_states + hidden_states, gate = self.input_layernorm(hidden_states, adarms_cond) + + # Self Attention + hidden_states, self_attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = _gated_residual(residual, hidden_states, gate) + + # Fully Connected + residual = hidden_states + hidden_states, gate = self.post_attention_layernorm(hidden_states, adarms_cond) + hidden_states = self.mlp(hidden_states) + hidden_states = _gated_residual(residual, hidden_states, gate) + + outputs = (hidden_states,) + if output_attentions: + outputs += (self_attn_weights,) + + return outputs + + +@auto_docstring +class GemmaPreTrainedModel(PreTrainedModel): + config_class = GemmaConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["GemmaDecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn_3 = True + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_flex_attn = True + _supports_cache_class = True + _supports_quantized_cache = True + _supports_static_cache = True + _supports_attention_backend = True + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + elif isinstance(module, GemmaRMSNorm): + if hasattr(module, 'weight'): + module.weight.data.fill_(1.0) + + +@auto_docstring +class GemmaModel(GemmaPreTrainedModel): + def __init__(self, config: GemmaConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [GemmaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + + cond_dim = getattr(config, 'adarms_cond_dim', None) if getattr(config, 'use_adarms', False) else None + self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps, cond_dim=cond_dim) + self.rotary_emb = GemmaRotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + adarms_cond: Optional[torch.Tensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> BaseModelOutputWithPast: + """ + adarms_cond (`torch.Tensor` of shape `(batch_size, cond_dim)`, *optional*): + Condition for ADARMS. + """ + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if self.gradient_checkpointing and self.training and use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." + ) + use_cache = False + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache() + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device + ) + + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + # embed positions + hidden_states = inputs_embeds + # Convert to bfloat16 if the first layer uses bfloat16 + if len(self.layers) > 0 and self.layers[0].self_attn.q_proj.weight.dtype == torch.bfloat16: + hidden_states = hidden_states.to(torch.bfloat16) + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # normalized + # Gemma downcasts the below to float16, causing sqrt(3072)=55.4256 to become 55.5 + # See https://github.com/huggingface/transformers/pull/29402 + normalizer = torch.tensor(self.config.hidden_size**0.5, dtype=hidden_states.dtype) + #hidden_states = hidden_states * normalizer + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + adarms_cond=adarms_cond, + **kwargs, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states, _ = self.norm(hidden_states, adarms_cond) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values if use_cache else None, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + +class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ... + + +@auto_docstring +class GemmaForCausalLM(GemmaPreTrainedModel, GenerationMixin): + _tied_weights_keys = ["lm_head.weight"] + _tp_plan = {"lm_head": "colwise_rep"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + + def __init__(self, config): + super().__init__(config) + self.model = GemmaModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + adarms_cond: Optional[torch.Tensor] = None, + **kwargs: Unpack[KwargsForCausalLM], + ) -> CausalLMOutputWithPast: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + adarms_cond (`torch.Tensor` of shape `(batch_size, cond_dim)`, *optional*): + Condition for ADARMS. + + Example: + + ```python + >>> from transformers import AutoTokenizer, GemmaForCausalLM + + >>> model = GemmaForCausalLM.from_pretrained("google/gemma-7b") + >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b") + + >>> prompt = "What is your favorite condiment?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "What is your favorite condiment?" + ```""" + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs: BaseModelOutputWithPast = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + cache_position=cache_position, + adarms_cond=adarms_cond, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + The Gemma Model transformer with a sequence classification head on top (linear layer). + + [`GemmaForSequenceClassification`] uses the last token in order to do the classification, as other causal models + (e.g. GPT-2) do. + + Since it does classification on the last token, it requires to know the position of the last token. If a + `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If + no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the + padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in + each row of the batch). + """ +) +class GemmaForSequenceClassification(GemmaPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.model = GemmaModel(config) + self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + adarms_cond: Optional[torch.Tensor] = None, + ) -> SequenceClassifierOutputWithPast: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + + adarms_cond (`torch.Tensor` of shape `(batch_size, cond_dim)`, *optional*): + Condition for ADARMS. + """ + + transformer_outputs: BaseModelOutputWithPast = self.model( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + adarms_cond=adarms_cond, + ) + hidden_states = transformer_outputs.last_hidden_state + logits = self.score(hidden_states) + + if input_ids is not None: + batch_size = input_ids.shape[0] + else: + batch_size = inputs_embeds.shape[0] + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") + if self.config.pad_token_id is None: + last_non_pad_token = -1 + elif input_ids is not None: + # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id + non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32) + token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32) + last_non_pad_token = (token_indices * non_pad_mask).argmax(-1) + else: + last_non_pad_token = -1 + logger.warning_once( + f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be " + "unexpected if using padding tokens in conjunction with `inputs_embeds.`" + ) + + pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token] + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config) + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@auto_docstring +class GemmaForTokenClassification(GemmaPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.model = GemmaModel(config) + if getattr(config, "classifier_dropout", None) is not None: + classifier_dropout = config.classifier_dropout + elif getattr(config, "hidden_dropout", None) is not None: + classifier_dropout = config.hidden_dropout + else: + classifier_dropout = 0.1 + self.dropout = nn.Dropout(classifier_dropout) + self.score = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + adarms_cond: Optional[torch.Tensor] = None, + ) -> TokenClassifierOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + + adarms_cond (`torch.Tensor` of shape `(batch_size, cond_dim)`, *optional*): + Condition for ADARMS. + """ + + outputs: BaseModelOutputWithPast = self.model( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + adarms_cond=adarms_cond, + ) + sequence_output = outputs.last_hidden_state + sequence_output = self.dropout(sequence_output) + logits = self.score(sequence_output) + + loss = None + if labels is not None: + loss = self.loss_function(logits, labels, self.config) + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = [ + "GemmaModel", + "GemmaForCausalLM", + "GemmaForSequenceClassification", + "GemmaForTokenClassification", + "GemmaPreTrainedModel", +] diff --git a/lightx2v/models/networks/openpi/transformers_replace/models/paligemma/modeling_paligemma.py b/lightx2v/models/networks/openpi/transformers_replace/models/paligemma/modeling_paligemma.py new file mode 100644 index 000000000..fbf0e9489 --- /dev/null +++ b/lightx2v/models/networks/openpi/transformers_replace/models/paligemma/modeling_paligemma.py @@ -0,0 +1,622 @@ +# coding=utf-8 +# Copyright 2024 the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch PaliGemmamodel.""" + +from dataclasses import dataclass +from typing import Optional, Union + +import torch +import torch.utils.checkpoint +from torch import nn + +from ...cache_utils import Cache, HybridCache, StaticCache +from ...generation import GenerationMixin +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_outputs import BaseModelOutputWithPast +from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack +from ...utils import LossKwargs, ModelOutput, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging +from ..auto import AutoModel +from .configuration_paligemma import PaliGemmaConfig + + +logger = logging.get_logger(__name__) + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for Paligemma outputs, with hidden states and attentions. + """ +) +class PaligemmaModelOutputWithPast(BaseModelOutputWithPast): + r""" + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + image_hidden_states (`torch.FloatTensor`, *optional*): + A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`. + image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state. + """ + + image_hidden_states: Optional[torch.FloatTensor] = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for PaliGemma causal language model (or autoregressive) outputs. + """ +) +class PaliGemmaCausalLMOutputWithPast(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.text_config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + image_hidden_states (`torch.FloatTensor`, *optional*): + A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`. + image_hidden_states of the model produced by the vision encoder after projecting last hidden state. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + past_key_values: Optional[Union[list[torch.FloatTensor], Cache]] = None + hidden_states: Optional[tuple[torch.FloatTensor]] = None + attentions: Optional[tuple[torch.FloatTensor]] = None + image_hidden_states: Optional[torch.FloatTensor] = None + + +class PaliGemmaMultiModalProjector(nn.Module): + def __init__(self, config: PaliGemmaConfig): + super().__init__() + self.linear = nn.Linear(config.vision_config.hidden_size, config.vision_config.projection_dim, bias=True) + + def forward(self, image_features): + hidden_states = self.linear(image_features) + + return hidden_states + + +@auto_docstring +class PaliGemmaPreTrainedModel(PreTrainedModel): + config_class = PaliGemmaConfig + base_model_prefix = "" + supports_gradient_checkpointing = True + _no_split_modules = ["PaliGemmaMultiModalProjector"] + _skip_keys_device_placement = "past_key_values" + _supports_cache_class = True + _supports_quantized_cache = True + _supports_static_cache = True + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_flex_attn = True + _supports_attention_backend = True + + def _init_weights(self, module): + # important: this ported version of PaliGemmaisn't meant for training from scratch - only + # inference and fine-tuning + std = getattr(self.config, "initializer_range", self.config.get_text_config().initializer_range) + + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + + +@auto_docstring( + custom_intro=""" + The Base Paligemma model which consists of a vision backbone and a language model withou language modeling head., + """ +) +class PaliGemmaModel(PaliGemmaPreTrainedModel): + _checkpoint_conversion_mapping = {"language_model.model": "language_model"} + # we are filtering the logits/labels so we shouldn't divide the loss based on num_items_in_batch + accepts_loss_kwargs = False + + def __init__(self, config: PaliGemmaConfig): + super().__init__(config) + self.vision_tower = AutoModel.from_config(config=config.vision_config) + self.multi_modal_projector = PaliGemmaMultiModalProjector(config) + self.vocab_size = config.text_config.vocab_size + + language_model = AutoModel.from_config(config=config.text_config) + self.language_model = language_model + + self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1 + self.post_init() + + # Copied from transformers.models.llava.modeling_llava.LlavaModel.get_input_embeddings with Llava->PaliGemma + def get_input_embeddings(self): + return self.language_model.get_input_embeddings() + + # Copied from transformers.models.llava.modeling_llava.LlavaModel.set_input_embeddings with Llava->PaliGemma + def set_input_embeddings(self, value): + self.language_model.set_input_embeddings(value) + + def set_decoder(self, decoder): + self.language_model = decoder + + def get_decoder(self): + return self.language_model + + def _update_causal_mask( + self, + attention_mask, + token_type_ids=None, + past_key_values=None, + cache_position=None, + input_tensor=None, + is_training: Optional[bool] = None, + ): + if self.config.text_config._attn_implementation == "flash_attention_2": + if attention_mask is not None and 0.0 in attention_mask: + return attention_mask + return None + is_training = is_training if is_training is not None else self.training + using_static_cache = isinstance(past_key_values, StaticCache) + min_dtype = torch.finfo(self.dtype).min + if input_tensor is None: + input_tensor = attention_mask + + inputs_lead_dim, sequence_length = input_tensor.shape[:2] + if using_static_cache: + target_length = past_key_values.get_max_cache_shape() + elif isinstance(past_key_values, HybridCache): + target_length = past_key_values.get_max_cache_shape() + else: + target_length = ( + attention_mask.shape[-1] + if isinstance(attention_mask, torch.Tensor) + else cache_position[0] + sequence_length + 1 + ) + + if attention_mask is not None and attention_mask.dim() == 4: + # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. + return attention_mask + + causal_mask = torch.full( + (sequence_length, target_length), fill_value=min_dtype, dtype=self.dtype, device=cache_position.device + ) + # Causal diagonal mask only if training, otherwise attend to the whole prefix. Training-specific attn for prefix is handled below + if sequence_length != 1: + if is_training: + causal_mask = torch.triu(causal_mask, diagonal=1) + else: + causal_mask[:, :sequence_length] = 0.0 + + causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) + causal_mask = causal_mask[None, None, :, :].expand(inputs_lead_dim, 1, -1, -1) + if attention_mask is not None: + causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit + mask_length = attention_mask.shape[-1] + + # First unmask prefix tokens during training + if is_training: + if token_type_ids is None: + raise ValueError("Token type ids must be provided during training") + causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( + token_type_ids[:, None, None, :].to(causal_mask.device) == 0, 0 + ) + + # Then apply padding mask (will mask pad tokens) + padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(causal_mask.device) + padding_mask = padding_mask == 0 + causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( + padding_mask, min_dtype + ) + + return causal_mask + + def get_image_features(self, pixel_values: torch.FloatTensor): + """ + Obtains image last hidden states from the vision tower and apply multimodal projection. + + Args: + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) + The tensors corresponding to the input images. + Returns: + image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). + """ + image_outputs = self.vision_tower(pixel_values) + selected_image_feature = image_outputs.last_hidden_state + image_features = self.multi_modal_projector(selected_image_feature) + return image_features + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor = None, + pixel_values: torch.FloatTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[list[torch.FloatTensor], Cache]] = None, + token_type_ids: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Union[tuple, PaligemmaModelOutputWithPast]: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.text_config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.text_config.vocab_size]`. + + Example: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, PaliGemmaForConditionalGeneration + + >>> model = PaliGemmaForConditionalGeneration.from_pretrained("google/paligemma2-3b-mix-224") + >>> processor = AutoProcessor.from_pretrained("google/paligemma2-3b-mix-224") + + >>> prompt = "Where is the cat standing?" + >>> url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(images=image, text=prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(**inputs,) + >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Where is the cat standing?\nsnow" + ```""" + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + is_training = token_type_ids is not None and labels is not None + + # Replace image id woth PAD if the image token if OOV, to avoid index-errors + if input_ids is not None and self.config.image_token_id >= self.vocab_size: + special_image_mask = input_ids == self.config.image_token_id + llm_input_ids = input_ids.clone() + llm_input_ids[special_image_mask] = 0 + else: + llm_input_ids = input_ids + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(llm_input_ids) + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device + ) + + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + 1 # Paligemma positions are 1-indexed + + # Merge text and images + if pixel_values is not None: + image_features = self.get_image_features(pixel_values) + + if input_ids is None: + special_image_mask = inputs_embeds == self.get_input_embeddings()( + torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) + ) + else: + special_image_mask = (input_ids == self.config.image_token_id).unsqueeze(-1) + special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device) + + if not is_torchdynamo_compiling() and inputs_embeds[special_image_mask].numel() != image_features.numel(): + image_tokens_in_text = (special_image_mask).sum(dim=1).sum(dim=0)[0] + raise ValueError( + f"Number of images does not match number of special image tokens in the input text. " + f"Got {image_tokens_in_text} image tokens in the text but {image_features.shape[0] * image_features.shape[1]} " + "tokens from image embeddings." + ) + image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) + + causal_mask = self._update_causal_mask( + attention_mask, token_type_ids, past_key_values, cache_position, inputs_embeds, is_training + ) + outputs = self.language_model( + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=True, + cache_position=cache_position, + **kwargs, + ) + + return PaligemmaModelOutputWithPast( + last_hidden_state=outputs.last_hidden_state, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=image_features if pixel_values is not None else None, + ) + + +class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ... + + +@auto_docstring( + custom_intro=""" + The Base Paligemma model which consists of a vision backbone and a language model without language modeling head., + """ +) +class PaliGemmaForConditionalGeneration(PaliGemmaPreTrainedModel, GenerationMixin): + _checkpoint_conversion_mapping = { + "^language_model.model": "model.language_model", + "^vision_tower": "model.vision_tower", + "^multi_modal_projector": "model.multi_modal_projector", + "^language_model.lm_head": "lm_head", + } + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config: PaliGemmaConfig): + super().__init__(config) + self.model = PaliGemmaModel(config) + self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) + self.post_init() + + def get_input_embeddings(self): + return self.model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.model.set_input_embeddings(value) + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model.set_decoder(decoder) + + def get_decoder(self): + return self.model.get_decoder() + + def get_image_features(self, pixel_values): + return self.model.get_image_features(pixel_values) + + # Make modules available throught conditional class for BC + @property + def language_model(self): + return self.model.language_model + + @property + def vision_tower(self): + return self.model.vision_tower + + @property + def multi_modal_projector(self): + return self.model.multi_modal_projector + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor = None, + pixel_values: torch.FloatTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[list[torch.FloatTensor], Cache]] = None, + token_type_ids: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + **kwargs: Unpack[KwargsForCausalLM], + ) -> Union[tuple, PaliGemmaCausalLMOutputWithPast]: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.text_config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.text_config.vocab_size]`. + + Example: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, PaliGemmaForConditionalGeneration + + >>> model = PaliGemmaForConditionalGeneration.from_pretrained("google/paligemma2-3b-mix-224") + >>> processor = AutoProcessor.from_pretrained("google/paligemma2-3b-mix-224") + + >>> prompt = "Where is the cat standing?" + >>> url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(images=image, text=prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(**inputs,) + >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Where is the cat standing?\nsnow" + ```""" + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + token_type_ids=token_type_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + labels=labels, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=True, + cache_position=cache_position, + **kwargs, + ) + + hidden_states = outputs[0] + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function( + logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs + ) + + return PaliGemmaCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=outputs.image_hidden_states, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + inputs_embeds=None, + cache_position=None, + position_ids=None, + pixel_values=None, + attention_mask=None, + token_type_ids=None, + use_cache=True, + logits_to_keep=None, + labels=None, + **kwargs, + ): + # Overwritten -- custom `position_ids` and `pixel_values` handling + model_inputs = super().prepare_inputs_for_generation( + input_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + cache_position=cache_position, + use_cache=use_cache, + logits_to_keep=logits_to_keep, + token_type_ids=token_type_ids, + **kwargs, + ) + + # position_ids in Paligemma are 1-indexed + if model_inputs.get("position_ids") is not None: + model_inputs["position_ids"] += 1 + # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore + # Otherwise we need pixel values to be passed to model. NOTE: use_cache=False needs pixel_values always + if cache_position[0] == 0: + model_inputs["pixel_values"] = pixel_values + is_training = token_type_ids is not None and labels is not None + if cache_position[0] == 0 and isinstance(past_key_values, HybridCache): + input_tensor = inputs_embeds if inputs_embeds is not None else input_ids + causal_mask = self.model._update_causal_mask( + attention_mask, token_type_ids, past_key_values, cache_position, input_tensor, is_training + ) + model_inputs["attention_mask"] = causal_mask + + return model_inputs + + @staticmethod + # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position + def _prepare_4d_causal_attention_mask_with_cache_position( + attention_mask: torch.Tensor, + sequence_length: int, + target_length: int, + dtype: torch.dtype, + cache_position: torch.Tensor, + batch_size: int, + **kwargs, + ): + """ + Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. + + Args: + attention_mask (`torch.Tensor`): + A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape + `(batch_size, 1, query_length, key_value_length)`. + sequence_length (`int`): + The sequence length being processed. + target_length (`int`): + The target length: when generating with static cache, the mask should be as long as the static cache, + to account for the 0 padding, the part of the cache that is not filled yet. + dtype (`torch.dtype`): + The dtype to use for the 4D attention mask. + cache_position (`torch.Tensor`): + Indices depicting the position of the input sequence tokens in the sequence. + batch_size (`torch.Tensor`): + Batch size. + """ + if attention_mask is not None and attention_mask.dim() == 4: + # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. + causal_mask = attention_mask + else: + min_dtype = torch.finfo(dtype).min + causal_mask = torch.full( + (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device + ) + if sequence_length != 1: + causal_mask = torch.triu(causal_mask, diagonal=1) + causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) + causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) + if attention_mask is not None: + causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit + mask_length = attention_mask.shape[-1] + padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( + causal_mask.device + ) + padding_mask = padding_mask == 0 + causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( + padding_mask, min_dtype + ) + + return causal_mask + + +__all__ = ["PaliGemmaForConditionalGeneration", "PaliGemmaPreTrainedModel", "PaliGemmaModel"] diff --git a/lightx2v/models/networks/openpi/transformers_replace/models/siglip/check.py b/lightx2v/models/networks/openpi/transformers_replace/models/siglip/check.py new file mode 100644 index 000000000..a4572df79 --- /dev/null +++ b/lightx2v/models/networks/openpi/transformers_replace/models/siglip/check.py @@ -0,0 +1,4 @@ +import transformers + +def check_whether_transformers_replace_is_installed_correctly(): + return transformers.__version__ == "4.53.2" diff --git a/lightx2v/models/networks/openpi/transformers_replace/models/siglip/modeling_siglip.py b/lightx2v/models/networks/openpi/transformers_replace/models/siglip/modeling_siglip.py new file mode 100644 index 000000000..81b97ed5c --- /dev/null +++ b/lightx2v/models/networks/openpi/transformers_replace/models/siglip/modeling_siglip.py @@ -0,0 +1,1237 @@ +# coding=utf-8 +# Copyright 2024 Google AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch Siglip model.""" + +import math +import warnings +from dataclasses import dataclass +from typing import Any, Callable, Optional, Union + +import numpy as np +import torch +import torch.utils.checkpoint +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss +from torch.nn.init import _calculate_fan_in_and_fan_out + +from ...activations import ACT2FN +from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ImageClassifierOutput +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...utils import ModelOutput, auto_docstring, can_return_tuple, logging, torch_int +from .configuration_siglip import SiglipConfig, SiglipTextConfig, SiglipVisionConfig + + +logger = logging.get_logger(__name__) + + +def _trunc_normal_(tensor, mean, std, a, b): + # Cut & paste from PyTorch official master until it's in a few official releases - RW + # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf + def norm_cdf(x): + # Computes standard normal cumulative distribution function + return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0 + + if (mean < a - 2 * std) or (mean > b + 2 * std): + warnings.warn( + "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. " + "The distribution of values may be incorrect.", + stacklevel=2, + ) + + # Values are generated by using a truncated uniform distribution and + # then using the inverse CDF for the normal distribution. + # Get upper and lower cdf values + l = norm_cdf((a - mean) / std) + u = norm_cdf((b - mean) / std) + + # Uniformly fill tensor with values from [l, u], then translate to + # [2l-1, 2u-1]. + tensor.uniform_(2 * l - 1, 2 * u - 1) + + # Use inverse cdf transform for normal distribution to get truncated + # standard normal + tensor.erfinv_() + + # Transform to proper mean, std + tensor.mul_(std * math.sqrt(2.0)) + tensor.add_(mean) + + # Clamp to ensure it's in the proper range + tensor.clamp_(min=a, max=b) + + +def trunc_normal_tf_( + tensor: torch.Tensor, mean: float = 0.0, std: float = 1.0, a: float = -2.0, b: float = 2.0 +) -> torch.Tensor: + """Fills the input Tensor with values drawn from a truncated + normal distribution. The values are effectively drawn from the + normal distribution :math:`\\mathcal{N}(\text{mean}, \text{std}^2)` + with values outside :math:`[a, b]` redrawn until they are within + the bounds. The method used for generating the random values works + best when :math:`a \\leq \text{mean} \\leq b`. + + NOTE: this 'tf' variant behaves closer to Tensorflow / JAX impl where the + bounds [a, b] are applied when sampling the normal distribution with mean=0, std=1.0 + and the result is subsequently scaled and shifted by the mean and std args. + + Args: + tensor: an n-dimensional `torch.Tensor` + mean: the mean of the normal distribution + std: the standard deviation of the normal distribution + a: the minimum cutoff value + b: the maximum cutoff value + """ + with torch.no_grad(): + _trunc_normal_(tensor, 0, 1.0, a, b) + tensor.mul_(std).add_(mean) + + +def variance_scaling_(tensor, scale=1.0, mode="fan_in", distribution="normal"): + fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor) + if mode == "fan_in": + denom = fan_in + elif mode == "fan_out": + denom = fan_out + elif mode == "fan_avg": + denom = (fan_in + fan_out) / 2 + + variance = scale / denom + + if distribution == "truncated_normal": + # constant is stddev of standard normal truncated to (-2, 2) + trunc_normal_tf_(tensor, std=math.sqrt(variance) / 0.87962566103423978) + elif distribution == "normal": + with torch.no_grad(): + tensor.normal_(std=math.sqrt(variance)) + elif distribution == "uniform": + bound = math.sqrt(3 * variance) + with torch.no_grad(): + tensor.uniform_(-bound, bound) + else: + raise ValueError(f"invalid distribution {distribution}") + + +def lecun_normal_(tensor): + variance_scaling_(tensor, mode="fan_in", distribution="truncated_normal") + + +def default_flax_embed_init(tensor): + variance_scaling_(tensor, mode="fan_in", distribution="normal") + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for vision model's outputs that also contains image embeddings of the pooling of the last hidden states. + """ +) +# Copied from transformers.models.clip.modeling_clip.CLIPVisionModelOutput with CLIP->Siglip +class SiglipVisionModelOutput(ModelOutput): + r""" + image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`): + The image embeddings obtained by applying the projection layer to the pooler_output. + """ + + image_embeds: Optional[torch.FloatTensor] = None + last_hidden_state: Optional[torch.FloatTensor] = None + hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None + attentions: Optional[tuple[torch.FloatTensor, ...]] = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for text model's outputs that also contains a pooling of the last hidden states. + """ +) +# Copied from transformers.models.clip.modeling_clip.CLIPTextModelOutput with CLIP->Siglip +class SiglipTextModelOutput(ModelOutput): + r""" + text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`): + The text embeddings obtained by applying the projection layer to the pooler_output. + """ + + text_embeds: Optional[torch.FloatTensor] = None + last_hidden_state: Optional[torch.FloatTensor] = None + hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None + attentions: Optional[tuple[torch.FloatTensor, ...]] = None + + +@dataclass +@auto_docstring +# Copied from transformers.models.clip.modeling_clip.CLIPOutput with CLIP->Siglip +class SiglipOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`): + Contrastive loss for image-text similarity. + logits_per_image (`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`): + The scaled dot product scores between `image_embeds` and `text_embeds`. This represents the image-text + similarity scores. + logits_per_text (`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`): + The scaled dot product scores between `text_embeds` and `image_embeds`. This represents the text-image + similarity scores. + text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`): + The text embeddings obtained by applying the projection layer to the pooled output of [`SiglipTextModel`]. + image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`): + The image embeddings obtained by applying the projection layer to the pooled output of [`SiglipVisionModel`]. + text_model_output (`BaseModelOutputWithPooling`): + The output of the [`SiglipTextModel`]. + vision_model_output (`BaseModelOutputWithPooling`): + The output of the [`SiglipVisionModel`]. + """ + + loss: Optional[torch.FloatTensor] = None + logits_per_image: Optional[torch.FloatTensor] = None + logits_per_text: Optional[torch.FloatTensor] = None + text_embeds: Optional[torch.FloatTensor] = None + image_embeds: Optional[torch.FloatTensor] = None + text_model_output: BaseModelOutputWithPooling = None + vision_model_output: BaseModelOutputWithPooling = None + + def to_tuple(self) -> tuple[Any]: + return tuple( + self[k] if k not in ["text_model_output", "vision_model_output"] else getattr(self, k).to_tuple() + for k in self.keys() + ) + + +class SiglipVisionEmbeddings(nn.Module): + def __init__(self, config: SiglipVisionConfig): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.image_size = config.image_size + self.patch_size = config.patch_size + + self.patch_embedding = nn.Conv2d( + in_channels=config.num_channels, + out_channels=self.embed_dim, + kernel_size=self.patch_size, + stride=self.patch_size, + padding="valid", + ) + + self.num_patches = (self.image_size // self.patch_size) ** 2 + self.num_positions = self.num_patches + self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) + self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False) + + def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor: + """ + This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution + images. This method is also adapted to support torch.jit tracing and no class embeddings. + + Adapted from: + - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and + - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211 + """ + + num_patches = embeddings.shape[1] + num_positions = self.position_embedding.weight.shape[0] + + # always interpolate when tracing to ensure the exported model works for dynamic input shapes + if not torch.jit.is_tracing() and num_patches == num_positions and height == width: + return self.position_embedding(self.position_ids) + + patch_pos_embed = self.position_embedding.weight.unsqueeze(0) + + dim = embeddings.shape[-1] + + new_height = height // self.patch_size + new_width = width // self.patch_size + + sqrt_num_positions = torch_int(num_positions**0.5) + patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim) + patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) + + patch_pos_embed = nn.functional.interpolate( + patch_pos_embed, + size=(new_height, new_width), + mode="bicubic", + align_corners=False, + ) + + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) + return patch_pos_embed + + def forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding=False) -> torch.Tensor: + _, _, height, width = pixel_values.shape + target_dtype = self.patch_embedding.weight.dtype + patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) # shape = [*, width, grid, grid] + embeddings = patch_embeds.flatten(2).transpose(1, 2) + + if interpolate_pos_encoding: + embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width) + else: + embeddings = embeddings + self.position_embedding(self.position_ids) + return embeddings + + +# Copied from transformers.models.clip.modeling_clip.CLIPTextEmbeddings with CLIP->Siglip +class SiglipTextEmbeddings(nn.Module): + def __init__(self, config: SiglipTextConfig): + super().__init__() + embed_dim = config.hidden_size + + self.token_embedding = nn.Embedding(config.vocab_size, embed_dim) + self.position_embedding = nn.Embedding(config.max_position_embeddings, embed_dim) + + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.register_buffer( + "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False + ) + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + ) -> torch.Tensor: + seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2] + max_position_embedding = self.position_embedding.weight.shape[0] + + if seq_length > max_position_embedding: + raise ValueError( + f"Sequence length must be less than max_position_embeddings (got `sequence length`: " + f"{seq_length} and max_position_embeddings: {max_position_embedding}" + ) + + if position_ids is None: + position_ids = self.position_ids[:, :seq_length] + + if inputs_embeds is None: + inputs_embeds = self.token_embedding(input_ids) + + position_embeddings = self.position_embedding(position_ids) + embeddings = inputs_embeds + position_embeddings + + return embeddings + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + scaling: float, + dropout: float = 0.0, + **kwargs, +): + attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + + attn_output = torch.matmul(attn_weights, value) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +class SiglipAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.embed_dim // self.num_heads + if self.head_dim * self.num_heads != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:" + f" {self.num_heads})." + ) + self.scale = self.head_dim**-0.5 + self.dropout = config.attention_dropout + self.is_causal = False + + self.k_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = False, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Input shape: Batch x Time x Channel""" + + batch_size, seq_length, embed_dim = hidden_states.shape + + queries = self.q_proj(hidden_states) + keys = self.k_proj(hidden_states) + values = self.v_proj(hidden_states) + + queries = queries.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) + keys = keys.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) + values = values.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation == "sdpa" and output_attentions: + logger.warning_once( + "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to " + 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + else: + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, attn_weights = attention_interface( + self, + queries, + keys, + values, + attention_mask, + is_causal=self.is_causal, + scaling=self.scale, + dropout=0.0 if not self.training else self.dropout, + ) + + attn_output = attn_output.reshape(batch_size, seq_length, embed_dim).contiguous() + attn_output = self.out_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights + + +# Copied from transformers.models.clip.modeling_clip.CLIPMLP with CLIP->Siglip +class SiglipMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.activation_fn = ACT2FN[config.hidden_act] + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = self.activation_fn(hidden_states) + hidden_states = self.fc2(hidden_states) + return hidden_states + + +class SiglipEncoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Union[SiglipVisionConfig, SiglipTextConfig]): + super().__init__() + self.embed_dim = config.hidden_size + self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + self.self_attn = SiglipAttention(config) + self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + self.mlp = SiglipMLP(config) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + output_attentions: Optional[bool] = False, + ) -> tuple[torch.FloatTensor]: + """ + Args: + hidden_states (`torch.FloatTensor`): + Input to the layer of shape `(batch, seq_len, embed_dim)`. + attention_mask (`torch.FloatTensor`): + Attention mask of shape `(batch, 1, q_len, k_v_seq_len)` where padding elements are indicated by very large negative values. + output_attentions (`bool`, *optional*, defaults to `False`): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + """ + residual = hidden_states + + hidden_states = self.layer_norm1(hidden_states) + hidden_states, attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + output_attentions=output_attentions, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.layer_norm2(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (attn_weights,) + + return outputs + + +@auto_docstring +class SiglipPreTrainedModel(PreTrainedModel): + config_class = SiglipConfig + base_model_prefix = "siglip" + supports_gradient_checkpointing = True + + _no_split_modules = [ + "SiglipTextEmbeddings", + "SiglipEncoderLayer", + "SiglipVisionEmbeddings", + "SiglipEncoderLayer", + "SiglipMultiheadAttentionPoolingHead", + ] + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_flex_attn = True + _supports_attention_backend = True + + def _init_weights(self, module): + """Initialize the weights""" + if isinstance(module, SiglipVisionEmbeddings): + width = ( + self.config.vision_config.hidden_size + if isinstance(self.config, SiglipConfig) + else self.config.hidden_size + ) + nn.init.normal_(module.position_embedding.weight, std=1 / np.sqrt(width)) + elif isinstance(module, nn.Embedding): + default_flax_embed_init(module.weight) + elif isinstance(module, SiglipAttention): + nn.init.xavier_uniform_(module.q_proj.weight) + nn.init.xavier_uniform_(module.k_proj.weight) + nn.init.xavier_uniform_(module.v_proj.weight) + nn.init.xavier_uniform_(module.out_proj.weight) + nn.init.zeros_(module.q_proj.bias) + nn.init.zeros_(module.k_proj.bias) + nn.init.zeros_(module.v_proj.bias) + nn.init.zeros_(module.out_proj.bias) + elif isinstance(module, SiglipMLP): + nn.init.xavier_uniform_(module.fc1.weight) + nn.init.xavier_uniform_(module.fc2.weight) + nn.init.normal_(module.fc1.bias, std=1e-6) + nn.init.normal_(module.fc2.bias, std=1e-6) + elif isinstance(module, SiglipMultiheadAttentionPoolingHead): + nn.init.xavier_uniform_(module.probe.data) + nn.init.xavier_uniform_(module.attention.in_proj_weight.data) + nn.init.zeros_(module.attention.in_proj_bias.data) + elif isinstance(module, SiglipModel): + logit_scale_init = torch.log(torch.tensor(1.0)) + module.logit_scale.data.fill_(logit_scale_init) + module.logit_bias.data.zero_() + elif isinstance(module, SiglipForImageClassification): + nn.init.normal_( + module.classifier.weight, + std=self.config.vision_config.hidden_size**-0.5 * self.config.initializer_factor, + ) + elif isinstance(module, (nn.Linear, nn.Conv2d)): + lecun_normal_(module.weight) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.LayerNorm): + module.bias.data.zero_() + module.weight.data.fill_(1.0) + + +# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoder with AltCLIP->Siglip +class SiglipEncoder(nn.Module): + """ + Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a + [`SiglipEncoderLayer`]. + + Args: + config: SiglipConfig + """ + + def __init__(self, config: SiglipConfig): + super().__init__() + self.config = config + self.layers = nn.ModuleList([SiglipEncoderLayer(config) for _ in range(config.num_hidden_layers)]) + self.gradient_checkpointing = False + + # Ignore copy + @can_return_tuple + def forward( + self, + inputs_embeds, + attention_mask: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + ) -> BaseModelOutput: + r""" + Args: + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. + This is useful if you want more control over how to convert `input_ids` indices into associated vectors + than the model's internal embedding lookup matrix. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors + for more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + """ + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + + encoder_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + hidden_states = inputs_embeds + for encoder_layer in self.layers: + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + + layer_outputs = encoder_layer( + hidden_states, + attention_mask, + output_attentions=output_attentions, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_attentions = all_attentions + (layer_outputs[1],) + + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + + return BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=encoder_states, + attentions=all_attentions, + ) + + +class SiglipTextTransformer(nn.Module): + def __init__(self, config: SiglipTextConfig): + super().__init__() + self.config = config + embed_dim = config.hidden_size + self.embeddings = SiglipTextEmbeddings(config) + self.encoder = SiglipEncoder(config) + self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + + self.head = nn.Linear(embed_dim, config.projection_size) + self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + ) -> BaseModelOutputWithPooling: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + + if input_ids is None: + raise ValueError("You have to specify input_ids") + + input_shape = input_ids.size() + input_ids = input_ids.view(-1, input_shape[-1]) + + hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids) + + # note: SigLIP's text model does not use a causal mask, unlike the original CLIP model. + # expand attention_mask + if attention_mask is not None and not self._use_flash_attention_2: + # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len] + attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) + + encoder_outputs: BaseModelOutput = self.encoder( + inputs_embeds=hidden_states, + attention_mask=attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + + last_hidden_state = encoder_outputs.last_hidden_state + last_hidden_state = self.final_layer_norm(last_hidden_state) + + # Assuming "sticky" EOS tokenization, last token is always EOS. + pooled_output = last_hidden_state[:, -1, :] + pooled_output = self.head(pooled_output) + + return BaseModelOutputWithPooling( + last_hidden_state=last_hidden_state, + pooler_output=pooled_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + The text model from SigLIP without any head or projection on top. + """ +) +class SiglipTextModel(SiglipPreTrainedModel): + config_class = SiglipTextConfig + + def __init__(self, config: SiglipTextConfig): + super().__init__(config) + self.text_model = SiglipTextTransformer(config) + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.text_model.embeddings.token_embedding + + def set_input_embeddings(self, value): + self.text_model.embeddings.token_embedding = value + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + ) -> BaseModelOutputWithPooling: + r""" + Examples: + + ```python + >>> from transformers import AutoTokenizer, SiglipTextModel + + >>> model = SiglipTextModel.from_pretrained("google/siglip-base-patch16-224") + >>> tokenizer = AutoTokenizer.from_pretrained("google/siglip-base-patch16-224") + + >>> # important: make sure to set padding="max_length" as that's how the model was trained + >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding="max_length", return_tensors="pt") + + >>> outputs = model(**inputs) + >>> last_hidden_state = outputs.last_hidden_state + >>> pooled_output = outputs.pooler_output # pooled (EOS token) states + ```""" + + return self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + + +class SiglipVisionTransformer(nn.Module): + def __init__(self, config: SiglipVisionConfig): + super().__init__() + self.config = config + embed_dim = config.hidden_size + + self.embeddings = SiglipVisionEmbeddings(config) + self.encoder = SiglipEncoder(config) + self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + self.use_head = True if not hasattr(config, "vision_use_head") else config.vision_use_head + if self.use_head: + self.head = SiglipMultiheadAttentionPoolingHead(config) + + @can_return_tuple + @auto_docstring + def forward( + self, + pixel_values, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + interpolate_pos_encoding: Optional[bool] = False, + ) -> BaseModelOutputWithPooling: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + + hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding) + # Convert to bfloat16 if the encoder uses bfloat16 + if len(self.encoder.layers) > 0 and self.encoder.layers[0].self_attn.q_proj.weight.dtype == torch.bfloat16: + hidden_states = hidden_states.to(torch.bfloat16) + + encoder_outputs: BaseModelOutput = self.encoder( + inputs_embeds=hidden_states, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + + last_hidden_state = encoder_outputs.last_hidden_state + last_hidden_state = self.post_layernorm(last_hidden_state) + + pooler_output = self.head(last_hidden_state) if self.use_head else None + + return BaseModelOutputWithPooling( + last_hidden_state=last_hidden_state, + pooler_output=pooler_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + +class SiglipMultiheadAttentionPoolingHead(nn.Module): + """Multihead Attention Pooling.""" + + def __init__(self, config: SiglipVisionConfig): + super().__init__() + + self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size)) + self.attention = torch.nn.MultiheadAttention(config.hidden_size, config.num_attention_heads, batch_first=True) + self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.mlp = SiglipMLP(config) + + def forward(self, hidden_state): + batch_size = hidden_state.shape[0] + probe = self.probe.repeat(batch_size, 1, 1) + + hidden_state = self.attention(probe, hidden_state, hidden_state)[0] + + residual = hidden_state + hidden_state = self.layernorm(hidden_state) + hidden_state = residual + self.mlp(hidden_state) + + return hidden_state[:, 0] + + +@auto_docstring( + custom_intro=""" + The vision model from SigLIP without any head or projection on top. + """ +) +class SiglipVisionModel(SiglipPreTrainedModel): + config_class = SiglipVisionConfig + main_input_name = "pixel_values" + + def __init__(self, config: SiglipVisionConfig): + super().__init__(config) + + self.vision_model = SiglipVisionTransformer(config) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.vision_model.embeddings.patch_embedding + + @can_return_tuple + @auto_docstring + def forward( + self, + pixel_values, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + interpolate_pos_encoding: bool = False, + ) -> BaseModelOutputWithPooling: + r""" + Examples: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, SiglipVisionModel + + >>> model = SiglipVisionModel.from_pretrained("google/siglip-base-patch16-224") + >>> processor = AutoProcessor.from_pretrained("google/siglip-base-patch16-224") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(images=image, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> last_hidden_state = outputs.last_hidden_state + >>> pooled_output = outputs.pooler_output # pooled features + ```""" + + return self.vision_model( + pixel_values=pixel_values, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + interpolate_pos_encoding=interpolate_pos_encoding, + ) + + +@auto_docstring +class SiglipModel(SiglipPreTrainedModel): + config_class = SiglipConfig + + def __init__(self, config: SiglipConfig): + super().__init__(config) + + if not isinstance(config.text_config, SiglipTextConfig): + raise TypeError( + "config.text_config is expected to be of type SiglipTextConfig but is of type" + f" {type(config.text_config)}." + ) + + if not isinstance(config.vision_config, SiglipVisionConfig): + raise TypeError( + "config.vision_config is expected to be of type SiglipVisionConfig but is of type" + f" {type(config.vision_config)}." + ) + + text_config = config.text_config + vision_config = config.vision_config + + # First, initialize the text and vision models with proper attention implementation + text_model = SiglipTextModel._from_config(text_config) + vision_model = SiglipVisionModel._from_config(vision_config) + + # Second, get the text and vision submodules (for backward compatibility) + self.text_model = text_model.text_model + self.vision_model = vision_model.vision_model + + self.logit_scale = nn.Parameter(torch.randn(1)) + self.logit_bias = nn.Parameter(torch.randn(1)) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def get_text_features( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + ) -> torch.FloatTensor: + r""" + Returns: + text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by + applying the projection layer to the pooled output of [`SiglipTextModel`]. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, AutoModel + >>> import torch + + >>> model = AutoModel.from_pretrained("google/siglip-base-patch16-224") + >>> tokenizer = AutoTokenizer.from_pretrained("google/siglip-base-patch16-224") + + >>> # important: make sure to set padding="max_length" as that's how the model was trained + >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding="max_length", return_tensors="pt") + >>> with torch.no_grad(): + ... text_features = model.get_text_features(**inputs) + ```""" + # Use SigLIP model's config for some fields (if specified) instead of those of vision & text components. + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + + text_outputs: BaseModelOutputWithPooling = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + + pooled_output = text_outputs.pooler_output + + return pooled_output + + @auto_docstring + def get_image_features( + self, + pixel_values: Optional[torch.FloatTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + interpolate_pos_encoding: bool = False, + ) -> torch.FloatTensor: + r""" + Returns: + image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by + applying the projection layer to the pooled output of [`SiglipVisionModel`]. + + Examples: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, AutoModel + >>> import torch + + >>> model = AutoModel.from_pretrained("google/siglip-base-patch16-224") + >>> processor = AutoProcessor.from_pretrained("google/siglip-base-patch16-224") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(images=image, return_tensors="pt") + + >>> with torch.no_grad(): + ... image_features = model.get_image_features(**inputs) + ```""" + # Use SiglipModel's config for some fields (if specified) instead of those of vision & text components. + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + + vision_outputs: BaseModelOutputWithPooling = self.vision_model( + pixel_values=pixel_values, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + interpolate_pos_encoding=interpolate_pos_encoding, + ) + + pooled_output = vision_outputs.pooler_output + + return pooled_output + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + pixel_values: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + return_loss: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + interpolate_pos_encoding: bool = False, + ) -> SiglipOutput: + r""" + return_loss (`bool`, *optional*): + Whether or not to return the contrastive loss. + + Examples: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, AutoModel + >>> import torch + + >>> model = AutoModel.from_pretrained("google/siglip-base-patch16-224") + >>> processor = AutoProcessor.from_pretrained("google/siglip-base-patch16-224") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> texts = ["a photo of 2 cats", "a photo of 2 dogs"] + >>> # important: we pass `padding=max_length` since the model was trained with this + >>> inputs = processor(text=texts, images=image, padding="max_length", return_tensors="pt") + + >>> with torch.no_grad(): + ... outputs = model(**inputs) + + >>> logits_per_image = outputs.logits_per_image + >>> probs = torch.sigmoid(logits_per_image) # these are the probabilities + >>> print(f"{probs[0][0]:.1%} that image 0 is '{texts[0]}'") + 31.9% that image 0 is 'a photo of 2 cats' + ```""" + # Use SigLIP model's config for some fields (if specified) instead of those of vision & text components. + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + + vision_outputs: BaseModelOutputWithPooling = self.vision_model( + pixel_values=pixel_values, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + interpolate_pos_encoding=interpolate_pos_encoding, + ) + + text_outputs: BaseModelOutputWithPooling = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + + image_embeds = vision_outputs.pooler_output + text_embeds = text_outputs.pooler_output + + # normalized features + image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True) + text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True) + + # cosine similarity as logits + logits_per_text = torch.matmul(text_embeds, image_embeds.t().to(text_embeds.device)) + + logit_scale, logit_bias = self.logit_scale.to(text_embeds.device), self.logit_bias.to(text_embeds.device) + logits_per_text = logits_per_text * logit_scale.exp() + logit_bias + + logits_per_image = logits_per_text.t() + + loss = None + if return_loss: + # Adapted from https://github.com/google-research/big_vision/blob/01edb81a4716f93a48be43b3a4af14e29cdb3a7f/big_vision/trainers/proj/image_text/siglip.py#L287 + eye = torch.eye(logits_per_text.size(0), device=logits_per_text.device) + m1_diag1 = -torch.ones_like(logits_per_text) + 2 * eye + loglik = torch.nn.functional.logsigmoid(m1_diag1 * logits_per_text) + nll = -torch.sum(loglik, dim=-1) + loss = nll.mean() + + return SiglipOutput( + loss=loss, + logits_per_image=logits_per_image, + logits_per_text=logits_per_text, + text_embeds=text_embeds, + image_embeds=image_embeds, + text_model_output=text_outputs, + vision_model_output=vision_outputs, + ) + + +@auto_docstring( + custom_intro=""" + SigLIP vision encoder with an image classification head on top (a linear layer on top of the pooled final hidden states of + the patch tokens) e.g. for ImageNet. + """ +) +class SiglipForImageClassification(SiglipPreTrainedModel): + main_input_name = "pixel_values" + + def __init__(self, config: SiglipConfig) -> None: + super().__init__(config) + + self.num_labels = config.num_labels + + # Create the vision model with proper attention + # and take only vision_model submodule (for backward compatibility) + vision_model = SiglipVisionModel._from_config(config.vision_config) + self.vision_model = vision_model.vision_model + + # Classifier head + self.classifier = ( + nn.Linear(config.vision_config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity() + ) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + pixel_values: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + interpolate_pos_encoding: bool = False, + ) -> ImageClassifierOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the image classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + + Examples: + + ```python + >>> from transformers import AutoImageProcessor, SiglipForImageClassification + >>> import torch + >>> from PIL import Image + >>> import requests + + >>> torch.manual_seed(3) # doctest: +IGNORE_RESULT + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> # note: we are loading a `SiglipModel` from the hub here, + >>> # so the head will be randomly initialized, hence the predictions will be random if seed is not set above. + >>> image_processor = AutoImageProcessor.from_pretrained("google/siglip-base-patch16-224") + >>> model = SiglipForImageClassification.from_pretrained("google/siglip-base-patch16-224") + + >>> inputs = image_processor(images=image, return_tensors="pt") + >>> outputs = model(**inputs) + >>> logits = outputs.logits + >>> # model predicts one of the two classes + >>> predicted_class_idx = logits.argmax(-1).item() + >>> print("Predicted class:", model.config.id2label[predicted_class_idx]) + Predicted class: LABEL_1 + ```""" + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + + outputs: BaseModelOutputWithPooling = self.vision_model( + pixel_values, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + interpolate_pos_encoding=interpolate_pos_encoding, + ) + + sequence_output = outputs.last_hidden_state + + # average pool the patch tokens + sequence_output = torch.mean(sequence_output, dim=1) + # apply classifier + logits = self.classifier(sequence_output) + + loss = None + if labels is not None: + # move labels to correct device to enable model parallelism + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + + return ImageClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = [ + "SiglipModel", + "SiglipPreTrainedModel", + "SiglipTextModel", + "SiglipVisionModel", + "SiglipForImageClassification", +] diff --git a/lightx2v/models/networks/openpi/weights/__init__.py b/lightx2v/models/networks/openpi/weights/__init__.py new file mode 100644 index 000000000..42d4cc820 --- /dev/null +++ b/lightx2v/models/networks/openpi/weights/__init__.py @@ -0,0 +1,3 @@ +from .loader import load_pi05_libero_weights + +__all__ = ["load_pi05_libero_weights"] diff --git a/lightx2v/models/networks/openpi/weights/loader.py b/lightx2v/models/networks/openpi/weights/loader.py new file mode 100644 index 000000000..7199e02de --- /dev/null +++ b/lightx2v/models/networks/openpi/weights/loader.py @@ -0,0 +1,59 @@ +"""Strict SafeTensors loader for the converted pi0.5-LIBERO checkpoint.""" + +from __future__ import annotations + +import logging +from pathlib import Path + +import torch +from safetensors.torch import load_model + +from ..config import Pi0Config + +LOGGER = logging.getLogger(__name__) + + +def _validate_transformers_runtime() -> None: + """Fail early unless the official patched Transformers runtime is active.""" + import transformers + + if transformers.__version__ != "4.53.2": + raise RuntimeError( + "OpenPI requires its private patched transformers==4.53.2 runtime; " + f"the current process imported transformers=={transformers.__version__}. " + "Launch with scripts/openpi/run_libero_*.sh or prepend " + "OPENPI_TRANSFORMERS_RUNTIME_PATH to PYTHONPATH." + ) + try: + from transformers.models.siglip import check + except ImportError as exc: + raise RuntimeError("OpenPI Transformers patches are missing (siglip/check.py not found)") from exc + if not check.check_whether_transformers_replace_is_installed_correctly(): + raise RuntimeError("OpenPI Transformers 4.53.2 is present but the official replacement patches are missing") + + +def load_pi05_libero_weights( + weight_path: str | Path, + config: Pi0Config, + device: torch.device | str, +): + """Build the exact official parameter tree and load it with strict key checks.""" + _validate_transformers_runtime() + config.validate_pi05_libero() + weight_path = Path(weight_path).expanduser().resolve() + if not weight_path.is_file(): + raise FileNotFoundError(f"Converted OpenPI SafeTensors file not found: {weight_path}") + + # Import only after selecting OpenPI's private Transformers runtime. + from ..pi0 import PI0Pytorch + + model = PI0Pytorch(config) + load_model(model, weight_path, strict=True, device="cpu") + + # Match upstream's mixed BF16/FP32 inference policy. + model.paligemma_with_expert.to_bfloat16_for_selected_params(config.dtype) + model.to(torch.device(device)) + model.eval() + parameter_count = sum(parameter.numel() for parameter in model.parameters()) + LOGGER.info("Loaded pi05_libero PyTorch weights strictly: %.3fB parameters", parameter_count / 1e9) + return model diff --git a/lightx2v/models/runners/openpi/__init__.py b/lightx2v/models/runners/openpi/__init__.py new file mode 100644 index 000000000..bb80255b9 --- /dev/null +++ b/lightx2v/models/runners/openpi/__init__.py @@ -0,0 +1 @@ +"""OpenPI policy and local LIBERO runner integration.""" diff --git a/lightx2v/models/runners/openpi/artifacts.py b/lightx2v/models/runners/openpi/artifacts.py new file mode 100644 index 000000000..ca8aee232 --- /dev/null +++ b/lightx2v/models/runners/openpi/artifacts.py @@ -0,0 +1,176 @@ +"""Durable, lightweight artifacts for local OpenPI LIBERO evaluation.""" + +from __future__ import annotations + +import json +import logging +import os +import uuid +from collections.abc import Iterable, Mapping +from pathlib import Path +from typing import Any + +import numpy as np + +LOGGER = logging.getLogger(__name__) + +EpisodeKey = tuple[str, int, int] + + +def atomic_write_text(path: Path, text: str) -> None: + """Replace ``path`` only after the complete payload reaches disk.""" + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.{os.getpid()}.{uuid.uuid4().hex}.tmp") + try: + with temporary.open("w", encoding="utf-8") as handle: + handle.write(text) + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + finally: + temporary.unlink(missing_ok=True) + + +def atomic_write_json(path: Path, payload: Mapping[str, Any]) -> None: + atomic_write_text(Path(path), json.dumps(dict(payload), indent=2, sort_keys=True) + "\n") + + +def atomic_save_numpy(path: Path, array: np.ndarray) -> None: + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.{os.getpid()}.{uuid.uuid4().hex}.tmp.npy") + try: + np.save(temporary, np.asarray(array)) + os.replace(temporary, path) + finally: + temporary.unlink(missing_ok=True) + + +def atomic_write_video(path: Path, frames: list[np.ndarray], fps: int) -> None: + if not frames: + raise ValueError("Cannot write a LIBERO video without policy-phase frames") + if fps < 1: + raise ValueError(f"Video fps must be positive, got {fps}") + + import imageio.v2 as imageio + + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.stem}.{os.getpid()}.{uuid.uuid4().hex}.tmp{path.suffix}") + try: + imageio.mimwrite(temporary, [np.asarray(frame) for frame in frames], fps=fps) + os.replace(temporary, path) + finally: + temporary.unlink(missing_ok=True) + + +def episode_directory(output_dir: Path, benchmark: str, task_id: int, init_state_id: int) -> Path: + return Path(output_dir) / "episodes" / benchmark / f"task_{task_id:02d}" / f"init_{init_state_id:02d}" + + +def episode_metrics_path(output_dir: Path, key: EpisodeKey) -> Path: + benchmark, task_id, init_state_id = key + return episode_directory(output_dir, benchmark, task_id, init_state_id) / "metrics.json" + + +def record_key(record: Mapping[str, Any]) -> EpisodeKey: + try: + return str(record["benchmark"]), int(record["task_id"]), int(record["init_state_id"]) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError(f"Invalid LIBERO episode record identity: {record!r}") from exc + + +def read_episode_records( + output_dir: Path, + expected_keys: Iterable[EpisodeKey], + protocol_id: str, +) -> dict[EpisodeKey, dict[str, Any]]: + """Read committed episode records; temporary files are deliberately ignored.""" + expected = tuple(expected_keys) + expected_set = set(expected) + if len(expected_set) != len(expected): + raise ValueError("Expected LIBERO episode keys contain duplicates") + + records: dict[EpisodeKey, dict[str, Any]] = {} + for key in expected: + path = episode_metrics_path(output_dir, key) + if not path.is_file(): + continue + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise RuntimeError(f"Cannot read committed LIBERO metrics: {path}") from exc + if not isinstance(payload, dict): + raise RuntimeError(f"LIBERO metrics must contain a JSON object: {path}") + if record_key(payload) != key: + raise RuntimeError(f"LIBERO metrics identity does not match its path: {path}") + if payload.get("protocol_id") != protocol_id: + raise RuntimeError(f"Existing episode uses a different evaluation protocol: {path}") + records[key] = payload + + episodes_root = Path(output_dir) / "episodes" + if episodes_root.is_dir(): + for path in episodes_root.rglob("metrics.json"): + try: + payload = json.loads(path.read_text(encoding="utf-8")) + key = record_key(payload) + except (OSError, json.JSONDecodeError, ValueError) as exc: + raise RuntimeError(f"Invalid committed LIBERO metrics outside the requested table: {path}") from exc + if key not in expected_set: + raise RuntimeError(f"Evaluation output contains an episode outside the requested selection: {path}") + + return records + + +def save_evaluation_episode( + output_dir: Path, + record: Mapping[str, Any], + frames: list[np.ndarray], + actions: np.ndarray, + *, + video_policy: str, + video_fps: int, + save_actions: bool, +) -> dict[str, Any]: + """Save optional payloads first and atomically commit ``metrics.json`` last.""" + saved = dict(record) + key = record_key(saved) + episode_dir = episode_directory(output_dir, *key) + episode_dir.mkdir(parents=True, exist_ok=True) + artifact_errors: list[str] = [] + + action_path: Path | None = None + if save_actions: + action_path = episode_dir / "actions.npy" + try: + atomic_save_numpy(action_path, actions) + except Exception as exc: # Artifact failures must not erase a completed rollout. + artifact_errors.append(f"actions: {type(exc).__name__}: {exc}") + LOGGER.exception("Failed to save LIBERO episode actions") + + save_video = video_policy == "all" or (video_policy == "failures" and not bool(saved.get("success"))) + video_path: Path | None = None + if save_video: + video_path = episode_dir / "rollout.mp4" + try: + atomic_write_video(video_path, frames, video_fps) + except Exception as exc: # Metrics remain the authoritative completion marker. + artifact_errors.append(f"video: {type(exc).__name__}: {exc}") + LOGGER.exception("Failed to save LIBERO episode video") + + saved["action_path"] = str(action_path) if action_path is not None and action_path.is_file() else None + saved["video_path"] = str(video_path) if video_path is not None and video_path.is_file() else None + saved["artifact_errors"] = artifact_errors + atomic_write_json(episode_dir / "metrics.json", saved) + return saved + + +def write_aggregate_outputs(output_dir: Path, records: Iterable[Mapping[str, Any]], summary: Mapping[str, Any]) -> None: + ordered = sorted( + (dict(record) for record in records), + key=lambda item: (int(item["benchmark_index"]), int(item["task_id"]), int(item["init_state_id"])), + ) + lines = "".join(json.dumps(record, sort_keys=True) + "\n" for record in ordered) + atomic_write_text(Path(output_dir) / "episodes.jsonl", lines) + atomic_write_json(Path(output_dir) / "summary.json", summary) diff --git a/lightx2v/models/runners/openpi/libero_evaluate.py b/lightx2v/models/runners/openpi/libero_evaluate.py new file mode 100644 index 000000000..d3480d864 --- /dev/null +++ b/lightx2v/models/runners/openpi/libero_evaluate.py @@ -0,0 +1,342 @@ +"""Quantitative, resumable local evaluation for the pi0.5-LIBERO policy.""" + +from __future__ import annotations + +import argparse +import json +import logging +import time +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import numpy as np + +from .artifacts import EpisodeKey, atomic_write_json, read_episode_records, save_evaluation_episode, write_aggregate_outputs +from .libero_protocol import ( + OFFICIAL_BENCHMARKS, + OFFICIAL_RESULTS, + EvaluationConfig, + TaskSpec, + build_task_inputs_manifest, + build_task_specs, + configure_libero, + create_environment, + ensure_task_inputs_manifest, + expected_episode_keys, + is_official_protocol, + load_evaluation_config, + load_policy_config, + load_task_init_states, + resolved_protocol, + run_episode, +) +from .openpi_runner import OpenPIPolicy + +LOGGER = logging.getLogger(__name__) + + +def _suite_keys(specs: list[TaskSpec], config: EvaluationConfig) -> tuple[EpisodeKey, ...]: + return tuple((spec.benchmark, spec.task_id, init_state_id) for spec in specs for init_state_id in range(config.num_trials_per_task)) + + +def _resume_prefixes( + task_specs: dict[str, list[TaskSpec]], + config: EvaluationConfig, + records: Mapping[EpisodeKey, Mapping[str, Any]], +) -> dict[str, int]: + """Require a committed prefix so one saved RNG state resumes the suite exactly.""" + prefixes: dict[str, int] = {} + for benchmark in config.benchmarks: + ordered = _suite_keys(task_specs[benchmark], config) + present = [key in records for key in ordered] + prefix = 0 + while prefix < len(present) and present[prefix]: + prefix += 1 + if any(present[prefix:]): + raise RuntimeError(f"Episode-level resume for {benchmark} requires a contiguous prefix; a later episode exists after the first missing metrics.json.") + if 0 < prefix < len(ordered) and not records[ordered[prefix - 1]].get("policy_rng_state_after"): + raise RuntimeError(f"Partial suite {benchmark} predates RNG-checkpoint resume. Keep it for analysis, but use a fresh output directory to continue evaluation.") + prefixes[benchmark] = prefix + return prefixes + + +def _task_metrics(records: list[Mapping[str, Any]], expected: int) -> dict[str, Any]: + completed = len(records) + successes = sum(bool(record.get("success")) for record in records) + errors = sum(record.get("termination_reason") == "exception" for record in records) + return { + "expected": expected, + "completed": completed, + "successes": successes, + "failures": completed - successes, + "errors": errors, + "success_rate": successes / expected * 100.0 if completed == expected and expected else None, + "success_rate_completed": successes / completed * 100.0 if completed else None, + } + + +def build_summary( + records: Mapping[EpisodeKey, Mapping[str, Any]], + task_specs: dict[str, list[TaskSpec]], + config: EvaluationConfig, + protocol_id: str, +) -> dict[str, Any]: + official = is_official_protocol(config) + per_benchmark: dict[str, Any] = {} + per_task: dict[str, Any] = {} + expected_total = 0 + + for benchmark in config.benchmarks: + specs = task_specs[benchmark] + expected = len(specs) * config.num_trials_per_task + expected_total += expected + suite_records = [record for key, record in records.items() if key[0] == benchmark] + suite_metrics = _task_metrics(suite_records, expected) + reference = OFFICIAL_RESULTS.get(benchmark) + rate = suite_metrics["success_rate"] + suite_metrics.update( + { + "official_reference": reference, + "official_comparable": official, + "delta_percentage_points": rate - reference["success_rate"] if official and rate is not None and reference else None, + } + ) + per_benchmark[benchmark] = suite_metrics + + for spec in specs: + task_records = [record for key, record in records.items() if key[0] == benchmark and key[1] == spec.task_id] + metrics = _task_metrics(task_records, config.num_trials_per_task) + metrics.update( + { + "task_name": str(spec.task.name), + "task_description": str(spec.task.language), + } + ) + per_task[f"{benchmark}/task_{spec.task_id:02d}"] = metrics + + completed_total = len(records) + successes = sum(bool(record.get("success")) for record in records.values()) + errors = sum(record.get("termination_reason") == "exception" for record in records.values()) + artifact_errors = sum(len(record.get("artifact_errors", [])) for record in records.values()) + complete = completed_total == expected_total + suite_rates = [entry["success_rate"] for entry in per_benchmark.values()] + mean_suite_rate = sum(suite_rates) / len(suite_rates) if complete and suite_rates and all(rate is not None for rate in suite_rates) else None + return { + "schema_version": 1, + "status": "complete_with_errors" if complete and (errors or artifact_errors) else "complete" if complete else "in_progress", + "protocol_id": protocol_id, + "official_protocol": official, + "resume_granularity": "episode_prefix", + "policy_rng_scope": "per_suite_continuous", + "full_official_table": official and set(config.benchmarks) == set(OFFICIAL_BENCHMARKS), + "expected_episodes": expected_total, + "completed_episodes": completed_total, + "successes": successes, + "failures": completed_total - successes, + "errors": errors, + "artifact_errors": artifact_errors, + "success_rate": successes / expected_total * 100.0 if complete and expected_total else None, + "success_rate_completed": successes / completed_total * 100.0 if completed_total else None, + "mean_suite_success_rate": mean_suite_rate, + "official_reference": {"comparable": official, "results": OFFICIAL_RESULTS}, + "per_benchmark": per_benchmark, + "per_task": per_task, + "updated_at_unix": time.time(), + } + + +def _write_current_outputs( + output_dir: Path, + records: Mapping[EpisodeKey, Mapping[str, Any]], + task_specs: dict[str, list[TaskSpec]], + config: EvaluationConfig, + protocol_id: str, +) -> dict[str, Any]: + summary = build_summary(records, task_specs, config, protocol_id) + write_aggregate_outputs(output_dir, records.values(), summary) + return summary + + +def _prepare_output_directory(output_dir: Path, resolved: Mapping[str, Any]) -> None: + output_dir.mkdir(parents=True, exist_ok=True) + resolved_path = output_dir / "resolved_eval_config.json" + if resolved_path.is_file(): + try: + existing = json.loads(resolved_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise RuntimeError(f"Cannot read existing resolved evaluation config: {resolved_path}") from exc + if existing.get("protocol_id") != resolved["protocol_id"]: + committed = next((output_dir / "episodes").rglob("metrics.json"), None) if (output_dir / "episodes").is_dir() else None + if committed is not None: + raise RuntimeError(f"Evaluation output contains records for a different protocol: {committed}") + LOGGER.warning("Replacing unused resolved evaluation config: %s", resolved_path) + atomic_write_json(resolved_path, resolved) + + +def run_evaluation(args: argparse.Namespace) -> dict[str, Any]: + config = load_evaluation_config(args) + output_dir = Path(args.output_dir).expanduser().resolve() + np.random.seed(config.env_seed) + runtime = configure_libero(config.libero_root, config.libero_config_dir) + task_specs = build_task_specs(runtime, config) + task_inputs_manifest = build_task_inputs_manifest(task_specs, config) + resolved, protocol_id = resolved_protocol( + config, + model_path=Path(args.model_path), + config_json=Path(args.config_json), + ) + _prepare_output_directory(output_dir, resolved) + + all_keys = expected_episode_keys(task_specs, config) + records = read_episode_records(output_dir, all_keys, protocol_id) + if records and not config.resume: + raise FileExistsError(f"Evaluation output already contains {len(records)} committed episodes; use --resume or a new output directory.") + ensure_task_inputs_manifest(output_dir, task_inputs_manifest, records, task_specs) + prefixes = _resume_prefixes(task_specs, config, records) + summary = _write_current_outputs(output_dir, records, task_specs, config, protocol_id) + if len(records) == len(all_keys): + LOGGER.info("All requested LIBERO episodes are already complete: %s", output_dir) + return summary + + policy = OpenPIPolicy( + load_policy_config( + Path(args.config_json), + Path(args.model_path), + seed=config.policy_seed, + actions_per_plan=config.actions_per_plan, + ) + ) + global_indices = {key: index for index, key in enumerate(all_keys)} + try: + for benchmark_index, benchmark in enumerate(config.benchmarks): + ordered_suite_keys = _suite_keys(task_specs[benchmark], config) + prefix = prefixes[benchmark] + if prefix == len(ordered_suite_keys): + LOGGER.info("Episode resume: skipping complete suite %s", benchmark) + continue + + policy.clear_action_queue() + policy.reset_rng() + # Official evaluation starts each suite in its own process. Mirror + # that process-level NumPy seed when several suites share a worker. + np.random.seed(config.env_seed) + if prefix: + policy.import_rng_state(str(records[ordered_suite_keys[prefix - 1]]["policy_rng_state_after"])) + LOGGER.info("Episode resume: %s continues after %d/%d episodes", benchmark, prefix, len(ordered_suite_keys)) + else: + LOGGER.info("Starting suite %s", benchmark) + + for spec in task_specs[benchmark]: + task_keys = [(benchmark, spec.task_id, init_state_id) for init_state_id in range(config.num_trials_per_task)] + if all(key in records for key in task_keys): + continue + + initial_states, init_states_loader = load_task_init_states(spec) + if config.num_trials_per_task > len(initial_states): + raise ValueError(f"{benchmark} task {spec.task_id} has {len(initial_states)} init states, but {config.num_trials_per_task} trials were requested") + env = create_environment(runtime, spec, config.render_size, config.env_seed) + try: + for init_state_id, key in enumerate(task_keys): + if key in records: + # Reproduce the official per-task reset count before a + # partially completed task's first pending episode. + env.reset() + env.set_init_state(initial_states[init_state_id]) + continue + + episode, frames, actions = run_episode( + policy=policy, + env=env, + initial_state=initial_states[init_state_id], + task_description=str(spec.task.language), + max_steps=config.max_steps[benchmark], + num_steps_wait=config.num_steps_wait, + collect_frames=config.video_policy != "none", + ) + record = { + "schema_version": 1, + "protocol_id": protocol_id, + "global_episode_index": global_indices[key], + "benchmark_index": benchmark_index, + "benchmark": benchmark, + "task_id": spec.task_id, + "task_name": str(spec.task.name), + "task_description": str(spec.task.language), + "bddl_file": str(spec.bddl_path), + "init_states_file": str(spec.init_states_path), + "init_states_loader": init_states_loader, + "init_state_id": init_state_id, + "env_seed": config.env_seed, + "policy_seed": config.policy_seed, + "policy_rng_scope": "per_suite_continuous", + "policy_rng_state_after": policy.export_rng_state(), + "max_policy_steps": config.max_steps[benchmark], + "actions_per_plan": config.actions_per_plan, + **episode, + } + saved = save_evaluation_episode( + output_dir, + record, + frames, + actions, + video_policy=config.video_policy, + video_fps=config.video_fps, + save_actions=config.save_actions, + ) + records[key] = saved + summary = _write_current_outputs(output_dir, records, task_specs, config, protocol_id) + LOGGER.info( + "suite=%s task=%02d init=%02d success=%s completed=%d/%d", + benchmark, + spec.task_id, + init_state_id, + saved["success"], + summary["completed_episodes"], + summary["expected_episodes"], + ) + if saved["termination_reason"] == "exception" and config.fail_fast: + raise RuntimeError(f"Fail-fast: {benchmark} task {spec.task_id} init {init_state_id} failed: {saved['exception_type']}: {saved['exception_message']}") + finally: + env.close() + finally: + policy.close() + summary = _write_current_outputs(output_dir, records, task_specs, config, protocol_id) + + LOGGER.info( + "LIBERO evaluation finished: successes=%d/%d success_rate=%s mean_suite_success_rate=%s", + summary["successes"], + summary["expected_episodes"], + summary["success_rate"], + summary["mean_suite_success_rate"], + ) + return summary + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Run local quantitative pi0.5-LIBERO evaluation") + parser.add_argument("--model-path", type=Path, required=True) + parser.add_argument("--config-json", type=Path, required=True) + parser.add_argument("--eval-config", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--libero-root", type=Path, required=True) + parser.add_argument("--libero-config-dir", type=Path, required=True) + parser.add_argument("--benchmarks") + parser.add_argument("--task-ids") + parser.add_argument("--num-trials-per-task", type=int) + parser.add_argument("--max-steps", type=int) + parser.add_argument("--video-policy", choices=("all", "failures", "none")) + parser.add_argument("--resume", action=argparse.BooleanOptionalAction, default=None) + parser.add_argument("--fail-fast", action=argparse.BooleanOptionalAction, default=None) + parser.add_argument("--save-actions", action=argparse.BooleanOptionalAction, default=None) + return parser + + +def main() -> None: + logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s") + summary = run_evaluation(build_parser().parse_args()) + print(json.dumps(summary, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/lightx2v/models/runners/openpi/libero_protocol.py b/lightx2v/models/runners/openpi/libero_protocol.py new file mode 100644 index 000000000..af2cce503 --- /dev/null +++ b/lightx2v/models/runners/openpi/libero_protocol.py @@ -0,0 +1,721 @@ +"""Official pi0.5-LIBERO protocol shared by rollout and quantitative evaluation.""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib +import json +import logging +import math +import os +import sys +import time +import traceback +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np + +from .artifacts import EpisodeKey, atomic_write_json, atomic_write_text + +LOGGER = logging.getLogger(__name__) + +LIBERO_BENCHMARKS = ("libero_spatial", "libero_object", "libero_goal", "libero_10", "libero_90") +OFFICIAL_BENCHMARKS = ("libero_spatial", "libero_object", "libero_goal", "libero_10") +MAX_STEPS_BY_BENCHMARK = { + "libero_spatial": 220, + "libero_object": 280, + "libero_goal": 300, + "libero_10": 520, + "libero_90": 400, +} +OFFICIAL_RESULTS = { + "libero_spatial": {"successes": 494, "episodes": 500, "success_rate": 98.8}, + "libero_object": {"successes": 491, "episodes": 500, "success_rate": 98.2}, + "libero_goal": {"successes": 490, "episodes": 500, "success_rate": 98.0}, + "libero_10": {"successes": 462, "episodes": 500, "success_rate": 92.4}, + "average": 96.85, +} +VIDEO_POLICIES = ("all", "failures", "none") +LIBERO_DUMMY_ACTION = np.asarray([0.0] * 6 + [-1.0], dtype=np.float32) +POLICY_IMAGE_SIZE = 224 +TASK_INPUTS_MANIFEST_FILENAME = "task_inputs_manifest.json" + + +@dataclass(frozen=True) +class EvaluationConfig: + benchmarks: tuple[str, ...] + task_ids: dict[str, tuple[int, ...] | None] + num_trials_per_task: int + env_seed: int + policy_seed: int + actions_per_plan: int + num_steps_wait: int + render_size: int + video_fps: int + video_policy: str + save_actions: bool + fail_fast: bool + resume: bool + max_steps: dict[str, int] + libero_root: Path + libero_config_dir: Path + + +@dataclass(frozen=True) +class LiberoRuntime: + benchmark_module: Any + get_libero_path: Any + env_type: Any + + +@dataclass(frozen=True) +class TaskSpec: + benchmark: str + task_id: int + suite: Any + task: Any + bddl_path: Path + init_states_path: Path + + +def _parse_csv(value: str) -> tuple[str, ...]: + return tuple(item.strip() for item in str(value).split(",") if item.strip()) + + +def parse_integer_selection(value: Any, label: str) -> tuple[int, ...] | None: + if value is None or value == "all": + return None + if isinstance(value, (list, tuple)): + values = [int(item) for item in value] + else: + values = [] + for segment in _parse_csv(str(value)): + if "-" in segment: + start_text, end_text = segment.split("-", 1) + start, end = int(start_text), int(end_text) + if start > end: + raise ValueError(f"{label} range starts after it ends: {segment!r}") + values.extend(range(start, end + 1)) + else: + values.append(int(segment)) + if not values: + raise ValueError(f"{label} cannot be empty") + if any(item < 0 for item in values): + raise ValueError(f"{label} must contain only non-negative integers") + if len(set(values)) != len(values): + raise ValueError(f"{label} contains duplicate IDs: {values}") + return tuple(values) + + +def load_evaluation_config(args: argparse.Namespace) -> EvaluationConfig: + eval_config_path = Path(args.eval_config).expanduser().resolve() + with eval_config_path.open("r", encoding="utf-8") as handle: + values = json.load(handle) + if not isinstance(values, dict): + raise ValueError(f"OpenPI evaluation config must be a JSON object: {eval_config_path}") + + benchmark_value = args.benchmarks if args.benchmarks is not None else values.get("benchmarks", OFFICIAL_BENCHMARKS) + benchmarks = _parse_csv(benchmark_value) if isinstance(benchmark_value, str) else tuple(str(item) for item in benchmark_value) + if not benchmarks: + raise ValueError("At least one LIBERO benchmark is required") + unknown = [name for name in benchmarks if name not in LIBERO_BENCHMARKS] + if unknown: + raise ValueError(f"Unknown LIBERO benchmarks: {unknown}; expected one of {LIBERO_BENCHMARKS}") + if len(set(benchmarks)) != len(benchmarks): + raise ValueError(f"Duplicate LIBERO benchmarks are not allowed: {benchmarks}") + + raw_task_ids: Any = args.task_ids if args.task_ids is not None else values.get("task_ids", "all") + if isinstance(raw_task_ids, dict): + task_ids = {name: parse_integer_selection(raw_task_ids.get(name, "all"), f"task_ids.{name}") for name in benchmarks} + else: + selection = parse_integer_selection(raw_task_ids, "task_ids") + task_ids = {name: selection for name in benchmarks} + + raw_max_steps = values.get("max_steps", MAX_STEPS_BY_BENCHMARK) + if args.max_steps is not None: + max_steps = {name: int(args.max_steps) for name in benchmarks} + elif isinstance(raw_max_steps, dict): + max_steps = {name: int(raw_max_steps[name]) for name in benchmarks} + else: + max_steps = {name: int(raw_max_steps) for name in benchmarks} + + def overridden(argument_name: str, config_name: str, default: Any) -> Any: + argument = getattr(args, argument_name) + return argument if argument is not None else values.get(config_name, default) + + config = EvaluationConfig( + benchmarks=benchmarks, + task_ids=task_ids, + num_trials_per_task=int(overridden("num_trials_per_task", "num_trials_per_task", 50)), + env_seed=int(values.get("env_seed", 7)), + policy_seed=int(values.get("policy_seed", 0)), + actions_per_plan=int(values.get("actions_per_plan", 5)), + num_steps_wait=int(values.get("num_steps_wait", 10)), + render_size=int(values.get("render_size", 256)), + video_fps=int(values.get("video_fps", 10)), + video_policy=str(overridden("video_policy", "video_policy", "none")), + save_actions=bool(overridden("save_actions", "save_actions", False)), + fail_fast=bool(overridden("fail_fast", "fail_fast", False)), + resume=bool(overridden("resume", "resume", True)), + max_steps=max_steps, + libero_root=Path(args.libero_root).expanduser().resolve(), + libero_config_dir=Path(args.libero_config_dir).expanduser().resolve(), + ) + if config.num_trials_per_task < 1: + raise ValueError("num_trials_per_task must be positive") + if config.actions_per_plan < 1: + raise ValueError("actions_per_plan must be positive") + if config.num_steps_wait < 0: + raise ValueError("num_steps_wait must be non-negative") + if config.env_seed < 0 or config.policy_seed < 0: + raise ValueError("env_seed and policy_seed must be non-negative") + if config.render_size < 1 or config.video_fps < 1: + raise ValueError("render_size and video_fps must be positive") + if config.video_policy not in VIDEO_POLICIES: + raise ValueError(f"video_policy must be one of {VIDEO_POLICIES}, got {config.video_policy!r}") + if any(limit < 1 for limit in config.max_steps.values()): + raise ValueError("Every max_steps value must be positive") + return config + + +def is_official_protocol(config: EvaluationConfig) -> bool: + return ( + all(name in OFFICIAL_BENCHMARKS for name in config.benchmarks) + and all(config.task_ids[name] is None for name in config.benchmarks) + and config.num_trials_per_task == 50 + and config.env_seed == 7 + and config.policy_seed == 0 + and config.actions_per_plan == 5 + and config.num_steps_wait == 10 + and config.render_size == 256 + and all(config.max_steps[name] == MAX_STEPS_BY_BENCHMARK[name] for name in config.benchmarks) + ) + + +def _file_descriptor(path: Path) -> dict[str, Any]: + path = path.expanduser().resolve() + if not path.is_file(): + raise FileNotFoundError(path) + stat = path.stat() + return {"path": str(path), "size_bytes": stat.st_size, "mtime_ns": stat.st_mtime_ns} + + +def _canonical_sha256(payload: Mapping[str, Any]) -> str: + encoded = json.dumps(dict(payload), sort_keys=True, separators=(",", ":")).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def _sha256_file(path: Path) -> tuple[str, int]: + """Hash one immutable view of a file and reject concurrent replacement.""" + path = path.expanduser().resolve() + before = path.stat() + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + after = path.stat() + before_identity = (before.st_dev, before.st_ino, before.st_size, before.st_mtime_ns) + after_identity = (after.st_dev, after.st_ino, after.st_size, after.st_mtime_ns) + if before_identity != after_identity: + raise RuntimeError(f"LIBERO task input changed while it was being hashed: {path}") + return digest.hexdigest(), after.st_size + + +def _manifest_path(path: Path, libero_root: Path) -> str: + path = path.expanduser().resolve() + try: + return str(path.relative_to(libero_root.expanduser().resolve())) + except ValueError: + return str(path) + + +def _file_manifest_entry(spec: TaskSpec, kind: str, path: Path, libero_root: Path) -> dict[str, Any]: + content_sha256, size_bytes = _sha256_file(path) + return { + "benchmark": spec.benchmark, + "task_id": spec.task_id, + "kind": kind, + "source": "file", + "path": _manifest_path(path, libero_root), + "size_bytes": size_bytes, + "content_sha256": content_sha256, + } + + +def _semantic_init_states_manifest_entry(spec: TaskSpec, libero_root: Path) -> dict[str, Any]: + states, loader = load_task_init_states(spec) + metadata = { + "dtype": states.dtype.str, + "shape": list(states.shape), + } + digest = hashlib.sha256() + digest.update(json.dumps(metadata, sort_keys=True, separators=(",", ":")).encode("utf-8")) + digest.update(states.tobytes(order="C")) + return { + "benchmark": spec.benchmark, + "task_id": spec.task_id, + "kind": "init_states", + "source": loader, + "path": _manifest_path(spec.init_states_path, libero_root), + "size_bytes": states.nbytes, + "dtype": metadata["dtype"], + "shape": metadata["shape"], + "content_sha256": digest.hexdigest(), + } + + +def build_task_inputs_manifest( + task_specs: Mapping[str, list[TaskSpec]], + config: EvaluationConfig, +) -> dict[str, Any]: + """Fingerprint the selected BDDL and initial-state content.""" + entries: list[dict[str, Any]] = [] + for benchmark in config.benchmarks: + for spec in task_specs[benchmark]: + entries.append(_file_manifest_entry(spec, "bddl", spec.bddl_path, config.libero_root)) + if spec.init_states_path.is_file(): + entries.append(_file_manifest_entry(spec, "init_states", spec.init_states_path, config.libero_root)) + else: + entries.append(_semantic_init_states_manifest_entry(spec, config.libero_root)) + + entries.sort(key=lambda entry: (str(entry["benchmark"]), int(entry["task_id"]), str(entry["kind"]))) + manifest: dict[str, Any] = { + "schema_version": 1, + "hash_algorithm": "sha256", + "scope": "selected_libero_task_inputs", + "task_count": sum(len(specs) for specs in task_specs.values()), + "input_count": len(entries), + "entries": entries, + } + manifest["manifest_sha256"] = _canonical_sha256(manifest) + return manifest + + +def _validate_legacy_record_input_paths( + records: Mapping[EpisodeKey, Mapping[str, Any]], + task_specs: Mapping[str, list[TaskSpec]], +) -> None: + specs_by_task = {(spec.benchmark, spec.task_id): spec for specs in task_specs.values() for spec in specs} + for key, record in records.items(): + if record.get("schema_version") != 1: + raise RuntimeError(f"Cannot adopt a missing task-input manifest for non-legacy episode {key}: schema_version={record.get('schema_version')!r}") + spec = specs_by_task.get(key[:2]) + if spec is None: + raise RuntimeError(f"Cannot match legacy episode {key} to a selected LIBERO task") + for field, expected in (("bddl_file", spec.bddl_path), ("init_states_file", spec.init_states_path)): + recorded = record.get(field) + if not recorded or Path(str(recorded)).expanduser().resolve() != expected: + raise RuntimeError(f"Cannot adopt task-input manifest: legacy episode {key} has {field}={recorded!r}, expected {expected}") + expected_loader = "direct_file" if spec.init_states_path.is_file() else "benchmark_api" + if record.get("init_states_loader") != expected_loader: + raise RuntimeError(f"Cannot adopt task-input manifest: legacy episode {key} has init_states_loader={record.get('init_states_loader')!r}, expected {expected_loader!r}") + + +def ensure_task_inputs_manifest( + output_dir: Path, + expected_manifest: Mapping[str, Any], + records: Mapping[EpisodeKey, Mapping[str, Any]], + task_specs: Mapping[str, list[TaskSpec]], +) -> Path: + """Create, verify, or safely adopt the task-input sidecar. + + The sidecar deliberately does not participate in ``protocol_id`` so that + schema-1 evaluations started before content fingerprints existed remain + resumable. Once present, the complete canonical manifest is compared on + every startup. + """ + manifest_path = Path(output_dir) / TASK_INPUTS_MANIFEST_FILENAME + expected = dict(expected_manifest) + if manifest_path.exists(): + if not manifest_path.is_file(): + raise RuntimeError(f"LIBERO task-input manifest is not a file: {manifest_path}") + try: + existing = json.loads(manifest_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise RuntimeError(f"Cannot read LIBERO task-input manifest: {manifest_path}") from exc + if not isinstance(existing, dict): + raise RuntimeError(f"LIBERO task-input manifest must contain a JSON object: {manifest_path}") + if existing != expected: + raise RuntimeError( + "LIBERO BDDL/init-state content differs from the committed evaluation manifest: " + f"{manifest_path} (existing={existing.get('manifest_sha256')!r}, " + f"current={expected.get('manifest_sha256')!r})" + ) + return manifest_path + + if records: + _validate_legacy_record_input_paths(records, task_specs) + LOGGER.warning( + "Adopting %s for %d legacy schema-1 episodes after validating their recorded BDDL/init-state paths. " + "The historical episodes predate content fingerprints and therefore cannot be verified retrospectively.", + manifest_path, + len(records), + ) + atomic_write_json(manifest_path, expected) + return manifest_path + + +def resolved_protocol( + config: EvaluationConfig, + *, + model_path: Path, + config_json: Path, +) -> tuple[dict[str, Any], str]: + """Build a cheap protocol identity without hashing model or source trees.""" + model_path = model_path.expanduser().resolve() + config_json = config_json.expanduser().resolve() + artifacts = { + "checkpoint": _file_descriptor(model_path / "model.safetensors"), + "model_config": _file_descriptor(config_json), + "norm_stats": _file_descriptor(model_path / "assets/physical-intelligence/libero/norm_stats.json"), + "tokenizer": _file_descriptor(model_path / "assets/paligemma_tokenizer.model"), + } + protocol_fields = { + "schema_version": 1, + "model_artifacts": artifacts, + "benchmarks": list(config.benchmarks), + "task_ids": {name: "all" if ids is None else list(ids) for name, ids in config.task_ids.items()}, + "num_trials_per_task": config.num_trials_per_task, + "env_seed": config.env_seed, + "policy_seed": config.policy_seed, + "policy_rng_scope": "per_suite_continuous", + "actions_per_plan": config.actions_per_plan, + "num_steps_wait": config.num_steps_wait, + "render_size": config.render_size, + "max_steps": config.max_steps, + "libero_root": str(config.libero_root), + } + encoded = json.dumps(protocol_fields, sort_keys=True, separators=(",", ":")).encode("utf-8") + protocol_id = hashlib.sha256(encoded).hexdigest() + resolved = { + **protocol_fields, + "protocol_id": protocol_id, + "protocol_name": "official_pi05_libero" if is_official_protocol(config) else "custom_pi05_libero", + "official_protocol": is_official_protocol(config), + "libero_config_dir": str(config.libero_config_dir), + "video_fps": config.video_fps, + "video_policy": config.video_policy, + "save_actions": config.save_actions, + "fail_fast": config.fail_fast, + "resume": config.resume, + } + return resolved, protocol_id + + +def load_policy_config( + config_json: Path, + model_path: Path, + *, + seed: int, + actions_per_plan: int, +) -> dict[str, Any]: + config_json = config_json.expanduser().resolve() + model_path = model_path.expanduser().resolve() + with config_json.open("r", encoding="utf-8") as handle: + values = json.load(handle) + if not isinstance(values, dict): + raise ValueError(f"OpenPI model config must be a JSON object: {config_json}") + values.update( + { + "model_cls": "openpi", + "task": "i2va", + "model_path": str(model_path), + "config_json": str(config_json), + "seed": int(seed), + "actions_per_plan": int(actions_per_plan), + } + ) + return values + + +def _assert_module_source(module: Any, libero_root: Path) -> None: + origin_value = getattr(module, "__file__", None) + if origin_value: + origins = (Path(origin_value).resolve(),) + else: + search_locations = getattr(module, "__path__", None) + if search_locations is None: + raise RuntimeError(f"Cannot verify LIBERO module source for {module.__name__!r}") + origins = tuple(Path(location).resolve() for location in search_locations) + if not origins: + raise RuntimeError(f"LIBERO namespace {module.__name__!r} has no search locations") + outside = [origin for origin in origins if not origin.is_relative_to(libero_root)] + if outside: + raise RuntimeError( + f"LIBERO module {module.__name__!r} resolves outside requested root {libero_root}: {outside}. Start a fresh worker and remove the conflicting LIBERO package from PYTHONPATH." + ) + + +def _constrain_libero_namespace(package: Any, libero_root: Path) -> None: + """Restrict the top-level namespace before resolving ``libero.libero``.""" + if getattr(package, "__file__", None): + _assert_module_source(package, libero_root) + return + + requested = (libero_root / "libero").resolve() + locations = tuple(Path(location).resolve() for location in getattr(package, "__path__", ())) + if requested not in locations: + raise RuntimeError(f"Requested LIBERO namespace root {requested} is absent; discovered search locations: {locations}") + package.__path__ = [str(requested)] + spec = getattr(package, "__spec__", None) + if spec is not None: + spec.submodule_search_locations = package.__path__ + _assert_module_source(package, libero_root) + + +def configure_libero(libero_root: Path, config_dir: Path) -> LiberoRuntime: + """Configure and import exactly the LIBERO checkout selected by the caller.""" + libero_root = libero_root.expanduser().resolve() + benchmark_root = libero_root / "libero" / "libero" + required = (benchmark_root / "bddl_files", benchmark_root / "init_files", benchmark_root / "assets") + missing = [str(path) for path in required if not path.is_dir()] + if missing: + raise FileNotFoundError(f"LIBERO checkout is incomplete under {libero_root}: {missing}") + + config_dir = config_dir.expanduser().resolve() + config_text = "\n".join( + ( + f"benchmark_root: {benchmark_root}", + f"bddl_files: {benchmark_root / 'bddl_files'}", + f"init_states: {benchmark_root / 'init_files'}", + f"datasets: {libero_root / 'libero' / 'datasets'}", + f"assets: {benchmark_root / 'assets'}", + "", + ) + ) + atomic_write_text(config_dir / "config.yaml", config_text) + os.environ["LIBERO_CONFIG_PATH"] = str(config_dir) + + root_text = str(libero_root) + if root_text in sys.path: + sys.path.remove(root_text) + sys.path.insert(0, root_text) + importlib.invalidate_caches() + + package = importlib.import_module("libero") + _constrain_libero_namespace(package, libero_root) + nested_package = importlib.import_module("libero.libero") + benchmark_module = importlib.import_module("libero.libero.benchmark") + envs_module = importlib.import_module("libero.libero.envs") + for module in (package, nested_package, benchmark_module, envs_module): + _assert_module_source(module, libero_root) + + return LiberoRuntime( + benchmark_module=benchmark_module, + get_libero_path=nested_package.get_libero_path, + env_type=envs_module.OffScreenRenderEnv, + ) + + +def build_task_specs(runtime: LiberoRuntime, config: EvaluationConfig) -> dict[str, list[TaskSpec]]: + factories = runtime.benchmark_module.get_benchmark_dict() + specs: dict[str, list[TaskSpec]] = {} + for benchmark_name in config.benchmarks: + if benchmark_name not in factories: + raise KeyError(f"LIBERO checkout does not provide benchmark {benchmark_name!r}") + suite = factories[benchmark_name]() + task_count = int(suite.n_tasks) + selected = config.task_ids[benchmark_name] + task_ids = tuple(range(task_count)) if selected is None else selected + invalid = [task_id for task_id in task_ids if not 0 <= task_id < task_count] + if invalid: + raise ValueError(f"Task IDs {invalid} are outside [0, {task_count}) for {benchmark_name}") + + suite_specs = [] + for task_id in task_ids: + task = suite.get_task(task_id) + bddl_path = Path(runtime.get_libero_path("bddl_files")) / task.problem_folder / task.bddl_file + init_states_path = Path(runtime.get_libero_path("init_states")) / task.problem_folder / task.init_states_file + if not bddl_path.is_file(): + raise FileNotFoundError(f"LIBERO BDDL file is missing: {bddl_path}") + suite_specs.append( + TaskSpec( + benchmark=benchmark_name, + task_id=task_id, + suite=suite, + task=task, + bddl_path=bddl_path.resolve(), + init_states_path=init_states_path.resolve(), + ) + ) + specs[benchmark_name] = suite_specs + return specs + + +def load_task_init_states(spec: TaskSpec) -> tuple[np.ndarray, str]: + if spec.init_states_path.is_file(): + import torch + + states = torch.load(spec.init_states_path, map_location="cpu", weights_only=False) + loader = "direct_file" + else: + states = spec.suite.get_task_init_states(spec.task_id) + loader = "benchmark_api" + array = np.asarray(states) + if array.ndim != 2 or len(array) == 0: + raise ValueError(f"LIBERO task {spec.benchmark}/{spec.task_id} returned invalid init states with shape {array.shape}") + if not np.isfinite(array).all(): + raise ValueError(f"LIBERO task {spec.benchmark}/{spec.task_id} returned non-finite init states") + return np.ascontiguousarray(array), loader + + +def create_environment(runtime: LiberoRuntime, spec: TaskSpec, render_size: int, seed: int) -> Any: + env = runtime.env_type( + bddl_file_name=str(spec.bddl_path), + camera_heights=render_size, + camera_widths=render_size, + ) + env.seed(seed) + return env + + +def _quat_to_axis_angle(quaternion: Any) -> np.ndarray: + # Preserve the simulator dtype exactly as the official helper does. + quat = np.asarray(quaternion).copy() + if quat.shape != (4,): + raise ValueError(f"Expected LIBERO quaternion shape (4,), got {quat.shape}") + quat[3] = np.clip(quat[3], -1.0, 1.0) + denominator = math.sqrt(max(0.0, 1.0 - float(quat[3]) ** 2)) + if math.isclose(denominator, 0.0): + return np.zeros(3) + return quat[:3] * (2.0 * math.acos(float(quat[3])) / denominator) + + +def rotate_rgb(observation: dict[str, Any], key: str) -> np.ndarray: + image = np.asarray(observation[key]) + if image.ndim != 3 or image.shape[-1] != 3: + raise ValueError(f"LIBERO camera {key!r} must be HxWx3, got {image.shape}") + return np.ascontiguousarray(image[::-1, ::-1], dtype=np.uint8) + + +def policy_rgb(observation: dict[str, Any], key: str) -> np.ndarray: + """Apply the official client-side rotate, PIL resize, and uint8 conversion.""" + image = rotate_rgb(observation, key) + if image.shape[:2] == (POLICY_IMAGE_SIZE, POLICY_IMAGE_SIZE): + return image + + from PIL import Image + + height, width = image.shape[:2] + ratio = max(width / POLICY_IMAGE_SIZE, height / POLICY_IMAGE_SIZE) + resized_height = int(height / ratio) + resized_width = int(width / ratio) + resized = Image.fromarray(image, mode="RGB").resize( + (resized_width, resized_height), + resample=Image.BILINEAR, + ) + canvas = Image.new("RGB", (POLICY_IMAGE_SIZE, POLICY_IMAGE_SIZE), 0) + canvas.paste( + resized, + ( + max(0, int((POLICY_IMAGE_SIZE - resized_width) / 2)), + max(0, int((POLICY_IMAGE_SIZE - resized_height) / 2)), + ), + ) + return np.asarray(canvas, dtype=np.uint8).copy() + + +def state_from_observation(observation: dict[str, Any]) -> np.ndarray: + return np.concatenate( + ( + np.asarray(observation["robot0_eef_pos"]), + _quat_to_axis_angle(observation["robot0_eef_quat"]), + np.asarray(observation["robot0_gripper_qpos"]), + ) + ) + + +def run_episode( + *, + policy: Any, + env: Any, + initial_state: Any, + task_description: str, + max_steps: int, + num_steps_wait: int, + collect_frames: bool, +) -> tuple[dict[str, Any], list[np.ndarray], np.ndarray]: + """Run one episode with the exact official warmup/policy done semantics.""" + frames: list[np.ndarray] = [] + actions: list[np.ndarray] = [] + success = False + warmup_steps = 0 + warmup_done_observed = 0 + action_chunk_calls = 0 + error_type: str | None = None + error_message: str | None = None + error_traceback: str | None = None + started = time.perf_counter() + + policy.clear_action_queue() + # Match the official evaluator: reset/init failures are infrastructure errors. + env.reset() + observation = env.set_init_state(initial_state) + try: + # Official examples/libero/main.py ignores done during settling. Both the + # one-off rollout and quantitative path use this same branch. + for _ in range(num_steps_wait): + observation, _reward, done, _info = env.step(LIBERO_DUMMY_ACTION.tolist()) + warmup_steps += 1 + warmup_done_observed += int(bool(done)) + + for _step in range(max_steps): + agentview = policy_rgb(observation, "agentview_image") + if collect_frames: + frames.append(agentview) + images = { + "agentview": agentview, + "wrist": policy_rgb(observation, "robot0_eye_in_hand_image"), + } + if policy.pending_action_count == 0: + action_chunk_calls += 1 + action = policy.next_action( + images=images, + state=state_from_observation(observation), + task_description=task_description, + ) + action = np.asarray(action).reshape(-1) + actions.append(action.copy()) + observation, _reward, done, _info = env.step(action.tolist()) + if bool(done): + success = True + break + except Exception as exc: # The official loop counts ordinary rollout errors as failures. + message = str(exc) + if isinstance(exc, MemoryError) or "out of memory" in message.lower(): + raise + error_type = type(exc).__name__ + error_message = message + error_traceback = traceback.format_exc(limit=20) + LOGGER.exception("LIBERO episode failed with an exception") + finally: + pending_actions_discarded = policy.pending_action_count + policy.clear_action_queue() + + action_dim = int(policy.output_action_dim) + action_array = np.stack(actions).reshape(-1, action_dim) if actions else np.empty((0, action_dim), dtype=np.float64) + termination_reason = "exception" if error_type is not None else "success" if success else "step_limit" + episode = { + "success": success, + "termination_reason": termination_reason, + "warmup_steps_requested": num_steps_wait, + "warmup_steps_executed": warmup_steps, + "warmup_done_observed": warmup_done_observed, + "policy_steps": int(action_array.shape[0]), + "total_env_steps": warmup_steps + int(action_array.shape[0]), + "action_chunk_calls": action_chunk_calls, + "pending_actions_discarded": pending_actions_discarded, + "elapsed_seconds": time.perf_counter() - started, + "exception_type": error_type, + "exception_message": error_message, + "exception_traceback": error_traceback, + } + return episode, frames, action_array + + +def expected_episode_keys(task_specs: dict[str, list[TaskSpec]], config: EvaluationConfig) -> tuple[EpisodeKey, ...]: + return tuple((benchmark, spec.task_id, init_state_id) for benchmark in config.benchmarks for spec in task_specs[benchmark] for init_state_id in range(config.num_trials_per_task)) diff --git a/lightx2v/models/runners/openpi/libero_rollout.py b/lightx2v/models/runners/openpi/libero_rollout.py new file mode 100644 index 000000000..5ac5f7359 --- /dev/null +++ b/lightx2v/models/runners/openpi/libero_rollout.py @@ -0,0 +1,183 @@ +"""Run one local pi0.5 policy episode in the selected LIBERO checkout.""" + +from __future__ import annotations + +import argparse +import json +import logging +import time +from pathlib import Path + +import numpy as np + +from .artifacts import atomic_save_numpy, atomic_write_json, atomic_write_video +from .libero_protocol import ( + LIBERO_BENCHMARKS, + MAX_STEPS_BY_BENCHMARK, + TaskSpec, + configure_libero, + create_environment, + load_policy_config, + load_task_init_states, + run_episode, +) +from .openpi_runner import OpenPIPolicy + +LOGGER = logging.getLogger(__name__) + + +def _output_path(path: Path, suffix: str, label: str) -> Path: + path = path.expanduser().resolve() + if path.suffix.lower() != suffix: + raise ValueError(f"{label} must end in {suffix}: {path}") + return path + + +def _task_spec(runtime, benchmark_name: str, task_id: int) -> TaskSpec: + factories = runtime.benchmark_module.get_benchmark_dict() + if benchmark_name not in factories: + raise KeyError(f"LIBERO checkout does not provide benchmark {benchmark_name!r}") + suite = factories[benchmark_name]() + task_count = int(suite.n_tasks) + if not 0 <= task_id < task_count: + raise ValueError(f"task_id must be in [0, {task_count}), got {task_id}") + task = suite.get_task(task_id) + return TaskSpec( + benchmark=benchmark_name, + task_id=task_id, + suite=suite, + task=task, + bddl_path=(Path(runtime.get_libero_path("bddl_files")) / task.problem_folder / task.bddl_file).resolve(), + init_states_path=(Path(runtime.get_libero_path("init_states")) / task.problem_folder / task.init_states_file).resolve(), + ) + + +def run_rollout(args: argparse.Namespace) -> dict: + if args.actions_per_plan < 1: + raise ValueError("--actions-per-plan must be positive") + if args.num_steps_wait < 0: + raise ValueError("--num-steps-wait must be non-negative") + if args.render_size < 1 or args.fps < 1: + raise ValueError("--render-size and --fps must be positive") + if args.env_seed < 0 or args.policy_seed < 0: + raise ValueError("--env-seed and --policy-seed must be non-negative") + max_steps = int(args.max_steps if args.max_steps is not None else MAX_STEPS_BY_BENCHMARK[args.benchmark]) + if max_steps < 1: + raise ValueError("--max-steps must be positive") + + video_path = _output_path(args.save_video_path, ".mp4", "save_video_path") + action_path = _output_path(args.save_action_path, ".npy", "save_action_path") + metrics_path = _output_path(args.save_metrics_path, ".json", "save_metrics_path") + np.random.seed(args.env_seed) + runtime = configure_libero(args.libero_root, args.libero_config_dir) + spec = _task_spec(runtime, args.benchmark, args.task_id) + if not spec.bddl_path.is_file(): + raise FileNotFoundError(f"LIBERO BDDL file is missing: {spec.bddl_path}") + initial_states, init_states_loader = load_task_init_states(spec) + if not 0 <= args.init_state_id < len(initial_states): + raise ValueError(f"init_state_id must be in [0, {len(initial_states)}), got {args.init_state_id}") + + task_description = args.task_description.strip() or str(spec.task.language) + env = create_environment(runtime, spec, args.render_size, args.env_seed) + policy = None + started = time.perf_counter() + try: + LOGGER.info("Loading local PyTorch OpenPI policy") + policy = OpenPIPolicy( + load_policy_config( + args.config_json, + args.model_path, + seed=args.policy_seed, + actions_per_plan=args.actions_per_plan, + ) + ) + policy.reset() + episode, frames, actions = run_episode( + policy=policy, + env=env, + initial_state=initial_states[args.init_state_id], + task_description=task_description, + max_steps=max_steps, + num_steps_wait=args.num_steps_wait, + collect_frames=True, + ) + rng_state_after = policy.export_rng_state() + finally: + env.close() + if policy is not None: + policy.close() + + artifact_errors: list[str] = [] + try: + atomic_save_numpy(action_path, actions) + except Exception as exc: + artifact_errors.append(f"actions: {type(exc).__name__}: {exc}") + LOGGER.exception("Failed to save single-rollout actions") + try: + atomic_write_video(video_path, frames, args.fps) + except Exception as exc: + artifact_errors.append(f"video: {type(exc).__name__}: {exc}") + LOGGER.exception("Failed to save single-rollout video") + + metrics = { + "schema_version": 1, + "mode": "single_rollout", + "benchmark": args.benchmark, + "task_id": args.task_id, + "task_name": str(spec.task.name), + "task_description": task_description, + "bddl_file": str(spec.bddl_path), + "init_states_file": str(spec.init_states_path), + "init_states_loader": init_states_loader, + "init_state_id": args.init_state_id, + "init_state_count": len(initial_states), + "env_seed": args.env_seed, + "policy_seed": args.policy_seed, + "policy_rng_state_after": rng_state_after, + "max_policy_steps": max_steps, + "actions_per_plan": args.actions_per_plan, + "video_frames": len(frames), + "video_fps": args.fps, + "video_path": str(video_path) if video_path.is_file() else None, + "action_path": str(action_path) if action_path.is_file() else None, + "artifact_errors": artifact_errors, + "total_elapsed_seconds": time.perf_counter() - started, + **episode, + } + # metrics.json is the final, atomic completion marker for the rollout. + atomic_write_json(metrics_path, metrics) + LOGGER.info("Single LIBERO rollout complete: success=%s metrics=%s", metrics["success"], metrics_path) + return metrics + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Run one local pi0.5-LIBERO rollout") + parser.add_argument("--model-path", type=Path, required=True) + parser.add_argument("--config-json", type=Path, required=True) + parser.add_argument("--libero-root", type=Path, required=True) + parser.add_argument("--libero-config-dir", type=Path, required=True) + parser.add_argument("--benchmark", choices=LIBERO_BENCHMARKS, default="libero_spatial") + parser.add_argument("--task-id", type=int, default=0) + parser.add_argument("--init-state-id", type=int, default=0) + parser.add_argument("--env-seed", type=int, default=7) + parser.add_argument("--policy-seed", type=int, default=0) + parser.add_argument("--task-description", default="") + parser.add_argument("--actions-per-plan", type=int, default=5) + parser.add_argument("--num-steps-wait", type=int, default=10) + parser.add_argument("--max-steps", type=int) + parser.add_argument("--render-size", type=int, default=256) + parser.add_argument("--fps", type=int, default=10) + parser.add_argument("--save-video-path", type=Path, required=True) + parser.add_argument("--save-action-path", type=Path, required=True) + parser.add_argument("--save-metrics-path", type=Path, required=True) + return parser + + +def main() -> None: + logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s") + metrics = run_rollout(build_parser().parse_args()) + print(json.dumps(metrics, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/lightx2v/models/runners/openpi/openpi_runner.py b/lightx2v/models/runners/openpi/openpi_runner.py new file mode 100644 index 000000000..2d15d6edd --- /dev/null +++ b/lightx2v/models/runners/openpi/openpi_runner.py @@ -0,0 +1,308 @@ +from __future__ import annotations + +import base64 +import os +import subprocess +import sys +from collections import deque +from pathlib import Path +from typing import Any + +import numpy as np + +from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.utils.registry_factory import RUNNER_REGISTER + +PROJECT_ROOT = Path(__file__).resolve().parents[4] + + +class OpenPIPolicy: + def __init__(self, config: Any): + self.action_horizon = int(config["action_horizon"]) + self.output_action_dim = int(config["output_action_dim"]) + self.actions_per_plan = int(config.get("actions_per_plan", 5)) + if not 1 <= self.actions_per_plan <= self.action_horizon: + raise ValueError(f"OpenPI actions_per_plan must be in [1, {self.action_horizon}], got {self.actions_per_plan}.") + + # Import only after the worker activates the patched Transformers path. + from lightx2v.models.networks.openpi import OpenPIModel + + self.model = OpenPIModel.from_config(config) + self.pending_actions: deque[np.ndarray] = deque() + + @property + def pending_action_count(self) -> int: + return len(self.pending_actions) + + def predict_action_chunk( + self, + images: dict[str, np.ndarray], + state: np.ndarray, + task_description: str, + seed: int | None = None, + ) -> np.ndarray: + actions = self.model.predict_action_chunk( + images=images, + state=state, + task_description=task_description, + seed=seed, + ) + # Quantile unnormalization returns float64; casting here changes simulator inputs. + actions = np.asarray(actions) + expected_shape = (self.action_horizon, self.output_action_dim) + if actions.shape != expected_shape: + raise ValueError(f"OpenPI expected action chunk shape {expected_shape}, got {actions.shape}.") + if not np.isfinite(actions).all(): + raise ValueError("OpenPI produced non-finite actions.") + return np.ascontiguousarray(actions) + + def next_action(self, images: dict[str, np.ndarray], state: np.ndarray, task_description: str) -> np.ndarray: + if not self.pending_actions: + chunk = self.predict_action_chunk(images, state, task_description, seed=None) + self.pending_actions.extend(action.copy() for action in chunk[: self.actions_per_plan]) + return self.pending_actions.popleft() + + def clear_action_queue(self) -> None: + self.pending_actions.clear() + + def reset_rng(self) -> None: + self.model.reset() + + def reset(self) -> None: + self.clear_action_queue() + self.reset_rng() + + def export_rng_state(self) -> str: + encoded = self.model.get_rng_state().cpu().numpy().tobytes() + return base64.b64encode(encoded).decode("ascii") + + def import_rng_state(self, encoded: str) -> None: + if not encoded: + raise ValueError("OpenPI RNG state is empty") + raw = base64.b64decode(encoded.encode("ascii"), validate=True) + state_array = np.frombuffer(raw, dtype=np.uint8).copy() + import torch + + state_tensor = torch.from_numpy(state_array) + self.model.set_rng_state(state_tensor) + + def close(self) -> None: + self.clear_action_queue() + + +@RUNNER_REGISTER("openpi") +class OpenPIRunner(BaseRunner): + def init_modules(self) -> None: + if self.config["task"] != "i2va": + raise ValueError(f"OpenPI currently supports only task='i2va', got {self.config['task']!r}.") + self.run_mode = str(self._value("run_mode", environment_name="OPENPI_RUN_MODE", default="rollout")) + if self.run_mode not in {"rollout", "evaluate"}: + raise ValueError(f"Unsupported OpenPI run mode {self.run_mode!r}; expected 'rollout' or 'evaluate'.") + self.config.lock() + + def warmup(self) -> None: + # Model initialization belongs to the isolated worker. + pass + + def _value(self, config_name: str, *, environment_name: str | None = None, default: Any = None) -> Any: + if environment_name: + environment_value = os.environ.get(environment_name) + if environment_value is not None and environment_value.strip() != "": + return environment_value + value = self.config.get(config_name) + return default if value is None or value == "" else value + + @staticmethod + def _path(value: Any, label: str, *, suffix: str | None = None, must_exist: bool = False) -> Path: + if value is None or value == "": + raise ValueError(f"OpenPI requires {label}.") + path = Path(str(value)).expanduser().resolve() + if suffix is not None and path.suffix.lower() != suffix: + raise ValueError(f"OpenPI {label} must end in {suffix}: {path}") + if must_exist and not path.exists(): + raise FileNotFoundError(f"OpenPI {label} does not exist: {path}") + return path + + def _configured_path( + self, + config_name: str, + label: str, + environment_name: str | None = None, + suffix: str | None = None, + must_exist: bool = False, + ) -> Path: + return self._path( + self._value(config_name, environment_name=environment_name), + label, + suffix=suffix, + must_exist=must_exist, + ) + + def _model_path(self) -> Path: + return self._configured_path("model_path", label="model_path", must_exist=True) + + def _model_config_path(self) -> Path: + return self._configured_path("config_json", label="config_json", suffix=".json", must_exist=True) + + def _libero_root(self) -> Path: + return self._configured_path( + "libero_root", + label="LIBERO root", + environment_name="OPENPI_LIBERO_ROOT", + must_exist=True, + ) + + def _libero_config_dir(self) -> Path: + return self._configured_path( + "libero_config_dir", + label="LIBERO config directory", + environment_name="OPENPI_LIBERO_CONFIG_DIR", + ) + + def _worker_environment(self) -> dict[str, str]: + runtime_path = self._configured_path( + "transformers_runtime_path", + label="patched Transformers runtime", + environment_name="OPENPI_TRANSFORMERS_RUNTIME_PATH", + must_exist=True, + ) + if not (runtime_path / "transformers").is_dir(): + raise FileNotFoundError(f"OpenPI patched Transformers package is missing: {runtime_path / 'transformers'}") + + child_env = os.environ.copy() + child_env["USE_FLAX"] = "0" + visible_devices = [item.strip() for item in child_env.get("CUDA_VISIBLE_DEVICES", "").split(",") if item.strip()] + if len(visible_devices) == 1 and visible_devices[0] != "0": + # robosuite 1.4.1 validates physical IDs as EGL ordinals. + visible_devices.append("0") + child_env["CUDA_VISIBLE_DEVICES"] = ",".join(visible_devices) + if visible_devices: + child_env["MUJOCO_EGL_DEVICE_ID"] = "0" + child_env["PYTHONPATH"] = os.pathsep.join((str(runtime_path), str(PROJECT_ROOT))) + return child_env + + def _option(self, environment_name: str, config_name: str, default: Any) -> str: + return str(self._value(config_name, environment_name=environment_name, default=default)) + + def _worker_command(self, module: str) -> list[str]: + return [ + "-m", + module, + "--model-path", + str(self._model_path()), + "--config-json", + str(self._model_config_path()), + "--libero-root", + str(self._libero_root()), + "--libero-config-dir", + str(self._libero_config_dir()), + ] + + def _rollout_output_paths(self, input_info: Any) -> tuple[Path, Path, Path]: + video_path = self._path(input_info.save_result_path, "save_result_path", suffix=".mp4") + action_value = input_info.save_action_path or self._value("save_action_path", environment_name="OPENPI_SAVE_ACTION_PATH") + action_path = self._path(action_value or video_path.with_suffix(".actions.npy"), "save_action_path", suffix=".npy") + metrics_value = self._value("save_metrics_path", environment_name="OPENPI_SAVE_METRICS_PATH") + metrics_path = self._path(metrics_value or video_path.with_suffix(".metrics.json"), "save_metrics_path", suffix=".json") + return video_path, action_path, metrics_path + + def _rollout_command(self, input_info: Any) -> list[str]: + video_path, action_path, metrics_path = self._rollout_output_paths(input_info) + command = [ + *self._worker_command("lightx2v.models.runners.openpi.libero_rollout"), + "--benchmark", + self._option("LIBERO_BENCHMARK", "libero_benchmark", "libero_spatial"), + "--task-id", + self._option("LIBERO_TASK_ID", "libero_task_id", 0), + "--init-state-id", + self._option("LIBERO_INIT_STATE_ID", "libero_init_state_id", 0), + "--env-seed", + self._option("OPENPI_ENV_SEED", "env_seed", 7), + "--policy-seed", + str(self._value("policy_seed", environment_name="OPENPI_POLICY_SEED", default=input_info.seed)), + "--actions-per-plan", + self._option("OPENPI_ACTIONS_PER_PLAN", "actions_per_plan", 5), + "--num-steps-wait", + self._option("OPENPI_NUM_STEPS_WAIT", "num_steps_wait", 10), + "--render-size", + self._option("OPENPI_RENDER_SIZE", "render_size", 256), + "--fps", + self._option("OPENPI_VIDEO_FPS", "video_fps", 10), + "--save-video-path", + str(video_path), + "--save-action-path", + str(action_path), + "--save-metrics-path", + str(metrics_path), + ] + task_description = str(input_info.prompt or "").strip() + if task_description: + command.extend(("--task-description", task_description)) + max_steps = self._value("max_steps", environment_name="OPENPI_MAX_STEPS") + if max_steps is not None: + command.extend(("--max-steps", str(max_steps))) + return command + + @staticmethod + def _boolean_argument(value: Any, enabled: str, disabled: str) -> str | None: + if value is None or value == "": + return None + if isinstance(value, bool): + is_enabled = value + else: + normalized = str(value).strip() + if normalized not in {"0", "1"}: + raise ValueError(f"Boolean OpenPI option must be 0 or 1, got {value!r}") + is_enabled = normalized == "1" + return enabled if is_enabled else disabled + + def _evaluate_command(self, input_info: Any) -> list[str]: + eval_config = self._configured_path( + "eval_config", + label="evaluation config", + environment_name="OPENPI_EVAL_CONFIG", + suffix=".json", + must_exist=True, + ) + output_dir = self._path(input_info.save_result_path, "evaluation output directory") + if output_dir.exists() and not output_dir.is_dir(): + raise NotADirectoryError(f"OpenPI evaluation output is not a directory: {output_dir}") + command = [ + *self._worker_command("lightx2v.models.runners.openpi.libero_evaluate"), + "--eval-config", + str(eval_config), + "--output-dir", + str(output_dir), + ] + for environment_name, config_name, argument in ( + ("OPENPI_EVAL_BENCHMARKS", "eval_benchmarks", "--benchmarks"), + ("OPENPI_EVAL_TASK_IDS", "eval_task_ids", "--task-ids"), + ("OPENPI_EVAL_NUM_TRIALS_PER_TASK", "eval_num_trials_per_task", "--num-trials-per-task"), + ("OPENPI_EVAL_MAX_STEPS", "eval_max_steps", "--max-steps"), + ("OPENPI_EVAL_VIDEO_POLICY", "eval_video_policy", "--video-policy"), + ): + value = self._value(config_name, environment_name=environment_name) + if value is not None and value != "": + command.extend((argument, str(value))) + for environment_name, config_name, enabled, disabled in ( + ("OPENPI_EVAL_RESUME", "eval_resume", "--resume", "--no-resume"), + ("OPENPI_EVAL_FAIL_FAST", "eval_fail_fast", "--fail-fast", "--no-fail-fast"), + ("OPENPI_EVAL_SAVE_ACTIONS", "eval_save_actions", "--save-actions", "--no-save-actions"), + ): + argument = self._boolean_argument(self._value(config_name, environment_name=environment_name), enabled, disabled) + if argument is not None: + command.append(argument) + return command + + def _run_local_worker(self, input_info: Any) -> dict[str, Any]: + arguments = self._rollout_command(input_info) if self.run_mode == "rollout" else self._evaluate_command(input_info) + command = [sys.executable, *arguments] + subprocess.run(command, env=self._worker_environment(), check=True) + + if input_info.return_result_tensor and self.run_mode == "rollout": + _, action_path, _ = self._rollout_output_paths(input_info) + return {"actions": np.load(action_path)} + return {"actions": None} + + def run_pipeline(self, input_info: Any) -> dict[str, Any]: + return self._run_local_worker(input_info) diff --git a/lightx2v/pipeline.py b/lightx2v/pipeline.py index ddb0eb46e..0e970c9c9 100755 --- a/lightx2v/pipeline.py +++ b/lightx2v/pipeline.py @@ -21,6 +21,7 @@ from lightx2v.models.runners.ltx2.ltx25_runner import LTX25Runner # noqa: F401 from lightx2v.models.runners.minimax_h3.minimax_h3_runner import MiniMaxH3Runner # noqa: F401 from lightx2v.models.runners.neopp.neopp_runner import NeoppRunner # noqa: F401 +from lightx2v.models.runners.openpi.openpi_runner import OpenPIRunner # noqa: F401 from lightx2v.models.runners.qwen_image.qwen_image_runner import QwenImageRunner # noqa: F401 from lightx2v.models.runners.seedvr.seedvr_runner import SeedVRRunner # noqa: F401 from lightx2v.models.runners.swiftvr.swiftvr_runner import SwiftVRRunner # noqa: F401 diff --git a/pyproject.toml b/pyproject.toml index 18794576e..4369690d2 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -99,6 +99,7 @@ exclude = [ ".cluade", ".cursor", "lightx2v_kernel", + "lightx2v/models/networks/openpi/transformers_replace", ] target-version = "py311" line-length = 200 diff --git a/scripts/openpi/1_convert_pi05_libero_to_pytorch.sh b/scripts/openpi/1_convert_pi05_libero_to_pytorch.sh new file mode 100755 index 000000000..f97e9f253 --- /dev/null +++ b/scripts/openpi/1_convert_pi05_libero_to_pytorch.sh @@ -0,0 +1,8 @@ +#!/usr/bin/env bash +set -euo pipefail + +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +openpi_root="${OPENPI_PATH:-$(cd -- "${script_dir}/../../.." && pwd)/openpi}" +python_bin="${OPENPI_CONVERT_PYTHON:-${openpi_root}/.venv/bin/python}" + +exec "${python_bin}" "${script_dir}/convert_jax_checkpoint.py" "$@" diff --git a/scripts/openpi/2_setup_pytorch_runtime.sh b/scripts/openpi/2_setup_pytorch_runtime.sh new file mode 100755 index 000000000..cc7d8da8b --- /dev/null +++ b/scripts/openpi/2_setup_pytorch_runtime.sh @@ -0,0 +1,14 @@ +#!/usr/bin/env bash +set -euo pipefail + +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +command="prepare" + +if [[ "${1:-}" == "check" || "${1:-}" == "--check" ]]; then + command="check" + shift +elif [[ "${1:-}" == "setup" || "${1:-}" == "prepare" ]]; then + shift +fi + +exec python "${script_dir}/runtime.py" "${command}" "$@" diff --git a/scripts/openpi/README.md b/scripts/openpi/README.md new file mode 100644 index 000000000..b9bee25b8 --- /dev/null +++ b/scripts/openpi/README.md @@ -0,0 +1,213 @@ +# OpenPI π0.5-LIBERO + +该目录提供 π0.5-LIBERO 的权重转换、运行环境准备、本地 rollout 和定量评测。 +所有推理都从 LightX2V 公共入口启动： + +```text +shell -> python -m lightx2v.infer -> OpenPIRunner + -> 本地 PyTorch policy -> LIBERO/MuJoCo -> 结果文件 +``` + +不启动 policy server，不使用 ROS，也不切换 Python 环境。OpenPI worker 与 +`lightx2v.infer` 使用同一个 base Python；仅任务特异的 Transformers 代码放在私有 +overlay 中。 + +以下命令默认在项目根目录执行： + +```bash +cd /data/liuhongda/lightx2v_openpi +conda activate base +``` + +启动脚本直接调用当前环境中的 `python`，因此运行前应确认 +`command -v python` 指向 base 环境。 + +## 默认路径 + +| 内容 | 路径 | +| --- | --- | +| Python | 当前已激活的 base 环境 | +| JAX checkpoint | `/data/liuhongda/openpi_data/openpi-assets/checkpoints/pi05_libero` | +| FP32 PyTorch checkpoint | `/data/liuhongda/openpi_data/openpi-assets/checkpoints/pi05_libero_pytorch_fp32` | +| OpenPI 源码 | `/data/liuhongda/openpi` | +| LIBERO | `/data/liuhongda/openpi/third_party/libero` | +| Transformers overlay | `/data/liuhongda/openpi_data/python_deps/openpi_official_pytorch_runtime` | + +路径都可以通过下文列出的环境变量覆盖。 + +## 1. 准备运行环境 + +setup 只安装或修复 OpenPI 所需的小包：base 环境中的 `mujoco==3.2.3`，以及 +私有 overlay 中的 `transformers==4.53.2` 和官方 replacement 文件。它不会修改 +Python、PyTorch 或 CUDA。 + +```bash +bash scripts/openpi/2_setup_pytorch_runtime.sh +``` + +训练或评测前可做只读检查： + +```bash +bash scripts/openpi/2_setup_pytorch_runtime.sh check +``` + +检查覆盖 checkpoint、Transformers replacement、MuJoCo 来源、LIBERO 来源和 +CUDA。启动脚本不会在每次运行前重复执行这项检查。 + +## 2. 转换官方权重 + +默认将官方 JAX checkpoint 转为 FP32 PyTorch checkpoint： + +```bash +bash scripts/openpi/1_convert_pi05_libero_to_pytorch.sh +``` + +选择输出精度或路径： + +```bash +bash scripts/openpi/1_convert_pi05_libero_to_pytorch.sh --precision bfloat16 + +OPENPI_JAX_CHECKPOINT=/path/to/pi05_libero \ +OPENPI_PYTORCH_CHECKPOINT=/path/to/pi05_libero_pytorch_fp32 \ +bash scripts/openpi/1_convert_pi05_libero_to_pytorch.sh --precision float32 +``` + +转换器使用 OpenPI 自身的环境，默认是 +`/data/liuhongda/openpi/.venv/bin/python`；本地推理仍使用 base Python。输出目录 +必须不存在或为空，脚本不会覆盖已有权重。 + +## 3. 单 episode rollout + +默认在 4 号卡运行 `libero_spatial/task_00/init_00`： + +```bash +bash scripts/openpi/run_libero_task_i2va.sh +``` + +指定 suite、task id 和 init-state id： + +```bash +CUDA_VISIBLE_DEVICES=6 \ +bash scripts/openpi/run_libero_task_i2va.sh libero_goal 3 0 +``` + +默认输出结构： + +```text +save_results/openpi_libero_tasks//task_XX/init_XX/ +├── rollout.mp4 +├── actions.npy +├── metrics.json +└── runtime/ +``` + +使用 `OPENPI_LIBERO_RESULT_ROOT=/path/to/results` 修改根目录。直接执行 +`run_libero_i2va.sh` 也会运行默认 episode，但 `run_libero_task_i2va.sh` 提供更清晰 +的任务目录组织。 + +## 4. 单卡定量评测 + +默认在一张卡上顺序评测四个 suite： + +```bash +CUDA_VISIBLE_DEVICES=6 \ +bash scripts/openpi/run_libero_evaluate_i2va.sh +``` + +只测一个 suite： + +```bash +CUDA_VISIBLE_DEVICES=6 \ +OPENPI_EVAL_BENCHMARKS=libero_goal \ +bash scripts/openpi/run_libero_evaluate_i2va.sh +``` + +最小 smoke test 应使用独立输出目录： + +```bash +CUDA_VISIBLE_DEVICES=4 \ +OPENPI_EVAL_BENCHMARKS=libero_goal \ +OPENPI_EVAL_TASK_IDS=3 \ +OPENPI_EVAL_NUM_TRIALS_PER_TASK=1 \ +OPENPI_EVAL_MAX_STEPS=2 \ +OPENPI_EVAL_RESUME=0 \ +OPENPI_EVAL_OUTPUT_DIR=/tmp/openpi-libero-smoke \ +bash scripts/openpi/run_libero_evaluate_i2va.sh +``` + +## 5. 按 suite 多卡评测 + +默认使用 4、5、6、7 号卡，每张卡运行一个 suite： + +```bash +bash scripts/openpi/run_libero_evaluate_parallel_i2va.sh +``` + +两卡运行完整 LIBERO-40： + +```bash +CUDA_VISIBLE_DEVICES=4,5 \ +bash scripts/openpi/run_libero_evaluate_parallel_i2va.sh +``` + +并行脚本固定评测四个官方 suite，支持 1、2 或 4 张卡。两卡运行时每张卡顺序执行 +两个 suite；每个 suite 都是独立的 +`python -m lightx2v.infer` 进程，并使用独立 cache 和 LIBERO config。 + +默认输出结构： + +```text +save_results/pi05_libero_pytorch_fp32_parallel_evaluation/ +├── logs/.log +├── runtime// +├── / +│ ├── episodes.jsonl +│ ├── episodes//task_XX/init_XX/metrics.json +│ └── summary.json +└── parallel_summary.json +``` + +`libero_summary.py` 只读取各 suite 的 JSON 结果并生成最终汇总，不是推理入口。 +输出锁会阻止两个并行任务同时写入同一目录。 + +## 评测协议 + +`configs/openpi/pi05_libero_eval.json` 是默认协议来源：每个 suite 包含 10 个任务， +每个任务测试 50 个 init states，共 500 episodes；完整 LIBERO-40 共 2000 个。 +环境/策略 seed 为 7/0，先执行 10 个 no-op，每次预测 10 个 action 并执行前 5 个。 +四个 suite 的最大步数分别为 220、280、300、520。 + +评测默认支持断点恢复。`protocol_id`、输入文件 manifest 和保存的 policy RNG state +用于避免混用不同协议，并保证前缀恢复后的随机数流与一次性运行一致。 + +## 常用覆盖项 + +| 环境变量 | 作用 | +| --- | --- | +| `OPENPI_MODEL_PATH` | PyTorch checkpoint | +| `OPENPI_CONFIG` | 模型 JSON | +| `OPENPI_EVAL_CONFIG` | 评测协议 JSON | +| `OPENPI_LIBERO_ROOT` | LIBERO checkout | +| `OPENPI_TRANSFORMERS_RUNTIME_PATH` | Transformers overlay | +| `CUDA_VISIBLE_DEVICES` | 单卡 GPU，或并行脚本的 GPU 列表 | +| `OPENPI_EVAL_BENCHMARKS` | 单卡评测的 suite 列表 | +| `OPENPI_EVAL_TASK_IDS` | task id 列表 | +| `OPENPI_EVAL_NUM_TRIALS_PER_TASK` | 每任务 trial 数 | +| `OPENPI_EVAL_MAX_STEPS` | 统一覆盖最大步数 | +| `OPENPI_EVAL_VIDEO_POLICY` | `none`、`failures` 或 `all` | +| `OPENPI_EVAL_SAVE_ACTIONS` | 是否保存 action，使用 `0/1` | +| `OPENPI_EVAL_RESUME` | 是否恢复已有结果，使用 `0/1` | +| `OPENPI_EVAL_OUTPUT_DIR` | 单卡评测输出目录 | +| `OPENPI_PARALLEL_OUTPUT_ROOT` | 多卡评测输出目录 | + +## 开发检查 + +```bash +bash -n scripts/openpi/run_libero_*.sh +python -m unittest discover -s scripts/openpi/tests -p 'test_*.py' -v +python scripts/openpi/tests/validate_pytorch_parity.py --self-check +pre-commit run --all-files +``` + +数值路径的关键约束是官方图像 resize/uint8 量化、FP64 动作反归一化、连续 policy +RNG 和 5-action replan queue；清理启动脚本时不应改变这些逻辑。 diff --git a/scripts/openpi/convert_jax_checkpoint.py b/scripts/openpi/convert_jax_checkpoint.py new file mode 100755 index 000000000..99ef8e96c --- /dev/null +++ b/scripts/openpi/convert_jax_checkpoint.py @@ -0,0 +1,207 @@ +#!/usr/bin/env python3 +"""Convert an OpenPI pi0.5 Orbax checkpoint with the upstream converter.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import shutil +import subprocess +import sys +import tempfile +from collections import Counter +from pathlib import Path + +from safetensors import safe_open + +PROJECT_ROOT = Path(__file__).resolve().parents[2] +WORKSPACE_ROOT = PROJECT_ROOT.parent +OPENPI_DATA_ROOT = WORKSPACE_ROOT / "openpi_data" +DEFAULT_OPENPI_ROOT = WORKSPACE_ROOT / "openpi" +DEFAULT_SOURCE = OPENPI_DATA_ROOT / "openpi-assets/checkpoints/pi05_libero" +DEFAULT_TOKENIZER = OPENPI_DATA_ROOT / "big_vision/paligemma_tokenizer.model" +DEFAULT_TRANSFORMERS_RUNTIME = OPENPI_DATA_ROOT / "python_deps/openpi_official_pytorch_runtime" + +EXPECTED_TENSORS = 812 +DTYPE_BY_PRECISION = {"float32": "F32", "bfloat16": "BF16"} + + +def _path(value: str | Path) -> Path: + return Path(value).expanduser().resolve() + + +def _default_output(source: Path, precision: str) -> Path: + suffix = "_pytorch_fp32" if precision == "float32" else "_pytorch" + return source.with_name(f"{source.name}{suffix}") + + +def _conversion_environment(args: argparse.Namespace) -> dict[str, str]: + python_paths = [] + if args.transformers_runtime.is_dir(): + python_paths.append(str(args.transformers_runtime)) + python_paths.append(str(args.openpi_root / "src")) + if current := os.environ.get("PYTHONPATH"): + python_paths.append(current) + environment = os.environ.copy() + environment.update( + { + "PYTHONPATH": os.pathsep.join(python_paths), + "PYTHONDONTWRITEBYTECODE": "1", + "PYTHONNOUSERSITE": "1", + "USE_FLAX": "0", + } + ) + return environment + + +def _check_inputs(args: argparse.Namespace) -> None: + required = ( + args.openpi_root / "examples/convert_jax_model_to_pytorch.py", + args.source / "params/_METADATA", + args.source / "assets/physical-intelligence/libero/norm_stats.json", + args.tokenizer, + ) + missing = [str(path) for path in required if not path.is_file()] + if missing: + raise RuntimeError("missing conversion input:\n- " + "\n- ".join(missing)) + if "pi05" not in str(args.source).lower(): + raise RuntimeError("the upstream converter selects pi0.5 layers from a source path containing 'pi05'") + if args.output == args.source or args.output.is_relative_to(args.source) or args.source.is_relative_to(args.output): + raise RuntimeError(f"source and output must not overlap: {args.source}, {args.output}") + if args.output.exists() and (not args.output.is_dir() or any(args.output.iterdir())): + if not args.dry_run: + raise RuntimeError(f"output must not exist or must be an empty directory: {args.output}") + print(f"note: output already exists and must be changed before conversion: {args.output}") + + probe = "import transformers; from transformers.models.siglip import check; assert transformers.__version__ == '4.53.2'; assert check.check_whether_transformers_replace_is_installed_correctly()" + try: + subprocess.run([sys.executable, "-c", probe], env=_conversion_environment(args), check=True) + except subprocess.CalledProcessError as error: + raise RuntimeError("the conversion Python cannot load OpenPI's patched transformers==4.53.2; run scripts/openpi/2_setup_pytorch_runtime.sh first") from error + + +def _run_converter(args: argparse.Namespace, output: Path) -> None: + command = [ + sys.executable, + str(args.openpi_root / "examples/convert_jax_model_to_pytorch.py"), + "--checkpoint-dir", + str(args.source), + "--config-name", + args.config_name, + "--output-path", + str(output), + "--precision", + args.precision, + ] + print("+", " ".join(command), flush=True) + subprocess.run(command, cwd=args.openpi_root, env=_conversion_environment(args), check=True) + + +def _copy_assets(args: argparse.Namespace, output: Path) -> None: + shutil.copytree(args.source / "assets", output / "assets", dirs_exist_ok=True) + tokenizer = output / "assets/paligemma_tokenizer.model" + tokenizer.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(args.tokenizer, tokenizer) + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _validate_output(args: argparse.Namespace, output: Path) -> dict[str, object]: + required = ( + output / "model.safetensors", + output / "config.json", + output / "assets/paligemma_tokenizer.model", + output / "assets/physical-intelligence/libero/norm_stats.json", + ) + missing = [str(path) for path in required if not path.is_file()] + if missing: + raise RuntimeError("converted checkpoint is incomplete:\n- " + "\n- ".join(missing)) + + config = json.loads((output / "config.json").read_text(encoding="utf-8")) + if config.get("precision") != args.precision: + raise RuntimeError(f"config precision is {config.get('precision')!r}, expected {args.precision!r}") + + with safe_open(output / "model.safetensors", framework="pt", device="cpu") as checkpoint: + keys = list(checkpoint.keys()) + dtypes = Counter(checkpoint.get_slice(key).get_dtype() for key in keys) + expected_dtypes = Counter({DTYPE_BY_PRECISION[args.precision]: EXPECTED_TENSORS}) + if len(keys) != EXPECTED_TENSORS or dtypes != expected_dtypes: + raise RuntimeError(f"expected {EXPECTED_TENSORS} {args.precision} tensors, got {dict(dtypes)}") + + rows = [f"{_sha256(path)} {path.relative_to(output).as_posix()}" for path in required] + (output / "SHA256SUMS").write_text("\n".join(rows) + "\n", encoding="utf-8") + return {"output": str(args.output), "precision": args.precision, "tensors": len(keys), "dtypes": dict(dtypes)} + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser.add_argument("--source", default=os.environ.get("OPENPI_JAX_CHECKPOINT", str(DEFAULT_SOURCE))) + parser.add_argument("--output", default=os.environ.get("OPENPI_PYTORCH_CHECKPOINT")) + parser.add_argument( + "--precision", + choices=tuple(DTYPE_BY_PRECISION), + default=os.environ.get("OPENPI_CONVERT_PRECISION", os.environ.get("OPENPI_OUTPUT_PRECISION", "float32")), + ) + parser.add_argument("--config-name", default=os.environ.get("OPENPI_CONFIG_NAME", "pi05_libero")) + parser.add_argument("--openpi-root", default=os.environ.get("OPENPI_PATH", str(DEFAULT_OPENPI_ROOT))) + parser.add_argument("--tokenizer", default=os.environ.get("OPENPI_TOKENIZER_PATH", str(DEFAULT_TOKENIZER))) + parser.add_argument( + "--transformers-runtime", + default=os.environ.get("OPENPI_TRANSFORMERS_RUNTIME_PATH", str(DEFAULT_TRANSFORMERS_RUNTIME)), + ) + parser.add_argument("--dry-run", action="store_true") + return parser + + +def main() -> int: + args = build_parser().parse_args() + args.source = _path(args.source) + args.output = _path(args.output) if args.output else _default_output(args.source, args.precision) + args.openpi_root = _path(args.openpi_root) + args.tokenizer = _path(args.tokenizer) + args.transformers_runtime = _path(args.transformers_runtime) + _check_inputs(args) + + plan = { + "source": str(args.source), + "output": str(args.output), + "precision": args.precision, + "converter": str(args.openpi_root / "examples/convert_jax_model_to_pytorch.py"), + "python": sys.executable, + } + print(json.dumps(plan, indent=2)) + if args.dry_run: + print("Dry run complete; no checkpoint was written.") + return 0 + + args.output.parent.mkdir(parents=True, exist_ok=True) + stage = Path(tempfile.mkdtemp(prefix=f".{args.output.name}.staging.", dir=args.output.parent)) + try: + _run_converter(args, stage) + _copy_assets(args, stage) + report = _validate_output(args, stage) + if args.output.exists(): + args.output.rmdir() + stage.rename(args.output) + except Exception: + shutil.rmtree(stage, ignore_errors=True) + raise + print(json.dumps(report, indent=2)) + print(f"OpenPI PyTorch checkpoint is ready: {args.output}") + return 0 + + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except (OSError, RuntimeError, subprocess.CalledProcessError) as error: + print(f"error: {error}", file=sys.stderr) + raise SystemExit(1) from None diff --git a/scripts/openpi/libero_summary.py b/scripts/openpi/libero_summary.py new file mode 100755 index 000000000..77d489876 --- /dev/null +++ b/scripts/openpi/libero_summary.py @@ -0,0 +1,145 @@ +#!/usr/bin/env python3 +"""Aggregate completed LIBERO suite summaries without launching inference.""" + +from __future__ import annotations + +import argparse +import json +import os +import time +from pathlib import Path + +SUITES = ("libero_spatial", "libero_object", "libero_goal", "libero_10") +EPISODES_PER_SUITE = 500 + + +def _parse_suites(value: str) -> tuple[str, ...]: + suites = tuple(item.strip() for item in value.split(",") if item.strip()) + unknown = [suite for suite in suites if suite not in SUITES] + if not suites or unknown or len(set(suites)) != len(suites): + raise ValueError(f"invalid LIBERO suite selection: {value!r}") + return suites + + +def _parse_status(values: list[str], suites: tuple[str, ...]) -> dict[str, int]: + statuses: dict[str, int] = {} + for value in values: + suite, separator, return_code = value.partition("=") + if not separator or suite not in suites or suite in statuses: + raise ValueError(f"invalid worker status: {value!r}") + statuses[suite] = int(return_code) + missing = set(suites) - statuses.keys() + if missing: + raise ValueError(f"missing worker status for: {sorted(missing)}") + return statuses + + +def aggregate(output_root: Path, suites: tuple[str, ...], statuses: dict[str, int], output_file: Path) -> tuple[dict, list[str]]: + errors: list[str] = [] + shards: dict[str, dict] = {} + successes = 0 + completed = 0 + rates: list[float] = [] + + for suite in suites: + summary_path = output_root / suite / "summary.json" + if not summary_path.is_file(): + errors.append(f"{suite}: missing {summary_path}") + continue + try: + summary = json.loads(summary_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + errors.append(f"{suite}: invalid summary: {exc}") + continue + + try: + shard_expected = int(summary["expected_episodes"]) + shard_completed = int(summary["completed_episodes"]) + shard_successes = int(summary["successes"]) + rate = float(summary["success_rate"]) + status = summary["status"] + official_protocol = summary["official_protocol"] + protocol_id = summary["protocol_id"] + rollout_errors = int(summary["errors"]) + artifact_errors = int(summary["artifact_errors"]) + except (KeyError, TypeError, ValueError) as exc: + errors.append(f"{suite}: incompatible summary schema: {exc}") + continue + if statuses[suite] != 0: + errors.append(f"{suite}: worker exit code {statuses[suite]}") + if status not in {"complete", "complete_with_errors"}: + errors.append(f"{suite}: summary status is {status!r}") + if official_protocol is not True: + errors.append(f"{suite}: evaluator did not mark the shard as official protocol") + if not isinstance(protocol_id, str) or not protocol_id: + errors.append(f"{suite}: missing protocol_id") + if rollout_errors or artifact_errors: + errors.append(f"{suite}: rollout_errors={rollout_errors} artifact_errors={artifact_errors}") + if shard_expected != EPISODES_PER_SUITE or shard_completed != EPISODES_PER_SUITE: + errors.append(f"{suite}: expected official {EPISODES_PER_SUITE}/{EPISODES_PER_SUITE} episodes, got {shard_completed}/{shard_expected}") + rates.append(rate) + completed += shard_completed + successes += shard_successes + shards[suite] = { + "output_dir": str(output_root / suite), + "expected_episodes": shard_expected, + "completed_episodes": shard_completed, + "successes": shard_successes, + "success_rate": rate, + "protocol_id": protocol_id, + } + + expected = EPISODES_PER_SUITE * len(suites) + payload = { + "schema_version": 1, + "status": "complete" if not errors else "invalid", + "expected_episodes": expected, + "completed_episodes": completed, + "successes": successes, + "failures": completed - successes, + "success_rate": successes / completed * 100.0 if completed else None, + "mean_suite_success_rate": sum(rates) / len(rates) if len(rates) == len(suites) else None, + "shards": shards, + "validation_errors": errors, + "updated_at_unix": time.time(), + } + output_file.parent.mkdir(parents=True, exist_ok=True) + temporary = output_file.with_name(f".{output_file.name}.{os.getpid()}.tmp") + temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + os.replace(temporary, output_file) + return payload, errors + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-root", type=Path, required=True) + parser.add_argument("--suites", required=True) + parser.add_argument("--worker-status", action="append", default=[]) + parser.add_argument("--output-file", type=Path) + args = parser.parse_args() + try: + output_root = args.output_root.expanduser().resolve() + suites = _parse_suites(args.suites) + statuses = _parse_status(args.worker_status, suites) + output_file = args.output_file.expanduser().resolve() if args.output_file else output_root / "parallel_summary.json" + payload, errors = aggregate(output_root, suites, statuses, output_file) + except (OSError, TypeError, ValueError) as exc: + parser.error(str(exc)) + + print("\nLIBERO suite summary") + for suite in suites: + shard = payload["shards"].get(suite, {}) + rate = shard.get("success_rate") + rate_text = "n/a" if rate is None else f"{float(rate):.2f}%" + print(f" {suite:<18} {shard.get('successes', 0):>3}/{shard.get('expected_episodes', 500):<3} {rate_text:>8}") + overall = payload["success_rate"] + overall_text = "n/a" if overall is None else f"{overall:.2f}%" + print(f" {'overall':<18} {payload['successes']:>3}/{payload['expected_episodes']:<3} {overall_text:>8}") + print(f" summary: {output_file}") + for error in errors: + print(f"error: {error}") + return 1 if errors else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/openpi/run_libero_evaluate_i2va.sh b/scripts/openpi/run_libero_evaluate_i2va.sh new file mode 100755 index 000000000..2306e579b --- /dev/null +++ b/scripts/openpi/run_libero_evaluate_i2va.sh @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +set -euo pipefail + +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +lightx2v_path="${LIGHTX2V_PATH:-$(cd -- "${script_dir}/../.." && pwd)}" +workspace_root="$(dirname -- "${lightx2v_path}")" +openpi_data_root="${OPENPI_DATA_ROOT:-${workspace_root}/openpi_data}" + +model_path="${OPENPI_MODEL_PATH:-${openpi_data_root}/openpi-assets/checkpoints/pi05_libero_pytorch_fp32}" +config_json="${OPENPI_CONFIG:-${lightx2v_path}/configs/openpi/pi05_libero.json}" +output_dir="${OPENPI_EVAL_OUTPUT_DIR:-${lightx2v_path}/save_results/pi05_libero_pytorch_fp32_evaluation}" +runtime_dir="${OPENPI_RUNTIME_DIR:-${output_dir}/runtime}" + +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-4}" +export MUJOCO_GL="${MUJOCO_GL:-egl}" +export PYOPENGL_PLATFORM="${PYOPENGL_PLATFORM:-egl}" +export OPENPI_RUN_MODE=evaluate +export OPENPI_LIBERO_ROOT="${OPENPI_LIBERO_ROOT:-${workspace_root}/openpi/third_party/libero}" +export OPENPI_LIBERO_CONFIG_DIR="${OPENPI_LIBERO_CONFIG_DIR:-${runtime_dir}/libero_config}" +export OPENPI_EVAL_CONFIG="${OPENPI_EVAL_CONFIG:-${lightx2v_path}/configs/openpi/pi05_libero_eval.json}" +export OPENPI_TRANSFORMERS_RUNTIME_PATH="${OPENPI_TRANSFORMERS_RUNTIME_PATH:-${openpi_data_root}/python_deps/openpi_official_pytorch_runtime}" +export NUMBA_CACHE_DIR="${NUMBA_CACHE_DIR:-${runtime_dir}/numba}" +export MPLCONFIGDIR="${MPLCONFIGDIR:-${runtime_dir}/matplotlib}" +export XDG_CACHE_HOME="${XDG_CACHE_HOME:-${runtime_dir}/cache}" +export PYTHONNOUSERSITE=1 +export TOKENIZERS_PARALLELISM=false +export PROFILING_DEBUG_LEVEL="${PROFILING_DEBUG_LEVEL:-0}" + +mkdir -p "${OPENPI_LIBERO_CONFIG_DIR}" "${NUMBA_CACHE_DIR}" "${MPLCONFIGDIR}" "${XDG_CACHE_HOME}" +cd "${lightx2v_path}" + +exec python -m lightx2v.infer \ + --model_cls openpi \ + --task i2va \ + --model_path "${model_path}" \ + --config_json "${config_json}" \ + --seed "${OPENPI_POLICY_SEED:-0}" \ + --save_result_path "${output_dir}" diff --git a/scripts/openpi/run_libero_evaluate_parallel_i2va.sh b/scripts/openpi/run_libero_evaluate_parallel_i2va.sh new file mode 100755 index 000000000..156d76d7a --- /dev/null +++ b/scripts/openpi/run_libero_evaluate_parallel_i2va.sh @@ -0,0 +1,137 @@ +#!/usr/bin/env bash +set -euo pipefail + +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +lightx2v_path="${LIGHTX2V_PATH:-$(cd -- "${script_dir}/../.." && pwd)}" +workspace_root="$(dirname -- "${lightx2v_path}")" +openpi_data_root="${OPENPI_DATA_ROOT:-${workspace_root}/openpi_data}" + +model_path="${OPENPI_MODEL_PATH:-${openpi_data_root}/openpi-assets/checkpoints/pi05_libero_pytorch_fp32}" +config_json="${OPENPI_CONFIG:-${lightx2v_path}/configs/openpi/pi05_libero.json}" +eval_config="${OPENPI_EVAL_CONFIG:-${lightx2v_path}/configs/openpi/pi05_libero_eval.json}" +libero_root="${OPENPI_LIBERO_ROOT:-${workspace_root}/openpi/third_party/libero}" +transformers_runtime="${OPENPI_TRANSFORMERS_RUNTIME_PATH:-${openpi_data_root}/python_deps/openpi_official_pytorch_runtime}" +gpu_list="${CUDA_VISIBLE_DEVICES:-4,5,6,7}" +suite_list="libero_spatial,libero_object,libero_goal,libero_10" +output_root="${OPENPI_PARALLEL_OUTPUT_ROOT:-${lightx2v_path}/save_results/pi05_libero_pytorch_fp32_parallel_evaluation}" + +IFS=',' read -ra gpus <<< "${gpu_list}" +suites=(libero_spatial libero_object libero_goal libero_10) +case "${#gpus[@]}" in + 1) assignments=("${suites[*]}") ;; + 2) assignments=("libero_10 libero_spatial" "libero_goal libero_object") ;; + 4) assignments=("${suites[@]}") ;; + *) + echo "LIBERO-40 parallel evaluation requires 1, 2, or 4 GPUs" >&2 + exit 2 + ;; +esac + +mkdir -p "${output_root}/logs" "${output_root}/runtime" +exec {lock_fd}>"${output_root}/.parallel.lock" +if ! flock -n "${lock_fd}"; then + echo "Another evaluation is using ${output_root}" >&2 + exit 2 +fi + +run_worker() { + local worker_index="$1" + local gpu="${gpus[worker_index]}" + local child_pid="" + local worker_status=0 + local suite suite_output suite_runtime log_path exit_code + + stop_suite() { + if [[ -n "${child_pid}" ]] && kill -0 "${child_pid}" 2>/dev/null; then + kill -TERM -- "-${child_pid}" 2>/dev/null || kill -TERM "${child_pid}" 2>/dev/null || true + wait "${child_pid}" 2>/dev/null || true + fi + exit 143 + } + trap stop_suite INT TERM + + for suite in ${assignments[worker_index]}; do + suite_output="${output_root}/${suite}" + suite_runtime="${output_root}/runtime/${suite}" + log_path="${output_root}/logs/${suite}.log" + mkdir -p "${suite_runtime}/libero_config" "${suite_runtime}/numba" "${suite_runtime}/matplotlib" "${suite_runtime}/cache" + printf '\n[%(%Y-%m-%dT%H:%M:%SZ)T] suite=%s gpu=%s\n' -1 "${suite}" "${gpu}" >> "${log_path}" + + ( + export CUDA_VISIBLE_DEVICES="${gpu}" + export MUJOCO_GL="${MUJOCO_GL:-egl}" + export PYOPENGL_PLATFORM="${PYOPENGL_PLATFORM:-egl}" + export OPENPI_RUN_MODE=evaluate + export OPENPI_LIBERO_ROOT="${libero_root}" + export OPENPI_LIBERO_CONFIG_DIR="${suite_runtime}/libero_config" + export OPENPI_EVAL_CONFIG="${eval_config}" + export OPENPI_EVAL_BENCHMARKS="${suite}" + export OPENPI_TRANSFORMERS_RUNTIME_PATH="${transformers_runtime}" + export NUMBA_CACHE_DIR="${suite_runtime}/numba" + export MPLCONFIGDIR="${suite_runtime}/matplotlib" + export XDG_CACHE_HOME="${suite_runtime}/cache" + export PYTHONNOUSERSITE=1 + export TOKENIZERS_PARALLELISM=false + export PROFILING_DEBUG_LEVEL="${PROFILING_DEBUG_LEVEL:-0}" + exec setsid python -m lightx2v.infer \ + --model_cls openpi \ + --task i2va \ + --model_path "${model_path}" \ + --config_json "${config_json}" \ + --seed "${OPENPI_POLICY_SEED:-0}" \ + --save_result_path "${suite_output}" + ) >> "${log_path}" 2>&1 & + child_pid=$! + if wait "${child_pid}"; then + exit_code=0 + else + exit_code=$? + worker_status=1 + fi + child_pid="" + printf '%s\n' "${exit_code}" > "${output_root}/runtime/${suite}.status" + echo "${suite}: exit=${exit_code}, log=${log_path}" + done + return "${worker_status}" +} + +worker_pids=() +stop_workers() { + local pid + for pid in "${worker_pids[@]}"; do + kill -TERM "${pid}" 2>/dev/null || true + done + for pid in "${worker_pids[@]}"; do + wait "${pid}" 2>/dev/null || true + done +} +trap stop_workers EXIT +trap 'exit 130' INT +trap 'exit 143' HUP TERM + +for index in "${!gpus[@]}"; do + echo "GPU ${gpus[index]}: ${assignments[index]}" + run_worker "${index}" & + worker_pids+=("$!") +done + +worker_failed=0 +for pid in "${worker_pids[@]}"; do + wait "${pid}" || worker_failed=1 +done +worker_pids=() +trap - EXIT HUP INT TERM + +summary_command=( + python "${script_dir}/libero_summary.py" + --output-root "${output_root}" + --suites "${suite_list}" +) +for suite in "${suites[@]}"; do + exit_code=1 + [[ ! -f "${output_root}/runtime/${suite}.status" ]] || read -r exit_code < "${output_root}/runtime/${suite}.status" + summary_command+=(--worker-status "${suite}=${exit_code}") +done + +"${summary_command[@]}" || worker_failed=1 +exit "${worker_failed}" diff --git a/scripts/openpi/run_libero_i2va.sh b/scripts/openpi/run_libero_i2va.sh new file mode 100755 index 000000000..32dbf779d --- /dev/null +++ b/scripts/openpi/run_libero_i2va.sh @@ -0,0 +1,41 @@ +#!/usr/bin/env bash +set -euo pipefail + +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +lightx2v_path="${LIGHTX2V_PATH:-$(cd -- "${script_dir}/../.." && pwd)}" +workspace_root="$(dirname -- "${lightx2v_path}")" +openpi_data_root="${OPENPI_DATA_ROOT:-${workspace_root}/openpi_data}" + +model_path="${OPENPI_MODEL_PATH:-${openpi_data_root}/openpi-assets/checkpoints/pi05_libero_pytorch_fp32}" +config_json="${OPENPI_CONFIG:-${lightx2v_path}/configs/openpi/pi05_libero.json}" +libero_root="${OPENPI_LIBERO_ROOT:-${workspace_root}/openpi/third_party/libero}" +transformers_runtime="${OPENPI_TRANSFORMERS_RUNTIME_PATH:-${openpi_data_root}/python_deps/openpi_official_pytorch_runtime}" +result_root="${OPENPI_LIBERO_RESULT_ROOT:-${lightx2v_path}/save_results/openpi_libero}" +runtime_dir="${OPENPI_RUNTIME_DIR:-${result_root}/runtime}" + +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-4}" +export MUJOCO_GL="${MUJOCO_GL:-egl}" +export PYOPENGL_PLATFORM="${PYOPENGL_PLATFORM:-egl}" +export OPENPI_RUN_MODE=rollout +export OPENPI_LIBERO_ROOT="${libero_root}" +export OPENPI_LIBERO_CONFIG_DIR="${OPENPI_LIBERO_CONFIG_DIR:-${runtime_dir}/libero_config}" +export OPENPI_TRANSFORMERS_RUNTIME_PATH="${transformers_runtime}" +export NUMBA_CACHE_DIR="${NUMBA_CACHE_DIR:-${runtime_dir}/numba}" +export MPLCONFIGDIR="${MPLCONFIGDIR:-${runtime_dir}/matplotlib}" +export XDG_CACHE_HOME="${XDG_CACHE_HOME:-${runtime_dir}/cache}" +export PYTHONNOUSERSITE=1 +export TOKENIZERS_PARALLELISM=false +export PROFILING_DEBUG_LEVEL="${PROFILING_DEBUG_LEVEL:-0}" + +mkdir -p "${OPENPI_LIBERO_CONFIG_DIR}" "${NUMBA_CACHE_DIR}" "${MPLCONFIGDIR}" "${XDG_CACHE_HOME}" +cd "${lightx2v_path}" + +exec python -m lightx2v.infer \ + --model_cls openpi \ + --task i2va \ + --model_path "${model_path}" \ + --config_json "${config_json}" \ + --seed "${OPENPI_POLICY_SEED:-0}" \ + --prompt "${OPENPI_TASK_DESCRIPTION:-}" \ + --save_result_path "${OPENPI_SAVE_VIDEO_PATH:-${result_root}/rollout.mp4}" \ + --save_action_path "${OPENPI_SAVE_ACTION_PATH:-${result_root}/actions.npy}" diff --git a/scripts/openpi/run_libero_task_i2va.sh b/scripts/openpi/run_libero_task_i2va.sh new file mode 100755 index 000000000..ec435ca23 --- /dev/null +++ b/scripts/openpi/run_libero_task_i2va.sh @@ -0,0 +1,42 @@ +#!/usr/bin/env bash +set -euo pipefail + +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +lightx2v_path="${LIGHTX2V_PATH:-$(cd -- "${script_dir}/../.." && pwd)}" + +if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then + echo "Usage: bash scripts/openpi/run_libero_task_i2va.sh [suite] [task_id] [init_state_id]" + exit 0 +fi +if [[ $# -gt 3 ]]; then + echo "Usage: bash scripts/openpi/run_libero_task_i2va.sh [suite] [task_id] [init_state_id]" >&2 + exit 2 +fi + +suite="${1:-${LIBERO_BENCHMARK:-libero_spatial}}" +task_id="${2:-${LIBERO_TASK_ID:-0}}" +init_state_id="${3:-${LIBERO_INIT_STATE_ID:-0}}" + +case "${suite}" in + libero_spatial|libero_object|libero_goal|libero_10) ;; + *) echo "Unknown LIBERO suite: ${suite}" >&2; exit 2 ;; +esac +if [[ ! "${task_id}" =~ ^[0-9]+$ || ! "${init_state_id}" =~ ^[0-9]+$ ]]; then + echo "task_id and init_state_id must be non-negative integers" >&2 + exit 2 +fi + +printf -v task_dir 'task_%02d' "$((10#${task_id}))" +printf -v init_dir 'init_%02d' "$((10#${init_state_id}))" +result_root="${OPENPI_LIBERO_RESULT_ROOT:-${lightx2v_path}/save_results/openpi_libero_tasks}" +episode_dir="${result_root}/${suite}/${task_dir}/${init_dir}" + +export LIBERO_BENCHMARK="${suite}" +export LIBERO_TASK_ID="$((10#${task_id}))" +export LIBERO_INIT_STATE_ID="$((10#${init_state_id}))" +export OPENPI_SAVE_VIDEO_PATH="${episode_dir}/rollout.mp4" +export OPENPI_SAVE_ACTION_PATH="${episode_dir}/actions.npy" +export OPENPI_SAVE_METRICS_PATH="${episode_dir}/metrics.json" +export OPENPI_RUNTIME_DIR="${episode_dir}/runtime" + +exec bash "${script_dir}/run_libero_i2va.sh" diff --git a/scripts/openpi/runtime.py b/scripts/openpi/runtime.py new file mode 100755 index 000000000..df7897610 --- /dev/null +++ b/scripts/openpi/runtime.py @@ -0,0 +1,498 @@ +#!/usr/bin/env python3 +"""Prepare the patched Transformers runtime and validate OpenPI/LIBERO. + +Transformers remains isolated because OpenPI carries task-specific replacement +files. MuJoCo is installed into and imported directly from the selected base +Python environment. +""" + +from __future__ import annotations + +import argparse +import filecmp +import importlib.metadata +import importlib.util +import json +import os +import shutil +import subprocess +import sys +import sysconfig +import tempfile +from pathlib import Path + +PROJECT_ROOT = Path(__file__).resolve().parents[2] +WORKSPACE_ROOT = PROJECT_ROOT.parent +OPENPI_DATA_ROOT = WORKSPACE_ROOT / "openpi_data" +DEFAULT_PYTHON = Path("/opt/conda/bin/python") +DEFAULT_MODEL = OPENPI_DATA_ROOT / "openpi-assets/checkpoints/pi05_libero_pytorch_fp32" +DEFAULT_TRANSFORMERS_RUNTIME = OPENPI_DATA_ROOT / "python_deps/openpi_official_pytorch_runtime" +DEFAULT_LIBERO_ROOT = WORKSPACE_ROOT / "openpi/third_party/libero" +DEFAULT_MODEL_CONFIG = PROJECT_ROOT / "configs/openpi/pi05_libero.json" +DEFAULT_EVAL_CONFIG = PROJECT_ROOT / "configs/openpi/pi05_libero_eval.json" + +TRANSFORMERS_EXPECTED = { + "transformers": ("transformers", "4.53.2"), + "huggingface-hub": ("huggingface_hub", "0.32.3"), + "tokenizers": ("tokenizers", "0.21.1"), +} +TRANSFORMERS_PACKAGES = tuple(f"{distribution}=={version}" for distribution, (_module, version) in TRANSFORMERS_EXPECTED.items()) + +# The official LIBERO client uses this MuJoCo version. All other simulator and +# model dependencies remain untouched in the base Python environment. +BASE_MUJOCO_VERSION = "3.2.3" +BASE_MUJOCO_PACKAGE = f"mujoco=={BASE_MUJOCO_VERSION}" + +PATCH_FILES = ( + "models/gemma/configuration_gemma.py", + "models/gemma/modeling_gemma.py", + "models/paligemma/modeling_paligemma.py", + "models/siglip/check.py", + "models/siglip/modeling_siglip.py", +) + + +def _resolved(path: str | Path) -> Path: + return Path(path).expanduser().resolve() + + +def _assert_safe_target(target: Path) -> None: + target = _resolved(target) + protected = { + Path.home().resolve(), + PROJECT_ROOT.resolve(), + WORKSPACE_ROOT.resolve(), + OPENPI_DATA_ROOT.resolve(), + Path(sys.prefix).resolve(), + Path(sys.base_prefix).resolve(), + Path(tempfile.gettempdir()).resolve(), + } + if target == Path(target.anchor) or target.parent == Path(target.anchor): + raise RuntimeError(f"refusing broad runtime target: {target}") + if target.is_relative_to(PROJECT_ROOT.resolve()): + raise RuntimeError(f"runtime target must stay outside the source checkout: {target}") + for path in protected: + if target == path or path.is_relative_to(target): + raise RuntimeError(f"runtime target is a protected directory or its parent: {target}") + + +def _patch_root() -> Path: + return PROJECT_ROOT / "lightx2v/models/networks/openpi/transformers_replace" + + +def _validate_patch_sources() -> None: + missing = [str(_patch_root() / relative) for relative in PATCH_FILES if not (_patch_root() / relative).is_file()] + if missing: + raise RuntimeError("missing OpenPI Transformers replacement files:\n- " + "\n- ".join(missing)) + + +def _run(command: list[str], *, env: dict[str, str] | None = None) -> None: + print("+", " ".join(command), flush=True) + subprocess.run(command, check=True, env=env) + + +def _pip_install(target: Path, requirements: tuple[str, ...], dry_run: bool) -> None: + command = [ + sys.executable, + "-m", + "pip", + "install", + "--disable-pip-version-check", + "--no-input", + "--no-deps", + "--no-compile", + "--target", + str(target), + *requirements, + ] + if dry_run: + print("+", " ".join(command)) + return + _run(command) + + +def _copy_transformers_patches(target: Path) -> None: + for relative in PATCH_FILES: + source = _patch_root() / relative + destination = target / "transformers" / relative + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source, destination) + + +def _check_patch_overlay(target: Path) -> None: + _validate_patch_sources() + mismatches = [] + for relative in PATCH_FILES: + source = _patch_root() / relative + installed = target / "transformers" / relative + if not installed.is_file() or not filecmp.cmp(source, installed, shallow=False): + mismatches.append(str(installed)) + if mismatches: + raise RuntimeError("OpenPI Transformers replacement mismatch:\n- " + "\n- ".join(mismatches)) + print("Transformers replacement set: exact OpenPI vendored copy") + + +PROBE = r""" +import importlib +import importlib.metadata +import json +import sys +import sysconfig +from pathlib import Path + +root = Path(sys.argv[1]).resolve() +expected = json.loads(sys.argv[2]) +result = {} +for distribution, values in expected.items(): + module_name, wanted = values + actual = importlib.metadata.version(distribution) + module = importlib.import_module(module_name) + origin = Path(module.__file__).resolve() + if actual != wanted: + raise RuntimeError(f"{distribution}: expected {wanted}, got {actual}") + if not origin.is_relative_to(root): + raise RuntimeError(f"{module_name} imported outside overlay: {origin}") + result[distribution] = {"version": actual, "origin": str(origin)} +print(json.dumps(result, sort_keys=True)) +""" + + +def _probe_overlay(target: Path, expected: dict[str, tuple[str, str]]) -> None: + env = os.environ.copy() + env.update( + { + "PYTHONPATH": str(target), + "PYTHONDONTWRITEBYTECODE": "1", + "MUJOCO_GL": "disable", + "USE_FLAX": "0", + } + ) + subprocess.run( + [sys.executable, "-c", PROBE, str(target), json.dumps(expected)], + check=True, + env=env, + ) + + +def _prepare_transformers(target: Path, dry_run: bool) -> None: + target = _resolved(target) + _assert_safe_target(target) + _validate_patch_sources() + if target.exists(): + try: + _probe_overlay(target, TRANSFORMERS_EXPECTED) + except (OSError, RuntimeError, subprocess.CalledProcessError) as exc: + raise RuntimeError(f"refusing to modify an invalid Transformers overlay: {target}") from exc + try: + _check_patch_overlay(target) + except RuntimeError: + pass + else: + print(f"transformers overlay already valid: {target}") + return + print(f"repair Transformers replacements: {target}") + if dry_run: + for relative in PATCH_FILES: + print(f"copy {_patch_root() / relative} -> {target / 'transformers' / relative}") + return + _copy_transformers_patches(target) + _check_patch_overlay(target) + _probe_overlay(target, TRANSFORMERS_EXPECTED) + print(f"transformers overlay ready: {target}") + return + + print(f"prepare transformers overlay: {target}") + if dry_run: + _pip_install(target, TRANSFORMERS_PACKAGES, dry_run=True) + for relative in PATCH_FILES: + print(f"copy {_patch_root() / relative} -> {target / 'transformers' / relative}") + return + + target.parent.mkdir(parents=True, exist_ok=True) + stage = Path(tempfile.mkdtemp(prefix=f".{target.name}.stage.", dir=target.parent)) + try: + _pip_install(stage, TRANSFORMERS_PACKAGES, dry_run=False) + _copy_transformers_patches(stage) + _check_patch_overlay(stage) + _probe_overlay(stage, TRANSFORMERS_EXPECTED) + if target.exists(): + raise RuntimeError(f"runtime target appeared while preparing it: {target}") + stage.replace(target) + finally: + if stage.exists(): + shutil.rmtree(stage) + print(f"transformers overlay ready: {target}") + + +def _base_site_packages() -> Path: + return _resolved(sysconfig.get_paths()["purelib"]) + + +def _base_mujoco_ready() -> bool: + try: + if importlib.metadata.version("mujoco") != BASE_MUJOCO_VERSION: + return False + spec = importlib.util.find_spec("mujoco") + if spec is None or spec.origin is None: + return False + return _resolved(spec.origin).is_relative_to(_base_site_packages()) + except (ImportError, importlib.metadata.PackageNotFoundError, ValueError): + return False + + +def _prepare_base_mujoco(dry_run: bool) -> None: + if _base_mujoco_ready(): + print(f"base MuJoCo already valid: {BASE_MUJOCO_VERSION} ({_base_site_packages() / 'mujoco'})") + return + command = [ + sys.executable, + "-m", + "pip", + "install", + "--disable-pip-version-check", + "--no-input", + "--no-deps", + "--upgrade", + BASE_MUJOCO_PACKAGE, + ] + if dry_run: + print("+", " ".join(command)) + return + _run(command) + if not _base_mujoco_ready(): + raise RuntimeError(f"MuJoCo {BASE_MUJOCO_VERSION} was not installed in the base site-packages") + + +def _load_json(path: Path, label: str) -> dict: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise RuntimeError(f"cannot read {label} {path}: {exc}") from exc + if not isinstance(value, dict): + raise RuntimeError(f"{label} must contain a JSON object: {path}") + return value + + +def _check_static_inputs(args: argparse.Namespace) -> None: + model = _resolved(args.model_path) + required_model_files = ( + model / "model.safetensors", + model / "config.json", + model / "assets/paligemma_tokenizer.model", + model / "assets/physical-intelligence/libero/norm_stats.json", + ) + missing = [str(path) for path in required_model_files if not path.is_file()] + if missing: + raise RuntimeError("missing checkpoint artifacts:\n- " + "\n- ".join(missing)) + checkpoint_config = _load_json(model / "config.json", "checkpoint config") + precision = checkpoint_config.get("precision") + expected_dtype = {"float32": "F32", "bfloat16": "BF16"}.get(precision) + if expected_dtype is None: + raise RuntimeError(f"checkpoint precision must be float32 or bfloat16, got {precision!r}") + try: + from safetensors import safe_open + except ImportError as exc: + raise RuntimeError("base environment is missing safetensors") from exc + tensor_dtypes: dict[str, str] = {} + with safe_open(model / "model.safetensors", framework="pt", device="cpu") as checkpoint: + for name in checkpoint.keys(): + dtype = checkpoint.get_slice(name).get_dtype() + tensor_dtypes[dtype] = tensor_dtypes.get(dtype, 0) + 1 + expected_tensors = {expected_dtype: 812} + if tensor_dtypes != expected_tensors: + raise RuntimeError(f"expected 812 {precision} checkpoint tensors, got {tensor_dtypes}") + print(f"checkpoint tensor manifest: 812/812 {expected_dtype}") + + _load_json(_resolved(args.model_config), "model config") + _load_json(_resolved(args.eval_config), "evaluation config") + + libero_root = _resolved(args.libero_root) + required_libero = ( + libero_root / "libero/libero/bddl_files", + libero_root / "libero/libero/init_files", + libero_root / "libero/libero/assets", + ) + missing = [str(path) for path in required_libero if not path.is_dir()] + if missing: + raise RuntimeError("incomplete official LIBERO checkout:\n- " + "\n- ".join(missing)) + + +COMBINED_PROBE = r""" +import importlib.metadata +import json +import sys +import sysconfig +from pathlib import Path + +transformers_root = Path(sys.argv[1]).resolve() +libero_root = Path(sys.argv[2]).resolve() +base_site = Path(sysconfig.get_paths()["purelib"]).resolve() + +import numpy +import torch +import transformers +import mujoco +import robosuite + +transformers_origin = Path(transformers.__file__).resolve() +if not transformers_origin.is_relative_to(transformers_root): + raise RuntimeError(f"transformers imported outside overlay: {transformers_origin}") +mujoco_origin = Path(mujoco.__file__).resolve() +if importlib.metadata.version("mujoco") != "3.2.3" or mujoco.__version__ != "3.2.3": + raise RuntimeError( + f"base MuJoCo must be 3.2.3, got distribution={importlib.metadata.version('mujoco')} module={mujoco.__version__}" + ) +if not mujoco_origin.is_relative_to(base_site): + raise RuntimeError(f"mujoco must be imported from base site-packages {base_site}, got {mujoco_origin}") +torch_origin = Path(torch.__file__).resolve() +if torch_origin.is_relative_to(transformers_root): + raise RuntimeError(f"torch was shadowed by an overlay: {torch_origin}") +if importlib.metadata.version("robosuite") != "1.4.1": + raise RuntimeError(f"validated base robosuite must remain 1.4.1, got {importlib.metadata.version('robosuite')}") +for module in (numpy, robosuite): + origin = Path(module.__file__).resolve() + if origin.is_relative_to(transformers_root): + raise RuntimeError(f"base package {module.__name__} was shadowed by an overlay: {origin}") + +sys.path.insert(0, str(libero_root)) +import libero +import libero.libero as libero_package +from libero.libero import benchmark + +namespace_paths = [str(Path(path).resolve()) for path in libero.__path__] +for module in (libero_package, benchmark): + origin = Path(module.__file__).resolve() + if not origin.is_relative_to(libero_root): + raise RuntimeError(f"{module.__name__} imported outside the official LIBERO root: {origin}") +namespace_extras = [path for path in namespace_paths if not Path(path).is_relative_to(libero_root)] + +payload = { + "python": sys.version.split()[0], + "torch": torch.__version__, + "torch_origin": str(torch_origin), + "cuda_available": torch.cuda.is_available(), + "transformers": transformers.__version__, + "numpy": numpy.__version__, + "mujoco": mujoco.__version__, + "mujoco_origin": str(mujoco_origin), + "base_site_packages": str(base_site), + "pillow": importlib.metadata.version("Pillow"), + "pyopengl": importlib.metadata.version("PyOpenGL"), + "glfw": importlib.metadata.version("glfw"), + "robosuite": importlib.metadata.version("robosuite"), + "libero_namespace": namespace_paths, + "libero_namespace_extras_ignored_by_source_guard": namespace_extras, + "libero_suites": sorted(benchmark.get_benchmark_dict()), +} +print(json.dumps(payload, sort_keys=True)) +if sys.argv[3] == "1" and not payload["cuda_available"]: + raise RuntimeError("CUDA is not available to the base interpreter") +""" + + +def _check_runtime(args: argparse.Namespace) -> None: + expected_python = _resolved(args.expected_python) + if _resolved(sys.executable) != expected_python: + raise RuntimeError(f"runtime check must use {expected_python}, got {_resolved(sys.executable)}") + + transformers_runtime = _resolved(args.transformers_runtime) + _probe_overlay(transformers_runtime, TRANSFORMERS_EXPECTED) + _check_patch_overlay(transformers_runtime) + _check_static_inputs(args) + + env = os.environ.copy() + with tempfile.TemporaryDirectory(prefix="lightx2v-openpi-runtime-check-") as cache_dir: + libero_config_dir = Path(cache_dir) / "libero_config" + libero_config_dir.mkdir() + benchmark_root = _resolved(args.libero_root) / "libero/libero" + (libero_config_dir / "config.yaml").write_text( + "\n".join( + ( + f"benchmark_root: {benchmark_root}", + f"bddl_files: {benchmark_root / 'bddl_files'}", + f"init_states: {benchmark_root / 'init_files'}", + f"datasets: {_resolved(args.libero_root) / 'libero/datasets'}", + f"assets: {benchmark_root / 'assets'}", + "", + ) + ), + encoding="utf-8", + ) + env.update( + { + "PYTHONPATH": os.pathsep.join((str(transformers_runtime), str(PROJECT_ROOT))), + "PYTHONDONTWRITEBYTECODE": "1", + "PYTHONNOUSERSITE": "1", + "MUJOCO_GL": "disable", + "LIBERO_CONFIG_PATH": str(libero_config_dir), + "NUMBA_CACHE_DIR": cache_dir, + "USE_FLAX": "0", + "TOKENIZERS_PARALLELISM": "false", + } + ) + command = [ + sys.executable, + "-c", + COMBINED_PROBE, + str(transformers_runtime), + str(_resolved(args.libero_root)), + "0" if args.no_cuda else "1", + ] + completed = subprocess.run(command, env=env, text=True, capture_output=True) + if completed.returncode != 0: + if completed.stdout: + print(completed.stdout, end="", file=sys.stderr) + if completed.stderr: + print(completed.stderr, end="", file=sys.stderr) + raise RuntimeError(f"combined OpenPI runtime probe failed with exit code {completed.returncode}") + if completed.stderr: + print(completed.stderr, end="", file=sys.stderr) + print(completed.stdout.strip()) + print("OpenPI runtime check: OK") + + +def _add_paths(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--transformers-runtime", + default=os.environ.get("OPENPI_TRANSFORMERS_RUNTIME_PATH", str(DEFAULT_TRANSFORMERS_RUNTIME)), + ) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(dest="command", required=True) + + prepare = subparsers.add_parser("prepare", help="prepare patched Transformers and base MuJoCo runtimes") + _add_paths(prepare) + prepare.add_argument("--component", choices=("all", "transformers", "mujoco"), default="all") + prepare.add_argument("--dry-run", action="store_true") + + check = subparsers.add_parser("check", help="validate runtime, checkpoint, configs, and official LIBERO paths") + _add_paths(check) + check.add_argument("--expected-python", default=os.environ.get("OPENPI_PYTHON", str(DEFAULT_PYTHON))) + check.add_argument("--model-path", default=os.environ.get("OPENPI_MODEL_PATH", str(DEFAULT_MODEL))) + check.add_argument("--model-config", default=os.environ.get("OPENPI_CONFIG", str(DEFAULT_MODEL_CONFIG))) + check.add_argument("--eval-config", default=os.environ.get("OPENPI_EVAL_CONFIG", str(DEFAULT_EVAL_CONFIG))) + check.add_argument("--libero-root", default=os.environ.get("OPENPI_LIBERO_ROOT", str(DEFAULT_LIBERO_ROOT))) + check.add_argument("--no-cuda", action="store_true", help="allow validation on a host without a visible CUDA device") + return parser + + +def main() -> int: + args = build_parser().parse_args() + try: + if args.command == "prepare": + if args.component in {"all", "transformers"}: + _prepare_transformers(_resolved(args.transformers_runtime), args.dry_run) + if args.component in {"all", "mujoco"}: + _prepare_base_mujoco(args.dry_run) + else: + _check_runtime(args) + except (OSError, RuntimeError, subprocess.CalledProcessError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/openpi/tests/test_task_inputs_manifest.py b/scripts/openpi/tests/test_task_inputs_manifest.py new file mode 100644 index 000000000..2de820865 --- /dev/null +++ b/scripts/openpi/tests/test_task_inputs_manifest.py @@ -0,0 +1,166 @@ +#!/usr/bin/env python3 +"""Focused tests for the LIBERO task-input manifest migration.""" + +from __future__ import annotations + +import hashlib +import json +import sys +import tempfile +import types +import unittest +from pathlib import Path +from types import SimpleNamespace + +PROJECT_ROOT = Path(__file__).resolve().parents[3] + + +def _namespace_package(name: str, path: Path) -> None: + package = types.ModuleType(name) + package.__path__ = [str(path)] + sys.modules[name] = package + + +# Import only the evaluator protocol modules. Importing the public lightx2v +# package initializes the selected accelerator, which is unrelated to these +# CPU-only filesystem tests. +_namespace_package("lightx2v", PROJECT_ROOT / "lightx2v") +_namespace_package("lightx2v.models", PROJECT_ROOT / "lightx2v/models") +_namespace_package("lightx2v.models.runners", PROJECT_ROOT / "lightx2v/models/runners") +_namespace_package("lightx2v.models.runners.openpi", PROJECT_ROOT / "lightx2v/models/runners/openpi") + +from lightx2v.models.runners.openpi.libero_protocol import ( # noqa: E402 + EvaluationConfig, + TaskSpec, + build_task_inputs_manifest, + ensure_task_inputs_manifest, + resolved_protocol, +) + + +class TaskInputsManifestTest(unittest.TestCase): + def setUp(self) -> None: + self.temporary_directory = tempfile.TemporaryDirectory() + self.root = Path(self.temporary_directory.name).resolve() + self.bddl_path = self.root / "libero/libero/bddl_files/suite/task.bddl" + self.init_states_path = self.root / "libero/libero/init_files/suite/task.pruned_init" + self.bddl_path.parent.mkdir(parents=True) + self.init_states_path.parent.mkdir(parents=True) + self.bddl_path.write_bytes(b"bddl-v1") + self.init_states_path.write_bytes(b"init-v1") + self.spec = TaskSpec( + benchmark="suite", + task_id=0, + suite=None, + task=None, + bddl_path=self.bddl_path.resolve(), + init_states_path=self.init_states_path.resolve(), + ) + self.task_specs = {"suite": [self.spec]} + self.manifest_config = SimpleNamespace(benchmarks=("suite",), libero_root=self.root) + + def tearDown(self) -> None: + self.temporary_directory.cleanup() + + def _legacy_record(self, **updates: object) -> dict[str, object]: + record: dict[str, object] = { + "schema_version": 1, + "bddl_file": str(self.bddl_path), + "init_states_file": str(self.init_states_path), + "init_states_loader": "direct_file", + } + record.update(updates) + return record + + def test_new_manifest_is_stable_and_content_sensitive(self) -> None: + first = build_task_inputs_manifest(self.task_specs, self.manifest_config) + second = build_task_inputs_manifest(self.task_specs, self.manifest_config) + self.assertEqual(first, second) + self.assertEqual(first["task_count"], 1) + self.assertEqual(first["input_count"], 2) + + self.bddl_path.write_bytes(b"bddl-v2") + changed = build_task_inputs_manifest(self.task_specs, self.manifest_config) + self.assertNotEqual(first["manifest_sha256"], changed["manifest_sha256"]) + + def test_existing_manifest_is_verified_strictly(self) -> None: + output_dir = self.root / "new-output" + output_dir.mkdir() + manifest = build_task_inputs_manifest(self.task_specs, self.manifest_config) + path = ensure_task_inputs_manifest(output_dir, manifest, {}, self.task_specs) + self.assertEqual(json.loads(path.read_text(encoding="utf-8")), manifest) + self.assertEqual(ensure_task_inputs_manifest(output_dir, manifest, {}, self.task_specs), path) + + self.init_states_path.write_bytes(b"init-v2") + changed = build_task_inputs_manifest(self.task_specs, self.manifest_config) + with self.assertRaisesRegex(RuntimeError, "content differs"): + ensure_task_inputs_manifest(output_dir, changed, {}, self.task_specs) + + def test_schema_one_records_are_adopted_after_path_validation(self) -> None: + output_dir = self.root / "legacy-output" + output_dir.mkdir() + manifest = build_task_inputs_manifest(self.task_specs, self.manifest_config) + records = {("suite", 0, 0): self._legacy_record()} + with self.assertLogs("lightx2v.models.runners.openpi.libero_protocol", level="WARNING"): + path = ensure_task_inputs_manifest(output_dir, manifest, records, self.task_specs) + self.assertTrue(path.is_file()) + + invalid_output = self.root / "invalid-legacy-output" + invalid_output.mkdir() + invalid_records = {("suite", 0, 0): self._legacy_record(bddl_file=str(self.root / "other.bddl"))} + with self.assertRaisesRegex(RuntimeError, "legacy episode"): + ensure_task_inputs_manifest(invalid_output, manifest, invalid_records, self.task_specs) + self.assertFalse((invalid_output / "task_inputs_manifest.json").exists()) + + def test_protocol_id_retains_the_schema_one_field_set(self) -> None: + model_path = self.root / "model" + norm_path = model_path / "assets/physical-intelligence/libero/norm_stats.json" + tokenizer_path = model_path / "assets/paligemma_tokenizer.model" + norm_path.parent.mkdir(parents=True) + (model_path / "model.safetensors").write_bytes(b"model") + norm_path.write_bytes(b"norm") + tokenizer_path.write_bytes(b"tokenizer") + config_json = self.root / "model.json" + config_json.write_bytes(b"{}") + config = EvaluationConfig( + benchmarks=("suite",), + task_ids={"suite": (0,)}, + num_trials_per_task=1, + env_seed=7, + policy_seed=0, + actions_per_plan=5, + num_steps_wait=10, + render_size=256, + video_fps=10, + video_policy="none", + save_actions=False, + fail_fast=False, + resume=True, + max_steps={"suite": 2}, + libero_root=self.root, + libero_config_dir=self.root / "runtime", + ) + resolved, protocol_id = resolved_protocol(config, model_path=model_path, config_json=config_json) + protocol_fields = { + key: value + for key, value in resolved.items() + if key + not in { + "protocol_id", + "protocol_name", + "official_protocol", + "libero_config_dir", + "video_fps", + "video_policy", + "save_actions", + "fail_fast", + "resume", + } + } + expected = hashlib.sha256(json.dumps(protocol_fields, sort_keys=True, separators=(",", ":")).encode("utf-8")).hexdigest() + self.assertEqual(protocol_id, expected) + self.assertNotIn("task_inputs_manifest", resolved) + + +if __name__ == "__main__": + unittest.main() diff --git a/scripts/openpi/tests/validate_pytorch_parity.py b/scripts/openpi/tests/validate_pytorch_parity.py new file mode 100644 index 000000000..066aef94b --- /dev/null +++ b/scripts/openpi/tests/validate_pytorch_parity.py @@ -0,0 +1,375 @@ +"""Compare upstream OpenPI PyTorch and LightX2V with identical input and noise. + +Run this with OpenPI's patched conversion environment. The deployed LightX2V +runtime itself does not import OpenPI or JAX. +""" + +from __future__ import annotations + +import argparse +import dataclasses +import gc +import json +import sys +import types +from collections import Counter +from pathlib import Path + +import h5py +import numpy as np +import torch + +PROJECT_ROOT = Path(__file__).resolve().parents[3] +WORKSPACE_ROOT = PROJECT_ROOT.parent +OPENPI_DATA_ROOT = WORKSPACE_ROOT / "openpi_data" +PRECISIONS = ("bfloat16", "float32") + + +def load_sample(path: Path) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + with h5py.File(path, "r") as handle: + demo = handle["data/demo_0"] + image = np.asarray(demo["obs/agentview_rgb"][0], dtype=np.uint8) + wrist = np.asarray(demo["obs/eye_in_hand_rgb"][0], dtype=np.uint8) + state = np.concatenate([demo["obs/ee_pos"][0], demo["obs/ee_ori"][0], demo["obs/gripper_states"][0]]) + return image, wrist, state + + +def _install_lightx2v_package_stub() -> None: + """Avoid LightX2V platform initialization in this model-only validator.""" + package = types.ModuleType("lightx2v") + package.__path__ = [str(PROJECT_ROOT / "lightx2v")] + sys.modules["lightx2v"] = package + + +def _create_upstream_policy(train_config, checkpoint: Path, device: str, precision: str): + """Create the official policy while honoring its selectable torch precision. + + ``policy_config.create_trained_policy`` in OpenPI 15a9616 always applies + BF16 after loading. This is the same construction with that one choice made + explicit, so the validator can cover both official model modes. + """ + from openpi import transforms + from openpi.policies import policy as policy_module + from openpi.training import checkpoints + + weight_path = checkpoint / "model.safetensors" + model = train_config.model.load_pytorch(train_config, str(weight_path)) + model.paligemma_with_expert.to_bfloat16_for_selected_params(precision) + + data_config = train_config.data.create(train_config.assets_dirs, train_config.model) + if data_config.asset_id is None: + raise ValueError("The upstream LIBERO data config has no normalization asset id") + norm_stats = checkpoints.load_norm_stats(checkpoint / "assets", data_config.asset_id) + + return policy_module.Policy( + model, + transforms=[ + transforms.InjectDefaultPrompt(None), + *data_config.data_transforms.inputs, + transforms.Normalize(norm_stats, use_quantiles=data_config.use_quantile_norm), + *data_config.model_transforms.inputs, + ], + output_transforms=[ + *data_config.model_transforms.outputs, + transforms.Unnormalize(norm_stats, use_quantiles=data_config.use_quantile_norm), + *data_config.data_transforms.outputs, + ], + metadata=train_config.policy_metadata, + is_pytorch=True, + pytorch_device=device, + ) + + +def _parameter_dtype_counts(model: torch.nn.Module) -> dict[str, int]: + counts = Counter(str(parameter.dtype).removeprefix("torch.") for parameter in model.parameters() if parameter.is_floating_point()) + return dict(sorted(counts.items())) + + +def _validate_selected_precision(counts: dict[str, int], precision: str, label: str) -> None: + if precision == "float32": + if set(counts) != {"float32"}: + raise RuntimeError(f"{label} FP32 mode contains unexpected parameter dtypes: {counts}") + return + if "bfloat16" not in counts or "float32" not in counts: + raise RuntimeError(f"{label} BF16 mode must retain the official selected FP32 parameters: {counts}") + + +def _to_numpy(value) -> np.ndarray: + if isinstance(value, torch.Tensor): + return value.detach().cpu().numpy() + return np.asarray(value) + + +def _array_comparison(reference, candidate) -> dict[str, object]: + reference_array = _to_numpy(reference) + candidate_array = _to_numpy(candidate) + same_shape = reference_array.shape == candidate_array.shape + if same_shape and reference_array.size: + difference = reference_array.astype(np.float64) - candidate_array.astype(np.float64) + max_abs_error = float(np.max(np.abs(difference))) + else: + max_abs_error = 0.0 if same_shape else None + return { + "upstream_shape": list(reference_array.shape), + "lightx2v_shape": list(candidate_array.shape), + "upstream_dtype": str(reference_array.dtype), + "lightx2v_dtype": str(candidate_array.dtype), + "exact": bool(same_shape and np.array_equal(reference_array, candidate_array)), + "max_abs_error": max_abs_error, + } + + +def _observation_comparison(upstream, local) -> dict[str, object]: + components = { + **{f"images/{key}": _array_comparison(upstream.images[key], local.images[key]) for key in upstream.images}, + **{f"image_masks/{key}": _array_comparison(upstream.image_masks[key], local.image_masks[key]) for key in upstream.image_masks}, + "state": _array_comparison(upstream.state, local.state), + "tokenized_prompt": _array_comparison(upstream.tokenized_prompt, local.tokenized_prompt), + "tokenized_prompt_mask": _array_comparison( + upstream.tokenized_prompt_mask, + local.tokenized_prompt_mask, + ), + } + expected_image_keys = set(local.images) + expected_mask_keys = set(local.image_masks) + keys_match = set(upstream.images) == expected_image_keys and set(upstream.image_masks) == expected_mask_keys + exact = keys_match and all(bool(component["exact"]) for component in components.values()) + return { + "exact": exact, + "image_keys_match": keys_match, + "components": components, + } + + +def _to_torch_batch(data, device: str): + if isinstance(data, dict): + return {key: _to_torch_batch(value, device) for key, value in data.items()} + return torch.from_numpy(np.array(data)).to(device)[None, ...] + + +def run_self_check() -> None: + """Exercise precision selection and the exact uint8 resize adapter.""" + from openpi_client import image_tools as upstream_image_tools + + _install_lightx2v_package_stub() + from lightx2v.models.networks.openpi.config import Pi0Config + from lightx2v.models.networks.openpi.infer.pre_infer import _resize_with_pad + + values = { + "action_dim": 32, + "action_horizon": 10, + "max_token_len": 200, + "paligemma_variant": "gemma_2b", + "action_expert_variant": "gemma_300m", + "pi05": True, + "discrete_state_input": False, + "pytorch_compile_mode": None, + } + for precision in PRECISIONS: + config = Pi0Config.from_mapping({**values, "dtype": precision}) + config.validate_pi05_libero() + if config.dtype != precision: + raise AssertionError(f"precision selection changed {precision!r} to {config.dtype!r}") + + image = np.arange(17 * 29 * 3, dtype=np.uint8).reshape(17, 29, 3) + upstream = upstream_image_tools.resize_with_pad(image, 24, 24) + local = _resize_with_pad(image, 24) + if not np.array_equal(local, upstream): + difference = np.abs(local.astype(np.int16) - upstream.astype(np.int16)) + raise AssertionError(f"resize-with-pad mismatch: max_abs={difference.max()}") + + print(json.dumps({"precision_selection": list(PRECISIONS), "uint8_resize_exact": True}, indent=2)) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser() + parser.add_argument( + "--checkpoint", + type=Path, + default=OPENPI_DATA_ROOT / "openpi-assets/checkpoints/pi05_libero_pytorch_fp32", + ) + parser.add_argument( + "--config", + type=Path, + default=PROJECT_ROOT / "configs/openpi/pi05_libero.json", + ) + parser.add_argument( + "--sample", + type=Path, + default=OPENPI_DATA_ROOT / "raw/huggingface/yifengzhu-hf/LIBERO-datasets/libero_spatial" / "pick_up_the_black_bowl_between_the_plate_and_the_ramekin_and_place_it_on_the_plate_demo.hdf5", + ) + parser.add_argument( + "--output", + type=Path, + default=OPENPI_DATA_ROOT / "results/pi05_libero_pytorch_parity.json", + ) + parser.add_argument("--device", default="cuda") + parser.add_argument("--precision", choices=PRECISIONS, help="Override the dtype selected by the model JSON") + parser.add_argument("--atol", type=float, default=1e-6) + parser.add_argument("--self-check", action="store_true", help="Run lightweight checks without loading a checkpoint") + return parser + + +def main() -> None: + args = build_parser().parse_args() + if args.self_check: + run_self_check() + return + if args.atol < 0: + raise ValueError("--atol must be non-negative") + + checkpoint = args.checkpoint.expanduser().resolve() + config_path = args.config.expanduser().resolve() + sample_path = args.sample.expanduser().resolve() + output_path = args.output.expanduser().resolve() + required = ( + checkpoint / "model.safetensors", + checkpoint / "assets/paligemma_tokenizer.model", + checkpoint / "assets/physical-intelligence/libero/norm_stats.json", + config_path, + sample_path, + ) + missing = [str(path) for path in required if not path.is_file()] + if missing: + raise FileNotFoundError(f"Parity inputs are missing: {missing}") + + with config_path.open("r", encoding="utf-8") as handle: + local_config = json.load(handle) + precision = args.precision or local_config["dtype"] + if precision not in PRECISIONS: + raise ValueError(f"Unsupported OpenPI precision {precision!r}; expected one of {PRECISIONS}") + local_config["dtype"] = precision + + image, wrist, state = load_sample(sample_path) + # Match the official LIBERO evaluation client: renderer frames are resized + # with PIL before they enter either policy. Both model adapters therefore + # receive the exact same 224x224 uint8 arrays. + from openpi_client import image_tools as client_image_tools + + image = client_image_tools.resize_with_pad(image, 224, 224) + wrist = client_image_tools.resize_with_pad(wrist, 224, 224) + prompt = "pick up the black bowl between the plate and the ramekin and place it on the plate" + noise = np.random.default_rng(0).standard_normal((10, 32)).astype(np.float32) + + from openpi.training import config as training_config + + upstream_config = training_config.get_config("pi05_libero") + upstream_config = dataclasses.replace( + upstream_config, + model=dataclasses.replace( + upstream_config.model, + dtype=precision, + pytorch_compile_mode=None, + ), + ) + upstream_policy = _create_upstream_policy(upstream_config, checkpoint, args.device, precision) + upstream_dtype_counts = _parameter_dtype_counts(upstream_policy._model) # noqa: SLF001 + _validate_selected_precision(upstream_dtype_counts, precision, "upstream") + raw_input = { + "observation/image": image.copy(), + "observation/wrist_image": wrist.copy(), + "observation/state": state.copy(), + "prompt": prompt, + } + upstream_inputs = upstream_policy._input_transform(raw_input) # noqa: SLF001 + upstream_torch_inputs = _to_torch_batch(upstream_inputs, args.device) + from openpi.models import model as upstream_model_module + + upstream_observation = upstream_model_module.Observation.from_dict(upstream_torch_inputs) + upstream_noise = torch.from_numpy(noise).to(args.device)[None, ...] + upstream_normalized = upstream_policy._sample_actions( # noqa: SLF001 + args.device, + upstream_observation, + noise=upstream_noise, + ) + upstream_outputs = { + "state": _to_numpy(upstream_torch_inputs["state"][0]), + "actions": _to_numpy(upstream_normalized[0]), + } + upstream_actions = np.asarray(upstream_policy._output_transform(upstream_outputs)["actions"]) # noqa: SLF001 + upstream_normalized_array = _to_numpy(upstream_normalized) + del upstream_policy, upstream_normalized, upstream_noise, upstream_torch_inputs + gc.collect() + if args.device.startswith("cuda"): + torch.cuda.empty_cache() + + _install_lightx2v_package_stub() + from lightx2v.models.networks.openpi import OpenPIModel + + local_config["model_path"] = str(checkpoint) + local_config["device"] = args.device + local_config["seed"] = 0 + model = OpenPIModel.from_config(local_config) + local_dtype_counts = _parameter_dtype_counts(model.core_model) + _validate_selected_precision(local_dtype_counts, precision, "LightX2V") + local_observation = model.pre_infer.infer( + images={"agentview": image, "wrist": wrist}, + state=state, + task_description=prompt, + ) + preprocessing = _observation_comparison(upstream_observation, local_observation) + local_normalized = model.transformer_infer.infer( + model.core_model, + local_observation, + model.device, + noise=torch.from_numpy(noise).to(model.device)[None, ...], + ) + local_actions = model.post_infer.infer(local_normalized) + + normalized_difference = _to_numpy(local_normalized).astype(np.float64) - upstream_normalized_array.astype(np.float64) + normalized_max_abs_error = float(np.max(np.abs(normalized_difference))) + normalized_passed = bool(np.allclose(_to_numpy(local_normalized), upstream_normalized_array, rtol=0.0, atol=args.atol)) + physical_difference = np.asarray(local_actions, dtype=np.float64) - np.asarray(upstream_actions, dtype=np.float64) + physical_max_abs_error = float(np.max(np.abs(physical_difference))) + physical_passed = bool(np.allclose(local_actions, upstream_actions, rtol=0.0, atol=args.atol)) + passed = bool(preprocessing["exact"] and normalized_passed and physical_passed) + report = { + "precision": precision, + "atol": args.atol, + "upstream_parameter_dtypes": upstream_dtype_counts, + "lightx2v_parameter_dtypes": local_dtype_counts, + "preprocessing": preprocessing, + "normalized_actions": { + "upstream_shape": list(upstream_normalized_array.shape), + "lightx2v_shape": list(local_normalized.shape), + "max_abs_error": normalized_max_abs_error, + "mean_abs_error": float(np.mean(np.abs(normalized_difference))), + "allclose": normalized_passed, + }, + "physical_actions": { + "upstream_shape": list(upstream_actions.shape), + "lightx2v_shape": list(local_actions.shape), + "max_abs_error": physical_max_abs_error, + "mean_abs_error": float(np.mean(np.abs(physical_difference))), + "allclose": physical_passed, + }, + "allclose": passed, + "upstream_normalized_actions": upstream_normalized_array[0].tolist(), + "lightx2v_normalized_actions": _to_numpy(local_normalized[0]).tolist(), + "upstream_actions": upstream_actions.tolist(), + "lightx2v_actions": local_actions.tolist(), + } + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") + raw_action_keys = { + "upstream_normalized_actions", + "lightx2v_normalized_actions", + "upstream_actions", + "lightx2v_actions", + } + print(json.dumps({key: value for key, value in report.items() if key not in raw_action_keys}, indent=2)) + del model + gc.collect() + if args.device.startswith("cuda"): + torch.cuda.empty_cache() + if not passed: + raise SystemExit( + "OpenPI PyTorch parity check failed: " + f"preprocessing_exact={preprocessing['exact']}, " + f"normalized_max_abs_error={normalized_max_abs_error}, " + f"physical_max_abs_error={physical_max_abs_error}, atol={args.atol}" + ) + + +if __name__ == "__main__": + main() From d27b5a9e7aac16967d8b0d17b1561660b06934f9 Mon Sep 17 00:00:00 2001 From: Chernobyllight Date: Fri, 4 Sep 2026 06:52:07 +0000 Subject: [PATCH 2/6] feat: add decoupled OpenPI ROS evaluation Integrate pi0.5-LIBERO inference with local PyTorch execution, official-compatible evaluation, multi-GPU suite scheduling, and decoupled ROS2 policy/simulator nodes. Simplify the integration while preserving preprocessing, FP64 actions, and continuous policy RNG semantics. --- lightx2v/models/networks/openpi/gemma.py | 4 +- lightx2v/models/networks/openpi/model.py | 3 - lightx2v/models/networks/openpi/pi0.py | 9 - .../models/networks/openpi/preprocessing.py | 8 +- .../models/runners/openpi/libero_evaluate.py | 6 +- .../models/runners/openpi/libero_protocol.py | 36 +- .../models/runners/openpi/libero_rollout.py | 5 +- .../models/runners/openpi/openpi_runner.py | 37 +- lightx2v_ros/src/common/common/contract.py | 4 + .../inference/openpi_node/__init__.py | 0 .../inference/inference/openpi_node/main.py | 231 ++++++++++ lightx2v_ros/src/inference/setup.py | 1 + lightx2v_ros/src/simulator/setup.py | 1 + .../simulator/simulator/libero_node/env.py | 54 ++- .../simulator/libero_node/evaluate.py | 250 +++++++++++ .../simulator/libero_node/observer.py | 43 +- .../src/simulator/simulator/sim/node.py | 75 +++- scripts/openpi/README.md | 233 +++++++++- scripts/openpi/convert_jax_checkpoint.py | 2 +- scripts/openpi/libero_summary.py | 36 +- .../run_libero_evaluate_parallel_i2va.sh | 31 +- scripts/openpi/run_libero_ros_i2va.sh | 75 ++++ scripts/openpi/runtime.py | 4 - .../openpi/tests/test_ros_openpi_contract.py | 418 ++++++++++++++++++ 24 files changed, 1379 insertions(+), 187 deletions(-) create mode 100644 lightx2v_ros/src/inference/inference/openpi_node/__init__.py create mode 100644 lightx2v_ros/src/inference/inference/openpi_node/main.py create mode 100644 lightx2v_ros/src/simulator/simulator/libero_node/evaluate.py create mode 100755 scripts/openpi/run_libero_ros_i2va.sh create mode 100644 scripts/openpi/tests/test_ros_openpi_contract.py diff --git a/lightx2v/models/networks/openpi/gemma.py b/lightx2v/models/networks/openpi/gemma.py index 21c1a05ea..e4181badf 100644 --- a/lightx2v/models/networks/openpi/gemma.py +++ b/lightx2v/models/networks/openpi/gemma.py @@ -107,7 +107,7 @@ def forward( position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, - adarms_cond=adarms_cond[0] if adarms_cond is not None else None, + adarms_cond=adarms_cond[0], ) prefix_past_key_values = prefix_output.past_key_values prefix_output = prefix_output.last_hidden_state @@ -119,7 +119,7 @@ def forward( position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, - adarms_cond=adarms_cond[1] if adarms_cond is not None else None, + adarms_cond=adarms_cond[1], ) suffix_output = suffix_output.last_hidden_state prefix_output = None diff --git a/lightx2v/models/networks/openpi/model.py b/lightx2v/models/networks/openpi/model.py index 89bc16759..0dc81528a 100644 --- a/lightx2v/models/networks/openpi/model.py +++ b/lightx2v/models/networks/openpi/model.py @@ -50,9 +50,6 @@ def from_config(cls, config: Mapping[str, Any]) -> "OpenPIModel": tokenizer_path = checkpoint_dir / "assets/paligemma_tokenizer.model" device = torch.device(values["device"]) - if device.type == "cuda" and not torch.cuda.is_available(): - raise RuntimeError("OpenPI config requests CUDA, but torch.cuda.is_available() is false") - core_model = load_pi05_libero_weights(weight_path, model_config, device) return cls( core_model=core_model, diff --git a/lightx2v/models/networks/openpi/pi0.py b/lightx2v/models/networks/openpi/pi0.py index 8513da856..37180fcaa 100644 --- a/lightx2v/models/networks/openpi/pi0.py +++ b/lightx2v/models/networks/openpi/pi0.py @@ -111,15 +111,6 @@ def __init__(self, config): self.gradient_checkpointing_enabled = False - msg = "OpenPI's patched transformers==4.53.2 runtime is not active. Run scripts/openpi/2_setup_pytorch_runtime.sh and prepend OPENPI_TRANSFORMERS_RUNTIME_PATH to PYTHONPATH." - try: - from transformers.models.siglip import check - - if not check.check_whether_transformers_replace_is_installed_correctly(): - raise ValueError(msg) - except ImportError: - raise ValueError(msg) from None - def gradient_checkpointing_enable(self): self.gradient_checkpointing_enabled = True self.paligemma_with_expert.paligemma.language_model.gradient_checkpointing = True diff --git a/lightx2v/models/networks/openpi/preprocessing.py b/lightx2v/models/networks/openpi/preprocessing.py index 07771a5fa..377d83f3c 100644 --- a/lightx2v/models/networks/openpi/preprocessing.py +++ b/lightx2v/models/networks/openpi/preprocessing.py @@ -7,6 +7,7 @@ import torch from . import image_tools +from .observation import Observation logger = logging.getLogger("openpi") @@ -114,12 +115,7 @@ def preprocess_observation_pytorch( else: out_masks[key] = observation.image_masks[key] - class SimpleProcessedObservation: - def __init__(self, **kwargs): - for key, value in kwargs.items(): - setattr(self, key, value) - - return SimpleProcessedObservation( + return Observation( images=out_images, image_masks=out_masks, state=observation.state, diff --git a/lightx2v/models/runners/openpi/libero_evaluate.py b/lightx2v/models/runners/openpi/libero_evaluate.py index d3480d864..ccd913814 100644 --- a/lightx2v/models/runners/openpi/libero_evaluate.py +++ b/lightx2v/models/runners/openpi/libero_evaluate.py @@ -204,7 +204,6 @@ def run_evaluation(args: argparse.Namespace) -> dict[str, Any]: Path(args.config_json), Path(args.model_path), seed=config.policy_seed, - actions_per_plan=config.actions_per_plan, ) ) global_indices = {key: index for index, key in enumerate(all_keys)} @@ -216,8 +215,7 @@ def run_evaluation(args: argparse.Namespace) -> dict[str, Any]: LOGGER.info("Episode resume: skipping complete suite %s", benchmark) continue - policy.clear_action_queue() - policy.reset_rng() + policy.reset() # Official evaluation starts each suite in its own process. Mirror # that process-level NumPy seed when several suites share a worker. np.random.seed(config.env_seed) @@ -252,6 +250,7 @@ def run_evaluation(args: argparse.Namespace) -> dict[str, Any]: task_description=str(spec.task.language), max_steps=config.max_steps[benchmark], num_steps_wait=config.num_steps_wait, + actions_per_plan=config.actions_per_plan, collect_frames=config.video_policy != "none", ) record = { @@ -300,7 +299,6 @@ def run_evaluation(args: argparse.Namespace) -> dict[str, Any]: finally: env.close() finally: - policy.close() summary = _write_current_outputs(output_dir, records, task_specs, config, protocol_id) LOGGER.info( diff --git a/lightx2v/models/runners/openpi/libero_protocol.py b/lightx2v/models/runners/openpi/libero_protocol.py index af2cce503..626ba7ded 100644 --- a/lightx2v/models/runners/openpi/libero_protocol.py +++ b/lightx2v/models/runners/openpi/libero_protocol.py @@ -12,6 +12,7 @@ import sys import time import traceback +from collections import deque from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path @@ -391,11 +392,12 @@ def resolved_protocol( } encoded = json.dumps(protocol_fields, sort_keys=True, separators=(",", ":")).encode("utf-8") protocol_id = hashlib.sha256(encoded).hexdigest() + official_protocol = is_official_protocol(config) resolved = { **protocol_fields, "protocol_id": protocol_id, - "protocol_name": "official_pi05_libero" if is_official_protocol(config) else "custom_pi05_libero", - "official_protocol": is_official_protocol(config), + "protocol_name": "official_pi05_libero" if official_protocol else "custom_pi05_libero", + "official_protocol": official_protocol, "libero_config_dir": str(config.libero_config_dir), "video_fps": config.video_fps, "video_policy": config.video_policy, @@ -411,7 +413,6 @@ def load_policy_config( model_path: Path, *, seed: int, - actions_per_plan: int, ) -> dict[str, Any]: config_json = config_json.expanduser().resolve() model_path = model_path.expanduser().resolve() @@ -426,7 +427,6 @@ def load_policy_config( "model_path": str(model_path), "config_json": str(config_json), "seed": int(seed), - "actions_per_plan": int(actions_per_plan), } ) return values @@ -453,7 +453,6 @@ def _assert_module_source(module: Any, libero_root: Path) -> None: def _constrain_libero_namespace(package: Any, libero_root: Path) -> None: """Restrict the top-level namespace before resolving ``libero.libero``.""" if getattr(package, "__file__", None): - _assert_module_source(package, libero_root) return requested = (libero_root / "libero").resolve() @@ -464,7 +463,6 @@ def _constrain_libero_namespace(package: Any, libero_root: Path) -> None: spec = getattr(package, "__spec__", None) if spec is not None: spec.submodule_search_locations = package.__path__ - _assert_module_source(package, libero_root) def configure_libero(libero_root: Path, config_dir: Path) -> LiberoRuntime: @@ -501,6 +499,9 @@ def configure_libero(libero_root: Path, config_dir: Path) -> LiberoRuntime: nested_package = importlib.import_module("libero.libero") benchmark_module = importlib.import_module("libero.libero.benchmark") envs_module = importlib.import_module("libero.libero.envs") + # robosuite 1.4.1 reads the physical CUDA mask as an EGL ordinal. Import it + # with the caller's mask intact, then select logical device 0 within it. + os.environ["MUJOCO_EGL_DEVICE_ID"] = "0" for module in (package, nested_package, benchmark_module, envs_module): _assert_module_source(module, libero_root) @@ -637,6 +638,7 @@ def run_episode( task_description: str, max_steps: int, num_steps_wait: int, + actions_per_plan: int, collect_frames: bool, ) -> tuple[dict[str, Any], list[np.ndarray], np.ndarray]: """Run one episode with the exact official warmup/policy done semantics.""" @@ -650,8 +652,8 @@ def run_episode( error_message: str | None = None error_traceback: str | None = None started = time.perf_counter() + pending_actions: deque[np.ndarray] = deque() - policy.clear_action_queue() # Match the official evaluator: reset/init failures are infrastructure errors. env.reset() observation = env.set_init_state(initial_state) @@ -671,13 +673,17 @@ def run_episode( "agentview": agentview, "wrist": policy_rgb(observation, "robot0_eye_in_hand_image"), } - if policy.pending_action_count == 0: + if not pending_actions: action_chunk_calls += 1 - action = policy.next_action( - images=images, - state=state_from_observation(observation), - task_description=task_description, - ) + chunk = policy.predict_action_chunk( + images=images, + state=state_from_observation(observation), + task_description=task_description, + ) + if len(chunk) < actions_per_plan: + raise ValueError(f"Policy returned {len(chunk)} actions; {actions_per_plan} are required per plan") + pending_actions.extend(action.copy() for action in chunk[:actions_per_plan]) + action = pending_actions.popleft() action = np.asarray(action).reshape(-1) actions.append(action.copy()) observation, _reward, done, _info = env.step(action.tolist()) @@ -693,8 +699,8 @@ def run_episode( error_traceback = traceback.format_exc(limit=20) LOGGER.exception("LIBERO episode failed with an exception") finally: - pending_actions_discarded = policy.pending_action_count - policy.clear_action_queue() + pending_actions_discarded = len(pending_actions) + pending_actions.clear() action_dim = int(policy.output_action_dim) action_array = np.stack(actions).reshape(-1, action_dim) if actions else np.empty((0, action_dim), dtype=np.float64) diff --git a/lightx2v/models/runners/openpi/libero_rollout.py b/lightx2v/models/runners/openpi/libero_rollout.py index 5ac5f7359..5d4e0e536 100644 --- a/lightx2v/models/runners/openpi/libero_rollout.py +++ b/lightx2v/models/runners/openpi/libero_rollout.py @@ -79,7 +79,6 @@ def run_rollout(args: argparse.Namespace) -> dict: task_description = args.task_description.strip() or str(spec.task.language) env = create_environment(runtime, spec, args.render_size, args.env_seed) - policy = None started = time.perf_counter() try: LOGGER.info("Loading local PyTorch OpenPI policy") @@ -88,7 +87,6 @@ def run_rollout(args: argparse.Namespace) -> dict: args.config_json, args.model_path, seed=args.policy_seed, - actions_per_plan=args.actions_per_plan, ) ) policy.reset() @@ -99,13 +97,12 @@ def run_rollout(args: argparse.Namespace) -> dict: task_description=task_description, max_steps=max_steps, num_steps_wait=args.num_steps_wait, + actions_per_plan=args.actions_per_plan, collect_frames=True, ) rng_state_after = policy.export_rng_state() finally: env.close() - if policy is not None: - policy.close() artifact_errors: list[str] = [] try: diff --git a/lightx2v/models/runners/openpi/openpi_runner.py b/lightx2v/models/runners/openpi/openpi_runner.py index 2d15d6edd..cbc9edf5c 100644 --- a/lightx2v/models/runners/openpi/openpi_runner.py +++ b/lightx2v/models/runners/openpi/openpi_runner.py @@ -4,7 +4,6 @@ import os import subprocess import sys -from collections import deque from pathlib import Path from typing import Any @@ -20,19 +19,11 @@ class OpenPIPolicy: def __init__(self, config: Any): self.action_horizon = int(config["action_horizon"]) self.output_action_dim = int(config["output_action_dim"]) - self.actions_per_plan = int(config.get("actions_per_plan", 5)) - if not 1 <= self.actions_per_plan <= self.action_horizon: - raise ValueError(f"OpenPI actions_per_plan must be in [1, {self.action_horizon}], got {self.actions_per_plan}.") # Import only after the worker activates the patched Transformers path. from lightx2v.models.networks.openpi import OpenPIModel self.model = OpenPIModel.from_config(config) - self.pending_actions: deque[np.ndarray] = deque() - - @property - def pending_action_count(self) -> int: - return len(self.pending_actions) def predict_action_chunk( self, @@ -56,29 +47,14 @@ def predict_action_chunk( raise ValueError("OpenPI produced non-finite actions.") return np.ascontiguousarray(actions) - def next_action(self, images: dict[str, np.ndarray], state: np.ndarray, task_description: str) -> np.ndarray: - if not self.pending_actions: - chunk = self.predict_action_chunk(images, state, task_description, seed=None) - self.pending_actions.extend(action.copy() for action in chunk[: self.actions_per_plan]) - return self.pending_actions.popleft() - - def clear_action_queue(self) -> None: - self.pending_actions.clear() - - def reset_rng(self) -> None: - self.model.reset() - def reset(self) -> None: - self.clear_action_queue() - self.reset_rng() + self.model.reset() def export_rng_state(self) -> str: encoded = self.model.get_rng_state().cpu().numpy().tobytes() return base64.b64encode(encoded).decode("ascii") def import_rng_state(self, encoded: str) -> None: - if not encoded: - raise ValueError("OpenPI RNG state is empty") raw = base64.b64decode(encoded.encode("ascii"), validate=True) state_array = np.frombuffer(raw, dtype=np.uint8).copy() import torch @@ -86,9 +62,6 @@ def import_rng_state(self, encoded: str) -> None: state_tensor = torch.from_numpy(state_array) self.model.set_rng_state(state_tensor) - def close(self) -> None: - self.clear_action_queue() - @RUNNER_REGISTER("openpi") class OpenPIRunner(BaseRunner): @@ -171,13 +144,7 @@ def _worker_environment(self) -> dict[str, str]: child_env = os.environ.copy() child_env["USE_FLAX"] = "0" - visible_devices = [item.strip() for item in child_env.get("CUDA_VISIBLE_DEVICES", "").split(",") if item.strip()] - if len(visible_devices) == 1 and visible_devices[0] != "0": - # robosuite 1.4.1 validates physical IDs as EGL ordinals. - visible_devices.append("0") - child_env["CUDA_VISIBLE_DEVICES"] = ",".join(visible_devices) - if visible_devices: - child_env["MUJOCO_EGL_DEVICE_ID"] = "0" + child_env.pop("MUJOCO_EGL_DEVICE_ID", None) child_env["PYTHONPATH"] = os.pathsep.join((str(runtime_path), str(PROJECT_ROOT))) return child_env diff --git a/lightx2v_ros/src/common/common/contract.py b/lightx2v_ros/src/common/common/contract.py index 3810a1659..5027e5173 100644 --- a/lightx2v_ros/src/common/common/contract.py +++ b/lightx2v_ros/src/common/common/contract.py @@ -53,6 +53,10 @@ def success_topic(self) -> str: def observation_ready_topic(self) -> str: return f"{self.namespace}/observation_ready" + @property + def observation_context_topic(self) -> str: + return f"{self.namespace}/observation_context" + @property def task_topic(self) -> str: return f"{self.namespace}/task_description" diff --git a/lightx2v_ros/src/inference/inference/openpi_node/__init__.py b/lightx2v_ros/src/inference/inference/openpi_node/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/lightx2v_ros/src/inference/inference/openpi_node/main.py b/lightx2v_ros/src/inference/inference/openpi_node/main.py new file mode 100644 index 000000000..0898f5056 --- /dev/null +++ b/lightx2v_ros/src/inference/inference/openpi_node/main.py @@ -0,0 +1,231 @@ +import json +from collections import deque + +import numpy as np +import rclpy +from common.contract import get_contract +from rclpy.node import Node +from sensor_msgs.msg import Image +from std_msgs.msg import Float32MultiArray, Float64MultiArray, MultiArrayDimension, String + +from lightx2v.models.runners.openpi.openpi_runner import OpenPIPolicy +from lightx2v.utils.set_config import auto_calc_config, get_default_config + + +class OpenPINode(Node): + def __init__(self): + super().__init__("openpi_node") + + self.declare_parameter("config_json", "") + self.declare_parameter("model_path", "") + self.declare_parameter("seed", 0) + self.declare_parameter("actions_per_plan", 5) + self.declare_parameter("numeric_precision", "float64") + + self.contract = get_contract("libero") + + precision = str(self.get_parameter("numeric_precision").value).strip().lower() + numeric_types = { + "float32": (Float32MultiArray, np.float32), + "float64": (Float64MultiArray, np.float64), + } + if precision not in numeric_types: + raise ValueError("numeric_precision must be 'float32' or 'float64'") + self.numeric_message_type, self.numeric_dtype = numeric_types[precision] + + config = self._policy_config() + self.actions_per_plan = int(self.get_parameter("actions_per_plan").value) + action_horizon = int(config["action_horizon"]) + if not 1 <= self.actions_per_plan <= action_horizon: + raise ValueError(f"actions_per_plan must be in [1, {action_horizon}]") + + self.get_logger().info("loading OpenPI policy") + self.policy = OpenPIPolicy(config) + self.get_logger().info("OpenPI policy loaded") + + self.images = {camera: None for camera in self.contract.policy_input_cameras} + self.state = None + self.task_description = "" + self.episode_index = None + self.plan_epoch = None + self.pending_observation = None + self.last_processed_observation = None + self.pending_actions = deque() + + self.action_pub = self.create_publisher(self.numeric_message_type, self.contract.action_topic, 10) + self._camera_subscriptions = [self.create_subscription(Image, self.contract.camera_topic(camera), self._image_callback(camera), 10) for camera in self.contract.policy_input_cameras] + self.state_sub = self.create_subscription(self.numeric_message_type, self.contract.state_topic, self._on_state, 10) + self.context_sub = self.create_subscription(String, self.contract.observation_context_topic, self._on_context, 10) + + self.get_logger().info(f"OpenPI ready on {self.contract.namespace}: horizon={action_horizon}, actions_per_plan={self.actions_per_plan}, precision={precision}") + + def _policy_config(self): + config_json = str(self.get_parameter("config_json").value).strip() + model_path = str(self.get_parameter("model_path").value).strip() + if not config_json or not model_path: + raise ValueError("OpenPI requires model_path and config_json") + seed = int(self.get_parameter("seed").value) + + config = get_default_config() + config.update( + { + "model_cls": "openpi", + "task": "i2va", + "model_path": model_path, + "config_json": config_json, + } + ) + config = auto_calc_config(config) + # ROS parameters take precedence over values loaded from the model JSON. + config["model_path"] = model_path + config["seed"] = seed + + expected = {"state_dim": self.contract.state_dim, "output_action_dim": self.contract.action_dim} + for name, dimension in expected.items(): + if int(config[name]) != dimension: + raise ValueError(f"OpenPI {name}={config[name]} does not match the LIBERO contract ({dimension})") + return config + + def _image_callback(self, camera): + def callback(msg): + identity = observation_identity(msg.header.frame_id) + if identity is None: + return + episode, observation = identity + self.images[camera] = (episode, observation, image_msg_to_rgb(msg)) + self._try_process_observation() + + return callback + + def _on_state(self, msg): + if not msg.layout.dim: + return + identity = state_identity(msg.layout.dim[0].label) + if identity is None: + return + episode, observation = identity + state = np.asarray(msg.data, dtype=self.numeric_dtype).reshape(-1) + if state.size != self.contract.state_dim: + self.get_logger().error(f"expected state length {self.contract.state_dim}, got {state.size}") + return + self.state = (episode, observation, state) + self._try_process_observation() + + def _on_context(self, msg): + try: + context = json.loads(msg.data) + episode = int(context["episode"]) + observation = int(context["observation"]) + plan_epoch = int(context["plan_epoch"]) + task_description = str(context["task_description"]).strip() + except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc: + self.get_logger().error(f"invalid observation context: {exc}") + return + + if episode != self.episode_index: + self.episode_index = episode + self.pending_actions.clear() + self.last_processed_observation = None + self.get_logger().info(f"episode {episode}: cleared ROS action queue") + if plan_epoch != self.plan_epoch: + if self.plan_epoch is not None: + self.pending_actions.clear() + self.get_logger().info(f"plan epoch {plan_epoch}: cleared ROS action queue") + self.plan_epoch = plan_epoch + identity = (episode, observation) + if identity == self.last_processed_observation: + return + self.task_description = task_description + self.pending_observation = identity + self._try_process_observation() + + def _try_process_observation(self): + if self.pending_observation is None or not self.task_description: + return + identity = self.pending_observation + if self.state is None or self.state[:2] != identity: + return + if any(value is None or value[:2] != identity for value in self.images.values()): + return + + episode, observation = identity + if not self.pending_actions: + chunk = self.policy.predict_action_chunk( + images={camera: self.images[camera][2] for camera in self.contract.policy_input_cameras}, + state=self.state[2], + task_description=self.task_description, + ) + self.pending_actions.extend(action.copy() for action in chunk[: self.actions_per_plan]) + self.get_logger().info(f"episode {episode} observation {observation}: predicted {len(chunk)} actions") + + action = self.pending_actions.popleft() + self.last_processed_observation = identity + self.pending_observation = None + self._publish_action(action, episode, observation) + + def _publish_action(self, action, episode, observation): + action = np.asarray(action, dtype=self.numeric_dtype).reshape(-1) + if action.size != self.contract.action_dim: + raise ValueError(f"expected action length {self.contract.action_dim}, got {action.size}") + if not np.isfinite(action).all(): + raise ValueError("OpenPI produced a non-finite action") + msg = self.numeric_message_type() + msg.layout.dim = [ + MultiArrayDimension( + label=f"episode={episode};observation={observation};plan_epoch={self.plan_epoch}", + size=action.size, + stride=action.size, + ) + ] + msg.data = action.tolist() + self.action_pub.publish(msg) + + +def observation_identity(frame_id): + parts = str(frame_id).rsplit("|", 2) + if len(parts) != 3: + return None + try: + return int(parts[1]), int(parts[2]) + except ValueError: + return None + + +def state_identity(label): + fields = {} + for item in str(label).split(";"): + name, separator, value = item.partition("=") + if separator: + fields[name] = value + try: + return int(fields["episode"]), int(fields["observation"]) + except (KeyError, ValueError): + return None + + +def image_msg_to_rgb(msg): + encoding = msg.encoding.lower() + if encoding not in {"rgb8", "bgr8"}: + raise ValueError(f"unsupported image encoding: {msg.encoding}") + row = np.frombuffer(msg.data, dtype=np.uint8).reshape(msg.height, msg.step) + image = row[:, : msg.width * 3].reshape(msg.height, msg.width, 3) + if encoding == "bgr8": + image = image[:, :, ::-1] + return np.ascontiguousarray(image.copy()) + + +def main(args=None): + rclpy.init(args=args) + node = OpenPINode() + try: + rclpy.spin(node) + except KeyboardInterrupt: + pass + finally: + node.destroy_node() + if rclpy.ok(): + rclpy.shutdown() + + +if __name__ == "__main__": + main() diff --git a/lightx2v_ros/src/inference/setup.py b/lightx2v_ros/src/inference/setup.py index ce367aec1..95ca933aa 100644 --- a/lightx2v_ros/src/inference/setup.py +++ b/lightx2v_ros/src/inference/setup.py @@ -21,6 +21,7 @@ "cosmos3_node = inference.cosmos3_node.main:main", "fastwam_node = inference.fastwam_node.main:main", "lingbot_va_node = inference.lingbot_va_node.main:main", + "openpi_node = inference.openpi_node.main:main", ], }, ) diff --git a/lightx2v_ros/src/simulator/setup.py b/lightx2v_ros/src/simulator/setup.py index 38f6b2dde..7d9217b7f 100644 --- a/lightx2v_ros/src/simulator/setup.py +++ b/lightx2v_ros/src/simulator/setup.py @@ -18,6 +18,7 @@ license="Apache-2.0", entry_points={ "console_scripts": [ + "libero_evaluate = simulator.libero_node.evaluate:main", "libero_node = simulator.libero_node.main:main", "robolab_node = simulator.robolab_node.main:main", "robodojo_node = simulator.robodojo_node.main:main", diff --git a/lightx2v_ros/src/simulator/simulator/libero_node/env.py b/lightx2v_ros/src/simulator/simulator/libero_node/env.py index 8356a2c57..b648616df 100644 --- a/lightx2v_ros/src/simulator/simulator/libero_node/env.py +++ b/lightx2v_ros/src/simulator/simulator/libero_node/env.py @@ -8,14 +8,23 @@ from ..sim.base_env import BaseSimEnv, Observation from .observer import LiberoActionObserver, build_task_catalog, default_libero_root +LIBERO_DUMMY_ACTION = np.asarray([0.0] * 6 + [-1.0]) +MAX_POLICY_STEPS = { + "libero_spatial": 220, + "libero_object": 280, + "libero_goal": 300, + "libero_10": 520, + "libero_90": 400, +} + def quat_to_axis_angle(quat): - quat = np.asarray(quat, dtype=np.float32).copy() + quat = np.asarray(quat).copy() quat[3] = np.clip(quat[3], -1.0, 1.0) den = np.sqrt(1.0 - quat[3] * quat[3]) if math.isclose(float(den), 0.0): - return np.zeros(3, dtype=np.float32) - return ((quat[:3] * 2.0 * math.acos(float(quat[3]))) / den).astype(np.float32) + return np.zeros(3) + return (quat[:3] * 2.0 * math.acos(float(quat[3]))) / den class LiberoEnv(BaseSimEnv): @@ -36,11 +45,16 @@ def __init__( init_state_id=0, image_size=224, seed=0, + settle_steps=0, libero_root=None, ): super().__init__(contract) self.image_size = int(image_size) + self.settle_steps = int(settle_steps) + if self.settle_steps < 0: + raise ValueError("settle_steps must be non-negative") self.libero_root = libero_root + np.random.seed(int(seed)) self.observer = LiberoActionObserver( benchmark_name=benchmark, task_id=int(task_id), @@ -66,6 +80,8 @@ def task_description(self) -> str: def reset(self) -> Observation: self.observer.reset() + for _ in range(self.settle_steps): + self.observer.step(LIBERO_DUMMY_ACTION) return self._observation() def step(self, action): @@ -80,10 +96,14 @@ def _observation(self) -> Observation: return Observation(images=images, state=self._state(obs)) def _state(self, obs) -> np.ndarray: - pos = np.asarray(obs["robot0_eef_pos"], dtype=np.float32) - axis_angle = quat_to_axis_angle(np.asarray(obs["robot0_eef_quat"], dtype=np.float32)) - gripper = np.asarray(obs["robot0_gripper_qpos"], dtype=np.float32) - return np.concatenate([pos, axis_angle, gripper]).astype(np.float32) + pos = np.asarray(obs["robot0_eef_pos"]) + axis_angle = quat_to_axis_angle(obs["robot0_eef_quat"]) + gripper = np.asarray(obs["robot0_gripper_qpos"]) + return np.concatenate([pos, axis_angle, gripper]) + + @property + def max_steps(self): + return MAX_POLICY_STEPS[self.benchmark] @property def supports_task_switch(self) -> bool: @@ -108,7 +128,15 @@ def set_task(self, task_name: str, task_config: str = "", seed=None) -> Observat raise ValueError(f"unknown LIBERO task {task_key!r}") init_state_id = self.init_state_id if str(task_config).strip() == "" else int(task_config) - new_seed = self.seed + 1 if seed is None or str(seed).strip() == "" else int(seed) + new_seed = self.seed if seed is None else int(seed) + + if task_key == self.task_name: + self.observer.select_init_state(init_state_id) + if seed is not None: + self.observer.seed = new_seed + self.observer.env.seed(new_seed) + self._sync_metadata() + return self.reset() # Construct the replacement first so an invalid task/config leaves the # currently displayed environment alive and usable. @@ -120,13 +148,13 @@ def set_task(self, task_name: str, task_config: str = "", seed=None) -> Observat seed=new_seed, libero_root=self.libero_root, ) + new_observer.reset() + for _ in range(self.settle_steps): + new_observer.step(LIBERO_DUMMY_ACTION) old_observer = self.observer self.observer = new_observer self._sync_metadata() - try: - old_observer.close() - except Exception: - pass + old_observer.close() return self._observation() def close(self) -> None: @@ -141,6 +169,7 @@ def build_libero_env(node) -> LiberoEnv: node.declare_parameter("init_state_id", 0) node.declare_parameter("image_size", contract.image_size) node.declare_parameter("seed", 0) + node.declare_parameter("settle_steps", 0) return LiberoEnv( contract, @@ -149,5 +178,6 @@ def build_libero_env(node) -> LiberoEnv: init_state_id=int(node.get_parameter("init_state_id").value), image_size=int(node.get_parameter("image_size").value), seed=int(node.get_parameter("seed").value), + settle_steps=int(node.get_parameter("settle_steps").value), libero_root=node.get_parameter("libero_root").value, ) diff --git a/lightx2v_ros/src/simulator/simulator/libero_node/evaluate.py b/lightx2v_ros/src/simulator/simulator/libero_node/evaluate.py new file mode 100644 index 000000000..9ef85ff36 --- /dev/null +++ b/lightx2v_ros/src/simulator/simulator/libero_node/evaluate.py @@ -0,0 +1,250 @@ +"""Run one official-order LIBERO suite through the ROS control plane.""" + +import json +import os +import time +from pathlib import Path + +import rclpy +from common.contract import get_contract +from rclpy.node import Node +from std_msgs.msg import String + +SUITES = ("libero_spatial", "libero_object", "libero_goal", "libero_10") +NUM_TASKS = 10 +NUM_INIT_STATES = 50 +FINISHED_STATES = ("success", "failure") + + +class LiberoSuiteEvaluator(Node): + def __init__(self): + super().__init__("libero_suite_evaluator") + self.declare_parameter("task_suite_name", "libero_spatial") + self.declare_parameter("output_dir", "data/libero/ros_evaluation") + self.declare_parameter("command_timeout", 180.0) + self.declare_parameter("overwrite", False) + self.contract = get_contract("libero") + + self.suite = str(self.get_parameter("task_suite_name").value).strip().lower() + if self.suite not in SUITES: + raise ValueError(f"task_suite_name must be one of {SUITES}, got {self.suite!r}") + + output_dir = Path(str(self.get_parameter("output_dir").value)).expanduser() + self.output_dir = output_dir / self.suite + self.episodes_path = self.output_dir / "episodes.jsonl" + self.summary_path = self.output_dir / "summary.json" + self.command_timeout = float(self.get_parameter("command_timeout").value) + if self.command_timeout <= 0: + raise ValueError("command_timeout must be positive") + self._prepare_output(bool(self.get_parameter("overwrite").value)) + + self.control_pub = self.create_publisher(String, self.contract.control_topic, 10) + self.status_sub = self.create_subscription(String, self.contract.status_topic, self._on_status, 10) + self.timer = self.create_timer(0.1, self._tick) + + self.status = None + self.phase = "waiting_for_simulator" + self.task_id = 0 + self.init_state_id = 0 + self.active_episode = None + self.set_task_after_episode = None + self.command_sent_at = None + self.results = [] + self.failure_reason = None + self.started_at = time.time() + self.get_logger().info(f"waiting for {self.contract.status_topic}; suite={self.suite}, output={self.output_dir}") + + @property + def task_key(self): + return f"{self.suite}/{self.task_id}" + + def _prepare_output(self, overwrite): + self.output_dir.mkdir(parents=True, exist_ok=True) + existing = [path for path in (self.episodes_path, self.summary_path) if path.exists()] + if existing and not overwrite: + paths = ", ".join(str(path) for path in existing) + raise FileExistsError(f"evaluation output already exists: {paths}; set overwrite:=true to replace it") + if overwrite: + for path in existing: + path.unlink() + + def _on_status(self, msg): + try: + self.status = json.loads(msg.data) + except (TypeError, ValueError) as exc: + self.get_logger().warning(f"ignoring invalid {self.contract.status_topic} message: {exc}") + + def _tick(self): + if self.status is None: + return + if bool(self.status.get("loop")): + self._abort("simulator parameter loop must be false for ordered suite evaluation") + return + + if self.phase == "waiting_for_simulator": + if self.control_pub.get_subscription_count() == 0: + return + if self.status.get("state") != "ready" or int(self.status.get("episode", -1)) != 0 or self.status.get("history"): + self._abort("simulator must be a fresh READY instance; launch libero_node with autostart:=false and loop:=false so no policy RNG is consumed before task 0/init 0") + return + if self._matches_target(): + self.active_episode = 0 + self._publish_control({"cmd": "start"}) + self.phase = "waiting_for_start" + else: + self._set_task() + elif self.phase == "waiting_for_ready": + episode = int(self.status.get("episode", -1)) + if self._matches_target() and self.status.get("state") == "ready" and episode > self.set_task_after_episode: + self.active_episode = episode + self._publish_control({"cmd": "start"}) + self.phase = "waiting_for_start" + else: + self._check_timeout("set_task") + elif self.phase == "waiting_for_start": + if self._matches_active_episode() and self.status.get("state") == "running": + self.phase = "running" + self.command_sent_at = None + else: + self._check_timeout("start") + elif self.phase == "running": + if not self._matches_active_episode(): + return + if self.status.get("state") in FINISHED_STATES: + self._record_episode() + if len(self.results) == NUM_TASKS * NUM_INIT_STATES: + self._finish() + else: + self._advance() + self._set_task() + + def _set_task(self): + self.set_task_after_episode = int(self.status.get("episode", -1)) + self._publish_control( + { + "cmd": "set_task", + "task_name": self.task_key, + "task_config": str(self.init_state_id), + } + ) + self.phase = "waiting_for_ready" + + def _publish_control(self, command): + msg = String() + msg.data = json.dumps(command, separators=(",", ":")) + self.control_pub.publish(msg) + self.command_sent_at = time.monotonic() + + def _matches_target(self): + return self.status.get("task_name") == self.task_key and str(self.status.get("task_config")) == str(self.init_state_id) + + def _matches_active_episode(self): + return self._matches_target() and self.status.get("episode") == self.active_episode + + def _check_timeout(self, command): + if self.command_sent_at is None: + return + if time.monotonic() - self.command_sent_at > self.command_timeout: + self._abort(f"timed out waiting for {command!r} acknowledgement") + + def _record_episode(self): + outcome = str(self.status["state"]) + result = { + "task_suite_name": self.suite, + "task_id": self.task_id, + "init_state_id": self.init_state_id, + "task_name": self.task_key, + "episode": self.active_episode, + "instruction": self.status.get("instruction", ""), + "seed": self.status.get("seed"), + "steps": int(self.status.get("episode_step", 0)), + "outcome": outcome, + "success": outcome == "success", + "timestamp": time.time(), + } + self.results.append(result) + with self.episodes_path.open("a", encoding="utf-8") as stream: + stream.write(json.dumps(result, separators=(",", ":")) + "\n") + stream.flush() + os.fsync(stream.fileno()) + self._write_summary(complete=False) + successes = sum(item["success"] for item in self.results) + self.get_logger().info(f"[{len(self.results)}/{NUM_TASKS * NUM_INIT_STATES}] task={self.task_id:02d} init={self.init_state_id:02d} {outcome}; success={successes / len(self.results):.2%}") + + def _advance(self): + self.init_state_id += 1 + if self.init_state_id == NUM_INIT_STATES: + self.task_id += 1 + self.init_state_id = 0 + self.active_episode = None + + def _summary(self, complete): + task_results = [] + for task_id in range(NUM_TASKS): + records = [item for item in self.results if item["task_id"] == task_id] + successes = sum(item["success"] for item in records) + task_results.append( + { + "task_id": task_id, + "episodes": len(records), + "successes": successes, + "success_rate": successes / len(records) if records else None, + } + ) + successes = sum(item["success"] for item in self.results) + return { + "protocol": "openpi_libero_official", + "task_suite_name": self.suite, + "task_order": "task_id_outer_init_state_id_inner", + "expected_episodes": NUM_TASKS * NUM_INIT_STATES, + "completed_episodes": len(self.results), + "successes": successes, + "success_rate": successes / len(self.results) if self.results else None, + "complete": complete, + "task_results": task_results, + "started_at": self.started_at, + "updated_at": time.time(), + } + + def _write_summary(self, complete): + temporary = self.summary_path.with_suffix(".json.tmp") + with temporary.open("w", encoding="utf-8") as stream: + json.dump(self._summary(complete), stream, indent=2) + stream.write("\n") + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, self.summary_path) + + def _finish(self): + self.phase = "done" + self._write_summary(complete=True) + successes = sum(item["success"] for item in self.results) + self.get_logger().info(f"suite complete: {successes}/{len(self.results)} ({successes / len(self.results):.2%}); summary={self.summary_path}") + self.timer.cancel() + rclpy.shutdown() + + def _abort(self, reason): + self.phase = "failed" + self.failure_reason = reason + self.get_logger().error(reason) + self.timer.cancel() + rclpy.shutdown() + + +def main(args=None): + rclpy.init(args=args) + node = LiberoSuiteEvaluator() + try: + rclpy.spin(node) + except KeyboardInterrupt: + pass + finally: + node.destroy_node() + if rclpy.ok(): + rclpy.shutdown() + if node.failure_reason: + raise RuntimeError(node.failure_reason) + + +if __name__ == "__main__": + main() diff --git a/lightx2v_ros/src/simulator/simulator/libero_node/observer.py b/lightx2v_ros/src/simulator/simulator/libero_node/observer.py index a26d5827b..a1182e582 100644 --- a/lightx2v_ros/src/simulator/simulator/libero_node/observer.py +++ b/lightx2v_ros/src/simulator/simulator/libero_node/observer.py @@ -28,7 +28,9 @@ def setup_libero_config(libero_root): if not (benchmark_root / "bddl_files").exists(): raise FileNotFoundError(f"LIBERO submodule is incomplete: {libero_root}") - config_dir = Path.home() / ".cache" / "lightx2v_ros" / "libero_config" + configured = os.environ.get("LIBERO_CONFIG_PATH") + cache_root = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) + config_dir = Path(configured).expanduser() if configured else cache_root / "lightx2v_ros" / "libero_config" config_file = config_dir / "config.yaml" config_dir.mkdir(parents=True, exist_ok=True) config_file.write_text( @@ -62,22 +64,6 @@ def load_libero(libero_root): return benchmark, get_libero_path, OffScreenRenderEnv -def load_init_states(get_libero_path, task, init_state_id): - state, _ = load_init_state(get_libero_path, task, init_state_id) - return state - - -def load_init_state(get_libero_path, task, init_state_id): - import torch - - init_states_path = Path(get_libero_path("init_states")) / task.problem_folder / task.init_states_file - init_states = torch.load(init_states_path, map_location="cpu", weights_only=False) - index = int(init_state_id) - if index < 0 or index >= len(init_states): - raise ValueError(f"init_state_id {index} is out of range for {task.name!r}; expected 0..{len(init_states) - 1}") - return init_states[index], len(init_states) - - def build_task_catalog(benchmark_module): """Return stable UI task ids mapped to their LIBERO suite/task metadata.""" factories = benchmark_module.get_benchmark_dict() @@ -128,19 +114,27 @@ def __init__( # Keep an owned copy: every restart must restore this exact MuJoCo state, # rather than returning the observation cached after the last action. self.init_state_id = int(init_state_id) - init_state, self.num_init_states = load_init_state(get_libero_path, task, self.init_state_id) - self.init_state = np.asarray(init_state).copy() + init_states_path = Path(get_libero_path("init_states")) / task.problem_folder / task.init_states_file + import torch + + self.init_states = np.asarray(torch.load(init_states_path, map_location="cpu", weights_only=False)) + self.num_init_states = len(self.init_states) + self.select_init_state(self.init_state_id) self.image_size = int(image_size) self.seed = int(seed) + # robosuite 1.4.1 mistakes a physical CUDA_VISIBLE_DEVICES value for an + # EGL ordinal. Imports must see the physical mask, while EGL must select + # logical device 0 inside that mask. + os.environ["MUJOCO_EGL_DEVICE_ID"] = "0" self.env = env_cls( bddl_file_name=str(bddl_file), camera_heights=self.image_size, camera_widths=self.image_size, camera_names=["robot0_eye_in_hand", "agentview", "frontview", "galleryview"], + render_gpu_device_id=0, ) self.env.seed(self.seed) - self.reset() @property def task_key(self): @@ -152,8 +146,15 @@ def reset(self): self.obs = self.env.set_init_state(self.init_state.copy()) return self.obs + def select_init_state(self, init_state_id): + index = int(init_state_id) + if index < 0 or index >= self.num_init_states: + raise ValueError(f"init_state_id {index} is out of range for {self.task.name!r}; expected 0..{self.num_init_states - 1}") + self.init_state_id = index + self.init_state = np.asarray(self.init_states[index]).copy() + def step(self, action): - action = np.asarray(action, dtype=np.float32) + action = np.asarray(action) self.obs, reward, success, info = self.env.step(action) return self.obs, reward, success, info diff --git a/lightx2v_ros/src/simulator/simulator/sim/node.py b/lightx2v_ros/src/simulator/simulator/sim/node.py index 7fd6347c9..cc216b117 100644 --- a/lightx2v_ros/src/simulator/simulator/sim/node.py +++ b/lightx2v_ros/src/simulator/simulator/sim/node.py @@ -35,7 +35,7 @@ from common.contract import EnvContract from rclpy.node import Node from sensor_msgs.msg import Image -from std_msgs.msg import Bool, Float32MultiArray, Int32, String +from std_msgs.msg import Bool, Float32MultiArray, Float64MultiArray, Int32, MultiArrayDimension, String from .base_env import BaseSimEnv @@ -50,11 +50,26 @@ FINISHED_STATES = (SUCCESS, FAILURE) -def rgb_to_image_msg(image, stamp, frame_id): +def parse_action_identity(label): + fields = {} + for item in str(label).split(";"): + name, separator, value = item.partition("=") + if separator: + fields[name] = value + try: + return int(fields["episode"]), int(fields["observation"]), int(fields["plan_epoch"]) + except (KeyError, ValueError): + return None + + +def rgb_to_image_msg(image, stamp, frame_id, episode_index=None, observation_index=None): image = np.ascontiguousarray(image) msg = Image() msg.header.stamp = stamp - msg.header.frame_id = frame_id + if episode_index is None or observation_index is None: + msg.header.frame_id = frame_id + else: + msg.header.frame_id = f"{frame_id}|{episode_index}|{observation_index}" msg.height = int(image.shape[0]) msg.width = int(image.shape[1]) msg.encoding = "rgb8" @@ -85,9 +100,19 @@ def __init__( # Per-episode step cap; <=0 means "use the env hint (env.max_steps)". # Hitting the cap ends the episode as FAILURE. self.declare_parameter("max_episode_steps", 0) + self.declare_parameter("numeric_precision", "float32") self.republish_period = float(self.get_parameter("republish_period").value) self.loop = bool(self.get_parameter("loop").value) + precision = str(self.get_parameter("numeric_precision").value).strip().lower() + numeric_types = { + "float32": (Float32MultiArray, np.float32), + "float64": (Float64MultiArray, np.float64), + } + if precision not in numeric_types: + raise ValueError("numeric_precision must be 'float32' or 'float64'") + self.numeric_message_type, self.numeric_dtype = numeric_types[precision] + # env_factory may declare/read its own parameters via `self`. self.env = env_factory(self) if self.env.contract is not contract: @@ -103,14 +128,15 @@ def __init__( else: self.max_episode_steps = 0 - self.state_pub = self.create_publisher(Float32MultiArray, contract.state_topic, 10) + self.state_pub = self.create_publisher(self.numeric_message_type, contract.state_topic, 10) self.image_pubs = {cam: self.create_publisher(Image, contract.camera_topic(cam), 10) for cam in contract.cameras} self.success_pub = self.create_publisher(Bool, contract.success_topic, 10) self.observation_ready_pub = self.create_publisher(Int32, contract.observation_ready_topic, 10) + self.observation_context_pub = self.create_publisher(String, contract.observation_context_topic, 10) self.task_pub = self.create_publisher(String, contract.task_topic, 10) self.episode_pub = self.create_publisher(Int32, contract.episode_topic, 10) self.status_pub = self.create_publisher(String, contract.status_topic, 10) - self.action_sub = self.create_subscription(Float32MultiArray, contract.action_topic, self.on_action, 10) + self.action_sub = self.create_subscription(self.numeric_message_type, contract.action_topic, self.on_action, 10) self.control_sub = self.create_subscription(String, contract.control_topic, self.on_control, 10) # `step_index` is a monotonic global observation counter (never reset), so the @@ -119,6 +145,7 @@ def __init__( # `episode_step` counts steps within the current episode (drives the step cap). self.episode_step = 0 self.episode_index = 0 + self.plan_epoch = 0 self.success = False self.state = READY self.history = [] # [{episode, task, config, seed, outcome, steps}] @@ -149,10 +176,18 @@ def publish_observation(self): for cam, pub in self.image_pubs.items(): image = self.obs.images.get(cam) if image is not None: - pub.publish(rgb_to_image_msg(image, stamp, cam)) - - state_msg = Float32MultiArray() - state_msg.data = np.asarray(self.obs.state, dtype=np.float32).reshape(-1).tolist() + pub.publish(rgb_to_image_msg(image, stamp, cam, self.episode_index, self.step_index)) + + state = np.asarray(self.obs.state, dtype=self.numeric_dtype).reshape(-1) + state_msg = self.numeric_message_type() + state_msg.layout.dim = [ + MultiArrayDimension( + label=f"episode={self.episode_index};observation={self.step_index}", + size=state.size, + stride=state.size, + ) + ] + state_msg.data = state.tolist() self.state_pub.publish(state_msg) task_msg = String() @@ -175,6 +210,17 @@ def publish_observation(self): ready_msg.data = self.step_index self.observation_ready_pub.publish(ready_msg) + context_msg = String() + context_msg.data = json.dumps( + { + "episode": self.episode_index, + "observation": self.step_index, + "plan_epoch": self.plan_epoch, + "task_description": self.env.task_description or "", + } + ) + self.observation_context_pub.publish(context_msg) + def publish_intermediate_frames(self, images): """Publish viewer-only frames rendered mid-action (no observation_ready).""" stamp = self.get_clock().now().to_msg() @@ -215,8 +261,14 @@ def publish_status(self): def on_action(self, msg): if self.state != RUNNING: return - - action = np.asarray(msg.data, dtype=np.float32).reshape(-1) + if msg.layout.dim: + action_identity = parse_action_identity(msg.layout.dim[0].label) + expected_identity = (self.episode_index, self.step_index, self.plan_epoch) + if action_identity != expected_identity: + self.get_logger().warning(f"dropping stale action {action_identity}; current observation is {expected_identity}") + return + + action = np.asarray(msg.data, dtype=self.numeric_dtype).reshape(-1) accepted_action_dims = tuple(int(dim) for dim in self.env.accepted_action_dims) if action.size not in accepted_action_dims: self.get_logger().error(f"expected action length in {accepted_action_dims}, got {action.size}") @@ -359,6 +411,7 @@ def _cmd_resume(self): self.state = RUNNING # Bump the counter: an action published for the pre-pause observation may # have been dropped, so re-advertise the current state as a new observation. + self.plan_epoch += 1 self.step_index += 1 self.publish_observation() self.publish_status() diff --git a/scripts/openpi/README.md b/scripts/openpi/README.md index b9bee25b8..47c647729 100644 --- a/scripts/openpi/README.md +++ b/scripts/openpi/README.md @@ -1,16 +1,16 @@ # OpenPI π0.5-LIBERO -该目录提供 π0.5-LIBERO 的权重转换、运行环境准备、本地 rollout 和定量评测。 -所有推理都从 LightX2V 公共入口启动： +该目录提供 π0.5-LIBERO 的权重转换、运行环境准备、本地 rollout、定量评测和 +ROS 交互。批量评测从 LightX2V 公共入口启动： ```text shell -> python -m lightx2v.infer -> OpenPIRunner -> 本地 PyTorch policy -> LIBERO/MuJoCo -> 结果文件 ``` -不启动 policy server，不使用 ROS，也不切换 Python 环境。OpenPI worker 与 -`lightx2v.infer` 使用同一个 base Python；仅任务特异的 Transformers 代码放在私有 -overlay 中。 +这条路径不启动 policy server。ROS 路径由 `openpi_node` 和 `libero_node` 通过 +topic 直连，同样不启动 OpenPI server。两条路径都使用当前 base Python；仅任务 +特异的 Transformers 代码放在私有 overlay 中。 以下命令默认在项目根目录执行： @@ -170,17 +170,220 @@ save_results/pi05_libero_pytorch_fp32_parallel_evaluation/ `libero_summary.py` 只读取各 suite 的 JSON 结果并生成最终汇总，不是推理入口。 输出锁会阻止两个并行任务同时写入同一目录。 -## 评测协议 +## 6. ROS 单 episode 交互 + +### 6.1 调用链与职责边界 + +```text +libero_node + ├─ 发布 256×256 RGB、8-D state、语言指令和 observation identity + ▼ +openpi_node + ├─ 同步 agentview / wrist / state / instruction + ├─ action queue 为空时调用一次 OpenPIPolicy.predict_action_chunk() + ├─ 得到 10×7 action，保存前 5 个 + └─ 每个新 observation 发布一个 7-D action + ▼ +libero_node + ├─ env.step(action) + ├─ 发布下一 observation + └─ 成功或达到 suite step cap 时结束 episode +``` + +| 组件 | 职责 | +| --- | --- | +| `OpenPIPolicy` | 加载模型、维护采样 RNG、一次调用生成完整 10-action chunk | +| `openpi_node` | 对齐跨 topic observation，维护 5-action 执行队列 | +| `libero_node` | 构造/重置环境、执行 warmup 和 action、判断成功和步数上限 | +| `SimulatorNode` | 提供 LightX2V ROS 通用状态机和 control/status topic | +| `libero_evaluate` | 按官方顺序调度 task/init state 并持久化指标 | + +主要实现文件： + +| 职责 | 文件 | +| --- | --- | +| ROS topic 契约 | `lightx2v_ros/src/common/common/contract.py` | +| OpenPI ROS 推理 node | `lightx2v_ros/src/inference/inference/openpi_node/main.py` | +| LIBERO 环境协议 | `lightx2v_ros/src/simulator/simulator/libero_node/env.py` | +| LIBERO observation 适配 | `lightx2v_ros/src/simulator/simulator/libero_node/observer.py` | +| ROS suite 调度与汇总 | `lightx2v_ros/src/simulator/simulator/libero_node/evaluate.py` | +| 通用 simulator 状态机 | `lightx2v_ros/src/simulator/simulator/sim/node.py` | +| 纯模型边界 | `lightx2v/models/runners/openpi/openpi_runner.py` | +| 三种进程的统一启动脚本 | `scripts/openpi/run_libero_ros_i2va.sh` | + +ROS 层不复制 PyTorch 模型、输入变换或 quantile 归一化实现，也不启动 OpenPI +policy server。 + +### 6.2 构建 ROS overlay + +首次使用或修改 ROS package 后构建 overlay： + +```bash +source /data/liuhongda/ros2_jazzy/install/setup.bash +cd /data/liuhongda/lightx2v_openpi/lightx2v_ros +colcon build --symlink-install --packages-select common simulator inference +source install/local_setup.bash +``` + +### 6.3 启动单 episode + +然后在两个 base 环境终端中启动同一个 episode。两个终端的 `ROS_DOMAIN_ID` 必须 +相同；下面让 MuJoCo 使用 4 号卡、模型使用 7 号卡。 + +Simulator 和 policy 也可以设置相同的 `CUDA_VISIBLE_DEVICES`，在显存允许时共享 +一张物理 GPU；这不会改变两个进程的 ROS 解耦关系，只可能带来少量资源竞争。例如 +两边都设置 `CUDA_VISIBLE_DEVICES=7` 时，进程内的 7 号卡会映射为逻辑设备 0。 +不要手动设置 `MUJOCO_EGL_DEVICE_ID=7`，LIBERO adapter 会自动选择逻辑 EGL 设备 0。 + +终端 1（LIBERO 仿真）： + +```bash +cd /data/liuhongda/lightx2v_openpi +conda activate base +ROS_DOMAIN_ID=77 CUDA_VISIBLE_DEVICES=4 \ +bash scripts/openpi/run_libero_ros_i2va.sh simulator libero_goal 3 0 +``` + +终端 2（OpenPI 推理）： + +```bash +cd /data/liuhongda/lightx2v_openpi +conda activate base +ROS_DOMAIN_ID=77 CUDA_VISIBLE_DEVICES=7 \ +bash scripts/openpi/run_libero_ros_i2va.sh policy +``` + +可在第三个已 source ROS overlay 的终端查看当前状态和 episode 结果： + +```bash +ros2 topic echo /libero/status +``` + +### 6.4 ROS 数据契约 + +OpenPI 使用 `/libero` namespace： + +| topic | 类型 | 内容 | +| --- | --- | --- | +| `/libero/agentview/image_raw` | `sensor_msgs/Image` | 第三人称 RGB | +| `/libero/wrist/image_raw` | `sensor_msgs/Image` | 腕部 RGB | +| `/libero/state` | `std_msgs/Float64MultiArray` | 8-D 末端/夹爪 state | +| `/libero/observation_context` | `std_msgs/String` | episode、observation、plan epoch 和语言指令 | +| `/libero/action` | `std_msgs/Float64MultiArray` | 带 observation identity 的 7-D action | +| `/libero/success` | `std_msgs/Bool` | 当前任务是否成功 | +| `/libero/control` | `std_msgs/String` | start/set_task 等 JSON 控制命令 | +| `/libero/status` | `std_msgs/String` | 状态机、步数和历史结果 JSON | + +相机消息、state 和 context 都携带 `(episode, observation)` 标识。推理 node 仅在 +两路图像、state 和 context 完全匹配时处理一次，避免 ROS topic 异步到达造成跨帧 +拼接。episode 切换时会丢弃未执行 action，但不会重置 policy RNG，从而保持官方 +suite 内连续 RNG。pause/resume 会递增 plan epoch；仿真端会拒绝旧 episode、旧 +observation 或旧 plan epoch 的迟到 action。 + +脚本默认使用 FP32 转换权重和官方 LIBERO checkout；可分别通过 +`OPENPI_MODEL_PATH`、`OPENPI_CONFIG`、`OPENPI_LIBERO_ROOT` 覆盖。 + +## 7. ROS 完整 suite 评测 + +完整 ROS 评测由 `libero_evaluate` ROS node 调度。它要求 simulator 以 +`autostart:=false`、`loop:=false` 启动，并从未运行过 episode 的初始 `ready` 状态 +接管，然后按 task id 0–9、init-state id 0–49 的官方顺序发送控制命令。单 suite +使用三个终端；前两个命令与上面相同，只需让 simulator 等待 evaluator 启动： + +```bash +ROS_DOMAIN_ID=81 CUDA_VISIBLE_DEVICES=4 OPENPI_ROS_AUTOSTART=false \ +bash scripts/openpi/run_libero_ros_i2va.sh simulator libero_spatial 0 0 + +ROS_DOMAIN_ID=81 CUDA_VISIBLE_DEVICES=7 \ +bash scripts/openpi/run_libero_ros_i2va.sh policy + +ROS_DOMAIN_ID=81 \ +OPENPI_ROS_OUTPUT_DIR=/data/liuhongda/lightx2v_openpi/save_results/openpi_ros_evaluation \ +bash scripts/openpi/run_libero_ros_i2va.sh evaluate libero_spatial +``` + +coordinator 固定以 task id 0–9 为外层、init-state id 0–49 为内层运行 500 个 +episode。每个 episode 会立即追加并 `fsync` 到 `episodes.jsonl`，`summary.json` +采用原子更新。已有输出默认拒绝覆盖；确认重跑时设置 +`OPENPI_ROS_OVERWRITE=true`。coordinator 只接受 episode 0、history 为空的初始 +`ready` simulator。如果初始环境已经是 task 0/init 0，它会直接发送 `start`，避免 +多做一次 reset。 + +四个 suite 并行时必须使用互不相同的 `ROS_DOMAIN_ID`，每个 domain 分别启动 +simulator、policy 和 evaluator。只有 4 张卡时，MuJoCo 与 policy 可以共享同一张卡： + +| suite | ROS domain | GPU | +| --- | ---: | ---: | +| `libero_spatial` | 81 | 4 | +| `libero_object` | 82 | 5 | +| `libero_goal` | 83 | 6 | +| `libero_10` | 84 | 7 | + +最终输出结构： + +```text +save_results/openpi_ros_evaluation/ +├── libero_spatial/{episodes.jsonl,summary.json} +├── libero_object/{episodes.jsonl,summary.json} +├── libero_goal/{episodes.jsonl,summary.json} +└── libero_10/{episodes.jsonl,summary.json} +``` + +非 ROS 的完整定量回归入口仍保留： + +```bash +CUDA_VISIBLE_DEVICES=4,5,6,7 \ +bash scripts/openpi/run_libero_evaluate_parallel_i2va.sh +``` + +该命令经过 `python -m lightx2v.infer`，不是 ROS 路径。ROS coordinator 只负责发送 +`set_task/start`、读取 `/libero/status` 和保存汇总，不会把环境循环放回模型 node。 + +## 8. 评测协议与定量证据 `configs/openpi/pi05_libero_eval.json` 是默认协议来源：每个 suite 包含 10 个任务， 每个任务测试 50 个 init states，共 500 episodes；完整 LIBERO-40 共 2000 个。 环境/策略 seed 为 7/0，先执行 10 个 no-op，每次预测 10 个 action 并执行前 5 个。 四个 suite 的最大步数分别为 220、280、300、520。 -评测默认支持断点恢复。`protocol_id`、输入文件 manifest 和保存的 policy RNG state -用于避免混用不同协议，并保证前缀恢复后的随机数流与一次性运行一致。 +| 项目 | 对齐设置 | +| --- | --- | +| 环境渲染 | 256×256 | +| 相机方向 | agentview 和 wrist 均旋转 180° | +| 模型图像输入 | PIL bilinear resize-with-pad 到 224×224，并保持官方 uint8 量化点 | +| state | 末端位置 3 + 四元数转 axis-angle 3 + gripper 2，共 8 维 | +| 数值链路 | ROS state/action 为 float64；quantile JSON 统计量保留 float64 | +| 环境 / policy seed | 7 / 0，policy RNG 在 suite 内连续 | +| episode warmup | 10 次 `[0, 0, 0, 0, 0, 0, -1]`，不送入模型 | +| action horizon / replan | 生成 10 个，执行前 5 个后重新规划 | +| 最大 policy steps | spatial 220、object 280、goal 300、LIBERO-10 520 | -## 常用覆盖项 +成功条件直接采用 LIBERO `env.step()` 返回的 `done`；warmup 不计入 policy step +上限。模型反归一化后的 7-D action 不做额外夹爪二值化，也不强制降为 float32。 + +直接评测默认支持断点恢复。`protocol_id`、输入文件 manifest 和保存的 policy RNG +state 用于避免混用不同协议，并保证前缀恢复后的随机数流与一次性运行一致。ROS +coordinator 当前按完整 suite 写新目录，不提供中途恢复。 + +已有的完整 2000-episode 数据来自直接 LightX2V 评测链路： + +| 实现 | spatial | object | goal | LIBERO-10 | 总计 | +| --- | ---: | ---: | ---: | ---: | ---: | +| 官方本地 JAX/OpenPI | 493/500 | 492/500 | 486/500 | 472/500 | 1943/2000（97.15%） | +| LightX2V PyTorch 直接评测 | 497/500 | 492/500 | 489/500 | 466/500 | 1944/2000（97.20%） | + +二者相差 `+0.05` 个百分点，说明 PyTorch 权重、模型数值路径和直接 LIBERO 协议 +处于同一水平。证据位置： + +- 官方日志：`/data/liuhongda/openpi/data/libero/official_pi05_libero/logs/` +- LightX2V 汇总： + `/data/liuhongda/lightx2v_openpi/save_results/pi05_libero_pytorch_fp32_parallel_evaluation_lhdtest/parallel_summary.json` + +已实测 `libero_spatial/task_00/init_00`：ROS 和直接评测均在第 79 个 policy step +成功，且都是 16 次模型前向。coordinator 随后正确切换到 `init_01`，并将第一局结果 +写入 JSONL/summary；完整 2000 局仍需实际跑完后再作为 ROS 定量结果引用。 + +## 9. 常用覆盖项 | 环境变量 | 作用 | | --- | --- | @@ -199,11 +402,14 @@ save_results/pi05_libero_pytorch_fp32_parallel_evaluation/ | `OPENPI_EVAL_RESUME` | 是否恢复已有结果，使用 `0/1` | | `OPENPI_EVAL_OUTPUT_DIR` | 单卡评测输出目录 | | `OPENPI_PARALLEL_OUTPUT_ROOT` | 多卡评测输出目录 | +| `OPENPI_ROS_AUTOSTART` | ROS simulator 是否直接开始；suite 评测必须为 `false` | +| `OPENPI_ROS_OUTPUT_DIR` | ROS suite 结果根目录 | +| `OPENPI_ROS_OVERWRITE` | 是否覆盖已有 ROS suite 结果，默认 `false` | -## 开发检查 +## 10. 开发检查与验收 ```bash -bash -n scripts/openpi/run_libero_*.sh +bash -n scripts/openpi/*.sh python -m unittest discover -s scripts/openpi/tests -p 'test_*.py' -v python scripts/openpi/tests/validate_pytorch_parity.py --self-check pre-commit run --all-files @@ -211,3 +417,8 @@ pre-commit run --all-files 数值路径的关键约束是官方图像 resize/uint8 量化、FP64 动作反归一化、连续 policy RNG 和 5-action replan queue；清理启动脚本时不应改变这些逻辑。 + +ROS 契约测试覆盖乱序/重复 topic、10→5 action queue、float64 传输、官方 +旋转/resize/state/action 变换和迟到 action 拒绝。修改 ROS 文件后还应重新构建 +`common simulator inference`；LIBERO adapter 会把可见的物理 GPU 映射为 EGL +逻辑设备 0。 diff --git a/scripts/openpi/convert_jax_checkpoint.py b/scripts/openpi/convert_jax_checkpoint.py index 99ef8e96c..870a5b68d 100755 --- a/scripts/openpi/convert_jax_checkpoint.py +++ b/scripts/openpi/convert_jax_checkpoint.py @@ -148,7 +148,7 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument( "--precision", choices=tuple(DTYPE_BY_PRECISION), - default=os.environ.get("OPENPI_CONVERT_PRECISION", os.environ.get("OPENPI_OUTPUT_PRECISION", "float32")), + default=os.environ.get("OPENPI_CONVERT_PRECISION", "float32"), ) parser.add_argument("--config-name", default=os.environ.get("OPENPI_CONFIG_NAME", "pi05_libero")) parser.add_argument("--openpi-root", default=os.environ.get("OPENPI_PATH", str(DEFAULT_OPENPI_ROOT))) diff --git a/scripts/openpi/libero_summary.py b/scripts/openpi/libero_summary.py index 77d489876..bcb595f51 100755 --- a/scripts/openpi/libero_summary.py +++ b/scripts/openpi/libero_summary.py @@ -13,35 +13,27 @@ EPISODES_PER_SUITE = 500 -def _parse_suites(value: str) -> tuple[str, ...]: - suites = tuple(item.strip() for item in value.split(",") if item.strip()) - unknown = [suite for suite in suites if suite not in SUITES] - if not suites or unknown or len(set(suites)) != len(suites): - raise ValueError(f"invalid LIBERO suite selection: {value!r}") - return suites - - -def _parse_status(values: list[str], suites: tuple[str, ...]) -> dict[str, int]: +def _parse_status(values: list[str]) -> dict[str, int]: statuses: dict[str, int] = {} for value in values: suite, separator, return_code = value.partition("=") - if not separator or suite not in suites or suite in statuses: + if not separator or suite not in SUITES or suite in statuses: raise ValueError(f"invalid worker status: {value!r}") statuses[suite] = int(return_code) - missing = set(suites) - statuses.keys() + missing = set(SUITES) - statuses.keys() if missing: raise ValueError(f"missing worker status for: {sorted(missing)}") return statuses -def aggregate(output_root: Path, suites: tuple[str, ...], statuses: dict[str, int], output_file: Path) -> tuple[dict, list[str]]: +def aggregate(output_root: Path, statuses: dict[str, int]) -> tuple[dict, list[str]]: errors: list[str] = [] shards: dict[str, dict] = {} successes = 0 completed = 0 rates: list[float] = [] - for suite in suites: + for suite in SUITES: summary_path = output_root / suite / "summary.json" if not summary_path.is_file(): errors.append(f"{suite}: missing {summary_path}") @@ -89,7 +81,7 @@ def aggregate(output_root: Path, suites: tuple[str, ...], statuses: dict[str, in "protocol_id": protocol_id, } - expected = EPISODES_PER_SUITE * len(suites) + expected = EPISODES_PER_SUITE * len(SUITES) payload = { "schema_version": 1, "status": "complete" if not errors else "invalid", @@ -98,12 +90,12 @@ def aggregate(output_root: Path, suites: tuple[str, ...], statuses: dict[str, in "successes": successes, "failures": completed - successes, "success_rate": successes / completed * 100.0 if completed else None, - "mean_suite_success_rate": sum(rates) / len(rates) if len(rates) == len(suites) else None, + "mean_suite_success_rate": sum(rates) / len(rates) if len(rates) == len(SUITES) else None, "shards": shards, "validation_errors": errors, "updated_at_unix": time.time(), } - output_file.parent.mkdir(parents=True, exist_ok=True) + output_file = output_root / "parallel_summary.json" temporary = output_file.with_name(f".{output_file.name}.{os.getpid()}.tmp") temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") os.replace(temporary, output_file) @@ -113,21 +105,17 @@ def aggregate(output_root: Path, suites: tuple[str, ...], statuses: dict[str, in def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--output-root", type=Path, required=True) - parser.add_argument("--suites", required=True) parser.add_argument("--worker-status", action="append", default=[]) - parser.add_argument("--output-file", type=Path) args = parser.parse_args() try: output_root = args.output_root.expanduser().resolve() - suites = _parse_suites(args.suites) - statuses = _parse_status(args.worker_status, suites) - output_file = args.output_file.expanduser().resolve() if args.output_file else output_root / "parallel_summary.json" - payload, errors = aggregate(output_root, suites, statuses, output_file) + statuses = _parse_status(args.worker_status) + payload, errors = aggregate(output_root, statuses) except (OSError, TypeError, ValueError) as exc: parser.error(str(exc)) print("\nLIBERO suite summary") - for suite in suites: + for suite in SUITES: shard = payload["shards"].get(suite, {}) rate = shard.get("success_rate") rate_text = "n/a" if rate is None else f"{float(rate):.2f}%" @@ -135,7 +123,7 @@ def main() -> int: overall = payload["success_rate"] overall_text = "n/a" if overall is None else f"{overall:.2f}%" print(f" {'overall':<18} {payload['successes']:>3}/{payload['expected_episodes']:<3} {overall_text:>8}") - print(f" summary: {output_file}") + print(f" summary: {output_root / 'parallel_summary.json'}") for error in errors: print(f"error: {error}") return 1 if errors else 0 diff --git a/scripts/openpi/run_libero_evaluate_parallel_i2va.sh b/scripts/openpi/run_libero_evaluate_parallel_i2va.sh index 156d76d7a..aded05256 100755 --- a/scripts/openpi/run_libero_evaluate_parallel_i2va.sh +++ b/scripts/openpi/run_libero_evaluate_parallel_i2va.sh @@ -3,16 +3,8 @@ set -euo pipefail script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" lightx2v_path="${LIGHTX2V_PATH:-$(cd -- "${script_dir}/../.." && pwd)}" -workspace_root="$(dirname -- "${lightx2v_path}")" -openpi_data_root="${OPENPI_DATA_ROOT:-${workspace_root}/openpi_data}" -model_path="${OPENPI_MODEL_PATH:-${openpi_data_root}/openpi-assets/checkpoints/pi05_libero_pytorch_fp32}" -config_json="${OPENPI_CONFIG:-${lightx2v_path}/configs/openpi/pi05_libero.json}" -eval_config="${OPENPI_EVAL_CONFIG:-${lightx2v_path}/configs/openpi/pi05_libero_eval.json}" -libero_root="${OPENPI_LIBERO_ROOT:-${workspace_root}/openpi/third_party/libero}" -transformers_runtime="${OPENPI_TRANSFORMERS_RUNTIME_PATH:-${openpi_data_root}/python_deps/openpi_official_pytorch_runtime}" gpu_list="${CUDA_VISIBLE_DEVICES:-4,5,6,7}" -suite_list="libero_spatial,libero_object,libero_goal,libero_10" output_root="${OPENPI_PARALLEL_OUTPUT_ROOT:-${lightx2v_path}/save_results/pi05_libero_pytorch_fp32_parallel_evaluation}" IFS=',' read -ra gpus <<< "${gpu_list}" @@ -33,6 +25,9 @@ if ! flock -n "${lock_fd}"; then echo "Another evaluation is using ${output_root}" >&2 exit 2 fi +for suite in "${suites[@]}"; do + rm -f "${output_root}/runtime/${suite}.status" +done run_worker() { local worker_index="$1" @@ -59,27 +54,14 @@ run_worker() { ( export CUDA_VISIBLE_DEVICES="${gpu}" - export MUJOCO_GL="${MUJOCO_GL:-egl}" - export PYOPENGL_PLATFORM="${PYOPENGL_PLATFORM:-egl}" - export OPENPI_RUN_MODE=evaluate - export OPENPI_LIBERO_ROOT="${libero_root}" + export OPENPI_EVAL_OUTPUT_DIR="${suite_output}" + export OPENPI_RUNTIME_DIR="${suite_runtime}" export OPENPI_LIBERO_CONFIG_DIR="${suite_runtime}/libero_config" - export OPENPI_EVAL_CONFIG="${eval_config}" export OPENPI_EVAL_BENCHMARKS="${suite}" - export OPENPI_TRANSFORMERS_RUNTIME_PATH="${transformers_runtime}" export NUMBA_CACHE_DIR="${suite_runtime}/numba" export MPLCONFIGDIR="${suite_runtime}/matplotlib" export XDG_CACHE_HOME="${suite_runtime}/cache" - export PYTHONNOUSERSITE=1 - export TOKENIZERS_PARALLELISM=false - export PROFILING_DEBUG_LEVEL="${PROFILING_DEBUG_LEVEL:-0}" - exec setsid python -m lightx2v.infer \ - --model_cls openpi \ - --task i2va \ - --model_path "${model_path}" \ - --config_json "${config_json}" \ - --seed "${OPENPI_POLICY_SEED:-0}" \ - --save_result_path "${suite_output}" + exec setsid bash "${script_dir}/run_libero_evaluate_i2va.sh" ) >> "${log_path}" 2>&1 & child_pid=$! if wait "${child_pid}"; then @@ -125,7 +107,6 @@ trap - EXIT HUP INT TERM summary_command=( python "${script_dir}/libero_summary.py" --output-root "${output_root}" - --suites "${suite_list}" ) for suite in "${suites[@]}"; do exit_code=1 diff --git a/scripts/openpi/run_libero_ros_i2va.sh b/scripts/openpi/run_libero_ros_i2va.sh new file mode 100755 index 000000000..71c1eaa85 --- /dev/null +++ b/scripts/openpi/run_libero_ros_i2va.sh @@ -0,0 +1,75 @@ +#!/usr/bin/env bash +set -euo pipefail + +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +lightx2v_path="${LIGHTX2V_PATH:-$(cd -- "${script_dir}/../.." && pwd)}" +workspace_root="$(dirname -- "${lightx2v_path}")" +openpi_data_root="${OPENPI_DATA_ROOT:-${workspace_root}/openpi_data}" +ros_setup="${ROS_SETUP:-${workspace_root}/ros2_jazzy/install/setup.bash}" +ros_overlay="${LIGHTX2V_ROS_SETUP:-${lightx2v_path}/lightx2v_ros/install/local_setup.bash}" +mode="${1:-}" +domain_id="${ROS_DOMAIN_ID:-0}" + +if [[ "${mode}" != "simulator" && "${mode}" != "policy" && "${mode}" != "evaluate" ]]; then + echo "Usage: bash scripts/openpi/run_libero_ros_i2va.sh {simulator [suite task_id init_state_id]|policy|evaluate [suite]}" >&2 + exit 2 +fi +if [[ ! -f "${ros_setup}" || ! -f "${ros_overlay}" ]]; then + echo "Build and source the ROS workspace first; missing ${ros_setup} or ${ros_overlay}" >&2 + exit 1 +fi + +set +u +source "${ros_setup}" +source "${ros_overlay}" +set -u + +export RMW_IMPLEMENTATION="${RMW_IMPLEMENTATION:-rmw_fastrtps_cpp}" +export NUMBA_CACHE_DIR="${NUMBA_CACHE_DIR:-/tmp/lightx2v-openpi-ros-${domain_id}-numba}" +export MPLCONFIGDIR="${MPLCONFIGDIR:-/tmp/lightx2v-openpi-ros-${domain_id}-matplotlib}" +export ROS_LOG_DIR="${ROS_LOG_DIR:-/tmp/lightx2v-openpi-ros-${domain_id}-log}" +export LIBERO_CONFIG_PATH="${LIBERO_CONFIG_PATH:-/tmp/lightx2v-openpi-ros-${domain_id}-libero-config}" +mkdir -p "${NUMBA_CACHE_DIR}" "${MPLCONFIGDIR}" "${ROS_LOG_DIR}" "${LIBERO_CONFIG_PATH}" + +if [[ "${mode}" == "simulator" ]]; then + suite="${2:-libero_spatial}" + task_id="${3:-0}" + init_state_id="${4:-0}" + export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-4}" + export MUJOCO_GL="${MUJOCO_GL:-egl}" + export PYOPENGL_PLATFORM="${PYOPENGL_PLATFORM:-egl}" + unset MUJOCO_EGL_DEVICE_ID + exec ros2 run simulator libero_node --ros-args \ + -p "autostart:=${OPENPI_ROS_AUTOSTART:-true}" \ + -p loop:=false \ + -p numeric_precision:=float64 \ + -p settle_steps:=10 \ + -p "libero_root:=${OPENPI_LIBERO_ROOT:-${workspace_root}/openpi/third_party/libero}" \ + -p "benchmark:=${suite}" \ + -p "task_id:=${task_id}" \ + -p "init_state_id:=${init_state_id}" \ + -p image_size:=256 \ + -p seed:=7 +fi + +if [[ "${mode}" == "evaluate" ]]; then + suite="${2:-libero_spatial}" + exec ros2 run simulator libero_evaluate --ros-args \ + -p "task_suite_name:=${suite}" \ + -p "output_dir:=${OPENPI_ROS_OUTPUT_DIR:-${lightx2v_path}/save_results/openpi_ros_evaluation}" \ + -p "command_timeout:=${OPENPI_ROS_COMMAND_TIMEOUT:-180.0}" \ + -p "overwrite:=${OPENPI_ROS_OVERWRITE:-false}" +fi + +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-7}" +export USE_FLAX=0 +export PYTHONNOUSERSITE=1 +export TOKENIZERS_PARALLELISM=false +export PYTHONPATH="${OPENPI_TRANSFORMERS_RUNTIME_PATH:-${openpi_data_root}/python_deps/openpi_official_pytorch_runtime}:${lightx2v_path}:${PYTHONPATH:-}" + +exec ros2 run inference openpi_node --ros-args \ + -p numeric_precision:=float64 \ + -p "model_path:=${OPENPI_MODEL_PATH:-${openpi_data_root}/openpi-assets/checkpoints/pi05_libero_pytorch_fp32}" \ + -p "config_json:=${OPENPI_CONFIG:-${lightx2v_path}/configs/openpi/pi05_libero.json}" \ + -p seed:=0 \ + -p actions_per_plan:=5 diff --git a/scripts/openpi/runtime.py b/scripts/openpi/runtime.py index df7897610..72e0d309f 100755 --- a/scripts/openpi/runtime.py +++ b/scripts/openpi/runtime.py @@ -137,7 +137,6 @@ def _check_patch_overlay(target: Path) -> None: import importlib.metadata import json import sys -import sysconfig from pathlib import Path root = Path(sys.argv[1]).resolve() @@ -364,8 +363,6 @@ def _check_static_inputs(args: argparse.Namespace) -> None: origin = Path(module.__file__).resolve() if not origin.is_relative_to(libero_root): raise RuntimeError(f"{module.__name__} imported outside the official LIBERO root: {origin}") -namespace_extras = [path for path in namespace_paths if not Path(path).is_relative_to(libero_root)] - payload = { "python": sys.version.split()[0], "torch": torch.__version__, @@ -381,7 +378,6 @@ def _check_static_inputs(args: argparse.Namespace) -> None: "glfw": importlib.metadata.version("glfw"), "robosuite": importlib.metadata.version("robosuite"), "libero_namespace": namespace_paths, - "libero_namespace_extras_ignored_by_source_guard": namespace_extras, "libero_suites": sorted(benchmark.get_benchmark_dict()), } print(json.dumps(payload, sort_keys=True)) diff --git a/scripts/openpi/tests/test_ros_openpi_contract.py b/scripts/openpi/tests/test_ros_openpi_contract.py new file mode 100644 index 000000000..df88e4062 --- /dev/null +++ b/scripts/openpi/tests/test_ros_openpi_contract.py @@ -0,0 +1,418 @@ +"""Offline contract tests for the OpenPI ROS bridge. + +Run after sourcing the ROS underlay and LightX2V overlay. The tests exercise +message conversion and callback ordering without starting DDS participants or +loading model weights. +""" + +from __future__ import annotations + +import json +import math +import sys +import tempfile +import types +import unittest +from collections import deque +from pathlib import Path + +import numpy as np + +PROJECT_ROOT = Path(__file__).resolve().parents[3] +ROS_SOURCE_ROOT = PROJECT_ROOT / "lightx2v_ros/src" +for package in ("common", "inference", "simulator"): + sys.path.insert(0, str(ROS_SOURCE_ROOT / package)) + + +def _stub_lightx2v_imports(): + """Keep helper tests independent of LightX2V CUDA initialization.""" + modules = { + "lightx2v": PROJECT_ROOT / "lightx2v", + "lightx2v.models": PROJECT_ROOT / "lightx2v/models", + "lightx2v.models.networks": PROJECT_ROOT / "lightx2v/models/networks", + "lightx2v.models.networks.openpi": PROJECT_ROOT / "lightx2v/models/networks/openpi", + "lightx2v.models.runners": PROJECT_ROOT / "lightx2v/models/runners", + "lightx2v.models.runners.openpi": PROJECT_ROOT / "lightx2v/models/runners/openpi", + "lightx2v.utils": PROJECT_ROOT / "lightx2v/utils", + } + for name, path in modules.items(): + module = types.ModuleType(name) + module.__path__ = [str(path)] + sys.modules[name] = module + + runner = types.ModuleType("lightx2v.models.runners.openpi.openpi_runner") + runner.OpenPIPolicy = object + sys.modules[runner.__name__] = runner + + config = types.ModuleType("lightx2v.utils.set_config") + config.auto_calc_config = lambda values: values + config.get_default_config = dict + sys.modules[config.__name__] = config + + +try: + from builtin_interfaces.msg import Time + from common.contract import LIBERO_CONTRACT + from sensor_msgs.msg import Image + from std_msgs.msg import Float64MultiArray, MultiArrayDimension, String + + _stub_lightx2v_imports() + from inference.openpi_node.main import OpenPINode, image_msg_to_rgb, observation_identity, state_identity + from simulator.libero_node.env import LiberoEnv, quat_to_axis_angle + from simulator.sim.node import RUNNING, SimulatorNode, parse_action_identity, rgb_to_image_msg + + from lightx2v.models.networks.openpi.infer.post_infer import OpenPIPostInfer + from lightx2v.models.networks.openpi.infer.pre_infer import _resize_with_pad +except ImportError as exc: # pragma: no cover - depends on sourced ROS setup + raise unittest.SkipTest(f"source the ROS underlay before running this test: {exc}") from exc + + +class _Logger: + def __init__(self): + self.errors = [] + self.warnings = [] + + def error(self, message): + self.errors.append(str(message)) + + def info(self, _message): + return None + + def warning(self, message): + self.warnings.append(str(message)) + + +class _Publisher: + def __init__(self): + self.messages = [] + + def publish(self, message): + self.messages.append(message) + + +class _ChunkPolicy: + output_action_dim = 7 + + def __init__(self): + self.calls = [] + + def predict_action_chunk(self, *, images, state, task_description): + call_index = len(self.calls) + self.calls.append( + { + "images": {name: image.copy() for name, image in images.items()}, + "state": state.copy(), + "task_description": task_description, + } + ) + return np.arange(70, dtype=np.float64).reshape(10, 7) + call_index * 1000.0 + + +def _policy_harness(): + node = types.SimpleNamespace( + contract=LIBERO_CONTRACT, + numeric_message_type=Float64MultiArray, + numeric_dtype=np.float64, + images={camera: None for camera in LIBERO_CONTRACT.policy_input_cameras}, + state=None, + task_description="", + episode_index=None, + plan_epoch=0, + pending_observation=None, + last_processed_observation=None, + pending_actions=deque(), + actions_per_plan=5, + policy=_ChunkPolicy(), + action_pub=_Publisher(), + get_logger=lambda: _Logger(), + ) + node._try_process_observation = types.MethodType(OpenPINode._try_process_observation, node) + node._publish_action = types.MethodType(OpenPINode._publish_action, node) + return node + + +def _context(episode, observation, prompt="pick up the bowl", plan_epoch=0): + message = String() + message.data = json.dumps( + { + "episode": episode, + "observation": observation, + "plan_epoch": plan_epoch, + "task_description": prompt, + } + ) + return message + + +def _state(episode, observation, values=None): + message = Float64MultiArray() + message.layout.dim = [MultiArrayDimension(label=f"episode={episode};observation={observation}", size=8, stride=8)] + message.data = (np.arange(8, dtype=np.float64) if values is None else np.asarray(values)).tolist() + return message + + +def _image(camera, episode, observation, value=0): + pixels = np.full((4, 5, 3), value, dtype=np.uint8) + return rgb_to_image_msg(pixels, Time(), camera, episode, observation) + + +def _deliver_observation(node, episode, observation, plan_epoch=0): + OpenPINode._on_context(node, _context(episode, observation, plan_epoch=plan_epoch)) + OpenPINode._image_callback(node, "wrist")(_image("wrist", episode, observation, 29)) + OpenPINode._on_state(node, _state(episode, observation)) + OpenPINode._image_callback(node, "agentview")(_image("agentview", episode, observation, 11)) + + +class RosMessageContractTest(unittest.TestCase): + def test_tagged_rgb_round_trip(self): + image = np.arange(5 * 7 * 3, dtype=np.uint8).reshape(5, 7, 3)[:, ::-1] + message = rgb_to_image_msg(image, Time(), "agentview", 13, 21) + + self.assertEqual(message.header.frame_id, "agentview|13|21") + self.assertEqual(observation_identity(message.header.frame_id), (13, 21)) + np.testing.assert_array_equal(image_msg_to_rgb(message), image) + + def test_bgr_message_with_row_padding(self): + rgb = np.arange(2 * 3 * 3, dtype=np.uint8).reshape(2, 3, 3) + row_bytes = 3 * 3 + 4 + encoded = np.full((2, row_bytes), 255, dtype=np.uint8) + encoded[:, :9] = rgb[:, :, ::-1].reshape(2, 9) + message = Image(height=2, width=3, encoding="bgr8", step=row_bytes, data=encoded.tobytes()) + + np.testing.assert_array_equal(image_msg_to_rgb(message), rgb) + + def test_identity_parser_rejects_untagged_and_malformed_frames(self): + self.assertIsNone(observation_identity("agentview")) + self.assertIsNone(observation_identity("agentview|episode|2")) + self.assertIsNone(observation_identity("agentview|1")) + self.assertEqual(state_identity("episode=5;observation=91"), (5, 91)) + self.assertIsNone(state_identity("episode=5")) + self.assertIsNone(state_identity("episode=5;observation=bad")) + + +class OpenPINodeSynchronizationTest(unittest.TestCase): + def test_waits_for_one_matching_set_and_ignores_duplicate_delivery(self): + node = _policy_harness() + OpenPINode._on_context(node, _context(3, 8)) + OpenPINode._on_state(node, _state(3, 7)) + OpenPINode._image_callback(node, "agentview")(_image("agentview", 3, 8)) + OpenPINode._image_callback(node, "wrist")(_image("wrist", 3, 8)) + self.assertEqual(len(node.action_pub.messages), 0) + + OpenPINode._on_state(node, _state(3, 8)) + self.assertEqual(len(node.policy.calls), 1) + self.assertEqual(len(node.action_pub.messages), 1) + + OpenPINode._on_context(node, _context(3, 8)) + OpenPINode._on_state(node, _state(3, 8)) + self.assertEqual(len(node.policy.calls), 1) + self.assertEqual(len(node.action_pub.messages), 1) + + def test_executes_five_actions_from_each_ten_action_chunk(self): + node = _policy_harness() + for observation in range(6): + _deliver_observation(node, 0, observation) + + self.assertEqual(len(node.policy.calls), 2) + published = np.asarray([message.data for message in node.action_pub.messages]) + first_chunk = np.arange(70, dtype=np.float64).reshape(10, 7) + expected = np.concatenate([first_chunk[:5], first_chunk[:1] + 1000.0]) + np.testing.assert_array_equal(published, expected) + + def test_new_episode_discards_the_old_action_tail_without_resetting_policy(self): + node = _policy_harness() + _deliver_observation(node, 0, 0) + self.assertEqual(len(node.pending_actions), 4) + + _deliver_observation(node, 1, 1) + self.assertEqual(len(node.policy.calls), 2) + self.assertEqual(len(node.pending_actions), 4) + np.testing.assert_array_equal(node.action_pub.messages[-1].data, np.arange(7) + 1000.0) + + def test_new_plan_epoch_discards_actions_queued_before_resume(self): + node = _policy_harness() + _deliver_observation(node, 0, 0, plan_epoch=0) + self.assertEqual(len(node.pending_actions), 4) + + _deliver_observation(node, 0, 1, plan_epoch=1) + self.assertEqual(len(node.policy.calls), 2) + self.assertEqual(len(node.pending_actions), 4) + np.testing.assert_array_equal(node.action_pub.messages[-1].data, np.arange(7) + 1000.0) + + def test_float64_state_and_action_are_not_quantized(self): + node = _policy_harness() + values = np.array( + [math.pi, np.nextafter(1.0, 2.0), 0.1234567890123, -0.987654321098, 0.1, -0.2, 0.3, -0.4], + dtype=np.float64, + ) + OpenPINode._on_state(node, _state(5, 9, values)) + self.assertEqual(node.state[2].dtype, np.float64) + np.testing.assert_array_equal(node.state[2], values) + + OpenPINode._publish_action(node, values[:7], 5, 9) + published = np.asarray(node.action_pub.messages[-1].data, dtype=np.float64) + np.testing.assert_array_equal(published, values[:7]) + self.assertEqual(parse_action_identity(node.action_pub.messages[-1].layout.dim[0].label), (5, 9, 0)) + + +class OfficialLiberoParityTest(unittest.TestCase): + def test_state_matches_official_quaternion_conversion_in_float64(self): + quaternion = np.array([0.11, -0.23, 0.37, 0.88], dtype=np.float64) + denominator = np.sqrt(1.0 - quaternion[3] * quaternion[3]) + expected_axis_angle = (quaternion[:3] * 2.0 * math.acos(quaternion[3])) / denominator + np.testing.assert_array_equal(quat_to_axis_angle(quaternion), expected_axis_angle) + + observation = { + "robot0_eef_pos": np.array([0.1, 0.2, 0.3], dtype=np.float64), + "robot0_eef_quat": quaternion, + "robot0_gripper_qpos": np.array([0.04, -0.04], dtype=np.float64), + } + state = LiberoEnv._state(object(), observation) + expected_state = np.concatenate([observation["robot0_eef_pos"], expected_axis_angle, observation["robot0_gripper_qpos"]]) + self.assertEqual(state.dtype, np.float64) + np.testing.assert_array_equal(state, expected_state) + + def test_simulator_consumes_float64_action_without_cast(self): + action = np.array( + [math.pi, np.nextafter(1.0, 2.0), 0.1234567890123, -0.987654321098, 0.1, -0.2, 0.3], + dtype=np.float64, + ) + + class Env: + accepted_action_dims = (7,) + + def step(self, received): + self.received = received.copy() + return object(), False, False + + env = Env() + node = types.SimpleNamespace( + state=RUNNING, + numeric_dtype=np.float64, + env=env, + _in_env_step=False, + step_index=4, + episode_step=2, + max_episode_steps=220, + success=False, + get_logger=lambda: _Logger(), + publish_observation=lambda: None, + publish_status=lambda: None, + _finish_episode=lambda _outcome: None, + ) + message = Float64MultiArray(data=action.tolist()) + SimulatorNode.on_action(node, message) + + self.assertEqual(env.received.dtype, np.float64) + np.testing.assert_array_equal(env.received, action) + + def test_simulator_rejects_action_from_an_old_plan_epoch(self): + class Env: + accepted_action_dims = (7,) + + def step(self, _received): + raise AssertionError("stale action must not reach the environment") + + logger = _Logger() + node = types.SimpleNamespace( + state=RUNNING, + numeric_dtype=np.float64, + env=Env(), + episode_index=2, + step_index=9, + plan_epoch=1, + get_logger=lambda: logger, + ) + message = Float64MultiArray(data=[0.0] * 7) + message.layout.dim = [MultiArrayDimension(label="episode=2;observation=9;plan_epoch=0", size=7, stride=7)] + + SimulatorNode.on_action(node, message) + + self.assertEqual(len(logger.warnings), 1) + + def test_simulator_rejects_malformed_tagged_action(self): + class Env: + accepted_action_dims = (7,) + + def step(self, _received): + raise AssertionError("malformed action must not reach the environment") + + logger = _Logger() + node = types.SimpleNamespace( + state=RUNNING, + numeric_dtype=np.float64, + env=Env(), + episode_index=2, + step_index=9, + plan_epoch=1, + get_logger=lambda: logger, + ) + message = Float64MultiArray(data=[0.0] * 7) + message.layout.dim = [MultiArrayDimension(label="episode=2;observation=9", size=7, stride=7)] + + SimulatorNode.on_action(node, message) + + self.assertEqual(len(logger.warnings), 1) + + def test_official_image_orientation_survives_ros_round_trip(self): + agentview = np.arange(9 * 13 * 3, dtype=np.uint8).reshape(9, 13, 3) + wrist = np.bitwise_xor(agentview, np.uint8(255)) + env = object.__new__(LiberoEnv) + env.contract = types.SimpleNamespace(cameras=("agentview", "wrist")) + env.observer = types.SimpleNamespace( + obs={ + "agentview_image": agentview, + "robot0_eye_in_hand_image": wrist, + "robot0_eef_pos": np.zeros(3), + "robot0_eef_quat": np.array([0.0, 0.0, 0.0, 1.0]), + "robot0_gripper_qpos": np.zeros(2), + } + ) + + observation = env._observation() + for camera, source in (("agentview", agentview), ("wrist", wrist)): + message = rgb_to_image_msg(observation.images[camera], Time(), camera, 0, 0) + np.testing.assert_array_equal(image_msg_to_rgb(message), source[::-1, ::-1]) + + def test_resize_matches_official_openpi_client(self): + client_source = PROJECT_ROOT.parent / "openpi/packages/openpi-client/src" + sys.path.insert(0, str(client_source)) + from openpi_client import image_tools + + raw = np.arange(157 * 256 * 3, dtype=np.uint8).reshape(157, 256, 3) + ros_image = image_msg_to_rgb(rgb_to_image_msg(np.ascontiguousarray(raw[::-1, ::-1]), Time(), "agentview", 2, 19)) + expected = image_tools.resize_with_pad(raw[::-1, ::-1], 224, 224) + actual = _resize_with_pad(ros_image, 224) + + np.testing.assert_array_equal(actual, expected) + + def test_action_unnormalization_and_ros_transport_match_official_formula(self): + q01 = np.linspace(-0.8, -0.2, 7, dtype=np.float64) + q99 = np.linspace(0.3, 1.1, 7, dtype=np.float64) + stats = { + "norm_stats": { + "state": {"q01": [0.0] * 8, "q99": [1.0] * 8}, + "actions": {"q01": q01.tolist(), "q99": q99.tolist()}, + } + } + normalized = np.linspace(-1.0, 1.0, 10 * 32, dtype=np.float32).reshape(1, 10, 32) + + with tempfile.TemporaryDirectory() as directory: + stats_path = Path(directory) / "norm_stats.json" + stats_path.write_text(json.dumps(stats), encoding="utf-8") + import torch + + actions = OpenPIPostInfer(stats_path).infer(torch.from_numpy(normalized)) + + expected = (normalized[0, :, :7] + 1.0) / 2.0 * (q99 - q01 + 1e-6) + q01 + self.assertEqual(actions.dtype, np.float64) + np.testing.assert_array_equal(actions, expected) + + node = _policy_harness() + OpenPINode._publish_action(node, actions[0], 2, 19) + transported = np.asarray(node.action_pub.messages[-1].data, dtype=np.float64) + np.testing.assert_array_equal(transported, expected[0]) + + +if __name__ == "__main__": + unittest.main() From 06e68787e35a5b339c33dc6a86329ac8e2989d5a Mon Sep 17 00:00:00 2001 From: Chernobyllight Date: Fri, 4 Sep 2026 08:20:12 +0000 Subject: [PATCH 3/6] refactor(openpi): streamline integration support files --- lightx2v/models/networks/openpi/NOTICE.md | 11 - .../1_convert_pi05_libero_to_pytorch.sh | 2 +- scripts/openpi/2_setup_pytorch_runtime.sh | 2 +- .../run_libero_evaluate_parallel_i2va.sh | 2 +- scripts/openpi/{ => support}/README.md | 16 +- .../{ => support}/convert_jax_checkpoint.py | 8 +- .../openpi/{ => support}/libero_summary.py | 2 +- scripts/openpi/{ => support}/runtime.py | 11 +- .../openpi/tests/test_ros_openpi_contract.py | 418 ------------------ .../openpi/tests/test_task_inputs_manifest.py | 166 ------- .../openpi/tests/validate_pytorch_parity.py | 375 ---------------- 11 files changed, 21 insertions(+), 992 deletions(-) delete mode 100644 lightx2v/models/networks/openpi/NOTICE.md rename scripts/openpi/{ => support}/README.md (96%) rename scripts/openpi/{ => support}/convert_jax_checkpoint.py (96%) rename scripts/openpi/{ => support}/libero_summary.py (98%) rename scripts/openpi/{ => support}/runtime.py (98%) delete mode 100644 scripts/openpi/tests/test_ros_openpi_contract.py delete mode 100644 scripts/openpi/tests/test_task_inputs_manifest.py delete mode 100644 scripts/openpi/tests/validate_pytorch_parity.py diff --git a/lightx2v/models/networks/openpi/NOTICE.md b/lightx2v/models/networks/openpi/NOTICE.md deleted file mode 100644 index ce898ed15..000000000 --- a/lightx2v/models/networks/openpi/NOTICE.md +++ /dev/null @@ -1,11 +0,0 @@ -# OpenPI attribution - -The `pi0.py`, `gemma.py`, and `preprocessing.py` implementation in this -directory is adapted from Physical Intelligence's OpenPI project at commit -`15a9616a00943ada6c20a0f158e3adb39df2ccac`. The files under -`transformers_replace/` are vendored from that revision without behavioral changes. - -OpenPI and the copied Hugging Face Transformers source files are distributed -under the Apache License 2.0. The localization changes replace OpenPI/JAX -imports with LightX2V-local, PyTorch-only modules while intentionally retaining -the official model parameter names for SafeTensors compatibility. diff --git a/scripts/openpi/1_convert_pi05_libero_to_pytorch.sh b/scripts/openpi/1_convert_pi05_libero_to_pytorch.sh index f97e9f253..7714802fb 100755 --- a/scripts/openpi/1_convert_pi05_libero_to_pytorch.sh +++ b/scripts/openpi/1_convert_pi05_libero_to_pytorch.sh @@ -5,4 +5,4 @@ script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" openpi_root="${OPENPI_PATH:-$(cd -- "${script_dir}/../../.." && pwd)/openpi}" python_bin="${OPENPI_CONVERT_PYTHON:-${openpi_root}/.venv/bin/python}" -exec "${python_bin}" "${script_dir}/convert_jax_checkpoint.py" "$@" +exec "${python_bin}" "${script_dir}/support/convert_jax_checkpoint.py" "$@" diff --git a/scripts/openpi/2_setup_pytorch_runtime.sh b/scripts/openpi/2_setup_pytorch_runtime.sh index cc7d8da8b..640fac7f4 100755 --- a/scripts/openpi/2_setup_pytorch_runtime.sh +++ b/scripts/openpi/2_setup_pytorch_runtime.sh @@ -11,4 +11,4 @@ elif [[ "${1:-}" == "setup" || "${1:-}" == "prepare" ]]; then shift fi -exec python "${script_dir}/runtime.py" "${command}" "$@" +exec python "${script_dir}/support/runtime.py" "${command}" "$@" diff --git a/scripts/openpi/run_libero_evaluate_parallel_i2va.sh b/scripts/openpi/run_libero_evaluate_parallel_i2va.sh index aded05256..5643b7103 100755 --- a/scripts/openpi/run_libero_evaluate_parallel_i2va.sh +++ b/scripts/openpi/run_libero_evaluate_parallel_i2va.sh @@ -105,7 +105,7 @@ worker_pids=() trap - EXIT HUP INT TERM summary_command=( - python "${script_dir}/libero_summary.py" + python "${script_dir}/support/libero_summary.py" --output-root "${output_root}" ) for suite in "${suites[@]}"; do diff --git a/scripts/openpi/README.md b/scripts/openpi/support/README.md similarity index 96% rename from scripts/openpi/README.md rename to scripts/openpi/support/README.md index 47c647729..4be24717b 100644 --- a/scripts/openpi/README.md +++ b/scripts/openpi/support/README.md @@ -12,6 +12,9 @@ shell -> python -m lightx2v.infer -> OpenPIRunner topic 直连，同样不启动 OpenPI server。两条路径都使用当前 base Python；仅任务 特异的 Transformers 代码放在私有 overlay 中。 +模型实现改编自 Physical Intelligence OpenPI（Apache-2.0）；随仓库提供的 +Transformers replacement 文件保留原始版权和许可证声明。 + 以下命令默认在项目根目录执行： ```bash @@ -167,7 +170,7 @@ save_results/pi05_libero_pytorch_fp32_parallel_evaluation/ └── parallel_summary.json ``` -`libero_summary.py` 只读取各 suite 的 JSON 结果并生成最终汇总，不是推理入口。 +`support/libero_summary.py` 只读取各 suite 的 JSON 结果并生成最终汇总，不是推理入口。 输出锁会阻止两个并行任务同时写入同一目录。 ## 6. ROS 单 episode 交互 @@ -387,6 +390,7 @@ coordinator 当前按完整 suite 写新目录，不提供中途恢复。 | 环境变量 | 作用 | | --- | --- | +| `OPENPI_DATA_ROOT` | checkpoint、tokenizer 和 Transformers overlay 的数据根目录 | | `OPENPI_MODEL_PATH` | PyTorch checkpoint | | `OPENPI_CONFIG` | 模型 JSON | | `OPENPI_EVAL_CONFIG` | 评测协议 JSON | @@ -406,19 +410,15 @@ coordinator 当前按完整 suite 写新目录，不提供中途恢复。 | `OPENPI_ROS_OUTPUT_DIR` | ROS suite 结果根目录 | | `OPENPI_ROS_OVERWRITE` | 是否覆盖已有 ROS suite 结果，默认 `false` | -## 10. 开发检查与验收 +## 10. 开发检查 ```bash bash -n scripts/openpi/*.sh -python -m unittest discover -s scripts/openpi/tests -p 'test_*.py' -v -python scripts/openpi/tests/validate_pytorch_parity.py --self-check pre-commit run --all-files ``` 数值路径的关键约束是官方图像 resize/uint8 量化、FP64 动作反归一化、连续 policy RNG 和 5-action replan queue；清理启动脚本时不应改变这些逻辑。 -ROS 契约测试覆盖乱序/重复 topic、10→5 action queue、float64 传输、官方 -旋转/resize/state/action 变换和迟到 action 拒绝。修改 ROS 文件后还应重新构建 -`common simulator inference`；LIBERO adapter 会把可见的物理 GPU 映射为 EGL -逻辑设备 0。 +修改 ROS 文件后需要重新构建 `common simulator inference`；LIBERO adapter 会把 +可见的物理 GPU 映射为 EGL 逻辑设备 0。 diff --git a/scripts/openpi/convert_jax_checkpoint.py b/scripts/openpi/support/convert_jax_checkpoint.py similarity index 96% rename from scripts/openpi/convert_jax_checkpoint.py rename to scripts/openpi/support/convert_jax_checkpoint.py index 870a5b68d..669ac9c4f 100755 --- a/scripts/openpi/convert_jax_checkpoint.py +++ b/scripts/openpi/support/convert_jax_checkpoint.py @@ -16,13 +16,13 @@ from safetensors import safe_open -PROJECT_ROOT = Path(__file__).resolve().parents[2] +PROJECT_ROOT = Path(__file__).resolve().parents[3] WORKSPACE_ROOT = PROJECT_ROOT.parent -OPENPI_DATA_ROOT = WORKSPACE_ROOT / "openpi_data" +OPENPI_DATA_ROOT = Path(os.environ.get("OPENPI_DATA_ROOT", str(WORKSPACE_ROOT / "openpi_data"))).expanduser().resolve() DEFAULT_OPENPI_ROOT = WORKSPACE_ROOT / "openpi" DEFAULT_SOURCE = OPENPI_DATA_ROOT / "openpi-assets/checkpoints/pi05_libero" DEFAULT_TOKENIZER = OPENPI_DATA_ROOT / "big_vision/paligemma_tokenizer.model" -DEFAULT_TRANSFORMERS_RUNTIME = OPENPI_DATA_ROOT / "python_deps/openpi_official_pytorch_runtime" +DEFAULT_TRANSFORMERS_OVERLAY = OPENPI_DATA_ROOT / "python_deps/openpi_official_pytorch_runtime" EXPECTED_TENSORS = 812 DTYPE_BY_PRECISION = {"float32": "F32", "bfloat16": "BF16"} @@ -155,7 +155,7 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--tokenizer", default=os.environ.get("OPENPI_TOKENIZER_PATH", str(DEFAULT_TOKENIZER))) parser.add_argument( "--transformers-runtime", - default=os.environ.get("OPENPI_TRANSFORMERS_RUNTIME_PATH", str(DEFAULT_TRANSFORMERS_RUNTIME)), + default=os.environ.get("OPENPI_TRANSFORMERS_RUNTIME_PATH", str(DEFAULT_TRANSFORMERS_OVERLAY)), ) parser.add_argument("--dry-run", action="store_true") return parser diff --git a/scripts/openpi/libero_summary.py b/scripts/openpi/support/libero_summary.py similarity index 98% rename from scripts/openpi/libero_summary.py rename to scripts/openpi/support/libero_summary.py index bcb595f51..93768ae8d 100755 --- a/scripts/openpi/libero_summary.py +++ b/scripts/openpi/support/libero_summary.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -"""Aggregate completed LIBERO suite summaries without launching inference.""" +"""Aggregate LIBERO suite summaries without launching inference.""" from __future__ import annotations diff --git a/scripts/openpi/runtime.py b/scripts/openpi/support/runtime.py similarity index 98% rename from scripts/openpi/runtime.py rename to scripts/openpi/support/runtime.py index 72e0d309f..dc0c201f7 100755 --- a/scripts/openpi/runtime.py +++ b/scripts/openpi/support/runtime.py @@ -21,12 +21,11 @@ import tempfile from pathlib import Path -PROJECT_ROOT = Path(__file__).resolve().parents[2] +PROJECT_ROOT = Path(__file__).resolve().parents[3] WORKSPACE_ROOT = PROJECT_ROOT.parent -OPENPI_DATA_ROOT = WORKSPACE_ROOT / "openpi_data" -DEFAULT_PYTHON = Path("/opt/conda/bin/python") +OPENPI_DATA_ROOT = Path(os.environ.get("OPENPI_DATA_ROOT", str(WORKSPACE_ROOT / "openpi_data"))).expanduser().resolve() DEFAULT_MODEL = OPENPI_DATA_ROOT / "openpi-assets/checkpoints/pi05_libero_pytorch_fp32" -DEFAULT_TRANSFORMERS_RUNTIME = OPENPI_DATA_ROOT / "python_deps/openpi_official_pytorch_runtime" +DEFAULT_TRANSFORMERS_OVERLAY = OPENPI_DATA_ROOT / "python_deps/openpi_official_pytorch_runtime" DEFAULT_LIBERO_ROOT = WORKSPACE_ROOT / "openpi/third_party/libero" DEFAULT_MODEL_CONFIG = PROJECT_ROOT / "configs/openpi/pi05_libero.json" DEFAULT_EVAL_CONFIG = PROJECT_ROOT / "configs/openpi/pi05_libero_eval.json" @@ -450,7 +449,7 @@ def _check_runtime(args: argparse.Namespace) -> None: def _add_paths(parser: argparse.ArgumentParser) -> None: parser.add_argument( "--transformers-runtime", - default=os.environ.get("OPENPI_TRANSFORMERS_RUNTIME_PATH", str(DEFAULT_TRANSFORMERS_RUNTIME)), + default=os.environ.get("OPENPI_TRANSFORMERS_RUNTIME_PATH", str(DEFAULT_TRANSFORMERS_OVERLAY)), ) @@ -465,7 +464,7 @@ def build_parser() -> argparse.ArgumentParser: check = subparsers.add_parser("check", help="validate runtime, checkpoint, configs, and official LIBERO paths") _add_paths(check) - check.add_argument("--expected-python", default=os.environ.get("OPENPI_PYTHON", str(DEFAULT_PYTHON))) + check.add_argument("--expected-python", default=os.environ.get("OPENPI_PYTHON", sys.executable)) check.add_argument("--model-path", default=os.environ.get("OPENPI_MODEL_PATH", str(DEFAULT_MODEL))) check.add_argument("--model-config", default=os.environ.get("OPENPI_CONFIG", str(DEFAULT_MODEL_CONFIG))) check.add_argument("--eval-config", default=os.environ.get("OPENPI_EVAL_CONFIG", str(DEFAULT_EVAL_CONFIG))) diff --git a/scripts/openpi/tests/test_ros_openpi_contract.py b/scripts/openpi/tests/test_ros_openpi_contract.py deleted file mode 100644 index df88e4062..000000000 --- a/scripts/openpi/tests/test_ros_openpi_contract.py +++ /dev/null @@ -1,418 +0,0 @@ -"""Offline contract tests for the OpenPI ROS bridge. - -Run after sourcing the ROS underlay and LightX2V overlay. The tests exercise -message conversion and callback ordering without starting DDS participants or -loading model weights. -""" - -from __future__ import annotations - -import json -import math -import sys -import tempfile -import types -import unittest -from collections import deque -from pathlib import Path - -import numpy as np - -PROJECT_ROOT = Path(__file__).resolve().parents[3] -ROS_SOURCE_ROOT = PROJECT_ROOT / "lightx2v_ros/src" -for package in ("common", "inference", "simulator"): - sys.path.insert(0, str(ROS_SOURCE_ROOT / package)) - - -def _stub_lightx2v_imports(): - """Keep helper tests independent of LightX2V CUDA initialization.""" - modules = { - "lightx2v": PROJECT_ROOT / "lightx2v", - "lightx2v.models": PROJECT_ROOT / "lightx2v/models", - "lightx2v.models.networks": PROJECT_ROOT / "lightx2v/models/networks", - "lightx2v.models.networks.openpi": PROJECT_ROOT / "lightx2v/models/networks/openpi", - "lightx2v.models.runners": PROJECT_ROOT / "lightx2v/models/runners", - "lightx2v.models.runners.openpi": PROJECT_ROOT / "lightx2v/models/runners/openpi", - "lightx2v.utils": PROJECT_ROOT / "lightx2v/utils", - } - for name, path in modules.items(): - module = types.ModuleType(name) - module.__path__ = [str(path)] - sys.modules[name] = module - - runner = types.ModuleType("lightx2v.models.runners.openpi.openpi_runner") - runner.OpenPIPolicy = object - sys.modules[runner.__name__] = runner - - config = types.ModuleType("lightx2v.utils.set_config") - config.auto_calc_config = lambda values: values - config.get_default_config = dict - sys.modules[config.__name__] = config - - -try: - from builtin_interfaces.msg import Time - from common.contract import LIBERO_CONTRACT - from sensor_msgs.msg import Image - from std_msgs.msg import Float64MultiArray, MultiArrayDimension, String - - _stub_lightx2v_imports() - from inference.openpi_node.main import OpenPINode, image_msg_to_rgb, observation_identity, state_identity - from simulator.libero_node.env import LiberoEnv, quat_to_axis_angle - from simulator.sim.node import RUNNING, SimulatorNode, parse_action_identity, rgb_to_image_msg - - from lightx2v.models.networks.openpi.infer.post_infer import OpenPIPostInfer - from lightx2v.models.networks.openpi.infer.pre_infer import _resize_with_pad -except ImportError as exc: # pragma: no cover - depends on sourced ROS setup - raise unittest.SkipTest(f"source the ROS underlay before running this test: {exc}") from exc - - -class _Logger: - def __init__(self): - self.errors = [] - self.warnings = [] - - def error(self, message): - self.errors.append(str(message)) - - def info(self, _message): - return None - - def warning(self, message): - self.warnings.append(str(message)) - - -class _Publisher: - def __init__(self): - self.messages = [] - - def publish(self, message): - self.messages.append(message) - - -class _ChunkPolicy: - output_action_dim = 7 - - def __init__(self): - self.calls = [] - - def predict_action_chunk(self, *, images, state, task_description): - call_index = len(self.calls) - self.calls.append( - { - "images": {name: image.copy() for name, image in images.items()}, - "state": state.copy(), - "task_description": task_description, - } - ) - return np.arange(70, dtype=np.float64).reshape(10, 7) + call_index * 1000.0 - - -def _policy_harness(): - node = types.SimpleNamespace( - contract=LIBERO_CONTRACT, - numeric_message_type=Float64MultiArray, - numeric_dtype=np.float64, - images={camera: None for camera in LIBERO_CONTRACT.policy_input_cameras}, - state=None, - task_description="", - episode_index=None, - plan_epoch=0, - pending_observation=None, - last_processed_observation=None, - pending_actions=deque(), - actions_per_plan=5, - policy=_ChunkPolicy(), - action_pub=_Publisher(), - get_logger=lambda: _Logger(), - ) - node._try_process_observation = types.MethodType(OpenPINode._try_process_observation, node) - node._publish_action = types.MethodType(OpenPINode._publish_action, node) - return node - - -def _context(episode, observation, prompt="pick up the bowl", plan_epoch=0): - message = String() - message.data = json.dumps( - { - "episode": episode, - "observation": observation, - "plan_epoch": plan_epoch, - "task_description": prompt, - } - ) - return message - - -def _state(episode, observation, values=None): - message = Float64MultiArray() - message.layout.dim = [MultiArrayDimension(label=f"episode={episode};observation={observation}", size=8, stride=8)] - message.data = (np.arange(8, dtype=np.float64) if values is None else np.asarray(values)).tolist() - return message - - -def _image(camera, episode, observation, value=0): - pixels = np.full((4, 5, 3), value, dtype=np.uint8) - return rgb_to_image_msg(pixels, Time(), camera, episode, observation) - - -def _deliver_observation(node, episode, observation, plan_epoch=0): - OpenPINode._on_context(node, _context(episode, observation, plan_epoch=plan_epoch)) - OpenPINode._image_callback(node, "wrist")(_image("wrist", episode, observation, 29)) - OpenPINode._on_state(node, _state(episode, observation)) - OpenPINode._image_callback(node, "agentview")(_image("agentview", episode, observation, 11)) - - -class RosMessageContractTest(unittest.TestCase): - def test_tagged_rgb_round_trip(self): - image = np.arange(5 * 7 * 3, dtype=np.uint8).reshape(5, 7, 3)[:, ::-1] - message = rgb_to_image_msg(image, Time(), "agentview", 13, 21) - - self.assertEqual(message.header.frame_id, "agentview|13|21") - self.assertEqual(observation_identity(message.header.frame_id), (13, 21)) - np.testing.assert_array_equal(image_msg_to_rgb(message), image) - - def test_bgr_message_with_row_padding(self): - rgb = np.arange(2 * 3 * 3, dtype=np.uint8).reshape(2, 3, 3) - row_bytes = 3 * 3 + 4 - encoded = np.full((2, row_bytes), 255, dtype=np.uint8) - encoded[:, :9] = rgb[:, :, ::-1].reshape(2, 9) - message = Image(height=2, width=3, encoding="bgr8", step=row_bytes, data=encoded.tobytes()) - - np.testing.assert_array_equal(image_msg_to_rgb(message), rgb) - - def test_identity_parser_rejects_untagged_and_malformed_frames(self): - self.assertIsNone(observation_identity("agentview")) - self.assertIsNone(observation_identity("agentview|episode|2")) - self.assertIsNone(observation_identity("agentview|1")) - self.assertEqual(state_identity("episode=5;observation=91"), (5, 91)) - self.assertIsNone(state_identity("episode=5")) - self.assertIsNone(state_identity("episode=5;observation=bad")) - - -class OpenPINodeSynchronizationTest(unittest.TestCase): - def test_waits_for_one_matching_set_and_ignores_duplicate_delivery(self): - node = _policy_harness() - OpenPINode._on_context(node, _context(3, 8)) - OpenPINode._on_state(node, _state(3, 7)) - OpenPINode._image_callback(node, "agentview")(_image("agentview", 3, 8)) - OpenPINode._image_callback(node, "wrist")(_image("wrist", 3, 8)) - self.assertEqual(len(node.action_pub.messages), 0) - - OpenPINode._on_state(node, _state(3, 8)) - self.assertEqual(len(node.policy.calls), 1) - self.assertEqual(len(node.action_pub.messages), 1) - - OpenPINode._on_context(node, _context(3, 8)) - OpenPINode._on_state(node, _state(3, 8)) - self.assertEqual(len(node.policy.calls), 1) - self.assertEqual(len(node.action_pub.messages), 1) - - def test_executes_five_actions_from_each_ten_action_chunk(self): - node = _policy_harness() - for observation in range(6): - _deliver_observation(node, 0, observation) - - self.assertEqual(len(node.policy.calls), 2) - published = np.asarray([message.data for message in node.action_pub.messages]) - first_chunk = np.arange(70, dtype=np.float64).reshape(10, 7) - expected = np.concatenate([first_chunk[:5], first_chunk[:1] + 1000.0]) - np.testing.assert_array_equal(published, expected) - - def test_new_episode_discards_the_old_action_tail_without_resetting_policy(self): - node = _policy_harness() - _deliver_observation(node, 0, 0) - self.assertEqual(len(node.pending_actions), 4) - - _deliver_observation(node, 1, 1) - self.assertEqual(len(node.policy.calls), 2) - self.assertEqual(len(node.pending_actions), 4) - np.testing.assert_array_equal(node.action_pub.messages[-1].data, np.arange(7) + 1000.0) - - def test_new_plan_epoch_discards_actions_queued_before_resume(self): - node = _policy_harness() - _deliver_observation(node, 0, 0, plan_epoch=0) - self.assertEqual(len(node.pending_actions), 4) - - _deliver_observation(node, 0, 1, plan_epoch=1) - self.assertEqual(len(node.policy.calls), 2) - self.assertEqual(len(node.pending_actions), 4) - np.testing.assert_array_equal(node.action_pub.messages[-1].data, np.arange(7) + 1000.0) - - def test_float64_state_and_action_are_not_quantized(self): - node = _policy_harness() - values = np.array( - [math.pi, np.nextafter(1.0, 2.0), 0.1234567890123, -0.987654321098, 0.1, -0.2, 0.3, -0.4], - dtype=np.float64, - ) - OpenPINode._on_state(node, _state(5, 9, values)) - self.assertEqual(node.state[2].dtype, np.float64) - np.testing.assert_array_equal(node.state[2], values) - - OpenPINode._publish_action(node, values[:7], 5, 9) - published = np.asarray(node.action_pub.messages[-1].data, dtype=np.float64) - np.testing.assert_array_equal(published, values[:7]) - self.assertEqual(parse_action_identity(node.action_pub.messages[-1].layout.dim[0].label), (5, 9, 0)) - - -class OfficialLiberoParityTest(unittest.TestCase): - def test_state_matches_official_quaternion_conversion_in_float64(self): - quaternion = np.array([0.11, -0.23, 0.37, 0.88], dtype=np.float64) - denominator = np.sqrt(1.0 - quaternion[3] * quaternion[3]) - expected_axis_angle = (quaternion[:3] * 2.0 * math.acos(quaternion[3])) / denominator - np.testing.assert_array_equal(quat_to_axis_angle(quaternion), expected_axis_angle) - - observation = { - "robot0_eef_pos": np.array([0.1, 0.2, 0.3], dtype=np.float64), - "robot0_eef_quat": quaternion, - "robot0_gripper_qpos": np.array([0.04, -0.04], dtype=np.float64), - } - state = LiberoEnv._state(object(), observation) - expected_state = np.concatenate([observation["robot0_eef_pos"], expected_axis_angle, observation["robot0_gripper_qpos"]]) - self.assertEqual(state.dtype, np.float64) - np.testing.assert_array_equal(state, expected_state) - - def test_simulator_consumes_float64_action_without_cast(self): - action = np.array( - [math.pi, np.nextafter(1.0, 2.0), 0.1234567890123, -0.987654321098, 0.1, -0.2, 0.3], - dtype=np.float64, - ) - - class Env: - accepted_action_dims = (7,) - - def step(self, received): - self.received = received.copy() - return object(), False, False - - env = Env() - node = types.SimpleNamespace( - state=RUNNING, - numeric_dtype=np.float64, - env=env, - _in_env_step=False, - step_index=4, - episode_step=2, - max_episode_steps=220, - success=False, - get_logger=lambda: _Logger(), - publish_observation=lambda: None, - publish_status=lambda: None, - _finish_episode=lambda _outcome: None, - ) - message = Float64MultiArray(data=action.tolist()) - SimulatorNode.on_action(node, message) - - self.assertEqual(env.received.dtype, np.float64) - np.testing.assert_array_equal(env.received, action) - - def test_simulator_rejects_action_from_an_old_plan_epoch(self): - class Env: - accepted_action_dims = (7,) - - def step(self, _received): - raise AssertionError("stale action must not reach the environment") - - logger = _Logger() - node = types.SimpleNamespace( - state=RUNNING, - numeric_dtype=np.float64, - env=Env(), - episode_index=2, - step_index=9, - plan_epoch=1, - get_logger=lambda: logger, - ) - message = Float64MultiArray(data=[0.0] * 7) - message.layout.dim = [MultiArrayDimension(label="episode=2;observation=9;plan_epoch=0", size=7, stride=7)] - - SimulatorNode.on_action(node, message) - - self.assertEqual(len(logger.warnings), 1) - - def test_simulator_rejects_malformed_tagged_action(self): - class Env: - accepted_action_dims = (7,) - - def step(self, _received): - raise AssertionError("malformed action must not reach the environment") - - logger = _Logger() - node = types.SimpleNamespace( - state=RUNNING, - numeric_dtype=np.float64, - env=Env(), - episode_index=2, - step_index=9, - plan_epoch=1, - get_logger=lambda: logger, - ) - message = Float64MultiArray(data=[0.0] * 7) - message.layout.dim = [MultiArrayDimension(label="episode=2;observation=9", size=7, stride=7)] - - SimulatorNode.on_action(node, message) - - self.assertEqual(len(logger.warnings), 1) - - def test_official_image_orientation_survives_ros_round_trip(self): - agentview = np.arange(9 * 13 * 3, dtype=np.uint8).reshape(9, 13, 3) - wrist = np.bitwise_xor(agentview, np.uint8(255)) - env = object.__new__(LiberoEnv) - env.contract = types.SimpleNamespace(cameras=("agentview", "wrist")) - env.observer = types.SimpleNamespace( - obs={ - "agentview_image": agentview, - "robot0_eye_in_hand_image": wrist, - "robot0_eef_pos": np.zeros(3), - "robot0_eef_quat": np.array([0.0, 0.0, 0.0, 1.0]), - "robot0_gripper_qpos": np.zeros(2), - } - ) - - observation = env._observation() - for camera, source in (("agentview", agentview), ("wrist", wrist)): - message = rgb_to_image_msg(observation.images[camera], Time(), camera, 0, 0) - np.testing.assert_array_equal(image_msg_to_rgb(message), source[::-1, ::-1]) - - def test_resize_matches_official_openpi_client(self): - client_source = PROJECT_ROOT.parent / "openpi/packages/openpi-client/src" - sys.path.insert(0, str(client_source)) - from openpi_client import image_tools - - raw = np.arange(157 * 256 * 3, dtype=np.uint8).reshape(157, 256, 3) - ros_image = image_msg_to_rgb(rgb_to_image_msg(np.ascontiguousarray(raw[::-1, ::-1]), Time(), "agentview", 2, 19)) - expected = image_tools.resize_with_pad(raw[::-1, ::-1], 224, 224) - actual = _resize_with_pad(ros_image, 224) - - np.testing.assert_array_equal(actual, expected) - - def test_action_unnormalization_and_ros_transport_match_official_formula(self): - q01 = np.linspace(-0.8, -0.2, 7, dtype=np.float64) - q99 = np.linspace(0.3, 1.1, 7, dtype=np.float64) - stats = { - "norm_stats": { - "state": {"q01": [0.0] * 8, "q99": [1.0] * 8}, - "actions": {"q01": q01.tolist(), "q99": q99.tolist()}, - } - } - normalized = np.linspace(-1.0, 1.0, 10 * 32, dtype=np.float32).reshape(1, 10, 32) - - with tempfile.TemporaryDirectory() as directory: - stats_path = Path(directory) / "norm_stats.json" - stats_path.write_text(json.dumps(stats), encoding="utf-8") - import torch - - actions = OpenPIPostInfer(stats_path).infer(torch.from_numpy(normalized)) - - expected = (normalized[0, :, :7] + 1.0) / 2.0 * (q99 - q01 + 1e-6) + q01 - self.assertEqual(actions.dtype, np.float64) - np.testing.assert_array_equal(actions, expected) - - node = _policy_harness() - OpenPINode._publish_action(node, actions[0], 2, 19) - transported = np.asarray(node.action_pub.messages[-1].data, dtype=np.float64) - np.testing.assert_array_equal(transported, expected[0]) - - -if __name__ == "__main__": - unittest.main() diff --git a/scripts/openpi/tests/test_task_inputs_manifest.py b/scripts/openpi/tests/test_task_inputs_manifest.py deleted file mode 100644 index 2de820865..000000000 --- a/scripts/openpi/tests/test_task_inputs_manifest.py +++ /dev/null @@ -1,166 +0,0 @@ -#!/usr/bin/env python3 -"""Focused tests for the LIBERO task-input manifest migration.""" - -from __future__ import annotations - -import hashlib -import json -import sys -import tempfile -import types -import unittest -from pathlib import Path -from types import SimpleNamespace - -PROJECT_ROOT = Path(__file__).resolve().parents[3] - - -def _namespace_package(name: str, path: Path) -> None: - package = types.ModuleType(name) - package.__path__ = [str(path)] - sys.modules[name] = package - - -# Import only the evaluator protocol modules. Importing the public lightx2v -# package initializes the selected accelerator, which is unrelated to these -# CPU-only filesystem tests. -_namespace_package("lightx2v", PROJECT_ROOT / "lightx2v") -_namespace_package("lightx2v.models", PROJECT_ROOT / "lightx2v/models") -_namespace_package("lightx2v.models.runners", PROJECT_ROOT / "lightx2v/models/runners") -_namespace_package("lightx2v.models.runners.openpi", PROJECT_ROOT / "lightx2v/models/runners/openpi") - -from lightx2v.models.runners.openpi.libero_protocol import ( # noqa: E402 - EvaluationConfig, - TaskSpec, - build_task_inputs_manifest, - ensure_task_inputs_manifest, - resolved_protocol, -) - - -class TaskInputsManifestTest(unittest.TestCase): - def setUp(self) -> None: - self.temporary_directory = tempfile.TemporaryDirectory() - self.root = Path(self.temporary_directory.name).resolve() - self.bddl_path = self.root / "libero/libero/bddl_files/suite/task.bddl" - self.init_states_path = self.root / "libero/libero/init_files/suite/task.pruned_init" - self.bddl_path.parent.mkdir(parents=True) - self.init_states_path.parent.mkdir(parents=True) - self.bddl_path.write_bytes(b"bddl-v1") - self.init_states_path.write_bytes(b"init-v1") - self.spec = TaskSpec( - benchmark="suite", - task_id=0, - suite=None, - task=None, - bddl_path=self.bddl_path.resolve(), - init_states_path=self.init_states_path.resolve(), - ) - self.task_specs = {"suite": [self.spec]} - self.manifest_config = SimpleNamespace(benchmarks=("suite",), libero_root=self.root) - - def tearDown(self) -> None: - self.temporary_directory.cleanup() - - def _legacy_record(self, **updates: object) -> dict[str, object]: - record: dict[str, object] = { - "schema_version": 1, - "bddl_file": str(self.bddl_path), - "init_states_file": str(self.init_states_path), - "init_states_loader": "direct_file", - } - record.update(updates) - return record - - def test_new_manifest_is_stable_and_content_sensitive(self) -> None: - first = build_task_inputs_manifest(self.task_specs, self.manifest_config) - second = build_task_inputs_manifest(self.task_specs, self.manifest_config) - self.assertEqual(first, second) - self.assertEqual(first["task_count"], 1) - self.assertEqual(first["input_count"], 2) - - self.bddl_path.write_bytes(b"bddl-v2") - changed = build_task_inputs_manifest(self.task_specs, self.manifest_config) - self.assertNotEqual(first["manifest_sha256"], changed["manifest_sha256"]) - - def test_existing_manifest_is_verified_strictly(self) -> None: - output_dir = self.root / "new-output" - output_dir.mkdir() - manifest = build_task_inputs_manifest(self.task_specs, self.manifest_config) - path = ensure_task_inputs_manifest(output_dir, manifest, {}, self.task_specs) - self.assertEqual(json.loads(path.read_text(encoding="utf-8")), manifest) - self.assertEqual(ensure_task_inputs_manifest(output_dir, manifest, {}, self.task_specs), path) - - self.init_states_path.write_bytes(b"init-v2") - changed = build_task_inputs_manifest(self.task_specs, self.manifest_config) - with self.assertRaisesRegex(RuntimeError, "content differs"): - ensure_task_inputs_manifest(output_dir, changed, {}, self.task_specs) - - def test_schema_one_records_are_adopted_after_path_validation(self) -> None: - output_dir = self.root / "legacy-output" - output_dir.mkdir() - manifest = build_task_inputs_manifest(self.task_specs, self.manifest_config) - records = {("suite", 0, 0): self._legacy_record()} - with self.assertLogs("lightx2v.models.runners.openpi.libero_protocol", level="WARNING"): - path = ensure_task_inputs_manifest(output_dir, manifest, records, self.task_specs) - self.assertTrue(path.is_file()) - - invalid_output = self.root / "invalid-legacy-output" - invalid_output.mkdir() - invalid_records = {("suite", 0, 0): self._legacy_record(bddl_file=str(self.root / "other.bddl"))} - with self.assertRaisesRegex(RuntimeError, "legacy episode"): - ensure_task_inputs_manifest(invalid_output, manifest, invalid_records, self.task_specs) - self.assertFalse((invalid_output / "task_inputs_manifest.json").exists()) - - def test_protocol_id_retains_the_schema_one_field_set(self) -> None: - model_path = self.root / "model" - norm_path = model_path / "assets/physical-intelligence/libero/norm_stats.json" - tokenizer_path = model_path / "assets/paligemma_tokenizer.model" - norm_path.parent.mkdir(parents=True) - (model_path / "model.safetensors").write_bytes(b"model") - norm_path.write_bytes(b"norm") - tokenizer_path.write_bytes(b"tokenizer") - config_json = self.root / "model.json" - config_json.write_bytes(b"{}") - config = EvaluationConfig( - benchmarks=("suite",), - task_ids={"suite": (0,)}, - num_trials_per_task=1, - env_seed=7, - policy_seed=0, - actions_per_plan=5, - num_steps_wait=10, - render_size=256, - video_fps=10, - video_policy="none", - save_actions=False, - fail_fast=False, - resume=True, - max_steps={"suite": 2}, - libero_root=self.root, - libero_config_dir=self.root / "runtime", - ) - resolved, protocol_id = resolved_protocol(config, model_path=model_path, config_json=config_json) - protocol_fields = { - key: value - for key, value in resolved.items() - if key - not in { - "protocol_id", - "protocol_name", - "official_protocol", - "libero_config_dir", - "video_fps", - "video_policy", - "save_actions", - "fail_fast", - "resume", - } - } - expected = hashlib.sha256(json.dumps(protocol_fields, sort_keys=True, separators=(",", ":")).encode("utf-8")).hexdigest() - self.assertEqual(protocol_id, expected) - self.assertNotIn("task_inputs_manifest", resolved) - - -if __name__ == "__main__": - unittest.main() diff --git a/scripts/openpi/tests/validate_pytorch_parity.py b/scripts/openpi/tests/validate_pytorch_parity.py deleted file mode 100644 index 066aef94b..000000000 --- a/scripts/openpi/tests/validate_pytorch_parity.py +++ /dev/null @@ -1,375 +0,0 @@ -"""Compare upstream OpenPI PyTorch and LightX2V with identical input and noise. - -Run this with OpenPI's patched conversion environment. The deployed LightX2V -runtime itself does not import OpenPI or JAX. -""" - -from __future__ import annotations - -import argparse -import dataclasses -import gc -import json -import sys -import types -from collections import Counter -from pathlib import Path - -import h5py -import numpy as np -import torch - -PROJECT_ROOT = Path(__file__).resolve().parents[3] -WORKSPACE_ROOT = PROJECT_ROOT.parent -OPENPI_DATA_ROOT = WORKSPACE_ROOT / "openpi_data" -PRECISIONS = ("bfloat16", "float32") - - -def load_sample(path: Path) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - with h5py.File(path, "r") as handle: - demo = handle["data/demo_0"] - image = np.asarray(demo["obs/agentview_rgb"][0], dtype=np.uint8) - wrist = np.asarray(demo["obs/eye_in_hand_rgb"][0], dtype=np.uint8) - state = np.concatenate([demo["obs/ee_pos"][0], demo["obs/ee_ori"][0], demo["obs/gripper_states"][0]]) - return image, wrist, state - - -def _install_lightx2v_package_stub() -> None: - """Avoid LightX2V platform initialization in this model-only validator.""" - package = types.ModuleType("lightx2v") - package.__path__ = [str(PROJECT_ROOT / "lightx2v")] - sys.modules["lightx2v"] = package - - -def _create_upstream_policy(train_config, checkpoint: Path, device: str, precision: str): - """Create the official policy while honoring its selectable torch precision. - - ``policy_config.create_trained_policy`` in OpenPI 15a9616 always applies - BF16 after loading. This is the same construction with that one choice made - explicit, so the validator can cover both official model modes. - """ - from openpi import transforms - from openpi.policies import policy as policy_module - from openpi.training import checkpoints - - weight_path = checkpoint / "model.safetensors" - model = train_config.model.load_pytorch(train_config, str(weight_path)) - model.paligemma_with_expert.to_bfloat16_for_selected_params(precision) - - data_config = train_config.data.create(train_config.assets_dirs, train_config.model) - if data_config.asset_id is None: - raise ValueError("The upstream LIBERO data config has no normalization asset id") - norm_stats = checkpoints.load_norm_stats(checkpoint / "assets", data_config.asset_id) - - return policy_module.Policy( - model, - transforms=[ - transforms.InjectDefaultPrompt(None), - *data_config.data_transforms.inputs, - transforms.Normalize(norm_stats, use_quantiles=data_config.use_quantile_norm), - *data_config.model_transforms.inputs, - ], - output_transforms=[ - *data_config.model_transforms.outputs, - transforms.Unnormalize(norm_stats, use_quantiles=data_config.use_quantile_norm), - *data_config.data_transforms.outputs, - ], - metadata=train_config.policy_metadata, - is_pytorch=True, - pytorch_device=device, - ) - - -def _parameter_dtype_counts(model: torch.nn.Module) -> dict[str, int]: - counts = Counter(str(parameter.dtype).removeprefix("torch.") for parameter in model.parameters() if parameter.is_floating_point()) - return dict(sorted(counts.items())) - - -def _validate_selected_precision(counts: dict[str, int], precision: str, label: str) -> None: - if precision == "float32": - if set(counts) != {"float32"}: - raise RuntimeError(f"{label} FP32 mode contains unexpected parameter dtypes: {counts}") - return - if "bfloat16" not in counts or "float32" not in counts: - raise RuntimeError(f"{label} BF16 mode must retain the official selected FP32 parameters: {counts}") - - -def _to_numpy(value) -> np.ndarray: - if isinstance(value, torch.Tensor): - return value.detach().cpu().numpy() - return np.asarray(value) - - -def _array_comparison(reference, candidate) -> dict[str, object]: - reference_array = _to_numpy(reference) - candidate_array = _to_numpy(candidate) - same_shape = reference_array.shape == candidate_array.shape - if same_shape and reference_array.size: - difference = reference_array.astype(np.float64) - candidate_array.astype(np.float64) - max_abs_error = float(np.max(np.abs(difference))) - else: - max_abs_error = 0.0 if same_shape else None - return { - "upstream_shape": list(reference_array.shape), - "lightx2v_shape": list(candidate_array.shape), - "upstream_dtype": str(reference_array.dtype), - "lightx2v_dtype": str(candidate_array.dtype), - "exact": bool(same_shape and np.array_equal(reference_array, candidate_array)), - "max_abs_error": max_abs_error, - } - - -def _observation_comparison(upstream, local) -> dict[str, object]: - components = { - **{f"images/{key}": _array_comparison(upstream.images[key], local.images[key]) for key in upstream.images}, - **{f"image_masks/{key}": _array_comparison(upstream.image_masks[key], local.image_masks[key]) for key in upstream.image_masks}, - "state": _array_comparison(upstream.state, local.state), - "tokenized_prompt": _array_comparison(upstream.tokenized_prompt, local.tokenized_prompt), - "tokenized_prompt_mask": _array_comparison( - upstream.tokenized_prompt_mask, - local.tokenized_prompt_mask, - ), - } - expected_image_keys = set(local.images) - expected_mask_keys = set(local.image_masks) - keys_match = set(upstream.images) == expected_image_keys and set(upstream.image_masks) == expected_mask_keys - exact = keys_match and all(bool(component["exact"]) for component in components.values()) - return { - "exact": exact, - "image_keys_match": keys_match, - "components": components, - } - - -def _to_torch_batch(data, device: str): - if isinstance(data, dict): - return {key: _to_torch_batch(value, device) for key, value in data.items()} - return torch.from_numpy(np.array(data)).to(device)[None, ...] - - -def run_self_check() -> None: - """Exercise precision selection and the exact uint8 resize adapter.""" - from openpi_client import image_tools as upstream_image_tools - - _install_lightx2v_package_stub() - from lightx2v.models.networks.openpi.config import Pi0Config - from lightx2v.models.networks.openpi.infer.pre_infer import _resize_with_pad - - values = { - "action_dim": 32, - "action_horizon": 10, - "max_token_len": 200, - "paligemma_variant": "gemma_2b", - "action_expert_variant": "gemma_300m", - "pi05": True, - "discrete_state_input": False, - "pytorch_compile_mode": None, - } - for precision in PRECISIONS: - config = Pi0Config.from_mapping({**values, "dtype": precision}) - config.validate_pi05_libero() - if config.dtype != precision: - raise AssertionError(f"precision selection changed {precision!r} to {config.dtype!r}") - - image = np.arange(17 * 29 * 3, dtype=np.uint8).reshape(17, 29, 3) - upstream = upstream_image_tools.resize_with_pad(image, 24, 24) - local = _resize_with_pad(image, 24) - if not np.array_equal(local, upstream): - difference = np.abs(local.astype(np.int16) - upstream.astype(np.int16)) - raise AssertionError(f"resize-with-pad mismatch: max_abs={difference.max()}") - - print(json.dumps({"precision_selection": list(PRECISIONS), "uint8_resize_exact": True}, indent=2)) - - -def build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser() - parser.add_argument( - "--checkpoint", - type=Path, - default=OPENPI_DATA_ROOT / "openpi-assets/checkpoints/pi05_libero_pytorch_fp32", - ) - parser.add_argument( - "--config", - type=Path, - default=PROJECT_ROOT / "configs/openpi/pi05_libero.json", - ) - parser.add_argument( - "--sample", - type=Path, - default=OPENPI_DATA_ROOT / "raw/huggingface/yifengzhu-hf/LIBERO-datasets/libero_spatial" / "pick_up_the_black_bowl_between_the_plate_and_the_ramekin_and_place_it_on_the_plate_demo.hdf5", - ) - parser.add_argument( - "--output", - type=Path, - default=OPENPI_DATA_ROOT / "results/pi05_libero_pytorch_parity.json", - ) - parser.add_argument("--device", default="cuda") - parser.add_argument("--precision", choices=PRECISIONS, help="Override the dtype selected by the model JSON") - parser.add_argument("--atol", type=float, default=1e-6) - parser.add_argument("--self-check", action="store_true", help="Run lightweight checks without loading a checkpoint") - return parser - - -def main() -> None: - args = build_parser().parse_args() - if args.self_check: - run_self_check() - return - if args.atol < 0: - raise ValueError("--atol must be non-negative") - - checkpoint = args.checkpoint.expanduser().resolve() - config_path = args.config.expanduser().resolve() - sample_path = args.sample.expanduser().resolve() - output_path = args.output.expanduser().resolve() - required = ( - checkpoint / "model.safetensors", - checkpoint / "assets/paligemma_tokenizer.model", - checkpoint / "assets/physical-intelligence/libero/norm_stats.json", - config_path, - sample_path, - ) - missing = [str(path) for path in required if not path.is_file()] - if missing: - raise FileNotFoundError(f"Parity inputs are missing: {missing}") - - with config_path.open("r", encoding="utf-8") as handle: - local_config = json.load(handle) - precision = args.precision or local_config["dtype"] - if precision not in PRECISIONS: - raise ValueError(f"Unsupported OpenPI precision {precision!r}; expected one of {PRECISIONS}") - local_config["dtype"] = precision - - image, wrist, state = load_sample(sample_path) - # Match the official LIBERO evaluation client: renderer frames are resized - # with PIL before they enter either policy. Both model adapters therefore - # receive the exact same 224x224 uint8 arrays. - from openpi_client import image_tools as client_image_tools - - image = client_image_tools.resize_with_pad(image, 224, 224) - wrist = client_image_tools.resize_with_pad(wrist, 224, 224) - prompt = "pick up the black bowl between the plate and the ramekin and place it on the plate" - noise = np.random.default_rng(0).standard_normal((10, 32)).astype(np.float32) - - from openpi.training import config as training_config - - upstream_config = training_config.get_config("pi05_libero") - upstream_config = dataclasses.replace( - upstream_config, - model=dataclasses.replace( - upstream_config.model, - dtype=precision, - pytorch_compile_mode=None, - ), - ) - upstream_policy = _create_upstream_policy(upstream_config, checkpoint, args.device, precision) - upstream_dtype_counts = _parameter_dtype_counts(upstream_policy._model) # noqa: SLF001 - _validate_selected_precision(upstream_dtype_counts, precision, "upstream") - raw_input = { - "observation/image": image.copy(), - "observation/wrist_image": wrist.copy(), - "observation/state": state.copy(), - "prompt": prompt, - } - upstream_inputs = upstream_policy._input_transform(raw_input) # noqa: SLF001 - upstream_torch_inputs = _to_torch_batch(upstream_inputs, args.device) - from openpi.models import model as upstream_model_module - - upstream_observation = upstream_model_module.Observation.from_dict(upstream_torch_inputs) - upstream_noise = torch.from_numpy(noise).to(args.device)[None, ...] - upstream_normalized = upstream_policy._sample_actions( # noqa: SLF001 - args.device, - upstream_observation, - noise=upstream_noise, - ) - upstream_outputs = { - "state": _to_numpy(upstream_torch_inputs["state"][0]), - "actions": _to_numpy(upstream_normalized[0]), - } - upstream_actions = np.asarray(upstream_policy._output_transform(upstream_outputs)["actions"]) # noqa: SLF001 - upstream_normalized_array = _to_numpy(upstream_normalized) - del upstream_policy, upstream_normalized, upstream_noise, upstream_torch_inputs - gc.collect() - if args.device.startswith("cuda"): - torch.cuda.empty_cache() - - _install_lightx2v_package_stub() - from lightx2v.models.networks.openpi import OpenPIModel - - local_config["model_path"] = str(checkpoint) - local_config["device"] = args.device - local_config["seed"] = 0 - model = OpenPIModel.from_config(local_config) - local_dtype_counts = _parameter_dtype_counts(model.core_model) - _validate_selected_precision(local_dtype_counts, precision, "LightX2V") - local_observation = model.pre_infer.infer( - images={"agentview": image, "wrist": wrist}, - state=state, - task_description=prompt, - ) - preprocessing = _observation_comparison(upstream_observation, local_observation) - local_normalized = model.transformer_infer.infer( - model.core_model, - local_observation, - model.device, - noise=torch.from_numpy(noise).to(model.device)[None, ...], - ) - local_actions = model.post_infer.infer(local_normalized) - - normalized_difference = _to_numpy(local_normalized).astype(np.float64) - upstream_normalized_array.astype(np.float64) - normalized_max_abs_error = float(np.max(np.abs(normalized_difference))) - normalized_passed = bool(np.allclose(_to_numpy(local_normalized), upstream_normalized_array, rtol=0.0, atol=args.atol)) - physical_difference = np.asarray(local_actions, dtype=np.float64) - np.asarray(upstream_actions, dtype=np.float64) - physical_max_abs_error = float(np.max(np.abs(physical_difference))) - physical_passed = bool(np.allclose(local_actions, upstream_actions, rtol=0.0, atol=args.atol)) - passed = bool(preprocessing["exact"] and normalized_passed and physical_passed) - report = { - "precision": precision, - "atol": args.atol, - "upstream_parameter_dtypes": upstream_dtype_counts, - "lightx2v_parameter_dtypes": local_dtype_counts, - "preprocessing": preprocessing, - "normalized_actions": { - "upstream_shape": list(upstream_normalized_array.shape), - "lightx2v_shape": list(local_normalized.shape), - "max_abs_error": normalized_max_abs_error, - "mean_abs_error": float(np.mean(np.abs(normalized_difference))), - "allclose": normalized_passed, - }, - "physical_actions": { - "upstream_shape": list(upstream_actions.shape), - "lightx2v_shape": list(local_actions.shape), - "max_abs_error": physical_max_abs_error, - "mean_abs_error": float(np.mean(np.abs(physical_difference))), - "allclose": physical_passed, - }, - "allclose": passed, - "upstream_normalized_actions": upstream_normalized_array[0].tolist(), - "lightx2v_normalized_actions": _to_numpy(local_normalized[0]).tolist(), - "upstream_actions": upstream_actions.tolist(), - "lightx2v_actions": local_actions.tolist(), - } - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") - raw_action_keys = { - "upstream_normalized_actions", - "lightx2v_normalized_actions", - "upstream_actions", - "lightx2v_actions", - } - print(json.dumps({key: value for key, value in report.items() if key not in raw_action_keys}, indent=2)) - del model - gc.collect() - if args.device.startswith("cuda"): - torch.cuda.empty_cache() - if not passed: - raise SystemExit( - "OpenPI PyTorch parity check failed: " - f"preprocessing_exact={preprocessing['exact']}, " - f"normalized_max_abs_error={normalized_max_abs_error}, " - f"physical_max_abs_error={physical_max_abs_error}, atol={args.atol}" - ) - - -if __name__ == "__main__": - main() From 5ebb0ff752335e5f6dab91363c1dfb483eb3471b Mon Sep 17 00:00:00 2001 From: Chernobyllight Date: Fri, 4 Sep 2026 08:38:51 +0000 Subject: [PATCH 4/6] docs(openpi): explain minimal data setup --- scripts/openpi/support/README.md | 82 ++++++++++++++++++++++++++++++++ 1 file changed, 82 insertions(+) diff --git a/scripts/openpi/support/README.md b/scripts/openpi/support/README.md index 4be24717b..53958ade1 100644 --- a/scripts/openpi/support/README.md +++ b/scripts/openpi/support/README.md @@ -38,6 +38,87 @@ conda activate base 路径都可以通过下文列出的环境变量覆盖。 +### 从零构造 `openpi_data` + +`/data/liuhongda/openpi_data` 不是 OpenPI 仓库自带的目录，也不是官方固定路径。 +OpenPI 官方下载器默认使用 `~/.cache/openpi`，并通过 `OPENPI_DATA_HOME` 修改缓存 +根目录；LightX2V 使用 `OPENPI_DATA_ROOT` 查找同一批资源。本指南将两个变量指向 +同一个目录： + +```bash +export OPENPI_DATA_HOME=/data/liuhongda/openpi_data +export OPENPI_DATA_ROOT=/data/liuhongda/openpi_data +``` + +两个变量的职责不同：`OPENPI_DATA_HOME` 只影响 OpenPI 官方下载器， +`OPENPI_DATA_ROOT` 只影响本目录下的 LightX2V 脚本。 + +如果已经有转换完成的 PyTorch checkpoint，只做推理或评测所需的最小结构是： + +```text +openpi_data/ +├── openpi-assets/checkpoints/pi05_libero_pytorch_fp32/ +│ ├── model.safetensors +│ ├── config.json +│ └── assets/ +│ ├── paligemma_tokenizer.model +│ └── physical-intelligence/libero/norm_stats.json +└── python_deps/openpi_official_pytorch_runtime/ # 由环境准备脚本生成 +``` + +如果要从官方 JAX checkpoint 开始转换，转换前还需要： + +```text +openpi_data/ +├── openpi-assets/checkpoints/pi05_libero/ +│ ├── params/ # 完整 Orbax checkpoint +│ └── assets/physical-intelligence/libero/norm_stats.json +└── big_vision/paligemma_tokenizer.model +``` + +在已经执行过 `uv sync` 的 OpenPI 仓库中，可以用官方 `maybe_download()` 下载这两项： + +```bash +cd /data/liuhongda/openpi + +export OPENPI_DATA_HOME=/data/liuhongda/openpi_data +export OPENPI_DATA_ROOT=/data/liuhongda/openpi_data + +uv run --no-sync python -c \ + 'from openpi.shared import download; print(download.maybe_download("gs://openpi-assets/checkpoints/pi05_libero"))' + +uv run --no-sync python -c \ + 'from openpi.shared import download; print(download.maybe_download("gs://big_vision/paligemma_tokenizer.model", gs={"token": "anon"}))' +``` + +下载器会自动创建 `openpi_data` 及其子目录，不需要手工逐级创建。 + +下载完成后回到 LightX2V，依次生成 Transformers overlay 和 FP32 PyTorch 权重： + +```bash +cd /data/liuhongda/lightx2v_openpi +conda activate base + +export OPENPI_DATA_ROOT=/data/liuhongda/openpi_data + +bash scripts/openpi/2_setup_pytorch_runtime.sh +bash scripts/openpi/1_convert_pi05_libero_to_pytorch.sh +bash scripts/openpi/2_setup_pytorch_runtime.sh check +``` + +对应的目录生成关系是： + +| 操作 | 生成内容 | +| --- | --- | +| OpenPI `maybe_download()` | `openpi-assets/checkpoints/pi05_libero` 和 `big_vision/paligemma_tokenizer.model` | +| `2_setup_pytorch_runtime.sh` | `python_deps/openpi_official_pytorch_runtime` | +| `1_convert_pi05_libero_to_pytorch.sh` | `openpi-assets/checkpoints/pi05_libero_pytorch_fp32` | + +LIBERO 仿真环境仍位于 `/data/liuhongda/openpi/third_party/libero`，不放在 +`openpi_data` 中。`lerobot/physical-intelligence/libero` 只在微调时需要； +`raw/`、`results/`、`manifests/`、`runtime_configs/` 和各种 cache 目录也都不是 +推理或评测的最小依赖。 + ## 1. 准备运行环境 setup 只安装或修复 OpenPI 所需的小包：base 环境中的 `mujoco==3.2.3`，以及 @@ -390,6 +471,7 @@ coordinator 当前按完整 suite 写新目录，不提供中途恢复。 | 环境变量 | 作用 | | --- | --- | +| `OPENPI_DATA_HOME` | OpenPI 官方下载器的缓存根目录，仅下载阶段使用 | | `OPENPI_DATA_ROOT` | checkpoint、tokenizer 和 Transformers overlay 的数据根目录 | | `OPENPI_MODEL_PATH` | PyTorch checkpoint | | `OPENPI_CONFIG` | 模型 JSON | From 87ee37ab3956730a7c99a82443128b405f3a2277 Mon Sep 17 00:00:00 2001 From: Chernobyllight Date: Fri, 4 Sep 2026 09:59:39 +0000 Subject: [PATCH 5/6] chore(openpi): order setup and conversion scripts --- ..._runtime.sh => 1_setup_pytorch_runtime.sh} | 6 +++--- ...sh => 2_convert_pi05_libero_to_pytorch.sh} | 4 ++-- scripts/openpi/support/README.md | 20 +++++++++---------- .../openpi/support/convert_jax_checkpoint.py | 2 +- 4 files changed, 16 insertions(+), 16 deletions(-) rename scripts/openpi/{2_setup_pytorch_runtime.sh => 1_setup_pytorch_runtime.sh} (88%) rename scripts/openpi/{1_convert_pi05_libero_to_pytorch.sh => 2_convert_pi05_libero_to_pytorch.sh} (54%) diff --git a/scripts/openpi/2_setup_pytorch_runtime.sh b/scripts/openpi/1_setup_pytorch_runtime.sh similarity index 88% rename from scripts/openpi/2_setup_pytorch_runtime.sh rename to scripts/openpi/1_setup_pytorch_runtime.sh index 640fac7f4..b7a41f3a0 100755 --- a/scripts/openpi/2_setup_pytorch_runtime.sh +++ b/scripts/openpi/1_setup_pytorch_runtime.sh @@ -5,10 +5,10 @@ script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" command="prepare" if [[ "${1:-}" == "check" || "${1:-}" == "--check" ]]; then - command="check" - shift + command="check" + shift elif [[ "${1:-}" == "setup" || "${1:-}" == "prepare" ]]; then - shift + shift fi exec python "${script_dir}/support/runtime.py" "${command}" "$@" diff --git a/scripts/openpi/1_convert_pi05_libero_to_pytorch.sh b/scripts/openpi/2_convert_pi05_libero_to_pytorch.sh similarity index 54% rename from scripts/openpi/1_convert_pi05_libero_to_pytorch.sh rename to scripts/openpi/2_convert_pi05_libero_to_pytorch.sh index 7714802fb..a0b2b1735 100755 --- a/scripts/openpi/1_convert_pi05_libero_to_pytorch.sh +++ b/scripts/openpi/2_convert_pi05_libero_to_pytorch.sh @@ -3,6 +3,6 @@ set -euo pipefail script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" openpi_root="${OPENPI_PATH:-$(cd -- "${script_dir}/../../.." && pwd)/openpi}" -python_bin="${OPENPI_CONVERT_PYTHON:-${openpi_root}/.venv/bin/python}" +convert_python="${OPENPI_CONVERT_PYTHON:-${openpi_root}/.venv/bin/python}" -exec "${python_bin}" "${script_dir}/support/convert_jax_checkpoint.py" "$@" +exec "${convert_python}" "${script_dir}/support/convert_jax_checkpoint.py" "$@" diff --git a/scripts/openpi/support/README.md b/scripts/openpi/support/README.md index 53958ade1..bbe0ced6d 100644 --- a/scripts/openpi/support/README.md +++ b/scripts/openpi/support/README.md @@ -101,9 +101,9 @@ conda activate base export OPENPI_DATA_ROOT=/data/liuhongda/openpi_data -bash scripts/openpi/2_setup_pytorch_runtime.sh -bash scripts/openpi/1_convert_pi05_libero_to_pytorch.sh -bash scripts/openpi/2_setup_pytorch_runtime.sh check +bash scripts/openpi/1_setup_pytorch_runtime.sh +bash scripts/openpi/2_convert_pi05_libero_to_pytorch.sh +bash scripts/openpi/1_setup_pytorch_runtime.sh check ``` 对应的目录生成关系是： @@ -111,8 +111,8 @@ bash scripts/openpi/2_setup_pytorch_runtime.sh check | 操作 | 生成内容 | | --- | --- | | OpenPI `maybe_download()` | `openpi-assets/checkpoints/pi05_libero` 和 `big_vision/paligemma_tokenizer.model` | -| `2_setup_pytorch_runtime.sh` | `python_deps/openpi_official_pytorch_runtime` | -| `1_convert_pi05_libero_to_pytorch.sh` | `openpi-assets/checkpoints/pi05_libero_pytorch_fp32` | +| `1_setup_pytorch_runtime.sh` | `python_deps/openpi_official_pytorch_runtime` | +| `2_convert_pi05_libero_to_pytorch.sh` | `openpi-assets/checkpoints/pi05_libero_pytorch_fp32` | LIBERO 仿真环境仍位于 `/data/liuhongda/openpi/third_party/libero`，不放在 `openpi_data` 中。`lerobot/physical-intelligence/libero` 只在微调时需要； @@ -126,13 +126,13 @@ setup 只安装或修复 OpenPI 所需的小包：base 环境中的 `mujoco==3.2 Python、PyTorch 或 CUDA。 ```bash -bash scripts/openpi/2_setup_pytorch_runtime.sh +bash scripts/openpi/1_setup_pytorch_runtime.sh ``` 训练或评测前可做只读检查： ```bash -bash scripts/openpi/2_setup_pytorch_runtime.sh check +bash scripts/openpi/1_setup_pytorch_runtime.sh check ``` 检查覆盖 checkpoint、Transformers replacement、MuJoCo 来源、LIBERO 来源和 @@ -143,17 +143,17 @@ CUDA。启动脚本不会在每次运行前重复执行这项检查。 默认将官方 JAX checkpoint 转为 FP32 PyTorch checkpoint： ```bash -bash scripts/openpi/1_convert_pi05_libero_to_pytorch.sh +bash scripts/openpi/2_convert_pi05_libero_to_pytorch.sh ``` 选择输出精度或路径： ```bash -bash scripts/openpi/1_convert_pi05_libero_to_pytorch.sh --precision bfloat16 +bash scripts/openpi/2_convert_pi05_libero_to_pytorch.sh --precision bfloat16 OPENPI_JAX_CHECKPOINT=/path/to/pi05_libero \ OPENPI_PYTORCH_CHECKPOINT=/path/to/pi05_libero_pytorch_fp32 \ -bash scripts/openpi/1_convert_pi05_libero_to_pytorch.sh --precision float32 +bash scripts/openpi/2_convert_pi05_libero_to_pytorch.sh --precision float32 ``` 转换器使用 OpenPI 自身的环境，默认是 diff --git a/scripts/openpi/support/convert_jax_checkpoint.py b/scripts/openpi/support/convert_jax_checkpoint.py index 669ac9c4f..7cdca9636 100755 --- a/scripts/openpi/support/convert_jax_checkpoint.py +++ b/scripts/openpi/support/convert_jax_checkpoint.py @@ -79,7 +79,7 @@ def _check_inputs(args: argparse.Namespace) -> None: try: subprocess.run([sys.executable, "-c", probe], env=_conversion_environment(args), check=True) except subprocess.CalledProcessError as error: - raise RuntimeError("the conversion Python cannot load OpenPI's patched transformers==4.53.2; run scripts/openpi/2_setup_pytorch_runtime.sh first") from error + raise RuntimeError("the conversion Python cannot load OpenPI's patched transformers==4.53.2; run scripts/openpi/1_setup_pytorch_runtime.sh first") from error def _run_converter(args: argparse.Namespace, output: Path) -> None: From 79086d1a2e8ac57da2a8a9e9c6815ca3f6113434 Mon Sep 17 00:00:00 2001 From: Chernobyllight Date: Wed, 9 Sep 2026 07:42:53 +0000 Subject: [PATCH 6/6] feat(openpi): add aligned LIBERO training support Add official-aligned PyTorch fine-tuning with FP32 master weights, BF16 compute, EMA, strict checkpoint resume, and LIBERO preprocessing. Streamline OpenPI inference, ROS evaluation, and runtime validation while keeping generated artifacts out of Git. --- .gitignore | 1 + lightx2v/models/networks/openpi/config.py | 23 +- lightx2v/models/networks/openpi/gemma.py | 120 ++- .../models/networks/openpi/image_tools.py | 2 +- lightx2v/models/networks/openpi/model.py | 26 +- lightx2v/models/networks/openpi/pi0.py | 127 ++- .../models/networks/openpi/preprocessing.py | 258 ++++-- .../models/networks/openpi/weights/loader.py | 23 +- .../models/runners/openpi/openpi_runner.py | 4 - lightx2v_ros/src/common/common/contract.py | 13 + .../inference/inference/openpi_node/main.py | 45 +- .../simulator/libero_node/observer.py | 19 +- .../src/simulator/simulator/sim/node.py | 16 +- .../configs/train/openpi/pi05_libero.yaml | 81 ++ .../lightx2v_train/data/openpi_libero.py | 403 ++++++++++ .../model_zoo/openpi/__init__.py | 5 + .../model_zoo/openpi/pi05_libero.py | 72 ++ .../lightx2v_train/trainers/openpi.py | 747 ++++++++++++++++++ .../lightx2v_train/utils/registry.py | 4 + .../scripts/openpi/run_pi05_finetune_ema.sh | 11 + .../scripts/openpi/run_pi05_resume_ema.sh | 33 + .../scripts/openpi/support/README.md | 470 +++++++++++ .../openpi/support/launch_pi05_libero.sh | 54 ++ scripts/openpi/1_setup_pytorch_runtime.sh | 4 +- scripts/openpi/run_libero_evaluate_i2va.sh | 1 - .../run_libero_evaluate_parallel_i2va.sh | 6 +- scripts/openpi/run_libero_ros_i2va.sh | 5 +- scripts/openpi/support/README.md | 11 +- scripts/openpi/support/runtime.py | 192 ++++- 29 files changed, 2477 insertions(+), 299 deletions(-) create mode 100644 lightx2v_train/configs/train/openpi/pi05_libero.yaml create mode 100644 lightx2v_train/lightx2v_train/data/openpi_libero.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/openpi/__init__.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/openpi/pi05_libero.py create mode 100644 lightx2v_train/lightx2v_train/trainers/openpi.py create mode 100755 lightx2v_train/scripts/openpi/run_pi05_finetune_ema.sh create mode 100755 lightx2v_train/scripts/openpi/run_pi05_resume_ema.sh create mode 100644 lightx2v_train/scripts/openpi/support/README.md create mode 100755 lightx2v_train/scripts/openpi/support/launch_pi05_libero.sh diff --git a/.gitignore b/.gitignore index d6fe813c1..090c2498e 100644 --- a/.gitignore +++ b/.gitignore @@ -28,6 +28,7 @@ server_cache/ app/.gradio/ *.pkl save_results/* +/output_train/ *.egg-info/ lightx2v_train/output_train/* lightx2v_train/output_infer/* diff --git a/lightx2v/models/networks/openpi/config.py b/lightx2v/models/networks/openpi/config.py index 9b39a391d..e87bb7b47 100644 --- a/lightx2v/models/networks/openpi/config.py +++ b/lightx2v/models/networks/openpi/config.py @@ -37,7 +37,9 @@ class Pi0Config: action_dim: int = 32 action_horizon: int = 10 max_token_len: int = 200 - dtype: Literal["bfloat16", "float32"] = "bfloat16" + compute_dtype: Literal["bfloat16", "float32"] = "bfloat16" + parameter_dtype: Literal["bfloat16", "float32"] | None = None + require_fp32_checkpoint: bool = False paligemma_variant: GemmaVariant = "gemma_2b" action_expert_variant: GemmaVariant = "gemma_300m" pi05: bool = True @@ -46,18 +48,25 @@ class Pi0Config: @classmethod def from_mapping(cls, config: Mapping[str, Any]) -> "Pi0Config": + compute_dtype = config.get("compute_dtype", config.get("dtype", "bfloat16")) return cls( action_dim=config["action_dim"], action_horizon=config["action_horizon"], max_token_len=config["max_token_len"], - dtype=config["dtype"], + compute_dtype=compute_dtype, + parameter_dtype=config.get("parameter_dtype"), + require_fp32_checkpoint=config.get("require_fp32_checkpoint", False), paligemma_variant=config["paligemma_variant"], action_expert_variant=config["action_expert_variant"], pi05=config["pi05"], discrete_state_input=config["discrete_state_input"], - pytorch_compile_mode=config["pytorch_compile_mode"], + pytorch_compile_mode=config.get("pytorch_compile_mode"), ) + @property + def resolved_parameter_dtype(self) -> Literal["bfloat16", "float32"]: + return self.parameter_dtype or self.compute_dtype + def validate_pi05_libero(self) -> None: expected = { "pi05": True, @@ -73,5 +82,9 @@ def validate_pi05_libero(self) -> None: if wrong: details = ", ".join(f"{name}={got!r} (expected {want!r})" for name, (got, want) in wrong.items()) raise ValueError(f"Configuration does not match the released pi05_libero checkpoint: {details}") - if self.dtype not in {"bfloat16", "float32"}: - raise ValueError(f"Unsupported OpenPI dtype: {self.dtype!r}") + if self.compute_dtype not in {"bfloat16", "float32"}: + raise ValueError(f"Unsupported OpenPI compute dtype: {self.compute_dtype!r}") + if self.resolved_parameter_dtype not in {"bfloat16", "float32"}: + raise ValueError(f"Unsupported OpenPI parameter dtype: {self.resolved_parameter_dtype!r}") + if self.require_fp32_checkpoint and self.resolved_parameter_dtype != "float32": + raise ValueError("require_fp32_checkpoint=true requires parameter_dtype='float32'") diff --git a/lightx2v/models/networks/openpi/gemma.py b/lightx2v/models/networks/openpi/gemma.py index e4181badf..d29c107bd 100644 --- a/lightx2v/models/networks/openpi/gemma.py +++ b/lightx2v/models/networks/openpi/gemma.py @@ -1,6 +1,7 @@ # Adapted from Physical Intelligence OpenPI (Apache-2.0), commit 15a9616. # Localized for the LightX2V OpenPI backend; no runtime OpenPI/JAX dependency. +from contextlib import nullcontext from typing import Literal import torch @@ -15,12 +16,14 @@ def __init__( self, vlm_config, action_expert_config, - use_adarms=None, - precision: Literal["bfloat16", "float32"] = "bfloat16", + use_adarms, + *, + parameter_precision: Literal["bfloat16", "float32"], + compute_precision: Literal["bfloat16", "float32"], ): - if use_adarms is None: - use_adarms = [False, False] super().__init__() + self.parameter_precision = parameter_precision + self.compute_precision = compute_precision vlm_config_hf = CONFIG_MAPPING["paligemma"]() vlm_config_hf._vocab_size = 257152 # noqa: SLF001 @@ -59,9 +62,9 @@ def __init__( self.gemma_expert = GemmaForCausalLM(config=action_expert_config_hf) self.gemma_expert.model.embed_tokens = None - self.to_bfloat16_for_selected_params(precision) + self._apply_parameter_precision(self.parameter_precision) - def to_bfloat16_for_selected_params(self, precision: Literal["bfloat16", "float32"] = "bfloat16"): + def _apply_parameter_precision(self, precision: Literal["bfloat16", "float32"]): if precision == "bfloat16": self.to(dtype=torch.bfloat16) elif precision == "float32": @@ -83,8 +86,44 @@ def to_bfloat16_for_selected_params(self, precision: Literal["bfloat16", "float3 if any(selector in name for selector in params_to_keep_float32): param.data = param.data.to(dtype=torch.float32) + @property + def uses_mixed_precision(self) -> bool: + return self.parameter_precision == "float32" and self.compute_precision == "bfloat16" + + def _compute_context(self, tensor: torch.Tensor): + if self.uses_mixed_precision and tensor.device.type == "cuda": + return torch.autocast(device_type="cuda", dtype=torch.bfloat16) + return nullcontext() + + def cast_transformer_activation(self, tensor: torch.Tensor | None): + if tensor is None: + return None + if self.compute_precision == "bfloat16" and not (self.uses_mixed_precision and tensor.device.type == "cpu"): + return tensor.to(dtype=torch.bfloat16) + return tensor.to(dtype=torch.float32) + def embed_image(self, image: torch.Tensor): - return self.paligemma.model.get_image_features(image) + if not self.uses_mixed_precision: + with self._compute_context(image): + return self.paligemma.model.get_image_features(image) + + # OpenPI keeps patch extraction and positional embedding in FP32, then + # switches the SigLIP encoder and projector to the configured matmul + # dtype. Wrapping get_image_features() as a whole in autocast would + # move the FP32/BF16 boundary ahead of the patch convolution. + vision_model = self.paligemma.model.vision_tower.vision_model + stem_context = torch.autocast(device_type="cuda", enabled=False) if image.device.type == "cuda" else nullcontext() + with stem_context: + hidden_states = vision_model.embeddings(image.float(), interpolate_pos_encoding=False) + hidden_states = self.cast_transformer_activation(hidden_states) + with self._compute_context(hidden_states): + encoder_outputs = vision_model.encoder( + inputs_embeds=hidden_states, + output_attentions=False, + output_hidden_states=False, + ) + hidden_states = vision_model.post_layernorm(encoder_outputs.last_hidden_state) + return self.paligemma.model.multi_modal_projector(hidden_states) def embed_language_tokens(self, tokens: torch.Tensor): return self.paligemma.language_model.embed_tokens(tokens) @@ -101,45 +140,45 @@ def forward( if adarms_cond is None: adarms_cond = [None, None] if inputs_embeds[1] is None: - prefix_output = self.paligemma.language_model.forward( - inputs_embeds=inputs_embeds[0], - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - use_cache=use_cache, - adarms_cond=adarms_cond[0], - ) + prefix_inputs = self.cast_transformer_activation(inputs_embeds[0]) + with self._compute_context(prefix_inputs): + prefix_output = self.paligemma.language_model.forward( + inputs_embeds=prefix_inputs, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + adarms_cond=adarms_cond[0], + ) prefix_past_key_values = prefix_output.past_key_values prefix_output = prefix_output.last_hidden_state suffix_output = None elif inputs_embeds[0] is None: - suffix_output = self.gemma_expert.model.forward( - inputs_embeds=inputs_embeds[1], - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - use_cache=use_cache, - adarms_cond=adarms_cond[1], - ) + suffix_inputs = self.cast_transformer_activation(inputs_embeds[1]) + with self._compute_context(suffix_inputs): + suffix_output = self.gemma_expert.model.forward( + inputs_embeds=suffix_inputs, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + adarms_cond=adarms_cond[1], + ) suffix_output = suffix_output.last_hidden_state prefix_output = None prefix_past_key_values = None else: + inputs_embeds = [self.cast_transformer_activation(tensor) for tensor in inputs_embeds] models = [self.paligemma.language_model, self.gemma_expert.model] num_layers = self.paligemma.config.text_config.num_hidden_layers - use_gradient_checkpointing = (hasattr(self.gemma_expert.model, "gradient_checkpointing") and self.gemma_expert.model.gradient_checkpointing and self.training) or ( - hasattr(self, "gradient_checkpointing") and self.gradient_checkpointing and self.training - ) - - if self.training and hasattr(self.gemma_expert.model, "gradient_checkpointing"): - if not self.gemma_expert.model.gradient_checkpointing: - self.gemma_expert.model.gradient_checkpointing = True - use_gradient_checkpointing = True + use_gradient_checkpointing = self.gemma_expert.model.gradient_checkpointing and self.training - def compute_layer_complete(layer_idx, inputs_embeds, attention_mask, position_ids, adarms_cond): - models = [self.paligemma.language_model, self.gemma_expert.model] + def compute_layer_with_precision(layer_idx, inputs_embeds, attention_mask, position_ids, adarms_cond): + with self._compute_context(inputs_embeds[0]): + return compute_layer(layer_idx, inputs_embeds, attention_mask, position_ids, adarms_cond) + def compute_layer(layer_idx, inputs_embeds, attention_mask, position_ids, adarms_cond): query_states = [] key_states = [] value_states = [] @@ -213,7 +252,7 @@ def compute_layer_complete(layer_idx, inputs_embeds, attention_mask, position_id for layer_idx in range(num_layers): if use_gradient_checkpointing: inputs_embeds = torch.utils.checkpoint.checkpoint( - compute_layer_complete, + compute_layer_with_precision, layer_idx, inputs_embeds, attention_mask, @@ -223,14 +262,15 @@ def compute_layer_complete(layer_idx, inputs_embeds, attention_mask, position_id preserve_rng_state=False, ) else: - inputs_embeds = compute_layer_complete(layer_idx, inputs_embeds, attention_mask, position_ids, adarms_cond) + inputs_embeds = compute_layer_with_precision(layer_idx, inputs_embeds, attention_mask, position_ids, adarms_cond) def compute_final_norms(inputs_embeds, adarms_cond): - outputs_embeds = [] - for i, hidden_states in enumerate(inputs_embeds): - out_emb, _ = models[i].norm(hidden_states, cond=adarms_cond[i]) - outputs_embeds.append(out_emb) - return outputs_embeds + with self._compute_context(inputs_embeds[0]): + outputs_embeds = [] + for i, hidden_states in enumerate(inputs_embeds): + out_emb, _ = models[i].norm(hidden_states, cond=adarms_cond[i]) + outputs_embeds.append(out_emb) + return outputs_embeds if use_gradient_checkpointing: outputs_embeds = torch.utils.checkpoint.checkpoint(compute_final_norms, inputs_embeds, adarms_cond, use_reentrant=False, preserve_rng_state=False) diff --git a/lightx2v/models/networks/openpi/image_tools.py b/lightx2v/models/networks/openpi/image_tools.py index 7d03e2957..1072d985e 100644 --- a/lightx2v/models/networks/openpi/image_tools.py +++ b/lightx2v/models/networks/openpi/image_tools.py @@ -1,4 +1,4 @@ -"""Image helpers copied from OpenPI's torch preprocessing path (Apache-2.0).""" +"""Image resizing used by OpenPI's inference path.""" import torch import torch.nn.functional as F # noqa: N812 diff --git a/lightx2v/models/networks/openpi/model.py b/lightx2v/models/networks/openpi/model.py index 0dc81528a..5ea6a6707 100644 --- a/lightx2v/models/networks/openpi/model.py +++ b/lightx2v/models/networks/openpi/model.py @@ -93,28 +93,6 @@ def _sample_noise(self, seed: int | None = None) -> torch.Tensor: generator=generator, ) - @torch.no_grad() - def predict_normalized_action_chunk( - self, - images: dict[str, np.ndarray], - state: np.ndarray, - task_description: str, - *, - seed: int | None = None, - noise: torch.Tensor | np.ndarray | None = None, - ) -> torch.Tensor: - observation = self.pre_infer.infer(images, state, task_description) - if noise is None: - noise_tensor = self._sample_noise(seed) - else: - noise_tensor = torch.as_tensor(noise, dtype=torch.float32, device=self.device) - if noise_tensor.ndim == 2: - noise_tensor = noise_tensor.unsqueeze(0) - expected = (1, self.model_config.action_horizon, self.model_config.action_dim) - if tuple(noise_tensor.shape) != expected: - raise ValueError(f"Noise must have shape {expected}, got {tuple(noise_tensor.shape)}") - return self.transformer_infer.infer(self.core_model, observation, self.device, noise=noise_tensor) - @torch.no_grad() def predict_action_chunk( self, @@ -123,5 +101,7 @@ def predict_action_chunk( task_description: str, seed: int | None = None, ) -> np.ndarray: - normalized = self.predict_normalized_action_chunk(images, state, task_description, seed=seed) + observation = self.pre_infer.infer(images, state, task_description) + noise = self._sample_noise(seed) + normalized = self.transformer_infer.infer(self.core_model, observation, self.device, noise=noise) return self.post_infer.infer(normalized) diff --git a/lightx2v/models/networks/openpi/pi0.py b/lightx2v/models/networks/openpi/pi0.py index 37180fcaa..a56ce65e7 100644 --- a/lightx2v/models/networks/openpi/pi0.py +++ b/lightx2v/models/networks/openpi/pi0.py @@ -22,7 +22,14 @@ def get_safe_dtype(target_dtype, device_type): return target_dtype -def create_sinusoidal_pos_embedding(time: torch.tensor, dimension: int, min_period: float, max_period: float, device="cpu") -> Tensor: +def create_sinusoidal_pos_embedding( + time: torch.Tensor, + dimension: int, + min_period: float, + max_period: float, + device="cpu", + dtype: torch.dtype | None = None, +) -> Tensor: """Computes sine-cosine positional embedding vectors for scalar positions.""" if dimension % 2 != 0: raise ValueError(f"dimension ({dimension}) must be divisible by 2") @@ -30,7 +37,9 @@ def create_sinusoidal_pos_embedding(time: torch.tensor, dimension: int, min_peri if time.ndim != 1: raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.") - dtype = get_safe_dtype(torch.float64, device.type) + if dtype is None: + dtype = get_safe_dtype(torch.float64, device.type) + time = time.to(dtype=dtype) fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device) period = min_period * (max_period / min_period) ** fraction @@ -91,7 +100,8 @@ def __init__(self, config): paligemma_config, action_expert_config, use_adarms=[False, True] if self.pi05 else [False, False], - precision=config.dtype, + parameter_precision=config.resolved_parameter_dtype, + compute_precision=config.compute_dtype, ) self.action_in_proj = nn.Linear(config.action_dim, action_expert_config.width) @@ -111,6 +121,12 @@ def __init__(self, config): self.gradient_checkpointing_enabled = False + def assert_fp32_parameters(self) -> None: + non_fp32 = [(name, parameter.dtype) for name, parameter in self.named_parameters() if parameter.requires_grad and parameter.is_floating_point() and parameter.dtype != torch.float32] + if non_fp32: + preview = ", ".join(f"{name}={dtype}" for name, dtype in non_fp32[:8]) + raise RuntimeError(f"OpenPI trainable parameters must remain FP32; found {preview}") + def gradient_checkpointing_enable(self): self.gradient_checkpointing_enabled = True self.paligemma_with_expert.paligemma.language_model.gradient_checkpointing = True @@ -119,17 +135,6 @@ def gradient_checkpointing_enable(self): logging.info("Enabled gradient checkpointing for PI0Pytorch model") - def gradient_checkpointing_disable(self): - self.gradient_checkpointing_enabled = False - self.paligemma_with_expert.paligemma.language_model.gradient_checkpointing = False - self.paligemma_with_expert.paligemma.vision_tower.gradient_checkpointing = False - self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = False - - logging.info("Disabled gradient checkpointing for PI0Pytorch model") - - def is_gradient_checkpointing_enabled(self): - return self.gradient_checkpointing_enabled - def _apply_checkpoint(self, func, *args, **kwargs): if self.gradient_checkpointing_enabled and self.training: return torch.utils.checkpoint.checkpoint(func, *args, use_reentrant=False, preserve_rng_state=False, **kwargs) @@ -172,11 +177,7 @@ def embed_prefix(self, images, img_masks, lang_tokens, lang_masks) -> tuple[torc att_masks = [] for img, img_mask in zip(images, img_masks, strict=True): - - def image_embed_func(img): - return self.paligemma_with_expert.embed_image(img) - - img_emb = self._apply_checkpoint(image_embed_func, img) + img_emb = self._apply_checkpoint(self.paligemma_with_expert.embed_image, img) bsize, num_img_embs = img_emb.shape[:2] @@ -185,12 +186,9 @@ def image_embed_func(img): att_masks += [0] * num_img_embs - def lang_embed_func(lang_tokens): - lang_emb = self.paligemma_with_expert.embed_language_tokens(lang_tokens) - lang_emb_dim = lang_emb.shape[-1] - return lang_emb * math.sqrt(lang_emb_dim) - - lang_emb = self._apply_checkpoint(lang_embed_func, lang_tokens) + lang_emb = self.paligemma_with_expert.embed_language_tokens(lang_tokens) + lang_emb = lang_emb * math.sqrt(lang_emb.shape[-1]) + lang_emb = self.paligemma_with_expert.cast_transformer_activation(lang_emb) embs.append(lang_emb) pad_masks.append(lang_masks) @@ -209,6 +207,8 @@ def lang_embed_func(lang_tokens): def embed_suffix(self, state, noisy_actions, timestep): """Embed state, noisy_actions, timestep to prepare for Expert Gemma processing.""" + noisy_actions = noisy_actions.float() + timestep = timestep.float() embs = [] pad_masks = [] att_masks = [] @@ -217,10 +217,7 @@ def embed_suffix(self, state, noisy_actions, timestep): if self.state_proj.weight.dtype == torch.float32: state = state.to(torch.float32) - def state_proj_func(state): - return self.state_proj(state) - - state_emb = self._apply_checkpoint(state_proj_func, state) + state_emb = self.state_proj(state) embs.append(state_emb[:, None, :]) bsize = state_emb.shape[0] @@ -231,34 +228,26 @@ def state_proj_func(state): att_masks += [1] - time_emb = create_sinusoidal_pos_embedding(timestep, self.action_in_proj.out_features, min_period=4e-3, max_period=4.0, device=timestep.device) + embedding_dtype = torch.float32 if self.paligemma_with_expert.uses_mixed_precision else None + time_emb = create_sinusoidal_pos_embedding( + timestep, + self.action_in_proj.out_features, + min_period=4e-3, + max_period=4.0, + device=timestep.device, + dtype=embedding_dtype, + ) time_emb = time_emb.type(dtype=timestep.dtype) - def action_proj_func(noisy_actions): - return self.action_in_proj(noisy_actions) - - action_emb = self._apply_checkpoint(action_proj_func, noisy_actions) + action_emb = self.action_in_proj(noisy_actions) if not self.pi05: time_emb = time_emb[:, None, :].expand_as(action_emb) action_time_emb = torch.cat([action_emb, time_emb], dim=2) - - def mlp_func(action_time_emb): - x = self.action_time_mlp_in(action_time_emb) - x = F.silu(x) - return self.action_time_mlp_out(x) - - action_time_emb = self._apply_checkpoint(mlp_func, action_time_emb) + action_time_emb = self.action_time_mlp_out(F.silu(self.action_time_mlp_in(action_time_emb))) adarms_cond = None else: - - def time_mlp_func(time_emb): - x = self.time_mlp_in(time_emb) - x = F.silu(x) - x = self.time_mlp_out(x) - return F.silu(x) - - time_emb = self._apply_checkpoint(time_mlp_func, time_emb) + time_emb = F.silu(self.time_mlp_out(F.silu(self.time_mlp_in(time_emb)))) action_time_emb = action_emb adarms_cond = time_emb @@ -279,13 +268,18 @@ def time_mlp_func(time_emb): def forward(self, observation, actions, noise=None, time=None) -> Tensor: """Do a full training forward pass and compute the loss (batch_size x num_steps x num_motors)""" + actions = actions.float() images, img_masks, lang_tokens, lang_masks, state = self._preprocess_observation(observation, train=True) if noise is None: noise = self.sample_noise(actions.shape, actions.device) + else: + noise = noise.float() if time is None: time = self.sample_time(actions.shape[0], actions.device) + else: + time = time.float() time_expanded = time[:, None, None] x_t = time_expanded * noise + (1 - time_expanded) * actions @@ -293,9 +287,8 @@ def forward(self, observation, actions, noise=None, time=None) -> Tensor: prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(images, img_masks, lang_tokens, lang_masks) suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = self.embed_suffix(state, x_t, time) - if self.paligemma_with_expert.paligemma.language_model.layers[0].self_attn.q_proj.weight.dtype == torch.bfloat16: - suffix_embs = suffix_embs.to(dtype=torch.bfloat16) - prefix_embs = prefix_embs.to(dtype=torch.bfloat16) + suffix_embs = self.paligemma_with_expert.cast_transformer_activation(suffix_embs) + prefix_embs = self.paligemma_with_expert.cast_transformer_activation(prefix_embs) pad_masks = torch.cat([prefix_pad_masks, suffix_pad_masks], dim=1) att_masks = torch.cat([prefix_att_masks, suffix_att_masks], dim=1) @@ -305,28 +298,24 @@ def forward(self, observation, actions, noise=None, time=None) -> Tensor: att_2d_masks_4d = self._prepare_attention_masks_4d(att_2d_masks) - def forward_func(prefix_embs, suffix_embs, att_2d_masks_4d, position_ids, adarms_cond): - (_, suffix_out), _ = self.paligemma_with_expert.forward( - attention_mask=att_2d_masks_4d, - position_ids=position_ids, - past_key_values=None, - inputs_embeds=[prefix_embs, suffix_embs], - use_cache=False, - adarms_cond=[None, adarms_cond], - ) - return suffix_out - - suffix_out = self._apply_checkpoint(forward_func, prefix_embs, suffix_embs, att_2d_masks_4d, position_ids, adarms_cond) + # The joint transformer checkpoints each Gemma block internally. + # Checkpointing this complete call as well would recompute the entire + # stack on top of the per-block recomputation. + (_, suffix_out), _ = self.paligemma_with_expert.forward( + attention_mask=att_2d_masks_4d, + position_ids=position_ids, + past_key_values=None, + inputs_embeds=[prefix_embs, suffix_embs], + use_cache=False, + adarms_cond=[None, adarms_cond], + ) suffix_out = suffix_out[:, -self.config.action_horizon :] suffix_out = suffix_out.to(dtype=torch.float32) - def action_out_proj_func(suffix_out): - return self.action_out_proj(suffix_out) - - v_t = self._apply_checkpoint(action_out_proj_func, suffix_out) + v_t = self.action_out_proj(suffix_out) - return F.mse_loss(u_t, v_t, reduction="none") + return F.mse_loss(u_t.float(), v_t.float(), reduction="none") @torch.no_grad() def sample_actions(self, device, observation, noise=None, num_steps=10) -> Tensor: diff --git a/lightx2v/models/networks/openpi/preprocessing.py b/lightx2v/models/networks/openpi/preprocessing.py index 377d83f3c..ea08efd8f 100644 --- a/lightx2v/models/networks/openpi/preprocessing.py +++ b/lightx2v/models/networks/openpi/preprocessing.py @@ -1,10 +1,15 @@ # Adapted from Physical Intelligence OpenPI (Apache-2.0), commit 15a9616. -# Localized for the LightX2V OpenPI backend. +"""OpenPI image preprocessing, including the released Augmax 0.4.1 recipe.""" + +from __future__ import annotations import logging +import math from collections.abc import Sequence +from dataclasses import dataclass import torch +import torch.nn.functional as F # noqa: N812 from . import image_tools from .observation import Observation @@ -20,6 +25,153 @@ IMAGE_RESOLUTION = (224, 224) +@dataclass(frozen=True) +class _GeometricParameters: + crop_offsets_yx: torch.Tensor + angles_radians: torch.Tensor + + +@dataclass(frozen=True) +class _ColorJitterParameters: + brightness: torch.Tensor + contrast: torch.Tensor + saturation: torch.Tensor + hue: torch.Tensor + apply: torch.Tensor + + +def _sample_geometric_parameters( + batch_size: int, + device: torch.device, + image_resolution: tuple[int, int], +) -> _GeometricParameters: + height, width = image_resolution + crop_height = int(height * 0.95) + crop_width = int(width * 0.95) + limits = torch.tensor( + ((height - crop_height) / 2.0, (width - crop_width) / 2.0), + dtype=torch.float32, + device=device, + ) + offsets = (torch.rand((batch_size, 2), device=device) * 2.0 - 1.0) * limits + angles = (torch.rand(batch_size, device=device) * 10.0 - 5.0) * (math.pi / 180.0) + return _GeometricParameters(offsets, angles) + + +def _sample_color_jitter_parameters(batch_size: int, device: torch.device) -> _ColorJitterParameters: + def sample(strength: float) -> torch.Tensor: + return (torch.rand(batch_size, device=device) * 2.0 - 1.0) * strength + + return _ColorJitterParameters( + brightness=sample(0.3), + contrast=sample(0.4), + hue=sample(0.1), + saturation=sample(0.5), + apply=torch.rand(batch_size, device=device) < 0.5, + ) + + +def _apply_fused_geometric_transform( + images: torch.Tensor, + parameters: _GeometricParameters, + image_resolution: tuple[int, int], +) -> torch.Tensor: + """Apply RandomCrop(95%), Resize, and Rotate as one Augmax transform.""" + batch_size, height, width, _ = images.shape + + crop_height = int(height * 0.95) + crop_width = int(width * 0.95) + output_y = torch.arange(height, dtype=torch.float32, device=images.device) - (height / 2.0 - 0.5) + output_x = torch.arange(width, dtype=torch.float32, device=images.device) - (width / 2.0 - 0.5) + output_y, output_x = torch.meshgrid(output_y, output_x, indexing="ij") + + cosine = torch.cos(parameters.angles_radians)[:, None, None] + sine = torch.sin(parameters.angles_radians)[:, None, None] + input_y = (cosine * output_y + sine * output_x) * (crop_height / height) + input_x = (-sine * output_y + cosine * output_x) * (crop_width / width) + input_y += parameters.crop_offsets_yx[:, 0, None, None] + (height / 2.0 - 0.5) + input_x += parameters.crop_offsets_yx[:, 1, None, None] + (width / 2.0 - 0.5) + + grid = torch.stack( + ( + input_x * (2.0 / (width - 1)) - 1.0, + input_y * (2.0 / (height - 1)) - 1.0, + ), + dim=-1, + ) + transformed = F.grid_sample( + images.permute(0, 3, 1, 2), + grid, + mode="bilinear", + padding_mode="zeros", + align_corners=True, + ) + return transformed.permute(0, 2, 3, 1) + + +def _rgb_to_hsv(images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + value, argmax = images.max(dim=-1) + minimum = images.min(dim=-1).values + value_range = value - minimum + safe_range = torch.where(value_range == 0.0, torch.ones_like(value_range), value_range) + second = (argmax + 1) % 3 + third = (argmax + 2) % 3 + second_value = torch.gather(images, -1, second.unsqueeze(-1)).squeeze(-1) + third_value = torch.gather(images, -1, third.unsqueeze(-1)).squeeze(-1) + hue = (2.0 * argmax.to(images.dtype) + (second_value - third_value) / safe_range) / 6.0 + hue = torch.where(value_range == 0.0, torch.zeros_like(hue), hue) + safe_value = torch.where(value == 0.0, torch.ones_like(value), value) + saturation = torch.where(value == 0.0, torch.zeros_like(value), value_range / safe_value) + return hue, saturation, value + + +def _hsv_to_rgb(hue: torch.Tensor, saturation: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + n = torch.tensor((5.0, 3.0, 1.0), dtype=value.dtype, device=value.device) + k = torch.remainder(n + hue.unsqueeze(-1) * 6.0, 6.0) + scale = torch.maximum(torch.zeros_like(k), torch.minimum(torch.minimum(k, 4.0 - k), torch.ones_like(k))) + return value.unsqueeze(-1) - value.unsqueeze(-1) * saturation.unsqueeze(-1) * scale + + +def _adjust_brightness(value: torch.Tensor, amount: torch.Tensor) -> torch.Tensor: + return torch.where(amount < 0.0, value * (1.0 + amount), value * (1.0 - amount) + amount) + + +def _adjust_contrast(value: torch.Tensor, amount: torch.Tensor) -> torch.Tensor: + slant = torch.tan((amount + 1.0) * (math.pi / 4.0)) + slant_squared = slant.square() + first_break = (slant - slant_squared) / (2.0 * (1.0 - slant_squared)) + second_break = 1.0 - first_break + return torch.where( + value < first_break, + value / slant, + torch.where(value > second_break, value / slant + 1.0 - 1.0 / slant, slant * (value - 0.5) + 0.5), + ) + + +def _apply_color_jitter(images: torch.Tensor, parameters: _ColorJitterParameters) -> torch.Tensor: + hue, saturation, value = _rgb_to_hsv(images) + value = _adjust_brightness(value, parameters.brightness[:, None, None]) + value = _adjust_contrast(value, parameters.contrast[:, None, None]) + hue += parameters.hue[:, None, None] + + # Augmax 0.4.1 samples saturation without applying it. + transformed = _hsv_to_rgb(hue, saturation, value) + return torch.where(parameters.apply[:, None, None, None], transformed, images) + + +def _augment_images( + images: torch.Tensor, + *, + geometric: _GeometricParameters | None, + color: _ColorJitterParameters, + image_resolution: tuple[int, int], +) -> torch.Tensor: + images = images / 2.0 + 0.5 + if geometric is not None: + images = _apply_fused_geometric_transform(images, geometric, image_resolution) + return _apply_color_jitter(images, color) * 2.0 - 1.0 + + def preprocess_observation_pytorch( observation, *, @@ -31,87 +183,61 @@ def preprocess_observation_pytorch( raise ValueError(f"images dict missing keys: expected {image_keys}, got {list(observation.images)}") batch_shape = observation.state.shape[:-1] + batch_size = observation.state.shape[0] + device = observation.state.device + geometric_parameters = None + base_color_parameters = None + wrist_color_parameters = None out_images = {} for key in image_keys: image = observation.images[key] - - is_channels_first = image.shape[1] == 3 - - if is_channels_first: + if image.ndim != 4 or (image.shape[1] != 3 and image.shape[-1] != 3): + raise ValueError(f"{key} must be BCHW or BHWC RGB, got {tuple(image.shape)}") + if image.shape[0] != batch_size: + raise ValueError(f"{key} batch size {image.shape[0]} does not match state batch size {batch_size}") + channels_first = image.shape[1] == 3 + if channels_first: image = image.permute(0, 2, 3, 1) if image.shape[1:3] != image_resolution: - logger.info(f"Resizing image {key} from {image.shape[1:3]} to {image_resolution}") + if train: + raise ValueError(f"Training image {key} has shape {tuple(image.shape[1:3])}; resize the uint8 image with PIL before converting it to [-1, 1]") + logger.info("Resizing image %s from %s to %s", key, image.shape[1:3], image_resolution) image = image_tools.resize_with_pad_torch(image, *image_resolution) if train: - image = image / 2.0 + 0.5 - - if "wrist" not in key: - height, width = image.shape[1:3] - - crop_height = int(height * 0.95) - crop_width = int(width * 0.95) - - max_h = height - crop_height - max_w = width - crop_width - if max_h > 0 and max_w > 0: - start_h = torch.randint(0, max_h + 1, (1,), device=image.device) - start_w = torch.randint(0, max_w + 1, (1,), device=image.device) - image = image[:, start_h : start_h + crop_height, start_w : start_w + crop_width, :] - - image = torch.nn.functional.interpolate( - image.permute(0, 3, 1, 2), - size=(height, width), - mode="bilinear", - align_corners=False, - ).permute(0, 2, 3, 1) - - angle = torch.rand(1, device=image.device) * 10 - 5 - if torch.abs(angle) > 0.1: - angle_rad = angle * torch.pi / 180.0 - cos_a = torch.cos(angle_rad) - sin_a = torch.sin(angle_rad) - grid_x = torch.linspace(-1, 1, width, device=image.device) - grid_y = torch.linspace(-1, 1, height, device=image.device) - grid_y, grid_x = torch.meshgrid(grid_y, grid_x, indexing="ij") - grid_x = grid_x.unsqueeze(0).expand(image.shape[0], -1, -1) - grid_y = grid_y.unsqueeze(0).expand(image.shape[0], -1, -1) - grid_x_rot = grid_x * cos_a - grid_y * sin_a - grid_y_rot = grid_x * sin_a + grid_y * cos_a - grid = torch.stack([grid_x_rot, grid_y_rot], dim=-1) - - image = torch.nn.functional.grid_sample( - image.permute(0, 3, 1, 2), - grid, - mode="bilinear", - padding_mode="zeros", - align_corners=False, - ).permute(0, 2, 3, 1) - - brightness_factor = 0.7 + torch.rand(1, device=image.device) * 0.6 - image = image * brightness_factor - - contrast_factor = 0.6 + torch.rand(1, device=image.device) * 0.8 - mean = image.mean(dim=[1, 2, 3], keepdim=True) - image = (image - mean) * contrast_factor + mean - - saturation_factor = 0.5 + torch.rand(1, device=image.device) * 1.0 - gray = image.mean(dim=-1, keepdim=True) - image = gray + (image - gray) * saturation_factor - image = torch.clamp(image, 0, 1) - image = image * 2.0 - 1.0 - - if is_channels_first: + if image.dtype != torch.float32: + raise ValueError(f"Training image {key} must be float32 in [-1, 1], got {image.dtype}") + if "wrist" in key: + if wrist_color_parameters is None: + wrist_color_parameters = _sample_color_jitter_parameters(batch_size, device) + image = _augment_images( + image, + geometric=None, + color=wrist_color_parameters, + image_resolution=image_resolution, + ) + else: + if geometric_parameters is None: + geometric_parameters = _sample_geometric_parameters(batch_size, device, image_resolution) + if base_color_parameters is None: + base_color_parameters = _sample_color_jitter_parameters(batch_size, device) + image = _augment_images( + image, + geometric=geometric_parameters, + color=base_color_parameters, + image_resolution=image_resolution, + ) + + if channels_first: image = image.permute(0, 3, 1, 2) - out_images[key] = image out_masks = {} for key in out_images: if key not in observation.image_masks: - out_masks[key] = torch.ones(batch_shape, dtype=torch.bool, device=observation.state.device) + out_masks[key] = torch.ones(batch_shape, dtype=torch.bool, device=device) else: out_masks[key] = observation.image_masks[key] diff --git a/lightx2v/models/networks/openpi/weights/loader.py b/lightx2v/models/networks/openpi/weights/loader.py index 7199e02de..6413c0f01 100644 --- a/lightx2v/models/networks/openpi/weights/loader.py +++ b/lightx2v/models/networks/openpi/weights/loader.py @@ -6,6 +6,7 @@ from pathlib import Path import torch +from safetensors import safe_open from safetensors.torch import load_model from ..config import Pi0Config @@ -13,6 +14,23 @@ LOGGER = logging.getLogger(__name__) +def _validate_checkpoint_precision(weight_path: Path, config: Pi0Config) -> None: + if not config.require_fp32_checkpoint: + return + + lower_precision = [] + with safe_open(weight_path, framework="pt", device="cpu") as checkpoint: + for name in checkpoint.keys(): + dtype = checkpoint.get_slice(name).get_dtype() + if dtype.startswith(("F", "BF")) and dtype != "F32": + lower_precision.append((name, dtype)) + if len(lower_precision) == 8: + break + if lower_precision: + preview = ", ".join(f"{name}={dtype}" for name, dtype in lower_precision) + raise RuntimeError(f"OpenPI fine-tuning requires an FP32 source checkpoint; found {preview}") + + def _validate_transformers_runtime() -> None: """Fail early unless the official patched Transformers runtime is active.""" import transformers @@ -43,6 +61,7 @@ def load_pi05_libero_weights( weight_path = Path(weight_path).expanduser().resolve() if not weight_path.is_file(): raise FileNotFoundError(f"Converted OpenPI SafeTensors file not found: {weight_path}") + _validate_checkpoint_precision(weight_path, config) # Import only after selecting OpenPI's private Transformers runtime. from ..pi0 import PI0Pytorch @@ -50,9 +69,9 @@ def load_pi05_libero_weights( model = PI0Pytorch(config) load_model(model, weight_path, strict=True, device="cpu") - # Match upstream's mixed BF16/FP32 inference policy. - model.paligemma_with_expert.to_bfloat16_for_selected_params(config.dtype) model.to(torch.device(device)) + if config.resolved_parameter_dtype == "float32": + model.assert_fp32_parameters() model.eval() parameter_count = sum(parameter.numel() for parameter in model.parameters()) LOGGER.info("Loaded pi05_libero PyTorch weights strictly: %.3fB parameters", parameter_count / 1e9) diff --git a/lightx2v/models/runners/openpi/openpi_runner.py b/lightx2v/models/runners/openpi/openpi_runner.py index cbc9edf5c..876670b1c 100644 --- a/lightx2v/models/runners/openpi/openpi_runner.py +++ b/lightx2v/models/runners/openpi/openpi_runner.py @@ -73,10 +73,6 @@ def init_modules(self) -> None: raise ValueError(f"Unsupported OpenPI run mode {self.run_mode!r}; expected 'rollout' or 'evaluate'.") self.config.lock() - def warmup(self) -> None: - # Model initialization belongs to the isolated worker. - pass - def _value(self, config_name: str, *, environment_name: str | None = None, default: Any = None) -> Any: if environment_name: environment_value = os.environ.get(environment_name) diff --git a/lightx2v_ros/src/common/common/contract.py b/lightx2v_ros/src/common/common/contract.py index 5027e5173..20d163f0a 100644 --- a/lightx2v_ros/src/common/common/contract.py +++ b/lightx2v_ros/src/common/common/contract.py @@ -12,6 +12,19 @@ from typing import Dict, Tuple +def parse_multiarray_label(label: str, *field_names: str) -> Tuple[int, ...] | None: + """Read integer fields from a ROS MultiArrayDimension label.""" + fields = {} + for item in str(label).split(";"): + name, separator, value = item.partition("=") + if separator: + fields[name] = value + try: + return tuple(int(fields[name]) for name in field_names) + except (KeyError, ValueError): + return None + + @dataclass(frozen=True) class EnvContract: # Logical environment id, e.g. "libero" or "robotwin". diff --git a/lightx2v_ros/src/inference/inference/openpi_node/main.py b/lightx2v_ros/src/inference/inference/openpi_node/main.py index 0898f5056..1f68d5d99 100644 --- a/lightx2v_ros/src/inference/inference/openpi_node/main.py +++ b/lightx2v_ros/src/inference/inference/openpi_node/main.py @@ -3,10 +3,10 @@ import numpy as np import rclpy -from common.contract import get_contract +from common.contract import get_contract, parse_multiarray_label from rclpy.node import Node from sensor_msgs.msg import Image -from std_msgs.msg import Float32MultiArray, Float64MultiArray, MultiArrayDimension, String +from std_msgs.msg import Float64MultiArray, MultiArrayDimension, String from lightx2v.models.runners.openpi.openpi_runner import OpenPIPolicy from lightx2v.utils.set_config import auto_calc_config, get_default_config @@ -20,19 +20,9 @@ def __init__(self): self.declare_parameter("model_path", "") self.declare_parameter("seed", 0) self.declare_parameter("actions_per_plan", 5) - self.declare_parameter("numeric_precision", "float64") self.contract = get_contract("libero") - precision = str(self.get_parameter("numeric_precision").value).strip().lower() - numeric_types = { - "float32": (Float32MultiArray, np.float32), - "float64": (Float64MultiArray, np.float64), - } - if precision not in numeric_types: - raise ValueError("numeric_precision must be 'float32' or 'float64'") - self.numeric_message_type, self.numeric_dtype = numeric_types[precision] - config = self._policy_config() self.actions_per_plan = int(self.get_parameter("actions_per_plan").value) action_horizon = int(config["action_horizon"]) @@ -52,12 +42,13 @@ def __init__(self): self.last_processed_observation = None self.pending_actions = deque() - self.action_pub = self.create_publisher(self.numeric_message_type, self.contract.action_topic, 10) - self._camera_subscriptions = [self.create_subscription(Image, self.contract.camera_topic(camera), self._image_callback(camera), 10) for camera in self.contract.policy_input_cameras] - self.state_sub = self.create_subscription(self.numeric_message_type, self.contract.state_topic, self._on_state, 10) - self.context_sub = self.create_subscription(String, self.contract.observation_context_topic, self._on_context, 10) + self.action_pub = self.create_publisher(Float64MultiArray, self.contract.action_topic, 10) + for camera in self.contract.policy_input_cameras: + self.create_subscription(Image, self.contract.camera_topic(camera), self._image_callback(camera), 10) + self.create_subscription(Float64MultiArray, self.contract.state_topic, self._on_state, 10) + self.create_subscription(String, self.contract.observation_context_topic, self._on_context, 10) - self.get_logger().info(f"OpenPI ready on {self.contract.namespace}: horizon={action_horizon}, actions_per_plan={self.actions_per_plan}, precision={precision}") + self.get_logger().info(f"OpenPI ready on {self.contract.namespace}: horizon={action_horizon}, actions_per_plan={self.actions_per_plan}, precision=float64") def _policy_config(self): config_json = str(self.get_parameter("config_json").value).strip() @@ -100,11 +91,11 @@ def callback(msg): def _on_state(self, msg): if not msg.layout.dim: return - identity = state_identity(msg.layout.dim[0].label) + identity = parse_multiarray_label(msg.layout.dim[0].label, "episode", "observation") if identity is None: return episode, observation = identity - state = np.asarray(msg.data, dtype=self.numeric_dtype).reshape(-1) + state = np.asarray(msg.data, dtype=np.float64).reshape(-1) if state.size != self.contract.state_dim: self.get_logger().error(f"expected state length {self.contract.state_dim}, got {state.size}") return @@ -164,12 +155,12 @@ def _try_process_observation(self): self._publish_action(action, episode, observation) def _publish_action(self, action, episode, observation): - action = np.asarray(action, dtype=self.numeric_dtype).reshape(-1) + action = np.asarray(action, dtype=np.float64).reshape(-1) if action.size != self.contract.action_dim: raise ValueError(f"expected action length {self.contract.action_dim}, got {action.size}") if not np.isfinite(action).all(): raise ValueError("OpenPI produced a non-finite action") - msg = self.numeric_message_type() + msg = Float64MultiArray() msg.layout.dim = [ MultiArrayDimension( label=f"episode={episode};observation={observation};plan_epoch={self.plan_epoch}", @@ -191,18 +182,6 @@ def observation_identity(frame_id): return None -def state_identity(label): - fields = {} - for item in str(label).split(";"): - name, separator, value = item.partition("=") - if separator: - fields[name] = value - try: - return int(fields["episode"]), int(fields["observation"]) - except (KeyError, ValueError): - return None - - def image_msg_to_rgb(msg): encoding = msg.encoding.lower() if encoding not in {"rgb8", "bgr8"}: diff --git a/lightx2v_ros/src/simulator/simulator/libero_node/observer.py b/lightx2v_ros/src/simulator/simulator/libero_node/observer.py index a1182e582..bd14cec9b 100644 --- a/lightx2v_ros/src/simulator/simulator/libero_node/observer.py +++ b/lightx2v_ros/src/simulator/simulator/libero_node/observer.py @@ -17,12 +17,6 @@ def default_libero_root(): return Path(__file__).resolve().parent / "LIBERO" -def add_python_path(path): - path = str(Path(path).expanduser()) - if path not in sys.path: - sys.path.insert(0, path) - - def setup_libero_config(libero_root): benchmark_root = libero_root / "libero" / "libero" if not (benchmark_root / "bddl_files").exists(): @@ -50,16 +44,13 @@ def setup_libero_config(libero_root): def load_libero(libero_root): - add_python_path(libero_root) + python_path = str(Path(libero_root).expanduser()) + if python_path not in sys.path: + sys.path.insert(0, python_path) setup_libero_config(libero_root) - try: - from libero.libero import benchmark, get_libero_path - from libero.libero.envs import OffScreenRenderEnv - except ModuleNotFoundError as exc: - if exc.name in {"robosuite", "bddl"}: - raise ModuleNotFoundError(f"Missing dependency '{exc.name}'. Activate the LIBERO runtime first.") from exc - raise + from libero.libero import benchmark, get_libero_path + from libero.libero.envs import OffScreenRenderEnv return benchmark, get_libero_path, OffScreenRenderEnv diff --git a/lightx2v_ros/src/simulator/simulator/sim/node.py b/lightx2v_ros/src/simulator/simulator/sim/node.py index cc216b117..8551d378a 100644 --- a/lightx2v_ros/src/simulator/simulator/sim/node.py +++ b/lightx2v_ros/src/simulator/simulator/sim/node.py @@ -32,7 +32,7 @@ import numpy as np import rclpy -from common.contract import EnvContract +from common.contract import EnvContract, parse_multiarray_label from rclpy.node import Node from sensor_msgs.msg import Image from std_msgs.msg import Bool, Float32MultiArray, Float64MultiArray, Int32, MultiArrayDimension, String @@ -50,18 +50,6 @@ FINISHED_STATES = (SUCCESS, FAILURE) -def parse_action_identity(label): - fields = {} - for item in str(label).split(";"): - name, separator, value = item.partition("=") - if separator: - fields[name] = value - try: - return int(fields["episode"]), int(fields["observation"]), int(fields["plan_epoch"]) - except (KeyError, ValueError): - return None - - def rgb_to_image_msg(image, stamp, frame_id, episode_index=None, observation_index=None): image = np.ascontiguousarray(image) msg = Image() @@ -262,7 +250,7 @@ def on_action(self, msg): if self.state != RUNNING: return if msg.layout.dim: - action_identity = parse_action_identity(msg.layout.dim[0].label) + action_identity = parse_multiarray_label(msg.layout.dim[0].label, "episode", "observation", "plan_epoch") expected_identity = (self.episode_index, self.step_index, self.plan_epoch) if action_identity != expected_identity: self.get_logger().warning(f"dropping stale action {action_identity}; current observation is {expected_identity}") diff --git a/lightx2v_train/configs/train/openpi/pi05_libero.yaml b/lightx2v_train/configs/train/openpi/pi05_libero.yaml new file mode 100644 index 000000000..1f282c1c3 --- /dev/null +++ b/lightx2v_train/configs/train/openpi/pi05_libero.yaml @@ -0,0 +1,81 @@ +model: + name: openpi_pi05_libero + checkpoint_dir: ${oc.env:OPENPI_INITIAL_CHECKPOINT} + pi05: true + discrete_state_input: false + paligemma_variant: gemma_2b + action_expert_variant: gemma_300m + action_dim: 32 + action_horizon: 10 + max_token_len: 200 + dtype: bfloat16 + parameter_dtype: float32 + require_fp32_checkpoint: true + pytorch_compile_mode: null + +distributed: + backend: nccl + timeout_minutes: 60 + sequence_parallel: + enabled: false + size: 1 + dp: + broadcast_buffers: false + find_unused_parameters: true + gradient_as_bucket_view: true + static_graph: false + +data: + train: + name: openpi_libero + repo_id: physical-intelligence/libero + root: ${oc.env:OPENPI_LEROBOT_ROOT} + norm_stats_path: ${oc.env:OPENPI_NORM_STATS_PATH} + tokenizer_path: ${oc.env:OPENPI_INITIAL_CHECKPOINT}/assets/paligemma_tokenizer.model + hf_cache_dir: ${oc.env:HF_HOME} + action_dim: ${model.action_dim} + action_horizon: ${model.action_horizon} + max_token_len: ${model.max_token_len} + global_batch_size: ${oc.env:OPENPI_GLOBAL_BATCH_SIZE,256} + gradient_accumulation_iters: ${training.gradient_accumulation_iters} + num_workers: ${oc.env:OPENPI_DATA_WORKERS,2} + pin_memory: true + shuffle: true + seed: ${training.seed} + +training: + method: openpi_flow_matching + max_train_iters: ${oc.env:OPENPI_MAX_TRAIN_ITERS,30000} + gradient_accumulation_iters: ${oc.env:OPENPI_GRADIENT_ACCUMULATION_ITERS,1} + gradient_checkpointing: true + max_grad_norm: 1.0 + optimizer: + learning_rate: ${training.lr_schedule.peak_lr} + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_epsilon: 0.00000001 + weight_decay: 0.0000000001 + lr_schedule: + warmup_steps: 10000 + peak_lr: 0.00005 + decay_steps: 1000000 + decay_lr: 0.00005 + ema: + decay: 0.999 + save_every_iters: ${oc.env:OPENPI_SAVE_EVERY_ITERS,1000} + save_total_limit: ${oc.env:OPENPI_SAVE_TOTAL_LIMIT,1} + keep_period: ${oc.env:OPENPI_KEEP_PERIOD,5000} + output_dir: ${oc.env:OPENPI_TRAIN_OUTPUT} + seed: ${oc.env:OPENPI_TRAIN_SEED,42} + +inference: + method: none + +logging: + rank_zero_only: true + train_log_every_iters: ${oc.env:OPENPI_LOG_EVERY_ITERS,10} + swanlab: + enable: false + +resume: + checkpoint_path: ${oc.env:OPENPI_RESUME_CHECKPOINT,''} diff --git a/lightx2v_train/lightx2v_train/data/openpi_libero.py b/lightx2v_train/lightx2v_train/data/openpi_libero.py new file mode 100644 index 000000000..2dc05a48e --- /dev/null +++ b/lightx2v_train/lightx2v_train/data/openpi_libero.py @@ -0,0 +1,403 @@ +"""Local LeRobot input pipeline for OpenPI pi0.5-LIBERO training.""" + +from __future__ import annotations + +import hashlib +import json +import os +import random +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import sentencepiece +import torch +from PIL import Image +from loguru import logger + +from lightx2v_train.runtime.distributed import get_rank, get_world_size, is_distributed +from lightx2v_train.utils.registry import DATA_REGISTER + +LIBERO_STATE_DIM = 8 +LIBERO_ACTION_DIM = 7 +MODEL_IMAGE_SIZE = 224 + +IMAGE_PIPELINE_CONTRACT = { + "target_size": MODEL_IMAGE_SIZE, + "resize": "openpi_client_pil_bilinear_uint8_resize_with_pad", + "quantization": "float01_to_uint8_truncate_before_resize", + "augmentation": "openpi_augmax_0.4.1_per_sample", +} + + +@dataclass +class OpenPIObservation: + images: dict[str, torch.Tensor] + image_masks: dict[str, torch.Tensor] + state: torch.Tensor + tokenized_prompt: torch.Tensor + tokenized_prompt_mask: torch.Tensor + token_ar_mask: torch.Tensor | None = None + token_loss_mask: torch.Tensor | None = None + + def to(self, device: torch.device, *, non_blocking: bool = False) -> "OpenPIObservation": + def move(value): + return None if value is None else value.to(device, non_blocking=non_blocking) + + return OpenPIObservation( + images={name: move(image) for name, image in self.images.items()}, + image_masks={name: move(mask) for name, mask in self.image_masks.items()}, + state=move(self.state), + tokenized_prompt=move(self.tokenized_prompt), + tokenized_prompt_mask=move(self.tokenized_prompt_mask), + token_ar_mask=move(self.token_ar_mask), + token_loss_mask=move(self.token_loss_mask), + ) + + def pin_memory(self) -> "OpenPIObservation": + def pin(value): + return None if value is None else value.pin_memory() + + return OpenPIObservation( + images={name: pin(image) for name, image in self.images.items()}, + image_masks={name: pin(mask) for name, mask in self.image_masks.items()}, + state=pin(self.state), + tokenized_prompt=pin(self.tokenized_prompt), + tokenized_prompt_mask=pin(self.tokenized_prompt_mask), + token_ar_mask=pin(self.token_ar_mask), + token_loss_mask=pin(self.token_loss_mask), + ) + + +class _PaligemmaTokenizer: + def __init__(self, path: Path, max_token_len: int): + if not path.is_file(): + raise FileNotFoundError(f"PaliGemma tokenizer not found: {path}") + self.max_token_len = max_token_len + self.processor = sentencepiece.SentencePieceProcessor(model_proto=path.read_bytes()) + + def tokenize(self, prompt: str) -> tuple[np.ndarray, np.ndarray]: + prompt = prompt.strip().replace("_", " ").replace("\n", " ") + tokens = self.processor.encode(prompt, add_bos=True) + self.processor.encode("\n") + if len(tokens) > self.max_token_len: + logger.warning("Prompt has {} tokens and will be truncated to {}", len(tokens), self.max_token_len) + tokens = tokens[: self.max_token_len] + mask = [True] * len(tokens) + padding = self.max_token_len - len(tokens) + tokens.extend([0] * padding) + mask.extend([False] * padding) + return np.asarray(tokens, dtype=np.int64), np.asarray(mask, dtype=np.bool_) + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _load_quantile_stats(path: Path) -> dict[str, dict[str, np.ndarray]]: + if not path.is_file(): + raise FileNotFoundError(f"OpenPI normalization statistics not found: {path}") + with path.open(encoding="utf-8") as stream: + stats = json.load(stream)["norm_stats"] + result = {} + for key, expected_dim in (("state", LIBERO_STATE_DIM), ("actions", LIBERO_ACTION_DIM)): + if key not in stats: + raise ValueError(f"Normalization statistics have no {key!r} entry: {path}") + values = {name: np.asarray(value) for name, value in stats[key].items()} + for name in ("q01", "q99"): + if name not in values or values[name].shape[-1] < expected_dim: + raise ValueError(f"Invalid {key}.{name} in {path}: expected at least {expected_dim} values") + if not np.isfinite(values[name][..., :expected_dim]).all(): + raise ValueError(f"Non-finite values in {key}.{name}: {path}") + if np.any(values["q99"][..., :expected_dim] < values["q01"][..., :expected_dim]): + raise ValueError(f"Every {key}.q99 value must be greater than or equal to q01: {path}") + result[key] = values + return result + + +def _normalize_quantile(value: np.ndarray, stats: dict[str, np.ndarray]) -> np.ndarray: + q01 = stats["q01"][..., : value.shape[-1]] + q99 = stats["q99"][..., : value.shape[-1]] + return ((value - q01) / (q99 - q01 + 1e-6) * 2.0 - 1.0).astype(np.float32) + + +def _as_numpy(value: Any) -> np.ndarray: + if torch.is_tensor(value): + return value.detach().cpu().numpy() + return np.asarray(value) + + +def _image_as_uint8_hwc(value: Any, key: str) -> np.ndarray: + image = _as_numpy(value) + if np.issubdtype(image.dtype, np.floating): + image = (255 * image).astype(np.uint8) + if image.ndim == 3 and image.shape[0] == 3: + image = np.transpose(image, (1, 2, 0)) + if image.ndim != 3 or image.shape[-1] != 3 or image.dtype != np.uint8: + raise ValueError(f"{key} must be an RGB image, got shape={image.shape}, dtype={image.dtype}") + return np.ascontiguousarray(image) + + +def _resize_with_pad_uint8(image: np.ndarray, height: int, width: int) -> np.ndarray: + if image.shape[:2] == (height, width): + return image + current_height, current_width = image.shape[:2] + ratio = max(current_width / width, current_height / height) + resized_height = int(current_height / ratio) + resized_width = int(current_width / ratio) + resized = Image.fromarray(image).resize((resized_width, resized_height), resample=Image.Resampling.BILINEAR) + canvas = Image.new(resized.mode, (width, height), 0) + canvas.paste(resized, ((width - resized_width) // 2, (height - resized_height) // 2)) + return np.asarray(canvas, dtype=np.uint8) + + +def _pad_last_dim(value: np.ndarray, size: int) -> np.ndarray: + if value.shape[-1] > size: + raise ValueError(f"Cannot pad dimension {value.shape[-1]} to the smaller size {size}") + if value.shape[-1] == size: + return value + widths = [(0, 0)] * value.ndim + widths[-1] = (0, size - value.shape[-1]) + return np.pad(value, widths) + + +class OpenPILiberoDataset(torch.utils.data.Dataset): + """Apply the official LIBERO repack/normalize/tokenize/pad contract locally.""" + + def __init__( + self, + dataset, + tasks: dict[int, str], + norm_stats_path: Path, + tokenizer_path: Path, + *, + action_horizon: int, + action_dim: int, + max_token_len: int, + ): + self.dataset = dataset + self.tasks = {int(index): prompt for index, prompt in tasks.items()} + self.norm_stats = _load_quantile_stats(norm_stats_path) + self.tokenizer = _PaligemmaTokenizer(tokenizer_path, max_token_len) + self.action_horizon = action_horizon + self.action_dim = action_dim + + def __len__(self) -> int: + return len(self.dataset) + + def __getitem__(self, index: int) -> dict[str, Any]: + item = self.dataset[index] + task_index = int(_as_numpy(item["task_index"]).reshape(())) + prompt = self.tasks[task_index] + + base_image = _resize_with_pad_uint8(_image_as_uint8_hwc(item["image"], "image"), MODEL_IMAGE_SIZE, MODEL_IMAGE_SIZE) + wrist_image = _resize_with_pad_uint8(_image_as_uint8_hwc(item["wrist_image"], "wrist_image"), MODEL_IMAGE_SIZE, MODEL_IMAGE_SIZE) + state = _as_numpy(item["state"]).astype(np.float32, copy=False) + actions = _as_numpy(item["actions"]).astype(np.float32, copy=False) + if state.shape != (LIBERO_STATE_DIM,): + raise ValueError(f"state must have shape ({LIBERO_STATE_DIM},), got {state.shape}") + if actions.shape != (self.action_horizon, LIBERO_ACTION_DIM): + raise ValueError(f"actions must have shape ({self.action_horizon}, {LIBERO_ACTION_DIM}), got {actions.shape}") + + state = _pad_last_dim(_normalize_quantile(state, self.norm_stats["state"]), self.action_dim) + actions = _pad_last_dim(_normalize_quantile(actions, self.norm_stats["actions"]), self.action_dim) + tokens, token_mask = self.tokenizer.tokenize(prompt) + return { + "images": { + "base_0_rgb": base_image, + "left_wrist_0_rgb": wrist_image, + "right_wrist_0_rgb": np.zeros_like(base_image), + }, + "image_masks": { + "base_0_rgb": True, + "left_wrist_0_rgb": True, + "right_wrist_0_rgb": False, + }, + "state": state, + "tokenized_prompt": tokens, + "tokenized_prompt_mask": token_mask, + "actions": actions, + } + + +def _collate(items: list[dict[str, Any]]) -> tuple[OpenPIObservation, torch.Tensor]: + image_names = tuple(items[0]["images"]) + images = {} + masks = {} + for name in image_names: + batch = torch.from_numpy(np.stack([item["images"][name] for item in items])) + images[name] = batch.permute(0, 3, 1, 2).to(torch.float32).div_(127.5).sub_(1.0) + masks[name] = torch.as_tensor([item["image_masks"][name] for item in items], dtype=torch.bool) + + observation = OpenPIObservation( + images=images, + image_masks=masks, + state=torch.from_numpy(np.stack([item["state"] for item in items])), + tokenized_prompt=torch.from_numpy(np.stack([item["tokenized_prompt"] for item in items])), + tokenized_prompt_mask=torch.from_numpy(np.stack([item["tokenized_prompt_mask"] for item in items])), + ) + actions = torch.from_numpy(np.stack([item["actions"] for item in items])) + return observation, actions + + +def _seed_worker(worker_id: int) -> None: + del worker_id + seed = torch.initial_seed() % 2**32 + np.random.seed(seed) + random.seed(seed) + + +class OpenPIDataLoader: + def __init__( + self, + loader: torch.utils.data.DataLoader, + generator: torch.Generator, + seed: int, + batches_per_epoch: int, + metadata: dict[str, Any], + ): + self.loader = loader + self.generator = generator + self.seed = seed + self.batches_per_epoch = batches_per_epoch + self.metadata = metadata + + def __len__(self) -> int: + return self.batches_per_epoch + + def __iter__(self): + iterator = iter(self.loader) + for _ in range(self.batches_per_epoch): + yield next(iterator) + + def set_epoch(self, epoch: int) -> None: + self.generator.manual_seed(self.seed + epoch) + if hasattr(self.loader.sampler, "set_epoch"): + self.loader.sampler.set_epoch(epoch) + + +@DATA_REGISTER("openpi_libero") +def build_openpi_libero(config: dict[str, Any], train_or_val: str): + if train_or_val != "train": + raise ValueError("OpenPI LIBERO integration currently provides a training split only") + + dataset_root = Path(config["root"]).expanduser().resolve() + norm_stats_path = Path(config["norm_stats_path"]).expanduser().resolve() + tokenizer_path = Path(config["tokenizer_path"]).expanduser().resolve() + cache_dir = Path(config["hf_cache_dir"]).expanduser().resolve() + if not (dataset_root / "meta/info.json").is_file(): + raise FileNotFoundError(f"LeRobot dataset is incomplete or missing: {dataset_root}") + cache_dir.mkdir(parents=True, exist_ok=True) + os.environ["HF_HOME"] = str(cache_dir) + os.environ["HF_DATASETS_CACHE"] = str(cache_dir / "datasets") + + from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata + + repo_id = str(config["repo_id"]) + metadata = LeRobotDatasetMetadata(repo_id, root=dataset_root) + action_horizon = int(config["action_horizon"]) + delta_timestamps = {"actions": [step / metadata.fps for step in range(action_horizon)]} + source = LeRobotDataset( + repo_id, + root=dataset_root, + delta_timestamps=delta_timestamps, + download_videos=False, + ) + dataset = OpenPILiberoDataset( + source, + metadata.tasks, + norm_stats_path, + tokenizer_path, + action_horizon=action_horizon, + action_dim=int(config["action_dim"]), + max_token_len=int(config["max_token_len"]), + ) + + global_batch_size = int(config["global_batch_size"]) + accumulation = int(config["gradient_accumulation_iters"]) + divisor = get_world_size() * accumulation + if global_batch_size % divisor: + raise ValueError(f"data.train.global_batch_size={global_batch_size} must be divisible by world_size * gradient_accumulation_iters={divisor}") + batch_size = global_batch_size // divisor + if batch_size < 1: + raise ValueError("Per-rank micro batch size must be at least one") + + sampler = None + shuffle = bool(config["shuffle"]) + seed = int(config["seed"]) + if is_distributed(): + sampler = torch.utils.data.DistributedSampler( + dataset, + num_replicas=get_world_size(), + rank=get_rank(), + shuffle=shuffle, + seed=seed, + drop_last=True, + ) + shuffle = False + generator = torch.Generator() + generator.manual_seed(seed) + num_workers = int(config["num_workers"]) + multiprocessing_options = {"multiprocessing_context": "spawn"} if num_workers > 0 else {} + loader = torch.utils.data.DataLoader( + dataset, + batch_size=batch_size, + shuffle=shuffle, + sampler=sampler, + num_workers=num_workers, + persistent_workers=num_workers > 0, + pin_memory=bool(config["pin_memory"]), + drop_last=True, + collate_fn=_collate, + worker_init_fn=_seed_worker, + generator=generator, + **multiprocessing_options, + ) + optimizer_batches_per_epoch = len(dataset) // global_batch_size + batches_per_epoch = optimizer_batches_per_epoch * accumulation + if batches_per_epoch < 1 or batches_per_epoch > len(loader): + raise RuntimeError( + f"Cannot form a complete optimizer batch from the dataset: dataset={len(dataset)}, global_batch={global_batch_size}, micro_batches={len(loader)}, required={batches_per_epoch}" + ) + logger.info( + "[openpi:data] root={} episodes={} frames={} language_tasks={} fps={} per_rank_batch={} world_size={} grad_accum={} global_batch={}", + dataset_root, + metadata.total_episodes, + metadata.total_frames, + metadata.total_tasks, + metadata.fps, + batch_size, + get_world_size(), + accumulation, + global_batch_size, + ) + return OpenPIDataLoader( + loader, + generator, + seed, + batches_per_epoch, + { + "root": str(dataset_root), + "repo_id": repo_id, + "norm_stats_path": str(norm_stats_path), + "tokenizer_path": str(tokenizer_path), + "episodes": metadata.total_episodes, + "frames": metadata.total_frames, + "language_tasks": metadata.total_tasks, + "fps": metadata.fps, + "global_batch_size": global_batch_size, + "per_rank_batch_size": batch_size, + "optimizer_batches_per_epoch": optimizer_batches_per_epoch, + "micro_batches_per_epoch": batches_per_epoch, + "shuffle": bool(config["shuffle"]), + "seed": seed, + "dataset_info_sha256": _sha256(dataset_root / "meta/info.json"), + "norm_stats_sha256": _sha256(norm_stats_path), + "tokenizer_sha256": _sha256(tokenizer_path), + "image_pipeline": IMAGE_PIPELINE_CONTRACT, + }, + ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/openpi/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/openpi/__init__.py new file mode 100644 index 000000000..a90b9159c --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/openpi/__init__.py @@ -0,0 +1,5 @@ +"""OpenPI training models.""" + +from .pi05_libero import OpenPIPi05LiberoModel + +__all__ = ["OpenPIPi05LiberoModel"] diff --git a/lightx2v_train/lightx2v_train/model_zoo/openpi/pi05_libero.py b/lightx2v_train/lightx2v_train/model_zoo/openpi/pi05_libero.py new file mode 100644 index 000000000..1b5571029 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/openpi/pi05_libero.py @@ -0,0 +1,72 @@ +"""Training wrapper around LightX2V's localized OpenPI pi0.5 network.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import torch +from loguru import logger + +from lightx2v_train.runtime.distributed import get_device +from lightx2v_train.utils.registry import MODEL_REGISTER + + +@MODEL_REGISTER("openpi_pi05_libero") +class OpenPIPi05LiberoModel: + """Load a converted SafeTensors checkpoint as model-only initialization.""" + + def __init__(self, config: dict[str, Any]): + self.model_config = dict(config["model"]) + self.initialization_path = Path(self.model_config["checkpoint_dir"]).expanduser().resolve() + self.device = get_device() + self.core_model: torch.nn.Module | None = None + + def load_components(self, *, load_transformer: bool, load_vae: bool, load_condition_encoder: bool) -> None: + del load_vae, load_condition_encoder + if not load_transformer: + raise ValueError("OpenPI training requires load_transformer=true") + if self.core_model is not None: + return + + weight_path = self.initialization_path / "model.safetensors" + if not weight_path.is_file(): + raise FileNotFoundError(f"Converted OpenPI weights not found: {weight_path}") + + from lightx2v.models.networks.openpi.config import Pi0Config + from lightx2v.models.networks.openpi.weights import load_pi05_libero_weights + + if self.model_config["parameter_dtype"] != "float32": + raise ValueError("OpenPI training requires model.parameter_dtype='float32'") + if not self.model_config["require_fp32_checkpoint"]: + raise ValueError("OpenPI training requires model.require_fp32_checkpoint=true") + pi0_config = Pi0Config.from_mapping(self.model_config) + pi0_config.validate_pi05_libero() + self.core_model = load_pi05_libero_weights(weight_path, pi0_config, self.device) + self.core_model.train() + logger.info("[openpi:model] model-only initialization checkpoint={}", self.initialization_path) + + def require_core_model(self) -> torch.nn.Module: + if self.core_model is None: + raise RuntimeError("OpenPI model components have not been loaded") + return self.core_model + + def enable_gradient_checkpointing(self) -> None: + model = self.require_core_model() + model.gradient_checkpointing_enable() + + def architecture_metadata(self) -> dict[str, Any]: + keys = ( + "pi05", + "discrete_state_input", + "paligemma_variant", + "action_expert_variant", + "action_dim", + "action_horizon", + "max_token_len", + ) + metadata = {key: self.model_config[key] for key in keys} + metadata["compute_dtype"] = self.model_config["dtype"] + metadata["parameter_dtype"] = self.model_config["parameter_dtype"] + metadata["require_fp32_checkpoint"] = self.model_config["require_fp32_checkpoint"] + return metadata diff --git a/lightx2v_train/lightx2v_train/trainers/openpi.py b/lightx2v_train/lightx2v_train/trainers/openpi.py new file mode 100644 index 000000000..637887f46 --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/openpi.py @@ -0,0 +1,747 @@ +"""PyTorch fine-tuning loop for the localized OpenPI pi0.5-LIBERO model.""" + +from __future__ import annotations + +import json +import math +import os +import random +import re +import shutil +import uuid +from contextlib import nullcontext +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +import numpy as np +import torch +import torch.distributed as dist +from loguru import logger +from safetensors import safe_open +from safetensors.torch import load_model, save_file, save_model +from torch.nn.parallel import DistributedDataParallel + +from lightx2v_train.runtime.distributed import ( + get_rank, + get_world_size, + is_distributed, + is_main_process, + reduce_mean, +) +from lightx2v_train.runtime.monitor import build_monitor +from lightx2v_train.utils.registry import TRAINER_REGISTER + +CHECKPOINT_SCHEMA_VERSION = 4 +CHECKPOINT_CREATOR = "lightx2v-openpi-trainer" +CHECKPOINT_PATTERN = re.compile(r"checkpoint-(\d{9})$") +REQUIRED_CHECKPOINT_FILES = ( + "model.safetensors", + "config.json", + "assets/paligemma_tokenizer.model", + "assets/physical-intelligence/libero/norm_stats.json", + "ema/model.safetensors", + "ema/config.json", + "ema/assets/paligemma_tokenizer.model", + "ema/assets/physical-intelligence/libero/norm_stats.json", + "training_state.pt", + "manifest.json", + "_SUCCESS", +) +NUMERICAL_POLICY = { + "parameter_dtype": "float32", + "gradient_dtype": "float32", + "optimizer_state_dtype": "float32", + "ema_dtype": "float32", + "loss_dtype": "float32", + "grad_scaler": False, + "tf32": False, +} + + +def _assert_fp32_gradients(model: torch.nn.Module) -> None: + wrong = { + name: str(parameter.grad.dtype) for name, parameter in model.named_parameters() if parameter.grad is not None and parameter.grad.is_floating_point() and parameter.grad.dtype != torch.float32 + } + if wrong: + preview = dict(list(wrong.items())[:20]) + raise RuntimeError(f"OpenPI gradients must accumulate in float32; found {len(wrong)} mismatches: {preview}") + + +def _assert_fp32_optimizer_state(optimizer: torch.optim.Optimizer) -> None: + wrong: dict[str, str] = {} + for parameter_index, state in enumerate(optimizer.state.values()): + for name, value in state.items(): + if torch.is_tensor(value) and value.is_floating_point() and value.dtype != torch.float32: + wrong[f"parameter_{parameter_index}.{name}"] = str(value.dtype) + if wrong: + preview = dict(list(wrong.items())[:20]) + raise RuntimeError(f"OpenPI optimizer floating-point state must be float32; found {len(wrong)} mismatches: {preview}") + + +def _validate_full_checkpoint_weight_precision(checkpoint_dir: Path) -> None: + """Reject precision-lossy model masters before mutating resume state.""" + weight_files = { + "online model": checkpoint_dir / "model.safetensors", + "EMA model": checkpoint_dir / "ema/model.safetensors", + } + for kind, path in weight_files.items(): + if not path.is_file(): + raise FileNotFoundError(f"OpenPI schema {CHECKPOINT_SCHEMA_VERSION} {kind} checkpoint is missing: {path}") + wrong: dict[str, str] = {} + with safe_open(path, framework="pt", device="cpu") as checkpoint: + for name in checkpoint.keys(): + dtype = checkpoint.get_slice(name).get_dtype() + if (dtype.startswith("F") or dtype.startswith("BF")) and dtype != "F32": + wrong[name] = dtype + if len(wrong) == 20: + break + if wrong: + raise RuntimeError(f"OpenPI schema {CHECKPOINT_SCHEMA_VERSION} {kind} checkpoint must keep every floating-point tensor in FP32; found lower-precision tensors in {path}: {wrong}") + + +def _shared_aliases(state_dict: dict[str, torch.Tensor]) -> dict[str, str]: + """Return alias-to-keeper names for exact tied tensors in a live model.""" + storage_groups: dict[tuple[torch.device, int, int], list[str]] = {} + for name, tensor in state_dict.items(): + if tensor.device.type == "meta" or tensor.numel() == 0: + continue + storage = tensor.untyped_storage() + identity = (tensor.device, storage.data_ptr(), storage.nbytes()) + storage_groups.setdefault(identity, []).append(name) + + aliases = {} + for names in storage_groups.values(): + if len(names) < 2: + continue + names.sort() + keeper = names[0] + keeper_tensor = state_dict[keeper] + keeper_storage = keeper_tensor.untyped_storage() + if keeper_tensor.data_ptr() != keeper_storage.data_ptr() or keeper_tensor.numel() * keeper_tensor.element_size() != keeper_storage.nbytes(): + raise RuntimeError(f"EMA tied tensor keeper {keeper!r} does not cover its complete storage") + keeper_signature = ( + keeper_tensor.dtype, + keeper_tensor.shape, + keeper_tensor.stride(), + keeper_tensor.storage_offset(), + ) + for name in names[1:]: + tensor = state_dict[name] + signature = (tensor.dtype, tensor.shape, tensor.stride(), tensor.storage_offset()) + if signature != keeper_signature: + raise RuntimeError(f"EMA checkpointing only supports exact tied tensor aliases; {keeper!r} and {name!r} share storage but have different views") + aliases[name] = keeper + return aliases + + +def _tensor_storage_bytes(value: Any) -> int: + """Count unique tensor storage bytes in a nested checkpoint payload.""" + seen: set[tuple[torch.device, int, int]] = set() + total = 0 + + def visit(item: Any) -> None: + nonlocal total + if torch.is_tensor(item): + if item.device.type == "meta" or item.numel() == 0: + return + storage = item.untyped_storage() + identity = (item.device, storage.data_ptr(), storage.nbytes()) + if identity not in seen: + seen.add(identity) + total += storage.nbytes() + elif isinstance(item, dict): + for nested in item.values(): + visit(nested) + elif isinstance(item, (list, tuple)): + for nested in item: + visit(nested) + + visit(value) + return total + + +def _write_json(path: Path, value: dict[str, Any]) -> None: + path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def _set_seed(seed: int) -> None: + seed += get_rank() + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + +def _capture_rng_state() -> dict[str, Any]: + numpy_state = np.random.get_state() + return { + "python": list(random.getstate()), + "numpy": { + "bit_generator": numpy_state[0], + "state": numpy_state[1].tolist(), + "position": numpy_state[2], + "has_gauss": numpy_state[3], + "cached_gaussian": numpy_state[4], + }, + "torch_cpu": torch.get_rng_state(), + "torch_cuda": torch.cuda.get_rng_state() if torch.cuda.is_available() else None, + } + + +def _gather_rng_states() -> list[dict[str, Any]]: + local_state = _capture_rng_state() + if not is_distributed(): + return [local_state] + states: list[dict[str, Any] | None] = [None] * get_world_size() + dist.all_gather_object(states, local_state) + if any(state is None for state in states): + raise RuntimeError("Failed to collect RNG state from every distributed rank") + return states # type: ignore[return-value] + + +def _nested_tuple(value): + if isinstance(value, list): + return tuple(_nested_tuple(item) for item in value) + return value + + +def _restore_rng_state(state: dict[str, Any]) -> None: + random.setstate(_nested_tuple(state["python"])) + numpy_state = state["numpy"] + np.random.set_state( + ( + numpy_state["bit_generator"], + np.asarray(numpy_state["state"], dtype=np.uint32), + int(numpy_state["position"]), + int(numpy_state["has_gauss"]), + float(numpy_state["cached_gaussian"]), + ) + ) + torch.set_rng_state(state["torch_cpu"]) + if torch.cuda.is_available(): + if state["torch_cuda"] is None: + raise RuntimeError("Checkpoint has no CUDA RNG state for a CUDA resume") + torch.cuda.set_rng_state(state["torch_cuda"]) + + +def _build_scheduler(optimizer: torch.optim.Optimizer, spec: dict[str, float | int]): + warmup_steps = int(spec["warmup_steps"]) + peak_lr = float(spec["peak_lr"]) + decay_steps = int(spec["decay_steps"]) + decay_lr = float(spec["decay_lr"]) + if peak_lr <= 0 or decay_lr < 0: + raise ValueError("Learning rates must be non-negative and peak_lr must be positive") + if warmup_steps < 0 or decay_steps < 1: + raise ValueError("warmup_steps must be non-negative and decay_steps must be positive") + + def factor(step: int) -> float: + if warmup_steps and step < warmup_steps: + initial_lr = peak_lr / (warmup_steps + 1) + value = initial_lr + (peak_lr - initial_lr) * step / warmup_steps + else: + progress = min(1.0, (step - warmup_steps) / max(1, decay_steps - warmup_steps)) + value = decay_lr + (peak_lr - decay_lr) * 0.5 * (1 + math.cos(math.pi * progress)) + return value / peak_lr + + return torch.optim.lr_scheduler.LambdaLR(optimizer, factor) + + +class ExponentialMovingAverage: + def __init__(self, model: torch.nn.Module, decay: float): + if not 0.0 < decay < 1.0: + raise ValueError(f"EMA decay must be between zero and one, got {decay}") + self.decay = decay + self.num_updates = 0 + self.shadow = {name: value.detach().to(dtype=torch.float32).clone() if value.is_floating_point() else value.detach().clone() for name, value in model.state_dict().items()} + + @torch.no_grad() + def update(self, model: torch.nn.Module) -> None: + current = model.state_dict() + for name, value in current.items(): + target = self.shadow[name] + if target.is_floating_point(): + target.mul_(self.decay).add_(value.detach(), alpha=1.0 - self.decay) + else: + target.copy_(value) + self.num_updates += 1 + + def inference_state(self, model: torch.nn.Module) -> tuple[dict[str, torch.Tensor], dict[str, str]]: + model_state = model.state_dict() + if model_state.keys() != self.shadow.keys(): + raise RuntimeError("Model state changed after EMA initialization") + aliases = _shared_aliases(model_state) + state = {name: value.detach().cpu().contiguous() for name, value in self.shadow.items() if name not in aliases} + return state, aliases + + def save(self, path: Path, model: torch.nn.Module) -> None: + state, aliases = self.inference_state(model) + metadata = {"format": "pt", "kind": "ema", "master_dtype": "float32", **aliases} + save_file(state, str(path), metadata=metadata) + + def load(self, path: Path, model: torch.nn.Module, *, num_updates: int) -> None: + expected = model.state_dict() + aliases = _shared_aliases(expected) + required = set(expected) - set(aliases) + if self.shadow.keys() != expected.keys(): + raise RuntimeError("Model state changed after EMA initialization") + + # Keep the existing FP32 shadow allocations and stream one tensor at a + # time from the CPU-mapped SafeTensors file. Loading the complete EMA on + # GPU and then cloning it would add roughly two model copies to resume's + # peak memory. + with safe_open(path, framework="pt", device="cpu") as checkpoint: + loaded = set(checkpoint.keys()) + missing = required - loaded + unexpected = loaded - set(expected) + wrong_shape = { + name: (tuple(checkpoint.get_slice(name).get_shape()), tuple(expected[name].shape)) + for name in required & loaded + if tuple(checkpoint.get_slice(name).get_shape()) != tuple(expected[name].shape) + } + if missing or unexpected or wrong_shape: + raise RuntimeError(f"EMA checkpoint does not match the model: missing={sorted(missing)}, unexpected={sorted(unexpected)}, wrong_shape={wrong_shape}") + for name in sorted(required): + self.shadow[name].copy_(checkpoint.get_tensor(name)) + + for alias, keeper in aliases.items(): + self.shadow[alias].copy_(self.shadow[keeper]) + self.num_updates = int(num_updates) + + +def _is_complete_checkpoint(path: Path) -> bool: + if not path.is_dir() or not CHECKPOINT_PATTERN.fullmatch(path.name): + return False + if any(not (path / name).is_file() for name in REQUIRED_CHECKPOINT_FILES): + return False + try: + manifest = json.loads((path / "manifest.json").read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return False + return manifest.get("created_by") == CHECKPOINT_CREATOR and manifest.get("schema_version") == CHECKPOINT_SCHEMA_VERSION + + +def _prune_checkpoints(output_dir: Path, total_limit: int | None, keep_period: int | None) -> None: + if total_limit is None or total_limit < 1 or not output_dir.is_dir(): + return + checkpoints = sorted( + (path for path in output_dir.iterdir() if _is_complete_checkpoint(path)), + key=lambda path: int(CHECKPOINT_PATTERN.fullmatch(path.name).group(1)), + ) + disposable = [] + for path in checkpoints: + step = int(CHECKPOINT_PATTERN.fullmatch(path.name).group(1)) + if keep_period is None or step % keep_period: + disposable.append(path) + for path in disposable[: max(0, len(disposable) - total_limit)]: + shutil.rmtree(path) + + +@TRAINER_REGISTER("openpi_flow_matching") +class OpenPIFlowMatchingTrainer: + """Fine-tune PI0Pytorch with strict full-state resume and EMA.""" + + def __init__(self, config: dict[str, Any]): + self.config = config + self.training_config = config["training"] + self.resume_config = config["resume"] + self.output_dir = Path(self.training_config["output_dir"]).expanduser().resolve() + self.max_train_iters = int(self.training_config["max_train_iters"]) + self.gradient_accumulation_iters = int(self.training_config["gradient_accumulation_iters"]) + self.max_grad_norm = float(self.training_config["max_grad_norm"]) + self.save_every_iters = int(self.training_config["save_every_iters"]) + self.save_total_limit = self.training_config["save_total_limit"] + keep_period = self.training_config["keep_period"] + self.keep_period = None if keep_period is None else int(keep_period) + if self.keep_period is not None and self.keep_period < 1: + raise ValueError("training.keep_period must be positive or null") + self.log_every_iters = int(config["logging"]["train_log_every_iters"]) + if self.log_every_iters < 1: + raise ValueError("logging.train_log_every_iters must be positive") + optimizer_config = self.training_config["optimizer"] + self.optimizer_spec = { + "learning_rate": float(optimizer_config["learning_rate"]), + "adam_beta1": float(optimizer_config["adam_beta1"]), + "adam_beta2": float(optimizer_config["adam_beta2"]), + "weight_decay": float(optimizer_config["weight_decay"]), + "adam_epsilon": float(optimizer_config["adam_epsilon"]), + } + schedule_config = self.training_config["lr_schedule"] + self.schedule_spec = { + "warmup_steps": int(schedule_config["warmup_steps"]), + "peak_lr": float(schedule_config["peak_lr"]), + "decay_steps": int(schedule_config["decay_steps"]), + "decay_lr": float(schedule_config["decay_lr"]), + } + ema_config = self.training_config["ema"] + self.ema_decay = float(ema_config["decay"]) + if self.gradient_accumulation_iters < 1: + raise ValueError("training.gradient_accumulation_iters must be positive") + data_accumulation = int(config["data"]["train"]["gradient_accumulation_iters"]) + if data_accumulation != self.gradient_accumulation_iters: + raise ValueError("data.train.gradient_accumulation_iters must equal training.gradient_accumulation_iters") + self.monitor = build_monitor(config) + + def set_model(self, model) -> None: + self.model = model + self.core_model = model.require_core_model() + + def set_data(self, dataloader_train, dataloader_eval=None) -> None: + if dataloader_eval is not None: + raise ValueError("OpenPI trainer does not run an in-training LIBERO rollout") + if not hasattr(dataloader_train, "metadata"): + raise TypeError("OpenPI trainer requires the openpi_libero data loader") + self.dataloader = dataloader_train + + def _wrap_ddp(self) -> torch.nn.Module: + if not is_distributed(): + return self.core_model + dp_config = self.config["distributed"]["dp"] + kwargs = { + "broadcast_buffers": bool(dp_config["broadcast_buffers"]), + "find_unused_parameters": bool(dp_config["find_unused_parameters"]), + "gradient_as_bucket_view": bool(dp_config["gradient_as_bucket_view"]), + "static_graph": bool(dp_config["static_graph"]), + } + if torch.cuda.is_available(): + kwargs.update(device_ids=[torch.cuda.current_device()], output_device=torch.cuda.current_device()) + logger.info("[openpi:ddp] wrapping model with {}", kwargs) + return DistributedDataParallel(self.core_model, **kwargs) + + def _resolve_resume_path(self) -> Path | None: + explicit = self.resume_config["checkpoint_path"] + if explicit: + path = Path(explicit).expanduser().resolve() + if not _is_complete_checkpoint(path): + raise RuntimeError(f"Not a complete LightX2V OpenPI checkpoint: {path}") + return path + candidates = [] + if self.output_dir.is_dir(): + candidates = [path for path in self.output_dir.iterdir() if CHECKPOINT_PATTERN.fullmatch(path.name)] + incomplete = [path for path in candidates if not _is_complete_checkpoint(path)] + if incomplete: + raise RuntimeError(f"Output directory contains incomplete or foreign checkpoints: {incomplete}") + if candidates: + raise RuntimeError(f"Resume is disabled but {self.output_dir} already contains checkpoints; use run_pi05_resume_ema.sh or choose a new OPENPI_TRAIN_OUTPUT") + return None + + def _checkpoint_contract(self) -> dict[str, Any]: + data = self.dataloader.metadata + image_pipeline = data["image_pipeline"] + architecture = self.model.architecture_metadata() + compute_dtype = architecture["compute_dtype"] + parameter_dtype = architecture["parameter_dtype"] + if compute_dtype != "bfloat16" or parameter_dtype != "float32": + raise RuntimeError( + f"The OpenPI training contract requires float32 canonical parameters and bfloat16 transformer compute; got parameter_dtype={parameter_dtype!r}, compute_dtype={compute_dtype!r}" + ) + return { + "architecture": architecture, + "numerical_policy": {**NUMERICAL_POLICY, "compute_dtype": compute_dtype}, + "optimizer": self.optimizer_spec, + "lr_schedule": self.schedule_spec, + "ema_decay": self.ema_decay, + "ema_master_dtype": "float32", + "gradient_accumulation_iters": self.gradient_accumulation_iters, + "gradient_checkpointing": bool(self.training_config["gradient_checkpointing"]), + "max_grad_norm": self.max_grad_norm, + "dataset": { + key: data[key] + for key in ( + "repo_id", + "episodes", + "frames", + "language_tasks", + "fps", + "global_batch_size", + "optimizer_batches_per_epoch", + "micro_batches_per_epoch", + "shuffle", + "seed", + "dataset_info_sha256", + "norm_stats_sha256", + "tokenizer_sha256", + ) + }, + "image_pipeline": image_pipeline, + } + + def _load_full_checkpoint(self, path: Path) -> tuple[int, int, int]: + state = torch.load(path / "training_state.pt", map_location="cpu", weights_only=True) + if state.get("schema_version") != CHECKPOINT_SCHEMA_VERSION: + raise RuntimeError(f"Unsupported training-state schema in {path}") + if state.get("contract") != self._checkpoint_contract(): + raise RuntimeError(f"Checkpoint training contract does not match the current config: {path}") + checkpoint_world_size = int(state["world_size"]) + if checkpoint_world_size != get_world_size(): + raise RuntimeError(f"Checkpoint world_size={checkpoint_world_size}, current world_size={get_world_size()}; strict resume requires the same world size") + + _validate_full_checkpoint_weight_precision(path) + load_model(self.core_model, str(path / "model.safetensors"), strict=True, device=str(self.model.device)) + self.core_model.assert_fp32_parameters() + self.optimizer.load_state_dict(state["optimizer"]) + _assert_fp32_optimizer_state(self.optimizer) + self.lr_scheduler.load_state_dict(state["lr_scheduler"]) + self.ema.load(path / "ema/model.safetensors", self.core_model, num_updates=int(state["ema_num_updates"])) + global_step = int(state["global_step"]) + if self.lr_scheduler.last_epoch != global_step: + raise RuntimeError(f"Checkpoint scheduler step {self.lr_scheduler.last_epoch} does not match global_step={global_step}") + rng_states = state.get("rng_states") + if not isinstance(rng_states, list) or len(rng_states) != checkpoint_world_size: + raise RuntimeError(f"Checkpoint does not contain one RNG state per rank: {path}") + self._resume_rng_state = rng_states[get_rank()] + logger.info("[openpi:resume] full state restored from {} at step={}", path, global_step) + return global_step, int(state["data_epoch"]), int(state["batches_in_epoch"]) + + def _write_inference_artifact(self, destination: Path) -> None: + tokenizer = Path(self.dataloader.metadata["tokenizer_path"]) + norm_stats = Path(self.dataloader.metadata["norm_stats_path"]) + tokenizer_target = destination / "assets/paligemma_tokenizer.model" + norm_target = destination / "assets/physical-intelligence/libero/norm_stats.json" + tokenizer_target.parent.mkdir(parents=True, exist_ok=True) + norm_target.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(tokenizer, tokenizer_target) + shutil.copy2(norm_stats, norm_target) + artifact_config = { + **self.model.architecture_metadata(), + "output_action_dim": 7, + "state_dim": 8, + "num_inference_steps": 10, + "pytorch_compile_mode": None, + } + _write_json(destination / "config.json", artifact_config) + + def _check_checkpoint_disk_space(self) -> None: + model_state = self.core_model.state_dict() + model_bytes = _tensor_storage_bytes(model_state) + aliases = _shared_aliases(model_state) + ema_state = {name: value for name, value in self.ema.shadow.items() if name not in aliases} + ema_bytes = _tensor_storage_bytes(ema_state) + optimizer_bytes = _tensor_storage_bytes(self.optimizer.state_dict()) + payload_bytes = model_bytes + ema_bytes + optimizer_bytes + required_bytes = int(payload_bytes * 1.15) + 512 * 1024**2 + free_bytes = shutil.disk_usage(self.output_dir).free + logger.info( + "[openpi:checkpoint] disk preflight payload={:.2f} GiB required_with_margin={:.2f} GiB free={:.2f} GiB", + payload_bytes / 1024**3, + required_bytes / 1024**3, + free_bytes / 1024**3, + ) + if free_bytes < required_bytes: + raise OSError(f"Not enough free space for an atomic OpenPI checkpoint: need at least {required_bytes / 1024**3:.2f} GiB, have {free_bytes / 1024**3:.2f} GiB in {self.output_dir}") + + def _save_checkpoint(self, global_step: int, data_epoch: int, batches_in_epoch: int) -> None: + rng_states = _gather_rng_states() + save_error: Exception | None = None + save_error_message: str | None = None + if is_main_process(): + stage: Path | None = None + try: + self.output_dir.mkdir(parents=True, exist_ok=True) + self._check_checkpoint_disk_space() + name = f"checkpoint-{global_step:09d}" + final = self.output_dir / name + if final.exists(): + raise FileExistsError(f"Refusing to replace an existing checkpoint: {final}") + stage = self.output_dir / f".{name}.tmp-{uuid.uuid4().hex}" + stage.mkdir() + save_model(self.core_model, str(stage / "model.safetensors"), metadata={"format": "pt"}) + ema_dir = stage / "ema" + ema_dir.mkdir() + self.ema.save(ema_dir / "model.safetensors", self.core_model) + contract = self._checkpoint_contract() + training_state = { + "schema_version": CHECKPOINT_SCHEMA_VERSION, + "global_step": global_step, + "world_size": get_world_size(), + "data_epoch": data_epoch, + "batches_in_epoch": batches_in_epoch, + "optimizer": self.optimizer.state_dict(), + "lr_scheduler": self.lr_scheduler.state_dict(), + "ema_num_updates": self.ema.num_updates, + "rng_states": rng_states, + "contract": contract, + } + torch.save(training_state, stage / "training_state.pt") + shutil.copy2(self.config["config_path"], stage / "config.yaml") + self._write_inference_artifact(stage) + self._write_inference_artifact(ema_dir) + manifest = { + "schema_version": CHECKPOINT_SCHEMA_VERSION, + "created_by": CHECKPOINT_CREATOR, + "created_at": datetime.now(timezone.utc).isoformat(), + "global_step": global_step, + "world_size": get_world_size(), + "initialization": { + "checkpoint_dir": str(self.model.initialization_path), + }, + "contains": ["model", "optimizer", "lr_scheduler", "global_step", "ema", "rng_per_rank"], + "contract": contract, + } + _write_json(stage / "manifest.json", manifest) + _write_json(stage / "_SUCCESS", {"global_step": global_step, "created_at": manifest["created_at"]}) + os.rename(stage, final) + logger.info("[openpi:checkpoint] saved complete checkpoint {}", final) + _prune_checkpoints( + self.output_dir, + None if self.save_total_limit is None else int(self.save_total_limit), + self.keep_period, + ) + except Exception as error: + save_error = error + save_error_message = f"{type(error).__name__}: {error}" + if stage is not None and stage.exists(): + try: + shutil.rmtree(stage) + except Exception as cleanup_error: + logger.error("[openpi:checkpoint] failed to remove staging directory {}: {}", stage, cleanup_error) + save_error_message += f"; staging cleanup also failed: {type(cleanup_error).__name__}: {cleanup_error}" + + if is_distributed(): + result = [save_error_message] + dist.broadcast_object_list(result, src=0) + if result[0] is not None: + raise RuntimeError(f"Rank 0 failed to save checkpoint: {result[0]}") from save_error + elif save_error is not None: + raise save_error + + def _build_data_iterator(self, epoch: int, batches_in_epoch: int): + self.dataloader.set_epoch(epoch) + iterator = iter(self.dataloader) + for _ in range(batches_in_epoch): + try: + next(iterator) + except StopIteration as error: + raise RuntimeError(f"Cannot restore data position epoch={epoch}, batch={batches_in_epoch}; dataset length changed") from error + return iterator + + def _next_batch(self, iterator, epoch: int, batches_in_epoch: int): + try: + batch = next(iterator) + except StopIteration: + epoch += 1 + batches_in_epoch = 0 + self.dataloader.set_epoch(epoch) + iterator = iter(self.dataloader) + batch = next(iterator) + return batch, iterator, epoch, batches_in_epoch + 1 + + def train(self) -> None: + _set_seed(int(self.training_config["seed"])) + torch.set_float32_matmul_precision("highest") + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + if self.training_config["gradient_checkpointing"]: + self.model.enable_gradient_checkpointing() + training_model = self._wrap_ddp() + training_model.train() + + self.optimizer = torch.optim.AdamW( + self.core_model.parameters(), + lr=self.optimizer_spec["learning_rate"], + betas=(self.optimizer_spec["adam_beta1"], self.optimizer_spec["adam_beta2"]), + weight_decay=self.optimizer_spec["weight_decay"], + eps=self.optimizer_spec["adam_epsilon"], + foreach=False, + ) + if not math.isclose(self.optimizer_spec["learning_rate"], float(self.schedule_spec["peak_lr"])): + raise ValueError("optimizer.learning_rate must equal lr_schedule.peak_lr") + self.lr_scheduler = _build_scheduler(self.optimizer, self.schedule_spec) + self.ema = ExponentialMovingAverage(self.core_model, self.ema_decay) + + resume_path = self._resolve_resume_path() + if resume_path is None: + global_step, data_epoch, batches_in_epoch = 0, 0, 0 + logger.info("[openpi:init] model-only warm start; optimizer, scheduler, step, and EMA are new") + else: + global_step, data_epoch, batches_in_epoch = self._load_full_checkpoint(resume_path) + if global_step > self.max_train_iters: + raise ValueError(f"Checkpoint step {global_step} exceeds max_train_iters={self.max_train_iters}") + + self.output_dir.mkdir(parents=True, exist_ok=True) + iterator = self._build_data_iterator(data_epoch, batches_in_epoch) + if resume_path is not None: + _restore_rng_state(self._resume_rng_state) + self.optimizer.zero_grad(set_to_none=True) + accumulated_loss: torch.Tensor | None = None + numerical_state_validated = resume_path is not None + last_saved_step = global_step if resume_path is not None else -1 + logger.info( + "[openpi:train] start step={}/{} world_size={} grad_accum={} global_batch={} ema_decay={}", + global_step, + self.max_train_iters, + get_world_size(), + self.gradient_accumulation_iters, + self.dataloader.metadata["global_batch_size"], + self.ema_decay, + ) + + try: + while global_step < self.max_train_iters: + for accumulation_index in range(self.gradient_accumulation_iters): + batch, iterator, data_epoch, batches_in_epoch = self._next_batch(iterator, data_epoch, batches_in_epoch) + observation, actions = batch + observation = observation.to(self.model.device, non_blocking=True) + actions = actions.to(self.model.device, dtype=torch.float32, non_blocking=True) + sync_gradients = accumulation_index + 1 == self.gradient_accumulation_iters + sync_context = nullcontext() if sync_gradients or not isinstance(training_model, DistributedDataParallel) else training_model.no_sync() + with sync_context: + loss = training_model(observation, actions).mean() + if loss.dtype != torch.float32: + raise RuntimeError(f"OpenPI flow-matching loss must be float32, got {loss.dtype}") + (loss / self.gradient_accumulation_iters).backward() + detached_loss = loss.detach() / self.gradient_accumulation_iters + accumulated_loss = detached_loss if accumulated_loss is None else accumulated_loss + detached_loss + + if not numerical_state_validated: + _assert_fp32_gradients(self.core_model) + grad_norm = torch.nn.utils.clip_grad_norm_(self.core_model.parameters(), self.max_grad_norm) + self.optimizer.step() + if not numerical_state_validated: + _assert_fp32_optimizer_state(self.optimizer) + numerical_state_validated = True + logger.info( + "[openpi:numerics] verified fp32 parameters, gradients, AdamW state, loss, and EMA masters; transformer compute={}", + self._checkpoint_contract()["numerical_policy"]["compute_dtype"], + ) + self.lr_scheduler.step() + self.optimizer.zero_grad(set_to_none=True) + self.ema.update(self.core_model) + global_step += 1 + + if global_step == 1 or global_step % self.log_every_iters == 0 or global_step == self.max_train_iters: + loss_value = reduce_mean(accumulated_loss) + grad_value = reduce_mean(grad_norm.detach()) + if torch.is_tensor(loss_value): + loss_value = loss_value.item() + if torch.is_tensor(grad_value): + grad_value = grad_value.item() + learning_rate = self.lr_scheduler.get_last_lr()[0] + logger.info( + "[openpi:train] step={}/{} loss={:.6f} grad_norm={:.6f} lr={:.8f}", + global_step, + self.max_train_iters, + loss_value, + grad_value, + learning_rate, + ) + self.monitor.log_metrics( + { + "train/loss": loss_value, + "train/grad_norm": grad_value, + "train/lr": learning_rate, + }, + step=global_step, + ) + accumulated_loss = None + + if self.save_every_iters and global_step % self.save_every_iters == 0: + self._save_checkpoint(global_step, data_epoch, batches_in_epoch) + last_saved_step = global_step + + if global_step != last_saved_step: + self._save_checkpoint(global_step, data_epoch, batches_in_epoch) + logger.info("[openpi:train] finished at step={}", global_step) + finally: + self.monitor.finish() diff --git a/lightx2v_train/lightx2v_train/utils/registry.py b/lightx2v_train/lightx2v_train/utils/registry.py index 2ef7acbe8..b456dfd26 100644 --- a/lightx2v_train/lightx2v_train/utils/registry.py +++ b/lightx2v_train/lightx2v_train/utils/registry.py @@ -81,6 +81,7 @@ def merge(self, other_register): "longcat_image": "lightx2v_train.model_zoo.longcat_image.longcat_image", "longcat_image_edit": "lightx2v_train.model_zoo.longcat_image.longcat_image_edit", "minimax_h3_t2av": "lightx2v_train.model_zoo.minimax_h3.minimax_h3_t2av", + "openpi_pi05_libero": "lightx2v_train.model_zoo.openpi.pi05_libero", "qwen_image": "lightx2v_train.model_zoo.qwen_image.qwen_image", "qwen_image_edit": "lightx2v_train.model_zoo.qwen_image.qwen_image_edit", "wan_t2v": "lightx2v_train.model_zoo.wan.wan_t2v", @@ -94,6 +95,7 @@ def merge(self, other_register): "consistency": "lightx2v_train.trainers.consistency.trainer", "dmd": "lightx2v_train.trainers.dmd.trainer", "flow_matching": "lightx2v_train.trainers.flow_matching", + "openpi_flow_matching": "lightx2v_train.trainers.openpi", "phased_dmd": "lightx2v_train.trainers.phased_dmd.trainer", "sgmd": "lightx2v_train.trainers.sgmd", "teacher_forcing": "lightx2v_train.trainers.teacher_forcing", @@ -140,6 +142,8 @@ def _ensure_data_registered(data_name): import lightx2v_train.data.training_cache_dataset # noqa: F401 elif data_name in {"prompt_dataset", "video_dataset"}: import lightx2v_train.data.video_dataset # noqa: F401 + elif data_name == "openpi_libero": + import lightx2v_train.data.openpi_libero # noqa: F401 def build_model(config): diff --git a/lightx2v_train/scripts/openpi/run_pi05_finetune_ema.sh b/lightx2v_train/scripts/openpi/run_pi05_finetune_ema.sh new file mode 100755 index 000000000..2570b6a90 --- /dev/null +++ b/lightx2v_train/scripts/openpi/run_pi05_finetune_ema.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ $# -ne 0 ]]; then + echo "Usage: bash lightx2v_train/scripts/openpi/run_pi05_finetune_ema.sh" >&2 + exit 2 +fi + +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +unset OPENPI_RESUME_CHECKPOINT +exec bash "${script_dir}/support/launch_pi05_libero.sh" diff --git a/lightx2v_train/scripts/openpi/run_pi05_resume_ema.sh b/lightx2v_train/scripts/openpi/run_pi05_resume_ema.sh new file mode 100755 index 000000000..3ab19728b --- /dev/null +++ b/lightx2v_train/scripts/openpi/run_pi05_resume_ema.sh @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ $# -ne 1 ]]; then + echo "Usage: bash lightx2v_train/scripts/openpi/run_pi05_resume_ema.sh CHECKPOINT" >&2 + exit 2 +fi + +checkpoint="$1" +if [[ ! -d "${checkpoint}" ]]; then + echo "Checkpoint directory not found: ${checkpoint}" >&2 + exit 1 +fi +checkpoint="$(cd -- "${checkpoint}" && pwd -P)" + +required_files=( + model.safetensors + ema/model.safetensors + training_state.pt + manifest.json + _SUCCESS +) +for path in "${required_files[@]}"; do + if [[ ! -f "${checkpoint}/${path}" ]]; then + echo "Not a complete LightX2V training checkpoint; missing ${checkpoint}/${path}" >&2 + exit 1 + fi +done + +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +export OPENPI_RESUME_CHECKPOINT="${checkpoint}" +export OPENPI_TRAIN_OUTPUT="${OPENPI_TRAIN_OUTPUT:-$(dirname -- "${checkpoint}")}" +exec bash "${script_dir}/support/launch_pi05_libero.sh" diff --git a/lightx2v_train/scripts/openpi/support/README.md b/lightx2v_train/scripts/openpi/support/README.md new file mode 100644 index 000000000..5c4f5e180 --- /dev/null +++ b/lightx2v_train/scripts/openpi/support/README.md @@ -0,0 +1,470 @@ +# OpenPI π0.5-LIBERO 训练复现 + +本文说明如何在 LightX2V 中从官方 π0.5 base 权重开始 LIBERO fine-tuning，如何 +恢复被中断的训练，以及如何用 EMA 权重完成定量评测。推理、LIBERO rollout 和 ROS +的完整说明见项目级 [OpenPI README](../../../../scripts/openpi/support/README.md)。 + +下面的命令默认在项目根目录执行： + +```bash +cd /data/liuhongda/lightx2v_openpi +conda activate base +``` + +训练直接使用当前环境的 `python`。开始前应确认： + +```bash +command -v python +``` + +## 1. 两个训练入口 + +训练脚本只分为两种互斥语义： + +| 脚本 | 输入权重 | 恢复 optimizer/step | 用途 | +| --- | --- | --- | --- | +| `run_pi05_finetune_ema.sh` | 官方 π0.5 base PyTorch 权重 | 否 | 开始一轮新的 LIBERO fine-tuning | +| `run_pi05_resume_ema.sh` | LightX2V 完整训练 checkpoint | 是 | 从中断位置严格续训 | + +这里的“新 fine-tuning”不是随机初始化。模型参数来自官方 π0.5 base；optimizer、 +scheduler 和 step 从零开始，EMA 从初始模型复制。当前接入不提供随机初始化训练。 + +调用链如下： + +```text +run_pi05_finetune_ema.sh / run_pi05_resume_ema.sh + -> support/launch_pi05_libero.sh + -> python -m torch.distributed.run + -> lightx2v_train/train.py + -> OpenPILiberoDataset + -> OpenPIPi05LiberoModel + -> OpenPIFlowMatchingTrainer +``` + +两个公开入口共用同一个 helper、同一份 YAML 和同一个 trainer，只有初始化/恢复方式 +不同。`support/launch_pi05_libero.sh` 是内部实现，不需要直接运行。 + +## 2. 默认文件组织 + +默认训练输入位于 `/data/liuhongda/openpi_data`： + +```text +openpi_data/ +├── openpi-assets/checkpoints/ +│ ├── pi05_base_pytorch_fp32/ +│ │ ├── model.safetensors +│ │ ├── config.json +│ │ └── assets/paligemma_tokenizer.model +│ └── pi05_libero/ +│ └── assets/physical-intelligence/libero/norm_stats.json +├── lerobot/physical-intelligence/libero/ +│ ├── data/chunk-*/episode_*.parquet +│ └── meta/ +│ ├── info.json +│ ├── episodes.jsonl +│ └── tasks.jsonl +└── python_deps/openpi_official_pytorch_runtime/ +``` + +各项用途如下： + +- `pi05_base_pytorch_fp32`：新 fine-tuning 的 model-only 初始权重，必须是 FP32。 +- `physical-intelligence/libero`：官方 LIBERO-40 LeRobot 训练集。 +- `norm_stats.json`：官方 LIBERO state/action q01、q99 归一化统计。 +- `openpi_official_pytorch_runtime`：包含 OpenPI replacement 的 Transformers 4.53.2 + 私有 overlay。 + +训练不读取 `pi05_libero_pytorch_fp32` specialist 权重，也不需要 MuJoCo、LIBERO +仿真或 ROS。仿真只在训练完成后的 rollout 评测阶段使用。 + +默认数据契约为： + +| 项目 | 值 | +| --- | ---: | +| episodes | 1693 | +| frames | 273465 | +| language tasks | 40 | +| FPS | 10 | +| state dimension | 8 | +| action dimension | 7（模型内部补到 32） | +| action horizon | 10 | + +### 2.1 下载 LeRobot 数据 + +如果本地还没有数据集，可以从 Hugging Face 下载： + +```bash +hf download physical-intelligence/libero \ + --repo-type dataset \ + --local-dir /data/liuhongda/openpi_data/lerobot/physical-intelligence/libero +``` + +不要将 HDF5 LIBERO demonstrations 直接填到 `OPENPI_LEROBOT_ROOT`。当前 loader +读取的是上述 LeRobot v2.0 parquet 结构。 + +### 2.2 准备 LIBERO norm stats + +默认从官方 `pi05_libero` checkpoint 读取已经用于该 recipe 的 quantile stats： + +```text +/data/liuhongda/openpi_data/openpi-assets/checkpoints/pi05_libero/ +└── assets/physical-intelligence/libero/norm_stats.json +``` + +如果缺少该文件，可使用 OpenPI 下载官方 checkpoint： + +```bash +cd /data/liuhongda/openpi +export OPENPI_DATA_HOME=/data/liuhongda/openpi_data + +uv run --no-sync python -c \ + 'from openpi.shared import download; print(download.maybe_download("gs://openpi-assets/checkpoints/pi05_libero"))' +``` + +如果其他 checkpoint 目录中已有同一份 stats，也可以通过 +`OPENPI_NORM_STATS_PATH=/path/to/norm_stats.json` 指定。不要用 LeRobot +`meta/stats.json` 中的全局 mean/std 代替这里的 state/action q01、q99。 + +### 2.3 准备 FP32 base 权重 + +如果已经有以下目录，可以跳过本节： + +```text +/data/liuhongda/openpi_data/openpi-assets/checkpoints/pi05_base_pytorch_fp32 +``` + +从 JAX 权重重新准备时，先用 OpenPI 下载官方 base checkpoint 和 tokenizer： + +```bash +cd /data/liuhongda/openpi +export OPENPI_DATA_HOME=/data/liuhongda/openpi_data + +uv run --no-sync python -c \ + 'from openpi.shared import download; print(download.maybe_download("gs://openpi-assets/checkpoints/pi05_base"))' + +uv run --no-sync python -c \ + 'from openpi.shared import download; print(download.maybe_download("gs://big_vision/paligemma_tokenizer.model", gs={"token": "anon"}))' +``` + +先准备 Transformers overlay，再调用 OpenPI 的转换器输出 FP32 参数： + +```bash +cd /data/liuhongda/lightx2v_openpi +bash scripts/openpi/1_setup_pytorch_runtime.sh prepare --component transformers + +cd /data/liuhongda/openpi +PYTHONPATH=/data/liuhongda/openpi_data/python_deps/openpi_official_pytorch_runtime:/data/liuhongda/openpi/src \ +.venv/bin/python examples/convert_jax_model_to_pytorch.py \ + --checkpoint-dir /data/liuhongda/openpi_data/openpi-assets/checkpoints/pi05_base \ + --config-name pi05_libero \ + --output-path /data/liuhongda/openpi_data/openpi-assets/checkpoints/pi05_base_pytorch_fp32 \ + --precision float32 + +mkdir -p /data/liuhongda/openpi_data/openpi-assets/checkpoints/pi05_base_pytorch_fp32/assets +cp -a /data/liuhongda/openpi_data/openpi-assets/checkpoints/pi05_base/assets/. \ + /data/liuhongda/openpi_data/openpi-assets/checkpoints/pi05_base_pytorch_fp32/assets/ +cp /data/liuhongda/openpi_data/big_vision/paligemma_tokenizer.model \ + /data/liuhongda/openpi_data/openpi-assets/checkpoints/pi05_base_pytorch_fp32/assets/paligemma_tokenizer.model +``` + +输出目录应事先不存在或为空。训练前检查会验证 `config.json` 中的精度、812 个 +safetensors tensor 的 dtype，以及 tokenizer 是否完整。 + +`scripts/openpi/2_convert_pi05_libero_to_pytorch.sh` 默认转换的是 fine-tuned +`pi05_libero` specialist，主要供推理复现使用，不要把它的默认输出误当成训练 base。 + +## 3. 准备和检查运行环境 + +训练只需要准备 Transformers overlay，不会修改 Python、PyTorch、CUDA 或 MuJoCo： + +```bash +cd /data/liuhongda/lightx2v_openpi +bash scripts/openpi/1_setup_pytorch_runtime.sh prepare --component transformers +``` + +可以单独执行训练前检查： + +```bash +python scripts/openpi/support/runtime.py train-check +``` + +没有可见 GPU、只想检查文件时使用： + +```bash +python scripts/openpi/support/runtime.py train-check --no-cuda +``` + +两个训练入口在启动 DDP 前都会自动执行 `train-check`，因此正式启动时不需要重复 +手工检查。检查范围包括： + +- Transformers 4.53.2 及 OpenPI replacement 文件； +- FP32 base checkpoint、config 和 tokenizer； +- LIBERO-40 数据规模及 parquet 文件； +- state/action q01、q99 norm stats； +- PyTorch、LeRobot、Pillow、SentencePiece、Augmax 和 CUDA。 + +## 4. 开始新的 EMA fine-tuning + +默认使用 4、5、6、7 号卡： + +```bash +cd /data/liuhongda/lightx2v_openpi +bash lightx2v_train/scripts/openpi/run_pi05_finetune_ema.sh +``` + +使用 0、1、2、3 号卡： + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 \ +bash lightx2v_train/scripts/openpi/run_pi05_finetune_ema.sh +``` + +默认输出到： + +```text +/data/liuhongda/lightx2v_openpi/output_train/openpi/pi05_libero +``` + +启动新的 fine-tuning 时，输出目录不能包含已有 `checkpoint-*`。如果目录里已有 +训练 checkpoint，应使用 resume 脚本或指定新的输出目录： + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 \ +OPENPI_TRAIN_OUTPUT=/data/liuhongda/lightx2v_openpi/output_train/openpi/pi05_libero_run2 \ +bash lightx2v_train/scripts/openpi/run_pi05_finetune_ema.sh +``` + +## 5. 官方对齐的默认训练参数 + +配置文件为 +[`pi05_libero.yaml`](../../../configs/train/openpi/pi05_libero.yaml)。默认值如下： + +| 参数 | 值 | +| --- | ---: | +| global batch size | 256 | +| optimizer updates | 30000 | +| warmup steps | 10000 | +| peak/end learning rate | 5e-5 / 5e-5 | +| AdamW betas | 0.9 / 0.95 | +| AdamW epsilon | 1e-8 | +| weight decay | 1e-10 | +| global gradient clip | 1.0 | +| EMA decay | 0.999 | +| seed | 42 | + +参数、梯度、AdamW state、loss 和 EMA master 保持 FP32；Gemma/SigLIP 的主要矩阵 +计算使用 BF16。不开 GradScaler，TF32 关闭。默认四卡、无梯度累计时： + +```text +per-GPU batch 64 × 4 GPUs × accumulation 1 = global batch 256 +``` + +若显存不足，可在保持 global batch 256 的情况下增加梯度累计。例如四卡累计 8 次： + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 \ +OPENPI_GRADIENT_ACCUMULATION_ITERS=8 \ +bash lightx2v_train/scripts/openpi/run_pi05_finetune_ema.sh +``` + +对应每卡 micro batch 为 8。梯度累计会改变浮点求和顺序；追求与官方设置尽量一致 +时优先使用 accumulation 1。 + +前 10000 step 的 learning rate 会从 0 线性增加到 `5e-5`，之后保持不变。因此训练 +早期看到 LR 持续变大是正常 warmup，不是异常发散。 + +## 6. 监控训练 + +默认日志路径： + +```bash +tail -f /data/liuhongda/lightx2v_openpi/output_train/openpi/pi05_libero/train.log +``` + +正常日志包括： + +```text +[openpi:numerics] verified fp32 parameters, gradients, AdamW state, loss, and EMA masters +[openpi:train] step=... loss=... grad_norm=... lr=... +[openpi:checkpoint] saved complete checkpoint ... +``` + +判断训练是否正常时重点看： + +- `loss`、`grad_norm` 和 `lr` 都是有限值； +- step 持续增加； +- step 1 完成 FP32 数值链检查； +- 每 1000 step 能成功生成完整 checkpoint； +- 不同 batch 的 loss 允许上下波动，不要求单调下降。 + +LeRobot 可能提示数据仍是 v2.0 global stats 格式。当前 loader 对该格式兼容，而且 +它正是本接入验证过的数据格式，不要在正式复现过程中临时转换数据版本。 + +## 7. Checkpoint 保存逻辑 + +默认每 1000 step 保存一次，5000 的整数倍长期保留，其他 checkpoint 只保留最新 +一个。完整 checkpoint 结构如下： + +```text +checkpoint-000030000/ +├── model.safetensors # online 模型参数 +├── ema/ +│ ├── model.safetensors # EMA 参数，用于推理和评测 +│ ├── config.json +│ └── assets/ +├── training_state.pt # optimizer、scheduler、step、RNG、数据位置 +├── manifest.json # 训练契约和 provenance +├── _SUCCESS # 完整写入标记 +├── config.yaml # 本次训练配置快照 +├── config.json # 推理模型配置 +└── assets/ # tokenizer 和 norm stats +``` + +当前 3.6B FP32 完整 checkpoint 约为 52 GB；一次写入前会要求约 60 GB 可用空间。 +使用默认保留策略跑满 30k 后，整个输出目录约为 359 GB。checkpoint 使用 staging +目录原子写入，保存期间训练会等待磁盘 I/O。 + +用途必须区分： + +- 继续训练：传 `checkpoint-XXXXXXXXX` 根目录。 +- 推理和成功率评测：使用 `checkpoint-XXXXXXXXX/ema`。 +- 不要把 `ema` 子目录传给 resume 脚本，它没有 optimizer 等训练状态。 + +## 8. 恢复训练 + +从 5000 step 的完整 checkpoint 恢复并继续到默认目标 30000： + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 \ +bash lightx2v_train/scripts/openpi/run_pi05_resume_ema.sh \ + /data/liuhongda/lightx2v_openpi/output_train/openpi/pi05_libero/checkpoint-000005000 +``` + +从 30000 step 继续到 35000： + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 \ +OPENPI_MAX_TRAIN_ITERS=35000 \ +bash lightx2v_train/scripts/openpi/run_pi05_resume_ema.sh \ + /data/liuhongda/lightx2v_openpi/output_train/openpi/pi05_libero/checkpoint-000030000 +``` + +`OPENPI_MAX_TRAIN_ITERS` 表示目标总步数，不是额外训练步数。如果从 30000 恢复且 +仍使用默认值 30000，程序会成功恢复后立即结束，不会更新参数。 + +严格恢复要求以下内容与保存时一致： + +- DDP world size； +- global batch 和 gradient accumulation； +- 模型结构与数值精度； +- optimizer、scheduler、EMA 参数； +- 数据集、tokenizer、norm stats、shuffle 和 seed 契约。 + +可以更换物理 GPU 编号，但 GPU 数量必须一致。建议始终从输出目录中最新的完整 +checkpoint 恢复；如果要从旧 checkpoint 创建分支训练，应同时指定新的 +`OPENPI_TRAIN_OUTPUT`，避免后续 step 与已有 checkpoint 重名。 + +## 9. 评测训练结果 + +本 recipe 没有离线 validation split。训练过程中观察 loss、gradient norm 和 LR； +最终模型质量由 LIBERO rollout 成功率衡量。 + +使用 30k EMA 权重在 4 张卡上评测 4 个 suite、每任务 50 trials，共 2000 episodes： + +```bash +cd /data/liuhongda/lightx2v_openpi + +CUDA_VISIBLE_DEVICES=4,5,6,7 \ +OPENPI_MODEL_PATH=/data/liuhongda/lightx2v_openpi/output_train/openpi/pi05_libero/checkpoint-000030000/ema \ +OPENPI_PARALLEL_OUTPUT_ROOT=/data/liuhongda/lightx2v_openpi/save_results/pi05_libero_trained_ema \ +bash scripts/openpi/run_libero_evaluate_parallel_i2va.sh +``` + +结果汇总位于： + +```text +save_results/pi05_libero_trained_ema/parallel_summary.json +``` + +官方公开结果为： + +| Suite | Success rate | +| --- | ---: | +| LIBERO-Spatial | 98.8% | +| LIBERO-Object | 98.2% | +| LIBERO-Goal | 98.0% | +| LIBERO-10 | 92.4% | +| Average | 96.85% | + +## 10. 常用环境变量 + +日常通常只需要修改 GPU、输出路径或恢复目标步数： + +| 环境变量 | 作用 | 默认值 | +| --- | --- | --- | +| `CUDA_VISIBLE_DEVICES` | 参与 DDP 的 GPU 列表 | `4,5,6,7` | +| `OPENPI_TRAIN_OUTPUT` | 日志和 checkpoint 根目录 | `output_train/openpi/pi05_libero` | +| `OPENPI_MAX_TRAIN_ITERS` | 目标 optimizer 总步数 | `30000` | +| `OPENPI_GRADIENT_ACCUMULATION_ITERS` | 梯度累计次数 | `1` | +| `OPENPI_DATA_WORKERS` | 每个 rank 的 loader workers | `2` | + +更换数据或资源路径时才需要： + +| 环境变量 | 作用 | 默认值 | +| --- | --- | --- | +| `OPENPI_DATA_ROOT` | OpenPI 数据根目录 | `/data/liuhongda/openpi_data` | +| `OPENPI_INITIAL_CHECKPOINT` | FP32 π0.5 base 权重 | `.../pi05_base_pytorch_fp32` | +| `OPENPI_LEROBOT_ROOT` | LIBERO-40 LeRobot 数据 | `.../lerobot/physical-intelligence/libero` | +| `OPENPI_NORM_STATS_PATH` | LIBERO quantile stats | `.../pi05_libero/assets/.../norm_stats.json` | +| `OPENPI_TRANSFORMERS_RUNTIME_PATH` | Transformers overlay | `.../openpi_official_pytorch_runtime` | +| `OPENPI_GLOBAL_BATCH_SIZE` | optimizer global batch | `256` | +| `OPENPI_NPROC_PER_NODE` | DDP 进程数 | 从可见 GPU 数量推导 | +| `OPENPI_TRAIN_SEED` | 训练、shuffle 和采样 seed | `42` | +| `OPENPI_SAVE_EVERY_ITERS` | checkpoint 保存间隔 | `1000` | +| `OPENPI_KEEP_PERIOD` | 永久 checkpoint 周期 | `5000` | +| `OPENPI_SAVE_TOTAL_LIMIT` | 非周期 checkpoint 保留数 | `1` | +| `OPENPI_LOG_EVERY_ITERS` | 训练日志间隔 | `10` | + +一般不需要手工设置 `OPENPI_NPROC_PER_NODE`。如果设置了，它应与实际可见 GPU 数量 +一致；global batch 必须能被 `NPROC × gradient accumulation` 整除。 + +## 11. 常见问题 + +### 为什么启动后 learning rate 一直增加？ + +前 10000 step 是官方设置的线性 warmup，LR 会逐步升到 `5e-5`，之后保持不变。 + +### 为什么保存 checkpoint 很慢？ + +完整 checkpoint 包含 online FP32、EMA FP32 和 AdamW state，单个约 52 GB。保存时 +需要把 staging 目录完整写入磁盘后再原子发布。 + +### 为什么官方 base 不能传给 resume 脚本？ + +官方 base 只有模型参数，没有 `training_state.pt`、EMA、optimizer、scheduler、 +逐 rank RNG 和数据位置。它只能传给新 fine-tuning 入口。 + +### 为什么 `/ema` 不能续训？ + +`ema` 是面向推理的模型 artifact，不是完整训练 checkpoint。续训必须传它的父目录。 + +### 为什么换成两张卡后 resume 失败？ + +完整 checkpoint 保存了一份每 rank RNG 状态，并记录了 `world_size=4`。严格复现时 +必须继续使用四个 DDP rank；GPU 编号可以改变。 + +### 如何确认 checkpoint 完整？ + +至少应存在： + +```text +model.safetensors +ema/model.safetensors +training_state.pt +manifest.json +_SUCCESS +``` + +resume 脚本和 trainer 都会再次验证这些文件及训练契约。 + +算法、数据处理、数值精度和验收结果都记录在本文中。 diff --git a/lightx2v_train/scripts/openpi/support/launch_pi05_libero.sh b/lightx2v_train/scripts/openpi/support/launch_pi05_libero.sh new file mode 100755 index 000000000..5b4625eb3 --- /dev/null +++ b/lightx2v_train/scripts/openpi/support/launch_pi05_libero.sh @@ -0,0 +1,54 @@ +#!/usr/bin/env bash +set -euo pipefail + +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +train_root="$(cd -- "${script_dir}/../../.." && pwd)" +lightx2v_root="$(dirname -- "${train_root}")" +workspace_root="$(dirname -- "${lightx2v_root}")" +openpi_data_root="${OPENPI_DATA_ROOT:-${workspace_root}/openpi_data}" +runtime="${OPENPI_TRANSFORMERS_RUNTIME_PATH:-${openpi_data_root}/python_deps/openpi_official_pytorch_runtime}" +runtime_tool="${lightx2v_root}/scripts/openpi/support/runtime.py" + +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-4,5,6,7}" +export OPENPI_INITIAL_CHECKPOINT="${OPENPI_INITIAL_CHECKPOINT:-${openpi_data_root}/openpi-assets/checkpoints/pi05_base_pytorch_fp32}" +export OPENPI_LEROBOT_ROOT="${OPENPI_LEROBOT_ROOT:-${openpi_data_root}/lerobot/physical-intelligence/libero}" +export OPENPI_NORM_STATS_PATH="${OPENPI_NORM_STATS_PATH:-${openpi_data_root}/openpi-assets/checkpoints/pi05_libero/assets/physical-intelligence/libero/norm_stats.json}" +export OPENPI_TRAIN_OUTPUT="${OPENPI_TRAIN_OUTPUT:-${lightx2v_root}/output_train/openpi/pi05_libero}" +export OPENPI_TRANSFORMERS_RUNTIME_PATH="${runtime}" +export OPENPI_GLOBAL_BATCH_SIZE="${OPENPI_GLOBAL_BATCH_SIZE:-256}" +export OPENPI_GRADIENT_ACCUMULATION_ITERS="${OPENPI_GRADIENT_ACCUMULATION_ITERS:-1}" +export HF_HOME="${HF_HOME:-${openpi_data_root}/hf_cache}" +export USE_FLAX=0 +export PYTHONDONTWRITEBYTECODE=1 +export PYTHONNOUSERSITE=1 +export TOKENIZERS_PARALLELISM=false +export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" + +nproc="${OPENPI_NPROC_PER_NODE:-$(awk -F, '{print NF}' <<<"${CUDA_VISIBLE_DEVICES}")}" +if [[ ! "${nproc}" =~ ^[1-9][0-9]*$ || ! "${OPENPI_GLOBAL_BATCH_SIZE}" =~ ^[1-9][0-9]*$ || ! "${OPENPI_GRADIENT_ACCUMULATION_ITERS}" =~ ^[1-9][0-9]*$ ]]; then + echo "GPU count, global batch size, and gradient accumulation must be positive integers." >&2 + exit 2 +fi + +batch_divisor=$((nproc * OPENPI_GRADIENT_ACCUMULATION_ITERS)) +if ((OPENPI_GLOBAL_BATCH_SIZE % batch_divisor != 0)); then + echo "OPENPI_GLOBAL_BATCH_SIZE must be divisible by GPUs * gradient accumulation." >&2 + exit 2 +fi +per_gpu_batch=$((OPENPI_GLOBAL_BATCH_SIZE / batch_divisor)) + +python "${runtime_tool}" train-check +echo "OpenPI train: GPUs=${nproc}, per-GPU batch=${per_gpu_batch}, accumulation=${OPENPI_GRADIENT_ACCUMULATION_ITERS}, global batch=${OPENPI_GLOBAL_BATCH_SIZE}" +if [[ -n "${OPENPI_RESUME_CHECKPOINT:-}" ]]; then + echo "Initialization: full resume from ${OPENPI_RESUME_CHECKPOINT}" +else + echo "Initialization: model-only warm start from ${OPENPI_INITIAL_CHECKPOINT}" +fi + +export PYTHONPATH="${runtime}:${lightx2v_root}:${train_root}${PYTHONPATH:+:${PYTHONPATH}}" +cd "${train_root}" +exec python -m torch.distributed.run \ + --standalone \ + --nproc_per_node="${nproc}" \ + train.py \ + --config configs/train/openpi/pi05_libero.yaml diff --git a/scripts/openpi/1_setup_pytorch_runtime.sh b/scripts/openpi/1_setup_pytorch_runtime.sh index b7a41f3a0..3b17d10ea 100755 --- a/scripts/openpi/1_setup_pytorch_runtime.sh +++ b/scripts/openpi/1_setup_pytorch_runtime.sh @@ -4,10 +4,10 @@ set -euo pipefail script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" command="prepare" -if [[ "${1:-}" == "check" || "${1:-}" == "--check" ]]; then +if [[ "${1:-}" == "check" ]]; then command="check" shift -elif [[ "${1:-}" == "setup" || "${1:-}" == "prepare" ]]; then +elif [[ "${1:-}" == "prepare" ]]; then shift fi diff --git a/scripts/openpi/run_libero_evaluate_i2va.sh b/scripts/openpi/run_libero_evaluate_i2va.sh index 2306e579b..103fd9a29 100755 --- a/scripts/openpi/run_libero_evaluate_i2va.sh +++ b/scripts/openpi/run_libero_evaluate_i2va.sh @@ -34,5 +34,4 @@ exec python -m lightx2v.infer \ --task i2va \ --model_path "${model_path}" \ --config_json "${config_json}" \ - --seed "${OPENPI_POLICY_SEED:-0}" \ --save_result_path "${output_dir}" diff --git a/scripts/openpi/run_libero_evaluate_parallel_i2va.sh b/scripts/openpi/run_libero_evaluate_parallel_i2va.sh index 5643b7103..9b55707aa 100755 --- a/scripts/openpi/run_libero_evaluate_parallel_i2va.sh +++ b/scripts/openpi/run_libero_evaluate_parallel_i2va.sh @@ -49,18 +49,14 @@ run_worker() { suite_output="${output_root}/${suite}" suite_runtime="${output_root}/runtime/${suite}" log_path="${output_root}/logs/${suite}.log" - mkdir -p "${suite_runtime}/libero_config" "${suite_runtime}/numba" "${suite_runtime}/matplotlib" "${suite_runtime}/cache" printf '\n[%(%Y-%m-%dT%H:%M:%SZ)T] suite=%s gpu=%s\n' -1 "${suite}" "${gpu}" >> "${log_path}" ( export CUDA_VISIBLE_DEVICES="${gpu}" export OPENPI_EVAL_OUTPUT_DIR="${suite_output}" export OPENPI_RUNTIME_DIR="${suite_runtime}" - export OPENPI_LIBERO_CONFIG_DIR="${suite_runtime}/libero_config" export OPENPI_EVAL_BENCHMARKS="${suite}" - export NUMBA_CACHE_DIR="${suite_runtime}/numba" - export MPLCONFIGDIR="${suite_runtime}/matplotlib" - export XDG_CACHE_HOME="${suite_runtime}/cache" + unset OPENPI_LIBERO_CONFIG_DIR NUMBA_CACHE_DIR MPLCONFIGDIR XDG_CACHE_HOME exec setsid bash "${script_dir}/run_libero_evaluate_i2va.sh" ) >> "${log_path}" 2>&1 & child_pid=$! diff --git a/scripts/openpi/run_libero_ros_i2va.sh b/scripts/openpi/run_libero_ros_i2va.sh index 71c1eaa85..f05bf2cef 100755 --- a/scripts/openpi/run_libero_ros_i2va.sh +++ b/scripts/openpi/run_libero_ros_i2va.sh @@ -68,8 +68,5 @@ export TOKENIZERS_PARALLELISM=false export PYTHONPATH="${OPENPI_TRANSFORMERS_RUNTIME_PATH:-${openpi_data_root}/python_deps/openpi_official_pytorch_runtime}:${lightx2v_path}:${PYTHONPATH:-}" exec ros2 run inference openpi_node --ros-args \ - -p numeric_precision:=float64 \ -p "model_path:=${OPENPI_MODEL_PATH:-${openpi_data_root}/openpi-assets/checkpoints/pi05_libero_pytorch_fp32}" \ - -p "config_json:=${OPENPI_CONFIG:-${lightx2v_path}/configs/openpi/pi05_libero.json}" \ - -p seed:=0 \ - -p actions_per_plan:=5 + -p "config_json:=${OPENPI_CONFIG:-${lightx2v_path}/configs/openpi/pi05_libero.json}" diff --git a/scripts/openpi/support/README.md b/scripts/openpi/support/README.md index bbe0ced6d..d43135d98 100644 --- a/scripts/openpi/support/README.md +++ b/scripts/openpi/support/README.md @@ -1,7 +1,7 @@ # OpenPI π0.5-LIBERO -该目录提供 π0.5-LIBERO 的权重转换、运行环境准备、本地 rollout、定量评测和 -ROS 交互。批量评测从 LightX2V 公共入口启动： +该目录提供 π0.5-LIBERO 的权重转换、运行环境准备、fine-tuning、本地 rollout、 +定量评测和 ROS 交互。批量评测从 LightX2V 公共入口启动： ```text shell -> python -m lightx2v.infer -> OpenPIRunner @@ -126,7 +126,7 @@ setup 只安装或修复 OpenPI 所需的小包：base 环境中的 `mujoco==3.2 Python、PyTorch 或 CUDA。 ```bash -bash scripts/openpi/1_setup_pytorch_runtime.sh +bash scripts/openpi/1_setup_pytorch_runtime.sh prepare --component transformers ``` 训练或评测前可做只读检查： @@ -504,3 +504,8 @@ RNG 和 5-action replan queue；清理启动脚本时不应改变这些逻辑。 修改 ROS 文件后需要重新构建 `common simulator inference`；LIBERO adapter 会把 可见的物理 GPU 映射为 EGL 逻辑设备 0。 + +## 11. π0.5-LIBERO 训练 + +训练数据、权重、环境准备、fine-tuning、checkpoint resume、数值对齐和评测方法见 +[训练复现文档](../../../lightx2v_train/scripts/openpi/support/README.md)。 diff --git a/scripts/openpi/support/runtime.py b/scripts/openpi/support/runtime.py index dc0c201f7..d00abacab 100755 --- a/scripts/openpi/support/runtime.py +++ b/scripts/openpi/support/runtime.py @@ -29,6 +29,10 @@ DEFAULT_LIBERO_ROOT = WORKSPACE_ROOT / "openpi/third_party/libero" DEFAULT_MODEL_CONFIG = PROJECT_ROOT / "configs/openpi/pi05_libero.json" DEFAULT_EVAL_CONFIG = PROJECT_ROOT / "configs/openpi/pi05_libero_eval.json" +DEFAULT_BASE_MODEL = OPENPI_DATA_ROOT / "openpi-assets/checkpoints/pi05_base_pytorch_fp32" +DEFAULT_TRAIN_DATASET = OPENPI_DATA_ROOT / "lerobot/physical-intelligence/libero" +DEFAULT_TRAIN_NORM_STATS = OPENPI_DATA_ROOT / "openpi-assets/checkpoints/pi05_libero/assets/physical-intelligence/libero/norm_stats.json" +DEFAULT_TRAIN_CONFIG = PROJECT_ROOT / "lightx2v_train/configs/train/openpi/pi05_libero.yaml" TRANSFORMERS_EXPECTED = { "transformers": ("transformers", "4.53.2"), @@ -49,6 +53,7 @@ "models/siglip/check.py", "models/siglip/modeling_siglip.py", ) +EXPECTED_TENSORS = 812 def _resolved(path: str | Path) -> Path: @@ -271,6 +276,20 @@ def _load_json(path: Path, label: str) -> dict: return value +def _checkpoint_dtypes(path: Path) -> dict[str, int]: + try: + from safetensors import safe_open + except ImportError as exc: + raise RuntimeError("base environment is missing safetensors") from exc + + dtypes: dict[str, int] = {} + with safe_open(path, framework="pt", device="cpu") as checkpoint: + for name in checkpoint.keys(): + dtype = checkpoint.get_slice(name).get_dtype() + dtypes[dtype] = dtypes.get(dtype, 0) + 1 + return dtypes + + def _check_static_inputs(args: argparse.Namespace) -> None: model = _resolved(args.model_path) required_model_files = ( @@ -287,19 +306,11 @@ def _check_static_inputs(args: argparse.Namespace) -> None: expected_dtype = {"float32": "F32", "bfloat16": "BF16"}.get(precision) if expected_dtype is None: raise RuntimeError(f"checkpoint precision must be float32 or bfloat16, got {precision!r}") - try: - from safetensors import safe_open - except ImportError as exc: - raise RuntimeError("base environment is missing safetensors") from exc - tensor_dtypes: dict[str, str] = {} - with safe_open(model / "model.safetensors", framework="pt", device="cpu") as checkpoint: - for name in checkpoint.keys(): - dtype = checkpoint.get_slice(name).get_dtype() - tensor_dtypes[dtype] = tensor_dtypes.get(dtype, 0) + 1 - expected_tensors = {expected_dtype: 812} + tensor_dtypes = _checkpoint_dtypes(model / "model.safetensors") + expected_tensors = {expected_dtype: EXPECTED_TENSORS} if tensor_dtypes != expected_tensors: - raise RuntimeError(f"expected 812 {precision} checkpoint tensors, got {tensor_dtypes}") - print(f"checkpoint tensor manifest: 812/812 {expected_dtype}") + raise RuntimeError(f"expected {EXPECTED_TENSORS} {precision} checkpoint tensors, got {tensor_dtypes}") + print(f"checkpoint tensor manifest: {EXPECTED_TENSORS}/{EXPECTED_TENSORS} {expected_dtype}") _load_json(_resolved(args.model_config), "model config") _load_json(_resolved(args.eval_config), "evaluation config") @@ -315,6 +326,134 @@ def _check_static_inputs(args: argparse.Namespace) -> None: raise RuntimeError("incomplete official LIBERO checkout:\n- " + "\n- ".join(missing)) +def _check_training_inputs(args: argparse.Namespace) -> None: + checkpoint = _resolved(args.initial_checkpoint) + required_checkpoint_files = ( + checkpoint / "model.safetensors", + checkpoint / "config.json", + checkpoint / "assets/paligemma_tokenizer.model", + ) + missing = [str(path) for path in required_checkpoint_files if not path.is_file()] + if missing: + raise RuntimeError("incomplete π0.5 base checkpoint:\n- " + "\n- ".join(missing)) + + checkpoint_config = _load_json(checkpoint / "config.json", "base checkpoint config") + if checkpoint_config.get("precision") != "float32": + raise RuntimeError("OpenPI training requires a lossless float32 base checkpoint") + tensor_dtypes = _checkpoint_dtypes(checkpoint / "model.safetensors") + expected_tensors = {"F32": EXPECTED_TENSORS} + if tensor_dtypes != expected_tensors: + raise RuntimeError(f"expected {EXPECTED_TENSORS} float32 base checkpoint tensors, got {tensor_dtypes}") + print(f"base checkpoint tensor manifest: {EXPECTED_TENSORS}/{EXPECTED_TENSORS} F32") + + dataset = _resolved(args.dataset_root) + info = _load_json(dataset / "meta/info.json", "LeRobot metadata") + expected = { + "codebase_version": "v2.0", + "total_episodes": 1693, + "total_frames": 273465, + "total_tasks": 40, + "fps": 10, + } + mismatches = {key: (info.get(key), value) for key, value in expected.items() if info.get(key) != value} + if mismatches: + raise RuntimeError(f"dataset is not the official LIBERO-40 training set: {mismatches}") + required_dataset_files = (dataset / "meta/episodes.jsonl", dataset / "meta/tasks.jsonl") + missing = [str(path) for path in required_dataset_files if not path.is_file()] + if missing or not any((dataset / "data").glob("chunk-*/episode_*.parquet")): + raise RuntimeError("LeRobot LIBERO data is incomplete: " + ", ".join(missing or [str(dataset / "data")])) + print("LeRobot dataset manifest: 1693 episodes, 273465 frames, 40 tasks, 10 FPS") + + norm_stats_path = _resolved(args.norm_stats_path) + norm_stats = _load_json(norm_stats_path, "LIBERO normalization statistics").get("norm_stats", {}) + for key, size in (("state", 8), ("actions", 7)): + values = norm_stats.get(key, {}) + for statistic in ("q01", "q99"): + if len(values.get(statistic, ())) != size: + raise RuntimeError(f"{norm_stats_path}: {key}.{statistic} must contain {size} values") + print(f"LIBERO quantile statistics: {norm_stats_path}") + + train_config = _resolved(args.train_config) + if not train_config.is_file(): + raise RuntimeError(f"training config is missing: {train_config}") + + +TRAIN_PROBE = r""" +import importlib +import importlib.metadata +import json +import sys +from pathlib import Path + +runtime = Path(sys.argv[1]).resolve() +modules = {} +for distribution, module_name in ( + ("torch", "torch"), + ("numpy", "numpy"), + ("Pillow", "PIL"), + ("safetensors", "safetensors"), + ("sentencepiece", "sentencepiece"), + ("lerobot", "lerobot"), + ("augmax", "augmax"), + ("omegaconf", "omegaconf"), +): + module = importlib.import_module(module_name) + origin = Path(module.__file__).resolve() + if origin.is_relative_to(runtime): + raise RuntimeError(f"base package {module_name} was shadowed by the Transformers overlay: {origin}") + modules[distribution] = {"version": importlib.metadata.version(distribution), "origin": str(origin)} + +from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata +import torch +import transformers + +modules["transformers"] = { + "version": transformers.__version__, + "origin": str(Path(transformers.__file__).resolve()), +} +modules["cuda_available"] = torch.cuda.is_available() +print(json.dumps(modules, sort_keys=True)) +if sys.argv[2] == "1" and not modules["cuda_available"]: + raise RuntimeError("CUDA is not available to the training interpreter") +""" + + +def _check_training_runtime(args: argparse.Namespace) -> None: + expected_python = _resolved(args.expected_python) + if _resolved(sys.executable) != expected_python: + raise RuntimeError(f"training check must use {expected_python}, got {_resolved(sys.executable)}") + + transformers_runtime = _resolved(args.transformers_runtime) + _probe_overlay(transformers_runtime, TRANSFORMERS_EXPECTED) + _check_patch_overlay(transformers_runtime) + _check_training_inputs(args) + + env = os.environ.copy() + env.update( + { + "PYTHONPATH": os.pathsep.join((str(transformers_runtime), str(PROJECT_ROOT), str(PROJECT_ROOT / "lightx2v_train"))), + "PYTHONDONTWRITEBYTECODE": "1", + "PYTHONNOUSERSITE": "1", + "USE_FLAX": "0", + "TOKENIZERS_PARALLELISM": "false", + } + ) + completed = subprocess.run( + [sys.executable, "-c", TRAIN_PROBE, str(transformers_runtime), "0" if args.no_cuda else "1"], + env=env, + text=True, + capture_output=True, + ) + if completed.returncode != 0: + if completed.stdout: + print(completed.stdout, end="", file=sys.stderr) + if completed.stderr: + print(completed.stderr, end="", file=sys.stderr) + raise RuntimeError(f"OpenPI training runtime probe failed with exit code {completed.returncode}") + print(completed.stdout.strip()) + print("OpenPI training runtime check: OK") + + COMBINED_PROBE = r""" import importlib.metadata import json @@ -453,6 +592,25 @@ def _add_paths(parser: argparse.ArgumentParser) -> None: ) +def _add_training_paths(parser: argparse.ArgumentParser) -> None: + _add_paths(parser) + parser.add_argument("--expected-python", default=os.environ.get("OPENPI_TRAIN_PYTHON", sys.executable)) + parser.add_argument( + "--initial-checkpoint", + default=os.environ.get("OPENPI_INITIAL_CHECKPOINT", str(DEFAULT_BASE_MODEL)), + ) + parser.add_argument( + "--dataset-root", + default=os.environ.get("OPENPI_LEROBOT_ROOT", str(DEFAULT_TRAIN_DATASET)), + ) + parser.add_argument( + "--norm-stats-path", + default=os.environ.get("OPENPI_NORM_STATS_PATH", str(DEFAULT_TRAIN_NORM_STATS)), + ) + parser.add_argument("--train-config", default=os.environ.get("OPENPI_TRAIN_CONFIG", str(DEFAULT_TRAIN_CONFIG))) + parser.add_argument("--no-cuda", action="store_true", help="allow validation without a visible CUDA device") + + def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description=__doc__) subparsers = parser.add_subparsers(dest="command", required=True) @@ -470,6 +628,12 @@ def build_parser() -> argparse.ArgumentParser: check.add_argument("--eval-config", default=os.environ.get("OPENPI_EVAL_CONFIG", str(DEFAULT_EVAL_CONFIG))) check.add_argument("--libero-root", default=os.environ.get("OPENPI_LIBERO_ROOT", str(DEFAULT_LIBERO_ROOT))) check.add_argument("--no-cuda", action="store_true", help="allow validation on a host without a visible CUDA device") + + train_check = subparsers.add_parser( + "train-check", + help="validate the OpenPI training runtime, FP32 base checkpoint, and LIBERO-40 data", + ) + _add_training_paths(train_check) return parser @@ -481,8 +645,10 @@ def main() -> int: _prepare_transformers(_resolved(args.transformers_runtime), args.dry_run) if args.component in {"all", "mujoco"}: _prepare_base_mujoco(args.dry_run) - else: + elif args.command == "check": _check_runtime(args) + else: + _check_training_runtime(args) except (OSError, RuntimeError, subprocess.CalledProcessError) as exc: print(f"error: {exc}", file=sys.stderr) return 1