groot/vla/data/dataset/robotwin.py

#!/usr/bin/env python3
"""DreamZero RoboTwin Dataset - reads qpos+videos+metas, outputs DreamTransform format."""
import os, random, json, logging
from pathlib import Path
import numpy as np
import torch
import torch.utils.data as data
import cv2

from groot.vla.model.dreamzero.transform.dreamzero_cotrain import DreamTransform
from groot.vla.data.schema.lerobot import (
    DatasetMetadata, DatasetStatistics, DatasetStatisticalValues,
    DatasetModalities, VideoMetadata, StateActionMetadata,
)
from groot.vla.data.schema.embodiment_tags import EmbodimentTag

logger = logging.getLogger(__name__)

# Per-dimension q01/q99 for RoboTwin velocity actions (14-dim)
# Computed across 8349 episodes
ACTION_Q01 = np.array([-0.002055, -0.003258, -0.002083, -0.002576, -0.002143, -0.004001, -0.006708, 0.0, -0.006258, -0.010985, -0.009753, -0.008991, -0.006935, -0.011197], dtype=np.float32)
ACTION_Q99 = np.array([0.004110, 0.004289, 0.002321, 0.006814, 0.002750, 0.002597, 0.006232, 0.0, 0.006063, 0.014186, 0.010926, 0.010506, 0.006669, 0.010019], dtype=np.float32)
STATE_Q01 = np.array([-0.2]*7 + [0.0]*7, dtype=np.float32)
STATE_Q99 = np.array([0.2]*7 + [0.0]*7, dtype=np.float32)

def _normalize_q99(x, q01, q99, eps=1e-8):
    return np.clip(2.0 * (x - q01) / (q99 - q01 + eps) - 1.0, -1.0, 1.0)


