import abc from collections.abc import Sequence import dataclasses import enum import logging import pathlib from typing import Generic, TypeVar import augmax from flax import nnx from flax import struct from flax import traverse_util import jax import jax.numpy as jnp import numpy as np import orbax.checkpoint as ocp import safetensors import torch from openpi_value.shared import image_tools import openpi_value.shared.array_typing as at from typing import Callable, Optional, Union import warnings logger = logging.getLogger("openpi") # Type variable for array types (JAX arrays, PyTorch tensors, or numpy arrays) ArrayT = TypeVar("ArrayT", bound=jax.Array | torch.Tensor | np.ndarray) class ModelType(enum.Enum): """Supported model types.""" PI0 = "pi0" PI0_FAST = "pi0_fast" PI05 = "pi05" # This may need change if we release a small model. IMAGE_RESOLUTION = (224, 224) # Data format # # Data transforms produce the model input as a nested dictionary which is later converted # into `Obesrvation` and `Actions` objects. See below. # # In the dictory form, this data should look like: # { # # Observation data. # "image": { # "base_0_rgb": (float32|uint8)[*b, h, w, 3], # RGB image in [-1, 1] or [0, 255] # ... # Additional camera views # }, # "image_mask": { # "base_0_rgb": bool[*b], # True if image is valid # ... # Masks for additional views # }, # "state": float32[*b, s], # Low-dimensional robot state # "tokenized_prompt": int32[*b, l], # Optional, tokenized language prompt # "tokenized_prompt_mask": bool[*b, l], # Optional, mask for tokenized prompt # "token_ar_mask": int32[*b, l], # Optional, autoregressive mask for FAST model # "token_loss_mask": bool[*b, l], # Optional, loss mask for FAST model # # # Actions data. # "actions": float32[*b ah ad] # } # where: # *b = batch dimensions # h,w = image height/width # s = state dimension # l = sequence length # @at.typecheck @struct.dataclass class Observation(Generic[ArrayT]): """Holds observations, i.e., inputs to the model. See `Observation.from_dict` to see the expected dictionary form. This is the format that should be produced by the data transforms. """ # Images, in [-1, 1] float32. # images: dict[str, at.Float[ArrayT, "*b h w c"]] images: dict[str, at.Float[ArrayT, "*b c h w"]] # Image masks, with same keys as images. image_masks: dict[str, at.Bool[ArrayT, "*b"]] # Low-dimensional robot state. state: at.Float[ArrayT, "*b s"] # Tokenized prompt. tokenized_prompt: at.Int[ArrayT, "*b l"] | None = None # Tokenized prompt mask. tokenized_prompt_mask: at.Bool[ArrayT, "*b l"] | None = None # * string # prompt: str | None = None # Optional, raw text prompt. Not used by the model. # pi0-fast model specific fields. # Token auto-regressive mask (for FAST autoregressive model). token_ar_mask: at.Int[ArrayT, "*b l"] | None = None # Token loss mask (for FAST autoregressive model). token_loss_mask: at.Bool[ArrayT, "*b l"] | None = None # * Custom episode_index: at.Int[ArrayT, "*b"] | None = None # frame_index: at.Int[ArrayT, "*b"] | None = None # * original frame_index frame_index: Union[at.Int[ArrayT, "*b"], at.Float[ArrayT, "*b"]] | None = None # * modified frame_index for progress estimation frame_index_progress: Union[at.Int[ArrayT, "*b"], at.Float[ArrayT, "*b"]] | None = None is_failure_data: at.Bool[ArrayT, "*b"] | None = None is_infer_data: at.Bool[ArrayT, "*b"] | None = None episode_length: Union[at.Int[ArrayT, "*b"], at.Float[ArrayT, "*b"]] | None = None action_advantage: Union[at.Int[ArrayT, "*b"], at.Float[ArrayT, "*b"]] | None = None action_advantage_original: Union[at.Int[ArrayT, "*b"], at.Float[ArrayT, "*b"]] | None = None image_original: dict[str, at.Float[ArrayT, "*b H W C"]] | None = None inferred_action: Union[at.Int[ArrayT, "*b d"], at.Float[ArrayT, "*b d"]] | None = None # noise: None | at.Float[ArrayT, "*b s"] = None noise: None | at.Float[ArrayT, "*b 1 50 d"] = None @classmethod def from_dict(cls, data: at.PyTree[ArrayT]) -> "Observation[ArrayT]": """This method defines the mapping between unstructured data (i.e., nested dict) to the structured Observation format.""" # Ensure that tokenized_prompt and tokenized_prompt_mask are provided together. if ("tokenized_prompt" in data) != ("tokenized_prompt_mask" in data): raise ValueError("tokenized_prompt and tokenized_prompt_mask must be provided together.") # If images are uint8, convert them to [-1, 1] float32. if "image" in data: img_key = "image" else: raise ValueError("No image key found in observation.") for key in data[img_key]: if data[img_key][key].dtype == np.uint8: data[img_key][key] = data[img_key][key].astype(np.float32) / 255.0 * 2.0 - 1.0 elif hasattr(data[img_key][key], "dtype") and data[img_key][key].dtype == torch.uint8: data[img_key][key] = data[img_key][key].to(torch.float32).permute(0, 3, 1, 2) / 255.0 * 2.0 - 1.0 if data.get("image_original", None) is not None: for key in data["image_original"]: if data["image_original"][key].dtype == np.uint8: data["image_original"][key] = data["image_original"][key].astype(np.float32) / 255.0 * 2.0 - 1.0 elif hasattr(data["image_original"][key], "dtype") and data["image_original"][key].dtype == torch.uint8: data["image_original"][key] = data["image_original"][key].to(torch.float32).permute(0, 3, 1, 2) / 255.0 * 2.0 - 1.0 # Handle noise reshaping: if noise is flattened [batch, 700], reshape to [batch, 1, 50, 14] if data.get("noise", None) is not None: # Check if noise is flattened (2D with shape [batch, 700]) if len(data["noise"].shape) == 2 and data["noise"].shape[1] == 700: # Reshape from [batch, 700] to [batch, 1, 50, 14] # 700 = 1 * 50 * 14 data["noise"] = data["noise"].reshape(data["noise"].shape[0], 1, 50, 14) elif len(data["noise"].shape) == 2 and data["noise"].shape[1] == 50: # * [batch, 50] # Reshape from [batch, 50] to [batch, 1, 50, 1] - might need action_dim # For now, assume action_dim=14 if shape[1] == 50*14 data["noise"] = data["noise"].unsqueeze(1) return cls( images=data["image"], image_masks=data["image_mask"], state=data["state"], tokenized_prompt=data.get("tokenized_prompt"), tokenized_prompt_mask=data.get("tokenized_prompt_mask"), token_ar_mask=data.get("token_ar_mask"), token_loss_mask=data.get("token_loss_mask"), frame_index=data.get("frame_index"), frame_index_progress=data.get("frame_index_progress"), is_failure_data=data.get("is_failure_data"), is_infer_data=data.get("is_infer_data"), episode_length=data.get("episode_length"), action_advantage=data.get("action_advantage"), action_advantage_original=data.get("action_advantage_original", None), image_original=data.get("image_original", None), episode_index=data.get("episode_index", None), inferred_action=data.get("inferred_action", None), noise=data.get("noise", None), ) def to_dict(self) -> at.PyTree[ArrayT]: """Convert the Observation to a nested dict.""" result = dataclasses.asdict(self) result["image"] = result.pop("images") result["image_mask"] = result.pop("image_masks") return result def drop_images(self, image_keys): for key in image_keys: self.images.pop(key, None) self.image_masks.pop(key, None) # Defines the format of the actions. This field is included as "actions" inside the dictionary # produced by the data transforms. Actions = at.Float[ArrayT, "*b ah ad"] def preprocess_observation( rng: at.KeyArrayLike | None, observation: Observation, *, train: bool = False, image_keys: Sequence[str] = [], # ! Not used anymore image_resolution: tuple[int, int] = IMAGE_RESOLUTION, ) -> Observation: """Preprocess the observations by performing image augmentations (if train=True), resizing (if necessary), and filling in a default image mask (if necessary). """ batch_shape = observation.state.shape[:-1] out_images = {} for key in observation.images: image = observation.images[key] 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(image, *image_resolution) if train: # Convert from [-1, 1] to [0, 1] for augmax. image = image / 2.0 + 0.5 transforms = [] if "wrist" not in key: height, width = image.shape[1:3] transforms += [ augmax.RandomCrop(int(width * 0.95), int(height * 0.95)), augmax.Resize(width, height), augmax.Rotate((-5, 5)), ] transforms += [ augmax.ColorJitter(brightness=0.3, contrast=0.4, saturation=0.5), ] sub_rngs = jax.random.split(rng, image.shape[0]) image = jax.vmap(augmax.Chain(*transforms))(sub_rngs, image) # Back to [-1, 1]. image = image * 2.0 - 1.0 out_images[key] = image # obtain mask out_masks = {} for key in out_images: if key not in observation.image_masks: # do not mask by default out_masks[key] = jnp.ones(batch_shape, dtype=jnp.bool) else: out_masks[key] = jnp.asarray(observation.image_masks[key]) return Observation( 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, frame_index=observation.frame_index, frame_index_progress=observation.frame_index_progress, is_failure_data=observation.is_failure_data, is_infer_data=observation.is_infer_data, episode_length=observation.episode_length, action_advantage=observation.action_advantage, action_advantage_original=observation.action_advantage_original, image_original=observation.image_original, episode_index=observation.episode_index, inferred_action=observation.inferred_action, noise=observation.noise, ) @dataclasses.dataclass(frozen=True) class BaseModelConfig(abc.ABC): """Configuration shared by all models. Specific models should inherit from this class, and implement the `create` method to create the corresponding model. """ # Action space dimension. action_dim: int # Action sequence length. action_horizon: int # Tokenized prompt maximum length. max_token_len: int @property @abc.abstractmethod def model_type(self) -> ModelType: """The model type.""" @abc.abstractmethod def create(self, rng: at.KeyArrayLike) -> "BaseModel": """Create a new model, initializing parameters.""" def load(self, params: at.Params, *, remove_extra_params: bool = True) -> "BaseModel": """Create a model with the given parameters.""" model = nnx.eval_shape(self.create, jax.random.key(0)) graphdef, state = nnx.split(model) def convert_keys_to_int(d): if isinstance(d, dict): return {int(k) if isinstance(k, str) and k.isdigit() else k: convert_keys_to_int(v) for k, v in d.items()} elif isinstance(d, (list, tuple)): return type(d)(convert_keys_to_int(x) for x in d) return d if remove_extra_params: params = ocp.transform_utils.intersect_trees(state.to_pure_dict(), params) params = convert_keys_to_int(params) at.check_pytree_equality(expected=state.to_pure_dict(), got=params, check_shapes=True, check_dtypes=False) state.replace_by_pure_dict(params) return nnx.merge(graphdef, state) def load_pytorch(self, train_config, weight_path: str): from openpi_value.models_pytorch import pi0_pytorch logger.info(f"train_config: {train_config}") model = pi0_pytorch.PI0Pytorch(config=train_config.model) if str(weight_path).endswith(".pt") or str(weight_path).endswith(".pth"): ckpt = torch.load(weight_path, map_location="cpu") model.load_state_dict(ckpt, strict=False) elif str(weight_path).endswith(".safetensors"): safetensors.torch.load_model(model, weight_path) return model @abc.abstractmethod def inputs_spec(self, *, batch_size: int = 1) -> tuple[Observation, Actions]: """Returns the input specification for the model. Values are jax.ShapeDtypeStruct.""" def fake_obs(self, batch_size: int = 1) -> Observation: observation_spec, _ = self.inputs_spec(batch_size=batch_size) return jax.tree.map(lambda x: jnp.ones(x.shape, x.dtype), observation_spec) def fake_act(self, batch_size: int = 1) -> Actions: _, action_spec = self.inputs_spec(batch_size=batch_size) return jax.tree.map(lambda x: jnp.ones(x.shape, x.dtype), action_spec) @dataclasses.dataclass class BaseModel(nnx.Module, abc.ABC): """Base class for all model implementations. Specific models should inherit from this class. They should call super().__init__() to initialize the shared attributes (action_dim, action_horizon, and max_token_len). """ action_dim: int action_horizon: int max_token_len: int @abc.abstractmethod def compute_loss( self, rng: at.KeyArrayLike, observation: Observation, actions: Actions, *, train: bool = False, ) -> at.Float[at.Array, "*b ah"]: ... @abc.abstractmethod def sample_actions(self, rng: at.KeyArrayLike, observation: Observation, **kwargs) -> Actions: ... def restore_params( params_path: pathlib.Path | str, *, restore_type: type[np.ndarray] | type[jax.Array] = jax.Array, dtype: jnp.dtype | None = None, sharding: jax.sharding.Sharding | None = None, ) -> at.Params: """Restores unstructured params PyTree from a checkpoint. This works with checkpoints saved with `save_state` during openpi training (see `training/checkpoints.py`) as well as pre-trained checkpoints released for openpi. Args: params_path: The local path to the checkpoint directory. restore_type: The type to restore the params as. Can be set to `np.ndarray` to load the params as a numpy array. dtype: The dtype to restore all params as. If not provided, will use the original dtype from the checkpoint. sharding: The sharding to use for the params. If not provided, the params will be replicated across all devices. Returns: The restored params. """ params_path = pathlib.Path(params_path).resolve() if not str(params_path).startswith("gs://") else params_path if restore_type is jax.Array and sharding is None: mesh = jax.sharding.Mesh(jax.devices(), ("x",)) sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec()) with ocp.PyTreeCheckpointer() as ckptr: metadata = ckptr.metadata(params_path) item = {"params": metadata["params"]} params = ckptr.restore( params_path, ocp.args.PyTreeRestore( item=item, restore_args=jax.tree.map( lambda _: ocp.ArrayRestoreArgs(sharding=sharding, restore_type=restore_type, dtype=dtype), item ), ), )["params"] # If the params were saved with `save_state` during openpi training, every key path will end with "value", which is # added by `nnx.State`. We remove the "value" suffix here and always return what NNX calls a "pure dict". flat_params = traverse_util.flatten_dict(params) if all(kp[-1] == "value" for kp in flat_params): flat_params = {kp[:-1]: v for kp, v in flat_params.items()} return traverse_util.unflatten_dict(flat_params)