File size: 520 Bytes
ad9aba4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 | import torch
from wm.dataset.dataset import RoboticsDatasetWrapper
from wm.dataset.data_config import get_config_by_name
import os
def test_franka_load():
dataset = RoboticsDatasetWrapper.get_dataset("franka", seq_len=10)
print(f"Dataset size: {len(dataset)}")
# Load first sample
sample = dataset[0]
print(f"Video shape: {sample['obs'].shape}") # (T, C, H, W)
print(f"Actions shape: {sample['action'].shape}") # (T, action_dim)
if __name__ == "__main__":
test_franka_load()
|