Upload folder using huggingface_hub

checkpoint_files/WMPO-mimicgen/coffee_eval/run_20251107_151454.log

@@ -0,0 +1,2 @@
+# Log started at 2025-11-07T15:14:54.076342
+# Format: {timestamp}	{step}	{data_json}

checkpoint_files/WMPO_models/coffee/P_128/config.json

@@ -3168,4 +3168,4 @@
   "transformers_version": "4.40.1",
   "use_fused_vision_backbone": true,
   "vision_backbone_id": "dinosiglip-vit-so-224px"
-}
+}

checkpoint_files/WMPO_models/coffee/P_128/dataset_statistics.json

@@ -0,0 +1,127 @@
+{
+  "coffee_d0_300_demos": {
+    "action": {
+      "mean": [
+        0.03910227492451668,
+        0.026549872010946274,
+        0.016112811863422394,
+        0.00617121160030365,
+        0.009802641347050667,
+        0.032283715903759,
+        -0.1772146075963974
+      ],
+      "std": [
+        0.27866998314857483,
+        0.36967650055885315,
+        0.35988616943359375,
+        0.058437686413526535,
+        0.06546443700790405,
+        0.1048748567700386,
+        0.9841791391372681
+      ],
+      "max": [
+        1.0,
+        1.0,
+        1.0,
+        0.30490386486053467,
+        0.3423435091972351,
+        0.9236063361167908,
+        1.0
+      ],
+      "min": [
+        -1.0,
+        -1.0,
+        -1.0,
+        -0.28000304102897644,
+        -0.5609405040740967,
+        -0.2864341139793396,
+        -1.0
+      ],
+      "q01": [
+        -0.5959250259399415,
+        -0.8261726045608521,
+        -0.6470922386646271,
+        -0.13291869014501573,
+        -0.13417640179395676,
+        -0.1444952118396759,
+        -1.0
+      ],
+      "q99": [
+        0.8507216429710388,
+        0.9586128246784209,
+        0.9927734541893004,
+        0.14416057765483808,
+        0.17207743942737572,
+        0.37156413376331276,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 65463,
+    "num_trajectories": 300
+  }
+}

checkpoint_files/WMPO_models/coffee/P_128/preprocessor_config.json

@@ -0,0 +1,114 @@
+{
+  "auto_map": {
+    "AutoImageProcessor": "processing_prismatic.PrismaticImageProcessor",
+    "AutoProcessor": "processing_prismatic.PrismaticProcessor"
+  },
+  "image_processor_type": "PrismaticImageProcessor",
+  "image_resize_strategy": "resize-naive",
+  "input_sizes": [
+    [
+      3,
+      224,
+      224
+    ],
+    [
+      3,
+      224,
+      224
+    ]
+  ],
+  "interpolations": [
+    "bicubic",
+    "bicubic"
+  ],
+  "means": [
+    [
+      0.485,
+      0.456,
+      0.406
+    ],
+    [
+      0.5,
+      0.5,
+      0.5
+    ]
+  ],
+  "processor_class": "PrismaticProcessor",
+  "stds": [
+    [
+      0.229,
+      0.224,
+      0.225
+    ],
+    [
+      0.5,
+      0.5,
+      0.5
+    ]
+  ],
+  "tvf_crop_params": [
+    {
+      "output_size": [
+        224,
+        224
+      ]
+    },
+    {
+      "output_size": [
+        224,
+        224
+      ]
+    }
+  ],
+  "tvf_do_letterbox": false,
+  "tvf_letterbox_fill": null,
+  "tvf_normalize_params": [
+    {
+      "inplace": false,
+      "mean": [
+        0.484375,
+        0.455078125,
+        0.40625
+      ],
+      "std": [
+        0.228515625,
+        0.2236328125,
+        0.224609375
+      ]
+    },
+    {
+      "inplace": false,
+      "mean": [
+        0.5,
+        0.5,
+        0.5
+      ],
+      "std": [
+        0.5,
+        0.5,
+        0.5
+      ]
+    }
+  ],
+  "tvf_resize_params": [
+    {
+      "antialias": true,
+      "interpolation": 3,
+      "max_size": null,
+      "size": [
+        224,
+        224
+      ]
+    },
+    {
+      "antialias": true,
+      "interpolation": 3,
+      "max_size": null,
+      "size": [
+        224,
+        224
+      ]
+    }
+  ],
+  "use_fused_vision_backbone": true
+}

checkpoint_files/WMPO_models/coffee/P_128/processing_prismatic.py

@@ -0,0 +1,257 @@
+"""
+processing_prismatic.py
+
+HuggingFace-style preprocessor definitions for Prismatic VLMs, inheriting from `ProcessorMixin`. Default configuration
+specifies `siglip-224px+7b`.
+"""
+
+from typing import Any, ClassVar, List, Optional, Tuple, Union
+
+import timm.data
+import torch
+import torchvision.transforms.functional as TVF
+from PIL import Image
+from torchvision.transforms import CenterCrop, Compose, Normalize, Resize, ToTensor
+from transformers import PreTrainedTokenizerBase
+from transformers.image_processing_utils import BatchFeature, ImageProcessingMixin
+from transformers.processing_utils import ProcessorMixin
+from transformers.tokenization_utils import PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
+from transformers.utils import TensorType
+
+
+# === Image Processing ===
+def letterbox_pad_transform(image: Image.Image, padding_fill_value: Tuple[int, int, int]) -> Image.Image:
+    """Given a PIL.Image, pad to square by adding a symmetric border around the height/width."""
+    (w, h), max_wh = image.size, max(image.size)
+    horizontal_pad, vertical_pad = int((max_wh - w) / 2), int((max_wh - h) / 2)
+    padding = (horizontal_pad, vertical_pad, horizontal_pad, vertical_pad)
+
+    return TVF.pad(image, padding, fill=padding_fill_value, padding_mode="constant")
+
+
+class PrismaticImageProcessor(ImageProcessingMixin):
+    model_input_names: ClassVar[List[str]] = ["pixel_values"]
+
+    def __init__(
+        self,
+        use_fused_vision_backbone: bool = False,
+        image_resize_strategy: str = "letterbox",
+        input_sizes: Optional[List[Tuple[int, int, int]]] = None,
+        interpolations: Optional[List[str]] = None,
+        means: Optional[List[Tuple[float, float, float]]] = None,
+        stds: Optional[List[Tuple[float, float, float]]] = None,
+        **kwargs: str,
+    ) -> None:
+        """
+        Initialize a PrismaticImageProcessor as a wrapper around a torchvision transform; this transform will be
+        created by TIMM, and edited to follow our custom `image_resize_strategy` logic.
+
+        @param use_fused_vision_backbone: Boolean indicating single or fused (dual) vision backbone
+        @param image_resize_strategy: Prismatic image resize strategy in < resize-naive | resize-crop | letterbox >
+        @param input_size: [TIMM :: `data_cfg`] Input image size as tuple (channels, width, height)
+        @param interpolation: [TIMM :: `data_cfg`] Interpolation as string (default: "bicubic")
+        @param mean: [TIMM :: `data_cfg`] Normalization mean as float tuple (or two-tuple if `fused_backbone`)
+        @param std: [TIMM :: `data_cfg`] Normalization std as float tuple (or two-tuple if `fused_backbone`)
+        """
+        self.use_fused_vision_backbone = use_fused_vision_backbone
+        self.image_resize_strategy = image_resize_strategy
+
+        # Handle `None` default values
+        input_sizes = [(3, 224, 224)] if input_sizes is None else input_sizes
+        means = [(0.5, 0.5, 0.5)] if means is None else means
+        stds = [(0.5, 0.5, 0.5)] if stds is None else stds
+
+        # TIMM `data_cfg` Parameters
+        self.input_sizes, self.interpolations, self.means, self.stds = input_sizes, interpolations, means, stds
+
+        # Grab torchvision transforms via TIMM =>> need to parse for specific "functional" transform values!
+        self.tvf_resize_params, self.tvf_crop_params, self.tvf_normalize_params = [], [], []
+        self.tvf_do_letterbox, self.tvf_letterbox_fill = False, None
+
+        for idx in range(len(input_sizes)):
+            transform = timm.data.create_transform(
+                input_size=self.input_sizes[idx],
+                interpolation=self.interpolations[idx],
+                mean=self.means[idx],
+                std=self.stds[idx],
+                crop_pct=1.0,  # Set to 1.0 to ignore cropping (initial Resize sets `input_size`)
+                crop_mode="center",  # Default crop mode -- no-op when `crop_pct == 1.0`
+                is_training=False,  # No image augmentations when loading the transform!
+            )
+
+            # [Validation] Ensure appropriate transform structure, expected sizes
+            if not (
+                isinstance(transform, Compose)
+                and (len(transform.transforms) == 4)
+                and isinstance(transform.transforms[0], Resize)
+                and isinstance(transform.transforms[1], CenterCrop)
+                and isinstance(transform.transforms[2], ToTensor)
+                and isinstance(transform.transforms[3], Normalize)
+                and (transform.transforms[0].size == self.input_sizes[idx][-1])
+                and (transform.transforms[1].size == self.input_sizes[idx][-2:])
+            ):
+                raise ValueError(f"Unexpected TIMM image transformation structure/sizes: `{transform}`")
+
+            # HF Image Processors *must* be JSON-serializable; as such, cannot have torchvision. as an attribute.
+            #   => Instead, we're going to parse the transform and call "torchvision.transforms.functional" (`tvf`)
+            resize_t, crop_t, norm_t = transform.transforms[0], transform.transforms[1], transform.transforms[3]
+            self.tvf_resize_params.append(
+                {
+                    "size": resize_t.size,
+                    "interpolation": TVF.pil_modes_mapping[resize_t.interpolation],
+                    "max_size": None,
+                    "antialias": True,
+                }
+            )
+            self.tvf_crop_params.append({"output_size": crop_t.size})
+            self.tvf_normalize_params.append(
+                {
+                    "mean": norm_t.mean.float().numpy().tolist(),
+                    "std": norm_t.std.float().numpy().tolist(),
+                    "inplace": False,
+                }
+            )
+            self.tvf_do_letterbox, self.tvf_letterbox_fill = False, None
+
+            # Handle Prismatic `image_resize_strategy`
+            if self.image_resize_strategy == "resize-naive":
+                self.tvf_resize_params[idx]["size"] = (resize_t.size, resize_t.size)
+            elif self.image_resize_strategy == "letterbox":
+                self.tvf_do_letterbox, self.tvf_letterbox_fill = True, tuple([int(x * 255) for x in self.means[idx]])
+            elif self.image_resize_strategy == "resize-crop":
+                pass
+            else:
+                raise ValueError(f"Image resize strategy `{self.image_resize_strategy}` is not supported!")
+
+        # Dispatch **kwargs to super()
+        super().__init__(**kwargs)
+
+    def apply_transform(self, img: Image.Image) -> torch.Tensor:
+        """Apply `functional` variant of TIMM's Transform = Compose([Resize -> CenterCrop -> ToTensor -> Normalize])"""
+        if self.tvf_do_letterbox:
+            img = letterbox_pad_transform(img, self.tvf_letterbox_fill)
+
+        # [Contract] Fused Backbones expect "channel-stacked" inputs; we'll unpack on the model side!
+        imgs_t = []
+        for idx in range(len(self.input_sizes)):
+            img_idx = TVF.resize(img, **self.tvf_resize_params[idx])
+            img_idx = TVF.center_crop(img_idx, **self.tvf_crop_params[idx])
+            img_idx_t = TVF.to_tensor(img_idx)
+            img_idx_t = TVF.normalize(img_idx_t, **self.tvf_normalize_params[idx])
+            imgs_t.append(img_idx_t)
+
+        # [Contract] `imgs_t` is a list of Tensors of shape [3, input_size, input_size]; stack along dim = 0
+        img_t = torch.vstack(imgs_t)
+
+        return img_t
+
+    def preprocess(
+        self,
+        images: Union[Image.Image, List[Image.Image]],
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        **_: str,
+    ) -> BatchFeature:
+        """
+        Preprocess an image (or batch of images); note that unlike the `transformers :: BaseImageProcessor` we
+        explicitly only handle PIL.Image.Image instances for simplicity.
+
+        @param images: A (batch of) PIL.Image.Image instance(s) to preprocess.
+        @param return_tensors: BatchFeature default Tensor format (e.g., "pt" for torch); if None, returns np.ndarray
+
+        @return: Instance of `transformers :: BatchFeature` with a single key "pixel_values"
+        """
+        if not isinstance(images, list):
+            images = [images]
+
+        # Apply `self.img_transform` to each image (will return list of torch.Tensors); stack into "batched" Tensor
+        pixel_values = torch.stack([self.apply_transform(img.convert("RGB")) for img in images])
+
+        # Return BatchFeature =>> note that for compatibility, constructor expects Dict[str, np.ndarray], so we convert
+        return BatchFeature(data={"pixel_values": pixel_values.float().numpy()}, tensor_type=return_tensors)
+
+    def __call__(self, images: Union[Image.Image, List[Image.Image]], **kwargs) -> BatchFeature:
+        return self.preprocess(images, **kwargs)
+
+
+# === PrismaticProcessor =>> Wraps both ImageProcessor and Tokenizer ===
+#   =>> https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava/processing_llava.py
+class PrismaticProcessor(ProcessorMixin):
+    attributes: ClassVar[List[str]] = ["image_processor", "tokenizer"]
+    image_processor_class: str = "AutoImageProcessor"
+    tokenizer_class: str = "AutoTokenizer"
+
+    def __init__(
+        self,
+        image_processor: Optional[ImageProcessingMixin] = None,
+        tokenizer: Optional[PreTrainedTokenizerBase] = None,
+    ) -> None:
+        super().__init__(image_processor, tokenizer)
+
+    def __call__(
+        self,
+        text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]],
+        images: Union[Image.Image, List[Image.Image]],
+        padding: Union[bool, str, PaddingStrategy] = False,
+        truncation: Optional[Union[bool, str, TruncationStrategy]] = None,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[Union[str, TensorType]] = TensorType.PYTORCH,
+    ) -> BatchFeature:
+        """
+        Preprocess a given (batch) of text/images for a Prismatic VLM; forwards text to the underlying LLM's tokenizer,
+        forwards images to PrismaticImageProcessor.
+
+        @param text: The (batch) of text to encode; must be a string or list of strings.
+        @param images: A (batch of) PIL.Image.Image instance(s) to preprocess.
+        @param padding: Sequence padding strategy (if multiple specified) in < True = "longest" | "max_length" | False >
+        @param truncation: Truncation strategy for the output sequences; requires `max_length` to be specified
+        @param max_length: Maximum length (in tokens) to truncate
+        @param return_tensors: Type of return tensors (usually "pt" or TensorType.PYTORCH)
+
+        @return: BatchFeature with keys for `input_ids`, `attention_mask` and `pixel_values`.
+        """
+        pixel_values = self.image_processor(images, return_tensors=return_tensors)["pixel_values"]
+        text_inputs = self.tokenizer(
+            text, return_tensors=return_tensors, padding=padding, truncation=truncation, max_length=max_length
+        )
+
+        # [Validate] Need same number of images and text inputs!
+        if pixel_values.shape[0] != text_inputs.input_ids.shape[0]:
+            raise ValueError("Batch is malformed; expected same number of images and text inputs!")
+
+        return BatchFeature(data={**text_inputs, "pixel_values": pixel_values})
+
+    # === Tokenizer Dispatch Utilities =>> check `PreTrainedTokenizerBase` for documentation ===
+    def batch_decode(
+        self,
+        sequences: Union[List[int], List[List[int]], torch.Tensor, Any],  # `Any` = np.ndarray | tf.Tensor
+        skip_special_tokens: bool = False,
+        clean_up_tokenization_spaces: Optional[bool] = None,
+        **kwargs: str,
+    ) -> List[str]:
+        return self.tokenizer.batch_decode(
+            sequences=sequences,
+            skip_special_tokens=skip_special_tokens,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            **kwargs,
+        )
+
+    def decode(
+        self,
+        token_ids: Union[int, List[int], torch.Tensor, Any],  # `Any` = np.ndarray | tf.Tensor
+        skip_special_tokens: bool = False,
+        clean_up_tokenization_spaces: Optional[bool] = None,
+        **kwargs: str,
+    ) -> str:
+        return self.tokenizer.decode(
+            token_ids=token_ids,
+            skip_special_tokens=skip_special_tokens,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            **kwargs,
+        )
+
+    @property
+    def model_input_names(self) -> List[str]:
+        tokenizer_input_names = self.tokenizer.model_input_names
+        image_processor_input_names = self.image_processor.model_input_names
+
+        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))

