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

Code before:
        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]
        if params is None:
            params = _load_params(self.ckpt_dir, input_spec[0])
        self.example_spec = input_spec[2]
        self.sample_spec = input_spec[3]

        # Extract the action properties from the output spec.
        output_spec = jax.tree.unflatten(self._exported.out_tree, self._exported.out_avals)
        if "actions" in output_spec:
            # This is a pi0 model.
            self._output_key = "actions"
            self._model_type = _model.ModelType.PI0
            action_horizon, action_dim = output_spec[self._output_key].shape
        elif "output_tokens" in output_spec:
            # This is a pi0-fast model.
            self._output_key = "output_tokens"
            self._model_type = _model.ModelType.PI0_FAST
            # The output is tokenized actions and so we have to infer the action properties
            # from other properties.
            action_dim = self.example_spec["state"].shape[-1]
            action_horizon = 15
        else:
            raise ValueError(f"Unknown output spec: {output_spec}")

        max_token_len = self.example_spec["prompt_tokens"].shape[-1]

        super().__init__(action_dim=action_dim, action_horizon=action_horizon, max_token_len=max_token_len)

        # wrap arrays in nnx.Param to make nnx happy
        self.params = jax.tree.map(lambda x: nnx.Param(x), params)

    @property
    def model_type(self) -> _model.ModelType:
        return self._model_type

    @override
    def sample_actions(self, rng: at.KeyArrayLike, observation: _model.Observation, **sample_kwargs) -> _model.Actions:
        if observation.state.ndim == 2 and observation.state.shape[0] != 1:
            raise ValueError("Only batch_size=1 is supported.")

        # Convert to the example format.
        example = _obs_to_example(observation, self.example_spec)
        example = _unbatch(example)

        # Resize the input images if needed.
        def resize_if_needed(key, image):
            target_shape = self.example_spec["image"][key].shape

Review comment (suggestion):
Is there anything else in the output spec we could use to infer the action horizon? or in the processor dict?
Hardcoding it to 15 will break for any model that's not DROID (since they all use different action chunks). Otherwise we somehow need to have users explicitly specify it alongside the checkpoint they want to load.

Code after:
        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]
        if params is None:
            params = _load_params(self.ckpt_dir, input_spec[0])
        self.example_spec = input_spec[2]
        self.sample_spec = input_spec[3]

        # Extract the action properties from the output spec.
        output_spec = jax.tree.unflatten(self._exported.out_tree, self._exported.out_avals)
        if "actions" in output_spec:
            # This is a pi0 model.
            self._output_key = "actions"
            self._model_type = _model.ModelType.PI0
            action_horizon, action_dim = output_spec[self._output_key].shape
        elif "output_tokens" in output_spec:
            # This is a pi0-fast model.
            self._output_key = "output_tokens"
            self._model_type = _model.ModelType.PI0_FAST
            # The output is tokenized actions and so we have to infer the action properties
            # from other properties.
            action_dim = self.example_spec["state"].shape[-1]
            # TODO(ury): Figure out how to get this information from the checkpoint.
            action_horizon = 15
        else:
            raise ValueError(f"Unknown output spec: {output_spec}")

        max_token_len = self.example_spec["prompt_tokens"].shape[-1]

        super().__init__(action_dim=action_dim, action_horizon=action_horizon, max_token_len=max_token_len)

        # wrap arrays in nnx.Param to make nnx happy
        self.params = jax.tree.map(lambda x: nnx.Param(x), params)

    @property
    def model_type(self) -> _model.ModelType:
        return self._model_type

    @override
    def sample_actions(self, rng: at.KeyArrayLike, observation: _model.Observation, **sample_kwargs) -> _model.Actions:
        if observation.state.ndim == 2 and observation.state.shape[0] != 1:
            raise ValueError("Only batch_size=1 is supported.")

        # Convert to the example format.
        example = _obs_to_example(observation, self.example_spec)
        example = _unbatch(example)

        # Resize the input images if needed.
        def resize_if_needed(key, image):