training_code/robot_video_dataset.py

# Author: Rui Heng Yang

import hashlib
import os
from pathlib import Path
from typing import Any, Optional
import numpy as np
import traceback
import torch
import torchvision.transforms.functional as transforms_F

from omegaconf import DictConfig, OmegaConf

from hydra.utils import instantiate
from .base_lerobot_dataset import BaseLerobotDataset
from ..video_latent_cache import VideoLatentCache, build_video_preprocess_spec
from .utils.normalizer import save_dataset_stats_to_json, load_dataset_stats_from_json
from ..dataset_utils import ResizeSmallestSideAspectPreserving, CenterCrop, Normalize
from fastwam.utils.logging_config import get_logger
from fastwam.utils import misc, pytorch_utils
logger = get_logger(__name__)


DEFAULT_PROMPT = "A video recorded from a robot's point of view executing the following instruction: {task}"


def _is_main_process_without_init() -> bool:
    return pytorch_utils._resolve_global_rank() == 0

class RobotVideoDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        dataset_dirs,
        shape_meta,
        num_frames=33,
        video_size=[384, 640],
        camera_key=None,
        processor=None,
        text_embedding_cache_dir=None,
        context_len=128,
        pretrained_norm_stats=None,
        val_set_proportion=0.05,
        is_training_set=False,
        episode_indices=None,
        episode_indices_path=None,
        episode_indices_key=None,
        global_sample_stride=1,
        action_video_freq_ratio: int = 1,
        skip_padding_as_possible: bool = False,
        max_padding_retry: int = 3,
        concat_multi_camera: str = "horizontal", # "horizontal", "vertical", "robotwin", or None
        video_backend: Optional[str] = None,
        video_latent_cache_dir: Optional[str] = None,
        video_latent_cache_vae_identity: Optional[dict[str, str]] = None,
        strict_getitem: bool = False,
        precompute_video_only: bool = False,
        override_instruction: Optional[str] = None, # whether to hardcode a specific instruction for all samples, for debugging
        eef_anchor_calibration_path: Optional[str] = None,
        eef_anchor_raw_resolution: int = 512,
    ):
        if video_latent_cache_dir is not None and not strict_getitem:
            raise ValueError(
                "`strict_getitem=true` is required with `video_latent_cache_dir` "
                "to prevent latent/action sample misalignment."
            )
        if video_latent_cache_dir is not None and video_backend is None:
            raise ValueError(
                "`video_backend` must be explicitly pinned when using a video latent cache."
            )
        if video_latent_cache_dir is not None and video_latent_cache_vae_identity is None:
            raise ValueError(
                "`video_latent_cache_vae_identity` is required with a video latent cache."
            )
        if strict_getitem and skip_padding_as_possible:
            raise ValueError(
                "`skip_padding_as_possible` must be false in strict/cache mode because "
                "padding retries substitute a random global sample index."
            )
        if precompute_video_only and (not strict_getitem or video_latent_cache_dir is not None):
            raise ValueError(
                "`precompute_video_only` requires strict_getitem=true and no latent cache."
            )

        if (num_frames - 1) % action_video_freq_ratio != 0:
            raise ValueError(
                "num_frames-1 must be divisible by action_video_freq_ratio, got "
                f"{num_frames - 1} and {action_video_freq_ratio}"
            )
        if ((num_frames - 1) // action_video_freq_ratio) % 4 != 0:
            raise ValueError(
                "video transitions must be divisible by 4 for tokenization, got "
                f"{(num_frames - 1) // action_video_freq_ratio}"
            )
        video_sample_indices = list(range(0, num_frames, action_video_freq_ratio))

        shape_meta_container = OmegaConf.to_container(shape_meta, resolve=True)
        processor_config: dict[str, Any] | None = None
        if isinstance(processor, DictConfig):
            processor_config = OmegaConf.to_container(processor, resolve=True)

        self.lerobot_dataset = BaseLerobotDataset(
            dataset_dirs=dataset_dirs,
            shape_meta=shape_meta_container,
            obs_size=num_frames,
            action_size=num_frames - 1,
            val_set_proportion=val_set_proportion,
            is_training_set=is_training_set,
            episode_indices=episode_indices,
            episode_indices_path=episode_indices_path,
            episode_indices_key=episode_indices_key,
            global_sample_stride=global_sample_stride,
            video_backend=video_backend,
            strict_getitem=strict_getitem,
            image_only=precompute_video_only,
            image_sample_indices=(video_sample_indices if precompute_video_only else None),
        )
    
        self.num_frames = num_frames
        self.action_video_freq_ratio = action_video_freq_ratio
        
        self.video_sample_indices = video_sample_indices

        self.camera_key = camera_key
        self.video_latent_cache_dir = video_latent_cache_dir
        self.strict_getitem = bool(strict_getitem)
        self.precompute_video_only = bool(precompute_video_only)
        self.lerobot_dataset._set_return_images(video_latent_cache_dir is None)

        self.video_size = video_size
        self.text_embedding_cache_dir = text_embedding_cache_dir
        self.context_len = context_len
        self.skip_padding_as_possible = skip_padding_as_possible
        self.max_padding_retry = max_padding_retry
        self.concat_multi_camera = concat_multi_camera
        self.override_instruction = override_instruction

        resolved_episode_indices_path = None
        episode_indices_path_sha256 = None
        if episode_indices_path is not None:
            resolved_path = Path(episode_indices_path).expanduser().resolve()
            if not resolved_path.is_file():
                raise FileNotFoundError(f"Episode split file not found: {resolved_path}")
            resolved_episode_indices_path = str(resolved_path)
            episode_indices_path_sha256 = hashlib.sha256(resolved_path.read_bytes()).hexdigest()

        preprocess_spec = build_video_preprocess_spec(
            dataset_dirs=dataset_dirs,
            shape_meta=shape_meta_container,
            processor_config=processor_config,
            num_frames=num_frames,
            action_video_freq_ratio=action_video_freq_ratio,
            video_size=video_size,
            concat_multi_camera=concat_multi_camera,
            video_backend=str(video_backend),
            global_sample_stride=global_sample_stride,
            sample_selection={
                "val_set_proportion": float(val_set_proportion),
                "is_training_set": bool(is_training_set),
                "episode_indices": (
                    None if episode_indices is None else [int(value) for value in episode_indices]
                ),
                "episode_indices_path": (
                    resolved_episode_indices_path
                ),
                "episode_indices_path_sha256": episode_indices_path_sha256,
                "episode_indices_key": episode_indices_key,
            },
        )
        self.video_preprocess_spec = preprocess_spec
        self.video_latent_cache = None
        if video_latent_cache_dir is not None:
            expected_vae_identity = OmegaConf.to_container(
                video_latent_cache_vae_identity, resolve=True
            ) if isinstance(video_latent_cache_vae_identity, DictConfig) else dict(
                video_latent_cache_vae_identity
            )
            self.video_latent_cache = VideoLatentCache(
                video_latent_cache_dir,
                expected_sample_count=len(self.lerobot_dataset),
                expected_preprocess_spec=preprocess_spec,
                expected_vae_identity=expected_vae_identity,
            )

        self.resize_transform = ResizeSmallestSideAspectPreserving(
            args={"img_w": self.video_size[1], "img_h": self.video_size[0]},
        )
        self.crop_transform = CenterCrop(
            args={"img_w": self.video_size[1], "img_h": self.video_size[0]},
        )
        self.normalize_transform = Normalize(
            args={"mean": 0.5, "std": 0.5},
        )
        if processor is not None:
            if isinstance(processor, DictConfig):
                processor = instantiate(processor)
            if not pretrained_norm_stats:
                if not is_training_set:
                    raise ValueError("pretrained_norm_stats must be provided for validation/test sets since we don't want to calculate stats on them.")
                dist_initialized = (
                    torch.distributed.is_available() and torch.distributed.is_initialized()
                )
                if _is_main_process_without_init() or not dist_initialized:
                    logger.info("Calculating dataset stats for normalization...")
                    dataset_stats = self.lerobot_dataset.get_dataset_stats(processor)
                    work_dir = misc.get_work_dir()
                    if _is_main_process_without_init():
                        save_dataset_stats_to_json(dataset_stats, os.path.join(work_dir, "dataset_stats.json"))
                else:
                    dataset_stats = None
                if dist_initialized:
                    obj_list = [dataset_stats]
                    torch.distributed.broadcast_object_list(obj_list, src=0)
                    dataset_stats = obj_list[0]
            else:
                dataset_stats = load_dataset_stats_from_json(pretrained_norm_stats)
                logger.info(f"Using dataset stats: {pretrained_norm_stats}")
                if _is_main_process_without_init():
                    work_dir = misc.get_work_dir()
                    save_dataset_stats_to_json(dataset_stats, os.path.join(work_dir, "dataset_stats.json"))

            processor.set_normalizer_from_stats(dataset_stats)
            self.lerobot_dataset.set_processor(processor)
        
        # EEF-relative camera RoPE: resolve every episode's anchors once, here,
        # and hold them in RAM. No sidecar exists (plan Section 15). Disabled
        # unless a calibration path is configured, so every other mode keeps its
        # existing sample keys and behavior byte for byte.
        self.eef_anchor_index = None
        if eef_anchor_calibration_path:
            from fastwam.datasets.eef_anchors import EpisodeAnchorIndex
            from fastwam.geometry import EEFProjector, load_calibration

            calibration = load_calibration(eef_anchor_calibration_path)
            self.eef_anchor_index = EpisodeAnchorIndex.build(
                [str(d) for d in dataset_dirs],
                calibration_path=eef_anchor_calibration_path,
                projector=EEFProjector(calibration, int(eef_anchor_raw_resolution)),
            )

    def __len__(self):
        return len(self.lerobot_dataset)

    def _get(self, idx):
        sample_idx = idx
        sample = None
        for attempt in range(self.max_padding_retry + 1):
            sample = self.lerobot_dataset[sample_idx]

            if not self.skip_padding_as_possible:
                break

            action_is_pad = sample["action_is_pad"]
            image_is_pad = sample["image_is_pad"]
            proprio_is_pad = sample["proprio_is_pad"]
            has_pad = False
            if bool(action_is_pad.any().item()):
                has_pad = True
            if bool(image_is_pad.any().item()):
                has_pad = True
            if bool(proprio_is_pad.any().item()):
                has_pad = True

            if not has_pad or attempt >= self.max_padding_retry:
                break

            sample_idx = np.random.randint(len(self.lerobot_dataset))
        
        image_is_pad = sample["image_is_pad"]

        if not self.precompute_video_only:
            image_is_pad = image_is_pad[self.video_sample_indices]
        video = None
        video_latents = None
        if self.video_latent_cache is None:
            video = sample["pixel_values"]  # [T, C, H, W] or [num_cameras, T, C, H, W]
            num_cameras = 1
            if video.ndim == 5:
                if not self.precompute_video_only:
                    video = video[:, self.video_sample_indices, :, :, :]
                num_cameras, T_video, C, H, W = video.shape
            else:
                assert video.ndim == 4, f"Expected video to have shape [T, C, H, W], but got {video.shape}"
                if not self.precompute_video_only:
                    video = video[self.video_sample_indices, :, :, :]
                T_video, C, H, W = video.shape

            video = video.view(num_cameras, T_video, C, H, W)  # [num_cameras, T_video, C, H, W]
            if self.concat_multi_camera == "robotwin":
                if num_cameras != 3:
                    raise ValueError(
                        f"`concat_multi_camera='robotwin'` requires exactly 3 cameras, got {num_cameras}"
                    )
                cam_top = transforms_F.resize(
                    video[0],
                    size=[256, 320],
                    interpolation=transforms_F.InterpolationMode.BILINEAR,
                    antialias=True,
                )  # [T_video, C, 256, 320]
                cam_left = transforms_F.resize(
                    video[1],
                    size=[128, 160],
                    interpolation=transforms_F.InterpolationMode.BILINEAR,
                    antialias=True,
                )  # [T_video, C, 128, 160]
                cam_right = transforms_F.resize(
                    video[2],
                    size=[128, 160],
                    interpolation=transforms_F.InterpolationMode.BILINEAR,
                    antialias=True,
                )  # [T_video, C, 128, 160]
                bottom = torch.cat([cam_left, cam_right], dim=-1)  # [T_video, C, 128, 320]
                video = torch.cat([cam_top, bottom], dim=-2)  # [T_video, C, 384, 320]
            elif num_cameras > 1:
                if self.concat_multi_camera == "horizontal":
                    video = torch.cat([video[i] for i in range(num_cameras)], dim=-1)  # [T_video, C, H, num_cameras*W]
                elif self.concat_multi_camera == "vertical":
                    video = torch.cat([video[i] for i in range(num_cameras)], dim=-2)  # [T_video, C, num_cameras*H, W]
                else:
                    raise ValueError(
                        f"Invalid concat_multi_camera: {self.concat_multi_camera}. "
                        "Expected one of: horizontal, vertical, robotwin."
                    )
            else:
                video = video.squeeze(0)  # [T_video, C, H, W]

            # final resize and normalization
            video = self.resize_transform(video)
            video = self.crop_transform(video)
            video = self.normalize_transform(video)  # [T_video, C, H, W]

            video = video.permute(1, 0, 2, 3) # [C, T_video, H, W], range [-1, 1]
        else:
            returned_index = int(sample["idx"])
            video_latents = self.video_latent_cache[returned_index]

        if self.precompute_video_only:
            if video is None:
                raise RuntimeError("precompute_video_only unexpectedly resolved cached latents.")
            return {"video": video}

        # Proxy (from lerobot): 
        #   action: [num_frames-1, action_dim] # start from t0, except the last frame
        #   proprio: [num_frames, proprio_dim] # start from t0 to the last frame, aligned with video frames
        action = sample["action"] # [T-1, action_dim]
        proprio = sample["proprio"][:-1, :] # [T-1, state_dim]， to align with action
        sampled_video_frames = len(self.video_sample_indices)
        if sampled_video_frames <= 1:
            raise ValueError(f"At least two sampled video frames are required, got {sampled_video_frames}")
        if action.shape[0] % (sampled_video_frames - 1) != 0:
            raise ValueError(
                f"`action` horizon must be divisible by `video` transitions, got {action.shape[0]} and {sampled_video_frames - 1}"
            )

        task = sample["instruction"]
        
        # FIXME
        if self.override_instruction is not None:
            task = self.override_instruction
        instruction = DEFAULT_PROMPT.format(task=task)

        context, context_mask = self._get_cached_text_context(instruction)
        # NOTE: to keep consistent with wan2.2's behavior
        context[~context_mask] = 0.0
        context_mask = torch.ones_like(context_mask)
        
        data = {
            "action": action,
            "proprio": proprio,
            "prompt": instruction,
            "context": context,
            "context_mask": context_mask,
            "image_is_pad": image_is_pad,
            "action_is_pad": sample["action_is_pad"],
            "proprio_is_pad": sample["proprio_is_pad"],
        }
        if video is not None:
            data["video"] = video
        else:
            data["video_latents"] = video_latents
            data["source_video_shape"] = torch.tensor(
                [3, sampled_video_frames, int(self.video_size[0]), int(self.video_size[1])],
                dtype=torch.int64,
            )
        if self.eef_anchor_index is not None:
            # Key off the RETURNED sample, never the requested idx. Both the
            # padding retry above and __getitem__'s exception handler resample
            # to a random global index, so an idx-keyed lookup would pair these
            # frames with another episode's gripper position -- silently.
            # Indexed, not `.get(..., 0)`: a missing `dataset_index` would
            # silently key every directory of a concatenated MultiLeRobotDataset
            # to directory 0, serving one suite's geometry for another's frames.
            data["eef_anchor_token"] = torch.from_numpy(
                self.eef_anchor_index.lookup(
                    dataset_index=int(sample["dataset_index"]),
                    episode_index=int(sample["episode_index"]),
                    frame_index=int(sample["frame_index"]),
                ).copy()
            )
        return data

    def _get_cached_text_context(self, prompt: str):
        if self.text_embedding_cache_dir is None:
            raise ValueError("text_embedding_cache_dir is not set.")
        cache_dir = self.text_embedding_cache_dir
        os.makedirs(cache_dir, exist_ok=True)
        hashed = hashlib.sha256(prompt.encode("utf-8")).hexdigest()
        cache_path = os.path.join(cache_dir, f"{hashed}.t5_len{self.context_len}.wan22ti2v5b.pt")
        if not os.path.exists(cache_path):
            raise FileNotFoundError(
                f"Missing text embedding cache: {cache_path}. "
                "Run scripts/precompute_text_embeds.py first."
            )
        payload = torch.load(cache_path, map_location="cpu")
        context = payload["context"]
        context_mask = payload["mask"].bool()
        if context.ndim != 2:
            raise ValueError(
                f"Cached `context` must be 2D [L, D], got shape {tuple(context.shape)} in {cache_path}"
            )
        if context_mask.ndim != 1:
            raise ValueError(
                f"Cached `mask` must be 1D [L], got shape {tuple(context_mask.shape)} in {cache_path}"
            )
        if context.shape[0] != self.context_len:
            raise ValueError(
                f"Cached context_len mismatch: expected {self.context_len}, got {context.shape[0]} in {cache_path}"
            )
        if context_mask.shape[0] != self.context_len:
            raise ValueError(
                f"Cached mask_len mismatch: expected {self.context_len}, got {context_mask.shape[0]} in {cache_path}"
            )

        return context, context_mask

    def __getitem__(self, idx):
        if self.strict_getitem:
            return self._get(idx)
        try:
            data = self._get(idx)
        except Exception as e:
            print(f"Error processing sample idx {idx}: {e}. Returning a random sample instead.")
            # trace back
            print(traceback.format_exc())
            random_idx = np.random.randint(len(self))
            data = self._get(random_idx)
        return data