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