diff --git a/conditioning.py b/conditioning.py
index 7f1d355c..db626656 100644
--- a/conditioning.py
+++ b/conditioning.py
@@ -11243,10 +11243,13 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner):
                 from value_function import OnlineValueFunction
             with _torch.inference_mode(False), _torch.enable_grad():
                 cond_tensor = serializable_to_tensor(payload).clone()
-                vf = OnlineValueFunction.load_or_create(vf_path, hidden_dim=cond_tensor.shape[-1])
+                with self._v2_stage("  vf: load"):
+                    vf = OnlineValueFunction.load_or_create(vf_path, hidden_dim=cond_tensor.shape[-1])
                 if vf is not None:
-                    vf.train_on(cond_tensor, float(reward))
-                    vf.save(vf_path)
+                    with self._v2_stage("  vf: train"):
+                        vf.train_on(cond_tensor, float(reward))
+                    with self._v2_stage("  vf: save"):
+                        vf.save(vf_path)
                     return vf.n_trained
         except Exception as e:
             print(f"[FunPackRefiner] Value function training failed: {e}")
@@ -11272,11 +11275,15 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner):
                 return None
             vf_path = refinement_state_path(refinement_key, "value_fn_x0", prefix="refine_v2", extension="pt")
             with _torch.inference_mode(False), _torch.enable_grad():
-                snapshot = _torch.load(snap_path, map_location="cpu", weights_only=False)
-                vf = LatentValueFunction.load_or_create(vf_path, hidden_dim=LatentValueFunction.DEFAULT_HIDDEN_DIM)
+                with self._v2_stage("  x0 vf: load snapshot"):
+                    snapshot = _torch.load(snap_path, map_location="cpu", weights_only=False)
+                with self._v2_stage("  x0 vf: load"):
+                    vf = LatentValueFunction.load_or_create(vf_path, hidden_dim=LatentValueFunction.DEFAULT_HIDDEN_DIM)
                 if vf is not None:
-                    vf.train_on(snapshot, float(reward))
-                    vf.save(vf_path)
+                    with self._v2_stage("  x0 vf: train"):
+                        vf.train_on(snapshot, float(reward))
+                    with self._v2_stage("  x0 vf: save"):
+                        vf.save(vf_path)
                     return vf.n_trained
         except Exception as e:
             print(f"[FunPackRefiner] Output value function training failed: {e}")
diff --git a/value_function.py b/value_function.py
index 6ac68c99..ad27d826 100644
--- a/value_function.py
+++ b/value_function.py
@@ -219,6 +219,9 @@ class OnlineValueFunction(nn.Module):
     def is_ready(self):
         return len(self.buffer_c) >= self.MIN_SAMPLES
 
+    def parameter_count(self):
+        return sum(p.numel() for p in self.parameters())
+
     def save(self, path):
         os.makedirs(os.path.dirname(path), exist_ok=True)
         buf = torch.stack(self.buffer_c) if self.buffer_c else torch.zeros(0, self.hidden_dim)