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()