#!/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))