Upload lerobot_ee6d_piper_real.py with huggingface_hub

lerobot_ee6d_piper_real.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+import numpy as np, torch, random
+from mmengine import fileio
+from scipy.interpolate import interp1d
+from ..utils import read_video_to_frames, read_parquet, quat_to_rotate6d, euler_to_rotate6d
+from PIL import Image
+from .base import DomainHandler
+import pandas as pd
+import json
+from typing import Iterable, Any
+from tqdm import tqdm
+import pickle
+
+DEBUG_TRIGGER = True 
+
+class PiperLeRobotEE6DHandler(DomainHandler):
+
+    def __init__(self, meta: dict, num_views: int):
+        super().__init__(meta, num_views)
+        stats_path = self.meta.get("action_stats_path", None)
+        self.stats_path = stats_path
+        if stats_path is not None and fileio.exists(stats_path):
+            with open(stats_path, 'rb') as f:
+                self.action_stats = pickle.load(f)
+                self.mean = self.action_stats['mean']
+                self.std = self.action_stats['std']
+                self.min = self.action_stats['min']
+                self.max = self.action_stats['max']
+
+
+    def index_candidates(self, T_left: int, training: bool) -> Iterable[int]:
+        return range(0, max(0, T_left - 10))
+    
+
+    def iter_episode(self, traj_idx: int, *, num_actions: int, training: bool,
+                     image_aug, lang_aug_map: dict | None, **kwargs):
+        datapath = self.meta["datalist"][traj_idx]
+        video_base_path = self.meta["video_base_path"]
+        tasks_path = self.meta.get("tasks_path", None)
+        with open(tasks_path, 'r') as f:
+            tasks = f.readlines()
+
+        df = pd.read_parquet(datapath)
+        task_index = df['task_index'].iloc[0]
+        cur_task = json.loads(tasks[task_index])["task"]
+        raw_action = np.stack(df['action_eepose'])
+
+        # 2. 提取左臂 (Arm 1) 的组件并转换
+        # 假设格式: [pos(3), euler(3), grip(1), pos(3), euler(3), grip(1)]
+        pos1 = raw_action[:, :3]
+        rot1_6d = euler_to_rotate6d(raw_action[:, 3:6], 'xyz') # 只传欧拉角部分 (N, 3)
+        grip1 = raw_action[:, 6:7]
+
+        # 3. 提取右臂 (Arm 2) 的组件并转换
+        pos2 = raw_action[:, 7:10]
+        rot2_6d = euler_to_rotate6d(raw_action[:, 10:13], 'xyz') # 只传欧拉角部分 (N, 3)
+        grip2 = raw_action[:, 13:14]
+
+        # 4. 重新拼接成 20 维的向量 (3+6+1 + 3+6+1)
+        action_ee6d = np.concatenate([pos1, rot1_6d, grip1, pos2, rot2_6d, grip2], axis=-1)
+        image_observations = ['observation.images.cam_high', 'observation.images.cam_left_wrist', 'observation.images.cam_right_wrist']
+        vid_paths = []
+        for image_type in image_observations:
+            image_type_base_path = video_base_path + image_type + "/"
+            vid_path = image_type_base_path + f"episode_{traj_idx:06d}.mp4"
+            vid_paths.append(vid_path)
+        
+        images = [read_video_to_frames(p) for p in vid_paths]
+        image_mask = torch.ones(self.num_views, dtype=torch.bool)
+        image_mask[:len(images)] = True
+
+        # stack 之后 shape 应该是 [T, 20]
+        all_actions = np.stack(action_ee6d) 
+
+        # --- 切分左右臂 ---
+        dim_per_arm = 10
+        left = all_actions[:, :dim_per_arm]
+        right = all_actions[:, dim_per_arm:]
+
+        # 4. 时间插值 (处理帧率对齐)
+        freq = 30.0 # 请确认你的录制频率
+        qdur = 1.0  # 预测未来 1 秒
+        t = np.arange(left.shape[0], dtype=np.float64) / freq
+
+        lt = np.arange(left.shape[0], dtype=np.float64) / float(freq)
+        rt = np.arange(right.shape[0], dtype=np.float64) / float(freq)
+
+        # Candidate indices (optionally shuffled)
+        idxs = list(self.index_candidates(left.shape[0], training))
+        if training: random.shuffle(idxs)
+
+        # Interpolators; clamp to endpoints
+        L = interp1d(lt, left, axis=0, bounds_error=False, fill_value=(left[0], left[-1]))
+        R = interp1d(rt, right, axis=0, bounds_error=False, fill_value=(right[0], right[-1]))
+        ref = (lt + rt) / 2.0
+
+        V = min(self.num_views, len(images))
+        for idx in idxs:
+            imgs = []
+            for v in range(min(self.num_views, len(images))):
+                imgs.append(image_aug(Image.fromarray(images[v][idx])))
+            while len(imgs) < self.num_views: imgs.append(torch.zeros_like(imgs[0]))
+            image_input = torch.stack(imgs, 0)
+            cur = t[idx]
+            q = np.linspace(cur, min(cur + qdur, float(t.max())), num_actions + 1, dtype=np.float32)
+            lseq, rseq = torch.tensor(L(q)), torch.tensor(R(q))
+            if (lseq[1]-lseq[0]).abs().max() < 1e-5 and (rseq[1]-rseq[0]).abs().max() < 1e-5:continue
+            if lang_aug_map is not None and ins in lang_aug_map: ins = random.choice(lang_aug_map[ins])
+
+            trajectory = torch.cat([lseq, rseq], -1).float()
+            # if self.stats_path is not None and hasattr(self, 'min') and hasattr(self, 'max'):
+            #     min_val = torch.tensor(self.min)
+            #     max_val = torch.tensor(self.max)
+            #     # 归一化到 [-1, 1]
+            #     norm_action = 2 * (trajectory - min_val) / (max_val - min_val + 1e-8) - 1
+            #     trajectory = norm_action
+
+            # if training:
+            #     # 注入微量高斯噪声 (例如 0.001 级别)
+            #     # 注意不要给最后的 gripper 维度加噪声（如果是 0/1）
+            #     noise = torch.randn_like(trajectory) * 0.001
+            #     # 第10维是左臂 gripper，第20维是右臂 gripper
+            #     noise[:, 9] = 0.0
+            #     noise[:, 19] = 0.0
+            #     trajectory += noise
+
+            yield {
+                "language_instruction": cur_task,
+                "image_input": image_input,
+                "image_mask": image_mask,
+                "abs_trajectory": trajectory
+            }
+
+            # yield {
+            #     "language_instruction": cur_task,
+            #     "image_input": image_input,
+            #     "image_mask": image_mask,
+            #     "abs_trajectory": torch.cat([lseq, rseq], -1).float()
+            # }