debug(inference): 添加推理阶段qpos归一化

roboimi/demos/vla_scripts/train_vla.py

@@ -1,6 +1,8 @@
 import sys
 import os
 import logging
+import json
+import pickle
 import hydra
 import torch
 from tqdm import tqdm
@@ -103,6 +105,46 @@ def main(cfg: DictConfig):
         log.error(f"❌ Failed to initialize agent: {e}")
         raise
 
+    # =========================================================================
+    # 2.5. Save Dataset Statistics as JSON
+    # =========================================================================
+    log.info("💾 Saving dataset statistics...")
+    try:
+        # Get dataset_dir from config
+        dataset_dir = cfg.data.get('dataset_dir', 'roboimi/demos/dataset/sim_transfer')
+        stats_path = Path(dataset_dir) / 'data_stats.pkl'
+
+        if stats_path.exists():
+            # Load pickle file
+            with open(stats_path, 'rb') as f:
+                stats = pickle.load(f)
+
+            # Extract action statistics
+            action_mean = stats['action']['mean'].tolist() if 'action' in stats else []
+            action_std = stats['action']['std'].tolist() if 'action' in stats else []
+            qpos_mean = stats['qpos']['mean'].tolist() if 'qpos' in stats else []
+            qpos_std = stats['qpos']['std'].tolist() if 'qpos' in stats else []
+
+            # Save as JSON
+            json_stats = {
+                'action_mean': action_mean,
+                'action_std': action_std,
+                'qpos_mean': qpos_mean,
+                'qpos_std': qpos_std
+            }
+            json_path = checkpoint_dir / 'dataset_stats.json'
+            with open(json_path, 'w') as f:
+                json.dump(json_stats, f, indent=2)
+
+            log.info(f"✅ Dataset statistics saved to {json_path}")
+        else:
+            log.warning(f"⚠️  Statistics file not found: {stats_path}")
+            log.warning("⚠️  Actions will not be denormalized during inference!")
+
+    except Exception as e:
+        log.warning(f"⚠️  Failed to save statistics as JSON: {e}")
+        log.warning("⚠️  Training will continue, but inference may not work correctly")
+
     # =========================================================================
     # 3. Setup Optimizer
     # =========================================================================