checkpoint_files/WMPO_models/coffee/P_1280/dataset_statistics.json

@@ -0,0 +1,127 @@
+{
+  "coffee_d0_300_demos": {
+    "action": {
+      "mean": [
+        0.03910227492451668,
+        0.026549872010946274,
+        0.016112811863422394,
+        0.00617121160030365,
+        0.009802641347050667,
+        0.032283715903759,
+        -0.1772146075963974
+      ],
+      "std": [
+        0.27866998314857483,
+        0.36967650055885315,
+        0.35988616943359375,
+        0.058437686413526535,
+        0.06546443700790405,
+        0.1048748567700386,
+        0.9841791391372681
+      ],
+      "max": [
+        1.0,
+        1.0,
+        1.0,
+        0.30490386486053467,
+        0.3423435091972351,
+        0.9236063361167908,
+        1.0
+      ],
+      "min": [
+        -1.0,
+        -1.0,
+        -1.0,
+        -0.28000304102897644,
+        -0.5609405040740967,
+        -0.2864341139793396,
+        -1.0
+      ],
+      "q01": [
+        -0.5959250259399415,
+        -0.8261726045608521,
+        -0.6470922386646271,
+        -0.13291869014501573,
+        -0.13417640179395676,
+        -0.1444952118396759,
+        -1.0
+      ],
+      "q99": [
+        0.8507216429710388,
+        0.9586128246784209,
+        0.9927734541893004,
+        0.14416057765483808,
+        0.17207743942737572,
+        0.37156413376331276,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 65463,
+    "num_trajectories": 300
+  }
+}

checkpoint_files/WMPO_models/square/P_128/dataset_statistics copy.json

@@ -0,0 +1,127 @@
+{
+  "square_d0_300_demos": {
+    "action": {
+      "mean": [
+        0.11140137165784836,
+        0.0891946405172348,
+        -0.3048413097858429,
+        -0.001657384680584073,
+        0.0032311847899109125,
+        -0.04843222722411156,
+        -0.014298464171588421
+      ],
+      "std": [
+        0.4996664524078369,
+        0.42365702986717224,
+        0.5243746042251587,
+        0.06765525043010712,
+        0.11461760103702545,
+        0.2895003855228424,
+        0.9998131394386292
+      ],
+      "max": [
+        1.0,
+        1.0,
+        1.0,
+        0.3858289420604706,
+        1.0,
+        1.0,
+        1.0
+      ],
+      "min": [
+        -1.0,
+        -1.0,
+        -1.0,
+        -0.7480839490890503,
+        -1.0,
+        -1.0,
+        -1.0
+      ],
+      "q01": [
+        -0.9899384248256683,
+        -1.0,
+        -1.0,
+        -0.17326428920030593,
+        -0.22381987452507018,
+        -0.9575464141368866,
+        -1.0
+      ],
+      "q99": [
+        1.0,
+        1.0,
+        0.9203536868095384,
+        0.15283892214298228,
+        0.43935394287109325,
+        1.0,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 52663,
+    "num_trajectories": 325
+  }
+}

checkpoint_files/WMPO_models/square/P_1280/dataset_statistics.json

@@ -0,0 +1,127 @@
+{
+  "square_d0_300_demos": {
+    "action": {
+      "mean": [
+        0.11140137165784836,
+        0.0891946405172348,
+        -0.3048413097858429,
+        -0.001657384680584073,
+        0.0032311847899109125,
+        -0.04843222722411156,
+        -0.014298464171588421
+      ],
+      "std": [
+        0.4996664524078369,
+        0.42365702986717224,
+        0.5243746042251587,
+        0.06765525043010712,
+        0.11461760103702545,
+        0.2895003855228424,
+        0.9998131394386292
+      ],
+      "max": [
+        1.0,
+        1.0,
+        1.0,
+        0.3858289420604706,
+        1.0,
+        1.0,
+        1.0
+      ],
+      "min": [
+        -1.0,
+        -1.0,
+        -1.0,
+        -0.7480839490890503,
+        -1.0,
+        -1.0,
+        -1.0
+      ],
+      "q01": [
+        -0.9899384248256683,
+        -1.0,
+        -1.0,
+        -0.17326428920030593,
+        -0.22381987452507018,
+        -0.9575464141368866,
+        -1.0
+      ],
+      "q99": [
+        1.0,
+        1.0,
+        0.9203536868095384,
+        0.15283892214298228,
+        0.43935394287109325,
+        1.0,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 52663,
+    "num_trajectories": 325
+  }
+}

checkpoint_files/WMPO_models/stack_three/P_128/dataset_statistics copy.json

@@ -0,0 +1,127 @@
+{
+  "stack_three_d0_300_demos": {
+    "action": {
+      "mean": [
+        0.05272146314382553,
+        -0.0009713830077089369,
+        -0.025267750024795532,
+        -0.002258762251585722,
+        0.0068724267184734344,
+        0.01087257731705904,
+        -0.0585864894092083
+      ],
+      "std": [
+        0.40330731868743896,
+        0.39849069714546204,
+        0.36804842948913574,
+        0.061656929552555084,
+        0.07684116065502167,
+        0.3334025740623474,
+        0.9981837868690491
+      ],
+      "max": [
+        1.0,
+        1.0,
+        1.0,
+        0.4257144033908844,
+        0.5358051061630249,
+        1.0,
+        1.0
+      ],
+      "min": [
+        -1.0,
+        -1.0,
+        -1.0,
+        -0.411845862865448,
+        -0.4935794174671173,
+        -1.0,
+        -1.0
+      ],
+      "q01": [
+        -1.0,
+        -1.0,
+        -0.6893969440460205,
+        -0.15048537731170655,
+        -0.17073330640792847,
+        -1.0,
+        -1.0
+      ],
+      "q99": [
+        1.0,
+        1.0,
+        1.0,
+        0.14510192692279816,
+        0.20136691212654115,
+        1.0,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 76929,
+    "num_trajectories": 300
+  }
+}

checkpoint_files/WMPO_models/stack_three/P_1280/dataset_statistics.json

@@ -0,0 +1,127 @@
+{
+  "stack_three_d0_300_demos": {
+    "action": {
+      "mean": [
+        0.05272146314382553,
+        -0.0009713830077089369,
+        -0.025267750024795532,
+        -0.002258762251585722,
+        0.0068724267184734344,
+        0.01087257731705904,
+        -0.0585864894092083
+      ],
+      "std": [
+        0.40330731868743896,
+        0.39849069714546204,
+        0.36804842948913574,
+        0.061656929552555084,
+        0.07684116065502167,
+        0.3334025740623474,
+        0.9981837868690491
+      ],
+      "max": [
+        1.0,
+        1.0,
+        1.0,
+        0.4257144033908844,
+        0.5358051061630249,
+        1.0,
+        1.0
+      ],
+      "min": [
+        -1.0,
+        -1.0,
+        -1.0,
+        -0.411845862865448,
+        -0.4935794174671173,
+        -1.0,
+        -1.0
+      ],
+      "q01": [
+        -1.0,
+        -1.0,
+        -0.6893969440460205,
+        -0.15048537731170655,
+        -0.17073330640792847,
+        -1.0,
+        -1.0
+      ],
+      "q99": [
+        1.0,
+        1.0,
+        1.0,
+        0.14510192692279816,
+        0.20136691212654115,
+        1.0,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 76929,
+    "num_trajectories": 300
+  }
+}

checkpoint_files/WMPO_models/three_piece_assembly/P_128/dataset_statistics copy.json

@@ -0,0 +1,127 @@
+{
+  "three_piece_assembly_d0_300_demos": {
+    "action": {
+      "mean": [
+        0.03906247764825821,
+        0.004677821882069111,
+        -0.012144365347921848,
+        -0.0002992207882925868,
+        0.004983036313205957,
+        0.00010933508019661531,
+        0.11926877498626709
+      ],
+      "std": [
+        0.4390011429786682,
+        0.39680296182632446,
+        0.3512018322944641,
+        0.05673661455512047,
+        0.08805105090141296,
+        0.05951535701751709,
+        0.9926273226737976
+      ],
+      "max": [
+        1.0,
+        1.0,
+        1.0,
+        0.27066391706466675,
+        0.6146722435951233,
+        0.6450624465942383,
+        1.0
+      ],
+      "min": [
+        -1.0,
+        -1.0,
+        -1.0,
+        -0.3985214829444885,
+        -0.6851584911346436,
+        -0.7204136848449707,
+        -1.0
+      ],
+      "q01": [
+        -1.0,
+        -1.0,
+        -0.6249198973178863,
+        -0.13705135881900787,
+        -0.19542053192853928,
+        -0.13806192874908446,
+        -1.0
+      ],
+      "q99": [
+        1.0,
+        1.0,
+        0.990310949087143,
+        0.133660526573658,
+        0.25193322896957465,
+        0.14019585102796567,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 101091,
+    "num_trajectories": 300
+  }
+}

checkpoint_files/WMPO_models/three_piece_assembly/P_1280/dataset_statistics.json

@@ -0,0 +1,127 @@
+{
+  "three_piece_assembly_d0_300_demos": {
+    "action": {
+      "mean": [
+        0.03906247764825821,
+        0.004677821882069111,
+        -0.012144365347921848,
+        -0.0002992207882925868,
+        0.004983036313205957,
+        0.00010933508019661531,
+        0.11926877498626709
+      ],
+      "std": [
+        0.4390011429786682,
+        0.39680296182632446,
+        0.3512018322944641,
+        0.05673661455512047,
+        0.08805105090141296,
+        0.05951535701751709,
+        0.9926273226737976
+      ],
+      "max": [
+        1.0,
+        1.0,
+        1.0,
+        0.27066391706466675,
+        0.6146722435951233,
+        0.6450624465942383,
+        1.0
+      ],
+      "min": [
+        -1.0,
+        -1.0,
+        -1.0,
+        -0.3985214829444885,
+        -0.6851584911346436,
+        -0.7204136848449707,
+        -1.0
+      ],
+      "q01": [
+        -1.0,
+        -1.0,
+        -0.6249198973178863,
+        -0.13705135881900787,
+        -0.19542053192853928,
+        -0.13806192874908446,
+        -1.0
+      ],
+      "q99": [
+        1.0,
+        1.0,
+        0.990310949087143,
+        0.133660526573658,
+        0.25193322896957465,
+        0.14019585102796567,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 101091,
+    "num_trajectories": 300
+  }
+}