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

Code before:
    print("Action dim:    ", model.action_dim)
    print("Action horizon:", model.action_horizon)
    print("Max token len: ", model.max_token_len)

    if processor is None:
        processor_names = model.processor_names()
        if len(processor_names) != 1:
            raise ValueError(
                f"Multiple processors found in the checkpoint. Please specify a processor name: {processor_names}"
            )
        processor = processor_names[0]

    norm_stats = _import_norm_stats(ckpt_dir, processor)

    if config.param_path:
        for part in config.param_path.split("/"):
            if part not in params:
                raise ValueError(f"{part} not found in the checkpoint. Available keys: {list(params)}")
            params = params[part]

    if config.transform is not None:
        params = _transforms.transform_dict(config.transform, params)

    # Make sure that the loaded params are compatible with the openpi model.
    openpi_model = nnx.eval_shape(model.openpi_model_config().create, jax.random.key(0))
    _, openpi_state = nnx.split(openpi_model)

    # Remove any extra params that are not part of the openpi model.
    params = ocp.transform_utils.intersect_trees(openpi_state.to_pure_dict(), params)
    # Make sure that all expected params are present.
    at.check_pytree_equality(expected=openpi_state.to_pure_dict(), got=params, check_shapes=True, check_dtypes=False)

    # Save params.
    ckpt = ocp.StandardCheckpointer()
    ckpt.save(out_dir / "params", {"params": params})
    ckpt.wait_until_finished()

    # Save norm stats.
    _normalize.save(out_dir / "assets", norm_stats)


class PiModel(_model.BaseModel):
    """A model loaded from an internal exported model directory."""

    def __init__(self, ckpt_dir: pathlib.Path | str, params: at.Params | None = None):
        """Load a model from the internal checkpoint directory. Must point at the "model" sub-directory."""
        self.ckpt_dir = download.maybe_download(str(ckpt_dir))
        with (self.ckpt_dir / "graph").open("rb") as f:
            self._exported = jax.export.deserialize(bytearray(f.read()))

        input_spec = jax.tree.unflatten(self._exported.in_tree, self._exported.in_avals)[0]

Review comment (suggestion):
```suggestion
    openpi_state = nnx.state(openpi_model)
```

Code after:
    print("Action dim:    ", model.action_dim)
    print("Action horizon:", model.action_horizon)
    print("Max token len: ", model.max_token_len)

    if processor is None:
        processor_names = model.processor_names()
        if len(processor_names) != 1:
            raise ValueError(
                f"Multiple processors found in the checkpoint. Please specify a processor name: {processor_names}"
            )
        processor = processor_names[0]

    norm_stats = _import_norm_stats(ckpt_dir, processor)

    if config.param_path:
        for part in config.param_path.split("/"):
            if part not in params:
                raise ValueError(f"{part} not found in the checkpoint. Available keys: {list(params)}")
            params = params[part]

    if config.transform is not None:
        params = _transforms.transform_dict(config.transform, params)

    # Make sure that the loaded params are compatible with the openpi model.
    openpi_model = nnx.eval_shape(model.openpi_model_config().create, jax.random.key(0))
    openpi_params = nnx.state(openpi_model).to_pure_dict()

    # Remove any extra params that are not part of the openpi model.
    params = ocp.transform_utils.intersect_trees(openpi_params, params)
    # Make sure that all expected params are present.
    at.check_pytree_equality(expected=openpi_params, got=params, check_shapes=True, check_dtypes=False)

    # Save params.
    ckpt = ocp.StandardCheckpointer()
    ckpt.save(out_dir / "params", {"params": params})
    ckpt.wait_until_finished()

    # Save norm stats.
    _normalize.save(out_dir / "assets", norm_stats)


class PiModel(_model.BaseModel):
    """A model loaded from an internal exported model directory."""

    def __init__(self, ckpt_dir: pathlib.Path | str, params: at.Params | None = None):
        """Load a model from the internal checkpoint directory. Must point at the "model" sub-directory."""
        self.ckpt_dir = download.maybe_download(str(ckpt_dir))
        with (self.ckpt_dir / "graph").open("rb") as f:
            self._exported = jax.export.deserialize(bytearray(f.read()))

        input_spec = jax.tree.unflatten(self._exported.in_tree, self._exported.in_avals)[0]