Repository: Physical-Intelligence/openpi
File: src/openpi/models/exported.py (Python)

Code before:
def _example_to_obs(example: dict) -> common.Observation:
    images, image_masks = {}, {}
    for k, v in example["image"].items():
        if k.endswith("_mask"):
            image_masks[k.removesuffix("_mask")] = v
        else:
            images[k] = v

    return common.Observation.from_dict(
        {
            "image": images,
            "image_mask": image_masks,
            "state": example["state"],
            "tokenized_prompt": example["prompt_tokens"],
            "tokenized_prompt_mask": example["mask_input"],
        }
    )


def import_norm_stats(ckpt_path: epath.Path, processor_name: str) -> dict[str, _normalize.NormStats]:
    path = ckpt_path / "processors" / processor_name
    if not (found_files := list(path.glob("*/norm_stats.msgpack"))):
        raise FileNotFoundError(f"No norm_stats.msgpack found in {path}")

    with found_files[0].open("rb") as f:
        norm_stats = flax.serialization.from_bytes(dict, f.read())

    input_state = norm_stats["input_norms"]["state"]
    state_stats = _normalize.NormStats(mean=input_state["mean"], std=input_state["std"])

    output_actions = norm_stats["output_norms"]["actions"]
    actions_stats = _normalize.NormStats(mean=output_actions["mean"], std=output_actions["std"])

    return {
        "state": state_stats,
        "actions": actions_stats,
    }


def _make_batch(data: dict) -> dict:
    return jax.tree.map(lambda x: x[jnp.newaxis, ...], data)


def _unbatch(data: dict) -> dict:
    return jax.tree.map(lambda x: x[0, ...], data)


Review comment (suggestion):
```suggestion
        norm_stats = flax.serialization.msgpack_restore(f.read())
```

Code after:
def _example_to_obs(example: dict) -> common.Observation:
    images, image_masks = {}, {}
    for k, v in example["image"].items():
        if k.endswith("_mask"):
            image_masks[k.removesuffix("_mask")] = v
        else:
            images[k] = v

    return common.Observation.from_dict(
        {
            "image": images,
            "image_mask": image_masks,
            "state": example["state"],
            "tokenized_prompt": example["prompt_tokens"],
            "tokenized_prompt_mask": example["mask_input"],
        }
    )


def import_norm_stats(ckpt_path: epath.Path, processor_name: str) -> dict[str, _normalize.NormStats]:
    path = ckpt_path / "processors" / processor_name
    if not (found_files := list(path.glob("*/norm_stats.msgpack"))):
        raise FileNotFoundError(f"No norm_stats.msgpack found in {path}")

    with found_files[0].open("rb") as f:
        norm_stats = flax.serialization.msgpack_restore(f.read())

    input_state = norm_stats["input_norms"]["state"]
    state_stats = _normalize.NormStats(mean=input_state["mean"], std=input_state["std"])

    output_actions = norm_stats["output_norms"]["actions"]
    actions_stats = _normalize.NormStats(mean=output_actions["mean"], std=output_actions["std"])

    return {
        "state": state_stats,
        "actions": actions_stats,
    }


def _make_batch(data: dict) -> dict:
    return jax.tree.map(lambda x: x[jnp.newaxis, ...], data)


def _unbatch(data: dict) -> dict:
    return jax.tree.map(lambda x: x[0, ...], data)