feat: add rollout trajectory image artifacts and swanlab logging

roboimi/demos/vla_scripts/train_vla.py

@@ -299,6 +299,45 @@ def _log_to_swanlab(swanlab_module, payload, step=None):
         log.warning(f"SwanLab log failed at step {step}: {exc}")
 
 
+def _log_rollout_trajectory_images_to_swanlab(
+    swanlab_module,
+    rollout_stats,
+    step=None,
+    context_label: str = 'rollout',
+):
+    if swanlab_module is None or not rollout_stats:
+        return
+
+    image_factory = getattr(swanlab_module, 'Image', None)
+    if image_factory is None:
+        return
+
+    payload = {}
+    for fallback_episode_index, episode in enumerate(rollout_stats.get('episodes', [])):
+        if not isinstance(episode, dict):
+            continue
+        artifact_paths = episode.get('artifact_paths', {})
+        if not isinstance(artifact_paths, dict):
+            continue
+        trajectory_image = artifact_paths.get('trajectory_image')
+        if not trajectory_image:
+            continue
+        episode_index = int(episode.get('episode_index', fallback_episode_index))
+        caption = f'{context_label} trajectory image - episode {episode_index} (front)'
+        try:
+            payload[f'rollout/trajectory_image_episode_{episode_index}'] = image_factory(
+                str(trajectory_image),
+                caption=caption,
+            )
+        except Exception as exc:
+            log.warning(
+                f"SwanLab rollout trajectory image upload prep failed at step {step}: {exc}"
+            )
+
+    if payload:
+        _log_to_swanlab(swanlab_module, payload, step=step)
+
+
 def _finish_swanlab(swanlab_module):
     if swanlab_module is None:
         return
@@ -661,6 +700,13 @@ def _run_training(cfg: DictConfig):
             rollout_cfg.eval.headless = True
             rollout_cfg.eval.device = 'cpu'
             rollout_cfg.eval.verbose_action = False
+            rollout_cfg.eval.record_video = False
+            rollout_cfg.eval.save_trajectory_image = True
+            rollout_cfg.eval.trajectory_image_camera_name = 'front'
+            rollout_cfg.eval.save_summary_json = True
+            rollout_cfg.eval.artifact_dir = str(
+                (run_output_dir / 'rollout_artifacts' / checkpoint_path.stem).resolve()
+            )
 
             log.info(
                 "🎯 开始 checkpoint rollout 验证: %s (episodes=%s, headless=True)",
@@ -867,6 +913,12 @@ def _run_training(cfg: DictConfig):
                         },
                         step=step,
                     )
+                    _log_rollout_trajectory_images_to_swanlab(
+                        swanlab_module,
+                        rollout_stats,
+                        step=step,
+                        context_label=f'epoch {completed_epoch} rollout',
+                    )
                     if rollout_avg_reward > best_rollout_reward:
                         best_rollout_reward = rollout_avg_reward
                         best_model_path = default_best_model_path