import abc from collections.abc import Sequence import dataclasses import enum import logging import pathlib from typing import Generic, TypeVar, TypeAlias 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 from openpi.shared import image_tools import openpi.shared.array_typing as at logger = logging.getLogger("openpi") ArrayT = TypeVar("ArrayT", at.Array, jax.ShapeDtypeStruct) ArrayOrStruct: TypeAlias = at.Array | jax.ShapeDtypeStruct class ModelType(enum.Enum): """Supported model types.""" PI0 = "pi0" PI0_FAST = "pi0_fast" # The model always expects these images IMAGE_KEYS = ( "base_0_rgb", "left_wrist_0_rgb", "right_wrist_0_rgb", ) # 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, in [-1, 1] float32. images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] # Image masks, with same keys as images. image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] # Low-dimensional robot state. state: at.Float[ArrayOrStruct, "*b s"] # Tokenized prompt. tokenized_prompt: at.Int[ArrayOrStruct, "*b l"] | None = None # Tokenized prompt mask. tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None # pi0-fast model specific fields. # Token auto-regressive mask (for FAST autoregressive model). token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None # Token loss mask (for FAST autoregressive model). token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = 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. for key in data["image"]: if data["image"][key].dtype == np.uint8: data["image"][key] = data["image"][key].astype(np.float32) / 255.0 * 2.0 - 1.0 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"), ) 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 @at.typecheck @struct.dataclass class ObservationPrefixPostfix(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, in [-1, 1] float32. images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] # Image masks, with same keys as images. image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] # Low-dimensional robot state. state: at.Float[ArrayOrStruct, "*b s"] # Tokenized prompt. tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_prefix"] | None = None tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_postfix"] | None = None # Tokenized prompt mask. tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None # pi0-fast model specific fields. # Token auto-regressive mask (for FAST autoregressive model). token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None # Token loss mask (for FAST autoregressive model). token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = 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_prefix, tokenized_prompt_postfix, and tokenized_prompt_mask are provided together. if ("tokenized_prompt_prefix" in data) != ("tokenized_prompt_postfix" in data) != ("tokenized_prompt_mask" in data): raise ValueError("tokenized_prompt_prefix, tokenized_prompt_postfix, and tokenized_prompt_mask must be provided together.") # If images are uint8, convert them to [-1, 1] float32. for key in data["image"]: if data["image"][key].dtype == np.uint8: data["image"][key] = data["image"][key].astype(np.float32) / 255.0 * 2.0 - 1.0 return cls( images=data["image"], image_masks=data["image_mask"], state=data["state"], tokenized_prompt_prefix=data.get("tokenized_prompt_prefix"), tokenized_prompt_postfix=data.get("tokenized_prompt_postfix"), tokenized_prompt_mask=data.get("tokenized_prompt_mask"), token_ar_mask=data.get("token_ar_mask"), token_loss_mask=data.get("token_loss_mask"), ) 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 @at.typecheck @struct.dataclass class RiclObservation(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. TODO: Make below brute forced (quick and dirty) class actually elegant. """ # Images, in [-1, 1] float32. query_images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] | None = None retrieved_0_images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] | None = None retrieved_1_images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] | None = None retrieved_2_images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] | None = None retrieved_3_images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] | None = None retrieved_4_images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] | None = None retrieved_5_images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] | None = None retrieved_6_images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] | None = None retrieved_7_images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] | None = None retrieved_8_images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] | None = None retrieved_9_images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] | None = None retrieved_10_images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] | None = None retrieved_11_images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] | None = None retrieved_12_images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] | None = None retrieved_13_images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] | None = None retrieved_14_images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] | None = None retrieved_15_images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] | None = None retrieved_16_images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] | None = None retrieved_17_images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] | None = None retrieved_18_images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] | None = None retrieved_19_images: dict[str, at.Float[ArrayOrStruct, "*b h w c"]] | None = None # Image masks, with same keys as images. query_image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] | None = None retrieved_0_image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] | None = None retrieved_1_image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] | None = None retrieved_2_image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] | None = None retrieved_3_image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] | None = None retrieved_4_image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] | None = None retrieved_5_image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] | None = None retrieved_6_image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] | None = None retrieved_7_image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] | None = None retrieved_8_image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] | None = None retrieved_9_image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] | None = None retrieved_10_image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] | None = None retrieved_11_image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] | None = None retrieved_12_image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] | None = None retrieved_13_image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] | None = None retrieved_14_image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] | None = None retrieved_15_image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] | None = None retrieved_16_image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] | None = None retrieved_17_image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] | None = None retrieved_18_image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] | None = None retrieved_19_image_masks: dict[str, at.Bool[ArrayOrStruct, "*b"]] | None = None # Distances between embeddings exp_lamda_distances: at.Float[ArrayOrStruct, "*b num_observations 1"] | None = None # Low-dimensional robot state. query_state: at.Float[ArrayOrStruct, "*b s"] | None = None retrieved_0_state: at.Float[ArrayOrStruct, "*b s"] | None = None retrieved_1_state: at.Float[ArrayOrStruct, "*b s"] | None = None retrieved_2_state: at.Float[ArrayOrStruct, "*b s"] | None = None retrieved_3_state: at.Float[ArrayOrStruct, "*b s"] | None = None retrieved_4_state: at.Float[ArrayOrStruct, "*b s"] | None = None retrieved_5_state: at.Float[ArrayOrStruct, "*b s"] | None = None retrieved_6_state: at.Float[ArrayOrStruct, "*b s"] | None = None retrieved_7_state: at.Float[ArrayOrStruct, "*b s"] | None = None retrieved_8_state: at.Float[ArrayOrStruct, "*b s"] | None = None retrieved_9_state: at.Float[ArrayOrStruct, "*b s"] | None = None retrieved_10_state: at.Float[ArrayOrStruct, "*b s"] | None = None retrieved_11_state: at.Float[ArrayOrStruct, "*b s"] | None = None retrieved_12_state: at.Float[ArrayOrStruct, "*b s"] | None = None retrieved_13_state: at.Float[ArrayOrStruct, "*b s"] | None = None retrieved_14_state: at.Float[ArrayOrStruct, "*b s"] | None = None retrieved_15_state: at.Float[ArrayOrStruct, "*b s"] | None = None retrieved_16_state: at.Float[ArrayOrStruct, "*b s"] | None = None retrieved_17_state: at.Float[ArrayOrStruct, "*b s"] | None = None retrieved_18_state: at.Float[ArrayOrStruct, "*b s"] | None = None retrieved_19_state: at.Float[ArrayOrStruct, "*b s"] | None = None # Tokenized prompt prefix query_tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_query_prefix"] | None = None retrieved_0_tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_prefix"] | None = None retrieved_1_tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_prefix"] | None = None retrieved_2_tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_prefix"] | None = None retrieved_3_tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_prefix"] | None = None retrieved_4_tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_prefix"] | None = None retrieved_5_tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_prefix"] | None = None retrieved_6_tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_prefix"] | None = None retrieved_7_tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_prefix"] | None = None retrieved_8_tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_prefix"] | None = None retrieved_9_tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_prefix"] | None = None retrieved_10_tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_prefix"] | None = None retrieved_11_tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_prefix"] | None = None retrieved_12_tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_prefix"] | None = None retrieved_13_tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_prefix"] | None = None retrieved_14_tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_prefix"] | None = None retrieved_15_tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_prefix"] | None = None retrieved_16_tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_prefix"] | None = None retrieved_17_tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_prefix"] | None = None retrieved_18_tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_prefix"] | None = None retrieved_19_tokenized_prompt_prefix: at.Int[ArrayOrStruct, "*b l_prefix"] | None = None # Tokenized prompt postfix query_tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_query_postfix"] | None = None retrieved_0_tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_postfix"] | None = None retrieved_1_tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_postfix"] | None = None retrieved_2_tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_postfix"] | None = None retrieved_3_tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_postfix"] | None = None retrieved_4_tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_postfix"] | None = None retrieved_5_tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_postfix"] | None = None retrieved_6_tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_postfix"] | None = None retrieved_7_tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_postfix"] | None = None retrieved_8_tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_postfix"] | None = None retrieved_9_tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_postfix"] | None = None retrieved_10_tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_postfix"] | None = None retrieved_11_tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_postfix"] | None = None retrieved_12_tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_postfix"] | None = None retrieved_13_tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_postfix"] | None = None retrieved_14_tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_postfix"] | None = None retrieved_15_tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_postfix"] | None = None retrieved_16_tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_postfix"] | None = None retrieved_17_tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_postfix"] | None = None retrieved_18_tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_postfix"] | None = None retrieved_19_tokenized_prompt_postfix: at.Int[ArrayOrStruct, "*b l_postfix"] | None = None # Tokenized prompt mask. query_tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l_query"] | None = None retrieved_0_tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_1_tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_2_tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_3_tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_4_tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_5_tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_6_tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_7_tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_8_tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_9_tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_10_tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_11_tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_12_tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_13_tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_14_tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_15_tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_16_tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_17_tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_18_tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_19_tokenized_prompt_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None # pi0-fast model specific fields. # Token auto-regressive mask (for FAST autoregressive model). query_token_ar_mask: at.Int[ArrayOrStruct, "*b l_query"] | None = None retrieved_0_token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None retrieved_1_token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None retrieved_2_token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None retrieved_3_token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None retrieved_4_token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None retrieved_5_token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None retrieved_6_token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None retrieved_7_token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None retrieved_8_token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None retrieved_9_token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None retrieved_10_token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None retrieved_11_token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None retrieved_12_token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None retrieved_13_token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None retrieved_14_token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None retrieved_15_token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None retrieved_16_token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None retrieved_17_token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None retrieved_18_token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None retrieved_19_token_ar_mask: at.Int[ArrayOrStruct, "*b l"] | None = None # Token loss mask (for FAST autoregressive model). query_token_loss_mask: at.Bool[ArrayOrStruct, "*b l_query"] | None = None retrieved_0_token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_1_token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_2_token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_3_token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_4_token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_5_token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_6_token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_7_token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_8_token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_9_token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_10_token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_11_token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_12_token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_13_token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_14_token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_15_token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_16_token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_17_token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_18_token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None retrieved_19_token_loss_mask: at.Bool[ArrayOrStruct, "*b l"] | None = None @classmethod def from_dict(cls, data: at.PyTree[ArrayT], num_retrieved_observations: int) -> "RiclObservation[ArrayT]": """This method defines the mapping between unstructured data (i.e., nested dict) to the structured Observation format.""" assert num_retrieved_observations <= 20, f"TODO: fix this brute force code" # Idenify all prefix all_prefix = [f"retrieved_{i}_" for i in range(num_retrieved_observations)] + ["query_"] # Ensure that tokenized_prompt_prefix, tokenized_prompt_postfix, and tokenized_prompt_mask are provided together. for prefix in all_prefix: if (f"{prefix}tokenized_prompt_prefix" in data) != (f"{prefix}tokenized_prompt_postfix" in data) != (f"{prefix}tokenized_prompt_mask" in data): raise ValueError(f"{prefix}tokenized_prompt_prefix, {prefix}tokenized_prompt_postfix, and {prefix}tokenized_prompt_mask must be provided together.") # If images are uint8, convert them to [-1, 1] float32. image_keys = list(data[f"query_image"].keys()) for prefix in all_prefix: for key in image_keys: if data[f"{prefix}image"][key].dtype == np.uint8: data[f"{prefix}image"][key] = data[f"{prefix}image"][key].astype(np.float32) / 255.0 * 2.0 - 1.0 return cls( query_images=data.get("query_image"), retrieved_0_images=data.get("retrieved_0_image"), retrieved_1_images=data.get("retrieved_1_image"), retrieved_2_images=data.get("retrieved_2_image"), retrieved_3_images=data.get("retrieved_3_image"), retrieved_4_images=data.get("retrieved_4_image"), retrieved_5_images=data.get("retrieved_5_image"), retrieved_6_images=data.get("retrieved_6_image"), retrieved_7_images=data.get("retrieved_7_image"), retrieved_8_images=data.get("retrieved_8_image"), retrieved_9_images=data.get("retrieved_9_image"), retrieved_10_images=data.get("retrieved_10_image"), retrieved_11_images=data.get("retrieved_11_image"), retrieved_12_images=data.get("retrieved_12_image"), retrieved_13_images=data.get("retrieved_13_image"), retrieved_14_images=data.get("retrieved_14_image"), retrieved_15_images=data.get("retrieved_15_image"), retrieved_16_images=data.get("retrieved_16_image"), retrieved_17_images=data.get("retrieved_17_image"), retrieved_18_images=data.get("retrieved_18_image"), retrieved_19_images=data.get("retrieved_19_image"), query_image_masks=data.get("query_image_mask"), retrieved_0_image_masks=data.get("retrieved_0_image_mask"), retrieved_1_image_masks=data.get("retrieved_1_image_mask"), retrieved_2_image_masks=data.get("retrieved_2_image_mask"), retrieved_3_image_masks=data.get("retrieved_3_image_mask"), retrieved_4_image_masks=data.get("retrieved_4_image_mask"), retrieved_5_image_masks=data.get("retrieved_5_image_mask"), retrieved_6_image_masks=data.get("retrieved_6_image_mask"), retrieved_7_image_masks=data.get("retrieved_7_image_mask"), retrieved_8_image_masks=data.get("retrieved_8_image_mask"), retrieved_9_image_masks=data.get("retrieved_9_image_mask"), retrieved_10_image_masks=data.get("retrieved_10_image_mask"), retrieved_11_image_masks=data.get("retrieved_11_image_mask"), retrieved_12_image_masks=data.get("retrieved_12_image_mask"), retrieved_13_image_masks=data.get("retrieved_13_image_mask"), retrieved_14_image_masks=data.get("retrieved_14_image_mask"), retrieved_15_image_masks=data.get("retrieved_15_image_mask"), retrieved_16_image_masks=data.get("retrieved_16_image_mask"), retrieved_17_image_masks=data.get("retrieved_17_image_mask"), retrieved_18_image_masks=data.get("retrieved_18_image_mask"), retrieved_19_image_masks=data.get("retrieved_19_image_mask"), exp_lamda_distances=data.get("exp_lamda_distances"), query_state=data.get("query_state"), retrieved_0_state=data.get("retrieved_0_state"), retrieved_1_state=data.get("retrieved_1_state"), retrieved_2_state=data.get("retrieved_2_state"), retrieved_3_state=data.get("retrieved_3_state"), retrieved_4_state=data.get("retrieved_4_state"), retrieved_5_state=data.get("retrieved_5_state"), retrieved_6_state=data.get("retrieved_6_state"), retrieved_7_state=data.get("retrieved_7_state"), retrieved_8_state=data.get("retrieved_8_state"), retrieved_9_state=data.get("retrieved_9_state"), retrieved_10_state=data.get("retrieved_10_state"), retrieved_11_state=data.get("retrieved_11_state"), retrieved_12_state=data.get("retrieved_12_state"), retrieved_13_state=data.get("retrieved_13_state"), retrieved_14_state=data.get("retrieved_14_state"), retrieved_15_state=data.get("retrieved_15_state"), retrieved_16_state=data.get("retrieved_16_state"), retrieved_17_state=data.get("retrieved_17_state"), retrieved_18_state=data.get("retrieved_18_state"), retrieved_19_state=data.get("retrieved_19_state"), query_tokenized_prompt_prefix=data.get("query_tokenized_prompt_prefix"), retrieved_0_tokenized_prompt_prefix=data.get("retrieved_0_tokenized_prompt_prefix"), retrieved_1_tokenized_prompt_prefix=data.get("retrieved_1_tokenized_prompt_prefix"), retrieved_2_tokenized_prompt_prefix=data.get("retrieved_2_tokenized_prompt_prefix"), retrieved_3_tokenized_prompt_prefix=data.get("retrieved_3_tokenized_prompt_prefix"), retrieved_4_tokenized_prompt_prefix=data.get("retrieved_4_tokenized_prompt_prefix"), retrieved_5_tokenized_prompt_prefix=data.get("retrieved_5_tokenized_prompt_prefix"), retrieved_6_tokenized_prompt_prefix=data.get("retrieved_6_tokenized_prompt_prefix"), retrieved_7_tokenized_prompt_prefix=data.get("retrieved_7_tokenized_prompt_prefix"), retrieved_8_tokenized_prompt_prefix=data.get("retrieved_8_tokenized_prompt_prefix"), retrieved_9_tokenized_prompt_prefix=data.get("retrieved_9_tokenized_prompt_prefix"), retrieved_10_tokenized_prompt_prefix=data.get("retrieved_10_tokenized_prompt_prefix"), retrieved_11_tokenized_prompt_prefix=data.get("retrieved_11_tokenized_prompt_prefix"), retrieved_12_tokenized_prompt_prefix=data.get("retrieved_12_tokenized_prompt_prefix"), retrieved_13_tokenized_prompt_prefix=data.get("retrieved_13_tokenized_prompt_prefix"), retrieved_14_tokenized_prompt_prefix=data.get("retrieved_14_tokenized_prompt_prefix"), retrieved_15_tokenized_prompt_prefix=data.get("retrieved_15_tokenized_prompt_prefix"), retrieved_16_tokenized_prompt_prefix=data.get("retrieved_16_tokenized_prompt_prefix"), retrieved_17_tokenized_prompt_prefix=data.get("retrieved_17_tokenized_prompt_prefix"), retrieved_18_tokenized_prompt_prefix=data.get("retrieved_18_tokenized_prompt_prefix"), retrieved_19_tokenized_prompt_prefix=data.get("retrieved_19_tokenized_prompt_prefix"), query_tokenized_prompt_postfix=data.get("query_tokenized_prompt_postfix"), retrieved_0_tokenized_prompt_postfix=data.get("retrieved_0_tokenized_prompt_postfix"), retrieved_1_tokenized_prompt_postfix=data.get("retrieved_1_tokenized_prompt_postfix"), retrieved_2_tokenized_prompt_postfix=data.get("retrieved_2_tokenized_prompt_postfix"), retrieved_3_tokenized_prompt_postfix=data.get("retrieved_3_tokenized_prompt_postfix"), retrieved_4_tokenized_prompt_postfix=data.get("retrieved_4_tokenized_prompt_postfix"), retrieved_5_tokenized_prompt_postfix=data.get("retrieved_5_tokenized_prompt_postfix"), retrieved_6_tokenized_prompt_postfix=data.get("retrieved_6_tokenized_prompt_postfix"), retrieved_7_tokenized_prompt_postfix=data.get("retrieved_7_tokenized_prompt_postfix"), retrieved_8_tokenized_prompt_postfix=data.get("retrieved_8_tokenized_prompt_postfix"), retrieved_9_tokenized_prompt_postfix=data.get("retrieved_9_tokenized_prompt_postfix"), retrieved_10_tokenized_prompt_postfix=data.get("retrieved_10_tokenized_prompt_postfix"), retrieved_11_tokenized_prompt_postfix=data.get("retrieved_11_tokenized_prompt_postfix"), retrieved_12_tokenized_prompt_postfix=data.get("retrieved_12_tokenized_prompt_postfix"), retrieved_13_tokenized_prompt_postfix=data.get("retrieved_13_tokenized_prompt_postfix"), retrieved_14_tokenized_prompt_postfix=data.get("retrieved_14_tokenized_prompt_postfix"), retrieved_15_tokenized_prompt_postfix=data.get("retrieved_15_tokenized_prompt_postfix"), retrieved_16_tokenized_prompt_postfix=data.get("retrieved_16_tokenized_prompt_postfix"), retrieved_17_tokenized_prompt_postfix=data.get("retrieved_17_tokenized_prompt_postfix"), retrieved_18_tokenized_prompt_postfix=data.get("retrieved_18_tokenized_prompt_postfix"), retrieved_19_tokenized_prompt_postfix=data.get("retrieved_19_tokenized_prompt_postfix"), query_tokenized_prompt_mask=data.get("query_tokenized_prompt_mask"), retrieved_0_tokenized_prompt_mask=data.get("retrieved_0_tokenized_prompt_mask"), retrieved_1_tokenized_prompt_mask=data.get("retrieved_1_tokenized_prompt_mask"), retrieved_2_tokenized_prompt_mask=data.get("retrieved_2_tokenized_prompt_mask"), retrieved_3_tokenized_prompt_mask=data.get("retrieved_3_tokenized_prompt_mask"), retrieved_4_tokenized_prompt_mask=data.get("retrieved_4_tokenized_prompt_mask"), retrieved_5_tokenized_prompt_mask=data.get("retrieved_5_tokenized_prompt_mask"), retrieved_6_tokenized_prompt_mask=data.get("retrieved_6_tokenized_prompt_mask"), retrieved_7_tokenized_prompt_mask=data.get("retrieved_7_tokenized_prompt_mask"), retrieved_8_tokenized_prompt_mask=data.get("retrieved_8_tokenized_prompt_mask"), retrieved_9_tokenized_prompt_mask=data.get("retrieved_9_tokenized_prompt_mask"), retrieved_10_tokenized_prompt_mask=data.get("retrieved_10_tokenized_prompt_mask"), retrieved_11_tokenized_prompt_mask=data.get("retrieved_11_tokenized_prompt_mask"), retrieved_12_tokenized_prompt_mask=data.get("retrieved_12_tokenized_prompt_mask"), retrieved_13_tokenized_prompt_mask=data.get("retrieved_13_tokenized_prompt_mask"), retrieved_14_tokenized_prompt_mask=data.get("retrieved_14_tokenized_prompt_mask"), retrieved_15_tokenized_prompt_mask=data.get("retrieved_15_tokenized_prompt_mask"), retrieved_16_tokenized_prompt_mask=data.get("retrieved_16_tokenized_prompt_mask"), retrieved_17_tokenized_prompt_mask=data.get("retrieved_17_tokenized_prompt_mask"), retrieved_18_tokenized_prompt_mask=data.get("retrieved_18_tokenized_prompt_mask"), retrieved_19_tokenized_prompt_mask=data.get("retrieved_19_tokenized_prompt_mask"), query_token_ar_mask=data.get("query_token_ar_mask"), retrieved_0_token_ar_mask=data.get("retrieved_0_token_ar_mask"), retrieved_1_token_ar_mask=data.get("retrieved_1_token_ar_mask"), retrieved_2_token_ar_mask=data.get("retrieved_2_token_ar_mask"), retrieved_3_token_ar_mask=data.get("retrieved_3_token_ar_mask"), retrieved_4_token_ar_mask=data.get("retrieved_4_token_ar_mask"), retrieved_5_token_ar_mask=data.get("retrieved_5_token_ar_mask"), retrieved_6_token_ar_mask=data.get("retrieved_6_token_ar_mask"), retrieved_7_token_ar_mask=data.get("retrieved_7_token_ar_mask"), retrieved_8_token_ar_mask=data.get("retrieved_8_token_ar_mask"), retrieved_9_token_ar_mask=data.get("retrieved_9_token_ar_mask"), retrieved_10_token_ar_mask=data.get("retrieved_10_token_ar_mask"), retrieved_11_token_ar_mask=data.get("retrieved_11_token_ar_mask"), retrieved_12_token_ar_mask=data.get("retrieved_12_token_ar_mask"), retrieved_13_token_ar_mask=data.get("retrieved_13_token_ar_mask"), retrieved_14_token_ar_mask=data.get("retrieved_14_token_ar_mask"), retrieved_15_token_ar_mask=data.get("retrieved_15_token_ar_mask"), retrieved_16_token_ar_mask=data.get("retrieved_16_token_ar_mask"), retrieved_17_token_ar_mask=data.get("retrieved_17_token_ar_mask"), retrieved_18_token_ar_mask=data.get("retrieved_18_token_ar_mask"), retrieved_19_token_ar_mask=data.get("retrieved_19_token_ar_mask"), query_token_loss_mask=data.get("query_token_loss_mask"), retrieved_0_token_loss_mask=data.get("retrieved_0_token_loss_mask"), retrieved_1_token_loss_mask=data.get("retrieved_1_token_loss_mask"), retrieved_2_token_loss_mask=data.get("retrieved_2_token_loss_mask"), retrieved_3_token_loss_mask=data.get("retrieved_3_token_loss_mask"), retrieved_4_token_loss_mask=data.get("retrieved_4_token_loss_mask"), retrieved_5_token_loss_mask=data.get("retrieved_5_token_loss_mask"), retrieved_6_token_loss_mask=data.get("retrieved_6_token_loss_mask"), retrieved_7_token_loss_mask=data.get("retrieved_7_token_loss_mask"), retrieved_8_token_loss_mask=data.get("retrieved_8_token_loss_mask"), retrieved_9_token_loss_mask=data.get("retrieved_9_token_loss_mask"), retrieved_10_token_loss_mask=data.get("retrieved_10_token_loss_mask"), retrieved_11_token_loss_mask=data.get("retrieved_11_token_loss_mask"), retrieved_12_token_loss_mask=data.get("retrieved_12_token_loss_mask"), retrieved_13_token_loss_mask=data.get("retrieved_13_token_loss_mask"), retrieved_14_token_loss_mask=data.get("retrieved_14_token_loss_mask"), retrieved_15_token_loss_mask=data.get("retrieved_15_token_loss_mask"), retrieved_16_token_loss_mask=data.get("retrieved_16_token_loss_mask"), retrieved_17_token_loss_mask=data.get("retrieved_17_token_loss_mask"), retrieved_18_token_loss_mask=data.get("retrieved_18_token_loss_mask"), retrieved_19_token_loss_mask=data.get("retrieved_19_token_loss_mask"), # Note: num_retrieved_observations is NOT included in the dataclass # It's metadata that can't be sharded and isn't needed during training ) # Defines the format of the actions. This field is included as "actions" inside the dictionary # produced by the data transforms. Actions = at.Float[ArrayOrStruct, "*b ah ad"] def preprocess_observation( rng: at.KeyArrayLike | None, observation: Observation, *, train: bool = False, image_keys: Sequence[str] = IMAGE_KEYS, 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). """ 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] 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, ) def preprocess_observation_prefix_postfix( rng: at.KeyArrayLike | None, observation: ObservationPrefixPostfix, *, train: bool = False, image_keys: Sequence[str] = IMAGE_KEYS, image_resolution: tuple[int, int] = IMAGE_RESOLUTION, ) -> ObservationPrefixPostfix: """Preprocess the observations by performing image augmentations (if train=True), resizing (if necessary), and filling in a default image mask (if necessary). """ 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] 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 ObservationPrefixPostfix( images=out_images, image_masks=out_masks, state=observation.state, tokenized_prompt_prefix=observation.tokenized_prompt_prefix, tokenized_prompt_postfix=observation.tokenized_prompt_postfix, tokenized_prompt_mask=observation.tokenized_prompt_mask, token_ar_mask=observation.token_ar_mask, token_loss_mask=observation.token_loss_mask, ) def extract_observation_from_ricl_observation( ricl_observation: RiclObservation, prefix: str, ) -> ObservationPrefixPostfix: return ObservationPrefixPostfix( images=getattr(ricl_observation, f"{prefix}images"), image_masks=getattr(ricl_observation, f"{prefix}image_masks"), state=getattr(ricl_observation, f"{prefix}state"), tokenized_prompt_prefix=getattr(ricl_observation, f"{prefix}tokenized_prompt_prefix"), tokenized_prompt_postfix=getattr(ricl_observation, f"{prefix}tokenized_prompt_postfix"), tokenized_prompt_mask=getattr(ricl_observation, f"{prefix}tokenized_prompt_mask"), token_ar_mask=getattr(ricl_observation, f"{prefix}token_ar_mask"), token_loss_mask=getattr(ricl_observation, f"{prefix}token_loss_mask"), ) @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) if remove_extra_params: params = ocp.transform_utils.intersect_trees(state.to_pure_dict(), 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) @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) -> 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 params_path.exists(): raise FileNotFoundError(f"Model params not found at: {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) # Handle StepMetadata object (orbax >= 0.1.0) or dict (older versions) if hasattr(metadata, "item_metadata"): params_meta = metadata.item_metadata["params"] else: params_meta = metadata["params"] item = {"params": params_meta} 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)