class RobotWinDataset(data.Dataset):
    """RoboTwin dataset for DreamZero. Reads converted qpos+videos+metas format."""

    def __init__(self, dataset_dir, num_frames=12, action_horizon=12,
                 state_horizon=1, num_views=1, video_height=160, video_width=320,
                 max_episodes=None, seed=42, **kwargs):
        self.dataset_dir = Path(dataset_dir)
        self.num_frames = num_frames
        self.action_horizon = action_horizon
        self.state_horizon = state_horizon
        self.num_views = num_views
        self.video_height = video_height
        self.video_width = video_width
        self.seed = seed
        self.max_episodes = max_episodes
        self.rng = random.Random(seed)

        # Scan task directories
        self.tasks = sorted([d for d in self.dataset_dir.iterdir() if d.is_dir()])
        if not self.tasks:
            raise FileNotFoundError("No task dirs in " + str(dataset_dir))

        # Build episode list: (task_dir, ep_idx, n_video_frames, n_state_frames)
        self.episodes = []
        for task_dir in self.tasks:
            video_dir = task_dir / "videos"
            qpos_dir = task_dir / "qpos"
            meta_dir = task_dir / "metas"
            if not video_dir.exists() or not qpos_dir.exists():
                continue

            video_files = sorted(video_dir.glob("episode*.mp4"))
            qpos_files = sorted(qpos_dir.glob("episode*.pt"))

            for vf in video_files:
                ep_name = vf.stem  # episode0, episode1, etc.
                qf = qpos_dir / (ep_name + ".pt")
                if not qf.exists():
                    continue
                # Get frame count via cv2
                cap = cv2.VideoCapture(str(vf))
                n_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
                cap.release()
                if n_frames < self.num_frames + self.action_horizon:
                    continue
                # Load qpos to get state count
                qpos_data = torch.load(str(qf))
                n_state = qpos_data.shape[0]
                if n_state < self.num_frames + self.action_horizon:
                    continue
                self.episodes.append((str(task_dir), str(vf), str(qf), n_frames, n_state, ep_name))

        if not self.episodes:
            raise FileNotFoundError("No valid episodes found (need >= {} frames)".format(
                self.num_frames + self.action_horizon))

        if self.max_episodes:
            self.rng.shuffle(self.episodes)
            self.episodes = self.episodes[:self.max_episodes]

        self._build_metadata()

        self.transform = DreamTransform(
            default_instruction="Perform the default behavior.",
            language_dropout_prob=0.0, always_use_default_instruction=False,
            max_state_dim=44, max_action_dim=32, max_length=512,
            state_horizon=self.state_horizon, action_horizon=self.action_horizon,
            num_views=self.num_views,
            embodiment_tag_mapping={"oxe_droid": 17},
            tokenizer_path=kwargs.get("tokenizer_path", "/root/autodl-tmp/checkpoints/umt5-xxl"),
        )
        self.transform.set_metadata(self.merged_metadata["oxe_droid"])
        self.transform.train()

        n_ep = len(self.episodes)
        logger.info("RobotWinDataset: {} episodes across {} tasks".format(n_ep, len(self.tasks)))

    def _build_metadata(self):
        ds = DatasetStatisticalValues(
            max=np.ones(1), min=np.zeros(1), mean=np.zeros(1),
            std=np.ones(1), q01=np.zeros(1), q99=np.ones(1))
        self.merged_metadata = {
            "oxe_droid": DatasetMetadata(
                statistics=DatasetStatistics(
                    state={"joint_position": ds}, action={"joint_position": ds}),
                modalities=DatasetModalities(
                    video={"image": VideoMetadata(
                        resolution=(self.video_width, self.video_height), channels=3, fps=30)},
                    state={"joint_position": StateActionMetadata(
                        absolute=True, shape=(14,), continuous=True)},
                    action={"joint_position": StateActionMetadata(
                        absolute=True, shape=(14,), continuous=True)},
                ),
                embodiment_tag=EmbodimentTag.OXE_DROID,
            )
        }

    def reset_seed(self, new_seed):
        self.seed = new_seed
        self.rng = random.Random(new_seed)

    def __len__(self):
        return max(len(self.episodes) * 10, 1)

    def _load_task_desc(self, task_dir, ep_name):
        """Load task description from metas directory."""
        meta_dir = task_dir / "metas"
        if not meta_dir.exists():
            return "Manipulate the object on the table"
        meta_files = sorted(meta_dir.glob("*.txt"))
        if not meta_files:
            return "Manipulate the object on the table"
        # Extract episode index
        ep_idx = int(ep_name.replace("episode", ""))
        # Try to find the matching meta file
        if ep_idx < len(meta_files):
            with open(meta_files[ep_idx]) as f:
                return f.read().strip()
        return "Manipulate the object on the table"

    def _compute_actions(self, qpos_data, t, video_fps=30):
        """Compute actions from state differences, matching video frame rate."""
        # qpos_data: [T_state, 14], sampled at 10Hz (from original 100Hz decimated)
        # Video at 30fps
        # We need 1 action per video frame
        # If qpos has more timesteps than video frames, subsample qpos
        # If qpos has fewer timesteps, interpolate

        # For simplicity: each video frame takes the action at the aligned state index
        # action[i] = qpos[t + i + 1] - qpos[t + i]  (velocity)
        # But we need to map from video frame index to state index

        # Qpos is typically at ~10Hz, video at ~30Hz
        # So ratio is state_per_video_frame = qpos_len / video_fps_per_segment
        # For simplicity with num_frames=12 and action_horizon=12:
        # just use state differences from qpos

        actions = np.zeros((self.action_horizon, 14), dtype=np.float32)
        for i in range(self.action_horizon):
            fi = min(t + i + 1, qpos_data.shape[0] - 1)
            si = min(t + i, qpos_data.shape[0] - 1)
            actions[i] = qpos_data[fi] - qpos_data[si]
        return actions

    def __getitem__(self, idx):
        task_dir, video_path, qpos_path, n_video_frames, n_qpos_frames, ep_name = \
            self.rng.choice(self.episodes)

        # Random start, ensuring enough frames for video + actions
        max_start = max(0, n_video_frames - self.num_frames - self.action_horizon)
        t = self.rng.randint(0, max_start) if max_start > 0 else 0

        # Read video frames
        cap = cv2.VideoCapture(video_path)
        frames = []
        for i in range(self.num_frames):
            fi = min(t + i, n_video_frames - 1)
            cap.set(cv2.CAP_PROP_POS_FRAMES, fi)
            ret, frame = cap.read()
            if not ret:
                frame = frames[-1] if frames else np.zeros((240, 320, 3), dtype=np.uint8)
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frames.append(frame)
        cap.release()

        video = np.stack(frames, axis=0)  # [T, H, W, C]
        video = video[:, np.newaxis, :, :, :]  # [T, V=1, H, W, C]

        # Resize if needed
        if video.shape[2] != self.video_height or video.shape[3] != self.video_width:
            resized = np.zeros((self.num_frames, 1, self.video_height, self.video_width, 3), dtype=np.uint8)
            for ti in range(self.num_frames):
                resized[ti, 0] = cv2.resize(video[ti, 0], (self.video_width, self.video_height))
            video = resized

        # Load qpos and compute state/action
        qpos_data = torch.load(qpos_path).numpy()  # [T_qpos, 14]

        # Map video frame index to qpos index
        # qpos sampling rate = n_qpos_frames / video_duration_in_frames (approx)
        video_to_qpos_ratio = n_qpos_frames / max(n_video_frames, 1)
        qpos_start = int(t * video_to_qpos_ratio)
        qpos_start = min(qpos_start, n_qpos_frames - 13)  # ensure enough room

        # State: use the state at the aligned timestep
        state_val = qpos_data[min(qpos_start, n_qpos_frames - 1)]
        state = np.tile(state_val, (self.state_horizon, 1))

        # Actions: velocity from consecutive qpos steps
        actions = np.zeros((self.action_horizon, 14), dtype=np.float32)
        for i in range(self.action_horizon):
            fi = min(qpos_start + i + 1, n_qpos_frames - 1)
            si = min(qpos_start + i, n_qpos_frames - 1)
            actions[i] = qpos_data[fi] - qpos_data[si]

        # Task description
        task_name = self._load_task_desc(Path(task_dir), ep_name)

        
        # Normalize state and action to [-1, 1]
        state = _normalize_q99(state, STATE_Q01, STATE_Q99)
        actions = _normalize_q99(actions, ACTION_Q01, ACTION_Q99)

        raw = {
        "video": video,
        "state": state,
            "action": actions,
            "annotation.human.action.task_description": task_name,
        }
        return dict(self.transform(raw))