feat(openpi): add aligned LIBERO training support

.gitignore

@@ -28,6 +28,7 @@ server_cache/
 app/.gradio/
 *.pkl
 save_results/*
+/output_train/
 *.egg-info/
 lightx2v_train/output_train/*
 lightx2v_train/output_infer/*

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
+}

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
+  }
+}

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",

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"]

lightx2v/models/networks/openpi/config.py

@@ -0,0 +1,90 @@
+"""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
+    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
+    discrete_state_input: bool = False
+    pytorch_compile_mode: str | None = None
+
+    @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"],
+            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.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,
+            "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.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'")