Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions configs/openpi/pi05_libero.json
Original file line number Diff line number Diff line change
@@ -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
}
26 changes: 26 additions & 0 deletions configs/openpi/pi05_libero_eval.json
Original file line number Diff line number Diff line change
@@ -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
}
}
2 changes: 2 additions & 0 deletions lightx2v/infer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -128,6 +129,7 @@ def main():
"seedvr2",
"swiftvr",
"neopp",
"openpi",
"motus",
"lingbot_world_fast",
"worldmirror",
Expand Down
11 changes: 11 additions & 0 deletions lightx2v/models/networks/openpi/NOTICE.md
Original file line number Diff line number Diff line change
@@ -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.
7 changes: 7 additions & 0 deletions lightx2v/models/networks/openpi/__init__.py
Original file line number Diff line number Diff line change
@@ -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"]
77 changes: 77 additions & 0 deletions lightx2v/models/networks/openpi/config.py
Original file line number Diff line number Diff line change
@@ -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}")
Loading
Loading