Implementation into openpi

[rohan-bansal]

Hey there, this is a great work!

What would it take to implement this into the openpi repository? I've currently created a new model file in src/openpi/models based off of BaseModel, except it calls Pi0Infererence.forward after processing observations in the sample_actions() function instead of proceeding with denoising steps.

    @override
    @torch.no_grad()
    def sample_actions(self, device, observation: _model.Observation, 
                    noise=None, num_steps=10) -> torch.Tensor:
        
        bsize = observation.state.shape[0]
        if noise is None:
            actions_shape = (bsize, self.config.action_horizon, self.config.action_dim)
            noise = self.sample_noise(actions_shape, device)

        # use pi0_pytorch preprocessing func
        images, img_masks, lang_tokens, lang_masks, state = self._preprocess_observation(observation, train=False)
        

        images = images[:self.num_views]
        images_converted = []
        for img in images:
            img = img[0]
            img = img.permute(1, 2, 0)  # Convert to (H, W, C)
            images_converted.append(img)
        
        images_stacked = torch.stack(images_converted, dim=0)  # (num_views, H, W, C)
        images_bf16 = images_stacked.to(dtype=torch.bfloat16, device='cuda')
        
        state_bf16 = state[0].to(dtype=torch.bfloat16, device='cuda')
        noise_input = noise[0].to(dtype=torch.bfloat16, device='cuda')
        
        # Run inference
        output_actions = self.inference_engine.forward(
            images_bf16,
            state_bf16,
            noise_input
        )
        
        # (1, action_horizon, action_dim)
        return output_actions.unsqueeze(0).to(dtype=torch.float32)

I'm noticing problems in simulation rollout (LIBERO), where robot motion is unintelligible. What am I doing wrong here?

[mayunchao1994]

Could you confirm that the shapes of images_bf16, state_bf16, and noise_input match mine?

Additionally, when initializing the model with Pi0Inference(checkpoint=weights, num_views=2, chunk_size=50), the num_views and chunk_size parameters must be consistent with your input.

[Dexmal]

Your code snippet looks fine to us. For further inspection, can you provide the following data files (by google drive?) so that we can work on this issue?

• Checkpoint generated by the openpi codebase
• The input tensors (images_bf16, state_bf16, noise_input)
• The desired output tensors (when you do not use our kernels)
  We will try to reproduce the mismatch and see what might go wrong.

[rohan-bansal]

Hi @Dexmal , sorry for the late reply. I have uploaded the requested data to this drive folder.

Some details:

• The input/output pkl files for the realtime kernel version are saved as torch tensors. The output pkl file for the non-realtime kernel version is saved in a jax tensor.
• I have converted the checkpoint using convert_from_jax.py in this repository with the prompt "put both moka pots on the stove".
• When testing on LIBERO, I use the libero-10 task suite and skip all tasks except task_id = 8 (which matches the above prompt)
• I am not sure how helpful the input tensors when using your kernels will be, since images and state will be wildly inaccurate due to weird robot motion. However, I have uploaded them anyway.

Let me know if you need anything else - would love to solve this problem as well.

[rohan-bansal]

@mayunchao1994 @Dexmal is there any update on this? 🙂

[mayunchao1994]

Sorry, I’ve been busy and am just now getting to this issue. Is the converted_checkpoint.pkl you provided already the Torch model converted using convert_from_jax.py? For the comparison, I need the checkpoint generated by the OpenPI codebase (the JAX model) so I can compare the differences between JAX and Torch. Could you share that again? I’ll prioritize resolving this.

[rohan-bansal]

The checkpoint I uploaded is the Torch model converted using your conversion script.

The JAX checkpoint from the OpenPI codebase is the default one. However, it is > 10GB and difficult to upload to GDrive. It would be easier to do the following:

Open src/openpi/shared/download.py, and set the DEFAULT_CACHE_DIR at the top to where you want the checkpoint downloaded.

Then run:

import openpi.shared.download as download

download.maybe_download("gs://openpi-assets/checkpoints/pi0_libero")

[mayunchao1994]

Okay, it's the official pi0_libero model. I mistakenly thought it was a model you finetuned yourself. Let me fix this issue.

[mayunchao1994]

