JiajunLI/roboimi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: 归一化从agent解耦到训练、推理脚本中

Browse Source
This commit is contained in:
gouhanke
2026-02-06 14:29:36 +08:00
parent a43a2e3d18
commit f4a5c77b7c
3 changed files with 72 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files
roboimi/demos/vla_scripts/train_vla.py
+18 -22
View File
@@ -106,43 +106,36 @@ def main(cfg: DictConfig):
raise
# =========================================================================
# 2.5. Save Dataset Statistics as JSON
# 2.5. Load Dataset Statistics (will be saved into checkpoints)
# =========================================================================
log.info("💾 Saving dataset statistics...")
log.info("💾 Loading dataset statistics...")
dataset_stats = None
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
dataset_stats = {
'normalization_type': cfg.data.get('normalization_type', 'gaussian'),
'action_mean': stats['action']['mean'].tolist(),
'action_std': stats['action']['std'].tolist(),
'action_min': stats['action']['min'].tolist(),
'action_max': stats['action']['max'].tolist(),
'qpos_mean': stats['qpos']['mean'].tolist(),
'qpos_std': stats['qpos']['std'].tolist(),
'qpos_min': stats['qpos']['min'].tolist(),
'qpos_max': stats['qpos']['max'].tolist(),
}
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}")
log.info(f"✅ Dataset statistics loaded (normalization: {dataset_stats['normalization_type']})")
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(f"⚠️ Failed to load statistics: {e}")
log.warning("⚠️ Training will continue, but inference may not work correctly")
# =========================================================================
@@ -234,6 +227,7 @@ def main(cfg: DictConfig):
'model_state_dict': agent.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss.item(),
'dataset_stats': dataset_stats,
}, checkpoint_path)
log.info(f"💾 Checkpoint saved: {checkpoint_path}")
@@ -246,6 +240,7 @@ def main(cfg: DictConfig):
'model_state_dict': agent.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss.item(),
'dataset_stats': dataset_stats,
}, best_model_path)
log.info(f"🌟 Best model updated: {best_model_path} (loss: {best_loss:.4f})")
@@ -258,6 +253,7 @@ def main(cfg: DictConfig):
'model_state_dict': agent.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss.item(),
'dataset_stats': dataset_stats,
}, final_model_path)
log.info(f"💾 Final model saved: {final_model_path}")