智元开源 AGIBOT WORLD 主题二：不只是成功轨迹，物理交互的完整数据长什么样

# 安装核心依赖
pip install huggingface_hub h5py numpy matplotlib

# 下载数据集（以 Hugging Face 为例，实际 repo 地址见官方公告）
huggingface-cli download Agibot/AGIBOT_WORLD_2026_RichInteraction \
  --repo-type dataset \
  --local-dir ./agibot_world_rich_interaction \
  --local-dir-use-symlinks False

import h5py
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path

data_dir = Path("./agibot_world_rich_interaction")

# 找到一条轨迹文件（实际目录结构以官方为准）
traj_files = sorted(data_dir.glob("**/*.h5"))
if not traj_files:
    print("未找到 HDF5 文件，请确认下载路径和目录结构")
    exit(1)

traj_path = traj_files[0]
print(f"解析轨迹: {traj_path.name}")

with h5py.File(traj_path, "r") as f:
    # 打印顶层结构，了解可用字段
    print("数据集顶层键:", list(f.keys()))

    # 典型结构：observations, actions, metadata
    # 力/力矩数据通常在 observations 下
    obs = f["observations"]

    # 提取末端力矩信号（字段名需对照官方文档）
    # 常见命名: force_torque, ft_sensor, wrist_ft 等
    ft_key = None
    for key in obs.keys():
        if "force" in key.lower() or "ft" in key.lower() or "torque" in key.lower():
            ft_key = key
            break

    if ft_key:
        ft_data = np.array(obs[ft_key])  # shape: (T, 6) — [Fx, Fy, Fz, Tx, Ty, Tz]
        print(f"力矩数据形状: {ft_data.shape}")
        print(f"力矩范围: Fx[{ft_data[:,0].min():.2f}, {ft_data[:,0].max():.2f}] N")

        # 可视化接触力变化
        time_steps = np.arange(ft_data.shape[0])
        fig, axes = plt.subplots(2, 1, figsize=(10, 6), sharex=True)

        axes[0].plot(time_steps, ft_data[:, 0], label="Fx", alpha=0.8)
        axes[0].plot(time_steps, ft_data[:, 1], label="Fy", alpha=0.8)
        axes[0].plot(time_steps, ft_data[:, 2], label="Fz", alpha=0.8)
        axes[0].set_ylabel("Force (N)")
        axes[0].legend()
        axes[0].set_title("末端接触力 — 轨迹中的物理对话")

        axes[1].plot(time_steps, ft_data[:, 3], label="Tx", alpha=0.8)
        axes[1].plot(time_steps, ft_data[:, 4], label="Ty", alpha=0.8)
        axes[1].plot(time_steps, ft_data[:, 5], label="Tz", alpha=0.8)
        axes[1].set_ylabel("Torque (N·m)")
        axes[1].set_xlabel("Time step")
        axes[1].legend()

        plt.tight_layout()
        plt.savefig("force_torque_profile.png", dpi=150)
        print("力矩曲线已保存到 force_torque_profile.png")
    else:
        print("未找到力矩字段，可用观测键:", list(obs.keys()))

    # 检查任务结果标注
    if "metadata" in f:
        meta = f["metadata"]
        if "result" in meta.attrs:
            print(f"任务结果: {meta.attrs['result']}")  # success / failure 类型
        if "failure_type" in meta.attrs:
            print(f"失败类型: {meta.attrs['failure_type']}")

import json

success_trajs = []
failure_trajs = []
failure_types = {}

for traj_path in traj_files[:200]:  # 先处理前 200 条做探索
    with h5py.File(traj_path, "r") as f:
        meta = f.get("metadata", None)
        if meta is None:
            continue

        result = meta.attrs.get("result", "unknown")
        if result == "success":
            success_trajs.append(traj_path)
        elif result == "failure":
            failure_trajs.append(traj_path)
            ft = meta.attrs.get("failure_type", "unclassified")
            failure_types[ft] = failure_types.get(ft, 0) + 1

print(f"成功轨迹: {len(success_trajs)}")
print(f"失败轨迹: {len(failure_trajs)}")
print(f"失败类型分布: {json.dumps(failure_types, indent=2)}")

# 关键洞察：失败轨迹的比例和类型分布
# 直接决定了你的策略修正模块需要覆盖哪些场景
ratio = len(failure_trajs) / max(len(success_trajs) + len(failure_trajs), 1)
print(f"失败占比: {ratio:.1%} — 如果低于 20%，数据多样性可能仍不足")