Sorry, it took so long to resolve this issue. The state embed is overwritten, resulting in inconsistent results with jax. This problem is now solved.
I used the provided images, states, noise, and "put both moka pots on the stove" as inputs. Since the provided images were preprocessed, I unnormalized them (img = ((img + 1.0) / 2.0 * 255.0).clip(0, 255)) before using them as input. I selected idx=20 for testing, used the model "gs://openpi-assets/checkpoints/pi0_libero", converted the model using convert_from_jax.py, and compared the outputs (return values of the sample_action function) between triton and jax. The results are shown below:

As you can see, the results are now consistent. Could you please pull the latest code and test it? Please let me know if you encounter any issues.

[rohan-bansal]

@mayunchao1994 can you please provide your rollout code? I seem to still be encountering issues rolling out. Additionally, you unnormalized the images but according to the pi0_infer input, it requires a normalized image?

[mayunchao1994]

I'm not calling the model within sample_actions, but instead using _policy_config.create_trained_policy, so I need the raw inputs without going through inputs = self._input_transform(inputs). Here is my test code for your reference.

import numpy as np
import torch
import argparse
from PIL import Image
import einops
import json
import cv2
import os
from pi0_infer import Pi0Inference
from openpi.training import config as _config
from openpi.policies import policy_config as _policy_config
import pickle

