OpenPI项目中UR5机器人微调时的数据配置问题解析

问题背景

在OpenPI项目中，用户尝试对UR5机器人进行微调时遇到了数据配置问题。具体表现为在数据转换过程中无法正确识别"actions"列，导致训练过程中出现KeyError错误。

核心问题分析

该问题的根源在于数据配置文件中键名映射的不一致性。原始数据集中的动作列名为"action"，而模型期望的输入键名为"actions"，这种命名差异导致了数据流的中断。

解决方案详解

经过技术验证，正确的解决方案需要从两个层面进行修改：

1. 数据重映射配置

在LeRobotUR5DataConfig类中，需要确保数据键名的正确映射。关键修改点在于RepackTransform部分：

_transforms.RepackTransform(
    {
        "observation.images.image": "image",
        "observation.images.wrist_image": "wrist_image",
        "observation.state": "state",
        "actions": "action",  # 将数据集中的"action"映射为模型需要的"actions"
        "prompt": "prompt",
    }
)

2. 输入输出处理逻辑

在UR5Inputs类中，需要正确处理状态和图像数据：

def __call__(self, data: dict) -> dict:
    state = transforms.pad_to_dim(data["observation.state"], self.action_dim)
    
    base_image = _parse_image(data["observation.images.test"])
    wrist_image = _parse_image(data["observation.images.hand_image"])
    
    inputs = {
        "state": state,
        "image": {
            "base_0_rgb": base_image,
            "left_wrist_0_rgb": wrist_image,
            "right_wrist_0_rgb": np.zeros_like(base_image),
        },
        # ...其他输入配置
    }
    
    if "actions" in data:
        actions = transforms.pad_to_dim(data["actions"], self.action_dim)
        inputs["actions"] = actions
        
    return inputs

技术要点解析

1. 数据维度对齐：使用pad_to_dim确保状态向量和动作向量的维度与模型期望的输入维度一致。

2. 图像数据处理：通过_parse_image函数统一图像数据的格式，确保无论是uint8还是float32类型都能被正确处理。

3. 动作序列处理：在输出层明确指定只返回前7个动作维度(6个自由度+1个夹爪控制)，符合UR5机器人的实际控制需求。

最佳实践建议

1. 在配置数据映射时，建议先打印原始数据集的列名，确保键名映射的准确性。

2. 对于机器人控制任务，建议在数据预处理阶段就完成绝对动作到相对动作的转换，这可以通过DeltaActions变换实现。

3. 当使用多相机系统时，需要明确指定每个相机的图像数据及其掩码，如示例中对右腕相机的处理方式。

通过以上配置调整和技术实现，可以有效解决UR5机器人在OpenPI项目中微调时的数据配置问题，为后续的模型训练和部署奠定基础。