code/src/datasets/droid_act.py

import os
import gc
import ctypes
import numpy as np
import torch
from torch.utils.data import IterableDataset, DataLoader
import tensorflow as tf
import tensorflow_datasets as tfds

tf.config.set_visible_devices([], 'GPU')

# Force glibc to return freed memory to the OS (Linux only)
try:
    _LIBC = ctypes.CDLL("libc.so.6")
    def _malloc_trim():
        _LIBC.malloc_trim(0)
except Exception:
    def _malloc_trim():
        pass

class DroidAct(IterableDataset):
    def __init__(
        self,
        droid_path,
        dataset_name='droid',
        length=None,
        history_len=15,
        future_len=15,
        full_sequence=False,
        input_modality="video",
        view_mode="single",
        load_future_image=False,
        future_image_mode="horizon",
        buffer_size=30000,
    ):
        super().__init__()
        self.droid_path = droid_path
        self.dataset_name = dataset_name
        self.length = length
        self.history_len = history_len
        self.future_len = future_len
        self.full_sequence = full_sequence
        self.input_modality = input_modality
        self.view_mode = view_mode
        self.load_future_image = load_future_image
        self.future_image_mode = future_image_mode
        self.buffer_size = buffer_size

    def __iter__(self):
        builder = tfds.builder_from_directory(builder_dir=self.droid_path)
        
        read_config = tfds.ReadConfig(shuffle_seed=42, shuffle_reshuffle_each_iteration=False)
        droid_ds = builder.as_dataset(split='train', shuffle_files=False, read_config=read_config)
        
        if self.length is not None:
            droid_ds = droid_ds.take(self.length)

        shuffle_buffer = []
        BUFFER_SIZE = self.buffer_size
        
        cam_key = 'exterior_image_1_left'
        wrist_key = 'wrist_image_left'

        if torch.distributed.is_available() and torch.distributed.is_initialized():
            rank = torch.distributed.get_rank()
            world_size = torch.distributed.get_world_size()
        else:
            rank = 0
            world_size = 1

        worker_info = torch.utils.data.get_worker_info()
        if worker_info is None:
            worker_id = 0
            num_workers = 1
        else:
            worker_id = worker_info.id
            num_workers = worker_info.num_workers

        total_shards = world_size * num_workers
        shard_index = rank * num_workers + worker_id
        ds_iterator = droid_ds.shard(num_shards=total_shards, index=shard_index)

        ds_iter = iter(ds_iterator)
        traj_id = -1
        while True:
            try:
                try:
                    traj_data = next(ds_iter)
                except StopIteration:
                    break
                except tf.errors.DataLossError as e:
                    traj_id += 1
                    print(f"[Warn] Skipping trajectory {traj_id}: TF DataLossError during iteration: {e}")
                    continue
                traj_id += 1
                
                traj_batch = next(iter(traj_data['steps'].batch(5000)))

                if traj_batch['reward'][-1].numpy() != 1:
                    del traj_batch
                    continue

                traj_len = traj_batch['action'].shape[0]
                
                images_np = traj_batch['observation'][cam_key].numpy()
                if images_np.dtype != np.uint8:
                    images_np = (images_np * 255).astype(np.uint8)

                wrist_np = None
                if self.view_mode == "multi":
                    obs = traj_batch['observation']
                    if wrist_key in obs:
                        wrist_np = obs[wrist_key].numpy()
                        if wrist_np.dtype != np.uint8:
                            wrist_np = (wrist_np * 255).astype(np.uint8)
                    else:
                        wrist_np = images_np.copy()
                    del obs

                # Process Proprioception: Cartesian + Gripper
                cart_pos = traj_batch['observation']['cartesian_position']
                
                # [Proprio Gripper]
                # Raw: 1D position in [0, 1]. 0 = Open, 1 = Closed.
                # Processed: Direct copy.
                # Result: Range [0, 1]. 0 = Open, 1 = Closed.
                gripper_pos = traj_batch['observation']['gripper_position']
                proprio_np = tf.concat([cart_pos, gripper_pos], axis=-1).numpy().astype(np.float32)

                # Process Actions: Cartesian Velocity + Gripper Command (Normalized)
                # DROID Actions: 6 (vel) + 1 (gripper)
                cart_vel = traj_batch['action_dict']['cartesian_velocity']
                cart_vel = tf.cast(cart_vel, tf.float32)
                cart_vel = tf.clip_by_value(cart_vel, -1.0, 1.0)
                
                # [Action Gripper]
                # Raw: 1D position in [0, 1]. 0 = Open, 1 = Closed.
                # Processed: Binarized to {-1, 1} based on 0.5 threshold.
                # Result: Range [-1, 1]. -1 = Open, 1 = Closed.
                grip_pos = traj_batch['action_dict']['gripper_position']
                grip_cmd = tf.where(grip_pos > 0.5, 1.0, -1.0) 
                grip_cmd = tf.cast(grip_cmd, tf.float32)

                actions_np = tf.concat([cart_vel, grip_cmd], axis=-1).numpy().astype(np.float32)
                
                instruction = traj_batch['language_instruction'][0].numpy().decode('utf-8')

                del traj_batch, cart_pos, gripper_pos, cart_vel, grip_pos, grip_cmd

                if self.full_sequence:
                    sample_indices = np.arange(traj_len)
                else:
                    num_samples = max(1, int(traj_len / (15 * 5)))
                    sample_indices = np.random.choice(traj_len, size=num_samples, replace=False)

                for t in sample_indices:
                    start_hist_obs = t - self.history_len + 1
                    hist_indices_obs = np.arange(start_hist_obs, t + 1)
                    hist_indices_obs = np.clip(hist_indices_obs, 0, traj_len - 1)
                    
                    start_hist_act = t - self.history_len
                    hist_indices_act = np.arange(start_hist_act, t)
                    
                    end_fut = t + self.future_len
                    fut_indices = np.arange(t, end_fut)

                    hist_imgs = images_np[hist_indices_obs]
                    hist_imgs_wrist = wrist_np[hist_indices_obs] if wrist_np is not None else None
                    hist_proprio = torch.from_numpy(proprio_np[hist_indices_obs])
                    
                    hist_actions = np.zeros((self.history_len, actions_np.shape[1]), dtype=np.float32)
                    valid_mask = hist_indices_act >= 0
                    if np.any(valid_mask):
                        valid_indices = hist_indices_act[valid_mask]
                        valid_indices = np.clip(valid_indices, 0, traj_len - 1)
                        hist_actions[valid_mask] = actions_np[valid_indices]
                    hist_actions = torch.from_numpy(hist_actions)
                    
                    fut_acts_np = np.zeros((self.future_len, actions_np.shape[1]), dtype=np.float32)
                    valid_mask_fut = fut_indices < traj_len
                    if np.any(valid_mask_fut):
                        valid_indices_fut = fut_indices[valid_mask_fut]
                        fut_acts_np[valid_mask_fut] = actions_np[valid_indices_fut]
                    fut_acts = torch.from_numpy(fut_acts_np)

                    sample = {
                        'proprioception': hist_proprio,
                        'history_actions': hist_actions,
                        'future_actions': fut_acts,
                        'instruction': instruction,
                    }

                    if self.load_future_image:
                        if self.future_image_mode == "last":
                            target_idx = traj_len - 1
                        else:
                            target_idx = min(t + self.future_len, traj_len - 1)
                        sample['future_image'] = images_np[target_idx].copy()

                    if self.input_modality == "video":
                        sample['video'] = hist_imgs
                        if self.view_mode == "multi":
                            sample['video_wrist'] = hist_imgs_wrist
                    elif self.input_modality == "image":
                        sample['image'] = images_np[t].copy()
                        if self.view_mode == "multi":
                            sample['image_wrist'] = wrist_np[t].copy() if wrist_np is not None else images_np[t].copy()
                    else:
                        raise ValueError(f"Unknown input_modality: {self.input_modality}")

                    shuffle_buffer.append(sample)
                    
                    if len(shuffle_buffer) >= BUFFER_SIZE:
                        idx = np.random.randint(len(shuffle_buffer))
                        shuffle_buffer[idx], shuffle_buffer[-1] = shuffle_buffer[-1], shuffle_buffer[idx]
                        yield shuffle_buffer.pop()

                del images_np, actions_np, proprio_np
                if wrist_np is not None:
                    del wrist_np

            except Exception as e:
                print(f"[Warn] Skipping trajectory {traj_id} due to error: {e}")
                continue
            finally:
                if traj_id % 50 == 0:
                    gc.collect()
                    _malloc_trim()
        
        np.random.shuffle(shuffle_buffer)
        for sample in shuffle_buffer:
            yield sample

