From 3c9167000c0b57180b6fb916120bdb91061b6268 Mon Sep 17 00:00:00 2001 From: dyf <1481321216@qq.com> Date: Tue, 19 May 2026 11:20:01 +0800 Subject: [PATCH] update yuanyou2 adapter --- src/openpi/policies/yuanyou2_policy.py | 185 +++++++++++++++++++++++++ 1 file changed, 185 insertions(+) create mode 100644 src/openpi/policies/yuanyou2_policy.py diff --git a/src/openpi/policies/yuanyou2_policy.py b/src/openpi/policies/yuanyou2_policy.py new file mode 100644 index 0000000000..8626741c5d --- /dev/null +++ b/src/openpi/policies/yuanyou2_policy.py @@ -0,0 +1,185 @@ +import dataclasses + +import einops +import numpy as np + +from openpi import transforms +from openpi.models import model as _model + + +YUANYOU2_STATE_DIM = 14 +YUANYOU2_ACTION_DIM = 14 + +# Yuanyou2 14-dim layout: +# [left_arm_6, left_gripper_1, right_arm_6, right_gripper_1] +LEFT_ARM_SLICE = slice(0, 6) +LEFT_GRIPPER_INDEX = 6 +RIGHT_ARM_SLICE = slice(7, 13) +RIGHT_GRIPPER_INDEX = 13 + + +def make_yuanyou2_example() -> dict: + """Creates a random input example for the Yuanyou2 policy.""" + return { + "observation/state": np.random.rand(YUANYOU2_STATE_DIM), + "observation/images/head": np.random.randint(256, size=(224, 224, 3), dtype=np.uint8), + "observation/images/left_wrist": np.random.randint(256, size=(224, 224, 3), dtype=np.uint8), + "observation/images/right_wrist": np.random.randint(256, size=(224, 224, 3), dtype=np.uint8), + # Only used during training / transform tests. + # Shape is usually [action_horizon, action_dim]. + "actions": np.random.rand(10, YUANYOU2_ACTION_DIM), + "prompt": "pick a cube and place it on another cube", + } + + +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 Yuanyou2Inputs(transforms.DataTransformFn): + """ + This class is used to convert inputs to the model to the expected format. It is used for both training and inference. + + For your own dataset, you can copy this class and modify the keys based on the comments below to pipe + the correct elements of your dataset into the model. + """ + + # Determines which model will be used. + # Do not change this for your own dataset. + model_type: _model.ModelType + + def __call__(self, data: dict) -> dict: + # Possibly need to parse images to uint8 (H,W,C) since LeRobot automatically + # stores as float32 (C,H,W), gets skipped for policy inference. + # Keep this for your own dataset, but if your dataset stores the images + # in a different key than "observation/image" or "observation/wrist_image", + # you should change it below. + # Pi0 models support three image inputs at the moment: one third-person view, + # and two wrist views (left and right). If your dataset does not have a particular type + # of image, e.g. wrist images, you can comment it out here and replace it with zeros like we do for the + # right wrist image below. + base_image = _parse_image(data["observation/images/head"]) + left_wrist_image = _parse_image(data["observation/images/left_wrist"]) + right_wrist_image = _parse_image(data["observation/images/right_wrist"]) + # base_image = _parse_image(data["observation/image"]) + # wrist_image = _parse_image(data["observation/wrist_image"]) + + state = np.asarray(data["observation/state"]) + + if state.shape[-1] != YUANYOU2_STATE_DIM: + raise ValueError( + f"Expected observation/state dim {YUANYOU2_STATE_DIM}, got {state.shape[-1]}. " + "Yuanyou2 state layout should be: " + "[left_arm_6, left_gripper, right_arm_6, right_gripper]." + ) + + # Create inputs dict. Do not change the keys in the dict below. + inputs = { + "state": state, + "image": { + "base_0_rgb": base_image, + "left_wrist_0_rgb": left_wrist_image, + # Pad any non-existent images with zero-arrays of the appropriate shape. + "right_wrist_0_rgb": right_wrist_image, + }, + "image_mask": { + "base_0_rgb": np.True_, + "left_wrist_0_rgb": np.True_, + # We only mask padding images for pi0 model, not pi0-FAST. Do not change this for your own dataset. + "right_wrist_0_rgb": np.True_, # if self.model_type == _model.ModelType.PI0_FAST else np.False_, + }, + } + + # Pad actions to the model action dimension. Keep this for your own dataset. + # Actions are only available during training. + # if "actions" in data: + # inputs["actions"] = data["actions"] + if "actions" in data: + actions = np.asarray(data["actions"]) + + if actions.shape[-1] != YUANYOU2_ACTION_DIM: + raise ValueError( + f"Expected actions dim {YUANYOU2_ACTION_DIM}, got {actions.shape[-1]}. " + "Yuanyou2 action layout should be: " + "[left_arm_6, left_gripper, right_arm_6, right_gripper]." + ) + + inputs["actions"] = actions + + # Pass the prompt (aka language instruction) to the model. + # Keep this for your own dataset (but modify the key if the instruction is not + # stored in "prompt"; the output dict always needs to have the key "prompt"). + if "prompt" in data: + inputs["prompt"] = data["prompt"] + + return inputs + +@dataclasses.dataclass(frozen=True) +class Yuanyou2Outputs(transforms.DataTransformFn): + """ + Converts model outputs back to Yuanyou2 action format. + + Yuanyou2 action layout: + [left_arm_6, left_gripper, right_arm_6, right_gripper] + """ + + def __call__(self, data: dict) -> dict: + actions = np.asarray(data["actions"][:, :YUANYOU2_ACTION_DIM]) + + if actions.shape[-1] != YUANYOU2_ACTION_DIM: + raise ValueError( + f"Expected output action dim {YUANYOU2_ACTION_DIM}, got {actions.shape[-1]}" + ) + + return {"actions": actions} + + +def decode_yuanyou2_action(action_14: np.ndarray) -> dict: + """ + Decode one Yuanyou2 14-dim action into arm and gripper commands. + + Input layout: + [left_arm_6, left_gripper, right_arm_6, right_gripper] + + Returns: + { + "left_arm": np.ndarray, shape (6,), + "left_gripper": float, + "right_arm": np.ndarray, shape (6,), + "right_gripper": float, + } + """ + action_14 = np.asarray(action_14) + + if action_14.shape[-1] != YUANYOU2_ACTION_DIM: + raise ValueError( + f"Expected action dim {YUANYOU2_ACTION_DIM}, got {action_14.shape[-1]}" + ) + + return { + "left_arm": action_14[LEFT_ARM_SLICE], + "left_gripper": float(action_14[LEFT_GRIPPER_INDEX]), + "right_arm": action_14[RIGHT_ARM_SLICE], + "right_gripper": float(action_14[RIGHT_GRIPPER_INDEX]), + } +# @dataclasses.dataclass(frozen=True) +# class Yuanyou2Outputs(transforms.DataTransformFn): +# """ +# This class is used to convert outputs from the model back the the dataset specific format. It is +# used for inference only. + +# For your own dataset, you can copy this class and modify the action dimension based on the comments below. +# """ + +# def __call__(self, data: dict) -> dict: +# # Only return the first N actions -- since we padded actions above to fit the model action +# # dimension, we need to now parse out the correct number of actions in the return dict. +# # For Libero, we only return the first 14 actions (since the rest is padding). +# # For your own dataset, replace `14` with the action dimension of your dataset. +# return {"actions": np.asarray(data["actions"][:, :YUANYOU2_ACTION_DIM])}