Upload folder using huggingface_hub

wm/config/fulltraj_dit/franka.yaml

@@ -0,0 +1,56 @@
+# Configuration for Franka (IsaacLab) World Model Training
+# Full Trajectory Generation (Bidirectional)
+
+# Dynamics Model Class
+dynamics_class: "Bidirectional_FullTrajectory"
+
+# Model identifier for DIT_CLASS_MAP
+model_name: "VideoDiT"
+
+# Configuration passed to the DiT model constructor
+model_config:
+  in_channels: 16            # Latent channels from Wan VAE
+  patch_size: 2
+  dim: 1024                  # Hidden dimension
+  num_layers: 16
+  num_heads: 16
+  action_dim: 7              # Franka action dimension (6-DoF + Gripper)
+  action_compress_rate: 4    # Compresses action sequence to latent sequence
+  max_frames: 33             # Franka sequence length (T=33, 1 + 4*8 windows)
+  action_dropout_prob: 0.1   # CFG for action conditioning
+  temporal_causal: false     # Bidirectional temporal attention for fulltraj
+  vae_name: "WanVAE"
+  vae_config:
+    - "/storage/ice-shared/ae8803che/hxue/data/checkpoint/wan_models/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth"
+  scheduler: "FlowMatch"     # Will be instantiated in dynamics class
+  training_timesteps: 1000
+
+# Dataset Configuration
+dataset:
+  name: "franka"
+  seq_len: 33                # Matches max_frames (e.g., T=33 for 8 steps in latent space)
+  train_test_split: 50       # 50:1 split
+
+# Training Hyperparameters  
+training:
+  batch_size: 4
+  learning_rate: 1e-4
+  num_epochs: 2000
+  grad_clip: 1.0
+  checkpoint_freq: 5000      # Numbered checkpoints for eval
+  latest_freq: 500           # Only updates latest.pt for resuming
+  val_freq: 1000             # Video Logging
+  eval_freq: 500             # MSE Rollout
+  log_freq: 10               # Steps
+  num_workers: 4
+
+# Distributed Training
+distributed:
+  use_ddp: true
+  use_fsdp: false
+
+# WandB Configuration
+wandb:
+  project: "world_model"
+  run_name: "franka_fulltraj_dit_v1"
+  api_key: "62da90010e5c8cc94a66361396c57cea8c2c1e21"

wm/dataset/dataset.py

@@ -100,7 +100,7 @@ class BaseRoboticsDataset(Dataset):
 
     def _get_action_slice(self, entry: Dict[str, Any], start: int, end: int) -> torch.Tensor:
         """Extract raw action slice without padding."""
-        if self.config.name in ["language_table", "rt1", "dreamer4"]:
+        if self.config.name in ["language_table", "rt1", "dreamer4", "pusht", "franka", "lang_table_50k"]:
             return entry['actions'][start:end]
         elif self.config.name == "recon":
             # RECON commands are linear_velocity and angular_velocity

wm/scripts/get_franka_stats.py

@@ -0,0 +1,53 @@
+import torch
+import os
+import matplotlib.pyplot as plt
+import numpy as np
+
+metadata_path = "/storage/ice-shared/ae8803che/hxue/data/dataset/franka/metadata.pt"
+if not os.path.exists(metadata_path):
+    print(f"Error: {metadata_path} not found.")
+    exit(1)
+
+metadata = torch.load(metadata_path)
+num_trajectories = len(metadata)
+
+lengths = []
+action_dims = set()
+
+# Handle both list and dict formats
+if isinstance(metadata, dict):
+    iterator = metadata.values()
+else:
+    iterator = metadata
+
+for info in iterator:
+    if 'num_frames' in info:
+        lengths.append(info['num_frames'])
+    elif 'actions' in info:
+        lengths.append(info['actions'].shape[0])
+    else:
+        print(f"Keys in info: {info.keys()}")
+        break
+    action_dims.add(info['actions'].shape[-1])
+
+avg_len = sum(lengths) / len(lengths)
+median_len = np.median(lengths)
+action_dim = list(action_dims)[0] if len(action_dims) == 1 else str(action_dims)
+
+print(f"Trajectories: {num_trajectories}")
+print(f"Action Dim: {action_dim}")
+print(f"Avg. Video Len: {avg_len:.1f}")
+print(f"Median Video Len: {median_len:.1f}")
+
+# Generate distribution plot
+plt.figure(figsize=(10, 6))
+plt.hist(lengths, bins=30, color='skyblue', edgecolor='black')
+plt.title(f"Franka Video Length Distribution (N={num_trajectories})")
+plt.xlabel("Number of Frames")
+plt.ylabel("Frequency")
+plt.grid(axis='y', alpha=0.75)
+
+save_path = "/storage/ice-shared/ae8803che/hxue/data/world_model/results/stats/franka_dist.png"
+os.makedirs(os.path.dirname(save_path), exist_ok=True)
+plt.savefig(save_path)
+print(f"Distribution plot saved to {save_path}")

wm/test/test_franka_load.py

@@ -0,0 +1,18 @@
+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()