def collate_fn(batch):
    return batch

if __name__ == "__main__":
    """
    Fast stats: count how many training samples DroidAct would yield.
    Now includes a progress bar (tqdm).
    """
    import argparse
    from tqdm import tqdm

    parser = argparse.ArgumentParser()
    parser.add_argument("--droid_path", type=str, default="/mnt/NTU_slab/draven/data/open_x_embodiment/droid/1.0.1")
    parser.add_argument("--split", type=str, default="train")
    parser.add_argument("--limit_traj", type=int, default=None)
    parser.add_argument("--full_sequence", action="store_true")
    args = parser.parse_args()

    builder = tfds.builder_from_directory(builder_dir=args.droid_path)
    read_config = tfds.ReadConfig(shuffle_seed=42, shuffle_reshuffle_each_iteration=False)
    ds = builder.as_dataset(split=args.split, shuffle_files=False, read_config=read_config)

    total_files = builder.info.splits[args.split].num_examples
    if args.limit_traj is not None:
        ds = ds.take(int(args.limit_traj))
        total_files = min(total_files, int(args.limit_traj))

    total_trajs = 0
    success_trajs = 0
    total_samples = 0
    SUBSAMPLE_DENOM = 15 * 5

    print(f"Scanning {total_files} trajectories from {args.droid_path}...")
    pbar = tqdm(enumerate(ds), total=total_files, unit="traj", desc="Scanning")

    for traj_id, traj_data in pbar:
        total_trajs += 1
        try:
            # Load only necessary data
            traj_batch = next(iter(traj_data["steps"].batch(5000)))

            if traj_batch["reward"][-1].numpy() != 1:
                continue

            success_trajs += 1
            traj_len = int(traj_batch["action"].shape[0])

            if args.full_sequence:
                total_samples += traj_len
            else:
                total_samples += max(1, int(traj_len / SUBSAMPLE_DENOM))

            pbar.set_postfix({"Succ": success_trajs, "Samples": total_samples})

        except Exception as e:
            pbar.write(f"[Warn] Skipping traj {traj_id}: {e}")
            continue

    print("\n" + "="*40)
    print(f"DONE. Split: {args.split} | FullSeq: {args.full_sequence}")
    print(f"Total Trajectories: {total_trajs}")
    print(f"Successful Trajs:   {success_trajs}")
    print(f"Total Samples:      {total_samples}")
    print("="*40)