class Pi0ModelEvaluator:
    def __init__(self, task, model_type: str, triton_path: str, jax_path: str, norm_stats_dir:str, config_name: str):
        self.triton_path = triton_path
        self.jax_path = jax_path
        self.task = task
        self.model_type = model_type
        self.norm_stats_dir = norm_stats_dir
        self.policy = None
        self.config_name = config_name
        self.norm_stats = None
        self.results = {
            'episode_results': []
        }
        if self.model_type == "triton":
            self.policy, self.norm_stats = self._load_model(model_type=self.model_type)
        elif self.model_type == "jax":
            self.policy = self._load_jax_model()

    def _load_jax_model(self):
        config = _config.get_config(self.config_name)
        policy = _policy_config.create_trained_policy(config, self.jax_path)
        return policy

    def _load_model(self, model_type="triton"):
        with open(self.triton_path, 'rb') as f:
            weights = pickle.load(f)
        policy = Pi0Inference(checkpoint=weights, num_views=2, chunk_size=50)
        norm_stats = self._load_norm_stats(self.norm_stats_dir) if self.norm_stats_dir else None
        return policy, norm_stats

    def _load_norm_stats(self, norm_stats_dir: str) -> dict:
        norm_stats_path = os.path.join(norm_stats_dir, "norm_stats.json")
        if os.path.exists(norm_stats_path):
            with open(norm_stats_path, 'r') as f:
                return json.load(f)['norm_stats']
        return None
    
    def _parse_image(self, 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

    def _pad_to_dim(self, x: np.ndarray, target_dim: int, axis: int = -1) -> np.ndarray:
        current_dim = x.shape[axis]
        if current_dim < target_dim:
            pad_width = [(0, 0)] * len(x.shape)
            pad_width[axis] = (0, target_dim - current_dim)
            return np.pad(x, pad_width)
        return x

    def _resize_with_pad(self, image: np.ndarray, height: int = 224, width: int = 224) -> np.ndarray:
        pil_image = Image.fromarray(image)
        cur_width, cur_height = pil_image.size
        if cur_width == width and cur_height == height:
            return image
        
        ratio = max(cur_width / width, cur_height / height)
        resized_height = int(cur_height / ratio)
        resized_width = int(cur_width / ratio)
        resized_image = pil_image.resize((resized_width, resized_height), resample=Image.BILINEAR)
        zero_image = Image.new(resized_image.mode, (width, height), 0)
        pad_height = max(0, int((height - resized_height) / 2))
        pad_width = max(0, int((width - resized_width) / 2))
        zero_image.paste(resized_image, (pad_width, pad_height))
        return np.array(zero_image)

    def _normalize_image(self, image: np.ndarray) -> np.ndarray:
        image = image.astype(np.float32) / 255.0 * 2.0 - 1.0
        return image

    def _normalize_state(self, state: np.ndarray, norm_stats: dict) -> np.ndarray:
        if norm_stats and "state" in norm_stats:
            state_mean = np.array(norm_stats["state"]["mean"])
            state_mean = self._pad_to_dim(state_mean, 32)
            state_std = np.array(norm_stats["state"]["std"])
            state_std = self._pad_to_dim(state_std, 32)
            return (state - state_mean) / (state_std + 1e-6)
        return None

    def _unnormalize_state(self, actions: np.ndarray, norm_stats: dict) -> np.ndarray:
        if norm_stats and "actions" in norm_stats:
            actions_mean = np.array(norm_stats["actions"]["mean"])
            actions_mean = self._pad_to_dim(actions_mean, 32)
            actions_std = np.array(norm_stats["actions"]["std"])
            actions_std = self._pad_to_dim(actions_std, 32)
            return actions * (actions_std + 1e-6) + actions_mean
        return None

    def _apply_input_transforms(self, data: dict, action_dim: int = 32, norm_stats: dict = None) -> dict:
        state = self._pad_to_dim(data["state"], action_dim)
        state = self._normalize_state(state, norm_stats)
        base_image = self._parse_image(data["base_0_rgb"])
        left_wrist_image = self._parse_image(data["left_wrist_0_rgb"])
        right_wrist_image = self._parse_image(data["right_wrist_0_rgb"])
        base_image = self._resize_with_pad(base_image, 224, 224)
        base_image = self._normalize_image(base_image)
        left_wrist_image = self._resize_with_pad(left_wrist_image, 224, 224)
        left_wrist_image = self._normalize_image(left_wrist_image)
        right_wrist_image = self._resize_with_pad(right_wrist_image, 224, 224)
        right_wrist_image = self._normalize_image(right_wrist_image)
        image_dict = {
            "base_0_rgb": base_image,
            "left_wrist_0_rgb": left_wrist_image,
            "right_wrist_0_rgb": right_wrist_image,
        }
        image_mask_dict = {
            "base_0_rgb": np.True_,
            "left_wrist_0_rgb": np.False_,
            "right_wrist_0_rgb": np.False_,
        }
        inputs = {
            "state": state,
            "image": image_dict,
            "image_mask": image_mask_dict,
        }
        if "prompt" in data:
            inputs["prompt"] = data["prompt"]
        return inputs

    def infer(self, inputs: dict, noise: np.ndarray) -> dict:
        if self.model_type == "triton":
            transformed_inputs = self._apply_input_transforms(inputs, action_dim=32, norm_stats=self.norm_stats)
            images = []
            for view in ["base_0_rgb", "left_wrist_0_rgb", "right_wrist_0_rgb"]:
                img = transformed_inputs["image"][view]
                images.append(torch.from_numpy(img))
            
            observation_images = torch.stack(images[0:2], dim=0).to(torch.float32)
            observation_state = torch.from_numpy(transformed_inputs["state"].astype(np.float32)).unsqueeze(0).to(torch.float32)
            normalized_observation_images = observation_images.cuda()
            normalized_observation_state = observation_state.cuda()

            diffusion_noise = torch.from_numpy(noise).to(torch.float32).cuda()
            self.policy.buffers['observation_images_normalized'].copy_(normalized_observation_images)
            self.policy.buffers['observation_state_normalized'].copy_(normalized_observation_state.squeeze(0))
            self.policy.buffers['diffusion_noise'].copy_(diffusion_noise)
            self.policy.record_run()
            actions = self.policy.buffers['diffusion_noise']
            actions = actions.cpu().float().numpy()
            actions = self._unnormalize_state(actions, self.norm_stats)[:, :7]
            actions[..., :6] = actions[..., :6] + inputs["state"][..., :6]
            # actions = actions[:, :7]
            return {
                "actions": actions
            }
        elif self.model_type == "jax":
            actions = self.policy.infer(inputs, noise=noise)
            # actions["actions"] = actions["actions"][:, :7]
            return actions

def main():
    parser = argparse.ArgumentParser(description="VLA Model Inference")
    parser.add_argument('--triton_path', type=str, default='pi0_libero_triton.pkl')
    parser.add_argument('--jax_path', type=str, default='pi0_libero')
    parser.add_argument('--norm_stats_dir', type=str, default='pi0_libero/assets/physical-intelligence/libero')
    parser.add_argument('--config_name', type=str, default='pi0_libero')

    args = parser.parse_args()

    example_image_global = cv2.imread("1.png")
    example_image_hand = cv2.imread("2.png")
    with open("noise_bf16_save.pkl", "rb") as f:
        noises = pickle.load(f)
    noise = noises[0].float().cpu().numpy()
    with open("states_bf16_save.pkl", "rb") as f:
        states = pickle.load(f)
    state = states[0].float().cpu().numpy()
    inputs_triton = {
        "base_0_rgb": example_image_global,
        "left_wrist_0_rgb": example_image_hand,
        "right_wrist_0_rgb": example_image_hand,
        "state": state,
        "prompt": "Sort bowls and paper cups into their designated placese"
    } 
    inputs_jax = {
        "observation/image": example_image_global,
        "observation/wrist_image": example_image_hand,
        "observation/state": state,
        "prompt": "Sort bowls and paper cups into their designated places"
    }

    pi0_triton = Pi0ModelEvaluator(task='check_consistency', model_type="triton", triton_path=args.triton_path, jax_path=args.jax_path,
                                    norm_stats_dir=args.norm_stats_dir, config_name=args.config_name)
    
    result_triton = pi0_triton.infer(inputs_triton, noise)
    del pi0_triton
    print("finish triton")
    pi0_jax = Pi0ModelEvaluator(task='check_consistency', model_type="jax", triton_path=args.triton_path, jax_path=args.jax_path,
                                    norm_stats_dir=args.norm_stats_dir, config_name=args.config_name)
    result_jax = pi0_jax.infer(inputs_jax, noise)
    del pi0_jax
    print("finish jax")
    print(f"Triton actions range: [{result_triton['actions'].min():.6f}, {result_triton['actions'].max():.6f}]")
    print(f"JAX actions range: [{result_jax['actions'].min():.6f}, {result_jax['actions'].max():.6f}]")

    joint_names = ['Joint 1', 'Joint 2', 'Joint 3', 'Joint 4', 'Joint 5', 'Joint 6', 'Gripper Width']
    triton_jax_mae = np.mean(np.abs(result_triton['actions'] - result_jax['actions']))
    print(f"Triton vs JAX MAE: {triton_jax_mae:.6f}")
    print("Per-dimension MAE:")
    for i in range(7):
        dim_mae = np.mean(np.abs(result_triton['actions'][:, i] - result_jax['actions'][:, i]))
        print(f"{joint_names[i]}: {dim_mae:.6f}")

if __name__ == "__main__":
    main()

[rohan-bansal]

Thanks for your code, and happy holidays! I spent a little bit playing with it, and did a quick integration into sample_actions - see this Gist for the policy.

From what I understand, you perform the input transform and the output transform in your own code, and disregard OpenPi's version. So after sample_actions is called in the linked file, I do not do an output transformation and immediately pass the resulting actions to the environment.

I implemented a deterministic replay system such that observations and noise are saved from a Jax model inference episode (blue), then replayed with the Triton model (red).

I'm noticing that some dimensions follow the original trajectory in a decent manner (x,y,z,gripper) but others do not (large offset for rx, ry, rz). Is this behavior explainable (either through a bug in my code, or a flaw in my reproduction logic?)

The norm_stats.json file is constant, I do not think it is the issue.

@mayunchao1994

[mayunchao1994]

Based on your visualization, it looks like extra_delta_transform is causing a translation offset in actions[..., :6] between JAX and Triton. The gripper values are reasonable since they're not delta. In the official pi0 code, after sample_actions, it calls outputs = self._output_transform(outputs), which includes actions[..., :6] = actions[..., :6] + observation["state"][..., :6]. I suspect delta_transform might be applied twice. Please also verify that outputs = self._output_transform(outputs) is consistent between the Triton and JAX versions. Additionally, you can try running the comparison results using the script I provided to see if it matches your script's results. You can also print the input to the model (output after input_transform), the model output, and the final output after output_transform to facilitate debugging.

[mayunchao1994]

I found a bug: decoder_x_buf should have shape (chunk_size, 1024), not (decoder_seq_len, 1024). I’ve updated the code. Please re-test the diff after aligning the input/output handling mentioned above.

[rohan-bansal]

Awesome thanks @mayunchao1994! After aligning the input/output transforms I was still getting weird results, but the decoder_x_buf issue was what fixed it for me. Here is the comparison graph now:

Is the flatlining between steps 200 and 250 for RY expected?

Additionally, now that this works, I intend to implement batching support.

[mayunchao1994]

Thanks a lot for raising this issue and providing a test case. It really helped me track down and fix the bug.
I think the differences at steps 200 and 250 for RY are expected, since the Triton version uses bf16 throughout, and some numerical precision differences vs. the JAX version are normal.
Implementing batching support would be very useful. We’re open to PRs — feel free to send a patch.