About openpi inference In SimplerEnv

[qjb-ldqmz]

Hi, dear author. The ckpt you provide in https://huggingface.co/HaomingSong/openpi0-bridge-lora
uv run scripts/serve_policy.py policy:checkpoint \ --policy.config=pi0_bridge_lora \ --policy.dir=$THE_MODEL_PATH
Is pi0_brige_lora implemented by yourselves, and the corresponding data_transform? Can you provide more details, Thank you.

[suhyeok-jang]

same issue

[DelinQu]

Hi suhyeok-jang,
HaomingSong has updated all the data config in Huggingface repos.
Let me know if you encounter any issues.

[qjb-ldqmz]

Thank you for your kind response. While the Trainconfig is relatively straightforward, what truly matters—and remains unclear—is the LeRobotBridgeDataConfig. I'm particularly eager to know what transformations it applies to the data.

[suhyeok-jang]

Thank you for sharing config. But i think important thing is missed

1. Action horizon (maybe 5 or 4?)
2. Batch size and train steps
3. Bridge input & output transform in the LeRobotBridgeDataConfig as @qjb-ldqmz said

[Luca-Wiehe]

@dataclasses.dataclass(frozen=True)
class LeRobotBridgeDataConfig(DataConfigFactory):
    use_quantile_norm: bool = True

    # Action keys that will be used to read the action sequence from the dataset.
    action_sequence_keys: Sequence[str] = ("action",)

    prompt_from_task: bool = True

    @override
    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
        # Make inputs look like they come from the Libero environment
        repack_transform = _transforms.Group(
            inputs=[
                _transforms.RepackTransform(
                    {
                        "observation/primary_image": "observation.images.image_0",
                        # "observation/left_yellow_image": "observation.images.image_1",
                        # "observation/right_blue_image": "observation.images.image_2",
                        # "observation/wirst_image": "observation.images.image_3",
                        "observation/state": "observation.state",
                        "actions": "action",
                        "prompt": "prompt",
                    }
                )
            ]
        )

        # Prepare data for policy training
        # Convert images to uint8 numpy arrays, add masks
        data_transforms = _transforms.Group(
            inputs=[
                ### THE FOLLOWING bridge_policy MODULE IS MISSING ###
                bridge_policy.BridgeInputs(
                    action_dim=model_config.action_dim,
                    model_type=model_config.model_type,
                )
            ],
            outputs=[bridge_policy.BridgeOutputs()],
        )

        # Model transforms include things like tokenizing the prompt and action targets
        model_transforms = ModelTransformFactory()(model_config)

        return dataclasses.replace(
            self.create_base_config(assets_dirs),
            repack_transforms=repack_transform,
            data_transforms=data_transforms,
            model_transforms=model_transforms,
            use_quantile_norm=self.use_quantile_norm,
            action_sequence_keys=self.action_sequence_keys,
            prompt_from_task=self.prompt_from_task,
        )

The DataConfig has been published as part of the HuggingFace repository. However, the bridge_policy.py is not posted anywhere.

Similar issue about the checkpoints provided for Fractal, i.e. Google Robot evaluation. The most important file (fractal_policy.py) that is used in the training configuration is never mentioned; not even in Haoming Song's fork of openpi.

Can you update us regarding this @DelinQu @HaomingSong?

[Luca-Wiehe]

Here's the fractal_policy.py for which I could make the fractal checkpoint work. Make sure to place it in openpi/src/openpi/policies/

# fractal_policy.py
import dataclasses

import einops
import numpy as np

from openpi import transforms
from openpi.models import model as _model

def _parse_image(image) -> np.ndarray:
    image = np.asarray(image)
    if np.issubdtype(image.dtype, np.floating):
        image = (255 * image).astype(np.uint8)
    if image.shape[0] == 3:
        image = einops.rearrange(image, "c h w -> h w c")
    return image


@dataclasses.dataclass(frozen=True)
class FractalInputs(transforms.DataTransformFn):
    action_dim: int

    model_type: _model.ModelType = _model.ModelType.PI0

    def __call__(self, data: dict) -> dict:
        mask_padding = self.model_type == _model.ModelType.PI0

        state = transforms.pad_to_dim(data["observation/state"], self.action_dim)
        base_image = _parse_image(data["observation/primary_image"])

        inputs = {
            "state": state,
            "image": {
                "base_0_rgb": base_image,
                "left_wrist_0_rgb": np.zeros_like(base_image),
                "right_wrist_0_rgb": np.zeros_like(base_image),
            },
            "image_mask": {
                "base_0_rgb": np.True_,
                "left_wrist_0_rgb": np.False_ if mask_padding else np.True_,
                "right_wrist_0_rgb": np.False_ if mask_padding else np.True_,
            },
        }

        if "actions" in data:
            actions = transforms.pad_to_dim(data["actions"], self.action_dim)
            inputs["actions"] = actions

        if "prompt" in data:
            inputs["prompt"] = data["prompt"]

        return inputs


@dataclasses.dataclass(frozen=True)
class FractalOutputs(transforms.DataTransformFn):
    def __call__(self, data: dict) -> dict:
        return {"actions": np.asarray(data["actions"][:, :7])}

Edit: While the code doesn't throw any errors, I cannot reproduce the performance that is reported in the spreadsheet. The success rate is close to zero. I am using @HaomingSong's openpi0-fast-fractal-fft checkpoint