Merge remote-tracking branch 'e/main'

Author: ryan-williams

apvd-cli/src/main.rs

@@ -792,7 +792,11 @@ fn run_train(
             };
 
             let train_result = if robust {
-                model.train_robust(max_steps)
+                model.train_robust_with_config(
+                    apvd_core::optimization::robust::OptimConfig::default(),
+                    max_steps,
+                    tiered_config.as_ref(),
+                )
             } else {
                 model.train(learning_rate, max_steps)
             };
@@ -812,36 +816,42 @@ fn run_train(
                 .map(|s| serde_json::to_value(s.v()).unwrap())
                 .collect();
 
-            // Collect BTD (Best To Date) step indices
+            // Collect BTD (Best To Date) step indices (using original indices if filtered)
             let mut btd_steps = Vec::new();
             let mut min_so_far = f64::INFINITY;
-            for (idx, step) in model.steps.iter().enumerate() {
+            for (i, step) in model.steps.iter().enumerate() {
                 let error = step.error.v();
                 if error < min_so_far {
                     min_so_far = error;
-                    btd_steps.push(idx);
+                    let abs_idx = model.step_indices.as_ref()
+                        .map_or(i, |indices| indices[i]);
+                    btd_steps.push(abs_idx);
                 }
             }
 
             // Build history if requested
+            let already_filtered = robust && tiered_config.is_some();
             let history = if include_history || include_checkpoints || tiered_config.is_some() {
                 // Determine which steps to include
-                let include_step: Box<dyn Fn(usize) -> bool> = if include_history {
-                    Box::new(|_| true) // all steps
+                let include_step: Box<dyn Fn(usize) -> bool> = if already_filtered || include_history {
+                    // When robust + tiered, steps are already filtered during training;
+                    // when include_history, output everything.
+                    Box::new(|_| true)
                 } else if include_checkpoints {
                     // Checkpoint steps: 0, 100, 500, 1000, best, final
+                    let total = model.total_steps;
                     let mut indices: std::collections::HashSet<usize> = [0, 100, 500, 1000]
                         .iter()
-                        .filter(|&&i| i < model.steps.len())
+                        .filter(|&&i| i < total)
                         .copied()
                         .collect();
                     indices.insert(model.min_idx); // best step
-                    indices.insert(model.steps.len() - 1); // final step
+                    indices.insert(total - 1); // final step
                     Box::new(move |idx| indices.contains(&idx))
                 } else if let Some(ref tc) = tiered_config {
-                    // Tiered keyframes
+                    // Tiered keyframes (non-robust path)
                     let tc = tc.clone();
-                    let final_idx = model.steps.len() - 1;
+                    let final_idx = model.total_steps - 1;
                     let min_idx = model.min_idx;
                     Box::new(move |idx| {
                         tc.is_keyframe(idx) || idx == final_idx || idx == min_idx
@@ -853,8 +863,16 @@ fn run_train(
                 let mut min_so_far = f64::INFINITY;
                 Some(
                     model.steps.iter().enumerate()
-                        .filter(|(idx, _)| include_step(*idx))
-                        .map(|(idx, step)| {
+                        .filter_map(|(i, step)| {
+                            let abs_idx = model.step_indices.as_ref()
+                                .map_or(i, |indices| indices[i]);
+                            if include_step(abs_idx) {
+                                Some((abs_idx, step))
+                            } else {
+                                None
+                            }
+                        })
+                        .map(|(abs_idx, step)| {
                             let error = step.error.v();
                             let is_best = if error < min_so_far {
                                 min_so_far = error;
@@ -863,7 +881,7 @@ fn run_train(
                                 None
                             };
                             TraceStep {
-                                step_idx: idx,
+                                step_idx: abs_idx,
                                 error,
                                 shapes: step.shapes.iter().map(|s| serde_json::to_value(s.v()).unwrap()).collect(),
                                 is_best,
@@ -880,7 +898,7 @@ fn run_train(
                 final_error: final_step.error.v(),
                 min_error: model.min_error,
                 min_step: model.min_idx,
-                total_steps: model.steps.len(),
+                total_steps: model.total_steps,
                 training_time_ms,
                 final_shapes,
                 history,

apvd-core/src/optimization/model.rs

@@ -3,15 +3,23 @@ use serde::{Deserialize, Serialize};
 use tsify::Tsify;
 
 use crate::{error::SceneError, step::Step, targets::TargetsMap, shape::InputSpec};
+use crate::tiered::TieredConfig;
 use super::adam::{AdamState, AdamConfig};
 use super::robust::{self, OptimConfig};
 
 #[derive(Debug, Clone, Tsify, Serialize, Deserialize)]
 pub struct Model {
     pub steps: Vec<Step>,
+    /// Original step indices when steps are filtered (e.g. tiered keyframes).
+    /// `None` means steps[i] corresponds to step i (no filtering).
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub step_indices: Option<Vec<usize>>,
     pub repeat_idx: Option<usize>,
     pub min_idx: usize,
     pub min_error: f64,
+    /// Actual number of training iterations (may differ from `steps.len()`
+    /// when step filtering is active).
+    pub total_steps: usize,
 }
 
 impl Model {
@@ -20,7 +28,7 @@ impl Model {
         let min_error = step.error.re;
         let steps = vec![step];
         let repeat_idx: Option<usize> = None;
-        Ok(Model { steps, min_idx: 0, repeat_idx, min_error })
+        Ok(Model { steps, step_indices: None, min_idx: 0, repeat_idx, min_error, total_steps: 1 })
     }
     pub fn train(&mut self, max_step_error_ratio: f64, max_steps: usize) -> Result<(), SceneError> {
         let num_steps = self.steps.len();
@@ -66,6 +74,7 @@ impl Model {
             }
             step = nxt;
         }
+        self.total_steps = self.steps.len();
         Ok(())
     }
     pub fn grad_size(&self) -> usize {
@@ -128,6 +137,7 @@ impl Model {
             }
             step = nxt;
         }
+        self.total_steps = self.steps.len();
         Ok(())
     }
 
@@ -139,27 +149,49 @@ impl Model {
     /// - Learning rate warmup for stability
     /// - Step rejection when error increases significantly
     pub fn train_robust(&mut self, max_steps: usize) -> Result<(), SceneError> {
-        self.train_robust_with_config(OptimConfig::default(), max_steps)
+        self.train_robust_with_config(OptimConfig::default(), max_steps, None)
     }
 
     /// Train using robust optimization with custom configuration.
-    pub fn train_robust_with_config(&mut self, config: OptimConfig, max_steps: usize) -> Result<(), SceneError> {
-        let num_steps = self.steps.len();
-        let initial_step = &self.steps[num_steps - 1];
-
-        let new_steps = robust::train_robust(initial_step, config, max_steps)?;
-
-        // Add new steps to history, tracking min error
-        for (idx, step) in new_steps.into_iter().skip(1).enumerate() {
-            let step_idx = num_steps + idx;
-            let err = step.error.re;
+    ///
+    /// When `tiered` is `Some`, only tiered keyframe steps (plus best and final)
+    /// are retained in memory, preventing OOM on long training runs.
+    pub fn train_robust_with_config(
+        &mut self,
+        config: OptimConfig,
+        max_steps: usize,
+        tiered: Option<&TieredConfig>,
+    ) -> Result<(), SceneError> {
+        let step_offset = self.steps.len() - 1;
+        let initial_step = &self.steps[step_offset];
+
+        let retain: Option<Box<dyn Fn(usize) -> bool>> = tiered.map(|tc| {
+            let tc = tc.clone();
+            Box::new(move |idx: usize| tc.is_keyframe(idx)) as Box<dyn Fn(usize) -> bool>
+        });
+
+        let result = robust::train_robust_filtered(
+            initial_step,
+            config,
+            max_steps,
+            step_offset,
+            retain.as_deref(),
+        )?;
+
+        // result.steps[0] is the initial step (already in self.steps), skip it
+        for step in result.steps.into_iter().skip(1) {
+            self.steps.push(step);
+        }
 
-            if err < self.min_error {
-                self.min_idx = step_idx;
-                self.min_error = err;
-            }
+        self.min_idx = result.min_idx;
+        self.min_error = result.min_error;
+        self.total_steps = result.total_steps;
 
-            self.steps.push(step);
+        // Store step indices when filtering is active
+        if tiered.is_some() {
+            let mut indices = vec![0usize]; // initial step
+            indices.extend(result.step_indices.iter().skip(1));
+            self.step_indices = Some(indices);
         }
 
         Ok(())

apvd-core/src/optimization/robust.rs

@@ -181,16 +181,56 @@ impl Optimizer {
     }
 }
 
-/// Train a model using robust optimization.
+/// Result of robust training with optional step filtering.
+pub struct RobustResult {
+    /// Retained steps (filtered by predicate + best + final).
+    pub steps: Vec<Step>,
+    /// Original (absolute) step index for each retained step.
+    pub step_indices: Vec<usize>,
+    /// Total iterations actually completed (accepted steps only).
+    pub total_steps: usize,
+    /// Absolute index of the best (min error) step.
+    pub min_idx: usize,
+    /// Best error achieved.
+    pub min_error: f64,
+}
+
+/// Train using robust optimization, retaining all steps.
 pub fn train_robust(
     initial_step: &Step,
     config: OptimConfig,
     max_steps: usize,
-) -> Result<Vec<Step>, SceneError> {
+) -> Result<RobustResult, SceneError> {
+    train_robust_filtered(initial_step, config, max_steps, 0, None)
+}
+
+/// Train using robust optimization with optional step filtering.
+///
+/// `step_offset`: added to internal indices so `retain` sees absolute step numbers.
+/// `retain`: if `Some`, only steps where `retain(abs_idx)` returns true are kept
+/// in the result (plus the best and final steps, which are always retained).
+pub fn train_robust_filtered(
+    initial_step: &Step,
+    config: OptimConfig,
+    max_steps: usize,
+    step_offset: usize,
+    retain: Option<&dyn Fn(usize) -> bool>,
+) -> Result<RobustResult, SceneError> {
     let grad_size = initial_step.grad_size();
     let mut optimizer = Optimizer::new(grad_size, config);
+
+    // Always retain the initial step
     let mut steps = vec![initial_step.clone()];
+    let mut step_indices = vec![step_offset];
+
     let mut current = initial_step.clone();
+    let mut min_idx = step_offset;
+    let mut min_error = initial_step.error.v();
+    // Track best step separately so we can ensure it's in the output
+    let mut best_step: Option<Step> = None;
+
+    // Count accepted steps (excluding initial)
+    let mut accepted_count: usize = 0;
 
     for step_idx in 0..max_steps {
         let error = current.error.clone();
@@ -235,11 +275,29 @@ pub fn train_robust(
         }
 
         optimizer.accept_step();
+        accepted_count += 1;
+        let abs_idx = step_offset + accepted_count;
 
         debug!("Step {}: error {:.6} -> {:.6} (lr={:.4})",
                step_idx, current_error, new_error, optimizer.effective_lr());
 
-        steps.push(new_step.clone());
+        // Track best
+        if new_error < min_error {
+            min_error = new_error;
+            min_idx = abs_idx;
+            best_step = Some(new_step.clone());
+        }
+
+        // Decide whether to retain this step
+        let should_retain = match retain {
+            Some(f) => f(abs_idx),
+            None => true,
+        };
+        if should_retain {
+            steps.push(new_step.clone());
+            step_indices.push(abs_idx);
+        }
+
         current = new_step;
 
         // Check for convergence (error + penalties very small)
@@ -249,7 +307,35 @@ pub fn train_robust(
         }
     }
 
-    Ok(steps)
+    let total_steps = step_offset + accepted_count + 1; // +1 for initial step
+
+    // Ensure best step is in the output
+    if let Some(best) = best_step {
+        if !step_indices.contains(&min_idx) {
+            steps.push(best);
+            step_indices.push(min_idx);
+        }
+    }
+
+    // Ensure final step is in the output
+    let final_abs_idx = step_offset + accepted_count;
+    if accepted_count > 0 && !step_indices.contains(&final_abs_idx) {
+        steps.push(current);
+        step_indices.push(final_abs_idx);
+    }
+
+    // Sort by step index to maintain order
+    let mut indexed: Vec<_> = step_indices.into_iter().zip(steps.into_iter()).collect();
+    indexed.sort_by_key(|(idx, _)| *idx);
+    let (step_indices, steps): (Vec<_>, Vec<_>) = indexed.into_iter().unzip();
+
+    Ok(RobustResult {
+        steps,
+        step_indices,
+        total_steps,
+        min_idx,
+        min_error,
+    })
 }
 
 #[cfg(test)]